gengliangwang commented on code in PR #59520: URL: https://github.com/apache/spark/pull/59520#discussion_r4497418476
########## examples/src/main/rust/native-data-source/tests/c_abi_smoke.rs: ########## @@ -0,0 +1,89 @@ +// 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. + +#![cfg(any(target_os = "linux", target_os = "macos"))] + +use std::error::Error; +use std::fs; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::time::{SystemTime, UNIX_EPOCH}; + +struct TestDirectory(PathBuf); + +impl Drop for TestDirectory { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } +} + +#[test] +fn c_abi_arrow_round_trip() -> Result<(), Box<dyn Error>> { + let unique = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos(); + let directory = TestDirectory( + std::env::temp_dir().join(format!("spark-native-c-{}-{unique}", std::process::id())), + ); + fs::create_dir(&directory.0)?; + + let manifest = Path::new(env!("CARGO_MANIFEST_DIR")); + // The header of the Spark source tree, as a library would copy it. + let header = + manifest.join("../../../../../sql/core/src/main/resources/org/apache/spark/sql/datasource"); + let harness = directory.0.join("c_abi_smoke"); + let compiler = std::env::var_os("CC").unwrap_or_else(|| "cc".into()); + let mut command = Command::new(compiler); + command + .args(["-std=c11", "-Wall", "-Wextra", "-Werror", "-pthread"]) + .arg("-I") + .arg(header) + .arg(manifest.join("tests/c_abi_smoke.c")) + .arg("-o") + .arg(&harness); + if cfg!(target_os = "linux") { + command.arg("-ldl"); + } + let compiled = command.output()?; + assert!( + compiled.status.success(), + "C ABI harness did not compile:\n{}", + String::from_utf8_lossy(&compiled.stderr) + ); + + let executable = std::env::current_exe()?; + let profile = executable + .parent() + .and_then(Path::parent) + .ok_or("Missing Cargo profile directory")?; + let library = profile.join(format!( + "{}spark_datasource_rust_sdk_range{}", + std::env::consts::DLL_PREFIX, + std::env::consts::DLL_SUFFIX + )); + assert!( + library.is_file(), + "Run cargo build --workspace first: {}", + library.display() + ); Review Comment: Done in e933ef2cf382: The C test now loads the library from the deps directory that cargo test rebuilds, and the build-first step and its assert message are gone from the READMEs. ########## 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; +pub use task::TaskContext; +pub use unsupported::Unsupported; + +pub type Error = Box<dyn std::error::Error + Send + Sync>; +pub type Result<T> = std::result::Result<T, Error>; +pub type Options = BTreeMap<String, String>; + +// Data sources: Spark creates a DataSource for each operation, and its associated types +// name the operations that the source implements. + +/// What Spark passes to create a data source. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct CreateInfo { + /// The name of the data source, as listed in the manifest of its package, or in + /// lower case for an installed library. A library that implements several data + /// sources, installed or listed under each of their names, can choose by this name. + pub name: String, + /// The options, with keys in lower case. A capability probe uses the options of the + /// table; an operation may override them with the options of its query. + pub options: Options, +} + +impl CreateInfo { + /// The information of a source created outside Spark, for example in a test. + pub fn new(name: impl Into<String>, options: Options) -> Self { + Self { + name: name.into(), + options, + } + } +} + +/// Driver-side source configuration, schema discovery, and the operations of a source. +/// +/// The associated types name the operations that the source implements, and +/// [`Unsupported`] stands for one that it does not implement: +/// +/// ```rust,ignore +/// impl DataSource for RangeSource { +/// type ScanBuilder = RangeScan; +/// type MicroBatchStream = RangeStream; +/// type BatchWrite = Unsupported; +/// type StreamingWrite = Unsupported; +/// ... +/// } +/// ``` +/// +/// The SDK installs native callbacks and reports capabilities only for the implemented +/// operations, and calls the `new_*` method of an operation only for them. +/// +/// Spark creates a source for each schema inference, scan, stream, and write, including +/// the write of each micro-batch, and may create one only to inspect its capabilities, +/// then drop it. Keep construction lightweight. Scan builders, streams, and writes cannot +/// borrow their source, which Spark drops first. +pub trait DataSource: Send + Sized + 'static { + fn new(info: CreateInfo) -> Result<Self>; + + /// Called once in each JVM when Spark loads the library, before it creates any + /// source. On failure, Spark does not use the library, and the error is logged. + fn init(_host: &HostInfo) -> Result<()> { + Ok(()) + } + + // Schema discovery. + + /// Whether [`schema`](Self::schema) infers the schema of the data. When false, the + /// SDK leaves the native function unset: Spark then requires a user-specified + /// schema, reports its standard error without one, and never calls `schema`. + const INFERS_SCHEMA: bool = true; + + /// Returns the schema of the data, when the user does not specify one. + fn schema(&self) -> Result<SchemaRef> { + Err( + "DataSource::schema is not implemented: implement it, or set \ + DataSource::INFERS_SCHEMA to false to require a user-specified schema" + .into(), + ) + } + + // Batch reads. + + /// The scan builder of batch reads, which [`new_scan`](Self::new_scan) creates, or + /// [`Unsupported`]. + type ScanBuilder: ScanBuilder; + + /// Creates an independent scan of [`ScanInfo::schema`], which the user may have + /// specified; reject unsupported schemas. + fn new_scan(&self, _info: ScanInfo) -> Result<Self::ScanBuilder> { + Err( + "DataSource::new_scan is not implemented: implement it, or set \ + DataSource::ScanBuilder to Unsupported" + .into(), + ) + } + + // Micro-batch streaming reads. + + /// The stream of micro-batch streaming reads, which + /// [`new_micro_batch_stream`](Self::new_micro_batch_stream) creates, or [`Unsupported`]. + type MicroBatchStream: MicroBatchStream; + + /// Creates the stream of a streaming query that reads the columns of + /// [`MicroBatchStreamInfo::schema`]; reject unsupported schemas. + fn new_micro_batch_stream( + &self, + _info: MicroBatchStreamInfo, + ) -> Result<Self::MicroBatchStream> { + Err( + "DataSource::new_micro_batch_stream is not implemented: implement it, or set \ + DataSource::MicroBatchStream to Unsupported" + .into(), + ) + } + + // Batch writes. + + /// The coordinator of batch writes, which [`new_write`](Self::new_write) creates, or + /// [`Unsupported`]. + type BatchWrite: BatchWrite; + + /// Validate the input schema and create an independent write. Construction may be + /// followed only by Drop, without commit or abort, so it must not make irreversible + /// changes. The SDK calls this only while the instance's capabilities include + /// `BATCH_WRITE`, and passes [`WriteMode::Truncate`] only when they include + /// `BATCH_TRUNCATE`, which the write declares with [`BatchWrite::SUPPORTS_TRUNCATE`]. + fn new_write(&self, _info: WriteInfo) -> Result<Self::BatchWrite> { + Err( + "DataSource::new_write is not implemented: implement it, or set \ + DataSource::BatchWrite to Unsupported" + .into(), + ) + } + + // Streaming writes. + + /// The coordinator of the write of each micro-batch of a streaming query, which + /// [`new_streaming_write`](Self::new_streaming_write) creates, or [`Unsupported`]. + type StreamingWrite: StreamingWrite; + + /// Validate the input schema and create the write of a micro-batch. As with + /// [`new_write`](Self::new_write), construction may be followed only by Drop, so it must + /// not make irreversible changes. The SDK calls this only while the instance's + /// capabilities include `STREAMING_WRITE`, and passes [`WriteMode::Truncate`] only when + /// they include `STREAMING_TRUNCATE`, which the write declares with + /// [`StreamingWrite::SUPPORTS_TRUNCATE`]. + fn new_streaming_write(&self, _info: StreamingWriteInfo) -> Result<Self::StreamingWrite> { + Err( + "DataSource::new_streaming_write is not implemented: implement it, or set \ + DataSource::StreamingWrite to Unsupported" + .into(), + ) + } + + // Instance restrictions. + + /// Returns the operations of this instance. The SDK keeps only those that the source + /// implements, so the default, [`Capabilities::all`], keeps them all. Override this to + /// restrict an instance, for example for a read-only destination. The SDK evaluates + /// it once when Spark creates the source. An instance without operations still + /// supports schema discovery. + fn capabilities(&self) -> Capabilities { + Capabilities::all() + } +} + +// Batch reads: a ScanBuilder negotiates pushdown and plans the scan on the driver, and a +// PartitionReaderFactory opens a reader of each partition on the executors. + +/// What Spark passes to create a scan. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct ScanInfo { + /// The schema to read: the schema of the source, or one that the user specified. + pub schema: SchemaRef, +} + +impl ScanInfo { + /// The information of a scan created outside Spark, for example in a test. + pub fn new(schema: SchemaRef) -> Self { + Self { schema } + } +} + +/// Negotiates pushdown on the driver, then plans the scan. +/// +/// Spark calls [`push_limit`](Self::push_limit) and [`prune_columns`](Self::prune_columns), +/// each at most once and in this order, then [`plan`](Self::plan), and drops the builder. +/// After a call fails, it makes no other call. It drops a builder without planning it when it +/// does not read the scan, for example to only explain a query. Only the serialized factory +/// and partitions reach the executors. +pub trait ScanBuilder: Send + Sized + 'static { + type Factory: PartitionReaderFactory; + + /// False only for [`Unsupported`], from which the SDK derives the capabilities. + #[doc(hidden)] + const SUPPORTED: bool = true; + + /// Count rows. Spark pushes a limit only when the query does not filter the rows, since + /// it filters them after reading. + fn push_limit(&mut self, _limit: usize) -> Result<LimitSupport> { + Ok(LimitSupport::Unsupported) + } + + /// Tells the scan the top-level columns that the query needs, in order, so that it can + /// plan to read only them. There may be none, for example for `SELECT count(*)`: the + /// batches then only have a number of rows. Readers get these columns in + /// [`ReaderContext::schema`] whether or not the scan implements this method. + fn prune_columns(&mut self, _columns: &[String]) -> Result<()> { + Ok(()) + } + + fn plan(self) -> Result<ScanPlan<Self::Factory>>; +} + +/// Whether a connector limits the rows produced by the whole scan. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum LimitSupport { + /// The scan may return more rows than the limit. A connector may still cap the + /// rows of each partition on its own: Spark applies its limit after reading. + Unsupported, + /// The whole scan returns at most the limit. Spark still applies its limit. + Global, +} + +/// A driver-created plan containing only serializable values. +/// +/// The SDK serializes the reader factory once and each partition independently. +/// Include snapshot identifiers when needed for repeatable task attempts. Do not +/// put pointers, open files, connections, or process-local resources here. +#[non_exhaustive] +pub struct ScanPlan<F: PartitionReaderFactory> { + pub factory: F, + pub partitions: Vec<F::Partition>, +} + +impl<F: PartitionReaderFactory> ScanPlan<F> { + pub fn new(factory: F, partitions: Vec<F::Partition>) -> Self { + Self { + factory, + partitions, + } + } +} + +/// Serializable configuration from which executors open partition readers. +/// +/// The factory carries configuration or a snapshot descriptor; open connections and +/// other live resources in [`open`](Self::open). Spark passes the schema to read with +/// each partition, so the factory need not carry it. +pub trait PartitionReaderFactory: Serialize + DeserializeOwned + Send + 'static { + type Partition: Serialize + DeserializeOwned + Send + 'static; + + /// Opens a lazy reader of the batches of `partition`, with the columns of + /// [`ReaderContext::schema`]. + /// + /// - Return batches lazily, with bounded memory. + /// - The reader owns its resources and cannot borrow this short-lived factory. It + /// may keep `context` until Spark releases the stream. + /// - Spark may call the reader on different threads, never concurrently, so it must + /// be `Send` but not necessarily `Sync`. + /// - Check [`TaskContext::is_interrupted`] during long operations, and use timeouts + /// for blocking I/O. + /// - Batches must own their buffers, because Spark may release the stream before + /// its batches. + /// + /// The SDK checks Spark-compatible Arrow types, the schema to read and its + /// nullability, and zero array offsets recursively. It reports incompatible layouts + /// instead of copying buffers. A batch may have more columns than the schema to read, + /// such as all the columns of the source: the SDK keeps those of the schema, by name. + /// Batches may use any supported Arrow encoding of a Spark type, such as `LargeUtf8` + /// for a string, but the first batch determines the Arrow schema of the stream, which + /// every later batch must have. + fn open<'task>( + &self, + partition: Self::Partition, + context: ReaderContext<'task>, + ) -> Result<PartitionReader<'task>>; +} + +/// What an executor reader receives with its partition. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct ReaderContext<'task> { + /// The schema to read: the columns that the scan reads, in order, which are those of + /// [`ScanBuilder::prune_columns`] for a batch scan, and of + /// [`MicroBatchStreamInfo::schema`] for a micro-batch. + pub schema: SchemaRef, + /// The task that reads the partition. + pub task: TaskContext<'task>, +} + +impl<'task> ReaderContext<'task> { + /// The context of a reader opened outside Spark, for example in a test. + pub fn new(schema: SchemaRef, task: TaskContext<'task>) -> Self { + Self { schema, task } + } +} + +/// The lazy batches of one partition, which [`PartitionReaderFactory::open`] returns. +/// Any Arrow `RecordBatchReader`, such as a Parquet reader, is such an iterator. +pub type PartitionReader<'task> = + Box<dyn Iterator<Item = std::result::Result<RecordBatch, ArrowError>> + Send + 'task>; + +// Micro-batch streaming reads: a MicroBatchStream finds new data and plans each micro-batch on +// the driver, and the PartitionReaderFactory of each plan reads its partitions on the executors, +// as in batch reads. + +/// What Spark passes to create the stream of a streaming query. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct MicroBatchStreamInfo { + /// The columns to read: those of the schema of the source, or of one that the user + /// specified, or only the top-level columns of it that the query needs. The readers of + /// the micro-batches get it in [`ReaderContext::schema`]. + pub schema: SchemaRef, + /// A directory in the checkpoint of the query that is reserved for this source, such as + /// `file:/tmp/checkpoint/sources/0`, where the stream may keep what it needs to recover + /// from a failure. It stays the same when the query restarts. It is a fully qualified + /// Hadoop path, whose characters are not escaped as in a URI, on a file system that the + /// library may not be able to access, such as `s3a`. Spark keeps what it needs there under + /// the names that start with `_spark`, which the library must not use. + pub checkpoint_location: String, +} + +impl MicroBatchStreamInfo { + /// The information of a stream created outside Spark, for example in a test. + pub fn new(schema: SchemaRef, checkpoint_location: impl Into<String>) -> Self { + Self { + schema, + checkpoint_location: checkpoint_location.into(), + } + } +} + +/// Finds new data and plans the micro-batches of a streaming query, on the driver. +/// +/// Spark creates a stream when a streaming query starts, or restarts from its checkpoint, +/// and drops it when the query stops. For each micro-batch, it asks for the +/// [`latest_offset`](Self::latest_offset), and when that differs from the end of the +/// previous micro-batch, [`plan`](Self::plan)s the data in between. Only the serialized +/// factory and partitions of each plan reach the executors, which read them as those of a +/// batch scan. +/// +/// # Offsets +/// +/// Spark stores offsets in the checkpoint of the query, as the JSON of +/// [`Offset`](Self::Offset), and compares them by that JSON, so serialize a position to the +/// same JSON each time: for example, use structs and `BTreeMap` rather than `HashMap`. A +/// restarted query passes the offsets of its checkpoint, which an earlier version of the +/// library may have written, to `latest_offset`, `plan`, and `commit`. +/// +/// # Admission control +/// +/// [`default_read_limit`](Self::default_read_limit) limits the data of each micro-batch, +/// typically with an option of the source. Spark passes the limit to `latest_offset`, except +/// when a micro-batch reads all the available data, as with `Trigger.Once`. With +/// `Trigger.AvailableNow`, a run reads the data that was available when it started: with a +/// limit, in several micro-batches, whose latest offsets must not go beyond +/// [`LatestOffsetInfo::run_end`]. +/// +/// # Lifecycle +/// +/// - Calls never run concurrently. +/// - After a call fails, including with a panic, Spark makes no other call: it stops the +/// query, and drops the stream. +/// - After a restart, Spark may plan the last range of the query again, so `plan` must +/// return partitions that read the same data for the same range. +pub trait MicroBatchStream: Send + 'static { + /// A position in the data, such as the next record of each partition of a log. + type Offset: Serialize + DeserializeOwned; + + /// The reader factory of the micro-batches, which can be that of batch scans. + type Factory: PartitionReaderFactory; + + /// False only for [`Unsupported`], from which the SDK derives the capabilities. + #[doc(hidden)] + const SUPPORTED: bool = true; + + /// Returns the offset to start from when the query has no checkpoint: the first + /// micro-batch reads the data after it. Spark keeps it in the checkpoint, and asks only + /// when the checkpoint does not have it yet, so a query that restarts before its first + /// micro-batch is committed starts from the same offset. + fn initial_offset(&mut self) -> Result<Self::Offset>; + + /// Returns the limit on the data of each micro-batch. Spark asks once, when the query + /// starts. The default, [`ReadLimit::AllAvailable`], reads all the available data in + /// each micro-batch. + fn default_read_limit(&self) -> Result<ReadLimit> { + Ok(ReadLimit::AllAvailable) + } + + /// Returns the end of the next micro-batch: the offset of the most recent data after + /// [`LatestOffsetInfo::start`], within its limit, and not beyond its end of the run. + /// Return `start` when there is no new data. When there is, return an offset after + /// `start`, even if the next unit of data, such as a file, is larger than the limit: + /// Spark sees no new data otherwise, and never reads further. + /// + /// With a limit, Spark also asks without a limit for some micro-batches, to report the + /// progress of the query, so do not count the calls as micro-batches, or reserve data in + /// them. + fn latest_offset(&mut self, info: LatestOffsetInfo<Self::Offset>) -> Result<Self::Offset>; + + /// Plans the micro-batch that reads the data after `start`, up to and including `end`. + fn plan(&mut self, start: Self::Offset, end: Self::Offset) -> Result<ScanPlan<Self::Factory>>; + + /// Tells the stream that Spark has processed the data up to and including `end` and + /// will not read it again, so that the stream can release it. + fn commit(&mut self, _end: Self::Offset) -> Result<()> { + Ok(()) + } +} + +/// A limit on the data of a micro-batch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum ReadLimit { + /// All the available data. + AllAvailable, + /// At most this number of rows, at least 1. + MaxRows(u64), +} + +/// What Spark passes to ask a stream for its latest offset. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct LatestOffsetInfo<O> { + /// The offset that the next micro-batch starts after: the end of the previous one, or + /// the initial offset. + pub start: O, + /// The limit on the data after `start`: that of + /// [`MicroBatchStream::default_read_limit`], or [`ReadLimit::AllAvailable`] when Spark + /// reads all the available data, such as for the progress of the query or with + /// `Trigger.Once`. + pub limit: ReadLimit, + /// With `Trigger.AvailableNow` and a limit, the end of the run: the latest offset must not + /// go beyond it, so that the run reads the data that was available when it started. It is + /// not the end of the next micro-batch. `None` otherwise. + pub run_end: Option<O>, +} + +impl<O> LatestOffsetInfo<O> { + /// A request without a limit, made outside Spark, for example in a test. Set + /// [`limit`](Self::limit) and [`run_end`](Self::run_end) for other requests. + pub fn new(start: O) -> Self { + Self { + start, + limit: ReadLimit::AllAvailable, + run_end: None, + } + } +} + +// Batch writes: a BatchWrite coordinates the write on the driver, and a DataWriterFactory +// opens a DataWriter for each task attempt on the executors. + +/// What Spark passes to create a batch write. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct WriteInfo { + /// The schema of the data to write. + pub schema: SchemaRef, + /// How the write changes the destination. + pub mode: WriteMode, +} + +impl WriteInfo { + /// The information of a write created outside Spark, for example in a test. + pub fn new(schema: SchemaRef, mode: WriteMode) -> Self { + Self { schema, mode } + } +} + +/// How a batch write changes the destination. +/// +/// New modes may be added. The SDK passes a mode only when the source instance +/// has the matching capability, so a source never receives a mode it did not +/// opt in to; a wildcard arm can return an error. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub enum WriteMode { + Append, + /// Replace existing data at driver commit, never at writer creation. + Truncate, +} + +/// A driver coordinator for one batch write, which also creates the serializable +/// configuration of its task writers. +/// +/// Keep irreversible publication, including removing overwritten data, in commit. +/// Drop releases resources; it must not publish data or perform fallible flushing. +/// +/// # Lifecycle +/// +/// Spark calls the methods in this order, so implementations need not track their own +/// phase or guard against repeated calls: +/// +/// - `writer_factory` runs once, right after [`DataSource::new_write`], before any task +/// opens a writer. When it fails, or the SDK cannot serialize its factory, the write +/// fails and the coordinator is dropped without `abort`, so Drop must also release +/// what only this write staged. +/// - `commit` runs at most once, when every task succeeded. The SDK fails the write +/// without calling it when a task message cannot be decoded. +/// - `abort` runs at most once, never after a successful `commit`. It follows a failed +/// `commit`, or a failed task, also when computing the data to write failed. +/// - Drop runs exactly once, possibly without `commit` or `abort`. +/// - A panic in any method is reported as that method's failure. +/// +/// Message `i` of `commit` and `abort` belongs to input partition `i`. +pub trait BatchWrite: Send + 'static { + /// The configuration from which each task attempt opens its writer. + type Factory: DataWriterFactory<Message = Self::Message>; + + /// The message with which a task attempt reports its staged output to the driver. + type Message: Serialize + DeserializeOwned + Send + 'static; + + /// Whether the write supports [`WriteMode::Truncate`], for the save mode "overwrite", + /// which adds the `BATCH_TRUNCATE` capability. + const SUPPORTS_TRUNCATE: bool = false; + + /// False only for [`Unsupported`], from which the SDK derives the capabilities. + #[doc(hidden)] + const SUPPORTED: bool = true; + + /// Returns the configuration of the task writers, which the SDK serializes for them. + /// Include a unique operation ID when needed; Spark supplies no batch query ID. + fn writer_factory(&self) -> Result<Self::Factory>; + + /// Publish only the winning attempt from each input partition. Task retries + /// and speculative attempts can stage other data that must not be published. + fn commit(&mut self, messages: Vec<Self::Message>) -> Result<()>; + + /// Abort after input, task, or driver commit failure. None stands for a partition + /// whose commit message did not reach the driver, or cannot be decoded. + /// Executor loss can leave orphaned staged data; plan for its cleanup. + fn abort(&mut self, messages: Vec<Option<Self::Message>>) -> Result<()>; +} + +/// Serializable configuration for independent task-attempt writers, which +/// [`BatchWrite::writer_factory`] or [`StreamingWrite::writer_factory`] creates. +/// +/// Open connections and files in open, not in the factory. A returned writer +/// cannot borrow this temporary factory, but may retain the context until +/// close. Callbacks can move between threads sequentially; writers must be Send. +pub trait DataWriterFactory: Serialize + DeserializeOwned + Send + 'static { + type Message: Serialize + DeserializeOwned + Send + 'static; + + fn open<'task>( + &self, + context: WriterContext<'task>, + ) -> Result<PartitionWriter<'task, Self::Message>>; +} + +/// What a task writer receives when it opens. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct WriterContext<'task> { + /// The schema of the batches to write, as in [`WriteInfo::schema`]. + pub schema: SchemaRef, + /// The task attempt that writes. + pub task: TaskContext<'task>, + /// The epoch of the micro-batch of a streaming write, or `None` for a batch write. + pub epoch_id: Option<i64>, +} + +impl<'task> WriterContext<'task> { + /// The context of a batch writer opened outside Spark, for example in a test. Set + /// [`epoch_id`](Self::epoch_id) for a streaming writer. + pub fn new(schema: SchemaRef, task: TaskContext<'task>) -> Self { + Self { + schema, + task, + epoch_id: None, + } + } +} + +/// One task attempt's writer. Spark drops it after commit or abort. +/// +/// # Lifecycle +/// +/// Spark calls the methods in this order, so implementations need not track their own +/// phase or guard against repeated calls: +/// +/// - `write` runs zero or more times, before `commit` or `abort`, and never after a +/// call has failed. +/// - `commit` runs at most once, never after a failed `write` or after `abort`. +/// - `abort` runs at most once, never after a `commit` whose message reached Spark. +/// It may follow a failed `write` or `commit`, including a `commit` that returned +/// a message that the SDK could not serialize, so it must handle a task attempt +/// that already finished its output. +/// - Drop runs exactly once when Spark closes the writer, possibly without +/// `commit` or `abort`. +/// - The writer of a task attempt that Spark kills, such as a speculative one, may still +/// write, abort, and drop after the driver committed or aborted the write. +/// - A panic in any method is reported as that method's failure. +pub trait DataWriter: Send { + type Message; + + /// Takes ownership of an imported Arrow batch, which may be retained after this call, + /// but not after the writer is dropped: its buffers belong to the data writer of Spark, + /// which frees them once it closes the writer, so do not pass the batch to a thread or a + /// value that outlives the writer. It has the schema of [`WriterContext::schema`], and at most + /// `spark.sql.execution.arrow.maxRecordsPerBatch` rows. 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. Process batches incrementally and check task cancellation. + fn write(&mut self, batch: RecordBatch) -> Result<()>; Review Comment: Done in e933ef2cf382: The `DataWriter::write` rustdoc and the README now give the ABI requirement (release the batch before Spark closes the writer) and what happens if it is broken: Spark logs an error, and the memory stays allocated until the batch is dropped. ########## sql/native/rust/src/abi.rs: ########## @@ -0,0 +1,733 @@ +// 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. + +//! Bindings and adapters for version 1 of `spark_datasource.h`. + +use std::cell::RefCell; +use std::ffi::{c_char, c_int, c_void, CStr, CString}; +use std::mem::{offset_of, size_of}; +use std::sync::Arc; + +pub use arrow_array::ffi::FFI_ArrowArray as ArrowArray; +pub use arrow_array::ffi_stream::FFI_ArrowArrayStream as ArrowArrayStream; +pub use arrow_schema::ffi::FFI_ArrowSchema as ArrowSchema; +use arrow_schema::{Schema, SchemaRef}; + +use crate::bridge::{c_message, provider_call}; +use crate::host::{self, LogEnabledFunction, LogFunction}; +use crate::{CreateInfo, HostInfo, LatestOffsetInfo, Options, ReadLimit, Result, WriteMode}; + +pub const ABI_VERSION: i32 = 1; +pub const BATCH_READ: u64 = 1; +pub const MICRO_BATCH_READ: u64 = 2; +pub const BATCH_WRITE: u64 = 4; +pub const STREAMING_WRITE: u64 = 8; +pub const BATCH_TRUNCATE: u64 = 16; +pub const STREAMING_TRUNCATE: u64 = 32; +pub const ADMISSION_CONTROL: u64 = 64; +pub const READ_BATCH: i32 = 0; +pub const READ_MICRO_BATCH: i32 = 1; +pub const READ_ALL_AVAILABLE: i32 = 0; +pub const READ_MAX_ROWS: i32 = 1; +pub const WRITE_APPEND: i32 = 0; +pub const WRITE_TRUNCATE: i32 = 1; +pub const LOG_ERROR: i32 = 1; + +pub type SparkDsSink = Option<unsafe extern "C" fn(*mut c_void, *const u8, i64)>; +pub type Interrupted = unsafe extern "C" fn(*const SparkDsTaskContext) -> bool; + +// What Spark passes to the library. + +#[repr(C)] +pub struct SparkDsHost { + pub size: i64, + pub abi_version: i32, + pub spark_version: *const c_char, + pub log: Option<LogFunction>, + pub is_log_enabled: Option<LogEnabledFunction>, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsCreateInfo { + pub size: i64, + pub name: *const c_char, + pub option_keys: *const *const c_char, + pub option_values: *const *const c_char, + pub num_options: i64, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsReaderInfo { + pub size: i64, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsPartitionsInfo { + pub size: i64, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsStreamReaderInfo { + pub size: i64, + pub checkpoint_location: *const c_char, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsInitialOffsetInfo { + pub size: i64, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsReadLimit { + pub size: i64, + pub kind: i32, + pub max_rows: i64, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsLatestOffsetInfo { + pub size: i64, + pub start: *const c_char, + pub limit: *const SparkDsReadLimit, + pub run_end: *const c_char, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsStreamPartitionsInfo { + pub size: i64, + pub start: *const c_char, + pub end: *const c_char, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsCommitOffsetInfo { + pub size: i64, + pub end: *const c_char, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsTaskContext { + pub size: i64, + pub stage_id: i32, + pub stage_attempt_number: i32, + pub partition_id: i32, + pub attempt_number: i32, + pub task_attempt_id: i64, + pub is_interrupted: Option<Interrupted>, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsReadInfo { + pub size: i64, + pub kind: i32, + pub schema: *mut ArrowSchema, + pub reader_state: *const u8, + pub reader_state_length: i64, + pub partition: *const u8, + pub partition_length: i64, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct SparkDsWriteInfo { Review Comment: Done in e933ef2cf382: abi.rs now has compile-time `size_of`/`offset_of!` assertions for every field of all 17 structs, on 64-bit targets. I checked the values against spark_datasource.h with a syntax-only C compile. ########## sql/native/rust/src/export.rs: ########## @@ -0,0 +1,519 @@ +// 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. + +//! The function table of an exported source: native callbacks that adapt the C ABI to the +//! bridge of the source. + +use std::ffi::{c_char, c_int, c_void}; + +use crate::abi::{ + self, ArrowArray, ArrowArrayStream, ArrowSchema, SparkDsCommitOffsetInfo, SparkDsCreateInfo, + SparkDsDataWriterInfo, SparkDsFinishDataWriterInfo, SparkDsFinishWriteInfo, + SparkDsInitialOffsetInfo, SparkDsLatestOffsetInfo, SparkDsLibrary, SparkDsPartitionsInfo, + SparkDsReadInfo, SparkDsReadLimit, SparkDsReaderInfo, SparkDsSink, SparkDsStreamPartitionsInfo, + SparkDsStreamReaderInfo, SparkDsTaskContext, SparkDsWriteInfo, +}; +use crate::bridge::{open_reader, Bridge, EncodedPlan}; +use crate::{ + host, write, BatchWrite, DataSource, MicroBatchStream, MicroBatchStreamInfo, ScanBuilder, + StreamingWrite, +}; + +/// Exports a data source through version 1 of Spark's C data source ABI. +/// +/// Call this macro once in a cdylib: +/// +/// ```rust,ignore +/// export_data_source!(RangeSource); +/// ``` +/// +/// The associated types of [`DataSource`](crate::DataSource) determine the operations: +/// the SDK installs the native callbacks of each operation whose type is not +/// [`Unsupported`](crate::Unsupported), and leaves the others null. Batch and streaming +/// operations share the executor callbacks, `read` and the data writer callbacks. The only exported +/// symbol is `spark_ds_init`; sources implement the safe Rust traits and do not handle +/// C pointers. The library must be built with `panic = "unwind"`, so that the SDK can +/// report a panic as an error. +#[macro_export] +macro_rules! export_data_source { + ($source:ty $(,)?) => { + #[cfg(panic = "abort")] + compile_error!( + "spark_native_source requires panic = \"unwind\": with panic = \"abort\", a panic \ + in the data source aborts the JVM of Spark instead of failing the operation" + ); + + const _: () = { + impl $crate::__private::Exported for $source { + fn bridge() -> &'static $crate::__private::Bridge<Self> { + static BRIDGE: std::sync::OnceLock<$crate::__private::Bridge<$source>> = + std::sync::OnceLock::new(); + BRIDGE.get_or_init(Default::default) + } + } + + static LIBRARY: $crate::__private::abi::SparkDsLibrary = + $crate::__private::library::<$source>(); + + #[export_name = "spark_ds_init"] + /// # Safety + /// The arguments must follow the spark_ds_init contract of spark_datasource.h. + unsafe extern "C" fn spark_ds_init( + host: *const $crate::__private::abi::SparkDsHost, + library: *mut *const $crate::__private::abi::SparkDsLibrary, + ) -> std::ffi::c_int { + unsafe { + $crate::__private::abi::initialize( + host, + library, + &LIBRARY, + <$source as $crate::DataSource>::init, + ) + } + } + }; + }; +} + +/// A data source that [`export_data_source!`] exports, with the registry of its handles. +pub trait Exported: DataSource { + fn bridge() -> &'static Bridge<Self>; +} + +// The factories that the executors restore for each kind of operation. +type ReaderFactory<S> = <<S as DataSource>::ScanBuilder as ScanBuilder>::Factory; +type StreamReaderFactory<S> = <<S as DataSource>::MicroBatchStream as MicroBatchStream>::Factory; +type WriterFactory<S> = <<S as DataSource>::BatchWrite as BatchWrite>::Factory; +type StreamingWriterFactory<S> = <<S as DataSource>::StreamingWrite as StreamingWrite>::Factory; + +/// The function table of a source, with the callbacks of the operations that it implements. +pub const fn library<S: Exported>() -> SparkDsLibrary { + let read = Bridge::<S>::READ; + let micro_batch_read = Bridge::<S>::MICRO_BATCH_READ; + // Batch and streaming writes share all the callbacks of writers. + let write = Bridge::<S>::WRITE || Bridge::<S>::STREAMING_WRITE; + SparkDsLibrary { + size: std::mem::size_of::<SparkDsLibrary>() as i64, + abi_version: abi::ABI_VERSION, + last_error: Some(abi::last_error), + create: Some(create::<S>), + schema: if S::INFERS_SCHEMA { + Some(schema::<S>) + } else { + None + }, + create_reader: if read { Some(create_reader::<S>) } else { None }, + create_stream_reader: if micro_batch_read { + Some(create_stream_reader::<S>) + } else { + None + }, + create_writer: if write { + Some(create_writer::<S>) + } else { + None + }, + close: Some(close::<S>), + push_limit: if read { Some(push_limit::<S>) } else { None }, + prune_columns: if read { Some(prune_columns::<S>) } else { None }, + partitions: if read { Some(partitions::<S>) } else { None }, + close_reader: if read { Some(close_reader::<S>) } else { None }, + initial_offset: if micro_batch_read { + Some(initial_offset::<S>) + } else { + None + }, + latest_offset: if micro_batch_read { + Some(latest_offset::<S>) + } else { + None + }, + stream_partitions: if micro_batch_read { + Some(stream_partitions::<S>) + } else { + None + }, + commit_offset: if micro_batch_read { + Some(commit_offset::<S>) + } else { + None + }, + close_stream_reader: if micro_batch_read { + Some(close_stream_reader::<S>) + } else { + None + }, + // Batch scans and micro-batches share the executor reader. + read: if read || micro_batch_read { + Some(read_partition::<S>) + } else { + None + }, + commit: if write { Some(commit::<S>) } else { None }, + abort: if write { Some(abort::<S>) } else { None }, + close_writer: if write { Some(close_writer::<S>) } else { None }, + create_data_writer: if write { + Some(create_data_writer::<S>) + } else { + None + }, + write: if write { Some(write_batch) } else { None }, + commit_data_writer: if write { + Some(commit_data_writer) + } else { + None + }, + abort_data_writer: if write { Some(abort_data_writer) } else { None }, + close_data_writer: if write { Some(close_data_writer) } else { None }, + default_read_limit: if micro_batch_read { + Some(default_read_limit::<S>) + } else { + None + }, + } +} + +/// The ID of the driver handle that Spark passes back, which the SDK returned as a pointer. +fn id(handle: *mut c_void) -> i64 { + handle as usize as i64 +} + +fn handle(id: i64) -> *mut c_void { + id as usize as *mut c_void +} + +/// Passes the partitions of a plan, and its state, to the sinks of Spark. +/// +/// # Safety +/// The sinks must follow the contract of `partitions` and `stream_partitions`. +unsafe fn emit_plan( + plan: EncodedPlan, + add_partition: SparkDsSink, + set_state: SparkDsSink, + sink: *mut c_void, +) -> crate::Result<()> { + for partition in &plan.partitions { + unsafe { abi::emit(add_partition, sink, partition)? }; + } + unsafe { abi::emit(set_state, sink, &plan.state) } +} + +unsafe extern "C" fn create<S: Exported>( + info: *const SparkDsCreateInfo, + capabilities: *mut u64, + source: *mut *mut c_void, +) -> c_int { + abi::call(|| { + abi::require_output(capabilities)?; + abi::require_output(source)?; + // An operation starts on the driver: follow the level of the log of Spark. + host::follow_level(); + let info = unsafe { abi::create_info(info)? }; + let (created, declared) = S::bridge().create(info)?; + unsafe { + capabilities.write(declared); + source.write(handle(created)); + } + Ok(()) + }) +} + +unsafe extern "C" fn schema<S: Exported>(source: *mut c_void, output: *mut ArrowSchema) -> c_int { + abi::call(|| { + let schema = S::bridge().schema(id(source))?; + unsafe { abi::export_schema(&schema, output) } + }) +} + +unsafe extern "C" fn create_reader<S: Exported>( + source: *mut c_void, + schema: *mut ArrowSchema, + _info: *const SparkDsReaderInfo, + reader: *mut *mut c_void, +) -> c_int { + abi::call(|| { + abi::require_output(reader)?; + let schema = unsafe { abi::import_schema(schema)? }; + let created = S::bridge().create_reader(id(source), schema)?; + unsafe { reader.write(handle(created)) }; + Ok(()) + }) +} + +unsafe extern "C" fn close<S: Exported>(source: *mut c_void) { + let _ = abi::call(|| S::bridge().close_source(id(source))); +} + +unsafe extern "C" fn push_limit<S: Exported>( + reader: *mut c_void, + limit: i32, + applied: *mut bool, +) -> c_int { + abi::call(|| { + abi::require_output(applied)?; + let result = S::bridge().push_limit(id(reader), limit.try_into()?)?; + unsafe { applied.write(result) }; + Ok(()) + }) +} + +unsafe extern "C" fn prune_columns<S: Exported>( + reader: *mut c_void, + columns: *const *const c_char, + count: i64, +) -> c_int { + abi::call(|| { + let columns = unsafe { abi::strings(columns, count)? }; + S::bridge().prune_columns(id(reader), &columns) + }) +} + +unsafe extern "C" fn partitions<S: Exported>( + reader: *mut c_void, + _info: *const SparkDsPartitionsInfo, + add_partition: SparkDsSink, + set_state: SparkDsSink, + sink: *mut c_void, +) -> c_int { + abi::call(|| { + let plan = S::bridge().plan(id(reader))?; + unsafe { emit_plan(plan, add_partition, set_state, sink) } + }) +} + +unsafe extern "C" fn close_reader<S: Exported>(reader: *mut c_void) { + let _ = abi::call(|| S::bridge().close_reader(id(reader))); +} + +unsafe extern "C" fn create_stream_reader<S: Exported>( + source: *mut c_void, + schema: *mut ArrowSchema, + info: *const SparkDsStreamReaderInfo, + reader: *mut *mut c_void, +) -> c_int { + abi::call(|| { + abi::require_output(reader)?; + let info = unsafe { abi::input(info)? }; + let checkpoint_location = unsafe { abi::string(info.checkpoint_location)? }; + let schema = unsafe { abi::import_schema(schema)? }; + let info = MicroBatchStreamInfo::new(schema, checkpoint_location); + let created = S::bridge().create_stream_reader(id(source), info)?; + unsafe { reader.write(handle(created)) }; + Ok(()) + }) +} + +unsafe extern "C" fn initial_offset<S: Exported>( + reader: *mut c_void, + _info: *const SparkDsInitialOffsetInfo, + set_offset: SparkDsSink, + sink: *mut c_void, +) -> c_int { + abi::call(|| { + let offset = S::bridge().initial_offset(id(reader))?; + unsafe { abi::emit(set_offset, sink, offset.as_bytes()) } + }) +} + +unsafe extern "C" fn default_read_limit<S: Exported>( + reader: *mut c_void, + limit: *mut SparkDsReadLimit, +) -> c_int { + abi::call(|| { + let value = S::bridge().default_read_limit(id(reader))?; + unsafe { abi::set_read_limit(limit, value) } + }) +} + +unsafe extern "C" fn latest_offset<S: Exported>( + reader: *mut c_void, + info: *const SparkDsLatestOffsetInfo, + set_offset: SparkDsSink, + sink: *mut c_void, +) -> c_int { + abi::call(|| { + let info = unsafe { abi::latest_offset_info(info)? }; + let offset = S::bridge().latest_offset(id(reader), info)?; + unsafe { abi::emit(set_offset, sink, offset.as_bytes()) } + }) +} + +unsafe extern "C" fn stream_partitions<S: Exported>( + reader: *mut c_void, + info: *const SparkDsStreamPartitionsInfo, + add_partition: SparkDsSink, + set_state: SparkDsSink, + sink: *mut c_void, +) -> c_int { + abi::call(|| { + let info = unsafe { abi::input(info)? }; + let start = unsafe { abi::string(info.start)? }; + let end = unsafe { abi::string(info.end)? }; + let plan = S::bridge().stream_partitions(id(reader), &start, &end)?; + unsafe { emit_plan(plan, add_partition, set_state, sink) } + }) +} + +unsafe extern "C" fn commit_offset<S: Exported>( + reader: *mut c_void, + info: *const SparkDsCommitOffsetInfo, +) -> c_int { + abi::call(|| { + let info = unsafe { abi::input(info)? }; + let end = unsafe { abi::string(info.end)? }; + S::bridge().commit_offset(id(reader), &end) + }) +} + +unsafe extern "C" fn close_stream_reader<S: Exported>(reader: *mut c_void) { + let _ = abi::call(|| S::bridge().close_stream_reader(id(reader))); +} + +unsafe extern "C" fn read_partition<S: Exported>( + task: *const SparkDsTaskContext, + info: *const SparkDsReadInfo, + output: *mut ArrowArrayStream, +) -> c_int { + abi::call(|| { + abi::require_output(output)?; + // A task starts on an executor: follow the level of the log of Spark. + host::follow_level(); + let info = unsafe { abi::input(info)? }; + let schema = unsafe { abi::import_schema(info.schema)? }; + let state = unsafe { abi::slice(info.reader_state, info.reader_state_length)? }; + let partition = unsafe { abi::slice(info.partition, info.partition_length)? }; + let stream = match info.kind { + abi::READ_BATCH => unsafe { + open_reader::<ReaderFactory<S>>(task, schema, state, partition)? + }, + abi::READ_MICRO_BATCH => unsafe { + open_reader::<StreamReaderFactory<S>>(task, schema, state, partition)? + }, + kind => return Err(format!("unsupported kind of read: {kind}").into()), + }; + unsafe { output.write(stream) }; + Ok(()) + }) +} + +unsafe extern "C" fn create_writer<S: Exported>( + source: *mut c_void, + schema: *mut ArrowSchema, + info: *const SparkDsWriteInfo, + set_state: SparkDsSink, + sink: *mut c_void, + writer: *mut *mut c_void, +) -> c_int { + abi::call(|| { + abi::require_output(writer)?; + set_state.ok_or("null result sink")?; + let (mode, query_id) = unsafe { abi::write_info(info)? }; + let schema = unsafe { abi::import_schema(schema)? }; + let (created, state) = S::bridge().create_writer(id(source), schema, mode, query_id)?; + if let Err(error) = unsafe { abi::emit(set_state, sink, &state) } { + S::bridge().close_writer(created)?; + return Err(error); + } + unsafe { writer.write(handle(created)) }; + Ok(()) + }) +} + +unsafe extern "C" fn commit<S: Exported>( + writer: *mut c_void, + info: *const SparkDsFinishWriteInfo, +) -> c_int { + abi::call(|| { + let info = unsafe { abi::input(info)? }; + let messages = + unsafe { abi::messages(info.messages, info.message_lengths, info.num_messages)? }; + S::bridge().commit(id(writer), info.epoch_id, &messages) + }) +} + +unsafe extern "C" fn abort<S: Exported>( + writer: *mut c_void, + info: *const SparkDsFinishWriteInfo, +) -> c_int { + abi::call(|| { + let info = unsafe { abi::input(info)? }; + // Spark rejected what create_writer returned, before any task ran. The SDK returns a + // single state, so this does not happen; the coordinator is dropped at close, as when + // its creation fails. Review Comment: Done in e933ef2cf382: The abort comment, the `BatchWrite` lifecycle rustdoc and the README now cover Spark rejecting the state when it cannot copy it, for example when it is too large. ########## 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 Review Comment: Done in e933ef2cf382: The README now names the structures that are rejected when shorter than in version 1: the task context, the info structures and the read limit that Spark passes to the calls. -- 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]
