(
tempdir, partitioning, null_fallback, infer_dictionary, partition_keys
)
| 2834 | ([None, 2, 3], [None, 2, 3]), |
| 2835 | ]) |
| 2836 | def test_partition_discovery( |
| 2837 | tempdir, partitioning, null_fallback, infer_dictionary, partition_keys |
| 2838 | ): |
| 2839 | # ARROW-9288 / ARROW-9476 |
| 2840 | table = pa.table({'a': range(9), 'b': [0.0] * 4 + [1.0] * 5}) |
| 2841 | |
| 2842 | has_null = None in partition_keys[0] or None in partition_keys[1] |
| 2843 | if partitioning == "directory" and has_null: |
| 2844 | # Directory partitioning can't handle the first part being null |
| 2845 | return |
| 2846 | |
| 2847 | if partitioning == "directory": |
| 2848 | partitioning = ds.DirectoryPartitioning.discover( |
| 2849 | ["part1", "part2"], infer_dictionary=infer_dictionary) |
| 2850 | fmt = "{0}/{1}" |
| 2851 | null_value = None |
| 2852 | else: |
| 2853 | if null_fallback: |
| 2854 | partitioning = ds.HivePartitioning.discover( |
| 2855 | infer_dictionary=infer_dictionary, null_fallback=null_fallback |
| 2856 | ) |
| 2857 | else: |
| 2858 | partitioning = ds.HivePartitioning.discover( |
| 2859 | infer_dictionary=infer_dictionary) |
| 2860 | fmt = "part1={0}/part2={1}" |
| 2861 | if null_fallback: |
| 2862 | null_value = null_fallback |
| 2863 | else: |
| 2864 | null_value = "__HIVE_DEFAULT_PARTITION__" |
| 2865 | |
| 2866 | basepath = tempdir / "dataset" |
| 2867 | basepath.mkdir() |
| 2868 | |
| 2869 | part_keys1, part_keys2 = partition_keys |
| 2870 | for part1 in part_keys1: |
| 2871 | for part2 in part_keys2: |
| 2872 | path = basepath / \ |
| 2873 | fmt.format(part1 or null_value, part2 or null_value) |
| 2874 | path.mkdir(parents=True) |
| 2875 | pq.write_table(table, path / "test.parquet") |
| 2876 | |
| 2877 | dataset = ds.dataset(str(basepath), partitioning=partitioning) |
| 2878 | |
| 2879 | def expected_type(key): |
| 2880 | if infer_dictionary: |
| 2881 | value_type = pa.string() if isinstance(key, str) else pa.int32() |
| 2882 | return pa.dictionary(pa.int32(), value_type) |
| 2883 | else: |
| 2884 | return pa.string() if isinstance(key, str) else pa.int32() |
| 2885 | expected_schema = table.schema.append( |
| 2886 | pa.field("part1", expected_type(part_keys1[0])) |
| 2887 | ).append( |
| 2888 | pa.field("part2", expected_type(part_keys2[0])) |
| 2889 | ) |
| 2890 | assert dataset.schema.equals(expected_schema) |
| 2891 | |
| 2892 | |
| 2893 | @pytest.mark.pandas |
nothing calls this directly
no test coverage detected