Repository navigation
[WIP][SQL] Add a Rust SDK for native data sources over the C ABI - #59520
Draft
gengliangwang wants to merge 6 commits into
Draft
gengliangwang wants to merge 6 commits into
gengliangwang wants to merge 6 commits into
Conversation
This was referenced Oct 15, 2026
Open
- Rename BatchWriterFactory/BatchWriter to DataWriterFactory/DataWriter, following DSv2 and avoiding a one-letter difference from BatchWrite. - Mark WriteMode, FilterSupport, LimitSupport, and Predicate as #[non_exhaustive] so new variants are not breaking changes. - Document the call-order guarantees the SDK enforces for BatchWrite and DataWriter, and remove the sink example's redundant phase tracking. Co-authored-by: Isaac <no-reply@databricks.com>
- Pass the data source name to DataSource::new in a non-exhaustive CreateInfo, and pass writes and executor readers and writers WriteInfo, ReaderContext, and WriterContext, which carry the negotiated schema, so ABI fields can be added. - Add DataSource::init with HostInfo, spark_native_source::log through the host log, and DataSource::INFERS_SCHEMA to leave the schema function unset. - Replace the capability list with a Capabilities set that sources can only narrow, declare truncation with BatchWriteSource::SUPPORTS_TRUNCATE, accept operations in any order, remove the one-argument macro form, and reject panic = "abort" at compile time. - Return a WritePlan from BatchWriteSource::new_write instead of a separate BatchWrite::factory, drop Copy and serde from WriteMode, and add constructors to ScanPlan and WritePlan. - Document that abort can follow a commit whose message fails to serialize, that Drop is the only cleanup when Spark closes a write without abort, and that message i belongs to partition i. - Pass commit messages that cannot be decoded to abort as None, parse each pushed predicate on its own, and include error sources in messages. - Rename BatchSource to BatchReadSource, BatchReaderFactory to PartitionReaderFactory, and the reader and writer aliases to PartitionReader and PartitionWriter, and remove LimitSupport::PerPartition. - Return plain TaskContext IDs, reject task contexts shorter than version 1 of the ABI, and add TaskContext::for_testing. - Add Predicate::literal_value to decode literals by their Spark SQL type. - Hide the macro's support modules in a private __private module. - Pin framework commit 74aa0df: accept only the types of version 1, without intervals or the null type, and check read batches against the nullability of the negotiated schema. Co-authored-by: Isaac <no-reply@databricks.com>
This branch has not been deployed
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
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.
What changes were proposed in this pull request?
Add an experimental Rust SDK for batch reads and writes over the native data source C ABI in #59424, pinned to
74aa0dffc53463034fa6549796c4dd599e9705de. The SDK includes a Range source and a write-only JSONL sink. This draft depends on #59424; the framework changes are not included here. #59442 remains the earlier JNI prototype.The API separates metadata, driver coordination, and executor resources:
DataSourcehandles construction from aCreateInfo(the source name and options), schema discovery, optional per-instance capability restrictions, and an optional per-JVMinithook.BatchReadSourcecreates aScanBuilder, whose consumingplan()returns typed partitions and a serializablePartitionReaderFactory. The factory opens a task-bound Arrow reader from aReaderContext, which carries the negotiated schema and the task.BatchWriteSourcecreates aWritePlan: a driverBatchWritecoordinator and a serializableDataWriterFactory, which opens task-boundDataWriters from aWriterContext. Writers receive owned Arrow batches and return typed commit messages. Driver commit receives the winning message of each partition at its position; abort receivesOption<Message>, withNonefor a message that is missing or cannot be decoded.Method inputs (
CreateInfo,WriteInfo,ReaderContext, andWriterContext) are#[non_exhaustive], andScanPlanandWritePlanhave constructors, so new ABI fields can reach connectors without breaking them.Registration installs matching native callbacks and generates default capabilities:
These are alternative registrations, one per library, with operations in any order. A
BatchWriteSourceopts into overwrite support withconst SUPPORTS_TRUNCATE: bool = true. Write-only connectors do not implement read traits.DataSource::capabilities(&self, registered)defaults to the registeredCapabilities; a connector can narrow them based on its options, and the SDK keeps only registered operations, so a source cannot enable an unregistered one. RemovingBatchWritealso removesTruncate. The macro rejects a library built withpanic = "abort", because the SDK can report a panic as an error only when it unwinds. Streaming callbacks remain absent.Reads retain the framework's planner/reader separation: driver planning and failures are cached, while executor reads restore only serialized values. Writes follow the existing native ABI: source handles can close immediately after writer creation, factory serialization is cached, failed task/driver commits remain abortable, and only close releases writers. Successful task messages and terminal failures are cached to prevent repeated side effects. Both executor readers and writers can retain a lifetime-bound
TaskContextfor cooperative cancellation and can move between threads sequentially. Every host provides the task IDs of version 1 of the ABI, so the SDK rejects a shorter task context and returns plain IDs;TaskContext::for_testing()builds contexts for tests outside Spark.The SDK owns native handles, typed serialization, Arrow import/export, schema and nullability validation, and error/panic conversion, including the messages of error sources. It passes the name of the source to
DataSource::new, so a library installed or listed under several names can choose the source to create.spark_native_source::logwrites to the driver or executor log through Spark, andDataSource::INFERS_SCHEMA = falseleaves the nativeschemafunction unset, so Spark reports its standard error without a user-specified schema. The SDK parses each pushed predicate on its own, so one that it cannot parse stays unsupported instead of failing the scan, andPredicate::literal_valuedecodes literals by their Spark SQL type. Each library exports onlyspark_ds_init; this change does not extend the C ABI or modify JVM classes. The standalone Cargo workspace does not change Maven/SBT or require Rust to build Spark.The Range example demonstrates partition planning, pushdowns, bounded batches, and empty projection. The new sink stages files per task attempt and publishes only winning messages through one atomic manifest replacement. Append preserves prior entries; truncate changes the manifest only at driver commit. Abort cleans the current operation while preserving previously published output. This example requires a local/shared POSIX filesystem, rejects concurrent writes, and documents its crash recovery and retained-file cleanup limitations.
Capability discovery follows the pinned framework's source-instance lifecycle. Native
BATCH_WRITEandTRUNCATEdetermine DSv2 batch-write support during analysis. The framework retains DSv2 read and streaming-write routing flags to avoid V1 fallback, then checks actual native capabilities with effective operation options.Why are the changes needed?
Rust connectors should implement planning, reading, and transactional write coordination without manually managing C function tables, opaque handles, serialized messages, Arrow ownership, or task-context lifetimes. Separating serializable factories from live resources supports independent task attempts, while typed commit messages make the driver/task write boundary explicit. Passing method inputs in non-exhaustive structs mirrors the size-prefixed structs of the ABI, so the API can grow with it.
Deriving capabilities from operation registration avoids maintaining a separate feature list. An optional source-instance restriction, which can only narrow the registered set, still supports cases such as a destination configured as read-only.
Does this PR introduce any user-facing change?
Yes. It adds an experimental, unpublished Rust SDK with batch read, append, and truncate APIs, plus runnable
rust_sdk_rangeandrust_sdk_sinkexamples. Construction, schema discovery, and capability restrictions belong toDataSource;BatchReadSourceandBatchWriteSourcecontain operation-specific methods. This is a prototype API, not a published compatibility commitment.Running the examples in Spark requires a distribution containing #59424 and Java 22 or later. Existing Spark behavior is unchanged. Structured Streaming is not implemented.
How was this patch tested?
From
sql/native/rust, using the available local dependency cache:All 64 tests passed: 38 SDK unit tests, 8 Range tests, 8 sink tests, 7 write ABI integration tests, one capability ABI integration test, one C integration test, and one compile-fail lifetime doctest.
The write ABI tests invoke
spark_ds_initand the actual function table. They cover registration and instance restrictions, the unset schema function, append/truncate, rejection of streaming requests, typed factory/message serialization, the input schema in the writer context, retained Arrow buffers, empty tasks, schema mismatches, live cancellation, write/commit errors and panics, serialization failures, missing and undecodable messages, and repeated-call suppression. The capability ABI test covers library initialization and its failure, host logging, the source name, and narrowed and unregistered capability declarations. The sink tests cover winning versus losing attempts, a message at the position of another partition, append/truncate publication, operation cleanup, concurrent-write rejection, and failed publication preserving the previous manifest.The C test compiles against the unmodified pinned framework header and loads the Range dynamic library. It covers function-table layout, capabilities, independent planners, serialized reads, rejection of task contexts shorter than version 1, cancellation, sequential cross-thread callbacks, and Arrow ownership. Read-side unit tests retain coverage for planning failures and executor retries after driver state is dropped, and cover predicates that do not parse, typed literals, and nullability checks of nested data. The macro's compile errors for missing, unknown, or removed operations and for
panic = "abort", and registration in any order, were checked in a scratch crate.Python syntax, ASCII, line-length, and whitespace checks passed. Both
.sparkpkgpackages were rebuilt and their manifests and embedded library bytes verified. Both dynamic libraries export onlyspark_ds_init. No new external crate dependency was added.Not run: the full Spark smoke scripts, because no local Spark build contains #59424's framework. The included sink smoke script checks append, overwrite, failed-overwrite cleanup, and analysis-time rejection with writes disabled. The Range script checks reading/pushdowns and early rejection of unsupported writes. Rust/C tests validate the native boundary; they do not validate Java FFM or Spark's planner/task integration.
Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-6) and Isaac (Claude Opus 5.5)