import json import os import posixpath from typing import Any, Callable, Dict, Iterator, List, Optional, Union import pandas as pd import pyarrow import pyarrow as pa import pytest from packaging.version import parse as parse_version from pyarrow.fs import FileType from pytest_lazy_fixtures import lf as lazy_fixture import ray from ray.data._internal.utils.arrow_utils import get_pyarrow_version from ray.data.block import Block from ray.data.dataset import Dataset from ray.data.datasource import FileBasedDatasource, PathPartitionParser from ray.data.datasource.file_based_datasource import _resolve_paths_and_filesystem from ray.data.datasource.partitioning import ( Partitioning, PartitionStyle, PathPartitionFilter, ) from ray.data.expressions import col from ray.data.tests.conftest import * # noqa from ray.tests.conftest import * # noqa class CSVDatasource(FileBasedDatasource): def __init__( self, paths, block_type: Union[pd.DataFrame, pa.Table], **file_based_datasource_kwargs, ): super().__init__(paths, **file_based_datasource_kwargs) self._block_type = block_type def _read_stream(self, f: pa.NativeFile, path: str) -> Iterator[Block]: assert self._block_type in {pd.DataFrame, pa.Table} if self._block_type is pa.Table: from pyarrow import csv yield csv.read_csv(f) if self._block_type is pd.DataFrame: yield pd.read_csv(f) def write_csv(data: Dict[str, List[Any]], path: str) -> None: df = pd.DataFrame(data) os.makedirs(os.path.dirname(path), exist_ok=True) df.to_csv(path, index=False, na_rep="NA") def read_csv( paths: Union[str, List[str]], *, partitioning: Partitioning, block_type: Union[pd.DataFrame, pa.Table], ) -> Dataset: datasource = CSVDatasource(paths, block_type=block_type, partitioning=partitioning) return ray.data.read_datasource(datasource) class PathPartitionEncoder: """Callable that generates directory path strings for path-based partition formats. Path-based partition formats embed all partition keys and values directly in their dataset file paths. Two path partition formats are currently supported - `HIVE` and `DIRECTORY`. For `HIVE` Partitioning, all partition directories will be generated using a `{key1}={value1}/{key2}={value2}` naming convention under the base directory. An accompanying ordered list of partition key field names must also be provided, where the order and length of all partition values must match the order and length of field names For `DIRECTORY` Partitioning, all directories will be generated from partition values using a `{value1}/{value2}` naming convention under the base directory. """ @staticmethod def of( style: PartitionStyle = PartitionStyle.HIVE, base_dir: Optional[str] = None, field_names: Optional[List[str]] = None, filesystem: Optional["pyarrow.fs.FileSystem"] = None, ) -> "PathPartitionEncoder": """Creates a new partition path encoder. Args: style: The partition style - may be either HIVE or DIRECTORY. base_dir: "/"-delimited base directory that all partition paths will be generated under (exclusive). field_names: The partition key field names (i.e. column names for tabular datasets). Required for HIVE partition paths, optional for DIRECTORY partition paths. When non-empty, the order and length of partition key field names must match the order and length of partition values. filesystem: Filesystem that will be used for partition path file I/O. Returns: The new partition path encoder. """ scheme = Partitioning(style, base_dir, field_names, None, filesystem) return PathPartitionEncoder(scheme) def __init__(self, partitioning: Partitioning): """Creates a new partition path encoder. Args: partitioning: The path-based partition scheme. All partition paths will be generated under this scheme's base directory. Field names are required for HIVE partition paths, optional for DIRECTORY partition paths. When non-empty, the order and length of partition key field names must match the order and length of partition values. """ style = partitioning.style field_names = partitioning.field_names if style == PartitionStyle.HIVE and not field_names: raise ValueError( "Hive partition path generation requires a corresponding list of " "partition key field names. Please retry your request with one " "or more field names specified." ) generators = { PartitionStyle.HIVE: self._as_hive_partition_dirs, PartitionStyle.DIRECTORY: self._as_directory_partition_dirs, } self._encoder_fn: Callable[[List[str]], List[str]] = generators.get(style) if self._encoder_fn is None: raise ValueError( f"Unsupported partition style: {style}. " f"Supported styles: {generators.keys()}" ) self._scheme = partitioning def __call__(self, partition_values: List[str]) -> str: """Returns the partition directory path for the given partition value strings. All files for this partition should be written to this directory. If a base directory is set, then the partition directory path returned will be rooted in this base directory. Args: partition_values: The partition value strings to include in the partition path. For HIVE partition paths, the order and length of partition values must match the order and length of partition key field names. Returns: Partition directory path for the given partition values. """ partition_dirs = self._as_partition_dirs(partition_values) return posixpath.join(self._scheme.normalized_base_dir, *partition_dirs) @property def scheme(self) -> Partitioning: """Returns the partitioning for this encoder.""" return self._scheme def _as_hive_partition_dirs(self, values: List[str]) -> List[str]: """Creates HIVE directory names for the given values.""" field_names = self._scheme.field_names return [f"{field_names[i]}={val}" for i, val in enumerate(values)] def _as_directory_partition_dirs(self, values: List[str]) -> List[str]: """Creates DIRECTORY partition directory names for the given values.""" return values def _as_partition_dirs(self, values: List[str]) -> List[str]: """Creates a list of partition directory names for the given values.""" field_names = self._scheme.field_names if field_names: assert len(values) == len(field_names), ( f"Expected {len(field_names)} partition value(s) but found " f"{len(values)}: {values}." ) return self._encoder_fn(values) @pytest.mark.parametrize("block_type", [pd.DataFrame, pa.Table]) class TestReadHivePartitionedFiles: def test_read_single_file(self, tmp_path, block_type, ray_start_regular_shared): path = os.path.join(tmp_path, "year=1970", "country=fr", "data.csv") write_csv({"number": [1, 2, 3]}, path) ds = read_csv(path, partitioning=Partitioning("hive"), block_type=block_type) df = ds.to_pandas() assert list(df.columns) == ["number", "year", "country"] assert list(df["number"]) == [1, 2, 3] assert list(df["year"]) == ["1970", "1970", "1970"] assert list(df["country"]) == ["fr", "fr", "fr"] def test_read_multiple_files(self, tmp_path, block_type, ray_start_regular_shared): path1 = os.path.join(tmp_path, "year=1970", "country=fr", "data.csv") write_csv({"number": [1, 2, 3]}, path1) path2 = os.path.join(tmp_path, "year=1971", "country=ir", "data.csv") write_csv({"number": [4, 5, 6]}, path2) ds = read_csv( [path1, path2], partitioning=Partitioning("hive"), block_type=block_type ) df = ds.to_pandas() assert list(df.columns) == ["number", "year", "country"] assert list(df[df["year"] == "1970"]["number"]) == [1, 2, 3] assert list(df[df["year"] == "1970"]["country"]) == ["fr", "fr", "fr"] assert list(df[df["year"] == "1971"]["number"]) == [4, 5, 6] assert list(df[df["year"] == "1971"]["country"]) == ["ir", "ir", "ir"] @pytest.mark.parametrize( "relative_paths", [ ["year=1970/country=fr/data.csv", "year=1971/language=ir/data.csv"], ["year=1970/country=fr/data.csv", "year=1971/ir/data.csv"], ["year=1970/country=fr/data.csv", "year=1971/data.csv"], ], ) @pytest.mark.skip # TODO: Unskip this test once #28869 is fixed. def test_read_files_with_mismatched_fields( self, relative_paths, tmp_path, block_type, ray_start_regular_shared ): paths = [ os.path.join(tmp_path, relative_path) for relative_path in relative_paths ] for path in paths: write_csv({"number": [0, 0, 0]}, path) with pytest.raises(ValueError): read_csv(paths, partitioning=Partitioning("hive"), block_type=block_type) def test_read_files_with_conflicting_key( self, tmp_path, block_type, ray_start_regular_shared ): path = os.path.join(tmp_path, "month=01", "data.csv") write_csv({"month": [1, 2, 3]}, path) with pytest.raises(ValueError): # `read_csv` should error because `month` is a field in both the CSV and # the path, and the data is different. ds = read_csv( path, partitioning=Partitioning("hive"), block_type=block_type ) ds.schema() @pytest.mark.parametrize("data", [[1, 1, 1], [1, None, 1]]) def test_read_files_with_legally_conflicting_key( self, data, tmp_path, block_type, ray_start_regular_shared ): # `month` is a field in both the path and the CSV, but because the data is # identical, we don't raise an error. path = os.path.join(tmp_path, "month=01", "data.csv") write_csv({"month": data}, path) ds = read_csv(path, partitioning=Partitioning("hive"), block_type=block_type) df = ds.to_pandas() assert list(df.columns) == ["month"] assert list(df["month"]) == [1, 1, 1] @pytest.mark.parametrize("block_type", [pd.DataFrame, pa.Table]) class TestReadUnpartitionedFiles: @pytest.mark.parametrize( "relative_path", ["year=1970/country=fr/data.csv", "1970/fr/data.csv"] ) def test_read_single_file( self, relative_path, tmp_path, block_type, ray_start_regular_shared ): path = os.path.join(tmp_path, relative_path) write_csv({"number": [1, 2, 3]}, path) ds = read_csv(path, partitioning=None, block_type=block_type) # `read_csv` shouldn't include fields like `year` and `country`.` assert list(ds.to_pandas().columns) == ["number"] @pytest.mark.parametrize( "relative_paths", [ ["year=1970/country=fr/data.csv", "year=1971/language=ir/data.csv"], ["year=1970/country=fr/data.csv", "year=1971/ir/data.csv"], ["year=1970/country=fr/data.csv", "year=1971/data.csv"], ["1970/fr/data.csv", "1971/data.csv"], ], ) @pytest.mark.skip # TODO: Unskip this test once #28869 is fixed. def test_read_files_with_mismatched_fields( self, relative_paths, tmp_path, block_type, ray_start_regular_shared ): paths = [ os.path.join(tmp_path, relative_path) for relative_path in relative_paths ] for path in paths: write_csv({"number": [0, 0, 0]}) # `read_csv` shouldn't raise an error if `partitioning` is set to `None`. read_csv(paths, partitioning=None, block_type=block_type) @pytest.mark.parametrize("block_type", [pd.DataFrame, pa.Table]) class TestReadDirPartitionedFiles: def test_read_single_file(self, tmp_path, block_type, ray_start_regular_shared): path = os.path.join(tmp_path, "1970", "fr", "data.csv") write_csv({"number": [1, 2, 3]}, path) ds = read_csv( path, partitioning=Partitioning( "dir", field_names=["year", "country"], base_dir=tmp_path ), block_type=block_type, ) df = ds.to_pandas() assert list(df.columns) == ["number", "year", "country"] assert list(df["number"]) == [1, 2, 3] assert list(df["year"]) == ["1970", "1970", "1970"] assert list(df["country"]) == ["fr", "fr", "fr"] def test_read_single_file_with_null_field( self, tmp_path, block_type, ray_start_regular_shared ): path = os.path.join(tmp_path, "1970", "data", "data.csv") write_csv({"number": [1, 2, 3]}, path) ds = read_csv( path, partitioning=Partitioning( "dir", field_names=["year", None], base_dir=tmp_path ), block_type=block_type, ) df = ds.to_pandas() assert list(df.columns) == ["number", "year"] assert list(df["number"]) == [1, 2, 3] assert list(df["year"]) == ["1970", "1970", "1970"] def test_read_single_file_with_missing_field( self, tmp_path, block_type, ray_start_regular_shared ): path = os.path.join(tmp_path, "1970", "data.csv") write_csv({"number": [0, 0, 0]}, path) # `read_csv` should error because `path` is missing the `country` field. with pytest.raises(ValueError): read_csv( path, partitioning=Partitioning( "dir", field_names=["year", "country"], base_dir=tmp_path ), block_type=block_type, ).schema() @pytest.mark.parametrize( "relative_path", ["1970/data.csv", "1970/us/94704/data.csv"] ) def test_read_single_file_with_invalid_field_names( self, relative_path, tmp_path, block_type, ray_start_regular_shared ): path = os.path.join(tmp_path, relative_path) write_csv({"number": [0, 0, 0]}, path) with pytest.raises(ValueError): read_csv( path, partitioning=Partitioning( "dir", field_names=["year", "country"], base_dir=tmp_path ), block_type=block_type, ).schema() def test_read_files_with_conflicting_key( self, tmp_path, block_type, ray_start_regular_shared ): path = os.path.join(tmp_path, "01", "data.csv") write_csv({"month": [1, 2, 3]}, path) with pytest.raises(ValueError): # `read_csv` should error because `month` is a field in both the CSV and # the path, and the data is different. read_csv( path, partitioning=Partitioning( "dir", field_names=["month"], base_dir=tmp_path ), block_type=block_type, ).schema() @pytest.mark.parametrize("data", [[1, 1, 1], [1, None, 1]]) def test_read_files_with_legally_conflicting_key( self, data, tmp_path, block_type, ray_start_regular_shared ): path = os.path.join(tmp_path, "01", "data.csv") write_csv({"month": data}, path) # `month` is a field in both the path and the CSV, but because the data is # identical, we don't raise an error. ds = read_csv( path, partitioning=Partitioning("dir", field_names=["month"], base_dir=tmp_path), block_type=block_type, ) df = ds.to_pandas() assert list(df.columns) == ["month"] assert list(df["month"]) == [1, 1, 1] def test_read_multiple_files(self, tmp_path, block_type, ray_start_regular_shared): path1 = os.path.join(tmp_path, "1970", "fr", "data.csv") write_csv({"number": [1, 2, 3]}, path1) path2 = os.path.join(tmp_path, "1971", "ir", "data.csv") write_csv({"number": [4, 5, 6]}, path2) ds = read_csv( [path1, path2], partitioning=Partitioning( "dir", field_names=["year", "country"], base_dir=tmp_path ), block_type=block_type, ) df = ds.to_pandas() assert list(df.columns) == ["number", "year", "country"] assert list(df[df["year"] == "1970"]["number"]) == [1, 2, 3] assert list(df[df["year"] == "1970"]["country"]) == ["fr", "fr", "fr"] assert list(df[df["year"] == "1971"]["number"]) == [4, 5, 6] assert list(df[df["year"] == "1971"]["country"]) == ["ir", "ir", "ir"] def _verify_resolved_paths_and_filesystem(scheme: Partitioning): assert scheme.base_dir is not None assert scheme.normalized_base_dir is not None paths, expected_fs = _resolve_paths_and_filesystem( scheme.base_dir, scheme.filesystem, ) path = paths[0] expected_path = f"{path}/" if path and not path.endswith("/") else path assert scheme.normalized_base_dir == expected_path assert isinstance(scheme.resolved_filesystem, type(expected_fs)) def test_partition_style_serde_round_trip(): for style in PartitionStyle: serialized = json.dumps(style) deserialized = PartitionStyle(json.loads(serialized)) assert deserialized == style def test_path_partition_base_properties(): style = PartitionStyle.DIRECTORY base_dir = "/foo/bar" field_names = ["baz", "qux"] scheme = Partitioning(style, base_dir, field_names, None) assert scheme.style == style assert scheme.base_dir == base_dir assert scheme.field_names == field_names _verify_resolved_paths_and_filesystem(scheme) scheme = Partitioning(style, None, field_names, None) assert scheme.style == style assert scheme.base_dir == "" assert scheme.field_names == field_names _verify_resolved_paths_and_filesystem(scheme) def test_path_partition_encoder_errors(): # no field names for default HIVE path partitioning with pytest.raises(ValueError): PathPartitionEncoder.of() # explicit no field names for HIVE path partitioning with pytest.raises(ValueError): PathPartitionEncoder.of(style=PartitionStyle.HIVE, field_names=[]) # invalid path partitioning style with pytest.raises(ValueError): PathPartitionEncoder.of(style=None) # partition field name and field value length mismatch for style in [PartitionStyle.HIVE, PartitionStyle.DIRECTORY]: path_partition_encoder = PathPartitionEncoder.of( style, field_names=["foo", "bar"], ) with pytest.raises(TypeError): path_partition_encoder(None) with pytest.raises(AssertionError): path_partition_encoder([]) with pytest.raises(AssertionError): path_partition_encoder(["1"]) with pytest.raises(AssertionError): path_partition_encoder(["1", "2", "3"]) @pytest.mark.parametrize( "fs,base_dir", [ (None, None), (lazy_fixture("local_fs"), lazy_fixture("local_path")), (lazy_fixture("s3_fs"), lazy_fixture("s3_path")), ( lazy_fixture("s3_fs_with_special_chars"), lazy_fixture("s3_path_with_special_chars"), ), ], ) def test_path_partition_encoder_hive(fs, base_dir): field_names = ["foo", "bar"] path_partition_encoder = PathPartitionEncoder.of( field_names=field_names, base_dir=base_dir, filesystem=fs, ) _verify_resolved_paths_and_filesystem(path_partition_encoder.scheme) partition_values = ["1", "2"] partition_path = path_partition_encoder(partition_values) assert partition_path == posixpath.join( path_partition_encoder.scheme.normalized_base_dir, "foo=1", "bar=2", ) if fs is not None: file_info = fs.get_file_info(partition_path) assert file_info.type == FileType.NotFound fs.create_dir(partition_path) file_info = fs.get_file_info(partition_path) assert file_info.type == FileType.Directory @pytest.mark.parametrize( "fs,base_dir", [ (None, None), (lazy_fixture("local_fs"), lazy_fixture("local_path")), (lazy_fixture("s3_fs"), lazy_fixture("s3_path")), ( lazy_fixture("s3_fs_with_special_chars"), lazy_fixture("s3_path_with_special_chars"), ), ], ) def test_path_partition_encoder_directory(fs, base_dir): path_partition_encoder = PathPartitionEncoder.of( style=PartitionStyle.DIRECTORY, field_names=["foo", "bar"], base_dir=base_dir, filesystem=fs, ) _verify_resolved_paths_and_filesystem(path_partition_encoder.scheme) partition_values = ["1", "2"] partition_path = path_partition_encoder(partition_values) assert partition_path == posixpath.join( path_partition_encoder.scheme.normalized_base_dir, *partition_values, ) if fs is not None: file_info = fs.get_file_info(partition_path) assert file_info.type == FileType.NotFound fs.create_dir(partition_path) file_info = fs.get_file_info(partition_path) assert file_info.type == FileType.Directory path_partition_encoder = PathPartitionEncoder.of( style=PartitionStyle.DIRECTORY, base_dir=base_dir, filesystem=fs, ) partition_path = path_partition_encoder([]) assert partition_path == path_partition_encoder.scheme.normalized_base_dir partition_path = path_partition_encoder(partition_values) assert partition_path == posixpath.join( path_partition_encoder.scheme.normalized_base_dir, *partition_values, ) def test_path_partition_parser_errors(): # no field names for DIRECTORY path partitioning with pytest.raises(ValueError): PathPartitionParser.of(style=PartitionStyle.DIRECTORY) # explicit no field names for DIRECTORY path partitioning with pytest.raises(ValueError): PathPartitionParser.of(style=PartitionStyle.DIRECTORY, field_names=[]) # invalid path partitioning style with pytest.raises(ValueError): PathPartitionParser.of(style=None) # HIVE partition field name and field value length or order mismatch path_partition_parser = PathPartitionParser.of( style=PartitionStyle.HIVE, field_names=["foo", "bar"], ) with pytest.raises(ValueError): path_partition_parser("foo=1/") with pytest.raises(ValueError): path_partition_parser("bar=1/foo=2/") with pytest.raises(ValueError): path_partition_parser("foo=1/bar=2/qux=3/") # ensure HIVE partition base directory is not considered a partition path_partition_parser = PathPartitionParser.of( style=PartitionStyle.HIVE, base_dir="foo=1", field_names=["foo", "bar"], ) with pytest.raises(ValueError): path_partition_parser("foo=1/bar=2/") # DIRECTORY partition field name and field value length mismatch path_partition_parser = PathPartitionParser.of( style=PartitionStyle.DIRECTORY, field_names=["foo", "bar"], ) with pytest.raises(ValueError): path_partition_parser("1/") with pytest.raises(ValueError): path_partition_parser("1/2/3/") # ensure DIRECTORY partition base directory is not considered a partition path_partition_parser = PathPartitionParser.of( style=PartitionStyle.DIRECTORY, base_dir="1", field_names=["foo", "bar"], ) with pytest.raises(ValueError): path_partition_parser("1/2/") @pytest.mark.parametrize( "fs,base_dir", [ (None, None), (lazy_fixture("local_fs"), lazy_fixture("local_path")), (lazy_fixture("s3_fs"), lazy_fixture("s3_path")), ( lazy_fixture("s3_fs_with_special_chars"), lazy_fixture("s3_path_with_special_chars"), ), ], ) def test_path_partition_parser_hive(fs, base_dir): partition_parser = PathPartitionParser.of(base_dir=base_dir, filesystem=fs) _verify_resolved_paths_and_filesystem(partition_parser.scheme) base_dir = partition_parser.scheme.normalized_base_dir # parse unpartitioned paths... partition_kvs = partition_parser("") assert partition_kvs == {} unpartitioned_paths = [ "", "foo/1", "bar/2", "baz/3", posixpath.join(base_dir, "test.txt"), posixpath.join(base_dir, "foo/test.txt"), posixpath.join(base_dir, "foo/bar/qux=3"), posixpath.join(base_dir, "test=1.txt"), ] for path in unpartitioned_paths: assert partition_parser(path) == {} partitioned_path = posixpath.join(base_dir, "foo=1/test.txt") assert partition_parser(partitioned_path) == {"foo": "1"} partitioned_path = posixpath.join(base_dir, " foo = 1 /test.txt") assert partition_parser(partitioned_path) == {" foo ": " 1 "} partitioned_path = posixpath.join(base_dir, "foo/bar=2/test.txt") assert partition_parser(partitioned_path) == {"bar": "2"} partitioned_path = posixpath.join(base_dir, "bar=2/foo=1/test") assert partition_parser(partitioned_path) == {"foo": "1", "bar": "2"} partitioned_path = posixpath.join(base_dir, "foo/bar/qux=3/") assert partition_parser(partitioned_path) == {"qux": "3"} @pytest.mark.parametrize( "path, expected_partitions", [ # '%2F' should decode to '/' ("bucket/key=partition%2Fvalue/file.txt", {"key": "partition/value"}), # '+' must remain literal when decoding path components. See # https://github.com/ray-project/ray/pull/57625#discussion_r2441360523. ("bucket/key=foo+bar/file.txt", {"key": "foo+bar"}), # '%2B' should decode to '+' ("bucket/key=foo%2Bbar/file.txt", {"key": "foo+bar"}), ], ) def test_path_partition_parser_decodes_special_characters( path: str, expected_partitions: Dict[str, str] ): partition_parser = PathPartitionParser.of(base_dir="bucket") assert partition_parser(path) == expected_partitions @pytest.mark.parametrize( "fs,base_dir", [ (None, None), (lazy_fixture("local_fs"), lazy_fixture("local_path")), (lazy_fixture("s3_fs"), lazy_fixture("s3_path")), ( lazy_fixture("s3_fs_with_special_chars"), lazy_fixture("s3_path_with_special_chars"), ), ], ) def test_path_partition_parser_dir(fs, base_dir): partition_parser = PathPartitionParser.of( PartitionStyle.DIRECTORY, base_dir=base_dir, field_names=["foo", "bar"], filesystem=fs, ) _verify_resolved_paths_and_filesystem(partition_parser.scheme) base_dir = partition_parser.scheme.normalized_base_dir # parse unpartitioned paths... partition_kvs = partition_parser("") assert partition_kvs == {} if base_dir: unpartitioned_paths = [ "", "foo/1", "bar/2", "baz/3", posixpath.join(base_dir, "test.txt"), ] for path in unpartitioned_paths: assert partition_parser(path) == {} partitioned_path = posixpath.join(base_dir, "1/2/test.txt") assert partition_parser(partitioned_path) == {"foo": "1", "bar": "2"} partitioned_path = posixpath.join(base_dir, " 1 / t w o /test.txt") assert partition_parser(partitioned_path) == {"foo": " 1 ", "bar": " t w o "} partitioned_path = posixpath.join(base_dir, "2/1/test.txt") assert partition_parser(partitioned_path) == {"foo": "2", "bar": "1"} partitioned_path = posixpath.join(base_dir, "1/2/") assert partition_parser(partitioned_path) == {"foo": "1", "bar": "2"} partitioned_path = posixpath.join(base_dir, "1/2/3") assert partition_parser(partitioned_path) == {"foo": "1", "bar": "2"} partition_parser = PathPartitionParser.of( PartitionStyle.DIRECTORY, base_dir=base_dir, field_names=["bar", "foo"], filesystem=fs, ) partitioned_path = posixpath.join(base_dir, "1/2/test") assert partition_parser(partitioned_path) == {"bar": "1", "foo": "2"} partitioned_path = posixpath.join(base_dir, "2/1/test") assert partition_parser(partitioned_path) == {"bar": "2", "foo": "1"} partition_parser = PathPartitionParser.of( PartitionStyle.DIRECTORY, base_dir=base_dir, field_names=["year", None, "country"], filesystem=fs, ) partitioned_path = posixpath.join(base_dir, "1970/countries/fr/products.csv") assert partition_parser(partitioned_path) == {"year": "1970", "country": "fr"} @pytest.mark.parametrize( "fs,base_dir", [ (None, None), (lazy_fixture("local_fs"), lazy_fixture("local_path")), (lazy_fixture("s3_fs"), lazy_fixture("s3_path")), ( lazy_fixture("s3_fs_with_special_chars"), lazy_fixture("s3_path_with_special_chars"), ), ], ) def test_path_partition_filter_hive(fs, base_dir): pass_through = PathPartitionFilter.of(None, base_dir=base_dir, filesystem=fs) _verify_resolved_paths_and_filesystem(pass_through.parser.scheme) paths = pass_through([]) assert paths == [] paths = pass_through(["foo/1", "bar/2", "baz/3"]) assert paths == ["foo/1", "bar/2", "baz/3"] filter_unpartitioned = PathPartitionFilter.of( base_dir=base_dir, filesystem=fs, filter_fn=lambda d: bool(d), ) _verify_resolved_paths_and_filesystem(filter_unpartitioned.parser.scheme) base_dir = filter_unpartitioned.parser.scheme.normalized_base_dir test_paths = [ posixpath.join(base_dir, "test.txt"), posixpath.join(base_dir, "foo/test.txt"), posixpath.join(base_dir, "foo=1/test.txt"), posixpath.join(base_dir, "foo/bar=2/test.txt"), posixpath.join(base_dir, "foo=1/bar=2/test"), posixpath.join(base_dir, "foo/bar/qux=3/"), posixpath.join(base_dir, "foo/bar/qux=3"), posixpath.join(base_dir, "test=1.txt"), ] if base_dir: test_paths.extend(["test.txt", "foo=1/test.txt"]) paths = filter_unpartitioned(test_paths) assert paths == [ posixpath.join(base_dir, "foo=1/test.txt"), posixpath.join(base_dir, "foo/bar=2/test.txt"), posixpath.join(base_dir, "foo=1/bar=2/test"), posixpath.join(base_dir, "foo/bar/qux=3/"), ] filter_values = PathPartitionFilter.of( base_dir=base_dir, filesystem=fs, filter_fn=lambda d: d and (d.get("qux") == "3" or (d.get("foo") == "1" and d.get("bar") == "2")), ) _verify_resolved_paths_and_filesystem(filter_values.parser.scheme) paths = filter_values(test_paths) assert paths == [ posixpath.join(base_dir, "foo=1/bar=2/test"), posixpath.join(base_dir, "foo/bar/qux=3/"), ] filter_field_name_values = PathPartitionFilter.of( base_dir=base_dir, field_names=["foo", "bar"], filesystem=fs, filter_fn=lambda d: d and d.get("foo") == "1" and d.get("bar") == "2", ) test_paths = [ posixpath.join(base_dir, "foo=1/bar=2/test"), posixpath.join(base_dir, "prefix/foo=1/padding/bar=2/test"), ] paths = filter_field_name_values(test_paths) assert paths == test_paths @pytest.mark.parametrize( "fs,base_dir", [ (None, None), (lazy_fixture("local_fs"), lazy_fixture("local_path")), (lazy_fixture("s3_fs"), lazy_fixture("s3_path")), ( lazy_fixture("s3_fs_with_special_chars"), lazy_fixture("s3_path_with_special_chars"), ), ], ) def test_path_partition_filter_directory(fs, base_dir): pass_through = PathPartitionFilter.of( None, style=PartitionStyle.DIRECTORY, base_dir=base_dir, field_names=["foo", "bar"], filesystem=fs, ) paths = pass_through([]) assert paths == [] paths = pass_through(["foo/1", "bar/2", "baz/3"]) assert paths == ["foo/1", "bar/2", "baz/3"] filter_unpartitioned = PathPartitionFilter.of( style=PartitionStyle.DIRECTORY, base_dir=base_dir, field_names=["foo", "bar"], filesystem=fs, filter_fn=lambda d: bool(d), ) _verify_resolved_paths_and_filesystem(filter_unpartitioned.parser.scheme) base_dir = filter_unpartitioned.parser.scheme.normalized_base_dir test_paths = [ posixpath.join(base_dir, "test.txt"), posixpath.join(base_dir, "1/2/test.txt"), posixpath.join(base_dir, "1/2/"), posixpath.join(base_dir, "2/1/"), posixpath.join(base_dir, "1/2/3"), ] if base_dir: # files outside of the base directory are implicitly unpartitioned test_paths.extend(["test.txt", "1/2/test.txt"]) paths = filter_unpartitioned(test_paths) assert paths == [ posixpath.join(base_dir, "1/2/test.txt"), posixpath.join(base_dir, "1/2/"), posixpath.join(base_dir, "2/1/"), posixpath.join(base_dir, "1/2/3"), ] filter_values = PathPartitionFilter.of( style=PartitionStyle.DIRECTORY, base_dir=base_dir, field_names=["foo", "bar"], filesystem=fs, filter_fn=lambda d: d and d["foo"] == "1" and d["bar"] == "2", ) _verify_resolved_paths_and_filesystem(filter_values.parser.scheme) paths = filter_values(test_paths) assert paths == [ posixpath.join(base_dir, "1/2/test.txt"), posixpath.join(base_dir, "1/2/"), posixpath.join(base_dir, "1/2/3"), ] @pytest.mark.parametrize( "partition_value,expected_type", [ ("1", int), ("1.0", float), ("spam", str), ("true", bool), ], ) def test_field_types(partition_value, expected_type): partitioning = Partitioning(style="hive", field_types={"key": expected_type}) parse = PathPartitionParser(partitioning) partitions = parse(f"key={partition_value}/data.parquet") assert set(partitions.keys()) == {"key"} assert isinstance(partitions["key"], expected_type) @pytest.mark.parametrize( "path,predicate,expected_result,description", [ # Simple equality matches ("country=US/file.parquet", col("country") == "US", True, "Exact match"), ("country=US/file.parquet", col("country") == "UK", False, "No match"), # AND predicates ( "country=US/year=2020/file.parquet", (col("country") == "US") & (col("year") == "2020"), True, "AND both match", ), ( "country=US/year=2020/file.parquet", (col("country") == "US") & (col("year") == "2021"), False, "AND one doesn't match", ), # OR predicates ( "country=US/file.parquet", (col("country") == "US") | (col("country") == "UK"), True, "OR first matches", ), ( "country=FR/file.parquet", (col("country") == "US") | (col("country") == "UK"), False, "OR neither matches", ), # Comparison operators ("year=2020/file.parquet", col("year") > "2019", True, "Greater than"), ("year=2020/file.parquet", col("year") < "2019", False, "Less than"), # NOT operator ("country=US/file.parquet", ~(col("country") == "UK"), True, "NOT false"), ("country=US/file.parquet", ~(col("country") == "US"), False, "NOT true"), # IS_IN operator ( "country=US/file.parquet", col("country").is_in(["US", "UK"]), True, "IS_IN matches", ), ( "country=FR/file.parquet", col("country").is_in(["US", "UK"]), False, "IS_IN no match", ), ( "year=2020/file.parquet", col("year").is_in(["2019", "2020", "2021"]), True, "IS_IN with multiple values", ), ], ) @pytest.mark.skipif( get_pyarrow_version() < parse_version("14.0.0"), reason="Partition predicate evaluation requires pyarrow >= 14.0.0", ) def test_evaluate_predicate_on_partition(path, predicate, expected_result, description): """Test partition predicate evaluation for automatic partition pruning.""" parser = PathPartitionParser(Partitioning("hive")) result = parser.evaluate_predicate_on_partition(path, predicate) assert ( result == expected_result ), f"{description}: Expected {expected_result}, got {result}" def test_evaluate_predicate_on_unpartitioned_file(): """Test that unpartitioned files are conservatively included.""" parser = PathPartitionParser(Partitioning("hive")) # Unpartitioned file should return False when filtering on partition columns # (we can't determine if it matches the predicate without partition values) result = parser.evaluate_predicate_on_partition( "data.parquet", col("country") == "US" ) assert result is False if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", __file__]))