wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/data/datasource/__init__.py
+78
View File
@@ -0,0 +1,78 @@
from ray.data._internal.datasource.delta_sharing_datasource import (
DeltaSharingDatasource,
)
from ray.data._internal.datasource.mcap_datasource import (
MCAPDatasource,
TimeRange,
)
from ray.data._internal.datasource.sql_datasource import Connection
from ray.data._internal.datasource.turbopuffer_datasink import (
TurbopufferDatasink,
)
from ray.data._internal.savemode import SaveMode
from ray.data.datasource.datasink import (
Datasink,
DummyOutputDatasink,
WriteResult,
WriteReturnType,
)
from ray.data.datasource.datasource import (
Datasource,
RandomIntRowDatasource,
Reader,
ReadTask,
)
from ray.data.datasource.file_based_datasource import (
FileBasedDatasource,
FileShuffleConfig,
_S3FileSystemWrapper,
)
from ray.data.datasource.file_datasink import (
BlockBasedFileDatasink,
RowBasedFileDatasink,
)
from ray.data.datasource.file_meta_provider import (
BaseFileMetadataProvider,
DefaultFileMetadataProvider,
FileMetadataProvider,
)
from ray.data.datasource.filename_provider import FilenameProvider
from ray.data.datasource.partitioning import (
Partitioning,
PartitionStyle,
PathPartitionFilter,
PathPartitionParser,
)
# Note: HuggingFaceDatasource should NOT be imported here, because
# we want to only import the Hugging Face datasets library when we use
# ray.data.from_huggingface() or HuggingFaceDatasource() directly.
__all__ = [
"BaseFileMetadataProvider",
"Connection",
"Datasink",
"Datasource",
"DefaultFileMetadataProvider",
"DeltaSharingDatasource",
"DummyOutputDatasink",
"FileBasedDatasource",
"FileShuffleConfig",
"FileMetadataProvider",
"FilenameProvider",
"MCAPDatasource",
"PartitionStyle",
"PathPartitionFilter",
"PathPartitionParser",
"Partitioning",
"RandomIntRowDatasource",
"ReadTask",
"Reader",
"RowBasedFileDatasink",
"TurbopufferDatasink",
"BlockBasedFileDatasink",
"_S3FileSystemWrapper",
"TimeRange",
"WriteResult",
"WriteReturnType",
"SaveMode",
]
python/ray/data/datasource/datasink.py
+206
View File
@@ -0,0 +1,206 @@
import itertools
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING, Generic, Iterable, List, Optional, TypeVar
import ray
from ray.data._internal.execution.interfaces import TaskContext
from ray.data.block import Block, BlockAccessor
from ray.util.annotations import DeveloperAPI
if TYPE_CHECKING:
import pyarrow as pa
logger = logging.getLogger(__name__)
WriteReturnType = TypeVar("WriteReturnType")
"""Generic type for the return value of `Datasink.write`."""
@dataclass
@DeveloperAPI
class WriteResult(Generic[WriteReturnType]):
"""Aggregated result of the Datasink write operations."""
# Total number of written rows.
num_rows: int
# Total size in bytes of written data.
size_bytes: int
# All returned values of `Datasink.write`.
write_returns: List[WriteReturnType]
@classmethod
def combine(cls, *wrs: "WriteResult") -> "WriteResult":
num_rows = sum(wr.num_rows for wr in wrs)
size_bytes = sum(wr.size_bytes for wr in wrs)
write_returns = list(itertools.chain(*[wr.write_returns for wr in wrs]))
return WriteResult(
num_rows=num_rows,
size_bytes=size_bytes,
write_returns=write_returns,
)
@DeveloperAPI
class Datasink(Generic[WriteReturnType]):
"""Interface for defining write-related logic.
If you want to write data to something that isn't built-in, subclass this class
and call :meth:`~ray.data.Dataset.write_datasink`.
"""
def on_write_start(self, schema: Optional["pa.Schema"] = None) -> None:
"""Callback for when a write job starts.
Use this method to perform setup for write tasks. For example, creating a
staging bucket in S3.
This is called on the driver when the first input bundle is ready, just
before write tasks are submitted. The schema is extracted from the first
input bundle, enabling schema-dependent initialization.
Args:
schema: The PyArrow schema of the data being written. This is
automatically extracted from the first input bundle. May be None
if the input data has no schema.
"""
pass
def write(
self,
blocks: Iterable[Block],
ctx: TaskContext,
) -> WriteReturnType:
"""Write blocks. This is used by a single write task.
Args:
blocks: Generator of data blocks.
ctx: ``TaskContext`` for the write task.
Returns:
Result of this write task. When the entire write operator finishes,
All returned values will be passed as `WriteResult.write_returns`
to `Datasink.on_write_complete`.
"""
raise NotImplementedError
def on_write_complete(self, write_result: WriteResult[WriteReturnType]):
"""Callback for when a write job completes.
This can be used to `commit` a write output. This method must
succeed prior to ``write_datasink()`` returning to the user. If this
method fails, then ``on_write_failed()`` is called.
Args:
write_result: Aggregated result of the
Write operator, containing write results and stats.
"""
pass
def on_write_failed(self, error: Exception) -> None:
"""Callback for when a write job fails.
This is called on a best-effort basis on write failures.
Args:
error: The first error encountered.
"""
pass
def get_name(self) -> str:
"""Return a human-readable name for this datasink.
This is used as the names of the write tasks.
"""
name = type(self).__name__
datasink_suffix = "Datasink"
if name.startswith("_"):
name = name[1:]
if name.endswith(datasink_suffix):
name = name[: -len(datasink_suffix)]
return name
@property
def supports_distributed_writes(self) -> bool:
"""If ``False``, only launch write tasks on the driver's node."""
return True
@property
def min_rows_per_write(self) -> Optional[int]:
"""The target number of rows to pass to each :meth:`~ray.data.Datasink.write` call.
If ``None``, Ray Data passes a system-chosen number of rows.
"""
return None
@DeveloperAPI
class DummyOutputDatasink(Datasink[None]):
"""An example implementation of a writable datasource for testing.
Examples:
>>> import ray
>>> from ray.data.datasource import DummyOutputDatasink
>>> output = DummyOutputDatasink()
>>> ray.data.range(10).write_datasink(output)
>>> assert output.num_ok == 1
"""
def __init__(self):
ctx = ray.data.DataContext.get_current()
# Setup a dummy actor to send the data. In a real datasource, write
# tasks would send data to an external system instead of a Ray actor.
@ray.remote(scheduling_strategy=ctx.scheduling_strategy)
class DataSink:
def __init__(self):
self.rows_written = 0
self.enabled = True
def write(self, block: Block) -> None:
block = BlockAccessor.for_block(block)
self.rows_written += block.num_rows()
def get_rows_written(self):
return self.rows_written
self.data_sink = DataSink.remote()
self.num_ok = 0
self.num_failed = 0
self.enabled = True
def write(
self,
blocks: Iterable[Block],
ctx: TaskContext,
) -> None:
tasks = []
if not self.enabled:
raise ValueError("disabled")
for b in blocks:
tasks.append(self.data_sink.write.remote(b))
ray.get(tasks)
def on_write_complete(self, write_result: WriteResult[None]):
self.num_ok += 1
def on_write_failed(self, error: Exception) -> None:
self.num_failed += 1
def _gen_datasink_write_result(
write_result_blocks: List[Block],
) -> WriteResult:
import pandas as pd
assert all(
isinstance(block, pd.DataFrame) and len(block) == 1
for block in write_result_blocks
)
total_num_rows = sum(result["num_rows"].sum() for result in write_result_blocks)
total_size_bytes = sum(result["size_bytes"].sum() for result in write_result_blocks)
write_returns = [result["write_return"][0] for result in write_result_blocks]
return WriteResult(total_num_rows, total_size_bytes, write_returns)
python/ray/data/datasource/datasource.py
+479
View File
@@ -0,0 +1,479 @@
import copy
from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Optional
import numpy as np
from ray.data._internal.util import _check_pyarrow_version
from ray.data.block import Block, BlockMetadata, Schema
from ray.data.datasource.util import _iter_sliced_blocks
from ray.data.expressions import Expr
from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI
if TYPE_CHECKING:
from ray.data.context import DataContext
class _DatasourceProjectionPushdownMixin:
"""Mixin for reading operators supporting projection pushdown.
The read stage only prunes columns; it never renames. Column renaming
is always carried by an ``AliasExpr`` in a ``Project`` operator above
the read. As a consequence, projection maps stored here are always
identity (``{name: name}``).
"""
def supports_projection_pushdown(self) -> bool:
"""Returns ``True`` in case ``Datasource`` supports projection operation
being pushed down into the reading layer"""
return False
def get_projection_map(self) -> Optional[Dict[str, str]]:
"""Return the projection map (always an identity mapping).
Returns:
Dict mapping selected column names to themselves. ``None``
means all columns are selected. Empty dict ``{}`` means no
columns are selected.
"""
return self._projection_map
def _get_data_columns(self) -> Optional[List[str]]:
"""Extract data columns from projection map.
Helper method for datasources that need to pass columns to legacy
read functions expecting a list of columns.
Returns:
List of column names, or None if all columns should be read.
Empty list [] means no columns.
"""
return (
list(self._projection_map.keys())
if self._projection_map is not None
else None
)
@staticmethod
def _combine_projection_map(
prev_projection_map: Optional[Dict[str, str]],
new_projection_map: Optional[Dict[str, str]],
) -> Optional[Dict[str, str]]:
"""Combine two projection maps. Identity-only; renames are not stored.
Args:
prev_projection_map: Previously-applied identity map.
new_projection_map: New identity map to compose.
Returns:
Combined identity map containing the columns present in both.
``None`` means "all columns" and acts as a passthrough.
"""
# Handle None cases (None means "all columns")
if prev_projection_map is None:
return new_projection_map
elif new_projection_map is None:
return prev_projection_map
# Both are identity maps; keep only columns present in both.
return {
name: name for name in prev_projection_map if name in new_projection_map
}
def apply_projection(
self,
projection_map: Optional[Dict[str, str]],
) -> "Datasource":
"""Apply a projection (column selection) to this datasource.
Args:
projection_map: Dict whose keys are the column names to select.
``None`` means select all columns. Any non-identity values
are ignored — the read stage does not rename.
Returns:
A new datasource instance with the projection applied.
"""
clone = copy.copy(self)
# Normalize any rename entries to identity — the read stage
# never renames.
normalized = None if projection_map is None else {k: k for k in projection_map}
clone._projection_map = self._combine_projection_map(
self._projection_map, normalized
)
return clone
class _DatasourcePredicatePushdownMixin:
"""Mixin for reading operators supporting predicate pushdown"""
def __init__(self):
self._predicate_expr: Optional[Expr] = None
def supports_predicate_pushdown(self) -> bool:
return False
def get_current_predicate(self) -> Optional[Expr]:
return self._predicate_expr
def apply_predicate(
self,
predicate_expr: Expr,
) -> "Datasource":
"""Apply a predicate to this datasource.
Default implementation that combines predicates using AND.
Subclasses that support predicate pushdown should have a _predicate_expr
attribute to store the predicate.
Note: Column rebinding is handled by the PredicatePushdown rule
before this method is called, so the predicate_expr should already
reference the correct column names.
"""
import copy
clone = copy.copy(self)
# Combine with existing predicate using AND
clone._predicate_expr = (
predicate_expr
if clone._predicate_expr is None
else clone._predicate_expr & predicate_expr
)
return clone
@PublicAPI
class Datasource(_DatasourceProjectionPushdownMixin, _DatasourcePredicatePushdownMixin):
"""Interface for defining a custom :class:`~ray.data.Dataset` datasource.
User may subclass this class to implement a custom datasource. The subclass should
implement :meth:`.get_read_tasks` and
:meth:`.estimate_inmemory_data_size` to read the data and estimate the in-memory data size, respectively.
To read a datasource into a dataset, use :meth:`~ray.data.read_datasource`.
Example:
>>> from ray.data.context import DataContext
>>> class MyDatasource(Datasource):
... def __init__(self, num_rows: int = 100):
... super().__init__()
... self.num_rows = num_rows
... def get_read_tasks(
... self,
... parallelism: int,
... per_task_row_limit: int | None = None,
... data_context: DataContext | None = None,
... ) -> List["ReadTask"]:
... # Split num_rows across parallelism tasks
... rows_per_task = self.num_rows // parallelism
... return [
... ReadTask(
... lambda: [pa.Table.from_pydict({"data": range(rows_per_task)})],
... BlockMetadata(rows_per_task, rows_per_task * 8, None, None),
... ) for _ in range(parallelism)
... ]
... def estimate_inmemory_data_size(self) -> Optional[int]:
... # Return total size for all data (independent of parallelism)
... return self.num_rows * 8
>>> ds = MyDatasource(num_rows=100)
>>> tasks = ds.get_read_tasks(parallelism=5)
>>> len(tasks) == 5
True
>>> tasks[0].metadata.num_rows == 20
True
>>> ds.estimate_inmemory_data_size() == sum(t.metadata.size_bytes for t in tasks)
True
""" # noqa: E501
def __init__(self):
"""Initialize the datasource and its mixins."""
_DatasourcePredicatePushdownMixin.__init__(self)
@Deprecated
def create_reader(self, **read_args) -> "Reader":
"""
Deprecated: Implement :meth:`~ray.data.Datasource.get_read_tasks` and
:meth:`~ray.data.Datasource.estimate_inmemory_data_size` instead.
"""
return _LegacyDatasourceReader(self, **read_args)
@Deprecated
def prepare_read(self, parallelism: int, **read_args) -> List["ReadTask"]:
"""
Deprecated: Implement :meth:`~ray.data.Datasource.get_read_tasks` and
:meth:`~ray.data.Datasource.estimate_inmemory_data_size` instead.
"""
raise NotImplementedError
def get_name(self) -> str:
"""Return a human-readable name for this datasource.
This will be used as the names of the read tasks.
"""
name = type(self).__name__
datasource_suffix = "Datasource"
if name.endswith(datasource_suffix):
name = name[: -len(datasource_suffix)]
return name
def estimate_inmemory_data_size(self) -> Optional[int]:
"""Return an estimate of the in-memory data size, or None if unknown.
Note that the in-memory data size may be larger than the on-disk data size.
"""
raise NotImplementedError
def get_read_tasks(
self,
parallelism: int,
per_task_row_limit: Optional[int] = None,
data_context: Optional["DataContext"] = None,
) -> List["ReadTask"]:
"""Execute the read and return read tasks.
Args:
parallelism: The requested read parallelism. The number of read
tasks should equal to this value if possible.
per_task_row_limit: The per-task row limit for the read tasks.
data_context: The data context to use to get read tasks.
Returns:
A list of read tasks that can be executed to read blocks from the
datasource in parallel.
"""
raise NotImplementedError
@property
def should_create_reader(self) -> bool:
"""Return True if the datasource should create a legacy reader"""
has_implemented_get_read_tasks = (
type(self).get_read_tasks is not Datasource.get_read_tasks
)
has_implemented_estimate_inmemory_data_size = (
type(self).estimate_inmemory_data_size
is not Datasource.estimate_inmemory_data_size
)
# False when both get_read_tasks and estimate_inmemory_data_size are implemented
return not (
has_implemented_get_read_tasks
and has_implemented_estimate_inmemory_data_size
)
@property
def supports_distributed_reads(self) -> bool:
"""If ``False``, only launch read tasks on the driver's node."""
return True
@Deprecated
class Reader:
"""A bound read operation for a :class:`~ray.data.Datasource`.
This is a stateful class so that reads can be prepared in multiple stages.
For example, it is useful for :class:`Datasets <ray.data.Dataset>` to know the
in-memory size of the read prior to executing it.
"""
def estimate_inmemory_data_size(self) -> Optional[int]:
"""Return an estimate of the in-memory data size, or None if unknown.
Note that the in-memory data size may be larger than the on-disk data size.
"""
raise NotImplementedError
def get_read_tasks(self, parallelism: int) -> List["ReadTask"]:
"""Execute the read and return read tasks.
Args:
parallelism: The requested read parallelism. The number of read
tasks should equal to this value if possible.
Returns:
A list of read tasks that can be executed to read blocks from the
datasource in parallel.
"""
raise NotImplementedError
class _LegacyDatasourceReader(Reader):
def __init__(self, datasource: Datasource, **read_args):
self._datasource = datasource
self._read_args = read_args
def estimate_inmemory_data_size(self) -> Optional[int]:
return None
def get_read_tasks(
self,
parallelism: int,
per_task_row_limit: Optional[int] = None,
data_context: Optional["DataContext"] = None,
) -> List["ReadTask"]:
"""Execute the read and return read tasks.
Args:
parallelism: The requested read parallelism. The number of read
tasks should equal to this value if possible.
per_task_row_limit: The per-task row limit for the read tasks.
data_context: The data context to use to get read tasks. Not used by this
legacy reader.
Returns:
A list of read tasks that can be executed to read blocks from the
datasource in parallel.
"""
return self._datasource.prepare_read(parallelism, **self._read_args)
@DeveloperAPI
class ReadTask(Callable[[], Iterable[Block]]):
"""A function used to read blocks from the :class:`~ray.data.Dataset`.
Read tasks are generated by :meth:`~ray.data.Datasource.get_read_tasks`,
and return a list of ``ray.data.Block`` when called. Initial metadata about the read
operation can be retrieved via the ``metadata`` attribute prior to executing the
read. Final metadata is returned after the read along with the blocks.
Ray will execute read tasks in remote functions to parallelize execution.
Note that the number of blocks returned can vary at runtime. For example,
if a task is reading a single large file it can return multiple blocks to
avoid running out of memory during the read.
The initial metadata should reflect all the blocks returned by the read,
e.g., if the metadata says ``num_rows=1000``, the read can return a single
block of 1000 rows, or multiple blocks with 1000 rows altogether.
The final metadata (returned with the actual block) reflects the exact
contents of the block itself.
"""
def __init__(
self,
read_fn: Callable[[], Iterable[Block]],
metadata: BlockMetadata,
schema: Optional["Schema"] = None,
per_task_row_limit: Optional[int] = None,
):
self._metadata = metadata
self._read_fn = read_fn
self._schema = schema
self._per_task_row_limit = per_task_row_limit
@property
def metadata(self) -> BlockMetadata:
return self._metadata
# TODO(justin): We want to remove schema from `ReadTask` later on
@property
def schema(self) -> Optional["Schema"]:
return self._schema
@property
def read_fn(self) -> Callable[[], Iterable[Block]]:
return self._read_fn
@property
def per_task_row_limit(self) -> Optional[int]:
"""Get the per-task row limit for this read task."""
return self._per_task_row_limit
def __call__(self) -> Iterable[Block]:
result = self._read_fn()
if not hasattr(result, "__iter__"):
DeprecationWarning(
"Read function must return Iterable[Block], got {}. "
"Probably you need to return `[block]` instead of "
"`block`.".format(result)
)
if self._per_task_row_limit is None:
yield from result
return
yield from _iter_sliced_blocks(result, self._per_task_row_limit)
@DeveloperAPI
class RandomIntRowDatasource(Datasource):
"""An example datasource that generates rows with random int64 columns.
Examples:
>>> import ray
>>> from ray.data.datasource import RandomIntRowDatasource
>>> source = RandomIntRowDatasource() # doctest: +SKIP
>>> ray.data.read_datasource( # doctest: +SKIP
... source, n=10, num_columns=2).take()
{'c_0': 1717767200176864416, 'c_1': 999657309586757214}
{'c_0': 4983608804013926748, 'c_1': 1160140066899844087}
"""
def __init__(self, n: int, num_columns: int):
"""Initialize the datasource that generates random-integer rows.
Args:
n: The number of rows to generate.
num_columns: The number of columns to generate.
"""
self._n = n
self._num_columns = num_columns
def estimate_inmemory_data_size(self) -> Optional[int]:
return self._n * self._num_columns * 8
def get_read_tasks(
self,
parallelism: int,
per_task_row_limit: Optional[int] = None,
data_context: Optional["DataContext"] = None,
) -> List[ReadTask]:
_check_pyarrow_version()
import pyarrow
read_tasks: List[ReadTask] = []
n = self._n
num_columns = self._num_columns
block_size = max(1, n // parallelism)
def make_block(count: int, num_columns: int) -> Block:
return pyarrow.Table.from_arrays(
np.random.randint(
np.iinfo(np.int64).max, size=(num_columns, count), dtype=np.int64
),
names=[f"c_{i}" for i in range(num_columns)],
)
schema = pyarrow.Table.from_pydict(
{f"c_{i}": [0] for i in range(num_columns)}
).schema
i = 0
while i < n:
count = min(block_size, n - i)
meta = BlockMetadata(
num_rows=count,
size_bytes=8 * count * num_columns,
input_files=None,
exec_stats=None,
)
read_tasks.append(
ReadTask(
lambda count=count, num_columns=num_columns: [
make_block(count, num_columns)
],
meta,
schema=schema,
per_task_row_limit=per_task_row_limit,
)
)
i += block_size
return read_tasks
def get_name(self) -> str:
"""Return a human-readable name for this datasource.
This will be used as the names of the read tasks.
Note: overrides the base `Datasource` method.
"""
return "RandomInt"
python/ray/data/datasource/file_based_datasource.py
+625
View File
@@ -0,0 +1,625 @@
import io
import logging
from dataclasses import dataclass
from typing import (
TYPE_CHECKING,
Any,
Callable,
Dict,
Iterable,
Iterator,
List,
Literal,
Optional,
Tuple,
TypeVar,
Union,
)
import numpy as np
import ray
from ray.data._internal.util import (
RetryingContextManager,
RetryingPyFileSystem,
_check_pyarrow_version,
_is_local_scheme,
infer_compression,
iterate_with_retry,
make_async_gen,
)
from ray.data.block import Block, BlockAccessor
from ray.data.context import DataContext
from ray.data.datasource.datasource import Datasource, ReadTask
from ray.data.datasource.file_meta_provider import (
BaseFileMetadataProvider,
DefaultFileMetadataProvider,
)
from ray.data.datasource.partitioning import (
Partitioning,
PathPartitionFilter,
PathPartitionParser,
)
from ray.data.datasource.path_util import (
_has_file_extension,
_resolve_paths_and_filesystem,
)
from ray.util.annotations import DeveloperAPI
if TYPE_CHECKING:
import pandas as pd
import pyarrow
logger = logging.getLogger(__name__)
# We should parallelize file size fetch operations beyond this threshold.
FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD = 16
# 16 file size fetches from S3 takes ~1.5 seconds with Arrow's S3FileSystem.
PATHS_PER_FILE_SIZE_FETCH_TASK = 16
@DeveloperAPI
@dataclass
class FileShuffleConfig:
"""Configuration for file shuffling.
This configuration object controls how files are shuffled while reading file-based
datasets. The random seed behavior is determined by the combination of ``seed``
and ``reseed_after_execution``:
- If ``seed`` is None, the random seed is always None (non-deterministic shuffling).
- If ``seed`` is not None and ``reseed_after_execution`` is False, the random seed is
constantly ``seed`` across executions.
- If ``seed`` is not None and ``reseed_after_execution`` is True, the random seed is
different for each execution.
.. note::
Even if you provided a seed, you might still observe a non-deterministic row
order. This is because tasks are executed in parallel and their completion
order might vary. If you need to preserve the order of rows, set
``DataContext.get_current().execution_options.preserve_order``.
Args:
seed: An optional integer seed for the file shuffler. If None, shuffling is
non-deterministic. If provided, shuffling is deterministic based on this
seed and the ``reseed_after_execution`` setting.
reseed_after_execution: If True, the random seed considers both ``seed`` and
``execution_idx``, resulting in different shuffling orders across executions.
If False, the random seed is constantly ``seed``, resulting in the same
shuffling order across executions. Only takes effect when ``seed`` is not None.
Defaults to True.
Example:
>>> import ray
>>> from ray.data import FileShuffleConfig
>>> # Fixed seed - same shuffle across executions
>>> shuffle = FileShuffleConfig(seed=42, reseed_after_execution=False)
>>> ds = ray.data.read_images("s3://anonymous@ray-example-data/batoidea", shuffle=shuffle)
>>>
>>> # Seed with reseed_after_execution - different shuffle per execution
>>> shuffle = FileShuffleConfig(seed=42, reseed_after_execution=True)
>>> ds = ray.data.read_images("s3://anonymous@ray-example-data/batoidea", shuffle=shuffle)
""" # noqa: E501
seed: Optional[int] = None
reseed_after_execution: bool = True
def __post_init__(self):
"""Ensure that the seed is either None or an integer."""
if self.seed is not None and not isinstance(self.seed, int):
raise ValueError("Seed must be an integer or None.")
def get_seed(self, execution_idx: int = 0) -> Optional[int]:
if self.seed is None:
return None
elif self.reseed_after_execution:
# Modulo ensures the result is in valid NumPy seed range [0, 2**32 - 1].
return hash((self.seed, execution_idx)) % (2**32)
else:
return self.seed
@DeveloperAPI
class FileBasedDatasource(Datasource):
"""File-based datasource for reading files.
Don't use this class directly. Instead, subclass it and implement `_read_stream()`.
"""
# If `_WRITE_FILE_PER_ROW` is `True`, this datasource calls `_write_row` and writes
# each row to a file. Otherwise, this datasource calls `_write_block` and writes
# each block to a file.
_WRITE_FILE_PER_ROW = False
_FILE_EXTENSIONS: Optional[Union[str, List[str]]] = None
# Number of threads for concurrent reading within each read task.
# If zero or negative, reading will be performed in the main thread.
_NUM_THREADS_PER_TASK = 0
def __init__(
self,
paths: Union[str, List[str]],
*,
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
schema: Optional[Union[type, "pyarrow.lib.Schema"]] = None,
open_stream_args: Optional[Dict[str, Any]] = None,
meta_provider: BaseFileMetadataProvider = DefaultFileMetadataProvider(),
partition_filter: PathPartitionFilter = None,
partitioning: Partitioning = None,
ignore_missing_paths: bool = False,
shuffle: Optional[Union[Literal["files"], FileShuffleConfig]] = None,
include_paths: bool = False,
file_extensions: Optional[List[str]] = None,
):
super().__init__()
_check_pyarrow_version()
self._supports_distributed_reads = not _is_local_scheme(paths)
if not self._supports_distributed_reads and ray.util.client.ray.is_connected():
raise ValueError(
"Because you're using Ray Client, read tasks scheduled on the Ray "
"cluster can't access your local files. To fix this issue, store "
"files in cloud storage or a distributed filesystem like NFS."
)
self._schema = schema
self._data_context = DataContext.get_current()
self._open_stream_args = open_stream_args
self._meta_provider = meta_provider
self._partition_filter = partition_filter
self._partitioning = partitioning
self._ignore_missing_paths = ignore_missing_paths
self._include_paths = include_paths
# Need this property for lineage tracking. We should not directly assign paths
# to self since it is captured every read_task_fn during serialization and
# causing this data being duplicated and excessive object store spilling.
self._source_paths_ref = ray.put(paths)
paths, self._filesystem = _resolve_paths_and_filesystem(paths, filesystem)
self._filesystem = RetryingPyFileSystem.wrap(
self._filesystem, retryable_errors=self._data_context.retried_io_errors
)
paths, file_sizes = map(
list,
zip(
*meta_provider.expand_paths(
paths,
self._filesystem,
partitioning,
ignore_missing_paths=ignore_missing_paths,
)
),
)
if ignore_missing_paths and len(paths) == 0:
raise ValueError(
"None of the provided paths exist. "
"The 'ignore_missing_paths' field is set to True."
)
if self._partition_filter is not None:
# Use partition filter to skip files which are not needed.
path_to_size = dict(zip(paths, file_sizes))
paths = self._partition_filter(paths)
file_sizes = [path_to_size[p] for p in paths]
if len(paths) == 0:
raise ValueError(
"No input files found to read. Please double check that "
"'partition_filter' field is set properly."
)
if file_extensions is not None:
path_to_size = dict(zip(paths, file_sizes))
paths = [p for p in paths if _has_file_extension(p, file_extensions)]
file_sizes = [path_to_size[p] for p in paths]
if len(paths) == 0:
raise ValueError(
"No input files found to read with the following file extensions: "
f"{file_extensions}. Please double check that "
"'file_extensions' field is set properly."
)
_validate_shuffle_arg(shuffle)
self._shuffle = shuffle
# Read tasks serialize `FileBasedDatasource` instances, and the list of paths
# can be large. To avoid slow serialization speeds, we store a reference to
# the paths rather than the paths themselves.
self._paths_ref = ray.put(paths)
self._file_sizes_ref = ray.put(file_sizes)
@property
def _source_paths(self) -> List[str]:
return ray.get(self._source_paths_ref)
def _paths(self) -> List[str]:
return ray.get(self._paths_ref)
def _file_sizes(self) -> List[float]:
return ray.get(self._file_sizes_ref)
def estimate_inmemory_data_size(self) -> Optional[int]:
total_size = 0
for sz in self._file_sizes():
if sz is not None:
total_size += sz
return total_size
def get_read_tasks(
self,
parallelism: int,
per_task_row_limit: Optional[int] = None,
data_context: Optional["DataContext"] = None,
) -> List[ReadTask]:
import numpy as np
open_stream_args = self._open_stream_args
partitioning = self._partitioning
paths = self._paths()
file_sizes = self._file_sizes()
execution_idx = data_context._execution_idx if data_context is not None else 0
paths, file_sizes = _shuffle_file_metadata(
paths, file_sizes, self._shuffle, execution_idx
)
filesystem = _wrap_s3_serialization_workaround(self._filesystem)
if open_stream_args is None:
open_stream_args = {}
def read_files(
read_paths: Iterable[str],
) -> Iterable[Block]:
nonlocal filesystem, open_stream_args, partitioning
fs = _unwrap_s3_serialization_workaround(filesystem)
for read_path in read_paths:
partitions: Dict[str, str] = {}
if partitioning is not None:
parse = PathPartitionParser(partitioning)
partitions = parse(read_path)
with RetryingContextManager(
self._open_input_source(fs, read_path, **open_stream_args),
context=self._data_context,
) as f:
for block in iterate_with_retry(
lambda: self._read_stream(f, read_path),
description="read stream iteratively",
match=self._data_context.retried_io_errors,
):
if partitions:
block = _add_partitions(block, partitions)
if self._include_paths:
block_accessor = BlockAccessor.for_block(block)
block = block_accessor.fill_column("path", read_path)
yield block
def create_read_task_fn(read_paths, num_threads):
def read_task_fn():
nonlocal num_threads, read_paths
# TODO: We should refactor the code so that we can get the results in
# order even when using multiple threads.
if self._data_context.execution_options.preserve_order:
num_threads = 0
if num_threads > 0:
num_threads = min(num_threads, len(read_paths))
logger.debug(
f"Reading {len(read_paths)} files with {num_threads} threads."
)
yield from make_async_gen(
iter(read_paths),
read_files,
num_workers=num_threads,
preserve_ordering=True,
)
else:
logger.debug(f"Reading {len(read_paths)} files.")
yield from read_files(read_paths)
return read_task_fn
# fix https://github.com/ray-project/ray/issues/24296
parallelism = min(parallelism, len(paths))
read_tasks = []
# Convert numpy arrays back to Python lists so downstream code
# (e.g. meta providers) doesn't receive numpy string types.
split_paths = [p.tolist() for p in np.array_split(paths, parallelism)]
split_file_sizes = [s.tolist() for s in np.array_split(file_sizes, parallelism)]
for read_paths, file_sizes in zip(split_paths, split_file_sizes):
if len(read_paths) <= 0:
continue
meta = self._meta_provider(
read_paths,
rows_per_file=self._rows_per_file(),
file_sizes=file_sizes,
)
read_task_fn = create_read_task_fn(read_paths, self._NUM_THREADS_PER_TASK)
read_task = ReadTask(
read_task_fn, meta, per_task_row_limit=per_task_row_limit
)
read_tasks.append(read_task)
return read_tasks
def resolve_compression(
self, path: str, open_args: Dict[str, Any]
) -> Optional[str]:
"""Resolves the compression format for a stream.
Args:
path: The file path to resolve compression for.
open_args: kwargs passed to
`pyarrow.fs.FileSystem.open_input_stream <https://arrow.apache.org/docs/python/generated/pyarrow.fs.FileSystem.html#pyarrow.fs.FileSystem.open_input_stream>`_
when opening input files to read.
Returns:
The compression format (e.g., "gzip", "snappy", "bz2") or None if
no compression is detected or specified.
"""
compression = open_args.get("compression", None)
if compression is None:
compression = infer_compression(path)
return compression
def _resolve_buffer_size(self, open_args: Dict[str, Any]) -> Optional[int]:
buffer_size = open_args.pop("buffer_size", None)
if buffer_size is None:
buffer_size = self._data_context.streaming_read_buffer_size
return buffer_size
def _file_to_snappy_stream(
self,
file: "pyarrow.NativeFile",
filesystem: "RetryingPyFileSystem",
) -> "pyarrow.PythonFile":
import pyarrow as pa
import snappy
from pyarrow.fs import HadoopFileSystem
stream = io.BytesIO()
if isinstance(filesystem.unwrap(), HadoopFileSystem):
snappy.hadoop_snappy.stream_decompress(src=file, dst=stream)
else:
snappy.stream_decompress(src=file, dst=stream)
stream.seek(0)
return pa.PythonFile(stream, mode="r")
def _open_input_source(
self,
filesystem: "RetryingPyFileSystem",
path: str,
**open_args,
) -> "pyarrow.NativeFile":
"""Opens a source path for reading and returns the associated Arrow NativeFile.
The default implementation opens the source path as a sequential input stream,
using self._data_context.streaming_read_buffer_size as the buffer size if none
is given by the caller.
Implementations that do not support streaming reads (e.g. that require random
access) should override this method.
"""
compression = self.resolve_compression(path, open_args)
buffer_size = self._resolve_buffer_size(open_args)
if compression == "snappy":
# Arrow doesn't support streaming Snappy decompression since the canonical
# C++ Snappy library doesn't natively support streaming decompression. We
# works around this by manually decompressing the file with python-snappy.
open_args["compression"] = None
file = filesystem.open_input_stream(
path, buffer_size=buffer_size, **open_args
)
return self._file_to_snappy_stream(file, filesystem)
open_args["compression"] = compression
return filesystem.open_input_stream(path, buffer_size=buffer_size, **open_args)
def _rows_per_file(self):
"""Returns the number of rows per file, or None if unknown."""
return None
def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
"""Streaming read a single file.
This method should be implemented by subclasses.
"""
raise NotImplementedError(
"Subclasses of FileBasedDatasource must implement _read_stream()."
)
@property
def supports_distributed_reads(self) -> bool:
return self._supports_distributed_reads
def _add_partitions(
data: Union["pyarrow.Table", "pd.DataFrame"], partitions: Dict[str, Any]
) -> Union["pyarrow.Table", "pd.DataFrame"]:
import pandas as pd
import pyarrow as pa
assert isinstance(data, (pa.Table, pd.DataFrame))
if isinstance(data, pa.Table):
return _add_partitions_to_table(data, partitions)
if isinstance(data, pd.DataFrame):
return _add_partitions_to_dataframe(data, partitions)
def _add_partitions_to_table(
table: "pyarrow.Table", partitions: Dict[str, Any]
) -> "pyarrow.Table":
import pyarrow as pa
import pyarrow.compute as pc
column_names = set(table.column_names)
for field, value in partitions.items():
column = pa.array([value] * len(table))
if field in column_names:
# TODO: Handle cast error.
column_type = table.schema.field(field).type
column = column.cast(column_type)
values_are_equal = pc.all(pc.equal(column, table[field]))
values_are_equal = values_are_equal.as_py()
if not values_are_equal:
raise ValueError(
f"Partition column {field} exists in table data, but partition "
f"value '{value}' is different from in-data values: "
f"{table[field].unique().to_pylist()}."
)
i = table.schema.get_field_index(field)
table = table.set_column(i, field, column)
else:
table = table.append_column(field, column)
return table
def _add_partitions_to_dataframe(
df: "pd.DataFrame", partitions: Dict[str, Any]
) -> "pd.DataFrame":
import pandas as pd
for field, value in partitions.items():
column = pd.Series(data=[value] * len(df), name=field)
if field in df:
column = column.astype(df[field].dtype)
mask = df[field].notna()
if not df[field][mask].equals(column[mask]):
raise ValueError(
f"Partition column {field} exists in table data, but partition "
f"value '{value}' is different from in-data values: "
f"{list(df[field].unique())}."
)
df[field] = column
return df
def _wrap_s3_serialization_workaround(filesystem: "pyarrow.fs.FileSystem"):
# This is needed because pa.fs.S3FileSystem assumes pa.fs is already
# imported before deserialization. See #17085.
import pyarrow as pa
import pyarrow.fs
base_fs = filesystem
if isinstance(filesystem, RetryingPyFileSystem):
base_fs = filesystem.unwrap()
if isinstance(base_fs, pa.fs.S3FileSystem):
return _S3FileSystemWrapper(filesystem)
return filesystem
def _unwrap_s3_serialization_workaround(
filesystem: Union["pyarrow.fs.FileSystem", "_S3FileSystemWrapper"],
):
if isinstance(filesystem, _S3FileSystemWrapper):
filesystem = filesystem.unwrap()
return filesystem
class _S3FileSystemWrapper:
"""pyarrow.fs.S3FileSystem wrapper that can be deserialized safely.
Importing pyarrow.fs during reconstruction triggers the pyarrow
S3 subsystem initialization.
NOTE: This is only needed for pyarrow<14.0.0 and should be removed
once the minimum supported pyarrow version exceeds that.
See https://github.com/apache/arrow/pull/38375 for context.
"""
def __init__(self, fs: "pyarrow.fs.FileSystem"):
self._fs = fs
def unwrap(self):
return self._fs
@classmethod
def _reconstruct(cls, fs_reconstruct, fs_args):
# Implicitly trigger S3 subsystem initialization by importing
# pyarrow.fs.
import pyarrow.fs # noqa: F401
return cls(fs_reconstruct(*fs_args))
def __reduce__(self):
return _S3FileSystemWrapper._reconstruct, self._fs.__reduce__()
def _resolve_kwargs(
kwargs_fn: Callable[[], Dict[str, Any]], **kwargs
) -> Dict[str, Any]:
if kwargs_fn:
kwarg_overrides = kwargs_fn()
kwargs.update(kwarg_overrides)
return kwargs
def _validate_shuffle_arg(
shuffle: Union[Literal["files"], FileShuffleConfig, None],
) -> None:
if not (
shuffle is None or shuffle == "files" or isinstance(shuffle, FileShuffleConfig)
):
raise ValueError(
f"Invalid value for 'shuffle': {shuffle}. "
"Valid values are None, 'files', `FileShuffleConfig`."
)
FileMetadata = TypeVar("FileMetadata")
def _shuffle_file_metadata(
paths: List[str],
file_metadata: List[FileMetadata],
shuffler: Union[Literal["files"], FileShuffleConfig, None],
execution_idx: int,
) -> Tuple[List[str], List[FileMetadata]]:
"""Shuffle file paths and sizes together using the given shuffler."""
if shuffler is None:
return paths, file_metadata
assert len(paths) == len(file_metadata), (
"Number of paths and file metadata must match. "
f"Got {len(paths)} paths and {len(file_metadata)} file metadata."
)
if len(paths) == 0:
return paths, file_metadata
if shuffler == "files":
seed = None
else:
assert isinstance(shuffler, FileShuffleConfig)
seed = shuffler.get_seed(execution_idx)
file_metadata_shuffler = np.random.default_rng(seed)
files_metadata = list(zip(paths, file_metadata))
file_metadata_shuffler.shuffle(files_metadata)
return list(map(list, zip(*files_metadata)))
python/ray/data/datasource/file_datasink.py
+308
View File
@@ -0,0 +1,308 @@
import logging
import posixpath
from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional
from urllib.parse import urlparse
from ray._common.retry import call_with_retry
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
from ray.data._internal.execution.interfaces import TaskContext
from ray.data._internal.planner.plan_write_op import WRITE_UUID_KWARG_NAME
from ray.data._internal.savemode import SaveMode
from ray.data._internal.util import (
RetryingPyFileSystem,
_is_local_scheme,
)
from ray.data._internal.utils.arrow_utils import add_creatable_buckets_param_if_s3_uri
from ray.data.block import Block, BlockAccessor
from ray.data.context import DataContext
from ray.data.datasource.datasink import Datasink, WriteResult
from ray.data.datasource.filename_provider import (
FilenameProvider,
_split_base_and_ext,
)
from ray.data.datasource.path_util import _resolve_paths_and_filesystem
from ray.util.annotations import DeveloperAPI
if TYPE_CHECKING:
import pyarrow
logger = logging.getLogger(__name__)
class _FileDatasink(Datasink[None]):
def __init__(
self,
path: str,
*,
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
try_create_dir: bool = True,
open_stream_args: Optional[Dict[str, Any]] = None,
filename_provider: Optional[FilenameProvider] = None,
dataset_uuid: Optional[str] = None,
file_format: Optional[str] = None,
mode: SaveMode = SaveMode.APPEND,
):
"""Initialize this datasink.
Args:
path: The folder to write files to.
filesystem: The filesystem to write files to. If not provided, the
filesystem is inferred from the path.
try_create_dir: Whether to create the directory to write files to.
open_stream_args: Arguments to pass to ``filesystem.open_output_stream``.
filename_provider: A :class:`ray.data.datasource.FilenameProvider` that
generates filenames for each row or block.
dataset_uuid: The UUID of the dataset being written. If specified, it's
included in the filename.
file_format: The file extension. If specified, files are written with this
extension.
mode: The save mode controlling behavior when the destination already
exists (e.g., append, overwrite, error, ignore).
"""
if open_stream_args is None:
open_stream_args = {}
if filename_provider is None:
filename_provider = FilenameProvider(
dataset_uuid=dataset_uuid, file_format=file_format
)
self._data_context = DataContext.get_current()
self.unresolved_path = path
paths, self.filesystem = _resolve_paths_and_filesystem(path, filesystem)
self.filesystem = RetryingPyFileSystem.wrap(
self.filesystem, retryable_errors=self._data_context.retried_io_errors
)
assert len(paths) == 1, len(paths)
self.path = paths[0]
self.try_create_dir = try_create_dir
self.open_stream_args = open_stream_args
self.filename_provider = filename_provider
self.dataset_uuid = dataset_uuid
self.file_format = file_format
self.mode = mode
self.has_created_dir = False
self._skip_write = False
self._write_started = False
def open_output_stream(self, path: str) -> "pyarrow.NativeFile":
return self.filesystem.open_output_stream(path, **self.open_stream_args)
def on_write_start(self, schema: Optional["pyarrow.Schema"] = None) -> None:
# Make idempotent - if already called, return early.
if self._write_started:
return
self._write_started = True
from pyarrow.fs import FileType
dir_exists = (
self.filesystem.get_file_info(self.path).type is not FileType.NotFound
)
if dir_exists:
if self.mode in {SaveMode.ERROR, SaveMode.CREATE}:
raise ValueError(
f"Path {self.path} already exists. "
"If this is unexpected, use mode='ignore' to ignore those files"
)
if self.mode == SaveMode.IGNORE:
logger.warning(f"[SaveMode={self.mode}] Skipping {self.path}")
self._skip_write = True
return
if self.mode == SaveMode.OVERWRITE:
logger.warning(f"[SaveMode={self.mode}] Replacing contents {self.path}")
self.filesystem.delete_dir_contents(self.path)
self.has_created_dir = self._create_dir(self.path)
def _create_dir(self, dest) -> bool:
"""Create a directory to write files to.
If ``try_create_dir`` is ``False``, this method is a no-op.
"""
from pyarrow.fs import FileType
# We should skip creating directories in s3 unless the user specifically
# overrides this behavior. PyArrow's s3fs implementation for create_dir
# will attempt to check if the parent directory exists before trying to
# create the directory (with recursive=True it will try to do this to
# all of the directories until the root of the bucket). An IAM Policy that
# restricts access to a subset of prefixes within the bucket might cause
# the creation of the directory to fail even if the permissions should
# allow the data can be written to the specified path. For example if a
# a policy only allows users to write blobs prefixed with s3://bucket/foo
# a call to create_dir for s3://bucket/foo/bar will fail even though it
# should not.
parsed_uri = urlparse(dest)
is_s3_uri = parsed_uri.scheme == "s3"
skip_create_dir_for_s3 = is_s3_uri and not self._data_context.s3_try_create_dir
if self.try_create_dir and not skip_create_dir_for_s3:
if self.filesystem.get_file_info(dest).type is FileType.NotFound:
# Arrow's S3FileSystem doesn't allow creating buckets by default, so we
# add a query arg enabling bucket creation if an S3 URI is provided.
tmp = add_creatable_buckets_param_if_s3_uri(dest)
self.filesystem.create_dir(tmp, recursive=True)
return True
return False
def write(
self,
blocks: Iterable[Block],
ctx: TaskContext,
) -> None:
builder = DelegatingBlockBuilder()
for block in blocks:
builder.add_block(block)
block = builder.build()
block_accessor = BlockAccessor.for_block(block)
if block_accessor.num_rows() == 0:
logger.warning(f"Skipped writing empty block to {self.path}")
return
self.write_block(block_accessor, 0, ctx)
def write_block(self, block: BlockAccessor, block_index: int, ctx: TaskContext):
raise NotImplementedError
def on_write_complete(self, write_result: WriteResult[None]):
# If no rows were written, we can delete the directory.
if self.has_created_dir and write_result.num_rows == 0:
self.filesystem.delete_dir(self.path)
@property
def supports_distributed_writes(self) -> bool:
return not _is_local_scheme(self.unresolved_path)
@DeveloperAPI
class RowBasedFileDatasink(_FileDatasink):
"""A datasink that writes one row to each file.
Subclasses must implement ``write_row_to_file`` and call the superclass constructor.
Examples:
.. testcode::
import io
from typing import Any, Dict
import pyarrow
from PIL import Image
from ray.data.datasource import RowBasedFileDatasink
class ImageDatasink(RowBasedFileDatasink):
def __init__(self, path: str, *, column: str, file_format: str = "png"):
super().__init__(path, file_format=file_format)
self._file_format = file_format
self._column = column
def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
image = Image.fromarray(row[self._column])
buffer = io.BytesIO()
image.save(buffer, format=self._file_format)
file.write(buffer.getvalue())
""" # noqa: E501
def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
"""Write a row to a file.
Args:
row: The row to write.
file: The file to write the row to.
"""
raise NotImplementedError
def write_block(self, block: BlockAccessor, block_index: int, ctx: TaskContext):
task_filename = self.filename_provider.get_filename_for_task(
ctx.kwargs[WRITE_UUID_KWARG_NAME],
ctx.task_idx,
)
base, ext = _split_base_and_ext(task_filename)
for row_index, row in enumerate(block.iter_rows(public_row_format=False)):
filename = f"{base}_{block_index:06}_{row_index:06}{ext}"
write_path = posixpath.join(self.path, filename)
logger.debug(f"Writing {write_path} file.")
def write_row_to_path():
with self.open_output_stream(write_path) as file:
self.write_row_to_file(row, file)
call_with_retry(
write_row_to_path,
description=f"write '{write_path}'",
match=self._data_context.retried_io_errors,
)
@DeveloperAPI
class BlockBasedFileDatasink(_FileDatasink):
"""A datasink that writes multiple rows to each file.
Subclasses must implement ``write_block_to_file`` and call the superclass
constructor.
Examples:
.. testcode::
class CSVDatasink(BlockBasedFileDatasink):
def __init__(self, path: str):
super().__init__(path, file_format="csv")
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
from pyarrow import csv
csv.write_csv(block.to_arrow(), file)
""" # noqa: E501
def __init__(
self,
path: str,
*,
min_rows_per_file: Optional[int] = None,
**file_datasink_kwargs,
):
"""Initialize this block-based file datasink.
Args:
path: The folder to write files to.
min_rows_per_file: The target minimum number of rows per file. When
``None``, rows are not buffered before being written.
**file_datasink_kwargs: Additional keyword arguments forwarded to
:class:`_FileDatasink`.
"""
super().__init__(path, **file_datasink_kwargs)
self._min_rows_per_file = min_rows_per_file
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
"""Write a block of data to a file.
Args:
block: The block to write.
file: The file to write the block to.
"""
raise NotImplementedError
def write_block(self, block: BlockAccessor, block_index: int, ctx: TaskContext):
filename = self.filename_provider.get_filename_for_task(
ctx.kwargs[WRITE_UUID_KWARG_NAME], ctx.task_idx
)
write_path = posixpath.join(self.path, filename)
def write_block_to_path():
with self.open_output_stream(write_path) as file:
self.write_block_to_file(block, file)
logger.debug(f"Writing {write_path} file.")
call_with_retry(
write_block_to_path,
description=f"write '{write_path}'",
match=self._data_context.retried_io_errors,
)
@property
def min_rows_per_write(self) -> Optional[int]:
return self._min_rows_per_file
python/ray/data/datasource/file_meta_provider.py
+491
View File
@@ -0,0 +1,491 @@
import itertools
import logging
import os
import pathlib
import re
from typing import (
TYPE_CHECKING,
Callable,
Iterator,
List,
Optional,
Tuple,
TypeVar,
Union,
)
import numpy as np
from ray.data._internal.execution.util import merge_label_selector
from ray.data._internal.progress.progress_bar import ProgressBar
from ray.data._internal.remote_fn import cached_remote_fn
from ray.data._internal.util import RetryingPyFileSystem
from ray.data.block import BlockMetadata
from ray.data.context import DataContext
from ray.data.datasource.partitioning import Partitioning, PathPartitionFilter
from ray.data.datasource.path_util import _has_file_extension
from ray.util.annotations import DeveloperAPI
if TYPE_CHECKING:
import pyarrow
logger = logging.getLogger(__name__)
@DeveloperAPI
class FileMetadataProvider:
"""Abstract callable that provides metadata for the files of a single dataset block.
Current subclasses:
- :class:`BaseFileMetadataProvider`
"""
def _get_block_metadata(
self,
paths: List[str],
**kwargs,
) -> BlockMetadata:
"""Resolves and returns block metadata for files in the given paths.
All file paths provided should belong to a single dataset block.
Args:
paths: The file paths for a single dataset block.
**kwargs: Additional kwargs used to determine block metadata.
Returns:
BlockMetadata aggregated across the given paths.
"""
raise NotImplementedError
def __call__(
self,
paths: List[str],
**kwargs,
) -> BlockMetadata:
return self._get_block_metadata(paths, **kwargs)
@DeveloperAPI
class BaseFileMetadataProvider(FileMetadataProvider):
"""Abstract callable that provides metadata for
:class:`~ray.data.datasource.file_based_datasource.FileBasedDatasource`
implementations that reuse the base :meth:`~ray.data.Datasource.prepare_read`
method.
Also supports file and file size discovery in input directory paths.
Current subclasses:
- :class:`DefaultFileMetadataProvider`
"""
def _get_block_metadata(
self,
paths: List[str],
*,
rows_per_file: Optional[int],
file_sizes: List[Optional[int]],
) -> BlockMetadata:
"""Resolves and returns block metadata for files of a single dataset block.
Args:
paths: The file paths for a single dataset block. These
paths will always be a subset of those previously returned from
:meth:`.expand_paths`.
rows_per_file: The fixed number of rows per input file, or None.
file_sizes: Optional file size per input file previously returned
from :meth:`.expand_paths`, where `file_sizes[i]` holds the size of
the file at `paths[i]`.
Returns:
BlockMetadata aggregated across the given file paths.
"""
raise NotImplementedError
def expand_paths(
self,
paths: List[str],
filesystem: Optional["RetryingPyFileSystem"],
partitioning: Optional[Partitioning] = None,
ignore_missing_paths: bool = False,
) -> Iterator[Tuple[str, int]]:
"""Expands all paths into concrete file paths by walking directories.
Also returns a sidecar of file sizes.
The input paths must be normalized for compatibility with the input
filesystem prior to invocation.
Args:
paths: A list of file and/or directory paths compatible with the
given filesystem.
filesystem: The filesystem implementation that should be used for
expanding all paths and reading their files.
partitioning: Partitioning describing how files under directories are
organized into partitions. If ``None``, paths are not interpreted as
partitioned.
ignore_missing_paths: If True, ignores any file paths in ``paths`` that
are not found. Defaults to False.
Returns:
An iterator of `(file_path, file_size)` pairs. None may be returned for the
file size if it is either unknown or will be fetched later by
`_get_block_metadata()`, but the length of
both lists must be equal.
"""
raise NotImplementedError
@DeveloperAPI
class DefaultFileMetadataProvider(BaseFileMetadataProvider):
"""Default metadata provider for
:class:`~ray.data.datasource.file_based_datasource.FileBasedDatasource`
implementations that reuse the base `prepare_read` method.
Calculates block size in bytes as the sum of its constituent file sizes,
and assumes a fixed number of rows per file.
"""
def _get_block_metadata(
self,
paths: List[str],
*,
rows_per_file: Optional[int],
file_sizes: List[Optional[int]],
) -> BlockMetadata:
if rows_per_file is None:
num_rows = None
else:
num_rows = len(paths) * rows_per_file
input_files = list(paths)
return BlockMetadata(
num_rows=num_rows,
size_bytes=None if None in file_sizes else int(sum(file_sizes)),
input_files=input_files,
exec_stats=None,
) # Exec stats filled in later.
def expand_paths(
self,
paths: List[str],
filesystem: "RetryingPyFileSystem",
partitioning: Optional[Partitioning] = None,
ignore_missing_paths: bool = False,
) -> Iterator[Tuple[str, int]]:
yield from _expand_paths(paths, filesystem, partitioning, ignore_missing_paths)
def _handle_read_os_error(error: OSError, paths: Union[str, List[str]]) -> str:
# NOTE: this is not comprehensive yet, and should be extended as more errors arise.
# NOTE: The latter patterns are raised in Arrow 10+, while the former is raised in
# Arrow < 10.
aws_error_pattern = (
r"^(?:(.*)AWS Error \[code \d+\]: No response body\.(.*))|"
r"(?:(.*)AWS Error UNKNOWN \(HTTP status 400\) during HeadObject operation: "
r"No response body\.(.*))|"
r"(?:(.*)AWS Error ACCESS_DENIED during HeadObject operation: No response "
r"body\.(.*))$"
)
if re.match(aws_error_pattern, str(error)):
# Specially handle AWS error when reading files, to give a clearer error
# message to avoid confusing users. The real issue is most likely that the AWS
# S3 file credentials have not been properly configured yet.
if isinstance(paths, str):
# Quote to highlight single file path in error message for better
# readability. List of file paths will be shown up as ['foo', 'boo'],
# so only quote single file path here.
paths = f'"{paths}"'
raise OSError(
(
f"Failing to read AWS S3 file(s): {paths}. "
"Please check that file exists and has properly configured access. "
"You can also run AWS CLI command to get more detailed error message "
"(e.g., aws s3 ls <file-name>). "
"See https://awscli.amazonaws.com/v2/documentation/api/latest/reference/s3/index.html " # noqa
"and https://docs.ray.io/en/latest/data/creating-datasets.html#reading-from-remote-storage " # noqa
"for more information."
)
)
else:
raise error
def _list_files(
paths: List[str],
filesystem: "RetryingPyFileSystem",
*,
partition_filter: Optional[PathPartitionFilter],
file_extensions: Optional[List[str]],
) -> List[Tuple[str, int]]:
return list(
_list_files_internal(
paths,
filesystem,
partition_filter=partition_filter,
file_extensions=file_extensions,
)
)
def _list_files_internal(
paths: List[str],
filesystem: "RetryingPyFileSystem",
*,
partition_filter: Optional[PathPartitionFilter],
file_extensions: Optional[List[str]],
) -> Iterator[Tuple[str, int]]:
default_meta_provider = DefaultFileMetadataProvider()
for path, file_size in default_meta_provider.expand_paths(paths, filesystem):
# HACK: PyArrow's `ParquetDataset` errors if input paths contain non-parquet
# files. To avoid this, we expand the input paths with the default metadata
# provider and then apply the partition filter or file extensions.
if (
partition_filter
and not partition_filter.apply(path)
or not _has_file_extension(path, file_extensions)
):
continue
yield path, file_size
def _expand_paths(
paths: List[str],
filesystem: "RetryingPyFileSystem",
partitioning: Optional[Partitioning],
ignore_missing_paths: bool = False,
) -> Iterator[Tuple[str, int]]:
"""Get the file sizes for all provided file paths."""
from pyarrow.fs import LocalFileSystem
from ray.data.datasource.file_based_datasource import (
FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD,
)
from ray.data.datasource.path_util import _is_http_url, _unwrap_protocol
# We break down our processing paths into a few key cases:
# 1. If len(paths) < threshold, fetch the file info for the individual files/paths
# serially.
# 2. If all paths are contained under the same parent directory (or base directory,
# if using partitioning), fetch all file infos at this prefix and filter to the
# provided paths on the client; this should be a single file info request.
# 3. If more than threshold requests required, parallelize them via Ray tasks.
# 1. Small # of paths case.
is_local = isinstance(filesystem, LocalFileSystem)
if isinstance(filesystem, RetryingPyFileSystem):
is_local = isinstance(filesystem.unwrap(), LocalFileSystem)
if (
len(paths) < FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD
# Local file systems are very fast to hit.
or is_local
):
yield from _get_file_infos_serial(paths, filesystem, ignore_missing_paths)
else:
# 2. Common path prefix case.
# Get longest common path of all paths.
common_path = os.path.commonpath(paths)
# If parent directory (or base directory, if using partitioning) is common to
# all paths, fetch all file infos at that prefix and filter the response to the
# provided paths.
if not _is_http_url(common_path) and (
(
partitioning is not None
and common_path == _unwrap_protocol(partitioning.base_dir)
)
or all(str(pathlib.Path(path).parent) == common_path for path in paths)
):
yield from _get_file_infos_common_path_prefix(
paths, common_path, filesystem, ignore_missing_paths
)
# 3. Parallelization case.
else:
# Parallelize requests via Ray tasks.
yield from _get_file_infos_parallel(paths, filesystem, ignore_missing_paths)
def _get_file_infos_serial(
paths: List[str],
filesystem: "RetryingPyFileSystem",
ignore_missing_paths: bool = False,
) -> Iterator[Tuple[str, int]]:
for path in paths:
yield from _get_file_infos(path, filesystem, ignore_missing_paths)
def _get_file_infos_common_path_prefix(
paths: List[str],
common_path: str,
filesystem: "pyarrow.fs.FileSystem",
ignore_missing_paths: bool = False,
) -> Iterator[Tuple[str, int]]:
path_to_size = {path: None for path in paths}
for path, file_size in _get_file_infos(
common_path, filesystem, ignore_missing_paths
):
if path in path_to_size:
path_to_size[path] = file_size
# Check if all `paths` have file size metadata.
# If any of paths has no file size, fall back to get files metadata in parallel.
# This can happen when path is a directory, but not a file.
have_missing_path = False
for path in paths:
if path_to_size[path] is None:
logger.debug(
f"Finding path {path} not have file size metadata. "
"Fall back to get files metadata in parallel for all paths."
)
have_missing_path = True
break
if have_missing_path:
# Parallelize requests via Ray tasks.
yield from _get_file_infos_parallel(paths, filesystem, ignore_missing_paths)
else:
# Iterate over `paths` to yield each path in original order.
# NOTE: do not iterate over `path_to_size` because the dictionary skips
# duplicated path, while `paths` might contain duplicated path if one wants
# to read same file multiple times.
for path in paths:
yield path, path_to_size[path]
def _get_file_infos_parallel(
paths: List[str],
filesystem: "RetryingPyFileSystem",
ignore_missing_paths: bool = False,
) -> Iterator[Tuple[str, int]]:
from ray.data.datasource.file_based_datasource import (
PATHS_PER_FILE_SIZE_FETCH_TASK,
_unwrap_s3_serialization_workaround,
_wrap_s3_serialization_workaround,
)
logger.warning(
f"Expanding {len(paths)} path(s). This may be a HIGH LATENCY "
f"operation on some cloud storage services. Moving all the "
"paths to a common parent directory will lead to faster "
"metadata fetching."
)
# Capture the filesystem in the fetcher func closure, but wrap it in our
# serialization workaround to make sure that the pickle roundtrip works as expected.
filesystem = _wrap_s3_serialization_workaround(filesystem)
def _file_infos_fetcher(paths: List[str]) -> List[Tuple[str, int]]:
fs = _unwrap_s3_serialization_workaround(filesystem)
return list(
itertools.chain.from_iterable(
_get_file_infos(path, fs, ignore_missing_paths) for path in paths
)
)
yield from _fetch_metadata_parallel(
paths, _file_infos_fetcher, PATHS_PER_FILE_SIZE_FETCH_TASK
)
Uri = TypeVar("Uri")
Meta = TypeVar("Meta")
def _fetch_metadata_parallel(
uris: List[Uri],
fetch_func: Callable[[List[Uri]], List[Meta]],
desired_uris_per_task: int,
**ray_remote_args,
) -> Iterator[Meta]:
"""Fetch file metadata in parallel using Ray tasks."""
remote_fetch_func = cached_remote_fn(fetch_func)
ray_remote_args = merge_label_selector(
dict(ray_remote_args),
DataContext.get_current().execution_options.label_selector,
)
if ray_remote_args:
remote_fetch_func = remote_fetch_func.options(**ray_remote_args)
# Choose a parallelism that results in a # of metadata fetches per task that
# dominates the Ray task overhead while ensuring good parallelism.
# Always launch at least 2 parallel fetch tasks.
parallelism = max(len(uris) // desired_uris_per_task, 2)
metadata_fetch_bar = ProgressBar(
"Metadata Fetch Progress", total=parallelism, unit="task"
)
fetch_tasks = []
for uri_chunk in np.array_split(uris, parallelism):
if len(uri_chunk) == 0:
continue
fetch_tasks.append(remote_fetch_func.remote(uri_chunk))
results = metadata_fetch_bar.fetch_until_complete(fetch_tasks)
yield from itertools.chain.from_iterable(results)
def _get_file_infos(
path: str, filesystem: "RetryingPyFileSystem", ignore_missing_path: bool = False
) -> List[Tuple[str, int]]:
"""Get the file info for all files at or under the provided path."""
from pyarrow.fs import FileType
file_infos = []
try:
file_info = filesystem.get_file_info(path)
except OSError as e:
_handle_read_os_error(e, path)
if file_info.type == FileType.Directory:
for file_path, file_size in _expand_directory(path, filesystem):
file_infos.append((file_path, file_size))
elif file_info.type == FileType.File:
file_infos.append((path, file_info.size))
elif file_info.type == FileType.NotFound and ignore_missing_path:
pass
else:
raise FileNotFoundError(path)
return file_infos
def _expand_directory(
path: str,
filesystem: "RetryingPyFileSystem",
exclude_prefixes: Optional[List[str]] = None,
ignore_missing_path: bool = False,
) -> List[Tuple[str, int]]:
"""
Expand the provided directory path to a list of file paths.
Args:
path: The directory path to expand.
filesystem: The filesystem implementation that should be used for
reading these files.
exclude_prefixes: The file relative path prefixes that should be
excluded from the returned file set. Default excluded prefixes are
"." and "_".
ignore_missing_path: If True, returns an empty list when ``path`` does not
exist instead of raising.
Returns:
An iterator of (file_path, file_size) tuples.
"""
if exclude_prefixes is None:
exclude_prefixes = [".", "_"]
from pyarrow.fs import FileSelector
selector = FileSelector(path, recursive=True, allow_not_found=ignore_missing_path)
files = filesystem.get_file_info(selector)
base_path = selector.base_dir
out = []
for file_ in files:
if not file_.is_file:
continue
file_path = file_.path
if not file_path.startswith(base_path):
continue
relative = file_path[len(base_path) :].lstrip("/")
if any(relative.startswith(prefix) for prefix in exclude_prefixes):
continue
out.append((file_path, file_.size))
# We sort the paths to guarantee a stable order.
return sorted(out)
python/ray/data/datasource/filename_provider.py
+171
View File
@@ -0,0 +1,171 @@
from typing import Any, Dict, Optional, Tuple
from ray.data.block import Block
from ray.util.annotations import Deprecated, PublicAPI
def _split_base_and_ext(filename: str) -> Tuple[str, str]:
"""Split a filename into (base, extension) where extension includes the dot.
Returns (base, ext) where ext includes the leading dot (e.g., ".parquet"),
or is empty string if the filename has no extension.
This is the single source of truth for separating a task filename's base
from its extension. Used by both row-filename derivation and checkpoint
base-filename extraction — these MUST agree for prefix-trie recovery.
"""
if "." in filename:
base, ext = filename.rsplit(".", 1)
return base, f".{ext}"
return filename, ""
@PublicAPI(stability="alpha")
class FilenameProvider:
"""Generates filenames when you write a :class:`~ray.data.Dataset`.
Use this class to customize the filenames used when writing a Dataset.
Override :meth:`~FilenameProvider.get_filename_for_task` to customize filenames.
For row-based writes (e.g., :meth:`~ray.data.Dataset.write_images`), row filenames
are automatically derived by appending ``_{block_index:06}_{row_index:06}`` to the
task filename.
Example:
This snippet shows you how to customize filenames with a prefix. For example,
a file might be named ``images_abc123_000000.png``.
.. testcode::
import ray
from ray.data.datasource import FilenameProvider
class ImageFilenameProvider(FilenameProvider):
def __init__(self, prefix: str, file_format: str):
super().__init__(file_format=file_format)
self.prefix = prefix
def get_filename_for_task(self, write_uuid, task_index):
return f"{self.prefix}_{write_uuid}_{task_index:06}.{self.file_format}"
ds = ray.data.read_parquet("s3://anonymous@ray-example-data/images.parquet")
ds.write_images(
"/tmp/results",
column="image",
filename_provider=ImageFilenameProvider("images", "png")
)
""" # noqa: E501
def __init__(
self,
dataset_uuid: Optional[str] = None,
file_format: Optional[str] = None,
) -> None:
"""Create a FilenameProvider.
Args:
dataset_uuid: An optional UUID to include as a prefix in generated
filenames. If provided, filenames will be formatted as
``{dataset_uuid}_{write_uuid}_{task_index}``.
file_format: An optional file extension (without the leading dot)
to append to generated filenames, e.g. ``"parquet"`` or
``"csv"``.
"""
self.dataset_uuid = dataset_uuid
self.file_format = file_format
def get_filename_for_task(self, write_uuid: str, task_index: int) -> str:
"""Generate a filename for a write task.
Override this method to customize filenames when writing a Dataset.
.. note::
Filenames must be unique and deterministic for a given write UUID and
task index.
Args:
write_uuid: The UUID of the write operation.
task_index: The index of the write task.
Returns:
The generated filename string.
"""
file_id = f"{write_uuid}_{task_index:06}"
filename = ""
if self.dataset_uuid is not None:
filename += f"{self.dataset_uuid}_"
filename += file_id
if self.file_format is not None:
filename += f".{self.file_format}"
return filename
@Deprecated(
message="Use get_filename_for_task() instead. The block and block_index "
"parameters are unused in practice because datasinks merge all blocks into "
"one before writing. These parameters will be removed in a future release. "
"Do not depend on block content or block_index in your FilenameProvider "
"implementation - filenames must be deterministic from (write_uuid, task_index) "
"alone to ensure checkpointing correctness."
)
def get_filename_for_block(
self, block: Optional[Block], write_uuid: str, task_index: int, block_index: int
) -> str:
"""Generate a filename for a block of data.
.. note::
Filenames must be unique and deterministic for a given write UUID and
task index. Do NOT depend on block content or block_index.
Checkpointing requires predicting the output filename BEFORE writing
data. This enables 2-phase commit: if a write fails after creating the
file but before committing the checkpoint, recovery can use the
predicted filename to delete orphaned files and retry cleanly. If
filenames depend on block content, this prediction is impossible and
checkpointing cannot guarantee exactly-once semantics.
Args:
block: Deprecated, unused. Do not depend on block content.
write_uuid: The UUID of the write operation.
task_index: The index of the write task.
block_index: Deprecated, always 0. Do not depend on this value.
Returns:
The filename to use for the block.
"""
raise NotImplementedError
@Deprecated(
message="Implement get_filename_for_task() instead. Row filenames are "
"automatically derived by appending _{block_index:06}_{row_index:06} to the "
"task filename. All files from the same task must share the task filename as "
"a prefix so that uncommitted data files can be identified and cleaned up "
"during checkpoint recovery."
)
def get_filename_for_row(
self,
row: Dict[str, Any],
write_uuid: str,
task_index: int,
block_index: int,
row_index: int,
) -> str:
"""Generate a filename for a row.
.. deprecated::
Implement :meth:`get_filename_for_task` instead. Row filenames are
automatically derived by appending ``_{block_index:06}_{row_index:06}``
to the task filename.
Args:
row: The row that will be written to a file.
write_uuid: The UUID of the write operation.
task_index: The index of the write task.
block_index: The index of the block *within* the write task.
row_index: The index of the row *within* the block.
Returns:
The filename to use for the row.
"""
raise NotImplementedError
python/ray/data/datasource/partitioning.py
+609
View File
@@ -0,0 +1,609 @@
import logging
import posixpath
import urllib.parse
from dataclasses import dataclass
from enum import Enum
from typing import (
TYPE_CHECKING,
Any,
Callable,
Dict,
List,
Optional,
Type,
Union,
)
import numpy as np
from ray.util.annotations import DeveloperAPI, PublicAPI
if TYPE_CHECKING:
import pyarrow
from ray.data.expressions import Expr
PartitionDataType = Type[Union[int, float, str, bool]]
logger = logging.getLogger(__name__)
@DeveloperAPI
class PartitionStyle(str, Enum):
"""Supported dataset partition styles.
Inherits from `str` to simplify plain text serialization/deserialization.
Examples:
>>> # Serialize to JSON text.
>>> json.dumps(PartitionStyle.HIVE) # doctest: +SKIP
'"hive"'
>>> # Deserialize from JSON text.
>>> PartitionStyle(json.loads('"hive"')) # doctest: +SKIP
<PartitionStyle.HIVE: 'hive'>
"""
HIVE = "hive"
DIRECTORY = "dir"
@DeveloperAPI
@dataclass
class Partitioning:
"""Partition scheme used to describe path-based partitions.
Path-based partition formats embed all partition keys and values directly in
their dataset file paths.
For example, to read a dataset with
`Hive-style partitions <https://athena.guide/articles/hive-style-partitioning>`_:
>>> import ray
>>> from ray.data.datasource.partitioning import Partitioning
>>> ds = ray.data.read_csv(
... "s3://anonymous@ray-example-data/iris.csv",
... partitioning=Partitioning("hive"),
... )
Instead, if your files are arranged in a directory structure such as:
.. code::
root/dog/dog_0.jpeg
root/dog/dog_1.jpeg
...
root/cat/cat_0.jpeg
root/cat/cat_1.jpeg
...
Then you can use directory-based partitioning:
>>> import ray
>>> from ray.data.datasource.partitioning import Partitioning
>>> root = "s3://anonymous@air-example-data/cifar-10/images"
>>> partitioning = Partitioning("dir", field_names=["class"], base_dir=root)
>>> ds = ray.data.read_images(root, partitioning=partitioning)
"""
#: The partition style - may be either HIVE or DIRECTORY.
style: PartitionStyle
#: "/"-delimited base directory that all partitioned paths should
#: exist under (exclusive). File paths either outside of, or at the first
#: level of, this directory will be considered unpartitioned. Specify
#: `None` or an empty string to search for partitions in all file path
#: directories.
base_dir: Optional[str] = None
#: The partition key field names (i.e. column names for tabular
#: datasets). When non-empty, the order and length of partition key
#: field names must match the order and length of partition values.
#: Required when parsing DIRECTORY partitioned paths or generating
#: HIVE partitioned paths.
field_names: Optional[List[str]] = None
#: A dictionary that maps partition key names to their desired data type. If not
#: provided, the data type defaults to string.
field_types: Optional[Dict[str, PartitionDataType]] = None
#: Filesystem that will be used for partition path file I/O.
filesystem: Optional["pyarrow.fs.FileSystem"] = None
def __post_init__(self):
if self.base_dir is None:
self.base_dir = ""
if self.field_types is None:
self.field_types = {}
self._normalized_base_dir = None
self._resolved_filesystem = None
@property
def normalized_base_dir(self) -> str:
"""Returns the base directory normalized for compatibility with a filesystem."""
if self._normalized_base_dir is None:
self._normalize_base_dir()
return self._normalized_base_dir
@property
def resolved_filesystem(self) -> "pyarrow.fs.FileSystem":
"""Returns the filesystem resolved for compatibility with a base directory."""
if self._resolved_filesystem is None:
self._normalize_base_dir()
return self._resolved_filesystem
def to_pyarrow(self) -> "pyarrow.dataset.Partitioning":
"""Convert to a PyArrow dataset Partitioning.
Returns:
Equivalent ``pyarrow.dataset.Partitioning`` instance.
Raises:
ValueError: If the partition style is not supported.
"""
import pyarrow.dataset as pds
schema = _partition_field_types_to_pa_schema(
self.field_names or [],
self.field_types or {},
)
if self.style == PartitionStyle.HIVE:
return pds.HivePartitioning(schema)
elif self.style == PartitionStyle.DIRECTORY:
return pds.DirectoryPartitioning(schema)
else:
raise ValueError(f"Unsupported partition style: {self.style}")
def _normalize_base_dir(self):
"""Normalizes the partition base directory for compatibility with the
given filesystem.
This should be called once a filesystem has been resolved to ensure that this
base directory is correctly discovered at the root of all partitioned file
paths.
"""
from ray.data.datasource.path_util import _resolve_paths_and_filesystem
paths, self._resolved_filesystem = _resolve_paths_and_filesystem(
self.base_dir,
self.filesystem,
)
assert (
len(paths) == 1
), f"Expected 1 normalized base directory, but found {len(paths)}"
normalized_base_dir = paths[0]
if len(normalized_base_dir) and not normalized_base_dir.endswith("/"):
normalized_base_dir += "/"
self._normalized_base_dir = normalized_base_dir
@DeveloperAPI
class PathPartitionParser:
"""Partition parser 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 under the base directory
will be discovered based on `{key1}={value1}/{key2}={value2}` naming
conventions. Key/value pairs do not need to be presented in the same
order across all paths. Directory names nested under the base directory that
don't follow this naming condition will be considered unpartitioned. If a
partition filter is defined, then it will be called with an empty input
dictionary for each unpartitioned file.
For `DIRECTORY` Partitioning, all directories under the base directory will
be interpreted as partition values of the form `{value1}/{value2}`. An
accompanying ordered list of partition field names must also be provided,
where the order and length of all partition values must match the order and
length of field names. Files stored directly in the base directory will
be considered unpartitioned. If a partition filter is defined, then it will
be called with an empty input dictionary for each unpartitioned file. For
example, if the base directory is `"foo"`, then `"foo.csv"` and `"foo/bar.csv"`
would be considered unpartitioned files but `"foo/bar/baz.csv"` would be associated
with partition `"bar"`. If the base directory is undefined, then `"foo.csv"` would
be unpartitioned, `"foo/bar.csv"` would be associated with partition `"foo"`, and
"foo/bar/baz.csv" would be associated with partition `("foo", "bar")`.
"""
@staticmethod
def of(
style: PartitionStyle = PartitionStyle.HIVE,
base_dir: Optional[str] = None,
field_names: Optional[List[str]] = None,
field_types: Optional[Dict[str, PartitionDataType]] = None,
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
) -> "PathPartitionParser":
"""Creates a path-based partition parser using a flattened argument list.
Args:
style: The partition style - may be either HIVE or DIRECTORY.
base_dir: "/"-delimited base directory to start searching for partitions
(exclusive). File paths outside of this directory will be considered
unpartitioned. Specify `None` or an empty string to search for
partitions in all file path directories.
field_names: The partition key names. Required for DIRECTORY partitioning.
Optional for HIVE partitioning. When non-empty, the order and length of
partition key field names must match the order and length of partition
directories discovered. Partition key field names are not required to
exist in the dataset schema.
field_types: A dictionary that maps partition key names to their desired
data type. If not provided, the data type default to string.
filesystem: Filesystem that will be used for partition path file I/O.
Returns:
The new path-based partition parser.
"""
scheme = Partitioning(style, base_dir, field_names, field_types, filesystem)
return PathPartitionParser(scheme)
def __init__(self, partitioning: Partitioning):
"""Creates a path-based partition parser.
Args:
partitioning: The path-based partition scheme. The parser starts
searching for partitions from this scheme's base directory. File paths
outside the base directory will be considered unpartitioned. If the
base directory is `None` or an empty string then this will search for
partitions in all file path directories. Field names are required for
DIRECTORY partitioning, and optional for HIVE partitioning. When
non-empty, the order and length of partition key field names must match
the order and length of partition directories discovered.
"""
style = partitioning.style
field_names = partitioning.field_names
if style == PartitionStyle.DIRECTORY and not field_names:
raise ValueError(
"Directory partitioning requires a corresponding list of "
"partition key field names. Please retry your request with one "
"or more field names specified."
)
parsers = {
PartitionStyle.HIVE: self._parse_hive_path,
PartitionStyle.DIRECTORY: self._parse_dir_path,
}
self._parser_fn: Callable[[str], Dict[str, str]] = parsers.get(style)
if self._parser_fn is None:
raise ValueError(
f"Unsupported partition style: {style}. "
f"Supported styles: {parsers.keys()}"
)
self._scheme = partitioning
def __call__(self, path: str) -> Dict[str, str]:
"""Parses partition keys and values from a single file path.
Args:
path: Input file path to parse.
Returns:
Dictionary mapping directory partition keys to values from the input file
path. Returns an empty dictionary for unpartitioned files.
"""
dir_path = self._dir_path_trim_base(path)
if dir_path is None:
return {}
partitions: Dict[str, str] = self._parser_fn(dir_path)
for field, data_type in self._scheme.field_types.items():
partitions[field] = _cast_value(partitions[field], data_type)
return partitions
def evaluate_predicate_on_partition(self, path: str, predicate: "Expr") -> bool:
"""Evaluate a predicate expression against partition values from a path.
This method enables partition pruning by evaluating predicates that reference
partition columns against the partition values parsed from file paths.
Args:
path: File path to parse partition values from.
predicate: Expression that references partition columns.
Returns:
True if the partition satisfies the predicate (should read the file),
False if it doesn't (can skip the file for partition pruning).
"""
import pyarrow as pa
from ray.data._internal.planner.plan_expression.expression_evaluator import (
NativeExpressionEvaluator,
)
# Parse partition values from the file path
partition_values = self(path)
if not partition_values:
# Unpartitioned file - exclude it when filtering on partition columns
# If the predicate references partition columns and the file doesn't have
# partition values in its path, we can't determine if it matches
return False
try:
# Create a single-row table with partition values
partition_table = pa.table(
{col: [val] for col, val in partition_values.items()}
)
# Evaluate using Ray Data's native evaluator
evaluator = NativeExpressionEvaluator(partition_table)
result = evaluator.visit(predicate)
# Extract boolean result from array-like types.
# NOTE: We must use ``.as_py()`` for PyArrow scalars because
# ``bool(pa.BooleanScalar(False))`` returns ``True`` (it
# checks validity/not-null, not the boolean value).
if isinstance(result, (pa.Array, pa.ChunkedArray)):
assert (
len(result) == 1
), f"Result expected to be of length 1 (got {result})"
return bool(result[0].as_py())
if isinstance(result, np.ndarray):
assert (
len(result) == 1
), f"Result expected to be of length 1 (got {result})"
return bool(result[0])
# Import pandas here to avoid circular dependencies
import pandas as pd
if isinstance(result, pd.Series):
assert (
len(result) == 1
), f"Result expected to be of length 1 (got {result})"
return bool(result.iloc[0])
# Scalar result
if isinstance(result, pa.Scalar):
return bool(result.as_py())
return bool(result)
except Exception:
logger.debug(
"Failed to evaluate predicate on partition for path %s, "
"conservatively including file.",
path,
exc_info=True,
)
return True
@property
def scheme(self) -> Partitioning:
"""Returns the partitioning for this parser."""
return self._scheme
def _dir_path_trim_base(self, path: str) -> Optional[str]:
"""Trims the normalized base directory and returns the directory path.
Returns None if the path does not start with the normalized base directory.
Simply returns the directory path if the base directory is undefined.
"""
if not path.startswith(self._scheme.normalized_base_dir):
return None
path = path[len(self._scheme.normalized_base_dir) :]
return posixpath.dirname(path)
def _parse_hive_path(self, dir_path: str) -> Dict[str, str]:
"""Hive partition path parser.
Returns a dictionary mapping partition keys to values given a hive-style
partition path of the form "{key1}={value1}/{key2}={value2}/..." or an empty
dictionary for unpartitioned files.
"""
dirs = [d for d in dir_path.split("/") if d and (d.count("=") == 1)]
kv_pairs = [d.split("=") for d in dirs] if dirs else []
# NOTE: PyArrow URL-encodes partition values when writing to cloud storage. To
# ensure the values are consistent when you read them back, we need to
# URL-decode them. See https://github.com/apache/arrow/issues/34905.
kv_pairs = [[key, urllib.parse.unquote(value)] for key, value in kv_pairs]
field_names = self._scheme.field_names
if field_names and kv_pairs:
if len(kv_pairs) != len(field_names):
raise ValueError(
f"Expected {len(field_names)} partition value(s) but found "
f"{len(kv_pairs)}: {kv_pairs}."
)
for i, field_name in enumerate(field_names):
if kv_pairs[i][0] != field_name:
raise ValueError(
f"Expected partition key {field_name} but found "
f"{kv_pairs[i][0]}"
)
return dict(kv_pairs)
def _parse_dir_path(self, dir_path: str) -> Dict[str, str]:
"""Directory partition path parser.
Returns a dictionary mapping directory partition keys to values from a
partition path of the form "{value1}/{value2}/..." or an empty dictionary for
unpartitioned files.
Requires a corresponding ordered list of partition key field names to map the
correct key to each value.
"""
dirs = [d for d in dir_path.split("/") if d]
field_names = self._scheme.field_names
if dirs and len(dirs) != len(field_names):
raise ValueError(
f"Expected {len(field_names)} partition value(s) but found "
f"{len(dirs)}: {dirs}."
)
if not dirs:
return {}
return {
field: directory
for field, directory in zip(field_names, dirs)
if field is not None
}
@PublicAPI(stability="beta")
class PathPartitionFilter:
"""Partition filter for path-based partition formats.
Used to explicitly keep or reject files based on a custom filter function that
takes partition keys and values parsed from the file's path as input.
"""
@staticmethod
def of(
filter_fn: Callable[[Dict[str, str]], bool],
style: PartitionStyle = PartitionStyle.HIVE,
base_dir: Optional[str] = None,
field_names: Optional[List[str]] = None,
field_types: Optional[Dict[str, PartitionDataType]] = None,
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
) -> "PathPartitionFilter":
"""Creates a path-based partition filter using a flattened argument list.
Args:
filter_fn: Callback used to filter partitions. Takes a dictionary mapping
partition keys to values as input. Unpartitioned files are denoted with
an empty input dictionary. Returns `True` to read a file for that
partition or `False` to skip it. Partition keys and values are always
strings read from the filesystem path. For example, this removes all
unpartitioned files:
.. code:: python
lambda d: True if d else False
This raises an assertion error for any unpartitioned file found:
.. code:: python
def do_assert(val, msg):
assert val, msg
lambda d: do_assert(d, "Expected all files to be partitioned!")
And this only reads files from January, 2022 partitions:
.. code:: python
lambda d: d["month"] == "January" and d["year"] == "2022"
style: The partition style - may be either HIVE or DIRECTORY.
base_dir: "/"-delimited base directory to start searching for partitions
(exclusive). File paths outside of this directory will be considered
unpartitioned. Specify `None` or an empty string to search for
partitions in all file path directories.
field_names: The partition key names. Required for DIRECTORY partitioning.
Optional for HIVE partitioning. When non-empty, the order and length of
partition key field names must match the order and length of partition
directories discovered. Partition key field names are not required to
exist in the dataset schema.
field_types: A dictionary that maps partition key names to their desired
data type. If not provided, the data type defaults to string.
filesystem: Filesystem that will be used for partition path file I/O.
Returns:
The new path-based partition filter.
"""
scheme = Partitioning(style, base_dir, field_names, field_types, filesystem)
path_partition_parser = PathPartitionParser(scheme)
return PathPartitionFilter(path_partition_parser, filter_fn)
def __init__(
self,
path_partition_parser: PathPartitionParser,
filter_fn: Callable[[Dict[str, str]], bool],
):
"""Creates a new path-based partition filter based on a parser.
Args:
path_partition_parser: The path-based partition parser.
filter_fn: Callback used to filter partitions. Takes a dictionary mapping
partition keys to values as input. Unpartitioned files are denoted with
an empty input dictionary. Returns `True` to read a file for that
partition or `False` to skip it. Partition keys and values are always
strings read from the filesystem path. For example, this removes all
unpartitioned files:
``lambda d: True if d else False``
This raises an assertion error for any unpartitioned file found:
``lambda d: assert d, "Expected all files to be partitioned!"``
And this only reads files from January, 2022 partitions:
``lambda d: d["month"] == "January" and d["year"] == "2022"``
"""
self._parser = path_partition_parser
self._filter_fn = filter_fn
def __call__(self, paths: List[str]) -> List[str]:
"""Returns all paths that pass this partition scheme's partition filter.
If no partition filter is set, then returns all input paths. If a base
directory is set, then only paths under this base directory will be parsed
for partitions. All paths outside of this base directory will automatically
be considered unpartitioned, and passed into the filter function as empty
dictionaries.
Also normalizes the partition base directory for compatibility with the
given filesystem before applying the filter.
Args:
paths: Paths to pass through the partition filter function. All
paths should be normalized for compatibility with the given
filesystem.
Returns:
List of paths that pass the partition filter, or all paths if no
partition filter is defined.
"""
filtered_paths = paths
if self._filter_fn is not None:
filtered_paths = [path for path in paths if self.apply(path)]
return filtered_paths
def apply(self, path: str) -> bool:
return self._filter_fn(self._parser(path))
@property
def parser(self) -> PathPartitionParser:
"""Returns the path partition parser for this filter."""
return self._parser
def _partition_field_types_to_pa_schema(
field_names: List[str],
field_types: Dict[str, PartitionDataType],
) -> "pyarrow.Schema":
"""Build a PyArrow schema from partition field names and Python types.
Args:
field_names: Ordered partition key names.
field_types: Mapping from field name to Python type. Fields not
present in the map default to ``str`` (``pa.string()``).
Returns:
A ``pyarrow.Schema`` with one field per partition key.
"""
import pyarrow as pa
type_map = {
int: pa.int64(),
float: pa.float64(),
bool: pa.bool_(),
str: pa.string(),
}
fields = []
for name in field_names:
py_type: PartitionDataType = field_types.get(name, str)
pa_type = type_map.get(py_type, pa.string())
fields.append(pa.field(name, pa_type))
return pa.schema(fields)
def _cast_value(value: str, data_type: PartitionDataType) -> Any:
if data_type is int:
return int(value)
elif data_type is float:
return float(value)
elif data_type is bool:
return value.lower() == "true"
else:
return value
python/ray/data/datasource/path_util.py
+492
View File
@@ -0,0 +1,492 @@
import logging
import pathlib
import sys
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
from urllib.parse import quote, unquote, urlparse
from ray.data._internal.util import (
RetryingPyFileSystem,
_normalize_paths_to_strings,
_resolve_custom_scheme,
)
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
import fsspec.spec
import pyarrow
def _get_fsspec_http_filesystem() -> "pyarrow.fs.PyFileSystem":
"""Get fsspec HTTPFileSystem wrapped in PyArrow PyFileSystem.
Returns:
PyFileSystem wrapping fsspec HTTPFileSystem.
Raises:
ImportError: If fsspec is not installed.
"""
try:
import fsspec # noqa: F401
from fsspec.implementations.http import HTTPFileSystem
except ModuleNotFoundError:
raise ImportError("Please install fsspec to read files from HTTP.") from None
from pyarrow.fs import FSSpecHandler, PyFileSystem
return PyFileSystem(FSSpecHandler(HTTPFileSystem()))
def _validate_and_wrap_filesystem(
filesystem: Optional[
Union["pyarrow.fs.FileSystem", "fsspec.spec.AbstractFileSystem"]
],
) -> Optional["pyarrow.fs.FileSystem"]:
"""Validate filesystem and wrap fsspec filesystems in PyArrow.
Args:
filesystem: Filesystem to validate and potentially wrap. Can be None,
a pyarrow.fs.FileSystem, or an fsspec.spec.AbstractFileSystem.
Returns:
None if filesystem is None, otherwise a pyarrow.fs.FileSystem
(either the original if already PyArrow, or wrapped if fsspec).
Raises:
TypeError: If filesystem is not None and not a valid pyarrow or fsspec filesystem.
"""
if filesystem is None:
return None
from pyarrow.fs import FileSystem
if isinstance(filesystem, FileSystem):
return filesystem
try:
import fsspec # noqa: F401
except ModuleNotFoundError:
raise TypeError("fsspec is not installed") from None
if not isinstance(filesystem, fsspec.spec.AbstractFileSystem):
raise TypeError(
f"Filesystem must conform to pyarrow.fs.FileSystem or "
f"fsspec.spec.AbstractFileSystem, got: {type(filesystem).__name__}"
)
from pyarrow.fs import FSSpecHandler, PyFileSystem
return PyFileSystem(FSSpecHandler(filesystem))
def _try_resolve_with_encoding(
path: str,
filesystem: Optional["pyarrow.fs.FileSystem"],
) -> Tuple["pyarrow.fs.FileSystem", str]:
"""Try resolving a path with URL encoding for special characters.
This handles paths with special characters like ';', '?', '#' that
may cause URI parsing errors.
Args:
path: The path to resolve.
filesystem: Optional filesystem to validate against.
Returns:
Tuple of (resolved_filesystem, resolved_path).
"""
from pyarrow.fs import _resolve_filesystem_and_path
encoded_path = quote(path, safe="/:", errors="ignore")
resolved_filesystem, resolved_path = _resolve_filesystem_and_path(
encoded_path, filesystem
)
return resolved_filesystem, unquote(resolved_path, errors="ignore")
def _has_file_extension(path: str, extensions: Optional[List[str]]) -> bool:
"""Check if a path has a file extension in the provided list.
Examples:
>>> _has_file_extension("foo.csv", ["csv"])
True
>>> _has_file_extension("foo.CSV", ["csv"])
True
>>> _has_file_extension("foo.CSV", [".csv"])
True
>>> _has_file_extension("foo.csv", ["json", "jsonl"])
False
>>> _has_file_extension("foo.csv", None)
True
Args:
path: The path to check.
extensions: A list of extensions to check against. If `None`, any extension is
considered valid.
Returns:
``True`` if ``path`` ends with one of the provided extensions (or
``extensions`` is ``None``), otherwise ``False``.
"""
assert extensions is None or isinstance(extensions, list), type(extensions)
if extensions is None:
return True
# If the user-specified extensions don't contain a leading dot, we add it here
extensions = [
f".{ext.lower()}" if not ext.startswith(".") else ext.lower()
for ext in extensions
]
# Ignore query components when checking extensions (for example,
# versioned object-store paths like `...parquet?versionId=...`).
# Keep `#` untouched because it can be part of object keys.
parsed_path = path.split("?", 1)[0]
return any(parsed_path.lower().endswith(ext) for ext in extensions)
# Mapping from URI schemes to compatible filesystem type_name values.
# Used to validate that a cached filesystem is compatible with a given URI scheme
# before attempting to use it, avoiding silent failures from PyArrow when the
# wrong filesystem type is passed to _resolve_filesystem_and_path.
_SCHEME_TO_FS_TYPE_NAMES = {
"": ("local",), # No scheme = local filesystem
"file": ("local",), # file:// = local filesystem
"s3": ("s3",), # s3:// = S3 filesystem
"s3a": ("s3",), # s3a:// = S3 filesystem (Hadoop compat)
"gs": ("gcs",), # gs:// = GCS filesystem
"gcs": ("gcs",), # gcs:// = GCS filesystem
"hdfs": ("hdfs",), # hdfs:// = Hadoop filesystem
"viewfs": ("hdfs",), # viewfs:// = Hadoop filesystem
"abfs": ("abfs",), # abfs:// = Azure Blob FileSystem
"abfss": ("abfs",), # abfss:// = Azure Blob FileSystem (TLS)
"http": ("py",), # http:// = fsspec HTTP (wrapped in PyFileSystem)
"https": ("py",), # https:// = fsspec HTTP (wrapped in PyFileSystem)
}
def _is_filesystem_compatible_with_scheme(
filesystem: "pyarrow.fs.FileSystem",
scheme: str,
) -> bool:
"""Check if a filesystem is compatible with a URI scheme.
Uses PyArrow's `type_name` property for reliable filesystem type detection.
This prevents silently using the wrong filesystem for a URI, which can result
in malformed paths or incorrect behavior.
Args:
filesystem: The PyArrow filesystem to check.
scheme: The URI scheme (e.g., 's3', 'gs', 'http', 'file', '').
Returns:
True if the filesystem can handle the scheme, False otherwise.
"""
# Get expected type names for this scheme
expected_types = _SCHEME_TO_FS_TYPE_NAMES.get(scheme.lower())
if expected_types is None:
# Unknown scheme (e.g., abfs://, az://, custom protocols) - trust user's filesystem
# This preserves backward compatibility for custom filesystems
return True
# Unwrap RetryingPyFileSystem to get the underlying filesystem's type
from ray.data._internal.util import RetryingPyFileSystem
unwrapped = (
filesystem.unwrap()
if isinstance(filesystem, RetryingPyFileSystem)
else filesystem
)
# Get the actual filesystem type
fs_type = unwrapped.type_name
# For PyFileSystem (fsspec wrappers), check the inner fsspec protocol
# rather than relying on type_name alone, since all fsspec wrappers
# share type_name "py" regardless of the underlying protocol.
if fs_type in ("py", "RetryingPyFileSystem") or fs_type.startswith("py::"):
from pyarrow.fs import FSSpecHandler, PyFileSystem
actual_fs = filesystem
if isinstance(actual_fs, RetryingPyFileSystem):
actual_fs = actual_fs.unwrap()
# After unwrapping, the inner filesystem may be a native PyArrow
# filesystem (e.g., S3FileSystem) rather than a PyFileSystem wrapper.
# Fall back to direct type_name matching in that case.
if not isinstance(actual_fs, PyFileSystem):
return actual_fs.type_name in expected_types
if isinstance(actual_fs.handler, FSSpecHandler):
inner_fs = actual_fs.handler.fs
protocol = getattr(inner_fs, "protocol", None)
if protocol is not None:
if isinstance(protocol, str):
protocol = (protocol,)
# Match scheme against fsspec protocol(s)
if scheme in protocol:
return True
# For bare paths (empty scheme), trust user-provided filesystem
if scheme == "":
return True
# Fallback: check HTTP
if scheme in ("http", "https"):
return _is_http_filesystem(filesystem)
return False
# Direct match for native PyArrow filesystems (s3, gcs, local, hdfs, etc.)
return fs_type in expected_types
def _resolve_single_path_with_fallback(
path: str,
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
) -> Tuple["pyarrow.fs.FileSystem", str]:
"""Resolve a single path with filesystem, with fallback to re-resolution on error.
This is a helper for lazy filesystem resolution. If a filesystem is provided,
it first validates that the filesystem type is compatible with the URI scheme,
then attempts to resolve the path. If the filesystem is incompatible or
resolution fails, it re-resolves without the cached filesystem.
Args:
path: A single file/directory path.
filesystem: Optional cached filesystem from previous resolution.
Returns:
Tuple of (resolved_filesystem, resolved_path).
Raises:
ValueError: If path resolution fails.
ImportError: If required dependencies are missing.
"""
import pyarrow as pa
from pyarrow.fs import _resolve_filesystem_and_path
path = _resolve_custom_scheme(path)
# Validate/wrap filesystem if needed
try:
filesystem = _validate_and_wrap_filesystem(filesystem)
except TypeError as e:
raise ValueError(f"Invalid filesystem provided: {e}") from e
# Parse scheme to validate filesystem compatibility
parsed = urlparse(path, allow_fragments=False)
scheme = parsed.scheme.lower() if parsed.scheme else ""
# Check HTTP scheme FIRST - PyArrow doesn't support HTTP/HTTPS natively
if scheme in ("http", "https"):
# If we have a compatible cached HTTP filesystem, use it
if filesystem is not None and _is_filesystem_compatible_with_scheme(
filesystem, scheme
):
return filesystem, path
# Otherwise create a new HTTP filesystem
try:
resolved_filesystem = _get_fsspec_http_filesystem()
resolved_path = path
return resolved_filesystem, resolved_path
except ImportError as import_error:
raise ImportError(
f"Cannot resolve HTTP path '{path}': {import_error}"
) from import_error
# Try with provided filesystem only if scheme is compatible (fast path for cached FS)
if filesystem is not None and _is_filesystem_compatible_with_scheme(
filesystem, scheme
):
try:
_, resolved_path = _resolve_filesystem_and_path(path, filesystem)
# Return the wrapped filesystem we passed in.
return filesystem, resolved_path
except Exception:
# Fall through to full resolution without cached filesystem
pass
# Full resolution without cached filesystem
try:
resolved_filesystem, resolved_path = _resolve_filesystem_and_path(path, None)
except (pa.lib.ArrowInvalid, ValueError) as original_error:
# Try URL encoding for paths with special characters that may cause parsing issues
try:
resolved_filesystem, resolved_path = _try_resolve_with_encoding(path, None)
except (pa.lib.ArrowInvalid, ValueError, TypeError) as encoding_error:
# If encoding doesn't help, raise with both errors for full context
raise ValueError(
f"Failed to resolve path '{path}'. Initial error: {original_error}. "
f"URL encoding fallback also failed: {encoding_error}"
) from original_error
except TypeError as e:
raise ValueError(f"The path: '{path}' has an invalid type {e}") from e
return resolved_filesystem, resolved_path
def _resolve_paths_and_filesystem(
paths: Union[str, List[str]],
filesystem: Optional["pyarrow.fs.FileSystem"] = None,
) -> Tuple[List[str], "pyarrow.fs.FileSystem"]:
"""
Resolves and normalizes all provided paths, infers a filesystem from the
paths and assumes that all paths use the same filesystem.
Args:
paths: A single file/directory path or a list of file/directory paths.
A list of paths can contain both files and directories.
filesystem: The filesystem implementation that should be used for
reading these files. If None, a filesystem will be inferred. If not
None, the provided filesystem will still be validated against all
filesystems inferred from the provided paths to ensure
compatibility.
Returns:
A pair ``(resolved_paths, filesystem)``. *resolved_paths* lists the
normalized paths for each input path that resolved successfully, in
order.
If *filesystem* was ``None``, the returned *filesystem* is set from
``resolved_filesystem`` on the first successful path and is left
unchanged on later iterations whenever it is already non-``None``.
If *filesystem* was not ``None``, the returned value is always that
same validated instance, even when ``_resolve_single_path_with_fallback``
inferred a different filesystem for a given path. Callers should pass
``None`` or a filesystem compatible with the path URIs so returned paths
and filesystem stay consistent.
All paths are assumed to use one storage backend; mixing unrelated URI
schemes in a single call is unsupported and may fail when reading.
"""
paths = _normalize_paths_to_strings(paths)
# Validate/wrap filesystem upfront so we return a proper PyArrow filesystem
filesystem = _validate_and_wrap_filesystem(filesystem)
resolved_paths = []
for path in paths:
try:
resolved_filesystem, resolved_path = _resolve_single_path_with_fallback(
path, filesystem
)
except (ValueError, ImportError) as e:
logger.warning(f"Failed to resolve path '{path}': {e}, skipping")
continue
if filesystem is None:
filesystem = resolved_filesystem
# If the PyArrow filesystem is handled by a fsspec HTTPFileSystem, the protocol/
# scheme of paths should not be unwrapped/removed, because HTTPFileSystem
# expects full file paths including protocol/scheme. This is different behavior
# compared to other file system implementation in pyarrow.fs.FileSystem.
if not _is_http_filesystem(resolved_filesystem):
resolved_path = _unwrap_protocol(resolved_path)
resolved_path = resolved_filesystem.normalize_path(resolved_path)
resolved_paths.append(resolved_path)
return resolved_paths, filesystem
def _split_uri(uri: str):
"""Split a URI into (store_url, path) for use with obstore.
e.g. "s3://my-bucket/a/b/c.jpg" -> ("s3://my-bucket", "a/b/c.jpg")
"https://host.com/a/b?X-Amz-Signature=x" -> ("https://host.com", "a/b?X-Amz-Signature=x")
The query string is preserved so signed URLs (e.g. pre-signed S3 HTTPS)
reach obstore intact. Semicolons in object keys normally appear in
``parsed.path`` (not ``parsed.params``) for typical ``urlparse`` output.
Only the first leading ``/`` after the authority (as reported in
``parsed.path``) is removed. Extra leading slashes belong to the object
key (e.g. ``s3://bucket//abs/key`` -> key ``/abs/key``), so
``str.lstrip("/")`` is not used.
"""
parsed = urlparse(uri, allow_fragments=False)
store_url = f"{parsed.scheme}://{parsed.netloc}"
raw_path = parsed.path
path = raw_path[1:] if raw_path.startswith("/") else raw_path
if parsed.query:
path = f"{path}?{parsed.query}"
return store_url, path
def _is_http_filesystem(fs: "pyarrow.fs.FileSystem") -> bool:
"""Return whether ``fs`` is a PyFileSystem handled by a fsspec HTTPFileSystem."""
from pyarrow.fs import FSSpecHandler, PyFileSystem
# Try to import HTTPFileSystem
try:
from fsspec.implementations.http import HTTPFileSystem
except ModuleNotFoundError:
return False
if isinstance(fs, RetryingPyFileSystem):
fs = fs.unwrap()
if not isinstance(fs, PyFileSystem):
return False
return isinstance(fs.handler, FSSpecHandler) and isinstance(
fs.handler.fs, HTTPFileSystem
)
def _unwrap_protocol(path):
"""
Slice off any protocol prefixes on path.
"""
if sys.platform == "win32" and _is_local_windows_path(path):
# Represent as posix path such that downstream functions properly handle it.
# This is executed when 'file://' is NOT included in the path.
return pathlib.Path(path).as_posix()
parsed = urlparse(path, allow_fragments=False) # support '#' in path
params = ";" + parsed.params if parsed.params else "" # support ';' in path
query = "?" + parsed.query if parsed.query else "" # support '?' in path
netloc = parsed.netloc
if parsed.scheme == "s3" and "@" in parsed.netloc:
# If the path contains an @, it is assumed to be an anonymous
# credentialed path, and we need to strip off the credentials.
netloc = parsed.netloc.split("@")[-1]
parsed_path = parsed.path
# urlparse prepends the path with a '/'. This does not work on Windows
# so if this is the case strip the leading slash.
if (
sys.platform == "win32"
and not netloc
and len(parsed_path) >= 3
and parsed_path[0] == "/" # The problematic leading slash
and parsed_path[1].isalpha() # Ensure it is a drive letter.
and parsed_path[2:4] in (":", ":/")
):
parsed_path = parsed_path[1:]
return netloc + parsed_path + params + query
def _is_http_url(path) -> bool:
parsed = urlparse(path)
return parsed.scheme in ("http", "https")
def _is_local_windows_path(path: str) -> bool:
"""Determines if path is a Windows file-system location."""
if sys.platform != "win32":
return False
if len(path) >= 1 and path[0] == "\\":
return True
if (
len(path) >= 3
and path[1] == ":"
and (path[2] == "/" or path[2] == "\\")
and path[0].isalpha()
):
return True
return False
python/ray/data/datasource/util.py
+109
View File
@@ -0,0 +1,109 @@
from typing import Any, Dict, Iterable, Tuple
import ray
from ray.data.block import Block
def _iter_sliced_blocks(
blocks: Iterable[Block], per_task_row_limit: int
) -> Iterable[Block]:
"""Iterate over blocks, accumulating rows up to the per-task row limit."""
rows_read = 0
for block in blocks:
if rows_read >= per_task_row_limit:
break
from ray.data.block import BlockAccessor
accessor = BlockAccessor.for_block(block)
block_rows = accessor.num_rows()
if rows_read + block_rows <= per_task_row_limit:
yield block
rows_read += block_rows
else:
# Slice the block to meet the limit exactly
remaining_rows = per_task_row_limit - rows_read
sliced_block = accessor.slice(0, remaining_rows, copy=True)
yield sliced_block
break
def _validate_head_node_resources_for_local_scheduling(
ray_remote_args: Dict[str, Any],
*,
op_description: str,
default_num_cpus: int = 1,
default_num_gpus: int = 0,
default_memory: int = 0,
) -> None:
"""Ensure the head node has enough resources before pinning work there.
Local paths (``local://``) and other driver-local I/O schedule tasks on the
head node via ``NodeAffinitySchedulingStrategy``. If the head node was
intentionally started with zero logical resources (a common practice to
avoid OOMs), those tasks become unschedulable. Detect this upfront and
raise a clear error with remediation steps.
"""
# Ray defaults to reserving 1 CPU per task when num_cpus isn't provided.
num_cpus = ray_remote_args.get("num_cpus", default_num_cpus)
num_gpus = ray_remote_args.get("num_gpus", default_num_gpus)
memory = ray_remote_args.get("memory", default_memory)
# Resource keys follow the Resources map of ray.nodes() (e.g., CPU, GPU, memory).
required_resources: Dict[str, float] = {}
required_resources["CPU"] = float(num_cpus)
required_resources["GPU"] = float(num_gpus)
required_resources["memory"] = float(memory)
# Include any additional custom resources requested.
custom_resources = ray_remote_args.get("resources", {})
for name, amount in custom_resources.items():
if amount is None:
continue
try:
amount = float(amount)
except (TypeError, ValueError) as err:
raise ValueError(f"Invalid resource amount for '{name}': {amount}") from err
required_resources[name] = amount
head_node = next(
(
node
for node in ray.nodes()
if node.get("Alive")
and "node:__internal_head__" in node.get("Resources", {})
),
None,
)
if not head_node:
# The head node metadata is unavailable (e.g., during shutdown). Fall back
# to the default behavior and let Ray surface its own error.
return
# Build a map of required vs available resources on the head node.
head_resources: Dict[str, float] = head_node.get("Resources", {})
# Map: resource name -> (required, available).
insufficient: Dict[str, Tuple[float, float]] = {}
for name, req in required_resources.items():
avail = head_resources.get(name, 0.0)
if avail < req:
insufficient[name] = (req, avail)
# If nothing is below the required amount, we are good to proceed.
if not insufficient:
return
details = "; ".join(
f"{name} required {req:g} but head has {avail:g}"
for name, (req, avail) in insufficient.items()
)
raise ValueError(
f"{op_description} must run on the head node (e.g., for local:// paths), "
f"but the head node doesn't have enough resources: {details}. "
"Add resources to the head node, switch to a shared filesystem instead "
"of local://, or set the resource requests on this operation to 0 "
"(for example, num_cpus=0) so it can run without head resources."
)