[data] PR 1 of 8: Introduce TableDatasink + TableAdapter abstraction

[abhishekverma-ray]

[data] Introduce TableDatasink + TableAdapter abstraction

Summary

This PR adds a generic framework, python/ray/data/_internal/datasource/table/, that unifies the distributed write plumbing shared by Ray Data's table-format datasinks (Iceberg today, Delta next, Hudi/Paimon plausible). Supporting a new format reduces from "write a new ~3000-line Datasink subclass" to "implement a TableAdapter."

This PR is strictly additive. No existing call sites are touched; no public API moves. IcebergDatasink continues to work unchanged. Subsequent PRs in the stack migrate Iceberg onto this abstraction (PR 2) and add write_delta on top of it (PRs 3–8).

Motivation

The pre-existing IcebergDatasink and the in-flight DeltaDatasink were structurally identical above the format layer:

• Driver-side: validate inputs → load table → coordinate workers → reconcile per-worker schemas → commit one transaction.
• Worker-side: write Parquet files → return metadata to driver → on failure, hand orphan paths back for cleanup.

Without an abstraction, each format would re-implement that lifecycle (and reviewers would have to chase the same plumbing bugs in two places). This PR lifts the lifecycle once, behind a narrow TableAdapter contract, so format-specific code shrinks to "load the table, write a Parquet file, commit a transaction."

What's in this PR

python/ray/data/_internal/datasource/table/
├── __init__.py            # public re-exports
├── modes.py               # SaveMode (re-export) + UpsertSemantics
├── result.py              # TableWriteTaskResult[FileAction]
├── adapter.py             # TableAdapter[FileAction, DeletePredicate] ABC
│                          #   + SupportsUpserts opt-in Protocol
├── file_writer.py         # DataFileWriter protocol + ParquetFileWriter
├── table_datasink.py      # TableDatasink — the Template Method
└── SEQUENCE.md            # Mermaid sequence diagram (mirrors the ASCII
                           #   diagram embedded in adapter.py)

python/ray/data/tests/datasource/
├── test_table_datasink.py             # framework unit tests (FakeAdapter)
└── test_table_datasink_integration.py # real-I/O tests (ToyParquetAdapter)

9 new files, ~2,400 LoC (~1,200 production + ~1,200 tests). Zero existing files modified.

Architecture

TableDatasink is a Template Method; TableAdapter is the Strategy it delegates to. The framework owns what happens; adapters own how for one format. The full call sequence is reproduced as an ASCII diagram in the adapter.py module docstring and as Mermaid in SEQUENCE.md.

Framework owns (TableDatasink)

Concern	Responsibility
Lifecycle	on_write_start → per-task write → on_write_complete / on_write_failed
Mode validation	Reject modes the adapter doesn't declare in supported_modes; coerce string modes; reject UPSERT for adapters that don't conform to SupportsUpserts
Schema unification	Concatenate per-worker emitted schemas and run unify_schemas(..., promote_types=True) so int32+int64→int64 etc. before commit
Per-mode commit dispatch	APPEND → commit_append; OVERWRITE → build_overwrite_predicate + commit_overwrite; UPSERT → build_upsert_predicate + commit_upsert
UPSERT key plumbing	Collect per-task upsert_keys tables, concat, hand to build_upsert_predicate
Orphan cleanup	Collect per-task written paths (via adapter.path_for_action); attach to exceptions so on_write_failed hands them back to the adapter
Duplicate-path guard	Reject if two workers report the same file path (path obtained via adapter.path_for_action)

Adapter contract — split in two

TableAdapter[FileAction, DeletePredicate] — the APPEND + OVERWRITE baseline every adapter implements:

