gengliangwang commented on code in PR #59520: URL: https://github.com/apache/spark/pull/59520#discussion_r4497420166
########## sql/native/rust/README.md: ########## @@ -0,0 +1,467 @@ +<!-- +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with +this work for additional information regarding copyright ownership. +The ASF licenses this file to You under the Apache License, Version 2.0 +(the "License"); you may not use this file except in compliance with +the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +--> + +# Rust native data source SDK prototype + +This experimental SDK implements **batch and micro-batch streaming reads, with admission +control, and batch and streaming writes** over version 1 of the C ABI of native data sources, +[`spark_datasource.h`][header], which Spark 4.4.0 introduces. Spark calls the library through +the Java Foreign Function and Memory (FFM) API, so running a library in Spark requires Java 22 +or newer on the driver and executors. The [guide of native data sources][guide] describes the +ABI. + +The Rust and C tests run independently of Spark and Java. The workspace is not +part of Maven/SBT and introduces no Rust requirement for building Spark. This is +an API prototype, not a published crate or a compatibility commitment. + +The [Range connector][range] is a separate workspace member under +`examples/src/main/rust/native-data-source`, with a C integration test, packaging +helper, and Spark smoke script. The [file sink example][sink] is a separate write-only +connector. All crates share this directory's lockfile and `target` directory. Cargo commands +select the SDK by default; use `--workspace` to include the examples. + +## Connector API and lifecycle + +A connector implements [`DataSource`](src/lib.rs) for construction, schema discovery, +its operations, and optional capability restrictions. `DataSource::new` receives a +`CreateInfo` with the name of the data source and its options. Spark can create and +immediately drop a source just to discover capabilities; schema inference or an +operation need not follow. + +Four associated types name the operations that the source implements, and +`Unsupported` stands for one that it does not implement: + +```rust +impl DataSource for RangeSource { + type ScanBuilder = RangeScan; // batch reads, created by new_scan + type MicroBatchStream = RangeStream; // streaming reads, by new_micro_batch_stream + type BatchWrite = Unsupported; // no batch writes + type StreamingWrite = Unsupported; // no streaming writes + // new, schema, new_scan, and new_micro_batch_stream +} +``` + +Read support uses three pieces: + +* `DataSource::new_scan`: create an independent scan builder from a `ScanInfo`, which + carries the schema to read. +* `ScanBuilder`: negotiate the limit and the columns to read, then consume itself + in `plan()` to return a typed `ScanPlan<Factory>`. +* `PartitionReaderFactory`: serialize executor configuration and open one lazy reader, + an iterator of Arrow batches, from a typed partition and a `ReaderContext`, which + carries the schema to read and a borrowed `TaskContext`. + +Streaming reads reuse the executor side: `DataSource::new_micro_batch_stream` creates a +`MicroBatchStream` from a `MicroBatchStreamInfo`, which carries the columns to read and the +location of the checkpoint of the source, and the stream plans each micro-batch as a +`ScanPlan` of a `PartitionReaderFactory`; see +[Streaming reads and writes](#streaming-reads-and-writes). + +Write support is independent: a write-only connector sets `ScanBuilder` and +`MicroBatchStream` to `Unsupported` and implements `DataSource::new_write`, `BatchWrite`, +`DataWriterFactory`, and `DataWriter`. Streaming writes add +`DataSource::new_streaming_write` and `StreamingWrite`, with the same task writers. The +write lifecycle is described below. + +The inputs of these methods are `#[non_exhaustive]` structs, so that new ABI fields can +reach connectors without breaking them: `*Info` structs on the driver (`CreateInfo`, +`ScanInfo`, `MicroBatchStreamInfo`, `LatestOffsetInfo`, `WriteInfo`, `StreamingWriteInfo`, +and `HostInfo`), and `*Context` structs, which carry the live `TaskContext`, on the executors +(`ReaderContext` and `WriterContext`). `ScanPlan` has a constructor for the same reason. + +The factory and partition types implement `Serialize` and `Deserialize`. +Connections, pointers, and open readers remain in their creating process. +Include a snapshot identifier in the factory or partitions when retrying a task +must read the same snapshot. Multiple task attempts open independent readers. + +```text +Driver: + DataSource::new_scan -> ScanBuilder -> ScanPlan { factory, partitions } + | | + mutable negotiation immutable serialized descriptors + | +Executor, per task: v + restored factory + partition + ReaderContext -> Arrow reader -> Arrow C Stream +``` + +The Rust API maps to the framework's driver-only `ColumnarBatchPlanner` and its +`ColumnarPlannedBatch` containing a serializable `PartitionReaderFactory`. +`ColumnarScanBuilder.build()` builds a single scan and transfers planner +ownership; it does not plan partitions. `ColumnarBatchScan` plans on the first +request for partitions or a reader factory and caches success or failure. + +For a native source, `NativeBatchPlanner` pushes the limit and the columns to read, +each at most once and in this order, then calls `partitions` once, which returns the +partitions and the state of the reader, and closes the native planner handle. The Rust +`ScanBuilder::plan()` consumes its builder in that call, and the SDK encodes the factory +as the state. After a call fails, Spark calls no other function on the handle than its +close function. Batch planner creation remains lazy for scans that do not read in batch. + +The JVM's `NativePartitionReaderFactory` carries only the library location, +source name, read schema, and serialized state. The native executor `read` entry +point gets the kind of the read, batch or micro-batch, and the schema to read with +the state and the partition, so it restores the Rust factory from those bytes without +creating or accessing the driver registry, and passes the schema to the reader in +`ReaderContext`. Closing the driver handles cannot invalidate factory or partition +bytes already copied to Spark, including for retried tasks. + +The reader owns its executor resources and cannot borrow its temporary factory. +It may borrow `TaskContext` until the returned Arrow stream is released. Every host +provides the task IDs of version 1 of the ABI, so the SDK rejects a shorter task +context, as it rejects any structure shorter than in version 1. `is_interrupted()` checks +Spark's live cancellation callback, not a snapshot of its initial value; for a reader, it +is also true once the task ended. Check it during long operations; blocking I/O still +requires timeouts. `TaskContext::for_testing()`, with its `with_*` methods and +`interrupted_by`, builds a context with chosen IDs and an optional interruption flag for +tests outside Spark. + +Export the source from the connector's `cdylib` crate: + +```rust +spark_native_source::export_data_source!(RangeSource); +``` + +The macro exports only `spark_ds_init` and returns Spark's versioned function +table. The connector uses safe Rust traits without C pointers. Rust trait objects +never cross the binary boundary between libraries. + +`DataSource::init` runs once in each JVM when Spark loads the library, with the +Spark version in `HostInfo`; a failure prevents Spark from using the library and is +logged. The macros of the `log` crate, which the SDK re-exports as +`spark_native_source::log`, write to the driver or executor log through Spark from any +thread, including in the library's dependencies, and drop messages outside Spark, such +as in unit tests. The SDK sets the level of the `log` crate to that of the log of Spark, +which it asks with `SparkDsHost.is_log_enabled` when Spark loads the library and when an +operation or a task starts, so it also follows a level that changes while Spark runs. + +A source that cannot infer its schema, such as a sink of any input, sets +`DataSource::INFERS_SCHEMA` to false. The SDK then leaves the native `schema` +function unset, so Spark requires a user-specified schema and reports its standard +error without one. + +A library that implements several data sources is installed or listed in its +package under each of their names. Spark passes the name in `CreateInfo::name`, so +one exported `DataSource` can choose the source to create; the macro exports a +single Rust type per library. + +The implementation separates public connector traits, driver handle state in +`bridge.rs`, Arrow ownership in `stream.rs`, the writers in `write.rs`, C bindings in +`abi.rs`, and the function table in `export.rs`. The macro reaches them through a +hidden `__private` module, which is not part of the API. + +## Source capabilities + +The associated types of `DataSource` determine the operations of the source: batch +reads unless `ScanBuilder` is `Unsupported`, micro-batch reads unless `MicroBatchStream` +is, batch writes unless `BatchWrite` is, and streaming writes unless `StreamingWrite` is. +The SDK installs the native callbacks of these operations, leaves the others null, and +derives the default capability set from them, so the macro takes only the source type. +A `BatchWrite` or a `StreamingWrite` opts into truncation with +`const SUPPORTS_TRUNCATE: bool = true`, which declares `BATCH_TRUNCATE`, for the save mode +"overwrite" of a batch write, or `STREAMING_TRUNCATE`, for the output mode "complete" of a +streaming write. Its `new_*` method then gets `WriteMode::Truncate` for them; without it +the write only appends. Every stream also declares the native capability +`SPARK_DS_ADMISSION_CONTROL`, whose default read limit, all the available data, reads as +a stream without it does. The macro also rejects a library built with `panic = "abort"`, +because the SDK can report a panic as an error only when it unwinds. + +Most connectors do not implement `capabilities`. A connector can restrict the +supported set for the current options by overriding this `DataSource` method: + +```rust +fn capabilities(&self) -> Capabilities { + if self.read_only { + // The SDK also drops the truncations, which require their writes. + Capabilities::all().without(Capabilities::BATCH_WRITE | Capabilities::STREAMING_WRITE) + } else { + Capabilities::all() + } +} +``` + +The SDK evaluates this once during `create` and keeps only implemented operations, +so a source can narrow its capabilities but never enable one that it does not +implement, and keeps each truncation only with its write. +An empty set permits schema discovery and close, but rejects operation creation +before calling the connector. The declaration stays fixed for that instance. +The default declaration is `Capabilities::all()`, every operation of the SDK. Limit and +column pruning remain per-scan contracts. + +Spark also calls `create` and `close` during query analysis to discover table +capabilities. A probe uses the table's options, while an operation may override +them with query options and checks capabilities again on a new source instance. +Keep `DataSource::new` and `capabilities` lightweight and safe to drop +immediately; do not assume that `schema` or `new_scan` will follow. Executor +resources belong in `PartitionReaderFactory::open`. + +Spark propagates the native `BATCH_WRITE` and `BATCH_TRUNCATE` flags to +`Table.capabilities()`, as `BATCH_WRITE` and `TRUNCATE`, so a read-only source rejects +batch writes during analysis, before Spark computes their input. Its DSv2 read and +streaming-write flags remain present to avoid V1 fallback; the native flags still reject +unsupported reads or streaming writes during planning. Thus the DSv2 table's flags are not +a literal copy of the native mask declared by the Rust source, and an unsupported +streaming read or write fails when its query starts or plans its first micro-batch, after +`start()` returns. + +## Batch write lifecycle + +```text +Driver: + DataSource::new_write -> BatchWrite + | writer_factory() ^ commit(messages) / abort(messages) + serialized factory | typed messages + | | +Executor, per task attempt: v | + DataWriterFactory::open(WriterContext) + -> DataWriter::write(RecordBatch)* + -> commit() / abort() -> Drop +``` + +`DataSource::new_write(WriteInfo)` validates the input schema and returns an +independent driver coordinator, a `BatchWrite`. Its `writer_factory` returns the +serializable executor configuration of its task writers, which the SDK encodes as the +state of the writer that native `create_writer` returns. Restoring an executor writer +uses only those bytes and the schema that Spark passes with them, without creating a +driver registry. + +Each `DataWriterFactory::open` receives a `WriterContext` with the input schema and +the task, and returns a task-bound writer. `write` takes ownership of a +`RecordBatch`, so its Arrow buffers can be retained after the callback, until the writer +is dropped: they belong to the data writer of Spark, which frees them when it closes the +writer. The SDK +moves the incoming Arrow array, which has the schema that the data writer got, and +releases imported resources on failures. Input slices need not have zero offsets; +the zero-offset restriction applies to arrays exported to Spark by readers. When Spark +writes the batches of a native scan without converting them to rows, it copies their +values as the scan returned them, without the checks of a conversion, such as the +precision of a decimal. Check cancellation during lengthy writes and use I/O timeouts. + +Task `commit()` returns a typed, serde-serializable message, the `Message` of the +`BatchWrite`. Driver `commit` +receives one winning message per partition, at the position of its partition; +driver `abort` receives `Vec<Option<Message>>`, with `None` for a partition whose +message did not arrive or cannot be decoded, so a corrupted message cannot prevent +cleanup. A failed task or driver commit leaves the same object available for abort. +A task commit whose message cannot be serialized counts as failed, so `abort` can +follow a `commit` that returned a message. Spark calls each method at most once, never +writes after a failure or a commit, and never aborts a successful commit, as the header +documents, so connectors do not need to track their own phase; the `# Lifecycle` +sections of `BatchWrite` and `DataWriter` list the exact call-order guarantees. Drop only +releases resources. + +Stage data per task attempt and publish it only at driver commit. In particular, +`WriteMode::Truncate` must leave existing data intact until that commit. Retries +and speculative attempts may stage data for the same partition; publish only the +attempts whose messages Spark passes to commit. Executor loss can prevent cleanup, +including after task commit, so connectors need a strategy for orphaned data. +Spark provides no batch query ID: generate and serialize an operation identity +when needed. Task-attempt IDs alone are not globally unique across applications. + +When `writer_factory` or the serialization of its factory fails, `create_writer` fails +and the SDK drops the coordinator without commit or abort, so Drop must also release what +only that write staged. Spark also creates writers late, after prerequisite AQE shuffle +stages, but input computation can still fail during write tasks. Constructors must +therefore avoid irreversible changes. The epoch ID tells the task writers of batch writes, +whose epoch ID is `-1`, from those of micro-batches. + +The [file sink][sink] demonstrates append/truncate using staged JSONL files and +atomic publication of a manifest on a local or shared filesystem. Its documented +single-writer assumption and orphan cleanup limitations are example constraints, +not transaction guarantees provided by the SDK. + +## Streaming reads and writes + +Micro-batch streaming reads reuse the executor side of batch scans: + +```text +Driver, for each run of a query: + DataSource::new_micro_batch_stream -> MicroBatchStream + initial_offset(), default_read_limit() + latest_offset(LatestOffsetInfo) -> Offset, stored as JSON in the checkpoint + plan(start, end) -> ScanPlan { factory, partitions }, for each micro-batch + commit(end) +Executor, per task: + restored factory + partition + ReaderContext -> Arrow reader -> Arrow C Stream +``` + +Spark creates a stream when a query starts or restarts from its checkpoint, and closes it +when the query stops. It keeps the initial offset in the checkpoint, and asks for it only +when the checkpoint does not have it yet. For each micro-batch, it asks for the latest +offset after the end of the previous micro-batch, and plans the data up to that offset. +The executors read the partitions of a micro-batch exactly as those of a batch scan: the +two share `read`, which gets the kind of the read. `MicroBatchStreamInfo::schema` has the +columns to read, which are all the columns of the schema until Spark prunes the columns of +streaming scans. + +`MicroBatchStream::default_read_limit` limits the data of each micro-batch, for example +to `ReadLimit::MaxRows` from an option of the source; the default reads all the available +data. Spark passes the limit to `latest_offset` in `LatestOffsetInfo::limit`, and with +`Trigger.AvailableNow`, the end of the run in `LatestOffsetInfo::run_end`, which the Review Comment: Done in e933ef2cf382: The README now says `run_end` is passed only with `Trigger.AvailableNow` and a `ReadLimit::MaxRows` limit. ########## examples/src/main/rust/native-data-source/src/lib.rs: ########## @@ -0,0 +1,596 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! A deterministic range source: `id` in [start, end), and `label = "row-{id}"`. +//! All C ABI and Arrow C Stream plumbing is supplied by the SDK's export macro. + +use std::collections::BTreeSet; +use std::sync::Arc; + +use arrow_array::StringArray; +use arrow_array::{ArrayRef, Int64Array, RecordBatch, RecordBatchOptions}; +use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; +use serde::{Deserialize, Serialize}; +use spark_native_source::{ + CreateInfo, DataSource, LatestOffsetInfo, LimitSupport, MicroBatchStream, MicroBatchStreamInfo, + Options, PartitionReader, PartitionReaderFactory, ReadLimit, ReaderContext, Result, + ScanBuilder, ScanInfo, ScanPlan, Unsupported, +}; + +pub struct RangeSource { + start: i64, + end: i64, + num_partitions: usize, + batch_size: usize, + max_rows_per_trigger: Option<u64>, +} + +pub struct RangeScan { + start: i64, + end: i64, + num_partitions: usize, + limit: Option<usize>, + factory: RangeReaderFactory, +} + +/// Reader configuration travels to each executor; it contains no live objects. +/// Spark passes the columns to read to the reader, so the factory omits them. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RangeReaderFactory { + batch_size: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RangePartition { + start: i64, + end: i64, +} + +/// A streaming read of the range. The range is fixed, so all its rows are available when +/// the query starts: the first micro-batch reads them, unless `maxRowsPerTrigger` limits each +/// micro-batch, and later ones find no new data. +pub struct RangeStream { + start: i64, + end: i64, + num_partitions: usize, + max_rows_per_trigger: Option<u64>, + factory: RangeReaderFactory, +} + +/// The offset of a streaming read: the next row to read. Spark stores it in the checkpoint +/// as `{"next":13}`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct RangeOffset { + next: i64, +} + +impl DataSource for RangeSource { + type ScanBuilder = RangeScan; + type MicroBatchStream = RangeStream; + // A read-only source. + type BatchWrite = Unsupported; + type StreamingWrite = Unsupported; + + fn new(info: CreateInfo) -> Result<Self> { + let options = info.options; + for key in options.keys() { + if !matches!( + key.as_str(), + "start" | "end" | "numpartitions" | "batchsize" | "maxrowspertrigger" + ) { + return Err(format!("Unknown range option: {key}").into()); + } + } + let start = option(&options, "start", 0_i64)?; + let end = option(&options, "end", 100_i64)?; + let num_partitions = option(&options, "numpartitions", 2_usize)?; + let batch_size = option(&options, "batchsize", 1024_usize)?; + let max_rows_per_trigger = options + .get("maxrowspertrigger") + .map(|_| option(&options, "maxrowspertrigger", 0_u64)) + .transpose()?; + if end < start { + return Err("end must be greater than or equal to start".into()); + } + if num_partitions == 0 || batch_size == 0 || max_rows_per_trigger == Some(0) { + return Err("numPartitions, batchSize and maxRowsPerTrigger must be positive".into()); + } + Ok(Self { + start, + end, + num_partitions, + batch_size, + max_rows_per_trigger, + }) + } + + fn schema(&self) -> Result<SchemaRef> { + output_schema(&["id".into(), "label".into()]) + } + + fn new_scan(&self, info: ScanInfo) -> Result<Self::ScanBuilder> { + read_columns(&info.schema)?; + Ok(RangeScan { + start: self.start, + end: self.end, + num_partitions: self.num_partitions, + limit: None, + factory: RangeReaderFactory { + batch_size: self.batch_size, + }, + }) + } + + fn new_micro_batch_stream(&self, info: MicroBatchStreamInfo) -> Result<Self::MicroBatchStream> { + read_columns(&info.schema)?; + Ok(RangeStream { + start: self.start, + end: self.end, + num_partitions: self.num_partitions, + max_rows_per_trigger: self.max_rows_per_trigger, + factory: RangeReaderFactory { + batch_size: self.batch_size, + }, + }) + } +} + +impl PartitionReaderFactory for RangeReaderFactory { + type Partition = RangePartition; + + fn open<'task>( + &self, + partition: RangePartition, + context: ReaderContext<'task>, + ) -> Result<PartitionReader<'task>> { + let batch_size = self.batch_size; + if batch_size == 0 || partition.start > partition.end { + return Err("Invalid serialized range plan".into()); + } + // The schema to read has the columns that the query needs, in order, with the + // types that the scan validated. + let schema = context.schema; + let columns: Vec<String> = schema.fields().iter().map(|f| f.name().clone()).collect(); + output_schema(&columns)?; + let task = context.task; + let mut current = partition.start; + let batches = std::iter::from_fn(move || { + if current == partition.end { + return None; + } + if task.is_interrupted() { + current = partition.end; + return Some(Err(ArrowError::ComputeError( + "Range read was cancelled".into(), + ))); + } + let next = + (i128::from(current) + batch_size as i128).min(i128::from(partition.end)) as i64; + let num_rows = (i128::from(next) - i128::from(current)) as usize; + let arrays: Vec<ArrayRef> = columns + .iter() + .map(|column| match column.as_str() { + "id" => Arc::new(Int64Array::from_iter_values(current..next)) as ArrayRef, + "label" => Arc::new(StringArray::from_iter_values( + (current..next).map(|id| format!("row-{id}")), + )) as ArrayRef, + _ => unreachable!("column names were validated before reading"), + }) + .collect(); + current = next; + // Explicit row counts also preserve COUNT(*) with a zero-column projection. + Some(RecordBatch::try_new_with_options( + Arc::clone(&schema), + arrays, + &RecordBatchOptions::new().with_row_count(Some(num_rows)), + )) + }); + Ok(Box::new(batches)) + } +} + +// The readers get the columns to read in their context, so the scan does not prune columns. +impl ScanBuilder for RangeScan { + type Factory = RangeReaderFactory; + + fn push_limit(&mut self, limit: usize) -> Result<LimitSupport> { + self.limit = Some(limit); + Ok(LimitSupport::Global) + } + + fn plan(self) -> Result<ScanPlan<RangeReaderFactory>> { + // Apply a global limit by shortening the range, then split it. + let lower = i128::from(self.start); + let upper = i128::from(self.end); + let upper = self + .limit + .map_or(upper, |limit| upper.min(lower + limit as i128)); + let partitions = split(lower, upper, self.num_partitions); + Ok(ScanPlan::new(self.factory, partitions)) + } +} + +impl MicroBatchStream for RangeStream { + type Offset = RangeOffset; + // Micro-batches are read as batch scans are. + type Factory = RangeReaderFactory; + + fn initial_offset(&mut self) -> Result<RangeOffset> { + Ok(RangeOffset { next: self.start }) + } + + fn default_read_limit(&self) -> Result<ReadLimit> { + Ok(self + .max_rows_per_trigger + .map_or(ReadLimit::AllAvailable, ReadLimit::MaxRows)) + } + + fn latest_offset(&mut self, info: LatestOffsetInfo<RangeOffset>) -> Result<RangeOffset> { + // All the rows of the range are available, within the limit and the end of the run. + let mut next = self.end; + if let ReadLimit::MaxRows(rows) = info.limit { + let rows = i64::try_from(rows).unwrap_or(i64::MAX); + next = next.min(info.start.next.saturating_add(rows)); + } + if let Some(run_end) = info.run_end { + next = next.min(run_end.next); + } + Ok(RangeOffset { next }) Review Comment: Done in e933ef2cf382: `RangeStream::latest_offset` now rejects a start offset outside [start, end] before Spark logs anything, and there is a test for it. ########## examples/src/main/rust/native-data-source/src/lib.rs: ########## @@ -0,0 +1,596 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! A deterministic range source: `id` in [start, end), and `label = "row-{id}"`. +//! All C ABI and Arrow C Stream plumbing is supplied by the SDK's export macro. + +use std::collections::BTreeSet; +use std::sync::Arc; + +use arrow_array::StringArray; +use arrow_array::{ArrayRef, Int64Array, RecordBatch, RecordBatchOptions}; +use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; +use serde::{Deserialize, Serialize}; +use spark_native_source::{ + CreateInfo, DataSource, LatestOffsetInfo, LimitSupport, MicroBatchStream, MicroBatchStreamInfo, + Options, PartitionReader, PartitionReaderFactory, ReadLimit, ReaderContext, Result, + ScanBuilder, ScanInfo, ScanPlan, Unsupported, +}; + +pub struct RangeSource { + start: i64, + end: i64, + num_partitions: usize, + batch_size: usize, + max_rows_per_trigger: Option<u64>, +} + +pub struct RangeScan { + start: i64, + end: i64, + num_partitions: usize, + limit: Option<usize>, + factory: RangeReaderFactory, +} + +/// Reader configuration travels to each executor; it contains no live objects. +/// Spark passes the columns to read to the reader, so the factory omits them. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RangeReaderFactory { + batch_size: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RangePartition { + start: i64, + end: i64, +} + +/// A streaming read of the range. The range is fixed, so all its rows are available when +/// the query starts: the first micro-batch reads them, unless `maxRowsPerTrigger` limits each +/// micro-batch, and later ones find no new data. +pub struct RangeStream { + start: i64, + end: i64, + num_partitions: usize, + max_rows_per_trigger: Option<u64>, + factory: RangeReaderFactory, +} + +/// The offset of a streaming read: the next row to read. Spark stores it in the checkpoint +/// as `{"next":13}`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct RangeOffset { + next: i64, +} + +impl DataSource for RangeSource { + type ScanBuilder = RangeScan; + type MicroBatchStream = RangeStream; + // A read-only source. + type BatchWrite = Unsupported; + type StreamingWrite = Unsupported; + + fn new(info: CreateInfo) -> Result<Self> { + let options = info.options; + for key in options.keys() { + if !matches!( + key.as_str(), + "start" | "end" | "numpartitions" | "batchsize" | "maxrowspertrigger" + ) { + return Err(format!("Unknown range option: {key}").into()); + } + } + let start = option(&options, "start", 0_i64)?; + let end = option(&options, "end", 100_i64)?; + let num_partitions = option(&options, "numpartitions", 2_usize)?; + let batch_size = option(&options, "batchsize", 1024_usize)?; + let max_rows_per_trigger = options + .get("maxrowspertrigger") + .map(|_| option(&options, "maxrowspertrigger", 0_u64)) + .transpose()?; + if end < start { + return Err("end must be greater than or equal to start".into()); + } + if num_partitions == 0 || batch_size == 0 || max_rows_per_trigger == Some(0) { Review Comment: Done in e933ef2cf382: The Range example now rejects `numPartitions` above 100000 and `batchSize` above 1000000, with tests and updated docs for the options. ########## sql/native/rust/src/bridge.rs: ########## @@ -0,0 +1,876 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Driver-only handles of sources, scans, streams, and writes, and executor readers restored +//! from serialized values. + +use std::any::Any; +use std::collections::HashMap; +use std::ffi::CString; +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; + +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_schema::SchemaRef; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; + +use crate::write::{Batch, DriverWrite, DriverWriter}; +use crate::{ + abi, BatchWrite, Capabilities, CreateInfo, DataSource, LatestOffsetInfo, LimitSupport, + MicroBatchStream, MicroBatchStreamInfo, PartitionReaderFactory, ReadLimit, ReaderContext, + Result, ScanBuilder, ScanInfo, ScanPlan, StreamingWrite, StreamingWriteInfo, TaskContext, + WriteInfo, WriteMode, +}; + +#[path = "stream.rs"] +pub(crate) mod stream; + +const STATE_VERSION: u32 = 3; + +type HandleMap<T> = Mutex<HashMap<i64, Arc<Mutex<T>>>>; + +/// One registry per loaded plugin. Provider code runs under a per-handle lock, +/// so independent sources, scans, and streams do not serialize behind a global lock. +pub struct Bridge<S: DataSource> { + next_handle: AtomicI64, + sources: HandleMap<SourceState<S>>, + // The builder of each scan, until Spark plans the scan. + scans: HandleMap<Option<S::ScanBuilder>>, + streams: HandleMap<S::MicroBatchStream>, + writers: HandleMap<Box<dyn DriverWrite>>, +} + +impl<S: DataSource> Default for Bridge<S> { + fn default() -> Self { + Self { + next_handle: AtomicI64::new(1), + sources: Mutex::new(HashMap::new()), + scans: Mutex::new(HashMap::new()), + streams: Mutex::new(HashMap::new()), + writers: Mutex::new(HashMap::new()), + } + } +} + +/// The partitions of a scan or a micro-batch, and the state of its reader: its factory. +pub struct EncodedPlan { + pub state: Vec<u8>, + pub partitions: Vec<Vec<u8>>, +} + +impl<S: DataSource> Bridge<S> { + // Whether the source implements each operation, whose native callbacks the SDK installs. + + pub const READ: bool = <S::ScanBuilder as ScanBuilder>::SUPPORTED; + + pub const MICRO_BATCH_READ: bool = <S::MicroBatchStream as MicroBatchStream>::SUPPORTED; + + pub const WRITE: bool = <S::BatchWrite as BatchWrite>::SUPPORTED; + + pub const STREAMING_WRITE: bool = <S::StreamingWrite as StreamingWrite>::SUPPORTED; + + const CAPABILITIES: Capabilities = Capabilities::implemented(&[ + (Capabilities::BATCH_READ, Self::READ), + (Capabilities::MICRO_BATCH_READ, Self::MICRO_BATCH_READ), + (Capabilities::BATCH_WRITE, Self::WRITE), + (Capabilities::STREAMING_WRITE, Self::STREAMING_WRITE), + ( + Capabilities::BATCH_TRUNCATE, + <S::BatchWrite as BatchWrite>::SUPPORTS_TRUNCATE, + ), + ( + Capabilities::STREAMING_TRUNCATE, + <S::StreamingWrite as StreamingWrite>::SUPPORTS_TRUNCATE, + ), + ]); + + fn insert<T>(&self, handles: &HandleMap<T>, value: T) -> Result<i64> { + let handle = self + .next_handle + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1)) + .map_err(|_| "native handle IDs exhausted")?; + lock(handles).insert(handle, Arc::new(Mutex::new(value))); + Ok(handle) + } + + pub fn create(&self, info: CreateInfo) -> Result<(i64, u64)> { + self.register_source(S::new(info)?) + } + + fn register_source(&self, source: S) -> Result<(i64, u64)> { + let capabilities = Self::effective(source.capabilities()); + let handle = self.insert( + &self.sources, + SourceState { + source, + capabilities, + }, + )?; + // Every stream supports admission control: its default read limit is all the + // available data unless it overrides MicroBatchStream::default_read_limit. + let admission_control = if capabilities.contains(Capabilities::MICRO_BATCH_READ) { + abi::ADMISSION_CONTROL + } else { + 0 + }; + Ok((handle, capabilities.bits() | admission_control)) + } + + /// Keeps the implemented operations of a declaration, so that a source can narrow its + /// operations but never add one, and each truncation only with its write. + fn effective(declared: Capabilities) -> Capabilities { + let mut capabilities = declared.intersection(Self::CAPABILITIES); + for (write, truncate) in [ + (Capabilities::BATCH_WRITE, Capabilities::BATCH_TRUNCATE), + ( + Capabilities::STREAMING_WRITE, + Capabilities::STREAMING_TRUNCATE, + ), + ] { + if !capabilities.contains(write) { + capabilities = capabilities.without(truncate); + } + } + capabilities + } + + pub fn schema(&self, handle: i64) -> Result<SchemaRef> { + let source = get(&self.sources, handle)?; + let schema = lock(&source).source.schema()?; + stream::validate_schema(&schema)?; + Ok(schema) + } + + pub fn create_reader(&self, handle: i64, schema: SchemaRef) -> Result<i64> { + stream::validate_schema(&schema)?; + let source = get(&self.sources, handle)?; + let builder = { + let source = lock(&source); + require(source.capabilities, Capabilities::BATCH_READ, "batch reads")?; + source.source.new_scan(ScanInfo::new(schema))? + }; + self.insert(&self.scans, Some(builder)) + } + + pub fn push_limit(&self, handle: i64, limit: usize) -> Result<bool> { + let scan = get(&self.scans, handle)?; + let mut scan = lock(&scan); + Ok(builder(&mut scan)?.push_limit(limit)? == LimitSupport::Global) + } + + pub fn prune_columns(&self, handle: i64, columns: &[String]) -> Result<()> { + let scan = get(&self.scans, handle)?; + let mut scan = lock(&scan); + builder(&mut scan)?.prune_columns(columns) + } + + /// Plans a scan, which consumes its builder. + pub fn plan(&self, handle: i64) -> Result<EncodedPlan> { + let scan = get(&self.scans, handle)?; + let builder = lock(&scan) + .take() + .ok_or("the scan has already been planned")?; + encode_plan(builder.plan()?) + } + + pub fn close_reader(&self, handle: i64) -> Result<()> { + remove(&self.scans, handle); + Ok(()) + } + + pub fn create_stream_reader(&self, handle: i64, info: MicroBatchStreamInfo) -> Result<i64> { + stream::validate_schema(&info.schema)?; + let source = get(&self.sources, handle)?; + let stream = { + let source = lock(&source); + require( + source.capabilities, + Capabilities::MICRO_BATCH_READ, + "micro-batch streaming reads", + )?; + source.source.new_micro_batch_stream(info)? + }; + self.insert(&self.streams, stream) + } + + pub fn initial_offset(&self, handle: i64) -> Result<String> { + let stream = get(&self.streams, handle)?; + let offset = lock(&stream).initial_offset()?; + encode_offset(&offset) + } + + pub fn default_read_limit(&self, handle: i64) -> Result<ReadLimit> { + let stream = get(&self.streams, handle)?; + let limit = lock(&stream).default_read_limit(); + limit + } + + /// Decodes the offsets of a request, and encodes the latest offset. + pub fn latest_offset(&self, handle: i64, info: LatestOffsetInfo<String>) -> Result<String> { + let stream = get(&self.streams, handle)?; + let info = LatestOffsetInfo { + start: decode_offset(&info.start)?, + limit: info.limit, + run_end: info.run_end.as_deref().map(decode_offset).transpose()?, + }; + let offset = lock(&stream).latest_offset(info)?; + encode_offset(&offset) + } + + pub fn stream_partitions(&self, handle: i64, start: &str, end: &str) -> Result<EncodedPlan> { + let stream = get(&self.streams, handle)?; + let plan = lock(&stream).plan(decode_offset(start)?, decode_offset(end)?)?; + encode_plan(plan) + } + + pub fn commit_offset(&self, handle: i64, end: &str) -> Result<()> { + let stream = get(&self.streams, handle)?; + let end = decode_offset(end)?; + let result = lock(&stream).commit(end); + result + } + + pub fn close_stream_reader(&self, handle: i64) -> Result<()> { + remove(&self.streams, handle); + Ok(()) + } + + /// Creates a batch write, or the write of a micro-batch of the streaming query `query_id`, + /// with the state of its task writers. + pub fn create_writer( + &self, + handle: i64, + schema: SchemaRef, + mode: WriteMode, + query_id: Option<String>, + ) -> Result<(i64, Vec<u8>)> { + stream::validate_schema(&schema)?; + let source = get(&self.sources, handle)?; + let (writer, state) = { + let source = lock(&source); + let capabilities = source.capabilities; + let truncate = mode == WriteMode::Truncate; + match query_id { + None => { + require(capabilities, Capabilities::BATCH_WRITE, "batch writes")?; + if truncate { + let overwrite = "the save mode overwrite"; + require(capabilities, Capabilities::BATCH_TRUNCATE, overwrite)?; + } + let write = source.source.new_write(WriteInfo::new(schema, mode))?; + DriverWriter::create(Batch(write))? + } + Some(query_id) => { + require( + capabilities, + Capabilities::STREAMING_WRITE, + "streaming writes", + )?; + if truncate { + let complete = "the output mode complete"; + require(capabilities, Capabilities::STREAMING_TRUNCATE, complete)?; + } + let info = StreamingWriteInfo::new(schema, mode, query_id); + let write = source.source.new_streaming_write(info)?; + DriverWriter::create(write)? + } + } + }; + Ok((self.insert(&self.writers, writer)?, state)) + } + + pub fn commit(&self, handle: i64, epoch_id: i64, messages: &[Option<&[u8]>]) -> Result<()> { + let writer = get(&self.writers, handle)?; + let result = lock(&writer).commit(epoch_id, messages); + result + } + + pub fn abort(&self, handle: i64, epoch_id: i64, messages: &[Option<&[u8]>]) -> Result<()> { + let writer = get(&self.writers, handle)?; + let result = lock(&writer).abort(epoch_id, messages); + result + } + + pub fn close_writer(&self, handle: i64) -> Result<()> { + remove(&self.writers, handle); + Ok(()) + } + + pub fn close_source(&self, handle: i64) -> Result<()> { + remove(&self.sources, handle); + Ok(()) + } +} + +/// Opens the executor reader of a partition with factory `F`, using only serialized values. +/// No source, scan, stream, or driver handle registry is created or consulted. +/// +/// # Safety +/// `task` must satisfy TaskContext::from_raw and remain valid until the +/// returned stream is released. The stream's consumer must follow the +/// Arrow C Stream ownership protocol and keep the plugin loaded. +pub unsafe fn open_reader<F: PartitionReaderFactory>( + task: *const abi::SparkDsTaskContext, + schema: SchemaRef, + state: &[u8], + partition: &[u8], +) -> Result<FFI_ArrowArrayStream> { + stream::validate_schema(&schema)?; + let factory: F = decode(state)?; + let task = unsafe { TaskContext::from_raw(task)? }; + let reader = factory.open(decode(partition)?, ReaderContext::new(schema.clone(), task))?; + Ok(unsafe { stream::export(reader, &schema) }) +} + +fn require(capabilities: Capabilities, capability: Capabilities, operation: &str) -> Result<()> { + if capabilities.contains(capability) { + Ok(()) + } else { + Err(format!("source does not support {operation}").into()) + } +} + +/// The builder of a scan that is not planned yet. +fn builder<B>(scan: &mut Option<B>) -> Result<&mut B> { + scan.as_mut() + .ok_or_else(|| "the scan has already been planned".into()) +} + +/// Encodes the factory of a plan, as the state of its reader, and each of its partitions. +fn encode_plan<F: PartitionReaderFactory>(plan: ScanPlan<F>) -> Result<EncodedPlan> { + Ok(EncodedPlan { + state: encode(&plan.factory)?, + partitions: plan.partitions.iter().map(encode).collect::<Result<_>>()?, + }) +} + +fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> { + // A provider panic becomes an ABI error. Recover the lock so close can + // still remove and release the handle after that failed call. + mutex.lock().unwrap_or_else(|error| error.into_inner()) +} + +fn get<T>(handles: &HandleMap<T>, handle: i64) -> Result<Arc<Mutex<T>>> { + lock(handles) + .get(&handle) + .cloned() + .ok_or_else(|| format!("unknown or closed native handle: {handle}").into()) +} + +fn remove<T>(handles: &HandleMap<T>, handle: i64) { + let value = lock(handles).remove(&handle); + // Provider destructors run after releasing the registry lock. Repeated + // closes are harmless, and IDs are never reused. + drop(value); +} + +struct SourceState<S> { + source: S, + capabilities: Capabilities, +} + +/// Offsets are compact JSON on a single line, which Spark stores in the checkpoint as they +/// are, so they have no envelope: the library owns their format. +fn encode_offset<T: Serialize>(offset: &T) -> Result<String> { + Ok(serde_json::to_string(offset)?) +} + +fn decode_offset<T: DeserializeOwned>(json: &str) -> Result<T> { + serde_json::from_str(json) + .map_err(|error| format!("cannot decode the offset {json}: {error}").into()) +} + +#[derive(Serialize, Deserialize)] +struct Envelope<T> { + version: u32, + value: T, +} + +pub(crate) fn encode<T: Serialize>(value: &T) -> Result<Vec<u8>> { + Ok(serde_json::to_vec(&Envelope { + version: STATE_VERSION, + value, + })?) +} + +pub(crate) fn decode<T: DeserializeOwned>(bytes: &[u8]) -> Result<T> { + let envelope: Envelope<T> = serde_json::from_slice(bytes)?; + if envelope.version != STATE_VERSION { + return Err(format!("unsupported native state version: {}", envelope.version).into()); + } + Ok(envelope.value) +} + +pub(crate) fn provider_call<T>( + action: impl FnOnce() -> Result<T>, +) -> std::result::Result<T, String> { + catch_unwind(AssertUnwindSafe(|| { + action().map_err(|error| describe(&*error)) + })) + .unwrap_or_else(|payload| Err(panic_message(payload))) +} + +pub(crate) fn drop_provider<T>(value: T) { + // Close has no error channel and must never unwind across the native ABI. + if let Err(payload) = catch_unwind(AssertUnwindSafe(|| drop(value))) { + let _ = panic_message(payload); + } +} Review Comment: Done in e933ef2cf382: `drop_provider` now logs a caught panic with `log::warn!`, and handle removal for sources, scans and streams now goes through `drop_provider`, so those panics are logged too. ########## sql/native/rust/src/lib.rs: ########## @@ -0,0 +1,760 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Experimental SDK for native data sources in Rust, over version 1 of the C ABI of Spark, +//! `spark_datasource.h`. +//! +//! Implement [`DataSource`], naming the types of the operations that it implements: a +//! [`ScanBuilder`] for batch reads, a [`MicroBatchStream`] for micro-batch streaming reads, a +//! [`BatchWrite`] for batch writes, and a [`StreamingWrite`] for streaming writes, or +//! [`Unsupported`] for an operation that it does not implement. Then invoke +//! [`export_data_source!`] once in a `cdylib` crate. The SDK +//! owns handles, plan/message serialization, C ABI adaptation, and Arrow import/export. +//! Rust traits stay inside the plugin; they are not an ABI between dynamic libraries. +//! +//! Methods receive their inputs in `#[non_exhaustive]` structs, so that new ABI fields +//! can reach connectors without breaking them: `*Info` structs on the driver +//! ([`CreateInfo`], [`ScanInfo`], [`MicroBatchStreamInfo`], [`LatestOffsetInfo`], +//! [`WriteInfo`], [`StreamingWriteInfo`], and [`HostInfo`]), and `*Context` structs, which +//! carry the live [`TaskContext`], on the executors ([`ReaderContext`] and [`WriterContext`]). + +use std::collections::BTreeMap; + +use arrow_array::RecordBatch; +use arrow_schema::{ArrowError, SchemaRef}; +use serde::de::DeserializeOwned; +use serde::Serialize; + +mod abi; +mod bridge; +mod capabilities; +mod export; +mod host; +mod task; +mod unsupported; +mod write; + +/// Implementation details of [`export_data_source!`] and of the SDK's own tests, not an API. +#[doc(hidden)] +pub mod __private { + pub use crate::bridge::Bridge; + pub use crate::export::{library, Exported}; + + pub mod abi { + pub use crate::abi::*; + } +} + +pub use capabilities::Capabilities; +pub use host::HostInfo; +/// The `log` crate, whose macros write to the log of the Spark driver or executor, with +/// the path of the library, once Spark initializes the library. +/// +/// Any thread can call them, including in the dependencies of the library. Outside +/// Spark, for example in unit tests, their messages are dropped. The SDK sets the level of +/// the `log` crate to that of the log of Spark, such as `Debug` when Spark logs debug +/// messages, so the library does not set it. +pub use log; Review Comment: Done in e933ef2cf382: The SDK now re-exports `arrow_array` and `arrow_schema`, and the README says a connector must use the same Arrow major and that an Arrow major upgrade is a breaking change of the SDK. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
