[data] Add catalog support and credential refresh to write_delta - #65079
Open
abhishekverma-ray wants to merge 3 commits into
Open
[data] Add catalog support and credential refresh to write_delta#65079abhishekverma-ray wants to merge 3 commits into
abhishekverma-ray wants to merge 3 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a prototype implementation of Dataset.write_delta in Ray Data, enabling users to write datasets to Delta Lake tables with support for APPEND and OVERWRITE modes, schema evolution, and credential refresh. Feedback on the changes highlights a few critical issues: the initial table existence check in on_write_start is not wrapped in the retry mechanism, making it vulnerable to transient network or credential errors; the relative path computation in _write_parquet is not cross-platform compatible and will fail on Windows due to backslash path separators; and the environment variable mutation for GCS credentials is not thread-safe and requires synchronization using a thread lock.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Ports a small, independently-authored prototype write_delta implementation for a side-by-side comparison against the standalone write_delta design on stack2/pr1-write-delta-core. Dataset.write_delta and the DeltaDatasink class are kept as close to the shared original as possible -- deliberate fixes are called out below; everything else is unmodified -- so the diff between the two stacks reflects genuine design differences, not incidental ones. Supports SaveMode.APPEND and SaveMode.OVERWRITE (a prototype scope, by design): workers write Parquet files via pyarrow.dataset.write_dataset and report AddAction metadata; the driver commits all files in one Delta transaction via the deltalake library, creating the table if it doesn't exist. Also supports partition_by (Hive-style), storage_options, and optional name/description table metadata. Bug fix 1 (found while building a write_delta release benchmark on top of this branch): a Dataset sourced from a file-listing-based read (e.g. read_parquet) failed every write with "Schema error: Invalid data type for Delta Lake: Null". Root cause: on_write_start's schema argument can be Ray Data's internal file-listing metadata schema (e.g. a null-typed __file_chunk_metadata column) rather than the dataset's actual row schema, and _resolve_schema unconditionally preferred that captured schema over the real schema workers observed while writing. Fixed by having write() return the schema of the Arrow table(s) it actually wrote, and having _resolve_schema prefer that over the driver-side capture, which now serves only its originally-intended purpose: a fallback for a genuinely empty dataset with no worker-observed schema at all. Bug fix 2 (found from an actual CI run of the release benchmark, build 101454): a large distributed write (TPCH SF1000, ~1000 concurrent write tasks against a brand-new table root) failed with "AWS Error SLOW_DOWN during PutObject operation" while writing Parquet files, before any commit was attempted. Root cause: pyarrow.dataset.write_dataset defaults to create_dir=True, which for an S3-backed filesystem issues a PutObject to a single marker key representing the table root directory (verified empirically -- S3 has no real directories, so this is purely a convention). With ~1000 tasks all racing to create that identical key against a cold prefix, the concurrent-PUT burst triggers AWS throttling for no functional benefit (object writes don't need a pre-existing directory marker). Fixed by passing create_dir=False specifically when the resolved filesystem is a cloud object store (S3FileSystem, GcsFileSystem, AzureFileSystem); local (and other real) filesystems still get create_dir=True since they genuinely need the directory to exist (verified separately that create_dir=False raises FileNotFoundError against a local path that doesn't exist yet). Addresses review feedback from cursor[bot] and gemini-code-assist[bot] on both PR ray-project#64923 (this branch) and PR ray-project#64925 (the release-benchmark branch stacked on top, whose cumulative diff includes this file), across four review rounds: - Worker tasks now honor storage_options when building their Parquet write filesystem (previously only the driver's DeltaTable/ create_table_with_add_actions calls received it, so a cloud write with explicit credentials could succeed on the driver's commit while every worker's file write silently used different, ambient credentials, or failed outright). - An empty-dataset APPEND against an already-existing table is now a pure no-op -- no schema resolution, no empty commit added to the Delta log. APPEND against a not-yet-existing table with no resolvable schema still raises (a brand-new table can't be created without knowing its shape). - An empty-dataset OVERWRITE against an already-existing table now falls back to that table's own current schema instead of raising: truncating an existing table is well-defined even when the incoming write has no rows (and therefore no worker-observed schema) to resolve one from. - Partition value/key parsing: a NULL partition value (written by PyArrow's Hive writer as the literal directory name "__HIVE_DEFAULT_PARTITION__") now round-trips back as Python None instead of that literal string; partition column names are now URL-unquoted the same way values already were. - write() now unifies the schema across every non-empty block a single task writes, instead of keeping only the last block's schema (a task can receive multiple blocks with promotable but non-identical schemas). - The worker-side Azure filesystem helper now reads the SAS token from AZURE_STORAGE_SAS_TOKEN, matching ray.data.catalog's vended-credential key, instead of a different name that would never match. - The worker-side GCS filesystem helper no longer leaves GOOGLE_APPLICATION_CREDENTIALS mutated after use -- the env var is a process-wide value that could otherwise leak a task's credential path to other tasks reusing the same warm Ray worker. - The written-file-to-table-root relative path computation now strips leading slashes from both sides before computing the relative path: some cloud filesystems can return a table root without a leading slash from FileSystem.from_uri while written-file paths from that same filesystem still carry one (verified against a real S3-compatible server), and that mismatch previously made the relative path walk upward instead of resolving correctly, corrupting the Delta log. - The worker-side S3 filesystem helper now also accepts AWS_DEFAULT_REGION (a common boto-style key) as a fallback when AWS_REGION isn't set, matching what the driver's commit path (which forwards the full storage_options dict to deltalake) can already honor. - Dataset.write_delta's ray_remote_args parameter is now typed Optional[Dict[str, Any]], matching its actual default of None. Adds a new test suite (test_write_delta.py) covering what this implementation actually supports: APPEND/OVERWRITE round-trips, unsupported-mode rejection, single- and multi-column partitioning, multi-block writes (including cross-block schema unification within a single task), table name/description metadata, the empty-dataset error path (and the empty-APPEND-against-existing-table and empty-OVERWRITE-against-existing-table no-op/truncate paths), storage_options passthrough to both the commit calls and the worker filesystem (S3, Azure, and GCS, including the GCS env-var restore behavior and the AWS_DEFAULT_REGION fallback), null/URL-encoded partition value and key handling, a simulated cross-filesystem leading-slash mismatch in the written-file relative path computation, a regression test for the read_parquet-sourced schema bug above, and a regression test that create_dir is only skipped for cloud object-store filesystems. One pyproject.toml per-file-ignore added (F401) to keep the ported file lint-clean without editing it for that. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Abhishek Verma <abhishek.verma@anyscale.com>
Adds a write_delta counterpart to the existing write_parquet release
benchmark, reusing the same real dataset and harness rather than
inventing new synthetic-data infrastructure:
- read_and_consume_benchmark.py: add a --write-delta consume option,
mirroring the existing --write (write_parquet) option exactly --
same read path, same Benchmark/BenchmarkMetric reporting, just a
different write call.
- release_data_tests.yaml: register a write_delta test entry alongside
write_parquet, reading the same TPCH SF1000 lineitem source
(s3://ray-benchmark-data/tpch/parquet/sf1000/lineitem), so the two
are a direct, apples-to-apples Parquet-vs-Delta write comparison on
identical input data.
- byod_install_deltalake.sh: new post-build script installing
deltalake==1.5.0 (matching the version pinned in
python/requirements_compiled*.txt), needed because -- unlike
parquet/image/tfrecords -- deltalake isn't in the base release-test
image. Mirrors byod_install_pyiceberg.sh's exact shape.
Bug fix (found from an actual CI run of this test, build 101449): the
write itself succeeded (6B rows, 815GiB written), but the subsequent
commit step failed with "OSError: Generic S3 error ... Received
redirect without LOCATION, this normally indicates an incorrectly
configured region". Root cause: the ray-data-write-benchmark bucket is
in us-west-2 (confirmed via its S3 HEAD response's x-amz-bucket-region
header), but deltalake's Rust S3 client defaults to us-east-1 and,
unlike PyArrow's S3FileSystem (which write_parquet uses and which
resolves a bucket's region automatically), doesn't follow a region
redirect -- it needs the region supplied explicitly via
storage_options. Fixed by passing storage_options={"AWS_REGION":
"us-west-2"} to the write_delta call, matching the us-west-2 region
convention already used elsewhere in this same directory
(training_ingest_benchmark.py's S3_IMAGE_AWS_REGION,
multi_node_train_benchmark.py's AWS_REGION env var).
Second bug fix (found from a subsequent CI run, build 101454): a large
distributed write (~1000 concurrent tasks against a brand-new table
root) failed with "AWS Error SLOW_DOWN during PutObject operation"
while writing Parquet files. Root cause (fixed on stack3/pr1-write-delta-core,
which this branch is stacked on): pyarrow.dataset.write_dataset's default
create_dir=True made every task independently PUT the same directory
marker key for the table root, and ~1000 tasks racing on that one key
triggered AWS throttling. Fixed there by skipping create_dir for cloud
object-store filesystems (S3/GCS/Azure), which have no real
directories.
Adds three more write_delta benchmark variants on top of the original
(all reading exclusively from our own internal TPCH test-data bucket,
s3://ray-benchmark-data/tpch/parquet/sf{N}/lineitem -- no external or
public dataset is used anywhere in this change), inspired by (but not
copying any data/code/infra from) delta-io/delta's own benchmark suite,
which tests data loading at both a small "smoke" scale and a large
scale:
- write_delta_smoke: SF10 (small, already used elsewhere in this file)
on a single-node fixed_size_1_cpu_compute.yaml cluster (also already
used elsewhere in this file as a smoke tier), for fast, cheap nightly
feedback on the basic write_delta path without waiting for the
expensive SF1000 run.
- write_delta_overwrite: SF100, exercising SaveMode.OVERWRITE against
an already-existing table (not just a brand-new one, which the
original entry already covers via APPEND). Populates the table once
first (a separately-reported "setup_populate" Benchmark.run_fn step)
so the measured "main" step genuinely replaces existing data of
realistic size, rather than trivially creating an empty table --
SF100 (not SF1000) keeps this variant's now-doubled read/write cost
proportionate to the existing SF1000 test rather than doubling it.
- write_delta_partitioned: SF1000, Hive-partitioned by column08 (this
S3 dataset stores lineitem columns positionally, not by TPCH name;
column08 is l_returnflag per tpch/common.py's TABLE_COLUMNS mapping
-- a 3-value, low-cardinality column, safe at this scale), exercising
the partitioned-write code path at full scale, which the original
entry doesn't touch at all.
read_and_consume_benchmark.py gains two new --write-delta modifier
flags (--write-delta-mode, --write-delta-partition-by; validated to
require --write-delta) to drive these variants without duplicating the
consume_fn logic. No changes needed to benchmark.py: Benchmark.run_fn
already supports multiple named calls per script run (its own docstring
documents exactly this pattern), which is what the overwrite variant's
setup_populate/main split relies on.
Verified locally end-to-end (local Parquet source, local write_delta
destination, no S3/cluster access): the original single-write path by
invoking the modified script's main() directly with a constructed
args.Namespace; the three new variants via a dedicated smoke test
confirming (a) the CLI validation guard rejects the new modifier flags
without --write-delta, (b) SaveMode.OVERWRITE against an
already-populated table results in the original row count, not double
(i.e. genuinely overwritten, not appended), and (c)
--write-delta-partition-by produces the expected Hive-style
column08=<value> directories and reads back correctly.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Abhishek Verma <abhishek.verma@anyscale.com>
abhishekverma-ray
force-pushed
the
stack3/pr3-write-delta-cloud-catalog
branch
from
July 28, 2026 14:33
a86f61c to
0a4603f
Compare
abhishekverma-ray
force-pushed
the
stack3/pr3-write-delta-cloud-catalog
branch
from
July 28, 2026 14:52
0a4603f to
f262cbb
Compare
abhishekverma-ray
force-pushed
the
stack3/pr3-write-delta-cloud-catalog
branch
from
July 28, 2026 15:03
f262cbb to
9531e17
Compare
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
Reviewed by Cursor Bugbot for commit 9531e17. Configure here.
…dential refresh Adds three things to the write_delta prototype, stacked on stack3/pr1 (core) and stack3/pr2 (release benchmark), both currently under review: 1. Deterministic write filenames (prerequisite bug fix). _write_parquet previously generated write_uuid = uuid.uuid4().hex fresh on every call, so any retry produced a different filename than the attempt it was replacing -- silently leaking orphan Parquet files instead of the intended overwrite. Fixed by reusing the same mechanism ParquetDatasink already relies on for exactly this reason: ctx.kwargs[WRITE_UUID_KWARG_NAME], a UUID baked into the task's remote-call kwargs once per write job (at MapOperator.create() time), which survives full Ray task retries, not just retries within one task attempt. 2. Credential refresh on an authentication error (expired session token, expired vended credentials), on both the driver's commit and each worker's Parquet write -- there was previously no retry logic anywhere in this file. Adds is_auth_error()/AUTH_ERROR_PATTERNS and a DeltaConfig (context.py) controlling retry attempts/backoff and a credential_refresh_enabled escape hatch, mirroring IcebergConfig's existing shape. Both _with_retry (driver) and _with_worker_retry (worker) wrap ray._common.retry.call_with_retry with the same catch-detect-refresh-reraise closure pattern, rather than reimplementing a retry loop. Without a catalog, a plain retry is already the refresh: PyArrow's S3FileSystem/AzureFileSystem and deltalake's own object_store both fall back to the standard cloud SDK credential chain (env vars, config files, instance metadata) when given no explicit credentials, and that chain is re-resolved fresh on every new filesystem/DeltaTable construction, which a retry naturally does. No separate credential-fetching helper code was needed or added for this case. 3. A catalog= parameter on Dataset.write_delta, reusing the shared, already-shipped ray.data.catalog module (Catalog, ResolvedSource, ReaderFormat, CatalogAccessMode) -- the same abstraction read_delta/ write_parquet/write_iceberg already consume. Mirrors their exact established pattern: resolve once on the driver, reject a conflicting user-supplied filesystem=, merge storage_options (user's original values always win, including across a later refresh -- see below). The new piece: on an auth error, the datasink re-calls catalog.resolve() for a fresh vended credential, on both the driver and, unlike a sibling branch's (stack2/pr2-write-delta-hardening) driver-only design, on workers too. This requires the Catalog object itself to travel to workers as part of the pickled datasink state (confirmed safe: DatabricksUnityCatalog is picklable, and calling resolve() from a worker is safe since its only ray.init(...) call is gated on Ray not yet being initialized, which is never true inside a running worker task). Scoped to AWS only in this PR (user-confirmed decision): Unity Catalog's resolve() returns an explicit, picklable pyarrow filesystem for AWS-backed Delta writes, which a worker can safely rebuild from. For Azure, resolve() instead delivers the vended SAS token only by mutating this process's own environment (via DatabricksUnityCatalog._apply_env), and never restores that mutation -- calling this from a worker would leave live temporary credentials in that worker process's environment indefinitely, visible to whatever unrelated task Ray reuses that process for next. Worker-side refresh therefore only proceeds when resolve() returns an explicit filesystem (resolved.filesystem is not None); otherwise it declines and the auth error propagates, identical to today's behavior. Fixing this properly for Azure means patching the shared catalog.py's _apply_env (scope/restore its env mutation, populate ResolvedSource.storage_options instead of relying on env vars) -- that's a separate, deliberately deferred follow-up PR, since it's shared infrastructure affecting every existing catalog= consumer (read_delta, write_parquet, write_iceberg, read_parquet), not something to bundle into a write_delta-specific change. Also fixes a genuine bug found while wiring the refresh-preserves- user-overrides behavior: a naive re-merge of catalog-resolved storage_options with the datasink's *current* (possibly already- refreshed-once, i.e. no-longer-original) storage_options would let a stale value for a key the fresh resolve() just updated (e.g. a session token) keep winning on every subsequent refresh, making the refresh a no-op for exactly the fields that needed to change. Fixed by threading the caller's original storage_options through separately (user_storage_options) so every refresh re-merges fresh catalog values with that frozen original, not with the drifting current state. Known, disclosed limitations (not solved in this PR): - No coordination/caching across concurrent workers independently calling catalog.resolve() on their own auth errors -- a burst of token expirations near the same time could all hit the catalog's REST API concurrently. - Shipping the Catalog object to workers means its credential material (token/credential_provider) travels through Ray's object store as part of the task's serialized args, rather than staying strictly driver-side as the sibling branch's design deliberately chose. Adds regression tests: deterministic filenames survive a retry (mocked partial failure then retry, asserting identical basename_template both attempts); is_auth_error classification; driver- and worker-side retry+refresh for both the ambient (no catalog) and catalog-vended paths, including the AWS-succeeds / Azure-shaped-declines scope boundary; the user_storage_options-survives-refresh regression; DeltaConfig.credential_refresh_enabled=False opt-out; catalog= wiring (resolve/merge/reject-conflicting-filesystem); and picklability of a datasink holding a live catalog reference. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Abhishek Verma <abhishek.verma@anyscale.com>
abhishekverma-ray
force-pushed
the
stack3/pr3-write-delta-cloud-catalog
branch
from
July 28, 2026 15:36
9531e17 to
c6b2f4c
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Third PR in the
write_deltastack, on top of #64923 (core) and #64925 (benchmark) —only the top commit
a86f61c896is new here. It adds:1. Deterministic write filenames (prerequisite fix).
_write_parquetgenerated a freshuuid4()per call, so a retry wrote under a different name and leaked orphan Parquetfiles instead of overwriting. Now reuses
ctx.kwargs[WRITE_UUID_KWARG_NAME]— the sameper-job UUID
ParquetDatasinkrelies on, fixed once atMapOperator.create()time, so itsurvives full task retries. This had to land first; everything below adds retries.
2. Credential refresh on an auth error, driver and worker. There was previously no retry
logic in this file, so a write outliving its credential TTL (~1h for STS / Azure AD /
vended tokens) failed partway with no recovery. Adds
is_auth_error()plus_with_retry(driver) and
_with_worker_retry(worker), both wrappingray._common.retry.call_with_retryrather than hand-rolling a loop. Without a catalog aplain retry is the refresh: PyArrow and
deltalake'sobject_storere-resolve thestandard cloud SDK credential chain on every new filesystem/
DeltaTable.3.
catalog=onDataset.write_delta, reusing the already-shippedray.data.catalogmodule (
Catalog,ResolvedSource,ReaderFormat,CatalogAccessMode) exactly asread_delta/write_parquet/write_icebergdo — no new catalog classes or exports. Thenew part: on an auth error the datasink re-calls
catalog.resolve()for a fresh vendedcredential on workers as well as the driver, which means shipping the (picklable)
Catalogin the datasink state. Also fixes a bug found while wiring that up — re-mergingresolved
storage_optionsagainst the current dict let a stale token keep winning forexactly the keys
resolve()had just updated, making every later refresh a no-op; eachrefresh now merges against the caller's original values.
API changes
write_deltagains mutually-exclusivecatalog=/filesystem=kwargs (withcatalog,pathis the table identifier, e.g."main.schema.table"). NewDataContext.delta_config(
DeltaConfig, mirroringIcebergConfig) for retry attempts/backoff plus acredential_refresh_enabledopt-out, withRAY_DATA_DELTA_COMMIT_*env defaults. Noexisting signature or behavior changes;
ray/data/catalog.pyis untouched.Scope and disclosed limitations
Worker-side catalog refresh is AWS-only, by design. Unity Catalog returns an explicit
picklable PyArrow filesystem for AWS, but delivers Azure's SAS token by mutating the
calling process's environment without restoring it — on a reused Ray worker that would leak
live credentials into whatever unrelated task runs there next. Worker refresh therefore
proceeds only when
resolve()returns an explicit filesystem; otherwise the auth errorpropagates, identical to today. Fixing Azure means changing the shared
catalog.py, whichaffects every existing
catalog=consumer (read_delta,write_parquet,write_iceberg,read_parquet) — deferred to its own PR rather than bundled here.Also disclosed: workers call
catalog.resolve()independently with no cross-workercoordination or caching, and the
Catalog's credential material travels through the objectstore as part of the task's serialized args.
Test plan
15 new tests (78 in the file): filename determinism across a retry,
is_auth_errorclassification, driver + worker retry/refresh for both the ambient and catalog-vended paths
including the AWS-succeeds / Azure-declines boundary, the
refresh-preserves-user-
storage_optionsregression, thecredential_refresh_enabled=Falseopt-out,catalog=resolve/merge/conflict-rejection,and datasink picklability with a live catalog.
pytest python/ray/data/tests/datasource/test_write_delta.py— 78 passed.pytest python/ray/data/tests/test_catalog.py— 30 passed (existing suite unaffected).pytest python/ray/data/tests/datasource/test_delta.py— 15 passed (read side, untouched).ci/lint/pyrefly-check.sh, ruff 0.8.4, black 22.10.0 — clean.No live-cloud or live-Unity-Catalog run: the catalog paths use a
_FakeCatalogdouble(subclassing the real
Catalog, like the one intest_catalog.py) plus monkeypatched authfailures.
Related
Stacked on #64923 and #64925. Since both are still open, the diff shown here is cumulative
— only the top commit is new. Not a duplicate: no other open PR adds
catalog=orcredential refresh to
write_delta.AI assistance disclosure
This PR was developed with AI assistance (Claude). All changes were reviewed and are
defended by the submitting human.