chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,999 @@
|
||||
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__]))
|
||||
Reference in New Issue
Block a user