Skip to content

[WIP][SQL] Add a Rust SDK for native data sources over the C ABI - #59520

Draft
gengliangwang wants to merge 6 commits into
apache:masterfrom
gengliangwang:codex/rust-native-ffm-sdk
Draft

gengliangwang wants to merge 6 commits into
apache:masterfrom
gengliangwang:codex/rust-native-ffm-sdk

Conversation

@gengliangwang

@gengliangwang gengliangwang commented Oct 15, 2026 •

Copy link
Copy Markdown
Member

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:

  • DataSource handles construction from a CreateInfo (the source name and options), schema discovery, optional per-instance capability restrictions, and an optional per-JVM init hook.
  • BatchReadSource creates a ScanBuilder, whose consuming plan() returns typed partitions and a serializable PartitionReaderFactory. The factory opens a task-bound Arrow reader from a ReaderContext, which carries the negotiated schema and the task.
  • BatchWriteSource creates a WritePlan: a driver BatchWrite coordinator and a serializable DataWriterFactory, which opens task-bound DataWriters from a WriterContext. Writers receive owned Arrow batches and return typed commit messages. Driver commit receives the winning message of each partition at its position; abort receives Option<Message>, with None for a message that is missing or cannot be decoded.

Method inputs (CreateInfo, WriteInfo, ReaderContext, and WriterContext) are #[non_exhaustive], and ScanPlan and WritePlan have constructors, so new ABI fields can reach connectors without breaking them.

Registration installs matching native callbacks and generates default capabilities:

export_data_source!(RangeSource, batch_read);
export_data_source!(JsonSink, batch_write);
export_data_source!(ReadWriteSource, batch_read, batch_write);

These are alternative registrations, one per library, with operations in any order. A BatchWriteSource opts into overwrite support with const SUPPORTS_TRUNCATE: bool = true. Write-only connectors do not implement read traits. DataSource::capabilities(&self, registered) defaults to the registered Capabilities; a connector can narrow them based on its options, and the SDK keeps only registered operations, so a source cannot enable an unregistered one. Removing BatchWrite also removes Truncate. The macro rejects a library built with panic = "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 TaskContext for 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::log writes to the driver or executor log through Spark, and DataSource::INFERS_SCHEMA = false leaves the native schema function 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, and Predicate::literal_value decodes literals by their Spark SQL type. Each library exports only spark_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_WRITE and TRUNCATE determine 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_range and rust_sdk_sink examples. Construction, schema discovery, and capability restrictions belong to DataSource; BatchReadSource and BatchWriteSource contain 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:

cargo build --workspace --locked
cargo test --workspace --locked
cargo clippy --workspace --all-targets --locked -- -D warnings
cargo fmt --all -- --check
RUSTDOCFLAGS='-D warnings' cargo doc --workspace --no-deps --locked

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_init and 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 .sparkpkg packages were rebuilt and their manifests and embedded library bytes verified. Both dynamic libraries export only spark_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)

gengliangwang and others added 5 commits October 15, 2026 23:29
- 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

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant