diff --git a/python/ray/data/_internal/datasource/table/SEQUENCE.md b/python/ray/data/_internal/datasource/table/SEQUENCE.md new file mode 100644 index 000000000000..49d23084ac05 --- /dev/null +++ b/python/ray/data/_internal/datasource/table/SEQUENCE.md @@ -0,0 +1,75 @@ +# Table-datasink write sequence + +This is the canonical sequence diagram for the table-format write +abstraction (`TableDatasink` + `TableAdapter`). An ASCII version is +reproduced inline in the `adapter.py` and `table_datasink.py` module +docstrings; this file is the higher-fidelity Mermaid source for the docs +site and GitHub-rendered Markdown. + +The diagram describes one driver-orchestrated write. Worker steps run +inside Ray write tasks, one task per data partition, in parallel. + +```mermaid +sequenceDiagram + autonumber + participant Ray as Ray Data engine + participant Sink as TableDatasink
(framework, generic) + participant Adapter as TableAdapter
(format-specific) + participant Workers as Ray write tasks
(N workers) + + Note over Sink: Driver: pre-write setup + Ray->>Sink: on_write_start(schema) + Sink->>Adapter: preflight(mode, partition_cols, declared_schema) + Sink->>Adapter: on_write_start(schema_from_first_bundle) + + Note over Workers: One task per partition, in parallel + par per task + Ray->>Sink: write(blocks, ctx) + Sink->>Adapter: start_task(ctx) + loop per Arrow table in blocks + Sink->>Adapter: write_block(arrow_table) + Adapter-->>Sink: (file_actions, emitted_schema, upsert_keys?) + end + Sink->>Adapter: finalize_task() + Sink->>Adapter: task_metadata() + Sink-->>Workers: TableWriteTaskResult + end + + Note over Sink: Driver: aggregate + commit + Ray->>Sink: on_write_complete(results) + Sink->>Adapter: gather_task_metadata(all_task_metadata) + Sink->>Adapter: reconcile_schema(unified_schema) + + alt mode == APPEND + Sink->>Adapter: commit_append(file_actions, unified_schema) + else mode == OVERWRITE + Sink->>Adapter: build_overwrite_predicate(overwrite_filter) + Adapter-->>Sink: delete_predicate + Sink->>Adapter: commit_overwrite(file_actions, unified_schema, delete_predicate) + else mode == UPSERT (adapter conforms to SupportsUpserts) + Sink->>Adapter: build_upsert_predicate(upsert_keys, join_cols) + Adapter-->>Sink: delete_predicate + Sink->>Adapter: commit_upsert(file_actions, unified_schema, delete_predicate) + end + + Note over Sink: On failure (any step above) + Ray->>Sink: on_write_failed(error) + Sink->>Adapter: on_failure(orphan_paths) +``` + +## Step glossary + +| Step | Owner | Purpose | +|---|---|---| +| `preflight` | adapter | Load table from catalog/log; validate mode legality + schema/partitions. | +| `on_write_start` | adapter | Optional pre-write hook fed the first bundle's schema. Iceberg evolves schema here; Delta no-ops and evolves at commit. | +| `start_task` | adapter | Per-task setup (e.g. open a per-task file writer). | +| `write_block` | adapter | Persist one Arrow table as object-store files. Returns per-file actions, the emitted schema, and the upsert-key projection (UPSERT mode only). | +| `finalize_task` | adapter | Flush any per-task buffer; return extra file actions and schemas. | +| `task_metadata` | adapter | Free-form per-task state the driver needs (e.g. Delta's per-write UUID). | +| `gather_task_metadata` | adapter | Driver-side merge of per-task metadata across workers. | +| `reconcile_schema` | adapter | Apply the worker-unified schema to the table state. | +| `build_overwrite_predicate` | adapter | Translate a user `overwrite_filter` into the format's predicate type. | +| `build_upsert_predicate` | adapter (`SupportsUpserts`) | Translate upsert-key table into the format's predicate type. | +| `commit_append` / `commit_overwrite` / `commit_upsert` | adapter | One atomic transaction per mode. | +| `on_failure` | adapter | Best-effort cleanup of files written by failed tasks. Never destroys committed data. | diff --git a/python/ray/data/_internal/datasource/table/__init__.py b/python/ray/data/_internal/datasource/table/__init__.py new file mode 100644 index 000000000000..177662ef3a31 --- /dev/null +++ b/python/ray/data/_internal/datasource/table/__init__.py @@ -0,0 +1,26 @@ +"""Generic table datasink framework. + +This module is the shared abstraction sitting under the Iceberg and Delta +datasinks. It owns the distributed write lifecycle and delegates every +format-specific concern (table loading, per-block file writing, transactional +commit) to a ``TableAdapter``. + +Supporting a new table format reduces to implementing a new adapter. +""" + +from .adapter import SupportsUpserts, TableAdapter +from .file_writer import DataFileWriter, ParquetFileWriter +from .modes import SaveMode, UpsertSemantics +from .result import TableWriteTaskResult +from .table_datasink import TableDatasink + +__all__ = [ + "DataFileWriter", + "ParquetFileWriter", + "SaveMode", + "SupportsUpserts", + "TableAdapter", + "TableDatasink", + "TableWriteTaskResult", + "UpsertSemantics", +] diff --git a/python/ray/data/_internal/datasource/table/adapter.py b/python/ray/data/_internal/datasource/table/adapter.py new file mode 100644 index 000000000000..4acbb1c2d866 --- /dev/null +++ b/python/ray/data/_internal/datasource/table/adapter.py @@ -0,0 +1,366 @@ +"""Format adapter protocol for table datasinks. + +A ``TableAdapter`` plugs format-specific behaviour (Iceberg, Delta, Hudi, …) +into the generic ``TableDatasink`` framework. The framework owns distributed +plumbing (lifecycle, schema unification, upsert-key concat, orphan cleanup); +the adapter owns table loading, per-block file writing, and transactional +commit. + +The contract is split in two: + +* :class:`TableAdapter` — APPEND + OVERWRITE baseline. Every adapter + implements this. +* :class:`SupportsUpserts` — opt-in :class:`typing.Protocol` for UPSERT + capability. Adapters that support UPSERT inherit (or duck-type-conform to) + this Protocol; the framework dispatches via ``isinstance(adapter, + SupportsUpserts)`` and refuses ``mode=UPSERT`` for adapters that don't. + +Why split? Adapters that don't support UPSERT (e.g. Delta in the current +build) shouldn't have to declare a fictional ``upsert_semantics`` value just +to satisfy a base class. The Protocol lets the type checker prove that +non-upsert adapters don't carry upsert state, and lets new adapters opt in +without inheritance gymnastics. + +Sequence diagram (ASCII; see ``./SEQUENCE.md`` for the Mermaid source):: + + ┌──────────────────────────────────────────────┐ + │ TableDatasink (framework, generic) │ + └──────────────────────────────────────────────┘ + │ + Ray Data ──── on_write_start ───────┤ + │ ┌─────────────────────────┐ + │──▶│ adapter.preflight │ (1) + │ │ adapter.on_write_start │ (2) + │ └─────────────────────────┘ + │ + (per write task, on workers) + │ ┌─────────────────────────┐ + Ray Data ──── write(blocks) ────────┤──▶│ adapter.start_task │ + │ │ adapter.write_block × N │ (3) + │ │ adapter.finalize_task │ + │ │ adapter.task_metadata │ + │ └─────────────────────────┘ + │ + Ray Data ──── on_write_complete ────┤ + │ ┌──────────────────────────────┐ + │──▶│ adapter.gather_task_metadata │ (4) + │ │ adapter.reconcile_schema │ (5) + │ │ │ + │ │ if APPEND: │ + │ │ adapter.commit_append │ (6a) + │ │ │ + │ │ if OVERWRITE: │ + │ │ adapter.build_overwrite_ │ (6b) + │ │ predicate │ + │ │ adapter.commit_overwrite │ + │ │ │ + │ │ if UPSERT (SupportsUpserts):│ + │ │ adapter.build_upsert_ │ (6c) + │ │ predicate │ + │ │ adapter.commit_upsert │ + │ └──────────────────────────────┘ + │ + Ray Data ──── on_write_failed ──────┤ ┌─────────────────────────┐ + │──▶│ adapter.on_failure │ (7) + │ └─────────────────────────┘ + +Steps: + (1) Load table from catalog/log; validate mode legality and + schema/partitions. + (2) Pre-write hook fed the first input bundle's schema (Iceberg evolves + schema here; Delta no-ops and evolves at commit time). + (3) Per Arrow table: write Parquet to storage; return + ``(file_actions, emitted_schema, optional upsert_keys)``. + (4) Driver receives per-task metadata dicts (e.g. Delta's per-write UUID). + (5) Apply the unified worker schema to the table state. + (6a/b/c) One mode-specific commit. All paths atomically apply + ``file_actions``. + (7) Best-effort cleanup of files that were written by tasks that + subsequently failed; never destroys committed data. +""" + +from abc import ABC, abstractmethod +from typing import ( + Any, + Dict, + Generic, + List, + Optional, + Protocol, + Set, + Tuple, + TypeVar, + runtime_checkable, +) + +import pyarrow as pa + +from .modes import SaveMode, UpsertSemantics +from ray.data._internal.execution.interfaces import TaskContext + +FileAction = TypeVar("FileAction") +DeletePredicate = TypeVar("DeletePredicate") + + +class TableAdapter(Generic[FileAction, DeletePredicate], ABC): + """Plug-in for one table format. APPEND + OVERWRITE baseline. + + Subclasses implement the abstract methods. The framework calls them in + the order shown in the module-level sequence diagram. Adapters that also + support UPSERT additionally implement the :class:`SupportsUpserts` + Protocol. + + Type parameters: + FileAction: Per-file metadata produced by ``write_block`` (e.g. + PyIceberg's ``DataFile``, deltalake's ``AddAction``). + DeletePredicate: Format-specific predicate type produced by + ``build_overwrite_predicate`` and consumed by ``commit_overwrite`` + (e.g. PyIceberg's ``BooleanExpression``, Delta's SQL ``str``). + """ + + # ------------------------------------------------------------------ + # Introspection — declared once per adapter class. + # ------------------------------------------------------------------ + @property + @abstractmethod + def supported_modes(self) -> Set[SaveMode]: + """SaveMode values this adapter supports. + + UPSERT may appear here only if the adapter also conforms to + :class:`SupportsUpserts`; the framework enforces this at + ``TableDatasink`` construction time. + """ + + # ------------------------------------------------------------------ + # Driver, before workers start. + # ------------------------------------------------------------------ + @abstractmethod + def preflight( + self, + mode: SaveMode, + partition_cols: List[str], + declared_schema: Optional[pa.Schema], + ) -> None: + """Load the underlying table and validate the requested write. + + Implementations typically: + * reach into the catalog / log to load the current table state, + * validate that the mode is legal against the table state + (e.g. UPSERT requires an existing table), + * validate partition columns / declared schema against the existing + table. + + Must raise a descriptive error on conflict. + """ + + def on_write_start( + self, schema_from_first_bundle: Optional[pa.Schema] = None + ) -> None: + """Optional pre-write hook fed the first input bundle's schema. + + Iceberg uses this hook to evolve the table schema before any files + land, avoiding name-mapping errors during writes. Delta no-ops here + and evolves at commit time. Default: no-op. + """ + return None + + # ------------------------------------------------------------------ + # Worker side, executed inside each Ray write task. + # ------------------------------------------------------------------ + def start_task(self, ctx: TaskContext) -> None: + """Called once per task before the first ``write_block``. + + Adapters that need per-task state (e.g. a file writer, a write UUID + pulled from ``ctx.kwargs``) should initialize it here. Default: + no-op. + """ + return None + + @abstractmethod + def write_block( + self, arrow_table: pa.Table + ) -> Tuple[List[FileAction], pa.Schema, Optional[pa.Table]]: + """Write a single Arrow table to the object store. + + Returns a 3-tuple of: + * the list of per-file actions produced (may be empty if the + adapter is buffering and didn't flush yet), + * the emitted schema for this block (used later by + ``reconcile_schema``), + * the projected upsert-key sub-table for this block, or ``None`` + if not in UPSERT mode. + """ + + def finalize_task(self) -> Tuple[List[FileAction], List[pa.Schema]]: + """Flush per-task buffers, if any. + + Returns extra ``(file_actions, schemas)`` produced when the buffer + is drained. Default: nothing buffered. + """ + return ([], []) + + def task_metadata(self) -> Dict[str, Any]: + """Adapter-defined free-form metadata to ship back to the driver. + + Called by the framework once per task after ``finalize_task`` and + embedded in the ``TableWriteTaskResult.task_metadata`` field. + Default: empty dict. + """ + return {} + + # ------------------------------------------------------------------ + # Driver side, after every worker finishes. + # ------------------------------------------------------------------ + def gather_task_metadata(self, task_metadata: List[Dict[str, Any]]) -> None: + """Receive the per-task metadata dicts produced by ``task_metadata``. + + Called once on the driver before ``reconcile_schema``. Adapters can + merge whatever they need (e.g. a shared write UUID, total row + counts). Default: no-op. + """ + return None + + def reconcile_schema(self, unified_schema: Optional[pa.Schema]) -> None: + """Driver-side schema reconciliation. + + Adapters that evolve the table schema at commit time (e.g. Delta) + use this hook. Adapters that already evolved in ``on_write_start`` + (e.g. Iceberg) typically just stash ``unified_schema`` for later + use. Default: no-op. + """ + return None + + @abstractmethod + def commit_append( + self, + file_actions: List[FileAction], + unified_schema: Optional[pa.Schema], + ) -> None: + """Apply an APPEND write atomically. + + Adapters may treat ``file_actions == []`` as a special "empty + commit" (e.g. to create an empty table); the framework always + invokes ``commit_append`` once, even when nothing was written. + """ + + @abstractmethod + def build_overwrite_predicate( + self, overwrite_filter: Optional[Any] + ) -> Optional[DeletePredicate]: + """Return a format-specific predicate covering rows OVERWRITE deletes. + + ``overwrite_filter`` is the user-supplied filter expression for + partial overwrites; ``None`` means "replace all rows". + Implementations return ``None`` for full overwrite (commit will drop + everything) or a predicate for partial overwrite. + """ + + @abstractmethod + def commit_overwrite( + self, + file_actions: List[FileAction], + unified_schema: Optional[pa.Schema], + delete_predicate: Optional[DeletePredicate], + ) -> None: + """Apply an OVERWRITE write atomically. + + Deletes rows matching ``delete_predicate`` (or everything when + ``delete_predicate is None``) and appends ``file_actions`` in the + same transaction. + """ + + # ------------------------------------------------------------------ + # File-action introspection. + # ------------------------------------------------------------------ + def path_for_action(self, action: FileAction) -> Optional[str]: + """Return the relative path the ``action`` represents, or ``None``. + + The framework calls this for two purposes: + * duplicate-file detection across tasks (``on_write_complete``), and + * orphan-file tracking for cleanup on failure (``write`` -> + ``on_write_failed``). + + The default reads ``action.path``. Adapters whose file-metadata type + names the path field differently **must** override this — e.g. + PyIceberg's ``DataFile`` exposes ``file_path`` (not ``path``), so the + Iceberg adapter overrides this to return ``action.file_path``. An + adapter that returns ``None`` here opts out of both duplicate + detection and orphan cleanup for its file actions. + """ + return getattr(action, "path", None) + + # ------------------------------------------------------------------ + # Failure handling. + # ------------------------------------------------------------------ + def on_failure(self, written_paths: List[str]) -> None: + """Best-effort orphan file cleanup. Default: no-op.""" + return None + + # ------------------------------------------------------------------ + # Optional introspection used by the framework for naming / scheduling. + # ------------------------------------------------------------------ + def get_name(self) -> str: + """Short human-readable name for write tasks. Default: class name.""" + return type(self).__name__ + + @property + def supports_distributed_writes(self) -> bool: + """If ``False``, the framework pins write tasks to the driver.""" + return True + + @property + def min_rows_per_write(self) -> Optional[int]: + """Target rows per write task; ``None`` lets Ray Data decide.""" + return None + + +@runtime_checkable +class SupportsUpserts(Protocol[FileAction, DeletePredicate]): + """Opt-in capability mixin for adapters that implement UPSERT. + + Adapters that support UPSERT either: + * inherit from this Protocol explicitly + (``class FooAdapter(TableAdapter[F, D], SupportsUpserts[F, D])``), + or + * structurally conform to the Protocol (define ``upsert_semantics``, + ``build_upsert_predicate``, ``commit_upsert``) — the + ``@runtime_checkable`` decorator lets ``isinstance(adapter, + SupportsUpserts)`` succeed in either case. + + The framework checks this Protocol at ``TableDatasink`` construction + time and refuses ``SaveMode.UPSERT`` for adapters that don't conform. + """ + + #: How this adapter implements UPSERT (e.g. COPY_ON_WRITE for Iceberg's + #: scan-merge approach, MERGE_ON_READ for Hudi-style overlays). + upsert_semantics: UpsertSemantics + + def build_upsert_predicate( + self, upsert_keys: Optional[pa.Table], join_cols: List[str] + ) -> Optional[DeletePredicate]: + """Return a delete predicate matching ``upsert_keys`` on ``join_cols``. + + ``upsert_keys`` is the framework-concatenated key table aggregated + across every worker, or ``None`` when no worker emitted any keys + (e.g. every input block was empty). Implementations must treat + ``None`` — and an empty table — as "no rows to match", returning + ``None`` so the write degrades to a pure insert. + + Implementations also typically filter out null keys before + constructing the predicate (SQL-style: NULL never matches), and + return ``None`` if nothing remains. + """ + ... + + def commit_upsert( + self, + file_actions: List[FileAction], + unified_schema: Optional[pa.Schema], + delete_predicate: Optional[DeletePredicate], + ) -> None: + """Apply an UPSERT write atomically. + + Deletes rows matching ``delete_predicate`` and appends + ``file_actions`` in the same transaction. + """ + ... diff --git a/python/ray/data/_internal/datasource/table/file_writer.py b/python/ray/data/_internal/datasource/table/file_writer.py new file mode 100644 index 000000000000..ea8ac0ab6c04 --- /dev/null +++ b/python/ray/data/_internal/datasource/table/file_writer.py @@ -0,0 +1,407 @@ +"""Format-agnostic file writer protocol. + +Adapters that emit physical data files (Parquet today; ORC / Avro / custom in +the future) delegate the on-disk writing to a ``DataFileWriter``. This keeps +the choice of file format independent from the choice of catalog format, which +matters because not every table uses Parquet (Hudi's MOR row-groups, custom +columnar stores, etc.). + +The Delta adapter composes a ``ParquetFileWriter`` whose implementation is the +mature, partition-aware buffered writer that previously lived in +``delta/writer.py``. The Iceberg adapter currently keeps using +``pyiceberg._dataframe_to_data_files`` because PyIceberg produces fully-formed +``DataFile`` records with format-specific metadata; migrating Iceberg to this +protocol is a follow-up. +""" + +from abc import ABC, abstractmethod +from typing import Any, Dict, Generic, List, Tuple, TypeVar + +import pyarrow as pa + +FileAction = TypeVar("FileAction") + + +class DataFileWriter(Generic[FileAction], ABC): + """Write one or more Arrow tables to physical files. + + Implementations are stateful (they may buffer per partition) and are + used inside a single Ray task. They are NOT expected to be pickled or + shared across tasks. + """ + + @abstractmethod + def add_table(self, table: pa.Table) -> List[FileAction]: + """Push an Arrow table through the writer. + + Returns any file actions that were flushed as a result of this call. + If the writer is buffering, the returned list may be empty. + """ + + @abstractmethod + def flush(self) -> List[FileAction]: + """Drain any remaining buffer and return its file actions.""" + + +# ---------------------------------------------------------------------- +# Parquet writer — generalised from ``delta/writer.py``. +# ---------------------------------------------------------------------- + +import logging # noqa: E402 +import time # noqa: E402 +import uuid # noqa: E402 +from collections import defaultdict # noqa: E402 +from typing import Callable, DefaultDict, Optional, Set # noqa: E402 + +import pyarrow.compute as pc # noqa: E402 +import pyarrow.parquet as pq # noqa: E402 + +logger = logging.getLogger(__name__) + +_MAX_PARTITIONS = 10_000 +_VALID_COMPRESSIONS = {"snappy", "gzip", "brotli", "zstd", "lz4", "none"} + +# Sentinel used to normalise NaN inside partition keys. Python's ``float('nan') +# != float('nan')`` makes raw NaN values unusable as dict keys (every NaN is a +# distinct key), which both fragments per-partition buffers across calls and +# corrupts grouping within a single table. ``_PartitionKey`` maps every NaN to +# this one sentinel for hashing/equality while preserving the original values +# (real NaN / None) for directory naming. +_NAN_SENTINEL = object() + + +class _PartitionKey: + """Hashable partition key with NaN-safe equality. + + Two keys are equal iff their values match element-wise, with NaN treated + as equal to NaN and None equal to None. ``values`` retains the original + Python values (real ``float('nan')`` / ``None``) so the adapter-supplied + ``path_builder`` still receives them verbatim for Hive-style directory + encoding. + """ + + __slots__ = ("values", "_norm") + + def __init__(self, values: Tuple): + self.values = values + self._norm = tuple( + _NAN_SENTINEL if (isinstance(v, float) and v != v) else v for v in values + ) + + def __hash__(self) -> int: + return hash(self._norm) + + def __eq__(self, other: object) -> bool: + return isinstance(other, _PartitionKey) and self._norm == other._norm + + def __repr__(self) -> str: + return f"_PartitionKey({self.values!r})" + + +class ParquetFileWriter(DataFileWriter[FileAction]): + """Partition-aware buffered Parquet writer. + + Generalises ``ray.data._internal.datasource.delta.writer.DeltaFileWriter`` + so it can be reused by any table adapter whose file format is Parquet. + The Delta-specific ``AddAction`` shape is supplied via the ``action_factory`` + callback so this writer stays format-agnostic. + + Args: + filesystem: PyArrow filesystem rooted at the table directory. + partition_cols: Hive-style partition columns (may be empty). + write_uuid: Stable identifier woven into filenames (8-char prefix). + compression: Parquet compression codec; one of + ``{"snappy", "gzip", "brotli", "zstd", "lz4", "none"}``. + write_statistics: Whether to write Parquet stats per file. + target_file_size_bytes: If set, the writer buffers per partition until + this many bytes accumulate before writing a file. ``None`` flushes + one file per ``add_table`` call. + path_builder: ``(partition_cols, partition_values_tuple) -> + (relative_dir, partition_values_dict)``. Builds the Hive-style + partition directory; supplied by the adapter so it can honour + format-specific encoding (NULL handling, NaN, URL-encoding). + action_factory: Callable that wraps file metadata into the adapter's + ``FileAction`` type. Receives keyword args + ``(path, size, partition_values, modification_time, stats)``. + stats_factory: Callable producing a stats blob (JSON string for + Delta; the adapter may pass ``lambda _t: None`` to disable). + written_files: Mutable set that the writer extends with every path it + writes; the adapter passes this in for orphan-cleanup tracking. + validate_path: Optional callable used to validate every output path + before write (e.g. reject ``..`` traversal). + retry: Optional callable wrapping the Parquet write with retry/backoff. + Signature: ``retry(func, description)``. Default: no retry. + """ + + def __init__( + self, + *, + filesystem: pa.fs.FileSystem, + partition_cols: List[str], + write_uuid: Optional[str], + compression: str = "snappy", + write_statistics: bool = True, + target_file_size_bytes: Optional[int] = None, + path_builder: Callable[ + [List[str], Tuple], Tuple[str, Dict[str, Optional[str]]] + ], + action_factory: Callable[..., FileAction], + stats_factory: Callable[[pa.Table], Any] = lambda _t: None, + written_files: Optional[Set[str]] = None, + validate_path: Optional[Callable[[str], None]] = None, + retry: Optional[Callable[[Callable[[], None], str], None]] = None, + ): + if compression not in _VALID_COMPRESSIONS: + raise ValueError( + f"Invalid compression '{compression}'. " + f"Supported: {sorted(_VALID_COMPRESSIONS)}" + ) + if target_file_size_bytes is not None and target_file_size_bytes <= 0: + raise ValueError("target_file_size_bytes must be > 0") + + self._filesystem = filesystem + self._partition_cols = partition_cols + self._write_uuid = write_uuid + self._compression = compression + self._write_statistics = write_statistics + self._target_file_size_bytes = target_file_size_bytes + self._path_builder = path_builder + self._action_factory = action_factory + self._stats_factory = stats_factory + self._written_files = written_files if written_files is not None else set() + self._validate_path = validate_path + self._retry = retry + + self._buffers: DefaultDict["_PartitionKey", List[pa.Table]] = defaultdict(list) + self._buffer_bytes: DefaultDict["_PartitionKey", int] = defaultdict(int) + self._file_seq = 0 + # The owning task index — adapters that care about task identity in + # filenames set this via ``set_task_idx`` before writing. + self._task_idx: int = 0 + + # ------------------------------------------------------------------ + + def set_task_idx(self, task_idx: int) -> None: + """Set the Ray task index used in generated filenames.""" + self._task_idx = task_idx + + @property + def written_files(self) -> Set[str]: + """Set of relative paths written so far; mutated by ``add_table``.""" + return self._written_files + + # ------------------------------------------------------------------ + # DataFileWriter interface. + # ------------------------------------------------------------------ + + def add_table(self, table: pa.Table) -> List[FileAction]: + if len(table) == 0: + return [] + + if not self._target_file_size_bytes: + # No buffering: one file per partition per add_table call. + self._file_seq += 1 + return self._write_table_immediate(table, block_idx=self._file_seq) + + # Buffered path. + if self._partition_cols: + parts = self._partition_table(table, self._partition_cols) + else: + parts = {_PartitionKey(()): table} + + actions: List[FileAction] = [] + for partition_key, partition_table in parts.items(): + self._buffers[partition_key].append(partition_table) + self._buffer_bytes[partition_key] += getattr(partition_table, "nbytes", 0) + if self._buffer_bytes[partition_key] >= self._target_file_size_bytes: + actions.extend(self._flush_partition(partition_key)) + return actions + + def flush(self) -> List[FileAction]: + actions: List[FileAction] = [] + for partition_key in list(self._buffers.keys()): + actions.extend(self._flush_partition(partition_key)) + return actions + + # ------------------------------------------------------------------ + # Internal helpers — kept structurally identical to the original + # ``DeltaFileWriter`` so behaviour matches commit-for-commit. + # ------------------------------------------------------------------ + + def _write_table_immediate( + self, table: pa.Table, block_idx: int + ) -> List[FileAction]: + if self._partition_cols: + parts = self._partition_table(table, self._partition_cols) + return [ + a + for a in ( + self._write_partition(t, k, block_idx) for k, t in parts.items() + ) + if a is not None + ] + a = self._write_partition(table, _PartitionKey(()), block_idx) + return [a] if a is not None else [] + + def _flush_partition(self, partition_key: "_PartitionKey") -> List[FileAction]: + tables = self._buffers.get(partition_key) + if not tables: + return [] + merged = pa.concat_tables(tables, promote_options="none") + self._buffers[partition_key].clear() + self._buffer_bytes[partition_key] = 0 + self._file_seq += 1 + action = self._write_partition(merged, partition_key, self._file_seq) + return [action] if action is not None else [] + + def _partition_table( + self, table: pa.Table, cols: List[str] + ) -> Dict["_PartitionKey", pa.Table]: + """Partition a table by ``cols`` via a single sort + boundary walk. + + Sorts the table once by the partition columns (O(N log N), vectorised) + and then makes one O(N) pass over the sorted key columns, slicing out a + contiguous run per distinct key. This replaces the previous + O(N * K) per-value filtering and the struct ``dictionary_encode`` + fallback. + + Keys are :class:`_PartitionKey` so NaN groups correctly (``NaN`` is + treated as equal to ``NaN``) both within this table and across buffered + ``add_table`` calls. ``key.values`` preserves the original Python + values (real ``NaN`` / ``None``) for the adapter's ``path_builder``. + """ + n = len(table) + if n == 0: + return {} + + # Sort once by all partition columns (nulls last is fine; grouping only + # needs equal keys to be contiguous). + sort_idx = pc.sort_indices( + table, sort_keys=[(c, "ascending") for c in cols] + ) + sorted_table = table.take(sort_idx) + + # Materialise just the key columns as Python lists for the boundary + # walk. Slicing the sorted table is zero-copy. + key_columns = [sorted_table.column(c).to_pylist() for c in cols] + + out: Dict[_PartitionKey, pa.Table] = {} + run_start = 0 + prev_key: Optional[_PartitionKey] = None + for i in range(n): + cur_values = tuple(key_columns[j][i] for j in range(len(cols))) + cur_key = _PartitionKey(cur_values) + if prev_key is None: + prev_key = cur_key + elif cur_key != prev_key: + out[prev_key] = sorted_table.slice(run_start, i - run_start) + if len(out) > _MAX_PARTITIONS: + raise ValueError( + f"Too many partition values ({len(out)}+). " + f"Max: {_MAX_PARTITIONS}" + ) + run_start = i + prev_key = cur_key + # Close the final run. + out[prev_key] = sorted_table.slice(run_start, n - run_start) + if len(out) > _MAX_PARTITIONS: + raise ValueError( + f"Too many partition values ({len(out)}). Max: {_MAX_PARTITIONS}" + ) + return out + + def _write_partition( + self, table: pa.Table, partition_key: "_PartitionKey", block_idx: int + ) -> Optional[FileAction]: + if len(table) == 0: + return None + + # Drop partition columns from the on-disk payload — both Iceberg and + # Delta encode partition values in the directory, not the file. + if self._partition_cols: + data_cols = [c for c in table.column_names if c not in self._partition_cols] + table = table.select(data_cols) + + filename = self._filename(block_idx) + # ``key.values`` carries the original Python values (real NaN / None) + # so the adapter's path_builder can apply its own encoding. + partition_path, partition_dict = self._path_builder( + self._partition_cols, partition_key.values + ) + rel_path = partition_path + filename + + if self._validate_path is not None: + self._validate_path(rel_path) + self._written_files.add(rel_path) + + size = self._write_parquet(table, rel_path) + stats = self._stats_factory(table) + + return self._action_factory( + path=rel_path, + size=size, + partition_values=partition_dict, + modification_time=int(time.time() * 1000), + stats=stats, + ) + + def _filename(self, block_idx: int) -> str: + uid = uuid.uuid4().hex[:16] + prefix = self._write_uuid or "00000000" + prefix = prefix[:8].ljust(8, "0") + return f"part-{prefix}-{self._task_idx:05d}-{block_idx:05d}-{uid}.parquet" + + def _write_parquet(self, table: pa.Table, rel_path: str) -> int: + parent = _safe_dirname(rel_path) + if parent: + try: + self._filesystem.create_dir(parent, recursive=True) + except Exception: + # Non-fatal: many filesystems (e.g. object stores) create + # parents implicitly on write, and recursive=True tolerates + # an existing dir. Log at debug so genuine FS errors + # (permissions, etc.) are still discoverable; the subsequent + # write will surface a hard failure if the path is unusable. + logger.debug( + "create_dir(%s) failed; relying on the write to validate " + "path accessibility.", + parent, + exc_info=True, + ) + + result: Dict[str, int] = {"size": 0} + + def _do_write() -> None: + pq.write_table( + table, + rel_path, + filesystem=self._filesystem, + compression=self._compression, + write_statistics=self._write_statistics, + ) + info = self._filesystem.get_file_info(rel_path) + if info.size == 0: + try: + self._filesystem.delete_file(rel_path) + except Exception: + pass + raise RuntimeError(f"Written file is empty: {rel_path}") + result["size"] = info.size + + if self._retry is not None: + self._retry(_do_write, f"write Parquet file '{rel_path}'") + else: + _do_write() + return result["size"] + + +def _safe_dirname(path: str) -> str: + """Mirror of ``delta.utils.safe_dirname`` to keep this module standalone.""" + import os + import posixpath + + if "://" in path: + scheme, rest = path.split("://", 1) + directory = posixpath.dirname(rest) + return f"{scheme}://{directory}" if directory else "" + return os.path.dirname(path) diff --git a/python/ray/data/_internal/datasource/table/modes.py b/python/ray/data/_internal/datasource/table/modes.py new file mode 100644 index 000000000000..ec77f9357975 --- /dev/null +++ b/python/ray/data/_internal/datasource/table/modes.py @@ -0,0 +1,29 @@ +"""Table-specific enums layered on top of the generic SaveMode. + +`SaveMode` is the generic, public-facing enum from ``ray.data._internal.savemode``. +The table layer re-exports it for convenience and adds ``UpsertSemantics``, +which a ``TableAdapter`` uses to declare whether it supports copy-on-write or +merge-on-read upserts. +""" + +from enum import Enum + +from ray.data._internal.savemode import SaveMode + +__all__ = ["SaveMode", "UpsertSemantics"] + + +class UpsertSemantics(str, Enum): + """How a table adapter implements UPSERT. + + COPY_ON_WRITE + Matched rows are physically deleted and the new data is appended. + Used by Delta Lake today and by Iceberg's default upsert path. + + MERGE_ON_READ + Matched rows are logically masked via delete-files / delete-vectors; + no physical rewrite. Future Iceberg/Hudi support. + """ + + COPY_ON_WRITE = "copy_on_write" + MERGE_ON_READ = "merge_on_read" diff --git a/python/ray/data/_internal/datasource/table/result.py b/python/ray/data/_internal/datasource/table/result.py new file mode 100644 index 000000000000..444042955ec7 --- /dev/null +++ b/python/ray/data/_internal/datasource/table/result.py @@ -0,0 +1,60 @@ +"""Per-task write result produced by every table adapter. + +A single Ray Data write task collects one of these from the framework and the +driver aggregates them in ``TableDatasink.on_write_complete``. + +``FileAction`` is the adapter-defined per-file metadata type (e.g. PyIceberg's +``DataFile`` or deltalake's ``AddAction``). The framework treats it as opaque, +with one exception: to support duplicate-file detection and orphan cleanup it +needs each action's path, which it obtains via ``TableAdapter.path_for_action`` +(not by introspecting a fixed attribute). Adapters whose metadata names the +path field differently override that method — e.g. Iceberg's ``DataFile`` uses +``file_path``. +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, Generic, List, Optional, TypeVar + +import pyarrow as pa + +FileAction = TypeVar("FileAction") + + +@dataclass +class TableWriteTaskResult(Generic[FileAction]): + """Result a worker returns to the driver after a write task completes. + + Attributes: + file_actions: One or more per-file metadata objects produced by the + adapter. Driver hands these back to ``adapter.commit``. + emitted_schemas: PyArrow schema of every non-empty Arrow table the + worker processed in this task. Driver uses these for type-promoted + schema unification across workers. + upsert_keys: Concatenated key-column table for UPSERT mode, or ``None``. + written_paths: Best-effort list of relative paths written by this task, + used by ``on_write_failed`` to clean up orphans. + task_id: Worker task index (for logging / debugging). + task_metadata: Adapter-defined free-form metadata produced by this + task. The framework concatenates these across tasks and forwards + the list to ``adapter.gather_task_metadata`` before ``commit``, + giving adapters a channel for worker→driver state such as + Delta's per-write app-transaction UUID. + """ + + file_actions: List[FileAction] = field(default_factory=list) + emitted_schemas: List[pa.Schema] = field(default_factory=list) + upsert_keys: Optional[pa.Table] = None + written_paths: List[str] = field(default_factory=list) + task_id: Optional[int] = None + task_metadata: Dict[str, Any] = field(default_factory=dict) + + @property + def data_files(self) -> List[FileAction]: + """Deprecated alias for ``file_actions``. + + Provided so code written against the pre-abstraction + ``IcebergWriteResult.data_files`` keeps working without modification. + Prefer ``file_actions`` going forward; this alias may be removed in + a future release. + """ + return self.file_actions diff --git a/python/ray/data/_internal/datasource/table/table_datasink.py b/python/ray/data/_internal/datasource/table/table_datasink.py new file mode 100644 index 000000000000..0cf414f05ca7 --- /dev/null +++ b/python/ray/data/_internal/datasource/table/table_datasink.py @@ -0,0 +1,324 @@ +"""Generic table datasink that drives a ``TableAdapter``. + +This is the "Table framework" participant in the design sequence diagram +(see :mod:`ray.data._internal.datasource.table.adapter` for the full +diagram, also reproduced as Mermaid in ``./SEQUENCE.md``). + +It owns the Ray Data write lifecycle (``on_write_start`` → workers' ``write`` +→ ``on_write_complete`` → ``on_write_failed``) and delegates every +format-specific decision to a ``TableAdapter``. APPEND and OVERWRITE go to +methods on ``TableAdapter`` itself; UPSERT goes to methods on +:class:`SupportsUpserts`, which the adapter must conform to (checked at +``__init__`` time via ``isinstance``). + +The orchestration is intentionally simple and identical for every adapter, +so that supporting a new table format reduces to implementing a new adapter. +""" + +import logging +from typing import Any, Dict, Generic, Iterable, List, Optional, TypeVar + +import pyarrow as pa + +from .adapter import SupportsUpserts, TableAdapter +from .modes import SaveMode +from .result import TableWriteTaskResult +from ray.data._internal.execution.interfaces import TaskContext +from ray.data.block import Block, BlockAccessor +from ray.data.datasource.datasink import Datasink, WriteResult + +logger = logging.getLogger(__name__) + +FileAction = TypeVar("FileAction") +DeletePredicate = TypeVar("DeletePredicate") + + +class TableDatasink( + Datasink[TableWriteTaskResult], + Generic[FileAction, DeletePredicate], +): + """Generic datasink for any table format with a pluggable adapter. + + Args: + adapter: ``TableAdapter`` providing the format-specific behaviour. + mode: One of ``SaveMode.{APPEND,OVERWRITE,UPSERT,ERROR,IGNORE}``. Must + be in ``adapter.supported_modes``. + partition_cols: Optional Hive-style partition columns. Forwarded to + ``adapter.preflight``. + declared_schema: Optional user-declared schema. Forwarded to + ``adapter.preflight``. + join_cols: Columns to match on in UPSERT mode (ignored otherwise). + overwrite_filter: Predicate for partial OVERWRITE (forwarded to + ``adapter.build_overwrite_predicate``). + name: Human-readable name override for write tasks. + """ + + def __init__( + self, + adapter: TableAdapter[FileAction, DeletePredicate], + mode: SaveMode, + *, + partition_cols: Optional[List[str]] = None, + declared_schema: Optional[pa.Schema] = None, + join_cols: Optional[List[str]] = None, + overwrite_filter: Optional[Any] = None, + name: Optional[str] = None, + ): + self._adapter = adapter + self._mode = self._coerce_mode(mode) + self._partition_cols = list(partition_cols or []) + self._declared_schema = declared_schema + self._join_cols = list(join_cols or []) + self._overwrite_filter = overwrite_filter + self._name_override = name + + self._validate_mode_against_adapter() + + # ------------------------------------------------------------------ + # Helpers. + # ------------------------------------------------------------------ + + @staticmethod + def _coerce_mode(mode: Any) -> SaveMode: + if isinstance(mode, SaveMode): + return mode + if isinstance(mode, str): + try: + return SaveMode(mode.lower()) + except ValueError as e: + raise ValueError( + f"Invalid mode '{mode}'. Supported: {[m.value for m in SaveMode]}" + ) from e + raise TypeError(f"Invalid mode type: {type(mode).__name__}") + + def _validate_mode_against_adapter(self) -> None: + supported = self._adapter.supported_modes + if self._mode not in supported: + raise ValueError( + f"{type(self._adapter).__name__} does not support mode " + f"{self._mode}. Supported: {sorted(m.value for m in supported)}" + ) + if self._mode == SaveMode.UPSERT and not isinstance( + self._adapter, SupportsUpserts + ): + raise ValueError( + f"{type(self._adapter).__name__} does not support UPSERT. " + "Adapters that support UPSERT must conform to the " + "SupportsUpserts Protocol (implement upsert_semantics, " + "build_upsert_predicate, commit_upsert)." + ) + + # ------------------------------------------------------------------ + # Datasink overrides — Ray Data write lifecycle. + # ------------------------------------------------------------------ + + def get_name(self) -> str: + if self._name_override: + return self._name_override + return self._adapter.get_name() + + @property + def supports_distributed_writes(self) -> bool: + return self._adapter.supports_distributed_writes + + @property + def min_rows_per_write(self) -> Optional[int]: + return self._adapter.min_rows_per_write + + def on_write_start(self, schema: Optional[pa.Schema] = None) -> None: + """Driver-side lifecycle: preflight then adapter pre-write hook. + + Matches steps 3 and 4 of the design sequence diagram. + """ + self._adapter.preflight( + mode=self._mode, + partition_cols=self._partition_cols, + declared_schema=self._declared_schema, + ) + self._adapter.on_write_start(schema_from_first_bundle=schema) + + def write( + self, blocks: Iterable[Block], ctx: TaskContext + ) -> TableWriteTaskResult[FileAction]: + """Worker-side lifecycle: per-Arrow-table ``write_block`` calls. + + Matches step 5 of the design sequence diagram. + """ + self._adapter.start_task(ctx) + + file_actions: List[FileAction] = [] + emitted_schemas: List[pa.Schema] = [] + key_chunks: List[pa.Table] = [] + written_paths: List[str] = [] + + try: + for block in blocks: + arrow_table = BlockAccessor.for_block(block).to_arrow() + if arrow_table.num_rows == 0: + continue + actions, emitted_schema, upsert_keys = self._adapter.write_block( + arrow_table + ) + if actions: + file_actions.extend(actions) + for action in actions: + path = self._adapter.path_for_action(action) + if isinstance(path, str): + written_paths.append(path) + if emitted_schema is not None: + emitted_schemas.append(emitted_schema) + if upsert_keys is not None: + key_chunks.append(upsert_keys) + + extra_actions, extra_schemas = self._adapter.finalize_task() + if extra_actions: + file_actions.extend(extra_actions) + for action in extra_actions: + path = self._adapter.path_for_action(action) + if isinstance(path, str): + written_paths.append(path) + if extra_schemas: + emitted_schemas.extend(extra_schemas) + + # Kept inside the try so that a failure in ``_concat_tables`` (e.g. + # incompatible upsert-key schemas) or in a buggy ``task_metadata`` + # override still attaches the orphan-path list to the exception, + # letting ``on_write_failed`` clean up files this task wrote. + upsert_keys = _concat_tables(key_chunks) + task_metadata = dict(self._adapter.task_metadata() or {}) + except Exception as e: + # Surface the orphan-path list to the driver via the exception so + # the framework can hand it to ``adapter.on_failure``. + existing = getattr(e, "_table_written_paths", None) or [] + e._table_written_paths = list(existing) + written_paths + raise + + return TableWriteTaskResult( + file_actions=file_actions, + emitted_schemas=emitted_schemas, + upsert_keys=upsert_keys, + written_paths=written_paths, + task_id=getattr(ctx, "task_idx", None), + task_metadata=task_metadata, + ) + + def on_write_complete( + self, write_result: WriteResult[TableWriteTaskResult[FileAction]] + ) -> None: + """Driver-side: aggregate, reconcile, commit. + + Matches steps 6, 7, 8 of the design sequence diagram. + """ + all_actions: List[FileAction] = [] + all_schemas: List[pa.Schema] = [] + all_key_chunks: List[pa.Table] = [] + all_task_metadata: List[Dict[str, Any]] = [] + seen_paths = set() + + for r in write_result.write_returns or []: + if r is None: + continue + for action in r.file_actions: + path = self._adapter.path_for_action(action) + if isinstance(path, str): + if path in seen_paths: + raise ValueError(f"Duplicate file paths detected: {path}") + seen_paths.add(path) + all_actions.append(action) + if r.emitted_schemas: + all_schemas.extend(r.emitted_schemas) + if r.upsert_keys is not None: + all_key_chunks.append(r.upsert_keys) + # One entry per task that ran, even when the metadata dict is + # empty (the base ``task_metadata()`` returns ``{}``). Guarding on + # truthiness here would drop those tasks and leave adapters that + # rely on "one entry per task" undercounting. + all_task_metadata.append(r.task_metadata) + + unified_schema = ( + _unify_schemas(all_schemas) if all_schemas else self._declared_schema + ) + upsert_keys = _concat_tables(all_key_chunks) + + # Hand the adapter every task's metadata before commit. + self._adapter.gather_task_metadata(all_task_metadata) + + # Step 7 — schema reconciliation. + self._adapter.reconcile_schema(unified_schema) + + # Step 8 — mode-specific commit. ``_validate_mode_against_adapter`` + # at __init__ time guarantees ``self._mode`` is supported and that + # the adapter conforms to SupportsUpserts for UPSERT, so the + # branches below are exhaustive for any legal mode. + # + # CREATE / ERROR / IGNORE share the APPEND commit path: by the time we + # reach commit, ``preflight`` has already enforced their precondition — + # CREATE and ERROR raise if the table existed, and IGNORE sets the + # adapter's skip flag if the table existed (so ``commit_append`` + # no-ops). When the table did NOT exist, all of them must still + # create-and-commit the written data, exactly like APPEND; routing + # them anywhere else (or to the ``else`` below) would silently drop + # those files as orphans after the workers already wrote them. + if self._mode in ( + SaveMode.APPEND, + SaveMode.CREATE, + SaveMode.ERROR, + SaveMode.IGNORE, + ): + self._adapter.commit_append(all_actions, unified_schema) + elif self._mode == SaveMode.OVERWRITE: + predicate = self._adapter.build_overwrite_predicate( + self._overwrite_filter + ) + self._adapter.commit_overwrite( + all_actions, unified_schema, predicate + ) + elif self._mode == SaveMode.UPSERT: + # _validate_mode_against_adapter guarantees SupportsUpserts. + predicate = self._adapter.build_upsert_predicate( + upsert_keys, self._join_cols + ) + self._adapter.commit_upsert( + all_actions, unified_schema, predicate + ) + else: + raise ValueError(f"Unsupported mode: {self._mode}") + + def on_write_failed(self, error: Exception) -> None: + """Driver-side: hand the worker's orphan-path list to the adapter.""" + paths = list(getattr(error, "_table_written_paths", None) or []) + if paths: + logger.warning( + "Table write failed; attempting cleanup of %d orphaned files.", + len(paths), + ) + try: + self._adapter.on_failure(paths) + except Exception as cleanup_error: # noqa: BLE001 + logger.warning( + "Adapter on_failure raised %s; ignoring to avoid masking the " + "primary error.", + cleanup_error, + ) + + +# ---------------------------------------------------------------------- +# Shared helpers, exposed for adapters that want to call them directly. +# ---------------------------------------------------------------------- + + +def _unify_schemas(schemas: List[pa.Schema]) -> Optional[pa.Schema]: + """Type-promoted ``pa.unify_schemas``; tolerant of older PyArrow versions.""" + if not schemas: + return None + from ray.data._internal.arrow_ops.transform_pyarrow import unify_schemas + + return unify_schemas(schemas, promote_types=True) + + +def _concat_tables(tables: List[pa.Table]) -> Optional[pa.Table]: + if not tables: + return None + from ray.data._internal.arrow_ops.transform_pyarrow import concat + + return concat(tables) diff --git a/python/ray/data/tests/datasource/test_table_datasink.py b/python/ray/data/tests/datasource/test_table_datasink.py new file mode 100644 index 000000000000..cd59db816bbe --- /dev/null +++ b/python/ray/data/tests/datasource/test_table_datasink.py @@ -0,0 +1,632 @@ +"""Tests for the generic ``TableDatasink`` framework. + +These tests exercise the framework in isolation by plugging in a +``FakeAdapter`` (APPEND + OVERWRITE only) or a ``FakeUpsertAdapter`` (also +conforming to ``SupportsUpserts``). They guarantee that the framework — +independent of any specific table format — preserves the contract shown in +the design sequence diagram: + + preflight -> on_write_start -> + (per worker: start_task, write_block*, finalize_task, task_metadata) -> + gather_task_metadata -> reconcile_schema -> + (APPEND) commit_append + (OVERWRITE) build_overwrite_predicate -> commit_overwrite + (UPSERT) build_upsert_predicate -> commit_upsert +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Set, Tuple + +import pyarrow as pa +import pytest + +from ray.data._internal.datasource.table import ( + SaveMode, + SupportsUpserts, + TableAdapter, + TableDatasink, + TableWriteTaskResult, + UpsertSemantics, +) +from ray.data.datasource.datasink import WriteResult + + +@dataclass +class _FakeFileAction: + path: str + rows: int + + +@dataclass +class _FakeCall: + """Record of a single framework -> adapter call.""" + + method: str + args: tuple = () + kwargs: Dict[str, Any] = field(default_factory=dict) + + +class FakeAdapter(TableAdapter[_FakeFileAction, str]): + """Minimal in-memory APPEND+OVERWRITE adapter that records every call. + + Does **not** support UPSERT — does not implement the ``SupportsUpserts`` + Protocol. ``FakeUpsertAdapter`` below extends this class with the upsert + methods. + """ + + def __init__( + self, + supported_modes: Optional[Set[SaveMode]] = None, + fail_at: Optional[str] = None, + ): + self._supported_modes = supported_modes or { + SaveMode.APPEND, + SaveMode.OVERWRITE, + } + self._fail_at = fail_at + self.calls: List[_FakeCall] = [] + self._next_block_idx = 0 + # State the framework will inject. + self._reconciled_schema: Optional[pa.Schema] = None + + # --- introspection -------------------------------------------------- + @property + def supported_modes(self) -> Set[SaveMode]: + return self._supported_modes + + def get_name(self) -> str: + return "Fake" + + # --- lifecycle ------------------------------------------------------ + def _record(self, method: str, *args, **kwargs) -> None: + self.calls.append(_FakeCall(method=method, args=args, kwargs=kwargs)) + if self._fail_at == method: + raise RuntimeError(f"injected failure at {method}") + + def preflight(self, mode, partition_cols, declared_schema): + self._record( + "preflight", + mode=mode, + partition_cols=tuple(partition_cols), + declared_schema=declared_schema, + ) + + def on_write_start(self, schema_from_first_bundle=None): + self._record("on_write_start", schema=schema_from_first_bundle) + + def start_task(self, ctx): + self._record("start_task", task_idx=getattr(ctx, "task_idx", None)) + + def write_block(self, arrow_table): + self._record("write_block", num_rows=arrow_table.num_rows) + self._next_block_idx += 1 + upsert_keys = ( + arrow_table.select([arrow_table.column_names[0]]) + if "id" in arrow_table.column_names + else None + ) + return ( + [ + _FakeFileAction( + path=f"file-{self._next_block_idx}.parquet", + rows=arrow_table.num_rows, + ) + ], + arrow_table.schema, + upsert_keys, + ) + + def finalize_task(self): + self._record("finalize_task") + return ([], []) + + def task_metadata(self) -> Dict[str, Any]: + self._record("task_metadata") + return {"write_uuid": "fake-uuid"} + + def gather_task_metadata(self, task_metadata: List[Dict[str, Any]]) -> None: + self._record("gather_task_metadata", task_metadata=task_metadata) + + def reconcile_schema(self, unified_schema): + self._record("reconcile_schema", schema=unified_schema) + self._reconciled_schema = unified_schema + + # --- mode-specific commit methods ----------------------------------- + def commit_append(self, file_actions, unified_schema): + self._record( + "commit_append", + num_files=len(file_actions), + unified_schema=unified_schema, + ) + + def build_overwrite_predicate(self, overwrite_filter): + self._record( + "build_overwrite_predicate", overwrite_filter=overwrite_filter + ) + return f"predicate-for-overwrite:{overwrite_filter!r}" + + def commit_overwrite(self, file_actions, unified_schema, delete_predicate): + self._record( + "commit_overwrite", + num_files=len(file_actions), + unified_schema=unified_schema, + delete_predicate=delete_predicate, + ) + + def on_failure(self, written_paths): + self._record("on_failure", written_paths=tuple(written_paths)) + + +class FakeUpsertAdapter(FakeAdapter): + """``FakeAdapter`` that also conforms to :class:`SupportsUpserts`.""" + + upsert_semantics: UpsertSemantics = UpsertSemantics.COPY_ON_WRITE + + def __init__( + self, + supported_modes: Optional[Set[SaveMode]] = None, + fail_at: Optional[str] = None, + ): + super().__init__( + supported_modes=supported_modes + or {SaveMode.APPEND, SaveMode.OVERWRITE, SaveMode.UPSERT}, + fail_at=fail_at, + ) + + def build_upsert_predicate(self, upsert_keys, join_cols): + self._record( + "build_upsert_predicate", + upsert_keys=upsert_keys, + join_cols=tuple(join_cols), + ) + return f"predicate-for-upsert:{tuple(join_cols)}" + + def commit_upsert(self, file_actions, unified_schema, delete_predicate): + self._record( + "commit_upsert", + num_files=len(file_actions), + unified_schema=unified_schema, + delete_predicate=delete_predicate, + ) + + +# ---------------------------------------------------------------------- +# Helpers. +# ---------------------------------------------------------------------- + + +def _table(rows: int, cols: Optional[List[Tuple[str, pa.DataType]]] = None) -> pa.Table: + cols = cols or [("id", pa.int64()), ("v", pa.string())] + data = { + name: pa.array( + list(range(rows)) + if pa.types.is_integer(t) + else [f"r{i}" for i in range(rows)], + type=t, + ) + for name, t in cols + } + return pa.table(data) + + +class _FakeCtx: + """Simulate ``TaskContext`` for the framework's ``write`` entry-point.""" + + def __init__(self, task_idx: int = 0): + self.task_idx = task_idx + self.kwargs: Dict[str, Any] = {} + + +def _drive_write(sink: TableDatasink, batches: List[pa.Table], task_idx: int = 0): + """Run the framework's ``write`` lifecycle on a single task's worth of + Arrow tables and return the resulting ``TableWriteTaskResult``.""" + return sink.write(iter(batches), _FakeCtx(task_idx=task_idx)) + + +def _wrap_result(returns: List[TableWriteTaskResult]) -> WriteResult: + num_rows = sum(sum(a.rows for a in r.file_actions) for r in returns) + size_bytes = num_rows * 8 # arbitrary; only the counters matter to the framework + return WriteResult( + num_rows=num_rows, size_bytes=size_bytes, write_returns=list(returns) + ) + + +# ---------------------------------------------------------------------- +# Tests. +# ---------------------------------------------------------------------- + + +def test_mode_validation_rejects_unsupported_mode(): + adapter = FakeAdapter(supported_modes={SaveMode.APPEND}) + with pytest.raises(ValueError, match="does not support mode"): + TableDatasink(adapter, SaveMode.OVERWRITE) + + +def test_mode_validation_accepts_supported_mode(): + adapter = FakeAdapter() + sink = TableDatasink(adapter, SaveMode.APPEND) + assert sink.get_name() == "Fake" + + +def test_str_mode_is_coerced(): + adapter = FakeAdapter() + sink = TableDatasink(adapter, "append") + assert sink._mode == SaveMode.APPEND + + +def test_invalid_mode_raises(): + adapter = FakeAdapter() + with pytest.raises(ValueError, match="Invalid mode"): + TableDatasink(adapter, "nope") + + +def test_upsert_rejected_for_non_upsert_capable_adapter(): + """An adapter that doesn't conform to ``SupportsUpserts`` must be + refused at ``TableDatasink`` construction time when mode is UPSERT, + even if it (incorrectly) lists UPSERT in supported_modes.""" + adapter = FakeAdapter( + supported_modes={SaveMode.APPEND, SaveMode.OVERWRITE, SaveMode.UPSERT} + ) + with pytest.raises(ValueError, match="does not support UPSERT"): + TableDatasink(adapter, SaveMode.UPSERT) + + +def test_isinstance_supports_upserts_runtime_check(): + """``@runtime_checkable`` makes the structural check work.""" + assert not isinstance(FakeAdapter(), SupportsUpserts) + assert isinstance(FakeUpsertAdapter(), SupportsUpserts) + + +def test_on_write_start_runs_preflight_then_adapter_hook(): + adapter = FakeAdapter() + sink = TableDatasink(adapter, SaveMode.APPEND, partition_cols=["p"]) + sink.on_write_start(schema=pa.schema([("a", pa.int64())])) + methods = [c.method for c in adapter.calls] + assert methods == ["preflight", "on_write_start"] + assert adapter.calls[0].kwargs["partition_cols"] == ("p",) + + +def test_write_invokes_write_block_per_nonempty_table_and_skips_empties(): + adapter = FakeAdapter() + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + result = _drive_write( + sink, + [_table(0), _table(3), _table(0), _table(2)], + task_idx=7, + ) + methods = [c.method for c in adapter.calls] + # preflight, on_write_start, start_task, write_block, write_block, + # finalize_task, task_metadata + assert methods == [ + "preflight", + "on_write_start", + "start_task", + "write_block", + "write_block", + "finalize_task", + "task_metadata", + ] + assert result.task_id == 7 + assert [a.rows for a in result.file_actions] == [3, 2] + assert result.task_metadata == {"write_uuid": "fake-uuid"} + + +def test_on_write_complete_runs_lifecycle_in_diagram_order_for_append(): + adapter = FakeAdapter() + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + r = _drive_write(sink, [_table(4)]) + sink.on_write_complete(_wrap_result([r])) + + methods = [c.method for c in adapter.calls] + # APPEND skips both build_overwrite_predicate and build_upsert_predicate. + assert methods == [ + "preflight", + "on_write_start", + "start_task", + "write_block", + "finalize_task", + "task_metadata", + "gather_task_metadata", + "reconcile_schema", + "commit_append", + ] + + +def test_schema_reconciliation_type_promotes_across_workers(): + adapter = FakeAdapter() + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + r1 = _drive_write(sink, [_table(1, cols=[("a", pa.int32())])], task_idx=0) + r2 = _drive_write(sink, [_table(1, cols=[("a", pa.int64())])], task_idx=1) + sink.on_write_complete(_wrap_result([r1, r2])) + + reconcile = next(c for c in adapter.calls if c.method == "reconcile_schema") + unified = reconcile.kwargs["schema"] + assert unified is not None + # int32 + int64 -> int64 + assert unified.field("a").type == pa.int64() + + +def test_upsert_predicate_built_and_passed_into_commit_upsert(): + adapter = FakeUpsertAdapter() + sink = TableDatasink(adapter, SaveMode.UPSERT, join_cols=["id"]) + sink.on_write_start() + r = _drive_write(sink, [_table(2)]) + sink.on_write_complete(_wrap_result([r])) + + build = next(c for c in adapter.calls if c.method == "build_upsert_predicate") + commit = next(c for c in adapter.calls if c.method == "commit_upsert") + assert build.kwargs["join_cols"] == ("id",) + assert commit.kwargs["delete_predicate"] == "predicate-for-upsert:('id',)" + + +def test_overwrite_predicate_built_and_passed_into_commit_overwrite(): + adapter = FakeAdapter() + sink = TableDatasink(adapter, SaveMode.OVERWRITE, overwrite_filter="my-filter") + sink.on_write_start() + r = _drive_write(sink, [_table(2)]) + sink.on_write_complete(_wrap_result([r])) + + build = next(c for c in adapter.calls if c.method == "build_overwrite_predicate") + commit = next(c for c in adapter.calls if c.method == "commit_overwrite") + assert build.kwargs["overwrite_filter"] == "my-filter" + assert commit.kwargs["delete_predicate"] == "predicate-for-overwrite:'my-filter'" + + +def test_append_invokes_only_commit_append_with_no_predicate(): + adapter = FakeAdapter() + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + r = _drive_write(sink, [_table(2)]) + sink.on_write_complete(_wrap_result([r])) + + methods = [c.method for c in adapter.calls] + # No build_*_predicate calls of any kind for APPEND. + assert "build_overwrite_predicate" not in methods + assert "build_upsert_predicate" not in methods + commit = next(c for c in adapter.calls if c.method == "commit_append") + assert commit.kwargs["num_files"] == 1 + + +def test_duplicate_file_paths_across_tasks_raises(): + adapter = FakeAdapter() + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + r1 = _drive_write(sink, [_table(1)], task_idx=0) + # Mint a second result that re-uses the same file path. + duplicate_action = r1.file_actions[0] + r2 = TableWriteTaskResult( + file_actions=[duplicate_action], emitted_schemas=r1.emitted_schemas + ) + with pytest.raises(ValueError, match="Duplicate file paths"): + sink.on_write_complete(_wrap_result([r1, r2])) + + +def test_commit_append_invoked_with_empty_actions(): + """Empty writes still get one commit_append call so adapters can create + empty tables, no-op IGNORE, etc.""" + adapter = FakeAdapter() + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + sink.on_write_complete(_wrap_result([])) + + methods = [c.method for c in adapter.calls] + # No worker ran, so no start_task / write_block / finalize_task / + # task_metadata. But the driver still drives the commit lifecycle. + assert methods == [ + "preflight", + "on_write_start", + "gather_task_metadata", + "reconcile_schema", + "commit_append", + ] + commit = adapter.calls[-1] + assert commit.kwargs["num_files"] == 0 + + +def test_on_write_failed_forwards_orphan_paths_to_adapter(): + adapter = FakeAdapter() + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + + error = RuntimeError("boom") + error._table_written_paths = ["a.parquet", "b.parquet"] + sink.on_write_failed(error) + + failure = next(c for c in adapter.calls if c.method == "on_failure") + assert failure.kwargs["written_paths"] == ("a.parquet", "b.parquet") + + +def test_on_write_failed_handles_adapter_cleanup_exception(): + adapter = FakeAdapter(fail_at="on_failure") + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + # Adapter raises inside on_failure; framework must swallow. + error = RuntimeError("boom") + error._table_written_paths = ["x.parquet"] + sink.on_write_failed(error) # must not raise + + +def test_write_attaches_orphan_paths_on_adapter_exception(): + adapter = FakeAdapter(fail_at="write_block") + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + with pytest.raises(RuntimeError, match="injected failure"): + _drive_write(sink, [_table(1)]) + + +# ---------------------------------------------------------------------- +# PR #63619 review fixes. +# ---------------------------------------------------------------------- + + +@pytest.mark.parametrize("mode", [SaveMode.ERROR, SaveMode.IGNORE]) +def test_error_ignore_modes_route_to_commit_append(mode): + """ERROR / IGNORE that survive preflight must commit via commit_append + (not silently drop the written files). Regression for PR #63619 #11.""" + adapter = FakeAdapter( + supported_modes={SaveMode.APPEND, SaveMode.ERROR, SaveMode.IGNORE} + ) + sink = TableDatasink(adapter, mode) + sink.on_write_start() + r = _drive_write(sink, [_table(3)]) + sink.on_write_complete(_wrap_result([r])) + + commit = next( + (c for c in adapter.calls if c.method == "commit_append"), None + ) + assert commit is not None, f"{mode} must route to commit_append" + assert commit.kwargs["num_files"] == 1 + # No overwrite/upsert commit paths should fire. + assert not any( + c.method in ("commit_overwrite", "commit_upsert") for c in adapter.calls + ) + + +def test_path_for_action_default_reads_path_attribute(): + """The default path_for_action reads ``action.path``.""" + adapter = FakeAdapter() + action = _FakeFileAction(path="foo/bar.parquet", rows=1) + assert adapter.path_for_action(action) == "foo/bar.parquet" + + +def test_path_for_action_override_drives_duplicate_detection(): + """An adapter whose action names the path field differently can override + path_for_action; the framework's dedup must use the override. Regression + for PR #63619 #1 (Iceberg DataFile uses file_path, not path).""" + + @dataclass + class _IcebergLikeFile: + file_path: str # note: NOT ``path`` + rows: int = 0 # for the _wrap_result row counter only + + class _IcebergLikeAdapter(FakeAdapter): + def write_block(self, arrow_table): + self._record("write_block", num_rows=arrow_table.num_rows) + self._next_block_idx += 1 + return ( + [ + _IcebergLikeFile( + file_path=f"file-{self._next_block_idx}.parquet", + rows=arrow_table.num_rows, + ) + ], + arrow_table.schema, + None, + ) + + def path_for_action(self, action): + return getattr(action, "file_path", None) + + adapter = _IcebergLikeAdapter() + # Sanity: the default impl would miss this field, the override catches it. + assert ( + TableAdapter.path_for_action( + adapter, _IcebergLikeFile(file_path="x.parquet") + ) + is None + ) + assert adapter.path_for_action(_IcebergLikeFile(file_path="x.parquet")) == ( + "x.parquet" + ) + + # Duplicate file_path across two tasks must now be detected. + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + r1 = _drive_write(sink, [_table(1)], task_idx=0) + dup = r1.file_actions[0] + r2 = TableWriteTaskResult( + file_actions=[dup], emitted_schemas=r1.emitted_schemas + ) + with pytest.raises(ValueError, match="Duplicate file paths"): + sink.on_write_complete(_wrap_result([r1, r2])) + + +def test_upsert_with_no_keys_passes_none_predicate_safely(): + """When no worker emits upsert keys, build_upsert_predicate receives + ``None`` and the lifecycle still completes. Regression for PR #63619 #14.""" + adapter = FakeUpsertAdapter() + sink = TableDatasink(adapter, SaveMode.UPSERT, join_cols=["id"]) + sink.on_write_start() + # No write tasks ran -> no key chunks -> upsert_keys is None. + sink.on_write_complete(_wrap_result([])) + + build = next(c for c in adapter.calls if c.method == "build_upsert_predicate") + assert build.kwargs["upsert_keys"] is None + commit = next(c for c in adapter.calls if c.method == "commit_upsert") + assert commit.kwargs["num_files"] == 0 + + +def test_create_mode_routes_to_commit_append(): + """SaveMode.CREATE must dispatch to commit_append (same commit semantics + as ERROR), not fall through to the unsupported-mode error after files are + written. Regression for PR #63619 (cursor, high).""" + adapter = FakeAdapter( + supported_modes={SaveMode.CREATE, SaveMode.APPEND, SaveMode.OVERWRITE} + ) + sink = TableDatasink(adapter, SaveMode.CREATE) + sink.on_write_start() + r = _drive_write(sink, [_table(2)]) + sink.on_write_complete(_wrap_result([r])) + + commit = next( + (c for c in adapter.calls if c.method == "commit_append"), None + ) + assert commit is not None, "CREATE must route to commit_append" + assert commit.kwargs["num_files"] == 1 + assert not any( + c.method in ("commit_overwrite", "commit_upsert") for c in adapter.calls + ) + + +def test_orphan_paths_attached_when_task_metadata_raises(): + """If a post-write step (task_metadata / upsert-key concat) raises, the + orphan-path list must still ride out on the exception so on_write_failed + can clean up files this task already wrote. Regression for PR #63619 + (cursor, low).""" + adapter = FakeAdapter(fail_at="task_metadata") + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + + with pytest.raises(RuntimeError, match="injected failure at task_metadata"): + _drive_write(sink, [_table(2)]) + + # The block wrote one file before task_metadata blew up; its path must be + # captured on the exception for cleanup. + try: + _drive_write(sink, [_table(2)]) + except RuntimeError as e: + assert getattr(e, "_table_written_paths", None), ( + "orphan paths must be attached to the exception" + ) + assert all(p.endswith(".parquet") for p in e._table_written_paths) + + +def test_gather_receives_one_entry_per_task_even_when_metadata_empty(): + """Tasks whose task_metadata() returns {} must still appear in the list + handed to gather_task_metadata (one entry per task that ran). Regression + for PR #63619 (cursor, medium).""" + + class _EmptyMetaAdapter(FakeAdapter): + def task_metadata(self): + self._record("task_metadata") + return {} + + adapter = _EmptyMetaAdapter() + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + r1 = _drive_write(sink, [_table(1)], task_idx=0) + r2 = _drive_write(sink, [_table(1)], task_idx=1) + sink.on_write_complete(_wrap_result([r1, r2])) + + gather = next(c for c in adapter.calls if c.method == "gather_task_metadata") + # Pre-fix the truthiness guard dropped both empty dicts -> length 0. + assert len(gather.kwargs["task_metadata"]) == 2 + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main(["-v", "-x", __file__])) diff --git a/python/ray/data/tests/datasource/test_table_datasink_integration.py b/python/ray/data/tests/datasource/test_table_datasink_integration.py new file mode 100644 index 000000000000..01b11a7f54a8 --- /dev/null +++ b/python/ray/data/tests/datasource/test_table_datasink_integration.py @@ -0,0 +1,561 @@ +"""Integration tests for the ``TableDatasink`` framework via a real adapter. + +These tests close the coverage gap left by ``test_table_datasink.py`` (which +uses a recorder-only ``FakeAdapter``). Here we plug in a minimal real-I/O +``ToyParquetAdapter`` that uses :class:`ParquetFileWriter` to write actual +Parquet files to a tmp directory, then assert behaviour end-to-end: + +* Files appear on disk. +* Content round-trips via ``pyarrow.parquet.read_table``. +* ``ParquetFileWriter`` buffering / flush semantics work. +* Schema unification across workers is visible on disk. +* ``on_failure`` cleans up orphan files written by failing tasks. + +All tests drive the framework in-process via ``sink.write(iter([table]), ctx)`` +— the same pattern ``test_table_datasink.py`` uses — so they don't require a +Ray cluster. +""" + +import pickle +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Optional, Set, Tuple + +import pyarrow as pa +import pyarrow.fs as pa_fs +import pyarrow.parquet as pq +import pytest + +from ray.data._internal.datasource.table import ( + ParquetFileWriter, + SaveMode, + TableAdapter, + TableDatasink, + TableWriteTaskResult, +) +from ray.data.datasource.datasink import WriteResult + + +# ---------------------------------------------------------------------- +# ToyParquetAdapter — minimal real-I/O reference adapter. +# ---------------------------------------------------------------------- + + +@dataclass +class _ToyFile: + """Per-file metadata returned by ``ToyParquetAdapter.write_block``.""" + + path: str + size: int + partition_values: Dict[str, Optional[str]] + + +_HIVE_DEFAULT_PARTITION = "__HIVE_DEFAULT_PARTITION__" + + +def _hive_path_builder( + cols: List[str], values: Tuple +) -> Tuple[str, Dict[str, Optional[str]]]: + """Mirror Delta's Hive-style partition encoding: None -> default sentinel, + float NaN -> "NaN". Returns (relative_dir_with_trailing_slash, value_map).""" + if not cols: + return ("", {}) + parts: List[str] = [] + value_map: Dict[str, Optional[str]] = {} + for col, val in zip(cols, values): + if val is None: + encoded = _HIVE_DEFAULT_PARTITION + elif isinstance(val, float) and val != val: # NaN + encoded = "NaN" + else: + encoded = str(val) + parts.append(f"{col}={encoded}") + value_map[col] = None if val is None else encoded + return ("/".join(parts) + "/", value_map) + + +class ToyParquetAdapter(TableAdapter[_ToyFile, str]): + """Minimal real-I/O adapter for framework integration tests. + + Writes Parquet files into ``root`` via :class:`ParquetFileWriter`. + Supports APPEND and OVERWRITE only — does NOT implement + :class:`SupportsUpserts` so UPSERT is rejected by the framework. + + OVERWRITE semantics: ``commit_overwrite`` deletes every ``.parquet`` + file in ``root`` that isn't in the freshly-written set. This is + intentionally non-transactional (it's a test adapter); just enough to + exercise the framework's per-mode dispatch. + """ + + def __init__( + self, + root: str, + *, + partition_cols: Optional[List[str]] = None, + target_file_size_bytes: Optional[int] = None, + compression: str = "snappy", + fail_after_n_writes: Optional[int] = None, + ): + self._root = Path(root) + self._partition_cols = list(partition_cols or []) + self._target_file_size_bytes = target_file_size_bytes + self._compression = compression + self._fail_after_n_writes = fail_after_n_writes + + # Lifecycle state. + self._mode: Optional[SaveMode] = None + self._task_written: Set[str] = set() + self._writer: Optional[ParquetFileWriter] = None + self._n_writes_seen = 0 + + # ------------------------------------------------------------------ + # Introspection. + # ------------------------------------------------------------------ + @property + def supported_modes(self) -> Set[SaveMode]: + return {SaveMode.APPEND, SaveMode.OVERWRITE} + + def get_name(self) -> str: + return "ToyParquet" + + # ------------------------------------------------------------------ + # Driver lifecycle. + # ------------------------------------------------------------------ + def preflight(self, mode, partition_cols, declared_schema) -> None: + self._mode = mode + self._root.mkdir(parents=True, exist_ok=True) + + # ------------------------------------------------------------------ + # Worker lifecycle. + # ------------------------------------------------------------------ + def start_task(self, ctx) -> None: + # Build a fresh per-task ParquetFileWriter pointed at self._root. + # SubTreeFileSystem strips the absolute tmpdir prefix so the writer + # sees / reports paths relative to the table root (matches the + # Delta / Iceberg convention). + fs = pa_fs.SubTreeFileSystem(str(self._root), pa_fs.LocalFileSystem()) + self._task_written = set() + self._writer = ParquetFileWriter( + filesystem=fs, + partition_cols=self._partition_cols, + write_uuid="toy00000", + compression=self._compression, + write_statistics=False, + target_file_size_bytes=self._target_file_size_bytes, + path_builder=_hive_path_builder, + action_factory=lambda *, path, size, partition_values, modification_time, stats: _ToyFile( + path=path, size=size, partition_values=partition_values + ), + stats_factory=lambda _t: None, + written_files=self._task_written, + ) + self._writer.set_task_idx(int(getattr(ctx, "task_idx", 0) or 0)) + + def write_block( + self, arrow_table: pa.Table + ) -> Tuple[List[_ToyFile], pa.Schema, Optional[pa.Table]]: + if self._writer is None: + return ([], arrow_table.schema, None) + # Check the fail-injection BEFORE writing so the failing call leaves + # no on-disk side effects. Previous successful calls' files are the + # orphans that ``on_failure`` must clean up. + self._n_writes_seen += 1 + if ( + self._fail_after_n_writes is not None + and self._n_writes_seen >= self._fail_after_n_writes + ): + raise RuntimeError("injected write_block failure") + actions = self._writer.add_table(arrow_table) + return (actions, arrow_table.schema, None) + + def finalize_task(self) -> Tuple[List[_ToyFile], List[pa.Schema]]: + if self._writer is None: + return ([], []) + return (self._writer.flush(), []) + + # ------------------------------------------------------------------ + # Driver commit — files are already on disk; commit is bookkeeping. + # ------------------------------------------------------------------ + def commit_append(self, file_actions, unified_schema) -> None: + return None + + def build_overwrite_predicate(self, overwrite_filter) -> Optional[str]: + return None + + def commit_overwrite(self, file_actions, unified_schema, delete_predicate) -> None: + new_names = {a.path for a in file_actions} + for p in self._root.glob("*.parquet"): + if p.name not in new_names: + p.unlink() + + # ------------------------------------------------------------------ + # Failure handling. + # ------------------------------------------------------------------ + def on_failure(self, written_paths: List[str]) -> None: + for rel in written_paths: + full = self._root / rel + if full.exists(): + full.unlink() + + +# ---------------------------------------------------------------------- +# Helpers — drive the framework in-process. +# ---------------------------------------------------------------------- + + +class _FakeCtx: + def __init__(self, task_idx: int = 0): + self.task_idx = task_idx + self.kwargs: Dict[str, Any] = {} + + +def _drive_one_task( + sink: TableDatasink, batches: List[pa.Table], task_idx: int = 0 +) -> TableWriteTaskResult: + return sink.write(iter(batches), _FakeCtx(task_idx=task_idx)) + + +def _wrap_write_result(returns: List[TableWriteTaskResult]) -> WriteResult: + num_rows = sum(sum(getattr(a, "size", 0) for a in r.file_actions) for r in returns) + return WriteResult( + num_rows=num_rows, size_bytes=num_rows, write_returns=list(returns) + ) + + +def _list_parquet_files(root: Path) -> List[Path]: + return sorted(root.glob("*.parquet")) + + +def _list_parquet_files_recursive(root: Path) -> List[Path]: + return sorted(root.rglob("*.parquet")) + + +def _partition_dirs(root: Path) -> List[str]: + """Names of top-level partition directories (e.g. 'k=1', 'k=NaN').""" + return sorted(p.name for p in root.iterdir() if p.is_dir()) + + +def _read_file_only(f: Path) -> pa.Table: + """Read a single Parquet file's own columns, WITHOUT Hive partition + inference (``pq.read_table`` would re-attach partition columns from the + directory path). ``ParquetFile`` is a single-file reader, so the result + contains exactly what was written to disk.""" + return pq.ParquetFile(str(f)).read() + + +# ---------------------------------------------------------------------- +# Tests. +# ---------------------------------------------------------------------- + + +def test_append_writes_real_parquet_and_round_trips(tmp_path): + """APPEND through the framework writes a real Parquet file whose + contents round-trip via pq.read_table.""" + adapter = ToyParquetAdapter(str(tmp_path)) + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + + table = pa.table({"id": [1, 2, 3], "v": ["a", "b", "c"]}) + r = _drive_one_task(sink, [table]) + sink.on_write_complete(_wrap_write_result([r])) + + files = _list_parquet_files(tmp_path) + assert len(files) == 1, f"expected 1 parquet file, got {files}" + round_trip = pq.read_table(files[0]) + assert round_trip.column("id").to_pylist() == [1, 2, 3] + assert round_trip.column("v").to_pylist() == ["a", "b", "c"] + + +def test_overwrite_replaces_existing_files(tmp_path): + """OVERWRITE deletes everything from the previous APPEND and writes + fresh files.""" + # First APPEND a batch. + adapter = ToyParquetAdapter(str(tmp_path)) + sink_a = TableDatasink(adapter, SaveMode.APPEND) + sink_a.on_write_start() + r1 = _drive_one_task(sink_a, [pa.table({"id": [1], "v": ["old"]})]) + sink_a.on_write_complete(_wrap_write_result([r1])) + initial_files = {p.name for p in _list_parquet_files(tmp_path)} + assert len(initial_files) == 1 + + # Now OVERWRITE with a new batch. + adapter2 = ToyParquetAdapter(str(tmp_path)) + sink_b = TableDatasink(adapter2, SaveMode.OVERWRITE) + sink_b.on_write_start() + r2 = _drive_one_task(sink_b, [pa.table({"id": [2], "v": ["new"]})]) + sink_b.on_write_complete(_wrap_write_result([r2])) + + final_files = _list_parquet_files(tmp_path) + assert len(final_files) == 1 + # The original file must have been deleted. + assert final_files[0].name not in initial_files + round_trip = pq.read_table(final_files[0]) + assert round_trip.column("v").to_pylist() == ["new"] + + +def test_parquet_file_writer_buffers_until_target_size(tmp_path): + """With target_file_size_bytes large, multiple add_table calls accumulate + into one file at flush time.""" + adapter = ToyParquetAdapter( + str(tmp_path), target_file_size_bytes=10**9 + ) # effectively unbounded buffer + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + + batches = [pa.table({"id": [i], "v": [f"r{i}"]}) for i in range(5)] + # Before driving the task: no files yet. + assert _list_parquet_files(tmp_path) == [] + + r = _drive_one_task(sink, batches) + # finalize_task ran inside _drive_one_task -> sink.write, flushing the buffer. + sink.on_write_complete(_wrap_write_result([r])) + + files = _list_parquet_files(tmp_path) + assert len(files) == 1, f"expected 1 buffered+flushed file, got {len(files)}" + round_trip = pq.read_table(files[0]) + assert sorted(round_trip.column("id").to_pylist()) == [0, 1, 2, 3, 4] + + +def test_parquet_file_writer_flushes_per_block_when_unbuffered(tmp_path): + """With target_file_size_bytes None (default), each non-empty add_table + flushes immediately -> one file per Arrow block.""" + adapter = ToyParquetAdapter(str(tmp_path)) # target=None -> immediate flush + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + + batches = [pa.table({"id": [i], "v": [f"r{i}"]}) for i in range(3)] + r = _drive_one_task(sink, batches) + sink.on_write_complete(_wrap_write_result([r])) + + files = _list_parquet_files(tmp_path) + assert len(files) == 3 + + +def test_schema_unification_int32_int64_round_trips_on_disk(tmp_path): + """Two tasks emit different int widths; framework unifies; disk content + is recoverable.""" + adapter = ToyParquetAdapter(str(tmp_path)) + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + + r1 = _drive_one_task( + sink, + [pa.table({"a": pa.array([1], type=pa.int32())})], + task_idx=0, + ) + r2 = _drive_one_task( + sink, + [pa.table({"a": pa.array([10_000_000_000], type=pa.int64())})], + task_idx=1, + ) + sink.on_write_complete(_wrap_write_result([r1, r2])) + + files = _list_parquet_files(tmp_path) + assert len(files) == 2 + # Concatenate file contents; framework should have unified to int64. + all_rows = sorted( + v for f in files for v in pq.read_table(f).column("a").to_pylist() + ) + assert all_rows == [1, 10_000_000_000] + + +def test_empty_writes_still_call_commit_append(tmp_path): + """When no worker emits any file, commit_append still fires once.""" + adapter = ToyParquetAdapter(str(tmp_path)) + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + # No task ran; pass an empty write_returns list. + sink.on_write_complete(_wrap_write_result([])) + + # No files written, and adapter survived the lifecycle (no exception). + assert _list_parquet_files(tmp_path) == [] + + +def test_on_failure_removes_orphan_files(tmp_path): + """When write_block fails partway through, the framework attaches the + orphan paths to the exception and on_write_failed forwards them to the + adapter, which deletes them from disk.""" + # fail_after_n_writes=2 -> the first block writes a file, the second + # raises before writing anything. + adapter = ToyParquetAdapter(str(tmp_path), fail_after_n_writes=2) + sink = TableDatasink(adapter, SaveMode.APPEND) + sink.on_write_start() + + batches = [ + pa.table({"id": [1], "v": ["good"]}), + pa.table({"id": [2], "v": ["fails"]}), + ] + with pytest.raises(RuntimeError, match="injected write_block failure"): + _drive_one_task(sink, batches) + + # The first block did land on disk before the second one failed. + files_before_cleanup = _list_parquet_files(tmp_path) + assert len(files_before_cleanup) == 1 + + # Simulate Ray Data calling on_write_failed with the captured exception. + err = RuntimeError("from worker") + err._table_written_paths = [p.name for p in files_before_cleanup] + sink.on_write_failed(err) + + # Orphans removed. + assert _list_parquet_files(tmp_path) == [] + + +def test_adapter_pickles_for_distributed_writes(tmp_path): + """The adapter must round-trip through pickle so Ray Data can ship it + to workers.""" + adapter = ToyParquetAdapter(str(tmp_path), target_file_size_bytes=4096) + blob = pickle.dumps(adapter) + restored = pickle.loads(blob) + + assert isinstance(restored, ToyParquetAdapter) + assert restored.supported_modes == {SaveMode.APPEND, SaveMode.OVERWRITE} + assert restored.get_name() == "ToyParquet" + + # Verify the restored adapter can still drive a full lifecycle. + sink = TableDatasink(restored, SaveMode.APPEND) + sink.on_write_start() + r = _drive_one_task(sink, [pa.table({"id": [42]})]) + sink.on_write_complete(_wrap_write_result([r])) + assert len(_list_parquet_files(tmp_path)) == 1 + + +# ---------------------------------------------------------------------- +# Partitioning — exercises the sort-based _partition_table rewrite +# (PR #63619 findings #3/#4/#5/#6). +# ---------------------------------------------------------------------- + + +def test_partition_single_column_groups_correctly(tmp_path): + """Single-column partitioning routes each distinct value to its own + Hive directory; rows round-trip with partition column dropped from the + payload.""" + adapter = ToyParquetAdapter(str(tmp_path), partition_cols=["k"]) + sink = TableDatasink(adapter, SaveMode.APPEND, partition_cols=["k"]) + sink.on_write_start() + + table = pa.table({"k": [1, 2, 1, 2, 1], "v": ["a", "b", "c", "d", "e"]}) + r = _drive_one_task(sink, [table]) + sink.on_write_complete(_wrap_write_result([r])) + + assert _partition_dirs(tmp_path) == ["k=1", "k=2"] + # Reassemble all rows from all partition files; partition col is dropped + # on disk, so re-attach it from the directory name. + rows = [] + for f in _list_parquet_files_recursive(tmp_path): + kval = int(f.parent.name.split("=")[1]) + t = _read_file_only(f) + assert "k" not in t.column_names, "partition col must be dropped on disk" + for v in t.column("v").to_pylist(): + rows.append((kval, v)) + assert sorted(rows) == [(1, "a"), (1, "c"), (1, "e"), (2, "b"), (2, "d")] + + +def test_partition_nan_rows_grouped_together_single_column(tmp_path): + """All NaN rows in a float partition column land in one ``k=NaN`` + partition (not fragmented one-file-per-row). Covers #6 (within-table + NaN grouping).""" + adapter = ToyParquetAdapter(str(tmp_path), partition_cols=["k"]) + sink = TableDatasink(adapter, SaveMode.APPEND, partition_cols=["k"]) + sink.on_write_start() + + nan = float("nan") + table = pa.table( + {"k": [1.0, nan, 2.0, nan, nan], "v": ["a", "b", "c", "d", "e"]} + ) + r = _drive_one_task(sink, [table]) + sink.on_write_complete(_wrap_write_result([r])) + + assert _partition_dirs(tmp_path) == ["k=1.0", "k=2.0", "k=NaN"] + nan_files = list((tmp_path / "k=NaN").glob("*.parquet")) + assert len(nan_files) == 1, "all NaN rows must coalesce into one file" + nan_rows = _read_file_only(nan_files[0]).column("v").to_pylist() + assert sorted(nan_rows) == ["b", "d", "e"] + + +def test_partition_nan_buffer_coalesces_across_calls(tmp_path): + """With a large target_file_size_bytes, NaN-partition rows from multiple + add_table calls accumulate into a single buffered file rather than + fragmenting (every NaN being a distinct dict key). Covers #4.""" + adapter = ToyParquetAdapter( + str(tmp_path), partition_cols=["k"], target_file_size_bytes=10**9 + ) + sink = TableDatasink(adapter, SaveMode.APPEND, partition_cols=["k"]) + sink.on_write_start() + + nan = float("nan") + # Three separate blocks, each with a NaN-partition row. + batches = [ + pa.table({"k": [nan], "v": [f"r{i}"]}) for i in range(3) + ] + r = _drive_one_task(sink, batches) + sink.on_write_complete(_wrap_write_result([r])) + + nan_files = list((tmp_path / "k=NaN").glob("*.parquet")) + assert len(nan_files) == 1, ( + f"NaN buffer should coalesce to one file, got {len(nan_files)}" + ) + assert sorted(_read_file_only(nan_files[0]).column("v").to_pylist()) == [ + "r0", + "r1", + "r2", + ] + + +def test_partition_multi_column_with_nulls_and_nan(tmp_path): + """Multi-column partitioning groups correctly, including None and NaN + in the key (covers #6 fallback-path corruption + #5 dead-code removal).""" + adapter = ToyParquetAdapter(str(tmp_path), partition_cols=["a", "b"]) + sink = TableDatasink(adapter, SaveMode.APPEND, partition_cols=["a", "b"]) + sink.on_write_start() + + nan = float("nan") + table = pa.table( + { + "a": ["x", "x", "y", None, None], + "b": [1.0, 1.0, nan, nan, nan], + "v": ["r0", "r1", "r2", "r3", "r4"], + } + ) + r = _drive_one_task(sink, [table]) + sink.on_write_complete(_wrap_write_result([r])) + + # Expected groups: (x,1.0)->{r0,r1}, (y,NaN)->{r2}, (None,NaN)->{r3,r4} + files = _list_parquet_files_recursive(tmp_path) + group_to_rows = {} + for f in files: + # parent dir is "b=...", grandparent is "a=..." + b_dir = f.parent.name + a_dir = f.parent.parent.name + key = (a_dir, b_dir) + group_to_rows.setdefault(key, []).extend( + _read_file_only(f).column("v").to_pylist() + ) + got = {k: sorted(v) for k, v in group_to_rows.items()} + assert got == { + ("a=x", "b=1.0"): ["r0", "r1"], + ("a=y", "b=NaN"): ["r2"], + ("a=__HIVE_DEFAULT_PARTITION__", "b=NaN"): ["r3", "r4"], + } + + +def test_partition_high_cardinality_smoke(tmp_path): + """A few hundred distinct partition values complete quickly and route + correctly (sanity for the sort-based O(N log N) path, #3).""" + adapter = ToyParquetAdapter(str(tmp_path), partition_cols=["k"]) + sink = TableDatasink(adapter, SaveMode.APPEND, partition_cols=["k"]) + sink.on_write_start() + + n = 300 + table = pa.table({"k": list(range(n)), "v": [f"r{i}" for i in range(n)]}) + r = _drive_one_task(sink, [table]) + sink.on_write_complete(_wrap_write_result([r])) + + assert len(_partition_dirs(tmp_path)) == n + assert len(_list_parquet_files_recursive(tmp_path)) == n + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main(["-v", "-x", __file__]))