Method	Where	Purpose
supported_modes	property	Which SaveModes the adapter supports
preflight	driver	Load table, validate mode legality / partitions / declared schema
on_write_start	driver	Optional pre-write hook fed the first bundle's schema (Iceberg evolves schema here; Delta no-ops)
start_task	worker	Optional per-task setup
write_block	worker	Persist one Arrow table; return (file_actions, emitted_schema, upsert_keys?)
finalize_task	worker	Flush per-task buffers
task_metadata	worker	Free-form worker→driver state (e.g. Delta's app-transaction UUID)
gather_task_metadata	driver	Receive per-task metadata before commit
reconcile_schema	driver	Apply the unified schema to the table
build_overwrite_predicate	driver	Translate overwrite_filter into the format's predicate type
commit_append / commit_overwrite	driver	Atomic, mode-specific transactions
path_for_action	introspection	Relative path of a file action (default reads .path); used for dedup + orphan tracking
on_failure	driver	Clean up orphans on failed writes

SupportsUpserts[FileAction, DeletePredicate] — an opt-in, @runtime_checkable Protocol. Adapters that support UPSERT inherit (or structurally conform to) it; the framework dispatches via isinstance(adapter, SupportsUpserts) and refuses mode="upsert" otherwise. This keeps the base contract free of upsert concerns for formats that don't support it (e.g. Delta in this stack):

Member	Purpose
upsert_semantics	COPY_ON_WRITE or MERGE_ON_READ
build_upsert_predicate	Build a delete predicate from the concatenated upsert keys (None ⇒ pure insert)
commit_upsert	Atomic UPSERT transaction

TableAdapter is parameterized on both FileAction (the opaque per-file metadata, e.g. Iceberg's DataFile / Delta's AddAction) and DeletePredicate (the format's predicate type, e.g. PyIceberg's BooleanExpression / Delta's SQL str), so predicate types flow precisely from build_* to commit_*.

ParquetFileWriter

A reusable DataFileWriter. Partition-path encoding and file-action shape are pluggable via callbacks, so Delta's AddAction and Iceberg's DataFile can share the same file-writer body. Partitioning uses a single sort + boundary-walk (O(N log N)) with a NaN-safe partition key, so NaN/None group correctly and buffered partitions coalesce across add_table calls.

Tests

test_table_datasink.py — framework-level unit tests via a FakeAdapter that records every lifecycle call. Locks down mode validation and SupportsUpserts gating; lifecycle ordering; per-mode commit dispatch (APPEND / OVERWRITE / UPSERT); schema type-promotion across workers; empty-write still commits; orphan-path forwarding on failure (and that secondary cleanup failures are swallowed); duplicate-path rejection; and path_for_action (default + override) driving dedup.

test_table_datasink_integration.py — real-I/O tests via a ToyParquetAdapter that drives ParquetFileWriter against a LocalFileSystem tmpdir: APPEND/OVERWRITE round-trips read back with pyarrow.parquet; buffering/flush semantics; on-disk schema unification; orphan-file cleanup; adapter pickling; and partitioning correctness (single/multi-column, NaN grouping + buffer coalescing, high-cardinality).

Test plan

• pytest python/ray/data/tests/datasource/test_table_datasink.py python/ray/data/tests/datasource/test_table_datasink_integration.py — 36 pass (23 unit + 13 integration).
• No existing test suite touched; test_iceberg.py collects/imports cleanly (nothing references the new module yet).
• CI green.

Risk

Low. Purely additive new package. No existing code paths reach this module. If reviewers want to defer the abstraction work entirely, the directory can be deleted in a single revert with zero collateral.

Follow-ups (separate PRs)

1. PR 2 — Refactor IcebergDatasink onto TableDatasink + IcebergAdapter (which conforms to SupportsUpserts and overrides path_for_action → DataFile.file_path). Existing test_iceberg.py is the regression net.
2. PR 3 — write_delta APPEND mode (new public API).
3. PRs 4–7 — Additional write_delta modes (OVERWRITE / ERROR / IGNORE), partitioning, schema evolution, cloud + idempotency.
4. PR 8 — Shared Parquet I/O primitives.

🤖 Generated with Claude Code

Co-Authored-By: Claude noreply@anthropic.com

[gemini-code-assist]

Code Review

This pull request introduces a generic table datasink framework for Ray Data, abstracting format-specific logic into a TableAdapter protocol and providing a generalized ParquetFileWriter for partitioned writes. The implementation includes the core TableDatasink orchestration and comprehensive unit tests. Feedback highlights the need to explicitly document or enforce the path attribute requirement on FileAction objects for reliable orphan cleanup. Additionally, it is recommended to log filesystem errors during directory creation and to optimize the partitioning logic, which currently exhibits O(N * K) complexity and may become a bottleneck for high-cardinality columns.

[goutamvenkat-anyscale]

what's the distinction between these 2 functions? Feels weird that there's iceberg specific logic in this abstraction.

[abhishekverma-ray]

preflight

• Purpose: validate the requested write against the current table state. Has access to the user's stated intent (mode, partition_cols, declared_schema) but not the actual incoming data.
• When it raises: synchronously fails the write before any worker tasks are scheduled. No files are written, no orphan cleanup needed.

on_write_start

• Purpose: an optional hook fed the first input bundle's actual schema — i.e. the framework sniffs one Arrow table from the upstream operator and hands its schema to the adapter so the adapter can do data-dependent setup before workers start writing.
• When it does nothing: the default base-class implementation is a no-op. Adapters that don't need pre-write schema evolution just don't override it (e.g. Delta)

[goutamvenkat-anyscale]

Not sure I understand the purpose of this function?

[abhishekverma-ray]

start_task — per-task setup on the worker

• Where it runs: inside each Ray write task, on a worker process. The framework calls it exactly once per task, immediately before the first write_block of that task

[goutamvenkat-anyscale]

FileAction = TypeVar("FileAction")
  DeletePredicate = TypeVar("DeletePredicate")


  class TableAdapter(Generic[FileAction, DeletePredicate], ABC):
      """Driver-only. APPEND + OVERWRITE."""

      @property
      @abstractmethod
      def supported_modes(self) -> Set[SaveMode]: ...

      @abstractmethod
      def preflight(
          self, mode: SaveMode, partition_cols: List[str],
          declared_schema: Optional[pa.Schema],
      ) -> None: ...

      def on_write_start(self, schema: Optional[pa.Schema] = None) -> None: ...

      @abstractmethod
      def build_worker(self) -> "TableWorker[FileAction]": ...

      @abstractmethod
      def commit_append(
          self, file_actions: List[FileAction],
          unified_schema: Optional[pa.Schema],
      ) -> None: ...

      @abstractmethod
      def build_overwrite_predicate(
          self, overwrite_filter: Optional[Any],
      ) -> Optional[DeletePredicate]: ...

      @abstractmethod
      def commit_overwrite(
          self, file_actions: List[FileAction],
          unified_schema: Optional[pa.Schema],
          delete_predicate: Optional[DeletePredicate],
      ) -> None: ...

      def on_failure(self, written_paths: List[str]) -> None: ...


  @runtime_checkable
  class UpsertCapable(Protocol[FileAction, DeletePredicate]):
      upsert_semantics: UpsertSemantics

      def build_upsert_predicate(
          self, upsert_keys: pa.Table, join_cols: List[str],
      ) -> DeletePredicate: ...

      def commit_upsert(
          self, file_actions: List[FileAction],

      def commit_upsert(
          self, file_actions: List[FileAction],
          unified_schema: Optional[pa.Schema],
          delete_predicate: DeletePredicate,
      ) -> None: ...

  Framework dispatch — build then commit, still one place:

  if mode == SaveMode.APPEND:
      self._adapter.commit_append(actions, unified)
  elif mode == SaveMode.OVERWRITE:
      predicate = self._adapter.build_overwrite_predicate(self._overwrite_filter)
      self._adapter.commit_overwrite(actions, unified, predicate)
  elif mode == SaveMode.UPSERT:
      predicate = self._adapter.build_upsert_predicate(keys, self._join_cols)
      self._adapter.commit_upsert(actions, unified, predicate)

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

^{Reviewed by Cursor Bugbot for commit 8c4eb05e17c6cff426b78d97ac78cacd81e9fb02. Configure here.}

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[abhishekverma-ray]

Please refer to #64300 instead.