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