| 858 | ], |
| 859 | ) |
| 860 | def test_dataset_partitioning_format( |
| 861 | flavor: str, |
| 862 | expected_defined_partition: tuple, |
| 863 | expected_undefined_partition: tuple, |
| 864 | ): |
| 865 | |
| 866 | partitioning_schema = pa.schema([("foo", pa.string()), ("bar", pa.string())]) |
| 867 | |
| 868 | partitioning = getattr(ds, flavor)(schema=partitioning_schema) |
| 869 | |
| 870 | # test forward transformation (format) |
| 871 | assert ( |
| 872 | partitioning.format((pc.field("bar") == "ant bee") & (pc.field("foo") == "A")) |
| 873 | == expected_defined_partition |
| 874 | ) |
| 875 | |
| 876 | # test backward transformation (parse) |
| 877 | assert partitioning.parse("/".join(expected_defined_partition)).equals( |
| 878 | (pc.field("foo") == "A") & (pc.field("bar") == "ant bee") |
| 879 | ) |
| 880 | |
| 881 | # test complex expression can still be parsed into useful directory/path |
| 882 | assert ( |
| 883 | partitioning.format( |
| 884 | ((pc.field("bar") == "ant bee") & (pc.field("foo") == "A")) |
| 885 | & ((pc.field("bar") == "ant bee") & (pc.field("foo") == "A")) |
| 886 | ) |
| 887 | == expected_defined_partition |
| 888 | ) |
| 889 | |
| 890 | # test a different complex expression cannot be parsed into directory/path |
| 891 | # and just returns the same value as if no filter were applied. |
| 892 | assert ( |
| 893 | partitioning.format( |
| 894 | ((pc.field("bar") == "ant bee") & (pc.field("foo") == "A")) |
| 895 | | ((pc.field("bar") == "ant bee") & (pc.field("foo") == "A")) |
| 896 | ) |
| 897 | == expected_undefined_partition |
| 898 | ) |
| 899 | |
| 900 | if flavor != "HivePartitioning": |
| 901 | # Raises error upon filtering for lower level partition without filtering for |
| 902 | # higher level partition |
| 903 | with pytest.raises( |
| 904 | pa.ArrowInvalid, |
| 905 | match=( |
| 906 | "No partition key for foo but a key was provided" |
| 907 | " subsequently for bar" |
| 908 | ) |
| 909 | ): |
| 910 | partitioning.format(((pc.field("bar") == "ant bee"))) |
| 911 | else: |
| 912 | # Hive partitioning allows this to pass |
| 913 | assert partitioning.format(((pc.field("bar") == "ant bee"))) == ( |
| 914 | r"bar=ant%20bee", |
| 915 | "", |
| 916 | ) |
| 917 | |