From f975870d0b925b193e24504507fcc4bec2699afe Mon Sep 17 00:00:00 2001 From: Gengliang Wang Date: Thu, 15 Oct 2026 18:05:46 +0000 Subject: [PATCH 1/6] Add experimental Rust SDK for the native data source C ABI --- .../main/rust/native-data-source/.gitignore | 2 + .../main/rust/native-data-source/Cargo.toml | 34 + .../main/rust/native-data-source/README.md | 124 ++++ .../main/rust/native-data-source/package.py | 70 ++ .../src/main/rust/native-data-source/smoke.py | 77 ++ .../main/rust/native-data-source/src/lib.rs | 608 ++++++++++++++++ .../native-data-source/tests/c_abi_smoke.c | 459 ++++++++++++ .../native-data-source/tests/c_abi_smoke.rs | 87 +++ sql/native/rust/Cargo.lock | 664 +++++++++++++++++ sql/native/rust/Cargo.toml | 51 ++ sql/native/rust/README.md | 200 +++++ sql/native/rust/src/abi.rs | 362 ++++++++++ sql/native/rust/src/bridge.rs | 677 +++++++++++++++++ sql/native/rust/src/export.rs | 217 ++++++ sql/native/rust/src/lib.rs | 149 ++++ sql/native/rust/src/predicate.rs | 83 +++ sql/native/rust/src/stream.rs | 517 +++++++++++++ sql/native/rust/src/task.rs | 187 +++++ .../rust/tests/fixtures/spark_datasource.h | 683 ++++++++++++++++++ 19 files changed, 5251 insertions(+) create mode 100644 examples/src/main/rust/native-data-source/.gitignore create mode 100644 examples/src/main/rust/native-data-source/Cargo.toml create mode 100644 examples/src/main/rust/native-data-source/README.md create mode 100644 examples/src/main/rust/native-data-source/package.py create mode 100644 examples/src/main/rust/native-data-source/smoke.py create mode 100644 examples/src/main/rust/native-data-source/src/lib.rs create mode 100644 examples/src/main/rust/native-data-source/tests/c_abi_smoke.c create mode 100644 examples/src/main/rust/native-data-source/tests/c_abi_smoke.rs create mode 100644 sql/native/rust/Cargo.lock create mode 100644 sql/native/rust/Cargo.toml create mode 100644 sql/native/rust/README.md create mode 100644 sql/native/rust/src/abi.rs create mode 100644 sql/native/rust/src/bridge.rs create mode 100644 sql/native/rust/src/export.rs create mode 100644 sql/native/rust/src/lib.rs create mode 100644 sql/native/rust/src/predicate.rs create mode 100644 sql/native/rust/src/stream.rs create mode 100644 sql/native/rust/src/task.rs create mode 100644 sql/native/rust/tests/fixtures/spark_datasource.h diff --git a/examples/src/main/rust/native-data-source/.gitignore b/examples/src/main/rust/native-data-source/.gitignore new file mode 100644 index 0000000000000..aee95f2858593 --- /dev/null +++ b/examples/src/main/rust/native-data-source/.gitignore @@ -0,0 +1,2 @@ +*.sparkpkg +__pycache__/ diff --git a/examples/src/main/rust/native-data-source/Cargo.toml b/examples/src/main/rust/native-data-source/Cargo.toml new file mode 100644 index 0000000000000..98f3201457d95 --- /dev/null +++ b/examples/src/main/rust/native-data-source/Cargo.toml @@ -0,0 +1,34 @@ +# 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. + +[package] +name = "spark-native-range-demo" +workspace = "../../../../../sql/native/rust" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true + +[lib] +name = "spark_datasource_rust_sdk_range" +path = "src/lib.rs" +crate-type = ["cdylib", "rlib"] + +[dependencies] +spark-native-source = { path = "../../../../../sql/native/rust" } +arrow-array.workspace = true +arrow-schema.workspace = true +serde.workspace = true +serde_json.workspace = true diff --git a/examples/src/main/rust/native-data-source/README.md b/examples/src/main/rust/native-data-source/README.md new file mode 100644 index 0000000000000..9405136309be7 --- /dev/null +++ b/examples/src/main/rust/native-data-source/README.md @@ -0,0 +1,124 @@ + + +# Rust native range data source + +This example implements batch reads using the [Rust native data source SDK][sdk]. +The connector in [src/lib.rs](src/lib.rs) exports `rust_sdk_range` through the SDK's +`export_data_source!` macro. The macro exports the C entry point `spark_ds_init`; the connector +does not use JNI, raw pointers, or unsafe code. + +The example targets C ABI version 1 from [Spark PR #59424][native-source-pr], pinned to commit +`51af65c63a9369d793aab1a198ddde4c90c8f85b`. The SDK and example build without Spark or a JVM. +Running the Spark smoke script requires a distribution built with that native data source +framework and Java 22 or later on the driver and executors. The JNI SDK prototype in PR #59442 +targets a different interface. + +The source advertises only batch reads. It does not implement writes or Structured Streaming; +the framework rejects those operations from its declared capabilities. + +## Data and options + +| Column | Spark type | Value | +| --- | --- | --- | +| `id` | `BIGINT`, non-null | Integers in `[start, end)` | +| `label` | `STRING`, non-null | `row-{id}`, such as `row-13` | + +| Option | Default | Meaning | +| --- | --- | --- | +| `start` | `0` | Inclusive start, as a signed 64-bit integer | +| `end` | `100` | Exclusive end; must be at least `start` | +| `numpartitions` | `2` | Maximum number of partitions; must be positive | +| `batchsize` | `1024` | Maximum rows per Arrow batch; must be positive | + +Spark normalizes option names to lower case. Unknown options are rejected. + +The source evaluates `=`, `<=>`, `<`, `<=`, `>`, and `>=` comparisons between `id` and non-null +integer literals, including conjunctions of supported predicates. It reports these predicates +as `Exact`; Spark still evaluates all predicates again after reading. Unknown expressions and +unsupported literal types remain `Unsupported`. + +Spark pushes a limit only when no residual filter remains, which excludes filtered scans in +the pinned framework. The connector can apply a global limit by shortening the selected range +before splitting it into partitions. Column pruning preserves the requested order and supports +zero-column batches for operations such as `COUNT(*)`. + +## Planning and reading + +`RangeSource` creates a driver-side `RangeScan`. Pushdown methods update that builder, and +`plan(self)` consumes it to produce a `ScanPlan`. The plan contains: + +- a serializable `RangeReaderFactory` with output columns and batch size; +- disjoint `RangePartition` values with inclusive starts and exclusive ends. + +The SDK serializes this data for the executors. Each factory opens a lazy Arrow reader using a +partition and a lifetime-bound `TaskContext`. The example checks task cancellation before +generating each batch and keeps memory bounded by the configured batch size. Neither the factory +nor a partition contains a driver handle or open connection. A storage connector can also put a +snapshot identifier in its factory to make retried reads refer to the same data. + +## Build and test + +Run from the repository root with Rust and a C compiler: + +```sh +cargo build --manifest-path sql/native/rust/Cargo.toml --workspace --locked +cargo test --manifest-path sql/native/rust/Cargo.toml --workspace --locked +``` + +Build before testing so the `c_abi_smoke` integration test can load the actual connector library. +That test calls its function table through the C header and consumes Arrow C streams without a +JVM. The Rust tests cover filters, global limits, projections, serialization, bounded batches, +and empty output projections. These checks do not substitute for the Spark smoke script. + +## Package and run with Spark + +Set `SPARK_HOME` to the distribution with #59424, then run: + +```sh +python3 examples/src/main/rust/native-data-source/package.py \ + --library sql/native/rust/target/debug/libspark_datasource_rust_sdk_range.so \ + --output sql/native/rust/target/rust_sdk_range.sparkpkg + +"$SPARK_HOME/bin/spark-submit" --master 'local[2]' \ + examples/src/main/rust/native-data-source/smoke.py \ + --package sql/native/rust/target/rust_sdk_range.sparkpkg +``` + +The library name above is for Linux. On macOS use `libspark_datasource_rust_sdk_range.dylib`; +on Windows use `spark_datasource_rust_sdk_range.dll`. The packaging script detects the host +platform. For a cross-compiled library, supply `--platform`, for example `linux-aarch64`. + +The package can also be used from a Python session backed by the same Spark build: + +```python +spark.addArtifact("sql/native/rust/target/rust_sdk_range.sparkpkg") +df = ( + spark.read.format("rust_sdk_range") + .option("end", 100) + .option("numpartitions", 4) + .load() +) +df.where("id >= 10 AND id < 20").select("label", "id").show() +``` + +[smoke.py](smoke.py) checks rows, projection, unfiltered and filtered limits, zero-column counts, +and residual predicates against Spark. Run it against the pinned framework before relying on +end-to-end compatibility. + +[sdk]: ../../../../../sql/native/rust/README.md +[native-source-pr]: https://github.com/apache/spark/pull/59424 diff --git a/examples/src/main/rust/native-data-source/package.py b/examples/src/main/rust/native-data-source/package.py new file mode 100644 index 0000000000000..09019d3943458 --- /dev/null +++ b/examples/src/main/rust/native-data-source/package.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +# 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. + +"""Package the compiled range connector for the C ABI in Spark PR #59424.""" + +import argparse +import json +from pathlib import Path +import platform +import zipfile + + +def current_platform(): + os_name = {"Linux": "linux", "Darwin": "osx", "Windows": "windows"}.get(platform.system()) + arch = { + "x86_64": "x86_64", + "amd64": "x86_64", + "aarch64": "aarch64", + "arm64": "aarch64", + }.get(platform.machine().lower()) + return f"{os_name}-{arch}" if os_name and arch else None + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--library", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--platform", default=current_platform(), help="Override for cross builds") + args = parser.parse_args() + if not args.library.is_file(): + parser.error(f"Library does not exist: {args.library}") + if args.output.suffix != ".sparkpkg": + parser.error("The output file must have the .sparkpkg extension") + if args.platform not in { + f"{os_name}-{arch}" + for os_name in ("linux", "osx", "windows") + for arch in ("x86_64", "aarch64") + }: + parser.error("Specify a supported --platform, such as linux-x86_64") + if args.library.resolve() == args.output.resolve(): + parser.error("The output must differ from the input library") + + library_path = f"{args.platform}/{args.library.name}" + manifest = { + "abiVersion": 1, + "dataSources": ["rust_sdk_range"], + "libraries": {args.platform: library_path}, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(args.output, "w", compression=zipfile.ZIP_DEFLATED) as package: + package.writestr("spark-native-datasource.json", json.dumps(manifest, indent=2) + "\n") + package.write(args.library, library_path) + print(args.output.resolve()) + + +if __name__ == "__main__": + main() diff --git a/examples/src/main/rust/native-data-source/smoke.py b/examples/src/main/rust/native-data-source/smoke.py new file mode 100644 index 0000000000000..6f78bd1ec2548 --- /dev/null +++ b/examples/src/main/rust/native-data-source/smoke.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +# 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. + +"""Run with Java 22+ and spark-submit from a distribution that includes PR #59424.""" + +import argparse +from pathlib import Path + +from pyspark.sql import SparkSession, functions as F + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--package", type=Path, required=True) + args = parser.parse_args() + package = args.package.resolve(strict=True) + + spark = SparkSession.builder.appName("Rust native data source smoke").getOrCreate() + try: + spark.addArtifact(str(package)) + df = ( + spark.read.format("rust_sdk_range") + .option("start", 0) + .option("end", 31) + .option("numpartitions", 4) + .option("batchsize", 3) + .load() + ) + assert [tuple(row) for row in df.orderBy("id").collect()] == [ + (i, f"row-{i}") for i in range(31) + ] + + # The native framework pushes limits only for scans without residual filters. + unfiltered_limit = df.limit(4).collect() + assert len(unfiltered_limit) == 4 + assert all(0 <= row.id < 31 for row in unfiltered_limit) + + projected = df.where((F.col("id") >= 7) & (F.col("id") < 17)).select("label", "id") + assert [tuple(row) for row in projected.orderBy("id").collect()] == [ + (f"row-{i}", i) for i in range(7, 17) + ] + limited = df.where((F.col("id") >= 10) & (F.col("id") < 20)).limit(4).collect() + assert len(limited) == 4 + assert all(10 <= row.id < 20 for row in limited) + + # Count needs row counts even when column pruning produces zero-column Arrow batches. + assert df.select().count() == 31 + assert df.where(F.col("id") < 0).count() == 0 + + # Spark rechecks every predicate, including those the connector evaluates exactly. + # These expressions also exercise filters the connector does not evaluate. + arithmetic = df.where((F.col("id") % 3) == 1).select("id").orderBy("id").collect() + assert [row.id for row in arithmetic] == list(range(1, 31, 3)) + labels = df.where(F.col("label").endswith("3")).select("id").orderBy("id").collect() + assert [row.id for row in labels] == [3, 13, 23] + residual_limit = df.where(F.col("label") == "row-21").limit(1).select("id").collect() + assert [row.id for row in residual_limit] == [21] + print("Rust native data source Spark smoke checks passed") + finally: + spark.stop() + + +if __name__ == "__main__": + main() diff --git a/examples/src/main/rust/native-data-source/src/lib.rs b/examples/src/main/rust/native-data-source/src/lib.rs new file mode 100644 index 0000000000000..360459dffefa9 --- /dev/null +++ b/examples/src/main/rust/native-data-source/src/lib.rs @@ -0,0 +1,608 @@ +// 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, RecordBatchIterator, RecordBatchOptions}; +use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; +use serde::{Deserialize, Serialize}; +use spark_native_source::{ + BatchReader, BatchReaderFactory, BatchSource, FilterSupport, LimitSupport, Options, Predicate, + Result, ScanBuilder, ScanPlan, TaskContext, +}; + +pub struct RangeSource { + start: i64, + end: i64, + num_partitions: usize, + batch_size: usize, +} + +pub struct RangeScan { + lower: i128, + upper: i128, + num_partitions: usize, + limit: Option, + factory: RangeReaderFactory, +} + +/// Reader configuration travels to each executor; it contains no live objects. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RangeReaderFactory { + columns: Vec, + batch_size: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RangePartition { + start: i64, + end: i64, +} + +impl BatchSource for RangeSource { + type Scan = RangeScan; + + fn new(options: Options) -> Result { + for key in options.keys() { + if !matches!( + key.as_str(), + "start" | "end" | "numpartitions" | "batchsize" + ) { + 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)?; + if end < start { + return Err("end must be greater than or equal to start".into()); + } + if num_partitions == 0 || batch_size == 0 { + return Err("numPartitions and batchSize must be positive".into()); + } + Ok(Self { + start, + end, + num_partitions, + batch_size, + }) + } + + fn schema(&self) -> Result { + output_schema(&["id".into(), "label".into()]) + } + + fn scan(&self, schema: SchemaRef) -> Result { + let columns: Vec = schema.fields().iter().map(|f| f.name().clone()).collect(); + let expected = output_schema(&columns)?; + for (field, expected) in schema.fields().iter().zip(expected.fields()) { + if field.data_type() != expected.data_type() { + return Err(format!( + "Unsupported type for {}: {}", + field.name(), + field.data_type() + ) + .into()); + } + } + Ok(RangeScan { + lower: i128::from(self.start), + upper: i128::from(self.end), + num_partitions: self.num_partitions, + limit: None, + factory: RangeReaderFactory { + columns, + batch_size: self.batch_size, + }, + }) + } +} + +impl BatchReaderFactory for RangeReaderFactory { + type Partition = RangePartition; + + fn open<'task>( + &self, + partition: RangePartition, + task: TaskContext<'task>, + ) -> Result> { + let factory = self.clone(); + if factory.batch_size == 0 || partition.start > partition.end { + return Err("Invalid serialized range plan".into()); + } + let schema = output_schema(&factory.columns)?; + let batch_schema = Arc::clone(&schema); + 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) + factory.batch_size as i128) + .min(i128::from(partition.end)) as i64; + let num_rows = (i128::from(next) - i128::from(current)) as usize; + let columns: Vec = factory + .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(&batch_schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(num_rows)), + )) + }); + Ok(Box::new(RecordBatchIterator::new(batches, schema))) + } +} + +impl ScanBuilder for RangeScan { + type Factory = RangeReaderFactory; + + fn push_filters(&mut self, filters: &[Predicate]) -> Result> { + Ok(filters + .iter() + .map(|filter| { + if let Some((lower, upper)) = predicate_bounds(filter) { + self.lower = self.lower.max(lower); + self.upper = self.upper.min(upper); + FilterSupport::Exact + } else { + FilterSupport::Unsupported + } + }) + .collect()) + } + + fn push_limit(&mut self, limit: usize) -> Result { + self.limit = Some(self.limit.map_or(limit, |previous| previous.min(limit))); + Ok(LimitSupport::Global) + } + + fn prune_columns(&mut self, columns: &[String]) -> Result { + if columns + .iter() + .any(|column| !self.factory.columns.contains(column)) + { + return Err("Projection contains a column outside the read schema".into()); + } + output_schema(columns)?; + self.factory.columns = columns.to_vec(); + Ok(true) + } + + fn plan(self) -> Result> { + // Filter before applying a global limit, then split the selected interval. + let upper = self.limit.map_or(self.upper, |limit| { + self.upper.min(self.lower + limit as i128) + }); + let rows = (upper - self.lower).max(0); + let count = (self.num_partitions as i128).min(rows) as usize; + let mut partitions = Vec::with_capacity(count); + if count > 0 { + let width = rows / count as i128; + let extra = rows % count as i128; + let mut start = self.lower; + for index in 0..count { + let end = start + width + i128::from((index as i128) < extra); + partitions.push(RangePartition { + start: start as i64, + end: end as i64, + }); + start = end; + } + } + Ok(ScanPlan { + factory: self.factory, + partitions, + }) + } +} + +fn option(options: &Options, name: &str, default: T) -> Result +where + T: std::str::FromStr, + T::Err: std::fmt::Display, +{ + options.get(name).map_or(Ok(default), |value| { + value + .parse() + .map_err(|error| format!("Invalid {name}: {error}").into()) + }) +} + +fn output_schema(columns: &[String]) -> Result { + let mut seen = BTreeSet::new(); + let fields: Result> = columns + .iter() + .map(|name| { + if !seen.insert(name) { + return Err(format!("Duplicate column: {name}").into()); + } + let data_type = match name.as_str() { + "id" => DataType::Int64, + "label" => DataType::Utf8, + _ => return Err(format!("Unknown column: {name}").into()), + }; + Ok(Field::new(name, data_type, false)) + }) + .collect(); + Ok(Arc::new(Schema::new(fields?))) +} + +/// Closed-open bounds use i128 so strict comparisons never overflow at i64 extrema. +fn predicate_bounds(predicate: &Predicate) -> Option<(i128, i128)> { + let Predicate::Function { name, children } = predicate else { + return None; + }; + if children.len() != 2 { + return None; + } + if name == "AND" { + // Do not mutate the scan when only one side of an AND is supported. + let left = predicate_bounds(&children[0])?; + let right = predicate_bounds(&children[1])?; + return Some((left.0.max(right.0), left.1.min(right.1))); + } + let (value, reversed) = if is_id(&children[0]) { + (integer_literal(&children[1])?, false) + } else if is_id(&children[1]) { + (integer_literal(&children[0])?, true) + } else { + return None; + }; + let name = match (name.as_str(), reversed) { + (">", true) => "<", + (">=", true) => "<=", + ("<", true) => ">", + ("<=", true) => ">=", + (name, _) => name, + }; + let min = i128::from(i64::MIN); + let max = i128::from(i64::MAX) + 1; + let value = i128::from(value); + match name { + "=" | "<=>" => Some((value, value + 1)), + ">" => Some((value + 1, max)), + ">=" => Some((value, max)), + "<" => Some((min, value)), + "<=" => Some((min, value + 1)), + _ => None, + } +} + +fn is_id(predicate: &Predicate) -> bool { + matches!(predicate, Predicate::Column { name } if name.len() == 1 && name[0] == "id") +} + +fn integer_literal(predicate: &Predicate) -> Option { + let Predicate::Literal { data_type, value } = predicate else { + return None; + }; + let value = value.as_i64()?; + match data_type.as_str() { + "tinyint" if i8::try_from(value).is_ok() => Some(value), + "smallint" if i16::try_from(value).is_ok() => Some(value), + "int" if i32::try_from(value).is_ok() => Some(value), + "bigint" => Some(value), + _ => None, + } +} + +spark_native_source::export_data_source!(RangeSource); + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn source(values: &[(&str, &str)]) -> RangeSource { + RangeSource::new( + values + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(), + ) + .unwrap() + } + + fn scan(source: &RangeSource) -> RangeScan { + source.scan(source.schema().unwrap()).unwrap() + } + + fn comparison(name: &str, value: i64, reversed: bool) -> Predicate { + let mut children = vec![ + Predicate::Column { + name: vec!["id".into()], + }, + Predicate::Literal { + data_type: "bigint".into(), + value: json!(value), + }, + ]; + if reversed { + children.reverse(); + } + Predicate::Function { + name: name.into(), + children, + } + } + + fn read(plan: ScanPlan) -> Vec { + plan.partitions + .into_iter() + .flat_map(|partition| { + plan.factory + .open(partition, TaskContext::detached()) + .unwrap() + .map(|batch| batch.unwrap()) + .collect::>() + }) + .collect() + } + + fn ids(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|batch| { + let index = batch.schema().index_of("id").unwrap(); + batch + .column(index) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect() + } + + #[test] + fn filters_before_global_limit_and_preserves_projected_order() { + let source = source(&[ + ("start", "10"), + ("end", "25"), + ("numpartitions", "3"), + ("batchsize", "2"), + ]); + let mut scan = scan(&source); + assert_eq!( + scan.push_filters(&[comparison(">=", 13, false), comparison("<", 22, false),]) + .unwrap(), + vec![FilterSupport::Exact; 2] + ); + assert_eq!(scan.push_limit(5).unwrap(), LimitSupport::Global); + assert!(scan.prune_columns(&["label".into(), "id".into()]).unwrap()); + let plan = scan.plan().unwrap(); + assert_eq!(plan.partitions.len(), 3); + let batches = read(plan); + assert_eq!(ids(&batches), (13..18).collect::>()); + for batch in &batches { + assert!(batch.num_rows() <= 2); + assert_eq!(batch.schema().field(0).name(), "label"); + let labels = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for (row, id) in ids(std::slice::from_ref(batch)).iter().enumerate() { + assert_eq!(labels.value(row), format!("row-{id}")); + } + } + } + + #[test] + fn comparisons_handle_reversed_operands_and_int64_boundaries() { + for (start, end) in [(-2, 3), (i64::MIN, i64::MIN + 3), (i64::MAX - 3, i64::MAX)] { + let source = source(&[("start", &start.to_string()), ("end", &end.to_string())]); + for literal in [i64::MIN, -1, 0, 1, i64::MAX] { + for op in ["=", "<=>", ">", ">=", "<", "<="] { + for reversed in [false, true] { + let mut scan = scan(&source); + assert_eq!( + scan.push_filters(&[comparison(op, literal, reversed)]) + .unwrap(), + vec![FilterSupport::Exact] + ); + let expected: Vec<_> = (start..end) + .filter(|id| { + let (left, right) = if reversed { + (literal, *id) + } else { + (*id, literal) + }; + match op { + "=" | "<=>" => left == right, + ">" => left > right, + ">=" => left >= right, + "<" => left < right, + "<=" => left <= right, + _ => unreachable!(), + } + }) + .collect(); + assert_eq!( + ids(&read(scan.plan().unwrap())), + expected, + "range={start}..{end}, op={op}, value={literal}, reverse={reversed}" + ); + } + } + } + } + } + + #[test] + fn unsupported_filters_remain_residual_without_mutating_the_scan() { + let source = source(&[("end", "10")]); + let mut scan = scan(&source); + let unsupported_and = Predicate::Function { + name: "AND".into(), + children: vec![comparison(">", 5, false), Predicate::Unknown], + }; + let null: Predicate = serde_json::from_value(json!({ + "type": "function", "name": "<=>", "children": [ + {"type": "column", "name": ["id"]}, + {"type": "literal", "dataType": "bigint", "value": null} + ] + })) + .unwrap(); + let double: Predicate = serde_json::from_value(json!({ + "type": "function", "name": "=", "children": [ + {"type": "column", "name": ["id"]}, + {"type": "literal", "dataType": "double", "value": 1.5} + ] + })) + .unwrap(); + assert_eq!( + scan.push_filters(&[ + unsupported_and, + null, + double, + comparison("FUTURE_COMPARISON", 0, false), + ]) + .unwrap(), + vec![FilterSupport::Unsupported; 4] + ); + assert_eq!( + ids(&read(scan.plan().unwrap())), + (0..10).collect::>() + ); + } + + #[test] + fn conjunctions_and_contradictory_predicates_form_exact_intersections() { + let source = source(&[("end", "10")]); + let mut first = scan(&source); + let and = Predicate::Function { + name: "AND".into(), + children: vec![comparison(">=", 3, false), comparison("<=", 6, false)], + }; + assert_eq!( + first.push_filters(&[and]).unwrap(), + vec![FilterSupport::Exact] + ); + assert_eq!(ids(&read(first.plan().unwrap())), vec![3, 4, 5, 6]); + let mut empty = scan(&source); + empty + .push_filters(&[comparison(">", 5, false), comparison("<", 3, false)]) + .unwrap(); + assert!(empty.plan().unwrap().partitions.is_empty()); + } + + #[test] + fn serialized_plans_preserve_zero_column_row_counts_and_batch_sizes() { + let source = source(&[("end", "7"), ("numpartitions", "1"), ("batchsize", "3")]); + let mut scan = scan(&source); + scan.prune_columns(&[]).unwrap(); + let plan = scan.plan().unwrap(); + let factory = serde_json::from_slice(&serde_json::to_vec(&plan.factory).unwrap()).unwrap(); + let partitions = plan + .partitions + .iter() + .map(|partition| { + serde_json::from_slice(&serde_json::to_vec(partition).unwrap()).unwrap() + }) + .collect(); + let batches = read(ScanPlan { + factory, + partitions, + }); + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![3, 3, 1] + ); + assert!(batches.iter().all(|batch| batch.num_columns() == 0)); + } + + #[test] + fn zero_limit_empty_range_and_repeated_limits_do_not_add_rows() { + let provider = source(&[("end", "10"), ("numpartitions", "20")]); + let mut zero = scan(&provider); + zero.push_limit(0).unwrap(); + assert!(zero.plan().unwrap().partitions.is_empty()); + let mut limited = scan(&provider); + limited.push_limit(2).unwrap(); + limited.push_limit(5).unwrap(); + let plan = limited.plan().unwrap(); + assert_eq!(plan.partitions.len(), 2); + assert_eq!(ids(&read(plan)), vec![0, 1]); + let empty = source(&[("start", "5"), ("end", "5")]); + assert!(scan(&empty).plan().unwrap().partitions.is_empty()); + } + + #[test] + fn rejects_invalid_options_and_incompatible_user_schemas() { + for (key, value) in [ + ("batchsize", "0"), + ("batchsize", "-1"), + ("numpartitions", "0"), + ("end", "-1"), + ("start", "not-an-integer"), + ("unknown", "1"), + ] { + assert!(RangeSource::new([(key.into(), value.into())].into()).is_err()); + } + let source = source(&[]); + let invalid = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + assert!(source.scan(invalid).is_err()); + let labels = output_schema(&["label".into()]).unwrap(); + let mut scan = source.scan(labels).unwrap(); + assert!(scan.prune_columns(&["id".into()]).is_err()); + assert!(output_schema(&["id".into(), "id".into()]).is_err()); + } + + #[test] + fn opens_a_huge_partition_lazily_without_materializing_it() { + let source = source(&[ + ("start", "-9223372036854775808"), + ("end", "9223372036854775807"), + ("numpartitions", "1"), + ("batchsize", "2"), + ]); + let plan = scan(&source).plan().unwrap(); + let mut reader = plan + .factory + .open(plan.partitions[0].clone(), TaskContext::detached()) + .unwrap(); + let first = reader.next().unwrap().unwrap(); + let second = reader.next().unwrap().unwrap(); + assert_eq!(ids(&[first]), vec![i64::MIN, i64::MIN + 1]); + assert_eq!(ids(&[second]), vec![i64::MIN + 2, i64::MIN + 3]); + } +} diff --git a/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c b/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c new file mode 100644 index 0000000000000..bdbdae8f039e8 --- /dev/null +++ b/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c @@ -0,0 +1,459 @@ +/* + * 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. + */ + +/* An independent C client: layouts and function signatures come from Spark's actual header. */ +#include "spark_datasource.h" + +#include +#include +#include +#include +#include +#include + +#define CHECK(condition, message) do { \ + if (!(condition)) { \ + fprintf(stderr, "%s:%d: %s\n", __FILE__, __LINE__, (message)); \ + exit(1); \ + } \ +} while (0) + +typedef struct Bytes { + uint8_t* data; + int64_t length; +} Bytes; + +typedef struct ByteList { + Bytes* values; + size_t length; +} ByteList; + +typedef struct Plan { + ByteList partitions; + Bytes state; +} Plan; + +static const SparkDsLibrary* library; +static const SparkDsTaskContext* expected_task; +static bool interrupted; +static int interruption_checks; + +static void check_status(int status) { + if (status != 0) { + const char* error = library->last_error(); + fprintf(stderr, "Native status %d: %s\n", status, error == NULL ? "no message" : error); + exit(1); + } +} + +static void log_message(int32_t level, const char* message) { + fprintf(stderr, "Native log %d: %s\n", level, message); +} + +static bool is_interrupted(const SparkDsTaskContext* task) { + CHECK(task == expected_task, "Cancellation callback must receive the original task pointer"); + CHECK(task->stage_id == 71 && task->stage_attempt_number == 2, "Unexpected stage identity"); + CHECK(task->partition_id == 3 && task->attempt_number == 4, "Unexpected task identity"); + CHECK(task->task_attempt_id == INT64_C(12345678901), "Task attempt ID was truncated"); + interruption_checks++; + return interrupted; +} + +static bool unavailable_callback(const SparkDsTaskContext* task) { + (void) task; + CHECK(false, "A task callback outside its declared struct size must not be read"); + return false; +} + +static SparkDsTaskContext task_context(bool short_context) { + SparkDsTaskContext task = { + .size = sizeof(SparkDsTaskContext), + .stage_id = 71, + .stage_attempt_number = 2, + .partition_id = 3, + .attempt_number = 4, + .task_attempt_id = INT64_C(12345678901), + .is_interrupted = is_interrupted + }; + if (short_context) { + task.size = offsetof(SparkDsTaskContext, is_interrupted); + task.is_interrupted = unavailable_callback; + } + return task; +} + +/* Sink callbacks copy opaque bytes during the call, as required by the ABI. */ +static void add_bytes(void* sink, const uint8_t* data, int64_t length) { + ByteList* list = sink; + CHECK(length >= 0, "A sink received a negative length"); + CHECK(length == 0 || data != NULL, "A nonempty sink value must have data"); + Bytes* values = realloc(list->values, (list->length + 1) * sizeof(Bytes)); + CHECK(values != NULL, "Cannot allocate sink values"); + list->values = values; + Bytes* value = &list->values[list->length++]; + value->length = length; + value->data = malloc(length == 0 ? 1 : (size_t) length); + CHECK(value->data != NULL, "Cannot allocate sink data"); + if (length != 0) { + memcpy(value->data, data, (size_t) length); + } +} + +static void free_list(ByteList* list) { + for (size_t i = 0; i < list->length; i++) { + free(list->values[i].data); + } + free(list->values); +} + +static void check_bytes(const Bytes* left, const Bytes* right) { + CHECK(left->length == right->length, "Repeated planning changed the serialized size"); + CHECK(memcmp(left->data, right->data, (size_t) left->length) == 0, + "Repeated planning changed the serialized data"); +} + +static void* create_source(const char* const* keys, const char* const* values, int64_t count) { + SparkDsCreateInfo info = { + .size = sizeof(SparkDsCreateInfo), + .name = "rust_sdk_range", + .option_keys = keys, + .option_values = values, + .num_options = count + }; + uint64_t capabilities = 0; + void* source = NULL; + check_status(library->create(&info, &capabilities, &source)); + CHECK(source != NULL, "create returned a null source"); + CHECK(capabilities == SPARK_DS_BATCH_READ, "The Range source must only declare batch reads"); + return source; +} + +static Plan plan_scan( + const char* const* keys, + const char* const* values, + int64_t option_count, + const char* const* predicates, + const int32_t* expected_uses, + int64_t predicate_count, + int32_t limit, + const char* const* projection, + int64_t projection_count) { + void* source = create_source(keys, values, option_count); + struct ArrowSchema schema = {0}; + check_status(library->schema(source, &schema)); + CHECK(schema.release != NULL && schema.n_children == 2, "Unexpected source schema"); + CHECK(strcmp(schema.format, "+s") == 0, "The schema must describe a struct"); + CHECK(strcmp(schema.children[0]->name, "id") == 0, "Missing id field"); + CHECK(strcmp(schema.children[0]->format, "l") == 0, "id must be int64"); + CHECK(strcmp(schema.children[1]->name, "label") == 0, "Missing label field"); + CHECK(strcmp(schema.children[1]->format, "u") == 0, "label must be UTF-8"); + + void* reader = NULL; + check_status(library->create_reader(source, &schema, &reader)); + CHECK(reader != NULL, "create_reader returned a null reader"); + if (schema.release != NULL) { + schema.release(&schema); + } + library->close(source); + + if (predicate_count != 0) { + int32_t* uses = calloc((size_t) predicate_count, sizeof(int32_t)); + CHECK(uses != NULL, "Cannot allocate predicate results"); + check_status(library->push_predicates(reader, predicates, predicate_count, uses)); + for (int64_t i = 0; i < predicate_count; i++) { + CHECK(uses[i] == expected_uses[i], "Incorrect C predicate-use enum value"); + } + free(uses); + } + if (limit >= 0) { + CHECK(predicate_count == 0, "Spark only pushes a limit without predicates"); + bool applied = false; + check_status(library->push_limit(reader, limit, &applied)); + CHECK(applied, "The Range source must apply its global limit"); + } + if (projection_count >= 0) { + bool pruned = false; + check_status(library->prune_columns(reader, projection, projection_count, &pruned)); + CHECK(pruned, "The Range source must apply its projection"); + } + + Plan plan = {0}; + check_status(library->partitions(reader, add_bytes, &plan.partitions)); + ByteList repeated = {0}; + check_status(library->partitions(reader, add_bytes, &repeated)); + CHECK(repeated.length == plan.partitions.length, "Repeated planning changed partition count"); + for (size_t i = 0; i < repeated.length; i++) { + check_bytes(&plan.partitions.values[i], &repeated.values[i]); + } + free_list(&repeated); + + ByteList state = {0}; + check_status(library->serialize_reader(reader, add_bytes, &state)); + check_status(library->serialize_reader(reader, add_bytes, &state)); + CHECK(state.length == 2, "Each serialization must emit exactly one state value"); + check_bytes(&state.values[0], &state.values[1]); + plan.state = state.values[0]; + free(state.values[1].data); + free(state.values); + library->close_reader(reader); + return plan; +} + +static void free_plan(Plan* plan) { + free_list(&plan->partitions); + free(plan->state.data); +} + +static struct ArrowArrayStream open_stream( + const Plan* plan, size_t partition, const SparkDsTaskContext* task) { + struct ArrowArrayStream stream = {0}; + const Bytes* bytes = &plan->partitions.values[partition]; + check_status(library->read(task, plan->state.data, plan->state.length, + bytes->data, bytes->length, &stream)); + CHECK(stream.release != NULL && stream.get_schema != NULL && stream.get_next != NULL, + "read returned an incomplete Arrow stream"); + return stream; +} + +static void check_batch( + const struct ArrowSchema* schema, + const struct ArrowArray* batch, + int64_t first_id, + int64_t max_batch_size) { + CHECK(batch->offset == 0, "The root array must have zero offset"); + CHECK(batch->length > 0 && batch->length <= max_batch_size, "Unexpected batch row count"); + CHECK(batch->n_children == schema->n_children, "Batch and schema have different column counts"); + for (int64_t column = 0; column < schema->n_children; column++) { + const struct ArrowSchema* field = schema->children[column]; + const struct ArrowArray* array = batch->children[column]; + CHECK(array->offset == 0 && array->length == batch->length, "Unexpected child array layout"); + if (strcmp(field->name, "id") == 0) { + CHECK(strcmp(field->format, "l") == 0, "Unexpected id type"); + const int64_t* ids = array->buffers[1]; + for (int64_t row = 0; row < batch->length; row++) { + CHECK(ids[row] == first_id + row, "Unexpected id value"); + } + } else { + CHECK(strcmp(field->name, "label") == 0, "Unexpected output field"); + CHECK(strcmp(field->format, "u") == 0, "Unexpected label type"); + const int32_t* offsets = array->buffers[1]; + const char* data = array->buffers[2]; + for (int64_t row = 0; row < batch->length; row++) { + char expected[64]; + int length = snprintf(expected, sizeof(expected), "row-%" PRId64, first_id + row); + CHECK(offsets[row + 1] - offsets[row] == length, "Unexpected label length"); + CHECK(memcmp(data + offsets[row], expected, (size_t) length) == 0, + "Unexpected label value"); + } + } + } +} + +static void consume_plan( + const Plan* plan, + int64_t first_id, + int64_t expected_rows, + const char* const* columns, + int64_t column_count, + bool short_context) { + int64_t rows = 0; + for (size_t partition = 0; partition < plan->partitions.length; partition++) { + SparkDsTaskContext task = task_context(short_context); + expected_task = &task; + interrupted = false; + struct ArrowArrayStream stream = open_stream(plan, partition, &task); + struct ArrowSchema schema = {0}; + CHECK(stream.get_schema(&stream, &schema) == 0, "Cannot read the stream schema"); + CHECK(schema.n_children == column_count, "Unexpected projected column count"); + for (int64_t i = 0; i < column_count; i++) { + CHECK(strcmp(schema.children[i]->name, columns[i]) == 0, + "Projection did not preserve the requested column order"); + } + + struct ArrowArray* batches = NULL; + size_t count = 0; + while (true) { + struct ArrowArray batch = {0}; + int status = stream.get_next(&stream, &batch); + if (status != 0) { + fprintf(stderr, "Arrow error: %s\n", stream.get_last_error(&stream)); + } + CHECK(status == 0, "Cannot read a batch"); + if (batch.release == NULL) { + break; + } + struct ArrowArray* resized = realloc(batches, (count + 1) * sizeof(struct ArrowArray)); + CHECK(resized != NULL, "Cannot retain batches"); + batches = resized; + batches[count++] = batch; + } + stream.release(&stream); + CHECK(stream.release == NULL, "Releasing the stream must clear its release callback"); + + /* The arrays retain their buffers across later reads and the release of their stream. */ + for (size_t i = 0; i < count; i++) { + check_batch(&schema, &batches[i], first_id + rows, 2); + rows += batches[i].length; + batches[i].release(&batches[i]); + } + free(batches); + schema.release(&schema); + } + CHECK(rows == expected_rows, "Unexpected total row count"); +} + +static void check_cancellation(const Plan* plan) { + CHECK(plan->partitions.length == 1, "Cancellation test requires one partition"); + SparkDsTaskContext task = task_context(false); + expected_task = &task; + interrupted = false; + interruption_checks = 0; + struct ArrowArrayStream stream = open_stream(plan, 0, &task); + struct ArrowSchema schema = {0}; + CHECK(stream.get_schema(&stream, &schema) == 0, "Cannot read the stream schema"); + struct ArrowArray first = {0}; + CHECK(stream.get_next(&stream, &first) == 0 && first.release != NULL, "Missing first batch"); + CHECK(interruption_checks > 0, "The reader did not check the task callback"); + int previous_checks = interruption_checks; + interrupted = true; + struct ArrowArray next = {0}; + CHECK(stream.get_next(&stream, &next) != 0, + "A cancelled task must fail before its next batch"); + CHECK(interruption_checks > previous_checks, "Cancellation did not call the task callback"); + const char* error = stream.get_last_error(&stream); + CHECK(error != NULL && strstr(error, "cancel") != NULL, "Missing cancellation error"); + if (next.release != NULL) { + next.release(&next); + } + stream.release(&stream); + /* Early stream release must not invalidate the batch already handed to its consumer. */ + check_batch(&schema, &first, 0, 2); + first.release(&first); + schema.release(&schema); + interrupted = false; +} + +static void check_error_message(void) { + const char* keys[] = {"batchsize"}; + const char* values[] = {"0"}; + SparkDsCreateInfo info = { + .size = sizeof(SparkDsCreateInfo), + .name = "rust_sdk_range", + .option_keys = keys, + .option_values = values, + .num_options = 1 + }; + uint64_t capabilities = 0; + void* source = NULL; + CHECK(library->create(&info, &capabilities, &source) != 0, "Invalid options must return an error"); + const char* error = library->last_error(); + CHECK(error != NULL && strstr(error, "batchSize must be positive") != NULL, + "The C error channel did not preserve the provider's message"); + CHECK(source == NULL, "Failed creation must not return a source handle"); +} + +static void check_function_table(void) { + CHECK(library != NULL, "spark_ds_init returned no library"); + CHECK(library->size == sizeof(SparkDsLibrary), "The function table differs from Spark's header"); + CHECK(library->abi_version == SPARK_DS_ABI_VERSION, "Unexpected ABI version"); + CHECK(library->last_error != NULL, "Missing last_error function"); + CHECK(library->create != NULL && library->schema != NULL && library->close != NULL, + "Missing source functions"); + CHECK(library->create_reader != NULL && library->push_predicates != NULL && + library->push_limit != NULL && library->prune_columns != NULL && + library->partitions != NULL && library->serialize_reader != NULL && + library->close_reader != NULL && library->read != NULL, "Missing batch read functions"); + CHECK(library->create_stream_reader == NULL && library->initial_offset == NULL && + library->latest_offset == NULL && library->stream_partitions == NULL && + library->serialize_stream_reader == NULL && library->commit_offset == NULL && + library->close_stream_reader == NULL, "The batch SDK must not export streaming reads"); + CHECK(library->create_writer == NULL && library->serialize_writer == NULL && + library->commit == NULL && library->abort == NULL && library->close_writer == NULL && + library->create_data_writer == NULL && library->write == NULL && + library->commit_data_writer == NULL && library->abort_data_writer == NULL && + library->close_data_writer == NULL, "The batch SDK must not export writes"); +} + +int main(int argc, char** argv) { + CHECK(argc == 2, "Expected the path to the Range shared library"); + void* module = dlopen(argv[1], RTLD_NOW | RTLD_LOCAL); + if (module == NULL) { + fprintf(stderr, "Cannot load the Range library: %s\n", dlerror()); + return 1; + } + void* symbol = dlsym(module, "spark_ds_init"); + CHECK(symbol != NULL, "The library does not export spark_ds_init"); + int (*init)(const SparkDsHost*, const SparkDsLibrary**); + CHECK(sizeof(init) == sizeof(symbol), "This platform has incompatible function pointers"); + memcpy(&init, &symbol, sizeof(init)); + SparkDsHost host = { + .size = sizeof(SparkDsHost), + .abi_version = SPARK_DS_ABI_VERSION, + .spark_version = "4.4.0-c-abi-test", + .log = log_message + }; + CHECK(init(&host, &library) == 0, "spark_ds_init failed"); + check_function_table(); + + const char* keys[] = {"start", "end", "numpartitions", "batchsize"}; + const char* values[] = {"10", "25", "3", "2"}; + const char* lower = "{\"type\":\"function\",\"name\":\">=\",\"children\":[" + "{\"type\":\"column\",\"name\":[\"id\"]}," + "{\"type\":\"literal\",\"dataType\":\"bigint\",\"value\":13}]}"; + const char* upper = "{\"type\":\"function\",\"name\":\"<\",\"children\":[" + "{\"type\":\"column\",\"name\":[\"id\"]}," + "{\"type\":\"literal\",\"dataType\":\"bigint\",\"value\":22}]}"; + const char* label = "{\"type\":\"function\",\"name\":\"=\",\"children\":[" + "{\"type\":\"column\",\"name\":[\"label\"]}," + "{\"type\":\"literal\",\"dataType\":\"string\",\"value\":\"row-21\"}]}"; + const char* predicates[] = {lower, upper, label}; + const int32_t uses[] = { + SPARK_DS_PREDICATE_COMPLETE, SPARK_DS_PREDICATE_COMPLETE, SPARK_DS_PREDICATE_UNUSED + }; + const char* reordered[] = {"label", "id"}; + const char* id_only[] = {"id"}; + const char* original[] = {"id", "label"}; + + Plan filtered = plan_scan(keys, values, 4, predicates, uses, 2, -1, reordered, 2); + consume_plan(&filtered, 13, 9, reordered, 2, false); + free_plan(&filtered); + + Plan residual = plan_scan(keys, values, 4, predicates, uses, 3, -1, id_only, 1); + consume_plan(&residual, 13, 9, id_only, 1, false); + free_plan(&residual); + + Plan limited = plan_scan(keys, values, 4, NULL, NULL, 0, 5, id_only, 1); + consume_plan(&limited, 10, 5, id_only, 1, false); + free_plan(&limited); + + const char* short_keys[] = {"end", "numpartitions", "batchsize"}; + const char* short_values[] = {"7", "2", "2"}; + Plan empty = plan_scan(short_keys, short_values, 3, NULL, NULL, 0, -1, NULL, 0); + consume_plan(&empty, 0, 7, NULL, 0, false); + free_plan(&empty); + + const char* one_partition[] = {"20", "1", "2"}; + Plan cancellable = plan_scan(short_keys, one_partition, 3, NULL, NULL, 0, -1, NULL, -1); + consume_plan(&cancellable, 0, 20, original, 2, true); + check_cancellation(&cancellable); + free_plan(&cancellable); + check_error_message(); + + CHECK(dlclose(module) == 0, "Cannot close the Range library"); + puts("C ABI, capabilities, planning, Arrow ownership, and cancellation checks passed."); + return 0; +} diff --git a/examples/src/main/rust/native-data-source/tests/c_abi_smoke.rs b/examples/src/main/rust/native-data-source/tests/c_abi_smoke.rs new file mode 100644 index 0000000000000..78d7b5e1657c8 --- /dev/null +++ b/examples/src/main/rust/native-data-source/tests/c_abi_smoke.rs @@ -0,0 +1,87 @@ +// 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> { + 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")); + let header = manifest.join("../../../../../sql/native/rust/tests/fixtures"); + 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"]) + .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() + ); + + let result = Command::new(harness).arg(library).output()?; + assert!( + result.status.success(), + "C ABI harness failed ({}):\n{}\n{}", + result.status, + String::from_utf8_lossy(&result.stdout), + String::from_utf8_lossy(&result.stderr) + ); + Ok(()) +} diff --git a/sql/native/rust/Cargo.lock b/sql/native/rust/Cargo.lock new file mode 100644 index 0000000000000..2e411f09d9e22 --- /dev/null +++ b/sql/native/rust/Cargo.lock @@ -0,0 +1,664 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +# 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. +version = 4 + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "const-random", + "getrandom 0.3.4", + "once_cell", + "version_check", + "zerocopy", +] + +[[package]] +name = "android_system_properties" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" +dependencies = [ + "libc", +] + +[[package]] +name = "arrow-array" +version = "58.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae33dad492b7df00a217563a7b0ef2874df68a0deea1b1a3acf628152f7f7a69" +dependencies = [ + "ahash", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "hashbrown", + "num-complex", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-buffer" +version = "58.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9552f96391c005e6ab449fa941420935e7e062489b12b8b1b08879b2163f5b5" +dependencies = [ + "bytes", + "half", + "num-bigint", + "num-traits", +] + +[[package]] +name = "arrow-data" +version = "58.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b24852db04738907e06c04ea61e42fe7fda962a34513022dc0d0e754fb7976b" +dependencies = [ + "arrow-buffer", + "arrow-schema", + "half", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-schema" +version = "58.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21ca356ad6425cecb6eb7b28e4f659f1ee7880fbb1a16127de7dd62901efee9e" +dependencies = [ + "bitflags", + "serde", + "serde_core", +] + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "cc" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f74872d07caf508b30a21f6836e7d7016a2eaf7d9ff4f48deaa58cd8a0407630" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "num-traits", + "windows-link", +] + +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom 0.2.17", + "once_cell", + "tiny-keccak", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "find-msvc-tools" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aedcfb3409746eddb02b9e19ebda1c3394f759a152e48ee875a0844d1b955484" + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7883d941dae510fb2d978fc3fe018c71c9e2892fd38854de3e8b92c2e5ad9cc5" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "libc" +version = "0.2.190" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce5d3ddc6d3fa000eb1536d85e147bfe31aacaba692ed6a876f95cb7c855be78" + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "log" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", + "libm", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "spark-native-range-demo" +version = "0.1.0" +dependencies = [ + "arrow-array", + "arrow-schema", + "serde", + "serde_json", + "spark-native-source", +] + +[[package]] +name = "spark-native-source" +version = "0.1.0" +dependencies = [ + "arrow-array", + "arrow-data", + "arrow-schema", + "serde", + "serde_json", +] + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8593e8e72159ed2257d083c7a454a85cbf854f37a0966d8d483aff8c8a3ebcee" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + +[[package]] +name = "unicode-ident" +version = "1.0.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bb54f33acc68fd454578d9820b0bde1a1a3d17aa17bb7b6595806d02886d409" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e29d0c35b16e224a7eeb5cd2d25e3e1968fbd65604117b44d3b789d00ee8535" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f501a8bc3719dba86ef8ae4728879c08001bea749eb1333ac5b91e040e2a6b7" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 3.0.6", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.129" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23f0c9c52aa7cd7d77769a4cfe2a9adb1b331f489a41d912ce14513d5ab995c6" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zerocopy" +version = "0.8.60" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6400fc4bb7426f6faac3fb07566122bb18622e3022bd2c36426a92ed489b1dbe" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.60" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "99c27cc7525bad4df59a45e073c45b2c2aa820efb588234cc23a547ec69a34db" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/sql/native/rust/Cargo.toml b/sql/native/rust/Cargo.toml new file mode 100644 index 0000000000000..6b325226ba0a0 --- /dev/null +++ b/sql/native/rust/Cargo.toml @@ -0,0 +1,51 @@ +# 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. + +[package] +name = "spark-native-source" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +description = "Experimental Rust batch data source SDK for Apache Spark PR #59424" + +[workspace] +members = ["../../../examples/src/main/rust/native-data-source"] +default-members = ["."] +resolver = "2" + +[workspace.package] +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +publish = false + +[workspace.dependencies] +arrow-array = { version = "58", features = ["ffi"] } +arrow-schema = { version = "58", features = ["ffi", "serde"] } +arrow-data = "58" +serde = { version = "1", features = ["derive"] } +serde_json = "1" + +[dependencies] +arrow-array.workspace = true +arrow-schema.workspace = true +arrow-data.workspace = true +serde.workspace = true +serde_json.workspace = true + +[profile.release] +# Provider panics are translated into C ABI / Arrow errors by the SDK. +panic = "unwind" diff --git a/sql/native/rust/README.md b/sql/native/rust/README.md new file mode 100644 index 0000000000000..4749102680f5d --- /dev/null +++ b/sql/native/rust/README.md @@ -0,0 +1,200 @@ + + +# Rust native data source SDK prototype + +This experimental SDK implements **batch reads** over the C ABI proposed in +[Spark PR #59424](https://github.com/apache/spark/pull/59424), pinned to commit +`51af65c63a9369d793aab1a198ddde4c90c8f85b`. Spark calls the library through the +Java Foreign Function and Memory (FFM) API. This PR contains the SDK and example, +not the Spark framework; running the Spark example requires that framework and +Java 22 or newer on the driver and executors. + +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. Both crates share this directory's lockfile and +`target` directory. Cargo commands select the SDK by default; use `--workspace` +to include the example. + +## Connector API and lifecycle + +A connector implements three traits from [`src/lib.rs`](src/lib.rs): + +* `BatchSource`: parse driver options, discover the schema, and create an + independent scan builder. Spark can drop the source immediately afterwards. +* `ScanBuilder`: negotiate predicates, limit, and projection, then consume itself + in `plan()` to return a typed `ScanPlan`. +* `BatchReaderFactory`: serialize executor configuration and open one lazy Arrow + reader from a typed partition and a borrowed `TaskContext`. + +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: + BatchSource -> ScanBuilder -> ScanPlan { factory, partitions } + | | + mutable negotiation immutable serialized descriptors + | +Executor, per task: v + restored factory + partition + TaskContext -> Arrow reader -> Arrow C Stream +``` + +`ScanBuilder::plan()` runs on the first native `partitions` or `serialize_reader` +call, after pushdown negotiation. It is not the Java DSv2 `ScanBuilder.build()` +callback. The driver runtime caches the encoded result and shares it across +repeated callbacks without cloning its byte buffers. A failed negotiation or +planning call is terminal and retains its original error, including a panic. +Closing a driver handle releases its local resources; it cannot invalidate +factory or partition bytes already copied to Spark. + +The reader owns its executor resources and cannot borrow its temporary factory. +It may borrow `TaskContext` until the returned Arrow stream is released. Task IDs +are optional for compatibility with shorter ABI structures. `is_interrupted()` +checks Spark's live cancellation callback, not a snapshot of its initial value. +Check it during long operations; blocking I/O still requires timeouts. + +Register one provider in 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. `create` declares only `SPARK_DS_BATCH_READ`; streaming, write, and +truncate capabilities are absent, and their function slots are null. Query-level +pushdown results are independent of these source capabilities. The connector +uses safe Rust traits without C pointers. Rust trait objects never cross the +binary boundary between libraries. + +The implementation separates public connector traits, driver handle/planning +state in `bridge.rs`, Arrow ownership in `stream.rs`, and C bindings/adapters in +`abi.rs` and `export.rs`. It does not change Spark's native ABI or JVM classes. + +## Arrow contract + +`BatchReaderFactory::open` returns a lifetime-bound `Box`. Return fixed-schema batches lazily with bounded memory. Exported arrays +own their buffers independently of the reader: Spark may release a stream before +releasing its previously returned batches. + +The SDK validates the restored negotiated schema before exporting a reader and +checks each batch's schema and recursive array offsets. It compares Spark +logical types, ignoring nullability and Arrow metadata, accepting the supported +string/binary encodings and equivalent timestamp time zones. It rejects +unsupported types, schema drift, and nonzero array offsets with explicit errors. +It does not silently copy arrays, cast types, or expand dictionary encoding. + +Supported types include signed integers, float32/64, boolean, strings and binary +(including Large and View variants), decimal128 up to precision 38, date32, +microsecond timestamps, supported Spark intervals, null, structs, lists, and +maps. Dictionary encoding, unsigned integers, decimal256, large/fixed-size +lists, and non-microsecond timestamps are outside this prototype's contract. +Normalize sliced arrays when necessary before returning them. + +An empty projection must retain the number of rows even with no columns. The +[Range implementation][range-source] demonstrates this together with bounded +batches, disjoint partitions, exact integer filtering, and cooperative +cancellation. + +## Pushdown contracts + +Results correspond to input predicates in order: + +| Result | C enum | Connector responsibility | +| --- | --- | --- | +| `Exact` | `SPARK_DS_PREDICATE_COMPLETE` | Return only matching rows under Spark SQL semantics | +| `PruningOnly` | `SPARK_DS_PREDICATE_PARTIAL` | Skip only data that cannot match | +| `Unsupported` | `SPARK_DS_PREDICATE_UNUSED` | No evaluation | + +The pinned framework rechecks **all** filters, including `Exact`. Unknown +expression names and unsupported types must remain unsupported. The SDK +preserves literal JSON values, including exact 64-bit integers, rather than +converting all numbers to floating point. It is not a generic Spark expression +evaluator; connectors remain responsible for Spark SQL semantics. + +| Limit result | ABI boolean | Meaning | +| --- | --- | --- | +| `Unsupported` | `false` | No limit applied | +| `PerPartition` | `false` | A local cap; the total may exceed the limit | +| `Global` | `true` | The entire scan returns at most the limit | + +Spark retains its final limit and currently pushes a limit only when the scan +has no filters. The SDK also declines limit pushdown after any predicate that is +not `Exact`. Limiting rows before a residual filter can omit matching rows, even +when Spark retains its final limit. + +Projection returns true only when the reader produces exactly the requested +top-level columns in order. Keep internally needed predicate columns even if +they are absent from the output projection. All pushdown methods have +conservative defaults so a connector can add support incrementally. + +## Build and test without Spark + +Use Rust 1.85 or newer (as required by Arrow 58) and a C compiler. The C integration +test uses `dlopen` on Linux or macOS. From this directory: + +```bash +cargo build --workspace --locked +cargo test --workspace --locked +cargo clippy --workspace --all-targets --locked -- -D warnings +cargo fmt --all -- --check +``` + +Build before testing so the C harness loads the current example dynamic library. +The harness compiles against the unmodified framework header in +[`tests/fixtures/spark_datasource.h`](tests/fixtures/spark_datasource.h), copied +from the pinned commit above. It calls the actual function table and validates +capabilities and optional slots, driver handle independence, pushdown and +planning, serialized executor reads, cancellation, and Arrow buffer ownership. +This checks the C boundary without requiring Java; it does not substitute for +Spark planner and task-lifecycle testing. + +Unit tests additionally cover error and panic containment, terminal planning +failures, schema/layout validation, sized ABI structures, task-context lifetimes, +integer boundaries, projections, and deterministic partition plans. + +## Package and run with Spark + +See the [Range example guide][range] for packaging and `spark-submit` commands. +Its `package.py` creates a single-platform `.sparkpkg`; `smoke.py` checks reading, +filtered reads, projection, and unfiltered limit pushdown against a Spark +installation containing the pinned #59424 framework. No full Spark integration +is implied by the standalone Rust/C tests. + +## Scope and limitations + +Native sources run trusted code in the driver and executors. The SDK translates +provider errors and unwinding panics into C/Arrow errors. Process aborts, +allocation failure, or panics in custom Arrow buffer destructors cannot be +recovered. A reader's custom background work must stop before its task context +expires. Native memory is not reserved through Spark's memory manager. + +This prototype does not include writes, Structured Streaming, nested projection, +custom metrics, memory accounting, a generic filter evaluator, C++ bindings, or +an ABI extension. The versioned JSON envelope is an SDK implementation detail; +use the same plugin build on the driver and executors. Only the pinned framework +revision is targeted until the native API stabilizes. + +[range]: ../../../examples/src/main/rust/native-data-source/README.md +[range-source]: ../../../examples/src/main/rust/native-data-source/src/lib.rs diff --git a/sql/native/rust/src/abi.rs b/sql/native/rust/src/abi.rs new file mode 100644 index 0000000000000..311ce26cde315 --- /dev/null +++ b/sql/native/rust/src/abi.rs @@ -0,0 +1,362 @@ +// 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::panic::{catch_unwind, AssertUnwindSafe}; +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::panic_message; +use crate::{FilterSupport, Options, Result}; + +pub const ABI_VERSION: i32 = 1; +pub const BATCH_READ: u64 = 1; +pub const LOG_ERROR: i32 = 1; + +pub type SparkDsSink = Option; +pub type Interrupted = unsafe extern "C" fn(*const SparkDsTaskContext) -> bool; + +#[repr(C)] +pub struct SparkDsHost { + pub size: i64, + pub abi_version: i32, + pub spark_version: *const c_char, + pub log: Option, +} + +#[repr(C)] +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)] +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, +} + +#[repr(C)] +pub struct SparkDsWriteInfo { + pub size: i64, + pub mode: i32, + pub query_id: *const c_char, +} + +/// The complete version 1 table. Unsupported operations have null function slots. +#[repr(C)] +pub struct SparkDsLibrary { + pub size: i64, + pub abi_version: i32, + pub last_error: Option *const c_char>, + pub create: + Option c_int>, + pub schema: Option c_int>, + pub create_reader: + Option c_int>, + pub create_stream_reader: + Option c_int>, + pub create_writer: Option< + unsafe extern "C" fn( + *mut c_void, + *mut ArrowSchema, + *const SparkDsWriteInfo, + *mut *mut c_void, + ) -> c_int, + >, + pub close: Option, + pub push_predicates: + Option c_int>, + pub push_limit: Option c_int>, + pub prune_columns: + Option c_int>, + pub partitions: Option c_int>, + pub serialize_reader: + Option c_int>, + pub close_reader: Option, + pub initial_offset: + Option c_int>, + pub latest_offset: Option c_int>, + pub stream_partitions: Option< + unsafe extern "C" fn( + *mut c_void, + *const c_char, + *const c_char, + SparkDsSink, + *mut c_void, + ) -> c_int, + >, + pub serialize_stream_reader: + Option c_int>, + pub commit_offset: Option c_int>, + pub close_stream_reader: Option, + pub read: Option< + unsafe extern "C" fn( + *const SparkDsTaskContext, + *const u8, + i64, + *const u8, + i64, + *mut ArrowArrayStream, + ) -> c_int, + >, + pub serialize_writer: + Option c_int>, + pub commit: + Option c_int>, + pub abort: + Option c_int>, + pub close_writer: Option, + pub create_data_writer: Option< + unsafe extern "C" fn( + *const SparkDsTaskContext, + *const u8, + i64, + i64, + *mut *mut c_void, + ) -> c_int, + >, + pub write: + Option c_int>, + pub commit_data_writer: + Option c_int>, + pub abort_data_writer: Option c_int>, + pub close_data_writer: Option, +} + +thread_local! { + static LAST_ERROR: RefCell> = const { RefCell::new(None) }; +} + +/// Converts provider failures and unwinding panics into the C error convention. +pub fn call(action: impl FnOnce() -> Result<()>) -> c_int { + LAST_ERROR.with(|error| *error.borrow_mut() = None); + let message = match catch_unwind(AssertUnwindSafe(|| action().map_err(|e| e.to_string()))) { + Ok(Ok(())) => return 0, + Ok(Err(message)) => message, + Err(payload) => panic_message(payload), + }; + let message = CString::new(message.replace('\0', "\\0")).expect("NULs removed"); + LAST_ERROR.with(|error| *error.borrow_mut() = Some(message)); + 1 +} + +/// Returns the calling thread's error until its next adapter call. +/// +/// # Safety +/// The returned pointer must not be freed or used after another adapter call on this thread. +pub unsafe extern "C" fn last_error() -> *const c_char { + LAST_ERROR.with(|error| { + error + .borrow() + .as_ref() + .map_or(std::ptr::null(), |message| message.as_ptr()) + }) +} + +/// Reads only a field that fits in a sized C structure. +/// +/// # Safety +/// `pointer` must address `size` readable bytes, and `offset` must identify a field of type `T`. +pub unsafe fn field(pointer: *const u8, size: i64, offset: usize) -> Option { + let end = offset.checked_add(size_of::())?; + if usize::try_from(size).ok()? < end { + return None; + } + Some(unsafe { pointer.add(offset).cast::().read_unaligned() }) +} + +/// # Safety +/// A non-null pointer must address the size prefix of a live C structure. +pub unsafe fn struct_size(pointer: *const T) -> Result { + if pointer.is_null() { + return Err("null sized structure".into()); + } + let size = unsafe { pointer.cast::().read_unaligned() }; + if size < size_of::() as i64 { + return Err("structure is smaller than its size field".into()); + } + Ok(size) +} + +pub fn require_output(pointer: *mut T) -> Result<()> { + if pointer.is_null() { + return Err("null output pointer".into()); + } + Ok(()) +} + +/// # Safety +/// `host` and `output` must follow the `spark_ds_init` contract of the C header. +pub unsafe fn initialize( + host: *const SparkDsHost, + output: *mut *const SparkDsLibrary, + library: &'static SparkDsLibrary, +) -> c_int { + let result = call(|| { + require_output(output)?; + unsafe { output.write(std::ptr::null()) }; + let size = unsafe { struct_size(host)? }; + let version: i32 = + unsafe { field(host.cast(), size, offset_of!(SparkDsHost, abi_version)) } + .ok_or("host does not provide an ABI version")?; + if version < ABI_VERSION { + return Err(format!("host ABI {version} does not support ABI {ABI_VERSION}").into()); + } + unsafe { output.write(library) }; + Ok(()) + }); + if result != 0 { + if let Ok(size) = unsafe { struct_size(host) } { + let log: Option> = + unsafe { field(host.cast(), size, offset_of!(SparkDsHost, log)) }; + if let Some(Some(log)) = log { + unsafe { log(LOG_ERROR, last_error()) }; + } + } + } + result +} + +/// # Safety +/// `data` must address `length` initialized values that remain valid for the returned borrow. +pub unsafe fn slice<'a, T>(data: *const T, length: i64) -> Result<&'a [T]> { + let length = usize::try_from(length)?; + if length == 0 { + return Ok(&[]); + } + if data.is_null() + || length + .checked_mul(size_of::()) + .is_none_or(|size| size > isize::MAX as usize) + { + return Err("invalid array pointer or length".into()); + } + Ok(unsafe { std::slice::from_raw_parts(data, length) }) +} + +/// # Safety +/// Each element must be a NUL-terminated UTF-8 string, readable during this call. +pub unsafe fn strings(data: *const *const c_char, length: i64) -> Result> { + unsafe { slice(data, length)? } + .iter() + .map(|pointer| { + if pointer.is_null() { + return Err("null string".into()); + } + Ok(unsafe { CStr::from_ptr(*pointer) }.to_str()?.to_owned()) + }) + .collect() +} + +/// # Safety +/// `info` must follow the `create` contract of the C header. +pub unsafe fn options(info: *const SparkDsCreateInfo) -> Result { + let size = unsafe { struct_size(info)? }; + let pointer = info.cast(); + let keys = unsafe { field(pointer, size, offset_of!(SparkDsCreateInfo, option_keys)) } + .ok_or("create info does not provide option keys")?; + let values = unsafe { field(pointer, size, offset_of!(SparkDsCreateInfo, option_values)) } + .ok_or("create info does not provide option values")?; + let count = unsafe { field(pointer, size, offset_of!(SparkDsCreateInfo, num_options)) } + .ok_or("create info does not provide an option count")?; + Ok(unsafe { strings(keys, count)? } + .into_iter() + .zip(unsafe { strings(values, count)? }) + .collect()) +} + +pub fn predicate_use(support: FilterSupport) -> i32 { + match support { + FilterSupport::Unsupported => 0, + FilterSupport::PruningOnly => 1, + FilterSupport::Exact => 2, + } +} + +/// # Safety +/// The sink must be valid on this thread for this call and synchronously copy the bytes. +pub unsafe fn emit(callback: SparkDsSink, sink: *mut c_void, bytes: &[u8]) -> Result<()> { + let callback = callback.ok_or("null result sink")?; + unsafe { callback(sink, bytes.as_ptr(), bytes.len().try_into()?) }; + Ok(()) +} + +/// # Safety +/// `input` must be a live ArrowSchema that this call is permitted to move. +pub unsafe fn import_schema(input: *mut ArrowSchema) -> Result { + require_output(input)?; + let schema = unsafe { ArrowSchema::from_raw(input) }; + Ok(Arc::new(Schema::try_from(&schema)?)) +} + +/// # Safety +/// `output` must be writable ArrowSchema storage with a null release callback. +pub unsafe fn export_schema(schema: &SchemaRef, output: *mut ArrowSchema) -> Result<()> { + require_output(output)?; + let schema = ArrowSchema::try_from(schema.as_ref())?; + unsafe { output.write(schema) }; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn errors_and_panics_are_thread_local_c_strings() { + assert_ne!(call(|| Err("failed\0operation".into())), 0); + let error = unsafe { CStr::from_ptr(last_error()) }.to_str().unwrap(); + assert_eq!(error, "failed\\0operation"); + std::thread::spawn(|| { + assert!(unsafe { last_error() }.is_null()); + assert_ne!(call(|| panic!("provider panic")), 0); + let error = unsafe { CStr::from_ptr(last_error()) }.to_str().unwrap(); + assert!(error.contains("provider panic")); + }) + .join() + .unwrap(); + assert_eq!( + unsafe { CStr::from_ptr(last_error()) }.to_str().unwrap(), + error + ); + assert_eq!(call(|| Ok(())), 0); + assert!(unsafe { last_error() }.is_null()); + } + + #[test] + fn empty_arrays_accept_null_and_truncated_create_info_fails() { + assert!(unsafe { strings(std::ptr::null(), 0) }.unwrap().is_empty()); + let prefix = size_of::() as i64; + let info = (&prefix as *const i64).cast::(); + assert!(unsafe { options(info) }.is_err()); + } +} diff --git a/sql/native/rust/src/bridge.rs b/sql/native/rust/src/bridge.rs new file mode 100644 index 0000000000000..a83a2be3c5e9a --- /dev/null +++ b/sql/native/rust/src/bridge.rs @@ -0,0 +1,677 @@ +// 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. + +//! Transport-independent driver handles and executor plan restoration. + +use std::any::Any; +use std::collections::HashMap; +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::{Schema, SchemaRef}; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; + +use crate::{ + BatchReaderFactory, BatchSource, FilterSupport, LimitSupport, Options, Predicate, Result, + ScanBuilder, TaskContext, +}; + +#[path = "stream.rs"] +mod stream; + +const STATE_VERSION: u32 = 1; + +type HandleMap = Mutex>>>; + +/// One registry per loaded plugin. Provider code runs under a per-handle lock, +/// so independent sources and scans do not serialize behind a global lock. +pub struct Bridge { + next_handle: AtomicI64, + sources: HandleMap, + readers: HandleMap>, +} + +impl Default for Bridge { + fn default() -> Self { + Self { + next_handle: AtomicI64::new(1), + sources: Mutex::new(HashMap::new()), + readers: Mutex::new(HashMap::new()), + } + } +} + +impl Bridge { + fn insert(&self, handles: &HandleMap, value: T) -> Result { + 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, options: Options) -> Result { + self.insert(&self.sources, S::new(options)?) + } + + pub fn schema(&self, handle: i64) -> Result { + let source = get(&self.sources, handle)?; + let schema = lock(&source).schema()?; + stream::validate_schema(&schema)?; + Ok(schema) + } + + pub fn create_reader(&self, handle: i64, schema: SchemaRef) -> Result { + stream::validate_schema(&schema)?; + let source = get(&self.sources, handle)?; + let builder = lock(&source).scan(schema.clone())?; + self.insert(&self.readers, Reader::new(builder, schema)) + } + + pub fn push_predicates( + &self, + handle: i64, + predicates: &[Predicate], + ) -> Result> { + let reader = get(&self.readers, handle)?; + let mut reader = lock(&reader); + let support = reader.negotiate(|builder| { + let support = builder.push_filters(predicates)?; + if support.len() != predicates.len() { + return Err("push_filters must return one result per predicate".into()); + } + Ok(support) + })?; + reader.has_residual |= support.iter().any(|s| *s != FilterSupport::Exact); + Ok(support) + } + + pub fn push_limit(&self, handle: i64, limit: usize) -> Result { + let reader = get(&self.readers, handle)?; + let mut reader = lock(&reader); + let has_residual = reader.has_residual; + reader.negotiate(|builder| { + if has_residual { + Ok(false) + } else { + Ok(builder.push_limit(limit)? == LimitSupport::Global) + } + }) + } + + pub fn prune_columns(&self, handle: i64, columns: &[String]) -> Result { + let reader = get(&self.readers, handle)?; + let mut reader = lock(&reader); + let indices: Vec<_> = columns + .iter() + .map(|column| reader.schema.index_of(column)) + .collect::>()?; + let schema = Arc::new(reader.schema.project(&indices)?); + let pruned = reader.negotiate(|builder| builder.prune_columns(columns))?; + if pruned { + reader.schema = schema; + } + Ok(pruned) + } + + pub fn planned(&self, handle: i64) -> Result> { + let reader = get(&self.readers, handle)?; + let result = lock(&reader).plan(); + result + } + + pub fn close_source(&self, handle: i64) -> Result<()> { + remove(&self.sources, handle); + Ok(()) + } + + pub fn close_reader(&self, handle: i64) -> Result<()> { + remove(&self.readers, handle); + Ok(()) + } + + /// Restores only serialized values; no driver-side handle is needed. + /// + /// # 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( + &self, + task: *const crate::abi::SparkDsTaskContext, + state: &[u8], + partition: &[u8], + ) -> Result { + let state: ReaderState<::Factory> = decode(state)?; + let task = unsafe { TaskContext::from_raw(task)? }; + let reader = state.factory.open(decode(partition)?, task)?; + unsafe { stream::export(reader, &Arc::new(state.schema)) } + } +} + +fn lock(mutex: &Mutex) -> 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(handles: &HandleMap, handle: i64) -> Result>> { + lock(handles) + .get(&handle) + .cloned() + .ok_or_else(|| format!("unknown or closed native handle: {handle}").into()) +} + +fn remove(handles: &HandleMap, 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 Reader { + phase: Phase, + schema: SchemaRef, + has_residual: bool, +} + +enum Phase { + Negotiating(B), + Planned(Arc), + Failed(String), +} + +pub struct EncodedPlan { + pub state: Vec, + pub partitions: Vec>, +} + +#[derive(Serialize, Deserialize)] +struct ReaderState { + factory: F, + schema: Schema, +} + +impl Reader { + fn new(builder: B, schema: SchemaRef) -> Self { + Self { + phase: Phase::Negotiating(builder), + schema, + has_residual: false, + } + } + + fn negotiate(&mut self, action: impl FnOnce(&mut B) -> Result) -> Result { + let builder = match &mut self.phase { + Phase::Negotiating(builder) => builder, + Phase::Planned(_) => return Err("scan has already been planned".into()), + Phase::Failed(error) => return Err(error.clone().into()), + }; + match provider_call(|| action(builder)) { + Ok(value) => Ok(value), + Err(error) => { + // A partially mutated builder must never be reused after failure. + let previous = std::mem::replace(&mut self.phase, Phase::Failed(error.clone())); + // Install the failure before dropping provider state. Preserve + // the original cause even if its destructor also panics. + let _ = provider_call(|| { + drop(previous); + Ok(()) + }); + Err(error.into()) + } + } + } + + fn plan(&mut self) -> Result> { + match &self.phase { + Phase::Planned(plan) => return Ok(plan.clone()), + Phase::Failed(error) => return Err(error.clone().into()), + Phase::Negotiating(_) => {} + } + let phase = std::mem::replace( + &mut self.phase, + Phase::Failed("scan planning did not complete".into()), + ); + let Phase::Negotiating(builder) = phase else { + unreachable!("only a negotiating scan reaches this transition") + }; + let result = provider_call(|| { + let plan = builder.plan()?; + Ok(Arc::new(EncodedPlan { + state: encode(&ReaderState { + factory: plan.factory, + schema: self.schema.as_ref().clone(), + })?, + partitions: plan.partitions.iter().map(encode).collect::>()?, + })) + }); + match result { + Ok(plan) => { + self.phase = Phase::Planned(plan.clone()); + Ok(plan) + } + Err(error) => { + self.phase = Phase::Failed(error.clone()); + Err(error.into()) + } + } + } +} + +#[derive(Serialize, Deserialize)] +struct Envelope { + version: u32, + value: T, +} + +fn encode(value: &T) -> Result> { + Ok(serde_json::to_vec(&Envelope { + version: STATE_VERSION, + value, + })?) +} + +fn decode(bytes: &[u8]) -> Result { + let envelope: Envelope = serde_json::from_slice(bytes)?; + if envelope.version != STATE_VERSION { + return Err(format!("unsupported native state version: {}", envelope.version).into()); + } + Ok(envelope.value) +} + +fn provider_call(action: impl FnOnce() -> Result) -> std::result::Result { + catch_unwind(AssertUnwindSafe(|| { + action().map_err(|error| error.to_string()) + })) + .unwrap_or_else(|payload| Err(panic_message(payload))) +} + +pub(crate) fn panic_message(payload: Box) -> String { + let message = if let Some(message) = payload.downcast_ref::() { + message.clone() + } else if let Some(message) = payload.downcast_ref::<&str>() { + (*message).to_owned() + } else { + // An arbitrary panic payload can itself panic when dropped. + std::mem::forget(payload); + return "native provider panicked with a non-string payload".to_owned(); + }; + format!("native provider panicked: {message}") +} + +pub fn parse_predicates(predicates: &[String]) -> Result> { + predicates + .iter() + .map(|text| Ok(serde_json::from_str(text)?)) + .collect() +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicUsize; + + use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator}; + use arrow_schema::{DataType, Field}; + + use super::*; + use crate::{BatchReader, ScanPlan}; + + #[derive(Clone, Copy)] + enum FailureMode { + None, + Error, + Panic, + } + + impl FailureMode { + fn fail(self, message: &str) -> Result<()> { + match self { + Self::None => Ok(()), + Self::Error => Err(message.into()), + Self::Panic => panic!("{message}"), + } + } + } + + struct Source { + builds: Arc, + limits: Arc, + drops: Arc, + builder_drops: Arc, + filter_support: FilterSupport, + limit_support: LimitSupport, + failure: FailureMode, + } + + impl Drop for Source { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } + } + + struct Builder { + builds: Arc, + limits: Arc, + drops: Arc, + filter_support: FilterSupport, + limit_support: LimitSupport, + failure: FailureMode, + value: u64, + } + + impl Drop for Builder { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } + } + + impl BatchSource for Source { + type Scan = Builder; + + fn new(_: Options) -> Result { + Ok(Self { + builds: Arc::default(), + limits: Arc::default(), + drops: Arc::default(), + builder_drops: Arc::default(), + filter_support: FilterSupport::Unsupported, + limit_support: LimitSupport::Unsupported, + failure: FailureMode::None, + }) + } + + fn schema(&self) -> Result { + Ok(Arc::new(Schema::empty())) + } + + fn scan(&self, _: SchemaRef) -> Result { + Ok(Builder { + builds: self.builds.clone(), + limits: self.limits.clone(), + drops: self.builder_drops.clone(), + filter_support: self.filter_support, + limit_support: self.limit_support, + failure: self.failure, + value: 42, + }) + } + } + + #[derive(Serialize, Deserialize)] + struct Factory(u64); + + impl BatchReaderFactory for Factory { + type Partition = u64; + + fn open<'task>(&self, partition: u64, _: TaskContext<'task>) -> Result> { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let values = Arc::new(Int64Array::from(vec![(self.0 + partition) as i64])); + let batch = RecordBatch::try_new(schema.clone(), vec![values])?; + Ok(Box::new(RecordBatchIterator::new([Ok(batch)], schema))) + } + } + + impl ScanBuilder for Builder { + type Factory = Factory; + + fn push_filters(&mut self, filters: &[Predicate]) -> Result> { + self.failure.fail("negotiation failed")?; + Ok(vec![self.filter_support; filters.len()]) + } + + fn push_limit(&mut self, limit: usize) -> Result { + self.limits.fetch_add(1, Ordering::SeqCst); + self.value = limit as u64; + self.failure.fail("negotiation failed")?; + Ok(self.limit_support) + } + + fn plan(self) -> Result> { + self.builds.fetch_add(1, Ordering::SeqCst); + self.failure.fail("planning failed")?; + Ok(ScanPlan { + factory: Factory(self.value), + partitions: vec![2, 7], + }) + } + } + + fn error_message(result: Result) -> String { + result + .err() + .expect("expected a provider failure") + .to_string() + } + + #[test] + fn source_can_close_before_planning_and_plan_is_built_once() { + let bridge = Bridge::::default(); + let source = Source::new(Options::new()).unwrap(); + let builds = source.builds.clone(); + let drops = source.drops.clone(); + let handle = bridge.insert(&bridge.sources, source).unwrap(); + let schema = bridge.schema(handle).unwrap(); + let reader = bridge.create_reader(handle, schema).unwrap(); + bridge.close_source(handle).unwrap(); + bridge.close_source(handle).unwrap(); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert!(bridge.schema(handle).is_err()); + + let plan = bridge.planned(reader).unwrap(); + assert!(Arc::ptr_eq(&bridge.planned(reader).unwrap(), &plan)); + assert_eq!(decode::(&plan.partitions[0]).unwrap(), 2); + assert_eq!(decode::(&plan.partitions[1]).unwrap(), 7); + assert_eq!( + decode::>(&plan.state) + .unwrap() + .factory + .0, + 42 + ); + assert_eq!(builds.load(Ordering::SeqCst), 1); + assert!(bridge.prune_columns(reader, &[]).is_err()); + bridge.close_reader(reader).unwrap(); + assert!(bridge.planned(reader).is_err()); + } + + #[test] + fn residual_filters_prevent_calling_the_providers_limit_method() { + for support in [FilterSupport::PruningOnly, FilterSupport::Unsupported] { + let bridge = Bridge::::default(); + let mut source = Source::new(Options::new()).unwrap(); + source.filter_support = support; + source.limit_support = LimitSupport::Global; + let limits = source.limits.clone(); + let handle = bridge.insert(&bridge.sources, source).unwrap(); + let reader = bridge + .create_reader(handle, bridge.schema(handle).unwrap()) + .unwrap(); + assert_eq!( + bridge + .push_predicates(reader, &[Predicate::Unknown]) + .unwrap(), + vec![support] + ); + assert!(!bridge.push_limit(reader, 1).unwrap()); + assert_eq!(limits.load(Ordering::SeqCst), 0); + } + } + + #[test] + fn only_global_limits_return_true_to_the_abi() { + for support in [LimitSupport::PerPartition, LimitSupport::Global] { + let bridge = Bridge::::default(); + let mut source = Source::new(Options::new()).unwrap(); + source.limit_support = support; + let limits = source.limits.clone(); + let handle = bridge.insert(&bridge.sources, source).unwrap(); + let reader = bridge + .create_reader(handle, bridge.schema(handle).unwrap()) + .unwrap(); + assert_eq!( + bridge.push_limit(reader, 4).unwrap(), + support == LimitSupport::Global + ); + assert_eq!(limits.load(Ordering::SeqCst), 1); + } + } + + #[test] + fn planning_errors_and_panics_preserve_the_first_failure_without_retrying() { + for (failure, expected) in [ + (FailureMode::Error, "planning failed"), + ( + FailureMode::Panic, + "native provider panicked: planning failed", + ), + ] { + let bridge = Bridge::::default(); + let mut source = Source::new(Options::new()).unwrap(); + source.failure = failure; + let builds = source.builds.clone(); + let drops = source.builder_drops.clone(); + let handle = bridge.insert(&bridge.sources, source).unwrap(); + let reader = bridge + .create_reader(handle, bridge.schema(handle).unwrap()) + .unwrap(); + bridge.close_source(handle).unwrap(); + + assert_eq!(error_message(bridge.planned(reader)), expected); + assert_eq!(error_message(bridge.planned(reader)), expected); + assert_eq!(error_message(bridge.push_limit(reader, 1)), expected); + assert_eq!(builds.load(Ordering::SeqCst), 1); + assert_eq!(drops.load(Ordering::SeqCst), 1); + bridge.close_reader(reader).unwrap(); + assert_eq!(drops.load(Ordering::SeqCst), 1); + } + } + + #[test] + fn failed_negotiation_discards_the_partially_mutated_builder() { + for (failure, expected) in [ + (FailureMode::Error, "negotiation failed"), + ( + FailureMode::Panic, + "native provider panicked: negotiation failed", + ), + ] { + let bridge = Bridge::::default(); + let mut source = Source::new(Options::new()).unwrap(); + source.failure = failure; + let builds = source.builds.clone(); + let limits = source.limits.clone(); + let drops = source.builder_drops.clone(); + let handle = bridge.insert(&bridge.sources, source).unwrap(); + let reader = bridge + .create_reader(handle, bridge.schema(handle).unwrap()) + .unwrap(); + + assert_eq!(error_message(bridge.push_limit(reader, 5)), expected); + assert_eq!(error_message(bridge.push_limit(reader, 1)), expected); + assert_eq!(error_message(bridge.push_predicates(reader, &[])), expected); + assert_eq!(error_message(bridge.planned(reader)), expected); + assert_eq!(limits.load(Ordering::SeqCst), 1); + assert_eq!(builds.load(Ordering::SeqCst), 0); + assert_eq!(drops.load(Ordering::SeqCst), 1); + bridge.close_reader(reader).unwrap(); + bridge.close_source(handle).unwrap(); + } + } + + #[test] + fn scans_of_one_source_keep_pushdown_and_cached_plans_independent() { + let bridge = Bridge::::default(); + let mut source = Source::new(Options::new()).unwrap(); + source.filter_support = FilterSupport::PruningOnly; + source.limit_support = LimitSupport::Global; + let builds = source.builds.clone(); + let limits = source.limits.clone(); + let handle = bridge.insert(&bridge.sources, source).unwrap(); + let first = bridge + .create_reader(handle, bridge.schema(handle).unwrap()) + .unwrap(); + let second = bridge + .create_reader(handle, bridge.schema(handle).unwrap()) + .unwrap(); + bridge.close_source(handle).unwrap(); + + bridge + .push_predicates(second, &[Predicate::Unknown]) + .unwrap(); + assert!(bridge.push_limit(first, 5).unwrap()); + assert!(!bridge.push_limit(second, 1).unwrap()); + let first_plan = bridge.planned(first).unwrap(); + let second_plan = bridge.planned(second).unwrap(); + assert_eq!( + decode::>(&first_plan.state) + .unwrap() + .factory + .0, + 5 + ); + assert_eq!( + decode::>(&second_plan.state) + .unwrap() + .factory + .0, + 42 + ); + assert!(!Arc::ptr_eq(&first_plan, &second_plan)); + assert_eq!(builds.load(Ordering::SeqCst), 2); + assert_eq!(limits.load(Ordering::SeqCst), 1); + bridge.close_reader(first).unwrap(); + assert!(Arc::ptr_eq(&second_plan, &bridge.planned(second).unwrap())); + bridge.close_reader(second).unwrap(); + } + + #[test] + fn executor_rejects_a_factory_schema_different_from_the_negotiated_schema() { + let bridge = Bridge::::default(); + let handle = bridge.create(Options::new()).unwrap(); + let reader = bridge + .create_reader(handle, bridge.schema(handle).unwrap()) + .unwrap(); + let plan = bridge.planned(reader).unwrap(); + bridge.close_reader(reader).unwrap(); + bridge.close_source(handle).unwrap(); + + let task = crate::abi::SparkDsTaskContext { + size: std::mem::size_of::() as i64, + stage_id: 1, + stage_attempt_number: 0, + partition_id: 0, + attempt_number: 0, + task_attempt_id: 1, + is_interrupted: None, + }; + // The factory returns one id column, but the driver negotiated an empty schema. + // The expected schema must survive serialization and the removal of driver handles. + let error = + error_message(unsafe { bridge.open_reader(&task, &plan.state, &plan.partitions[0]) }); + assert!(error.contains("schema"), "{error}"); + } + + #[test] + fn rejects_a_different_wire_version() { + let error = decode::(br#"{"version":2,"value":42}"#).unwrap_err(); + assert!(error + .to_string() + .contains("unsupported native state version")); + } +} diff --git a/sql/native/rust/src/export.rs b/sql/native/rust/src/export.rs new file mode 100644 index 0000000000000..7f0bf2b30b452 --- /dev/null +++ b/sql/native/rust/src/export.rs @@ -0,0 +1,217 @@ +// 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. + +/// Exports one batch-read provider through version 1 of Spark's C data source ABI. +/// +/// Call this macro once in a cdylib. The only exported symbol is spark_ds_init; +/// providers implement the safe Rust traits and do not handle C pointers. +#[macro_export] +macro_rules! export_data_source { + ($source:ty) => { + mod __spark_native_datasource_exports { + use super::*; + use std::ffi::{c_char, c_int, c_void}; + use $crate::abi::{self, ArrowArrayStream, ArrowSchema, SparkDsLibrary, SparkDsSink}; + use $crate::bridge::Bridge; + + fn bridge() -> &'static Bridge<$source> { + static BRIDGE: std::sync::OnceLock> = std::sync::OnceLock::new(); + BRIDGE.get_or_init(Bridge::default) + } + + static LIBRARY: SparkDsLibrary = SparkDsLibrary { + size: std::mem::size_of::() as i64, + abi_version: abi::ABI_VERSION, + last_error: Some(abi::last_error), + create: Some(create), + schema: Some(schema), + create_reader: Some(create_reader), + create_stream_reader: None, + create_writer: None, + close: Some(close), + push_predicates: Some(push_predicates), + push_limit: Some(push_limit), + prune_columns: Some(prune_columns), + partitions: Some(partitions), + serialize_reader: Some(serialize_reader), + close_reader: Some(close_reader), + initial_offset: None, + latest_offset: None, + stream_partitions: None, + serialize_stream_reader: None, + commit_offset: None, + close_stream_reader: None, + read: Some(read), + serialize_writer: None, + commit: None, + abort: None, + close_writer: None, + create_data_writer: None, + write: None, + commit_data_writer: None, + abort_data_writer: None, + close_data_writer: None, + }; + + #[export_name = "spark_ds_init"] + /// # Safety + /// The arguments must follow the spark_ds_init contract of spark_datasource.h. + pub unsafe extern "C" fn initialize( + host: *const abi::SparkDsHost, + library: *mut *const SparkDsLibrary, + ) -> c_int { + unsafe { abi::initialize(host, library, &LIBRARY) } + } + + unsafe extern "C" fn create( + info: *const abi::SparkDsCreateInfo, + capabilities: *mut u64, + source: *mut *mut c_void, + ) -> c_int { + abi::call(|| { + abi::require_output(capabilities)?; + abi::require_output(source)?; + let options = unsafe { abi::options(info)? }; + let handle = bridge().create(options)?; + unsafe { + capabilities.write(abi::BATCH_READ); + source.write(handle as usize as *mut c_void); + } + Ok(()) + }) + } + + unsafe extern "C" fn schema(source: *mut c_void, output: *mut ArrowSchema) -> c_int { + abi::call(|| { + let schema = bridge().schema(source as usize as i64)?; + unsafe { abi::export_schema(&schema, output) } + }) + } + + unsafe extern "C" fn create_reader( + source: *mut c_void, + schema: *mut ArrowSchema, + reader: *mut *mut c_void, + ) -> c_int { + abi::call(|| { + abi::require_output(reader)?; + let schema = unsafe { abi::import_schema(schema)? }; + let handle = bridge().create_reader(source as usize as i64, schema)?; + unsafe { reader.write(handle as usize as *mut c_void) }; + Ok(()) + }) + } + + unsafe extern "C" fn close(source: *mut c_void) { + let _ = abi::call(|| bridge().close_source(source as usize as i64)); + } + + unsafe extern "C" fn push_predicates( + reader: *mut c_void, + predicates: *const *const c_char, + count: i64, + uses: *mut i32, + ) -> c_int { + abi::call(|| { + if count > 0 { + abi::require_output(uses)?; + } + let values = unsafe { abi::strings(predicates, count)? }; + let predicates = $crate::bridge::parse_predicates(&values)?; + let support = bridge().push_predicates(reader as usize as i64, &predicates)?; + for (index, support) in support.into_iter().enumerate() { + unsafe { uses.add(index).write(abi::predicate_use(support)) }; + } + Ok(()) + }) + } + + unsafe extern "C" fn push_limit( + reader: *mut c_void, + limit: i32, + applied: *mut bool, + ) -> c_int { + abi::call(|| { + abi::require_output(applied)?; + let result = bridge().push_limit(reader as usize as i64, limit.try_into()?)?; + unsafe { applied.write(result) }; + Ok(()) + }) + } + + unsafe extern "C" fn prune_columns( + reader: *mut c_void, + columns: *const *const c_char, + count: i64, + pruned: *mut bool, + ) -> c_int { + abi::call(|| { + abi::require_output(pruned)?; + let columns = unsafe { abi::strings(columns, count)? }; + let result = bridge().prune_columns(reader as usize as i64, &columns)?; + unsafe { pruned.write(result) }; + Ok(()) + }) + } + + unsafe extern "C" fn partitions( + reader: *mut c_void, + sink: SparkDsSink, + data: *mut c_void, + ) -> c_int { + abi::call(|| { + let plan = bridge().planned(reader as usize as i64)?; + for partition in &plan.partitions { + unsafe { abi::emit(sink, data, partition)? }; + } + Ok(()) + }) + } + + unsafe extern "C" fn serialize_reader( + reader: *mut c_void, + sink: SparkDsSink, + data: *mut c_void, + ) -> c_int { + abi::call(|| { + let plan = bridge().planned(reader as usize as i64)?; + unsafe { abi::emit(sink, data, &plan.state) } + }) + } + + unsafe extern "C" fn close_reader(reader: *mut c_void) { + let _ = abi::call(|| bridge().close_reader(reader as usize as i64)); + } + + unsafe extern "C" fn read( + task: *const abi::SparkDsTaskContext, + state: *const u8, + state_length: i64, + partition: *const u8, + partition_length: i64, + output: *mut ArrowArrayStream, + ) -> c_int { + abi::call(|| { + abi::require_output(output)?; + let state = unsafe { abi::slice(state, state_length)? }; + let partition = unsafe { abi::slice(partition, partition_length)? }; + let stream = unsafe { bridge().open_reader(task, state, partition)? }; + unsafe { output.write(stream) }; + Ok(()) + }) + } + } + }; +} diff --git a/sql/native/rust/src/lib.rs b/sql/native/rust/src/lib.rs new file mode 100644 index 0000000000000..2fa4343c96b08 --- /dev/null +++ b/sql/native/rust/src/lib.rs @@ -0,0 +1,149 @@ +// 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 batch-read SDK for the C ABI proposed in Spark PR #59424. +//! +//! Implement [`BatchSource`], [`ScanBuilder`] and [`BatchReaderFactory`], then +//! invoke [`export_data_source!`] once in a `cdylib` crate. The SDK owns handles, +//! plan serialization, C ABI adaptation, and Arrow C Stream export. Rust traits +//! stay inside the plugin; they are not an ABI between dynamic libraries. + +use std::collections::BTreeMap; + +use arrow_array::RecordBatchReader; +use arrow_schema::SchemaRef; +use serde::de::DeserializeOwned; +use serde::Serialize; + +#[doc(hidden)] +pub mod abi; +#[doc(hidden)] +pub mod bridge; +mod export; +mod predicate; +mod task; + +pub use predicate::Predicate; +pub use task::TaskContext; + +pub type Error = Box; +pub type Result = std::result::Result; +pub type Options = BTreeMap; +pub type BatchReader<'task> = Box; + +/// How completely a connector evaluates a pushed predicate. +/// Spark currently rechecks all predicates, including exact ones. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FilterSupport { + /// Implements Spark SQL semantics and returns only matching rows. + Exact, + /// May skip data that cannot match, but may also return nonmatching rows. + PruningOnly, + /// Does not evaluate this predicate. + Unsupported, +} + +/// How a connector limits the rows produced by the whole scan. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LimitSupport { + Unsupported, + /// Caps each partition; the ABI boolean is false because it promises a + /// bound on the whole scan. Spark retains its final limit in either case. + PerPartition, + /// Caps the whole scan; the ABI boolean is true. + 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. +pub struct ScanPlan { + pub factory: F, + pub partitions: Vec, +} + +/// Driver-side source configuration and schema discovery. +/// +/// Registering this trait declares only the batch-read capability. Writes and +/// streaming remain absent from the native function table. +pub trait BatchSource: Send + Sized + 'static { + type Scan: ScanBuilder; + + /// Spark passes option keys in lower case. + fn new(options: Options) -> Result; + + fn schema(&self) -> Result; + + /// The schema may be supplied by the user; reject unsupported schemas. + /// The returned builder must remain valid after this source is dropped. + fn scan(&self, schema: SchemaRef) -> Result; +} + +/// Serializable configuration for creating independent executor readers. +/// +/// This type carries the configuration or snapshot descriptor; connections and +/// other live resources are opened on the executor. The returned reader owns +/// those resources and cannot borrow this short-lived factory. +pub trait BatchReaderFactory: Serialize + DeserializeOwned + Send + 'static { + type Partition: Serialize + DeserializeOwned + Send + 'static; + + /// Return lazy batches with a fixed schema and bounded memory. The reader + /// may retain `task` until the stream is released. Check cancellation during + /// long operations and use timeouts for blocking I/O. + /// + /// The SDK checks Spark-compatible Arrow types, the negotiated schema, and + /// zero array offsets recursively. It reports incompatible layouts instead + /// of copying buffers. Exported batches must own their buffers independently + /// of the reader, because Spark may release the stream before its batches. + fn open<'task>( + &self, + partition: Self::Partition, + task: TaskContext<'task>, + ) -> Result>; +} + +/// Mutable driver-side negotiation, followed by immutable planning. +/// +/// The SDK consumes this builder on the first native `partitions` or +/// `serialize_reader` call and caches the result, including a planning failure. +/// This is later than Java DSv2 `ScanBuilder.build()`. Only serialized factory +/// and partition descriptors travel to executors; driver handles stay local. +pub trait ScanBuilder: Send + Sized + 'static { + type Factory: BatchReaderFactory; + + /// Return one result per input predicate in order. Unknown functions and + /// types must not be reported as exact. Spark currently rechecks all filters. + fn push_filters(&mut self, filters: &[Predicate]) -> Result> { + Ok(vec![FilterSupport::Unsupported; filters.len()]) + } + + /// Count rows after filtering. Spark currently pushes limits only for + /// scans with no filters; the SDK also declines a limit when any predicate + /// it received remains residual. + fn push_limit(&mut self, _limit: usize) -> Result { + Ok(LimitSupport::Unsupported) + } + + /// On true, output exactly these top-level columns in this order. On false, + /// keep the original schema. Empty projection must preserve row counts. + /// Retain internally needed filter columns without returning them. + fn prune_columns(&mut self, _columns: &[String]) -> Result { + Ok(false) + } + + fn plan(self) -> Result>; +} diff --git a/sql/native/rust/src/predicate.rs b/sql/native/rust/src/predicate.rs new file mode 100644 index 0000000000000..f0ff0b8391712 --- /dev/null +++ b/sql/native/rust/src/predicate.rs @@ -0,0 +1,83 @@ +// 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. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +/// The expression tree documented by `NativeBridge.pushPredicates`. +/// +/// Literal values keep their JSON representation so a connector can reject +/// unsupported types without lossy numeric conversions. In particular, +/// decimal values are strings and timestamps are integer microseconds. +/// Function names are extensible: unknown names must remain residual. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "lowercase")] +pub enum Predicate { + Column { + name: Vec, + }, + Literal { + #[serde(rename = "dataType")] + data_type: String, + value: Value, + }, + Function { + name: String, + children: Vec, + }, + #[serde(other)] + Unknown, +} + +#[cfg(test)] +mod tests { + use super::Predicate; + + #[test] + fn preserves_int64_literals_and_nested_column_paths() { + let json = r#"{"type":"function","name":">=","children":[ + {"type":"column","name":["outer","id"]}, + {"type":"literal","dataType":"bigint","value":9223372036854775807} + ]}"#; + let predicate: Predicate = serde_json::from_str(json).unwrap(); + let Predicate::Function { children, .. } = predicate else { + panic!("expected a function"); + }; + assert_eq!( + children[0], + Predicate::Column { + name: vec!["outer".into(), "id".into()] + } + ); + let Predicate::Literal { value, .. } = &children[1] else { + panic!("expected a literal"); + }; + assert_eq!(value.as_i64(), Some(i64::MAX)); + } + + #[test] + fn preserves_unknown_expressions_for_conservative_fallback() { + let json = r#"{"type":"future_expression","value":1}"#; + assert_eq!( + serde_json::from_str::(json).unwrap(), + Predicate::Unknown + ); + let json = r#"{"type":"function","name":"FUTURE_FUNCTION","children":[]}"#; + let Predicate::Function { name, .. } = serde_json::from_str(json).unwrap() else { + panic!("expected a function"); + }; + assert_eq!(name, "FUTURE_FUNCTION"); + } +} diff --git a/sql/native/rust/src/stream.rs b/sql/native/rust/src/stream.rs new file mode 100644 index 0000000000000..8a0958dc307e9 --- /dev/null +++ b/sql/native/rust/src/stream.rs @@ -0,0 +1,517 @@ +// 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. + +use std::ffi::{c_char, c_int, c_void, CString}; +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::sync::Arc; + +use arrow_array::ffi::FFI_ArrowArray; +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_array::{Array, StructArray}; +use arrow_data::ArrayData; +use arrow_schema::{ + ffi::FFI_ArrowSchema, DataType, Field, Fields, IntervalUnit, Schema, SchemaRef, TimeUnit, +}; + +use super::panic_message; +use crate::{BatchReader, Result}; + +struct StreamState<'task> { + reader: Option>, + schema: SchemaRef, + last_error: Option, + finished: bool, +} + +impl Drop for StreamState<'_> { + fn drop(&mut self) { + // release has no error return channel, but it must never unwind into C. + if let Some(reader) = self.reader.take() { + if let Err(payload) = catch_unwind(AssertUnwindSafe(|| drop(reader))) { + let _ = panic_message(payload); + } + } + } +} + +/// The reader may borrow the host task context until the consumer releases +/// this stream. Erasure into C private_data is confined to this boundary. +/// +/// # Safety +/// The caller must keep everything borrowed by the reader alive until release, +/// and the consumer must obey the Arrow C Stream ownership protocol. +pub(super) unsafe fn export<'task>( + reader: BatchReader<'task>, + expected_schema: &SchemaRef, +) -> Result { + // Establish guarded destruction before calling the provider's schema(). + let mut state = Box::new(StreamState { + reader: Some(reader), + schema: Arc::new(Schema::empty()), + last_error: None, + finished: false, + }); + state.schema = catch_unwind(AssertUnwindSafe(|| { + state.reader.as_ref().expect("reader is present").schema() + })) + .map_err(panic_message)?; + if canonical_fields(state.schema.fields())? != canonical_fields(expected_schema.fields())? { + return Err("executor reader schema does not match the negotiated schema".into()); + } + Ok(FFI_ArrowArrayStream { + get_schema: Some(get_schema), + get_next: Some(get_next), + get_last_error: Some(get_last_error), + release: Some(release), + private_data: Box::into_raw(state).cast::(), + }) +} + +fn callback( + state: &mut StreamState<'_>, + action: impl FnOnce(&mut StreamState<'_>) -> Result<()>, +) -> c_int { + let result = catch_unwind(AssertUnwindSafe(|| { + action(state).map_err(|e| e.to_string()) + })); + let message = match result { + Ok(Ok(())) => return 0, + Ok(Err(message)) => message, + Err(payload) => panic_message(payload), + }; + state.finished = true; + // Error text is a C string. A provider error may contain embedded NULs. + state.last_error = Some(CString::new(message.replace('\0', "\\0")).expect("NULs removed")); + 5 // EIO; get_last_error supplies the provider's actual error message. +} + +unsafe extern "C" fn get_schema( + stream: *mut FFI_ArrowArrayStream, + output: *mut FFI_ArrowSchema, +) -> c_int { + let state = unsafe { &mut *(*stream).private_data.cast::>() }; + callback(state, |state| { + let schema = FFI_ArrowSchema::try_from(state.schema.as_ref())?; + unsafe { output.write(schema) }; + Ok(()) + }) +} + +unsafe extern "C" fn get_next( + stream: *mut FFI_ArrowArrayStream, + output: *mut FFI_ArrowArray, +) -> c_int { + let state = unsafe { &mut *(*stream).private_data.cast::>() }; + callback(state, |state| { + if state.finished { + unsafe { output.write(FFI_ArrowArray::empty()) }; + return Ok(()); + } + match state.reader.as_mut().expect("reader is present").next() { + Some(batch) => { + let batch = batch?; + if batch.schema() != state.schema { + return Err("reader returned a batch with a different schema".into()); + } + let data = StructArray::from(batch).to_data(); + validate_layout(&data)?; + let array = FFI_ArrowArray::new(&data); + unsafe { output.write(array) }; + } + None => { + state.finished = true; + // A null release callback denotes EOF, including for zero-column data. + unsafe { output.write(FFI_ArrowArray::empty()) }; + } + } + Ok(()) + }) +} + +unsafe extern "C" fn get_last_error(stream: *mut FFI_ArrowArrayStream) -> *const c_char { + let state = unsafe { &*(*stream).private_data.cast::>() }; + state + .last_error + .as_ref() + .map_or(std::ptr::null(), |error| error.as_ptr()) +} + +unsafe extern "C" fn release(stream: *mut FFI_ArrowArrayStream) { + let state = unsafe { (*stream).private_data.cast::>() }; + // Mark released before destruction. The consumer owns the stream structure. + unsafe { stream.write(FFI_ArrowArrayStream::empty()) }; + if let Err(payload) = catch_unwind(AssertUnwindSafe(|| unsafe { drop(Box::from_raw(state)) })) { + let _ = panic_message(payload); + } +} + +// Compare Spark logical types, ignoring Arrow nullability/metadata and the +// interchangeable physical encodings of strings and binary data. Spark also +// distinguishes only the presence of a time zone, not its spelling. +pub(super) fn validate_schema(schema: &SchemaRef) -> Result<()> { + canonical_fields(schema.fields())?; + Ok(()) +} + +fn canonical_fields(fields: &Fields) -> Result { + fields + .iter() + .map(|field| { + Ok(Arc::new(Field::new( + field.name(), + canonical_type(field.data_type())?, + false, + ))) + }) + .collect() +} + +fn canonical_type(data_type: &DataType) -> Result { + use DataType::*; + Ok(match data_type { + Null | Boolean | Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Date32 => { + data_type.clone() + } + Utf8 | LargeUtf8 | Utf8View => Utf8, + Binary | LargeBinary | BinaryView => Binary, + Decimal128(precision, _) if *precision > 0 && *precision <= 38 => data_type.clone(), + Timestamp(TimeUnit::Microsecond, tz) => { + Timestamp(TimeUnit::Microsecond, tz.as_ref().map(|_| "UTC".into())) + } + Duration(TimeUnit::Microsecond) + | Interval(IntervalUnit::YearMonth | IntervalUnit::MonthDayNano) => data_type.clone(), + Struct(fields) => Struct(canonical_fields(fields)?), + List(field) => List(Arc::new(Field::new( + "item", + canonical_type(field.data_type())?, + false, + ))), + Map(entries, _) => { + let Struct(fields) = entries.data_type() else { + return Err("Arrow map entries must be a struct".into()); + }; + if fields.len() != 2 { + return Err("Arrow map entries must have a key and a value".into()); + } + let entries = Struct( + vec![ + Field::new("key", canonical_type(fields[0].data_type())?, false), + Field::new("value", canonical_type(fields[1].data_type())?, false), + ] + .into(), + ); + Map(Arc::new(Field::new("entries", entries, false)), false) + } + _ => return Err(format!("unsupported Spark Arrow type: {data_type}").into()), + }) +} + +fn validate_layout(data: &ArrayData) -> Result<()> { + if data.offset() != 0 { + return Err(format!( + "Spark requires zero Arrow array offsets: {} has offset {}; normalize sliced arrays", + data.data_type(), + data.offset(), + ) + .into()); + } + for child in data.child_data() { + validate_layout(child)?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use arrow_array::ffi_stream::ArrowArrayStreamReader; + use arrow_array::{BooleanArray, Int32Array, Int64Array, RecordBatch, RecordBatchReader}; + use arrow_schema::{ArrowError, DataType, Field}; + + use super::*; + + struct Reader { + schema: SchemaRef, + batches: std::vec::IntoIter>, + panic_next: bool, + panic_schema: bool, + panic_drop: bool, + drops: Arc, + } + + impl Iterator for Reader { + type Item = std::result::Result; + + fn next(&mut self) -> Option { + assert!(!self.panic_next, "iterator failed"); + self.batches.next() + } + } + + impl RecordBatchReader for Reader { + fn schema(&self) -> SchemaRef { + assert!(!self.panic_schema, "schema failed"); + self.schema.clone() + } + } + + impl Drop for Reader { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + assert!(!self.panic_drop, "drop failed"); + } + } + + fn export(reader: BatchReader<'static>) -> Result { + let expected = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + // Test readers own all their state and cannot borrow a host context. + unsafe { super::export(reader, &expected) } + } + + fn reader() -> Reader { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let values = Arc::new(Int64Array::from(vec![3, 5])); + let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap(); + Reader { + schema, + batches: vec![Ok(batch)].into_iter(), + panic_next: false, + panic_schema: false, + panic_drop: false, + drops: Arc::default(), + } + } + + #[test] + fn imported_batch_survives_early_stream_release() { + let reader = reader(); + let drops = reader.drops.clone(); + let mut stream = + ArrowArrayStreamReader::try_new(export(Box::new(reader)).unwrap()).unwrap(); + let batch = stream.next().unwrap().unwrap(); + drop(stream); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert_eq!(batch.num_rows(), 2); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.values(), &[3, 5]); + } + + #[test] + fn empty_batch_is_not_eof() { + let mut reader = reader(); + let empty = RecordBatch::new_empty(reader.schema.clone()); + reader.batches = vec![Ok(empty), reader.batches.next().unwrap()].into_iter(); + let mut stream = + ArrowArrayStreamReader::try_new(export(Box::new(reader)).unwrap()).unwrap(); + assert_eq!(stream.next().unwrap().unwrap().num_rows(), 0); + assert_eq!(stream.next().unwrap().unwrap().num_rows(), 2); + assert!(stream.next().is_none()); + } + + #[test] + fn iterator_panics_and_nul_errors_become_stream_errors() { + let mut panicking = reader(); + panicking.panic_next = true; + let mut stream = + ArrowArrayStreamReader::try_new(export(Box::new(panicking)).unwrap()).unwrap(); + assert!(stream + .next() + .unwrap() + .unwrap_err() + .to_string() + .contains("iterator failed")); + + let mut failing = reader(); + failing.batches = vec![Err(ArrowError::ComputeError("bad\0value".into()))].into_iter(); + let mut stream = + ArrowArrayStreamReader::try_new(export(Box::new(failing)).unwrap()).unwrap(); + assert!(stream + .next() + .unwrap() + .unwrap_err() + .to_string() + .contains("bad\\0value")); + } + + #[test] + fn schema_and_drop_panics_do_not_escape_the_stream_boundary() { + let mut failing = reader(); + let drops = failing.drops.clone(); + failing.panic_schema = true; + failing.panic_drop = true; + assert!(export(Box::new(failing)) + .unwrap_err() + .to_string() + .contains("schema failed")); + assert_eq!(drops.load(Ordering::SeqCst), 1); + + let mut failing = reader(); + let drops = failing.drops.clone(); + failing.panic_drop = true; + drop(export(Box::new(failing)).unwrap()); + assert_eq!(drops.load(Ordering::SeqCst), 1); + } + + #[test] + fn negotiated_schema_mismatch_drops_the_reader() { + let reader = reader(); + let drops = reader.drops.clone(); + let expected = Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)])); + // The reader owns its state; no task context needs to outlive this call. + let result = unsafe { super::export(Box::new(reader), &expected) }; + assert!(result + .unwrap_err() + .to_string() + .contains("does not match the negotiated schema")); + assert_eq!(drops.load(Ordering::SeqCst), 1); + } + + #[test] + fn logical_schema_comparison_accepts_equivalent_arrow_representations() { + let actual = Schema::new(vec![ + Field::new("text", DataType::LargeUtf8, true), + Field::new("bytes", DataType::BinaryView, true), + Field::new( + "time", + DataType::Timestamp(TimeUnit::Microsecond, Some("Europe/Paris".into())), + true, + ), + Field::new( + "nested", + DataType::Struct(vec![Field::new("text", DataType::Utf8View, true)].into()), + true, + ), + ]); + let expected = Schema::new(vec![ + Field::new("text", DataType::Utf8, false), + Field::new("bytes", DataType::Binary, false), + Field::new( + "time", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + false, + ), + Field::new( + "nested", + DataType::Struct(vec![Field::new("text", DataType::Utf8, false)].into()), + false, + ), + ]); + assert_eq!( + canonical_fields(actual.fields()).unwrap(), + canonical_fields(expected.fields()).unwrap() + ); + // A timestamp without a time zone remains a distinct Spark logical type. + assert_ne!( + canonical_type(&DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap(), + canonical_type(&DataType::Timestamp( + TimeUnit::Microsecond, + Some("UTC".into()) + )) + .unwrap() + ); + } + + #[test] + fn unsupported_types_are_rejected_inside_nested_fields() { + let nested_types = [ + DataType::List(Arc::new(Field::new("item", DataType::UInt64, false))), + DataType::Struct( + vec![Field::new( + "dictionary", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + )] + .into(), + ), + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Decimal256(40, 2), true), + ] + .into(), + ), + false, + )), + false, + ), + ]; + for data_type in nested_types { + let schema = Arc::new(Schema::new(vec![Field::new("nested", data_type, true)])); + assert!(validate_schema(&schema) + .unwrap_err() + .to_string() + .contains("unsupported Spark Arrow type")); + } + } + + #[test] + fn a_valid_sliced_boolean_inside_a_struct_is_rejected_before_export() { + let values = BooleanArray::from(vec![true, false, true]).slice(1, 2); + assert_eq!(values.to_data().offset(), 1); + let fields: Fields = vec![Field::new("flag", DataType::Boolean, false)].into(); + let nested = StructArray::new(fields.clone(), vec![Arc::new(values)], None); + let schema = Arc::new(Schema::new(vec![Field::new( + "nested", + DataType::Struct(fields), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(nested)]).unwrap(); + let mut reader = reader(); + reader.schema = schema.clone(); + reader.batches = vec![Ok(batch)].into_iter(); + let drops = reader.drops.clone(); + // All arrays are valid, owned Arrow values, without a borrowed task context. + let exported = unsafe { super::export(Box::new(reader), &schema) }.unwrap(); + let mut stream = ArrowArrayStreamReader::try_new(exported).unwrap(); + let error = stream.next().unwrap().unwrap_err().to_string(); + assert!(error.contains("zero Arrow array offsets"), "{error}"); + assert!(error.contains("offset 1"), "{error}"); + assert!(stream.next().is_none()); + drop(stream); + assert_eq!(drops.load(Ordering::SeqCst), 1); + } + + #[test] + fn a_batch_cannot_change_the_readers_advertised_schema() { + let mut reader = reader(); + let first = reader.batches.next().unwrap(); + let changed_schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let changed = + RecordBatch::try_new(changed_schema, vec![Arc::new(Int32Array::from(vec![7, 9]))]) + .unwrap(); + reader.batches = vec![first, Ok(changed)].into_iter(); + let drops = reader.drops.clone(); + let mut stream = + ArrowArrayStreamReader::try_new(export(Box::new(reader)).unwrap()).unwrap(); + assert_eq!(stream.next().unwrap().unwrap().num_rows(), 2); + assert!(stream + .next() + .unwrap() + .unwrap_err() + .to_string() + .contains("different schema")); + assert!(stream.next().is_none()); + drop(stream); + assert_eq!(drops.load(Ordering::SeqCst), 1); + } +} diff --git a/sql/native/rust/src/task.rs b/sql/native/rust/src/task.rs new file mode 100644 index 0000000000000..2252f82c8e32a --- /dev/null +++ b/sql/native/rust/src/task.rs @@ -0,0 +1,187 @@ +// 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. + +use std::marker::PhantomData; +use std::mem::offset_of; + +use crate::abi::{self, Interrupted, SparkDsTaskContext}; +use crate::Result; + +/// Task identifiers and cooperative cancellation, borrowed until its Arrow stream is released. +/// +/// A reader may retain this value and check it during long operations. It cannot retain it +/// beyond the lifetime of its returned reader. Blocking I/O still needs its own timeouts. +/// +/// ```compile_fail +/// use spark_native_source::TaskContext; +/// +/// fn retain_forever(task: TaskContext<'_>) -> TaskContext<'static> { +/// task +/// } +/// ``` +#[derive(Clone, Copy)] +pub struct TaskContext<'task> { + pointer: *const SparkDsTaskContext, + stage_id: Option, + stage_attempt_number: Option, + partition_id: Option, + attempt_number: Option, + task_attempt_id: Option, + interrupted: Option, + lifetime: PhantomData<&'task SparkDsTaskContext>, +} + +// The C ABI permits task callbacks from any thread. The lifetime prevents use after the +// host releases its context; the remaining fields are copied values or immutable callbacks. +unsafe impl Send for TaskContext<'_> {} +unsafe impl Sync for TaskContext<'_> {} + +impl TaskContext<'static> { + /// A context without Spark task metadata, for testing a reader directly. + pub fn detached() -> Self { + Self { + pointer: std::ptr::null(), + stage_id: None, + stage_attempt_number: None, + partition_id: None, + attempt_number: None, + task_attempt_id: None, + interrupted: None, + lifetime: PhantomData, + } + } +} + +impl<'task> TaskContext<'task> { + /// Wraps a host context without reading fields beyond its declared size. + /// + /// # Safety + /// The host must keep the context and its callbacks valid for `'task`, including calls from + /// other threads. In the Spark ABI this lasts until the exported Arrow stream is released. + #[doc(hidden)] + pub unsafe fn from_raw(pointer: *const SparkDsTaskContext) -> Result { + let size = unsafe { abi::struct_size(pointer)? }; + let bytes = pointer.cast(); + Ok(Self { + pointer, + stage_id: unsafe { abi::field(bytes, size, offset_of!(SparkDsTaskContext, stage_id)) }, + stage_attempt_number: unsafe { + abi::field( + bytes, + size, + offset_of!(SparkDsTaskContext, stage_attempt_number), + ) + }, + partition_id: unsafe { + abi::field(bytes, size, offset_of!(SparkDsTaskContext, partition_id)) + }, + attempt_number: unsafe { + abi::field(bytes, size, offset_of!(SparkDsTaskContext, attempt_number)) + }, + task_attempt_id: unsafe { + abi::field(bytes, size, offset_of!(SparkDsTaskContext, task_attempt_id)) + }, + interrupted: unsafe { + abi::field(bytes, size, offset_of!(SparkDsTaskContext, is_interrupted)) + } + .flatten(), + lifetime: PhantomData, + }) + } + + pub fn stage_id(&self) -> Option { + self.stage_id + } + + pub fn stage_attempt_number(&self) -> Option { + self.stage_attempt_number + } + + pub fn partition_id(&self) -> Option { + self.partition_id + } + + pub fn attempt_number(&self) -> Option { + self.attempt_number + } + + pub fn task_attempt_id(&self) -> Option { + self.task_attempt_id + } + + /// Whether Spark killed this task. Older contexts without the callback return false. + pub fn is_interrupted(&self) -> bool { + self.interrupted + .is_some_and(|callback| unsafe { callback(self.pointer) }) + } +} + +#[cfg(test)] +mod tests { + use std::mem::size_of; + use std::sync::atomic::{AtomicBool, Ordering}; + + use super::*; + + #[repr(C)] + struct HostTask { + task: SparkDsTaskContext, + interrupted: AtomicBool, + } + + unsafe extern "C" fn interrupted(task: *const SparkDsTaskContext) -> bool { + unsafe { &*task.cast::() } + .interrupted + .load(Ordering::SeqCst) + } + + #[test] + fn reads_ids_and_live_cancellation_from_the_host_context() { + let host = HostTask { + task: SparkDsTaskContext { + size: size_of::() as i64, + stage_id: 3, + stage_attempt_number: 4, + partition_id: 5, + attempt_number: 6, + task_attempt_id: 7, + is_interrupted: Some(interrupted), + }, + interrupted: AtomicBool::new(false), + }; + let context = unsafe { TaskContext::from_raw(&host.task) }.unwrap(); + assert_eq!(context.stage_id(), Some(3)); + assert_eq!(context.stage_attempt_number(), Some(4)); + assert_eq!(context.partition_id(), Some(5)); + assert_eq!(context.attempt_number(), Some(6)); + assert_eq!(context.task_attempt_id(), Some(7)); + assert!(!context.is_interrupted()); + host.interrupted.store(true, Ordering::SeqCst); + std::thread::scope(|scope| { + scope.spawn(|| assert!(context.is_interrupted())); + }); + } + + #[test] + fn only_reads_fields_present_in_the_declared_prefix() { + let prefix = size_of::() as i64; + let pointer = (&prefix as *const i64).cast::(); + let context = unsafe { TaskContext::from_raw(pointer) }.unwrap(); + assert_eq!(context.stage_id(), None); + assert_eq!(context.task_attempt_id(), None); + assert!(!context.is_interrupted()); + assert!(!TaskContext::detached().is_interrupted()); + } +} diff --git a/sql/native/rust/tests/fixtures/spark_datasource.h b/sql/native/rust/tests/fixtures/spark_datasource.h new file mode 100644 index 0000000000000..3ff46d2871560 --- /dev/null +++ b/sql/native/rust/tests/fixtures/spark_datasource.h @@ -0,0 +1,683 @@ +/* + * 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. + */ + +/* + * spark_datasource.h: the binary interface between Spark and data sources implemented in native + * code, such as Rust or C++. Since Spark 4.4.0. + * + * A native data source is a shared library that exports a single C function, spark_ds_init, + * which returns the table of the functions that implement the data source, a SparkDsLibrary. + * Spark calls them with the Foreign Function and Memory API of Java, so native data sources + * require Java 22 or later. Each library has its own table, so any number of libraries can be + * loaded side by side. Spark adapts a native data source to Data Source V2, with batch and + * micro-batch streaming reads, batch and streaming writes, and predicate, column and limit + * pushdown. + * + * This header is self-contained: it only depends on the C standard library, and defines the + * structs of the Arrow C data interface unless they are already defined, as by Arrow C++. A library + * can copy it into its source tree, or use the copy in the include directory of a Spark + * distribution. A library built against it works with all the releases of Spark that support its + * SPARK_DS_ABI_VERSION. + * + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of my_source is + * spark_datasource_my_source, such as libspark_datasource_my_source.so in /usr/local/lib on Linux. + * Spark looks for it in the directories of java.library.path, which include LD_LIBRARY_PATH on + * Linux, and in the lib directory of each installation prefix in PATH. A library can also be + * distributed in a native data source package: a zip file with the extension .sparkpkg that + * contains a manifest named spark-native-datasource.json and the library built for one or more + * platforms, added to a session with spark.addArtifact or under the paths in + * spark.sql.dataSource.native.paths. See the "Native Data Sources" page of the Spark SQL guide for + * details and examples. + * + * Lifecycle + * --------- + * Spark loads a library once in each JVM, and calls spark_ds_init before any other function. It + * creates a data source with create on the driver for each schema inference, scan and write, and + * closes it right after. create also returns the capabilities of the data source: Spark reports + * the operations that a data source does not declare as unsupported, without calling the library. + * + * - Batch reads. On the driver, Spark creates a reader with create_reader, pushes operations down + * by calling, each at most once and in this order, push_predicates, push_limit and + * prune_columns, plans the partitions with partitions, serializes the state of the reader with + * serialize_reader, and closes it. Then it calls read on the executors, in one task for each + * partition. If Spark never plans the partitions, for example to only explain the plan of the + * query, it closes the reader once the reader is no longer referenced. + * - Streaming reads. On the driver, Spark creates a stream reader with create_stream_reader and + * keeps it until the streaming query stops. For each micro-batch, it asks for the latest_offset, + * plans the partitions with stream_partitions and serializes the state with + * serialize_stream_reader, then calls read on the executors. Offsets are defined by the library, + * as compact JSON on a single line: Spark stores each offset as a line of the offset log of the + * checkpoint, so it rejects an offset that is not valid JSON or contains a line break. When a + * streaming query restarts, from its checkpoint, Spark creates a new stream reader and may call + * stream_partitions again for the last range that it planned before, so the partitions of a + * range must read the same data each time. + * - Writes. On the driver, right before it starts the tasks of the write (or of a micro-batch), + * Spark creates a writer with create_writer, and serializes its state with serialize_writer. + * With adaptive query execution, Spark first runs the stages of the data to write that end with + * a shuffle, and creates no writer if they fail. The rest of the data, such as its last stage, + * is computed in the tasks of the write, so a writer may still be aborted because computing the + * data failed: do irreversible work, such as removing the existing data for an overwrite, in + * commit. On the executors, each task creates a data writer with create_data_writer, calls write + * for each batch of at most spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends + * once it holds spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls commit_data_writer, + * or abort_data_writer if the task or the commit fails, before it closes the data writer. Then + * the driver calls commit with the messages of the tasks if they all succeeded, or abort + * otherwise, including when commit fails, and closes the writer. Several attempts of a task may + * write the same partition, after a failure or with speculative execution, each with a data + * writer of its own. Spark lets the attempts of a partition commit one at a time, and passes to + * commit the message of one committed attempt for each partition, so the write consists of the + * data of these attempts only. It aborts the data writers of the other attempts, except on an + * executor that is lost, where a data writer is neither committed nor aborted, and may even have + * committed before its message reached the driver. A streaming write commits or aborts each + * micro-batch (epoch). Spark may commit the same epoch more than once, for example when a + * streaming query restarts after a failure, so committing an epoch that is already committed + * must have no effect, for exactly-once output. The ID of the query, which create_writer gets for + * a streaming write, stays the same when the query restarts from its checkpoint, so the ID of the + * query and the epoch ID identify the output of a micro-batch, also when other queries write to + * the same destination. A micro-batch streaming query creates a new stream writer for each + * micro-batch, and Spark closes the previous stream writer of the query when the query creates a + * new one. Spark is not told when a streaming query stops, so it closes the last stream writer of + * the query once the writer is no longer referenced, or when the query is restarted. + * + * Conventions + * ----------- + * - Errors. The functions that return an int return 0 on success, or any other value on failure. + * They then set the message that last_error returns, and Spark reports it as a + * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail, and they may run after the query, + * from another thread: do the work that can fail a query, such as flushing data, in commit. A + * C++ exception or a Rust panic must not cross the interface: catch it and return an error + * instead. + * - Handles. create, create_reader, create_stream_reader, create_writer and create_data_writer + * return an opaque, non-NULL handle in their last argument, typically a pointer: Spark reports a + * NULL handle as an error. Spark releases each handle exactly once, with the close function of + * its kind: close, close_reader, close_stream_reader, close_writer or close_data_writer, and no + * other function releases a handle. A handle is not always released right after its last use: a + * batch reader whose partitions are never planned and the last stream writer of a query are + * released once they are no longer referenced (see Lifecycle above). A reader or writer handle + * must not depend on the data source handle it was created from, because Spark closes the data + * source handle right after creating it. + * - Memory. What Spark passes to a function, such as strings, byte arrays, SparkDsCreateInfo and + * SparkDsWriteInfo, is only valid during the call, except for SparkDsHost and + * SparkDsTaskContext, which document how long they live. Strings are UTF-8 and NUL-terminated. + * The JSON of the predicates is ASCII: it escapes the other characters, which any JSON parser + * decodes. The library returns strings and byte arrays by passing them to a SparkDsSink, which + * copies them. + * - State. Partitions, the state of readers and writers, and commit messages are byte arrays in a + * format defined by the library. Spark sends them to the executors and back. + * - Arrow data. Schemas and data are exchanged with the Arrow C data interface + * (https://arrow.apache.org/docs/format/CDataInterface.html), as structs allocated by Spark. + * For an output struct, the library initializes it and Spark takes ownership. For an input + * struct, the library may move it to take ownership; otherwise, Spark releases it after the + * call returns, so the library must not use it afterwards. The arrays that the library exports + * must have an offset of 0, at every level: Arrow Java ignores the offsets of the arrays it + * imports, so Spark rejects an array with another offset, such as a zero-copy slice. Copy a + * sliced array before exporting it. Spark also rejects an array with fewer values than its + * parent refers to, such as a column shorter than its batch. The buffers of an array must hold + * all of its values, as the C data interface requires: it does not carry the sizes of the + * buffers, so Spark cannot check them, and reads past the end of a buffer that is too short. + * - Schemas. Spark types map to Arrow types as in spark.sql.execution.arrow: for example, bigint + * is int64, string is utf8, timestamp is timestamp[us] with a time zone and timestamp_ntz is + * timestamp[us] without one. Data read from a library may also use large_utf8 and utf8_view for + * strings, large_binary and binary_view for binary data, and any names for the fields of the + * entries of a map. Spark rejects the other Arrow types, such as unsigned integers, decimal256, + * large lists, fixed-size lists and list views, and dictionary encoding. + * - Options. Option keys are passed in lower case, because Spark matches them + * case-insensitively. + * - Optional functions. A library only sets the functions of the operations it supports in its + * table, and leaves the others NULL. A data source must not declare a capability whose + * functions are NULL: Spark then fails the operation. The pushdown functions and the serialize + * functions are optional: without them, Spark pushes nothing down, and the reader or writer has + * no state. So are schema, without which the user must specify the schema, commit_offset and + * last_error. + * - Threads. Spark calls the library from multiple threads, and calls on different handles may + * run concurrently. Calls on the same handle never run concurrently. A handle may be released + * from a thread other than the threads that used it, such as a JVM Cleaner thread when Spark + * releases a handle once it is no longer referenced, for example the last stream writer of a + * streaming query. A SparkDsSink must be called on the thread of the call that it was passed + * to, during that call. The functions of SparkDsHost and SparkDsTaskContext can be called from + * any thread. + * + * Compatibility + * ------------- + * The function table and the structs that Spark passes to the library start with their size, so + * that they can grow. Within a SPARK_DS_ABI_VERSION: + * - fields are only appended to the structs, including SparkDsLibrary, and the existing fields + * never change, nor do the semantics of the existing functions; + * - Spark only reads the functions that fit in the size of the table of a library, so a library + * built against an earlier header keeps working, with the functions that it does not know as + * NULL. A library must only read the fields that fit in the size of a struct that Spark passes; + * - new values may appear where this documentation says so, such as the names of pushed + * predicates. Spark only uses a new capability or write mode with a data source that declares + * it. + * An incompatible change uses a new SPARK_DS_ABI_VERSION: spark_ds_init gets the newest version + * that Spark supports, the library returns a table of a version that it implements, and Spark + * rejects a table of a version that it does not support. Spark supports version 1 at least in all + * the releases of Spark 4 from 4.4.0. + */ +#ifndef SPARK_DATASOURCE_H +#define SPARK_DATASOURCE_H + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* The version of this interface. */ +#define SPARK_DS_ABI_VERSION 1 + +/* Exports spark_ds_init from the library, also when the library hides its other symbols. */ +#if defined(_WIN32) +#define SPARK_DS_EXPORT __declspec(dllexport) +#else +#define SPARK_DS_EXPORT __attribute__((visibility("default"))) +#endif + +/* The Arrow C data interface: https://arrow.apache.org/docs/format/CDataInterface.html */ +#ifndef ARROW_C_DATA_INTERFACE +#define ARROW_C_DATA_INTERFACE + +#define ARROW_FLAG_DICTIONARY_ORDERED 1 +#define ARROW_FLAG_NULLABLE 2 +#define ARROW_FLAG_MAP_KEYS_SORTED 4 + +struct ArrowSchema { + const char* format; + const char* name; + const char* metadata; + int64_t flags; + int64_t n_children; + struct ArrowSchema** children; + struct ArrowSchema* dictionary; + void (*release)(struct ArrowSchema*); + void* private_data; +}; + +struct ArrowArray { + int64_t length; + int64_t null_count; + int64_t offset; + int64_t n_buffers; + int64_t n_children; + const void** buffers; + struct ArrowArray** children; + struct ArrowArray* dictionary; + void (*release)(struct ArrowArray*); + void* private_data; +}; + +#endif /* ARROW_C_DATA_INTERFACE */ + +#ifndef ARROW_C_STREAM_INTERFACE +#define ARROW_C_STREAM_INTERFACE + +struct ArrowArrayStream { + int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out); + int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out); + const char* (*get_last_error)(struct ArrowArrayStream*); + void (*release)(struct ArrowArrayStream*); + void* private_data; +}; + +#endif /* ARROW_C_STREAM_INTERFACE */ + +/* + * Receives a string or a byte array that the library returns, and copies it. `sink` is the + * argument passed with the sink to the function that returns it. + */ +typedef void (*SparkDsSink)(void* sink, const uint8_t* data, int64_t length); + +/* --------------------------------------------------------------------------------------------- + * What Spark passes to the library + * --------------------------------------------------------------------------------------------- */ + +/* The levels of the messages that a library logs with SparkDsHost.log. */ +#define SPARK_DS_LOG_ERROR 1 +#define SPARK_DS_LOG_WARN 2 +#define SPARK_DS_LOG_INFO 3 +#define SPARK_DS_LOG_DEBUG 4 +#define SPARK_DS_LOG_TRACE 5 + +/* Spark, as spark_ds_init gets it. It stays valid until the JVM exits. */ +typedef struct SparkDsHost { + /* + * The size of this struct in the header that Spark was built with. A later version of Spark may + * append fields, so a library must only read the fields that fit in it. + */ + int64_t size; + /* The newest version of this interface that Spark supports. */ + int32_t abi_version; + /* The version of Spark, such as "4.4.0". */ + const char* spark_version; + /* + * Logs a message, a NUL-terminated UTF-8 string, at one of the SPARK_DS_LOG_* levels: in the log + * of the driver or of the executor, with the path of the library. + */ + void (*log)(int32_t level, const char* message); +} SparkDsHost; + +/* A data source, as create gets it. */ +typedef struct SparkDsCreateInfo { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* + * The name of the data source, as listed in the package manifest, or in lower case for an + * installed library. A library can implement several data sources. + */ + const char* name; + /* + * The num_options options given by the user, such as spark.read.option("key", "value"), with + * the keys in lower case. + */ + const char* const* option_keys; + const char* const* option_values; + int64_t num_options; +} SparkDsCreateInfo; + +/* + * The task that calls read or create_data_writer, on an executor. It stays valid until the stream + * that read returns is released, or until the data writer is closed. + */ +typedef struct SparkDsTaskContext { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* The ID of the stage of the task, and the attempt number of the stage. */ + int32_t stage_id; + int32_t stage_attempt_number; + /* The partition that the task reads or writes, and the attempt number of the task. */ + int32_t partition_id; + int32_t attempt_number; + /* The ID of the task attempt, unique in the Spark application. */ + int64_t task_attempt_id; + /* + * Returns whether the task was killed, for example because its query was canceled. Spark cannot + * interrupt native code, so a read or a write that runs for long, or blocks, should check it + * regularly, and fail once it is true. + */ + bool (*is_interrupted)(const struct SparkDsTaskContext* task); +} SparkDsTaskContext; + +/* The modes of a write. */ +#define SPARK_DS_WRITE_APPEND 0 +#define SPARK_DS_WRITE_TRUNCATE 1 + +/* A write, as create_writer gets it. */ +typedef struct SparkDsWriteInfo { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* + * SPARK_DS_WRITE_APPEND to add the data to the existing data, or SPARK_DS_WRITE_TRUNCATE to + * replace the existing data: with the save mode "overwrite" of a batch write, and in each + * micro-batch with the output mode "complete" of a streaming write. Spark only passes + * SPARK_DS_WRITE_TRUNCATE to a data source with the capability SPARK_DS_TRUNCATE. + */ + int32_t mode; + /* + * The ID of the streaming query of a streaming write, or NULL for a batch write. It stays the + * same when the query restarts from its checkpoint, so with the epoch IDs, it identifies the + * commits of the query: see Writes above. + */ + const char* query_id; +} SparkDsWriteInfo; + +/* --------------------------------------------------------------------------------------------- + * What the library returns to Spark + * --------------------------------------------------------------------------------------------- */ + +/* The capabilities of a data source, which create returns: the operations that it supports. */ +#define SPARK_DS_BATCH_READ 1 /* batch reads, with create_reader */ +#define SPARK_DS_MICRO_BATCH_READ 2 /* micro-batch streaming reads, with create_stream_reader */ +#define SPARK_DS_BATCH_WRITE 4 /* batch writes, with create_writer */ +#define SPARK_DS_STREAMING_WRITE 8 /* streaming writes, with create_writer */ +#define SPARK_DS_TRUNCATE 16 /* writes with the mode SPARK_DS_WRITE_TRUNCATE */ + +/* How a reader uses a predicate pushed with push_predicates. */ +#define SPARK_DS_PREDICATE_UNUSED 0 /* not at all */ +/* only to skip data, so it may return rows for which the predicate is not true */ +#define SPARK_DS_PREDICATE_PARTIAL 1 +/* to only return the rows for which the predicate is true */ +#define SPARK_DS_PREDICATE_COMPLETE 2 + +/* + * The functions of a library, which spark_ds_init returns. A library sets the functions of the + * operations it supports, and leaves the others NULL: see "Optional functions" above. + */ +typedef struct SparkDsLibrary { + /* + * The size of this struct in the header that the library was built with. Spark only reads the + * functions that fit in it. + */ + int64_t size; + /* The version of this interface that the table implements, SPARK_DS_ABI_VERSION. */ + int32_t abi_version; + + /* + * Returns the message of the last error on the calling thread, as a NUL-terminated UTF-8 string + * that stays valid until the next call on that thread. Optional: without it, Spark only reports + * the value that the function returned. + */ + const char* (*last_error)(void); + + /* ------------------------------------------------------------------------------------------- + * Data source (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Creates a data source. + * + * info: the name of the data source, and the options given by the user + * capabilities: set to the capabilities of the data source, SPARK_DS_* flags combined with |, + * which may depend on its name and its options. It is 0 when create is called. + * source: set to the data source handle + */ + int (*create)(const SparkDsCreateInfo* info, uint64_t* capabilities, void** source); + + /* + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. + * + * out: initialized with a struct type whose fields are the columns + */ + int (*schema)(void* source, struct ArrowSchema* out); + + /* + * Creates a reader for a batch scan. + * + * schema: the schema to read: the one the user specified, or the one returned by schema + * reader: set to the reader handle + */ + int (*create_reader)(void* source, struct ArrowSchema* schema, void** reader); + + /* + * Creates a reader for a micro-batch streaming scan. + * + * schema: the schema to read: the one the user specified, or the one returned by schema + * reader: set to the stream reader handle + */ + int (*create_stream_reader)(void* source, struct ArrowSchema* schema, void** reader); + + /* + * Creates a writer for a batch write, or for a micro-batch of a streaming write. + * + * schema: the schema of the data to write + * info: the mode of the write, and the ID of the query of a streaming write + * writer: set to the writer handle + */ + int (*create_writer)(void* source, struct ArrowSchema* schema, const SparkDsWriteInfo* info, + void** writer); + + /* Releases a data source handle. */ + void (*close)(void* source); + + /* ------------------------------------------------------------------------------------------- + * Batch reader (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Pushes predicates down to the reader. The predicates are combined with AND. + * + * Each predicate is a JSON expression tree. An expression is one of: + * - {"type": "column", "name": ["a", "b"]}: a column. A nested field has a name part for each + * level. + * - {"type": "literal", "dataType": "bigint", "value": 1}: a literal. The data type is a Spark + * SQL type name such as int, string or decimal(10,2). The value is null, a JSON boolean or + * number for boolean and numeric types, a string for string, a base64 string for binary, a + * decimal string for decimal, the days since the epoch for date, and the microseconds since + * the epoch for timestamp and timestamp_ntz. + * - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in + * org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=> + * (equal, with null equal to null), <, <=, > and >= with two children, IN with the value to + * look for and then the values to look in, IS_NULL and IS_NOT_NULL with one child, + * STARTS_WITH, ENDS_WITH and CONTAINS with the string and the string to look for, AND and OR + * with two predicates, and NOT with one. The children of the other functions are columns, + * literals or predicates. + * Spark does not pass predicates that it cannot express this way, such as predicates that use + * arithmetic or other functions, nor predicates on columns that contain strings with a + * non-binary collation, such as UTF8_LCASE: strings are compared by their UTF-8 bytes. A later + * version of Spark may pass predicates with other names without changing SPARK_DS_ABI_VERSION, + * so a reader must ignore a predicate that contains a name that it does not know, or a data + * type that it does not support. + * + * Spark evaluates all the predicates again after reading, so a reader may use them only to skip + * data, for example the files whose statistics show that no row matches. But it must never drop + * or skip a row for which the predicates are true under the semantics of Spark SQL, which Spark + * cannot detect. In particular: + * - For floating-point columns, NaN is equal to NaN and greater than any other value, including + * positive infinity, and -0.0 is equal to 0.0. IEEE 754 comparisons treat NaN differently, and + * total orders, such as f64::total_cmp in Rust or the comparison kernels of arrow-rs, also + * order -0.0 before 0.0. For example, d > 1.0 is true for a NaN value of d, and d = 0.0 for + * -0.0. + * - Comparisons with null are null, and a row only matches if the predicates are true. For + * example, NOT (i IN (1, NULL)) matches no row. + * - Strings are compared by their UTF-8 bytes, which order the characters outside the Basic + * Multilingual Plane after all the others, unlike comparisons of UTF-16 strings. + * - To use only a part of a predicate, a reader may take the parts that it does not support as + * true, but not under a NOT, where that would drop rows. + * The file org/apache/spark/sql/datasource/predicate-cases.json in the jar of Spark SQL lists + * predicates as Spark passes them to this function, with the rows of a table for which they are + * true, to test a reader with. + * + * predicates: the count predicates, as JSON + * uses: an array of count values, all SPARK_DS_PREDICATE_UNUSED, to set to how the reader uses + * each predicate, a SPARK_DS_PREDICATE_* value. Spark shows the predicates that the reader + * uses as pushed in the plan of the query, and evaluates all the predicates again after + * reading anyway. It takes other values as SPARK_DS_PREDICATE_PARTIAL. + */ + int (*push_predicates)(void* reader, const char* const* predicates, int64_t count, + int32_t* uses); + + /* + * Pushes a LIMIT down to the reader. Spark only pushes a limit when the query does not filter + * the rows, since it evaluates the predicates after reading. + * + * applied: set to whether the reader returns at most `limit` rows in total. Spark applies the + * limit again after reading either way. + */ + int (*push_limit)(void* reader, int32_t limit, bool* applied); + + /* + * Prunes the columns to read. + * + * columns: the names of the count top-level columns that the query needs, which include the + * columns of the pushed predicates. There may be none, for example for + * SELECT count(*): then the batches only have a number of rows. + * pruned: set to whether read returns exactly these columns, in this order. If not, it returns + * all the columns of the schema to read. + */ + int (*prune_columns)(void* reader, const char* const* columns, int64_t count, bool* pruned); + + /* + * Plans the partitions of the scan. Spark reads each partition in a separate task. + * + * add_partition: called with each partition, serialized in a format defined by the library + */ + int (*partitions)(void* reader, SparkDsSink add_partition, void* sink); + + /* + * Serializes the state of the reader that read needs on the executors, such as the pushed + * predicates. Spark calls it after partitions. + * + * set_state: called once with the state, or not at all if there is none + */ + int (*serialize_reader)(void* reader, SparkDsSink set_state, void* sink); + + /* Releases a batch reader handle. */ + void (*close_reader)(void* reader); + + /* ------------------------------------------------------------------------------------------- + * Stream reader (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Returns the offset to start from when the streaming query has no checkpoint. + * + * set_offset: called once with the offset, as compact JSON on a single line + */ + int (*initial_offset)(void* reader, SparkDsSink set_offset, void* sink); + + /* + * Returns the most recent offset available. + * + * set_offset: called once with the offset, as compact JSON on a single line + */ + int (*latest_offset)(void* reader, SparkDsSink set_offset, void* sink); + + /* + * Plans the partitions that read the data after `start` up to and including `end`, two offsets + * as JSON. + * + * add_partition: called with each partition, serialized in a format defined by the library + */ + int (*stream_partitions)(void* reader, const char* start, const char* end, + SparkDsSink add_partition, void* sink); + + /* + * Serializes the state of the stream reader that read needs on the executors. Spark calls it + * after each call to stream_partitions. + * + * set_state: called once with the state, or not at all if there is none + */ + int (*serialize_stream_reader)(void* reader, SparkDsSink set_state, void* sink); + + /* + * Tells the library that Spark has processed all the data up to and including `end`, an offset + * as JSON, and will not ask for it again. + */ + int (*commit_offset)(void* reader, const char* end); + + /* Releases a stream reader handle, when the streaming query stops. */ + void (*close_stream_reader)(void* reader); + + /* ------------------------------------------------------------------------------------------- + * Partition reader (executor) + * ------------------------------------------------------------------------------------------- */ + + /* + * Reads a partition of a batch or streaming scan. + * + * task: the task that reads the partition, valid until the stream is released + * reader_state: the state returned by serialize_reader or serialize_stream_reader, empty if + * there is none + * partition: a partition returned by partitions or stream_partitions + * out: the stream to initialize. Each array of the stream is a struct array whose fields are the + * columns to read: the columns of the schema to read, or the pruned columns. Spark calls the + * stream from the thread of the task that reads the partition. get_next may block, for + * example on network I/O, but Spark cannot interrupt native code: when the task is killed, + * Spark only stops reading once get_next returns, and then releases the stream. So a + * get_next that blocks forever keeps the task running, unless the task reaper + * (spark.task.reaper.enabled and spark.task.reaper.killTimeout) kills the executor. Use + * timeouts for blocking operations, and check task->is_interrupted. + */ + int (*read)(const SparkDsTaskContext* task, const uint8_t* reader_state, + int64_t reader_state_length, const uint8_t* partition, int64_t partition_length, + struct ArrowArrayStream* out); + + /* ------------------------------------------------------------------------------------------- + * Writer (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Serializes the state of the writer that create_data_writer needs on the executors. Spark + * calls it once, after creating the writer. + * + * set_state: called once with the state, or not at all if there is none + */ + int (*serialize_writer)(void* writer, SparkDsSink set_state, void* sink); + + /* + * Commits a write, when all the tasks succeeded. + * + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * messages, message_lengths: the num_messages messages returned by commit_data_writer of each + * task + */ + int (*commit)(void* writer, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages); + + /* + * Aborts a write, when it failed. + * + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * messages, message_lengths: the num_messages messages returned by commit_data_writer of each + * task. A message is NULL if its task did not commit. + */ + int (*abort)(void* writer, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages); + + /* Releases a writer handle. */ + void (*close_writer)(void* writer); + + /* ------------------------------------------------------------------------------------------- + * Data writer (executor) + * ------------------------------------------------------------------------------------------- */ + + /* + * Creates the writer for one task. + * + * task: the task, with the partition that it writes, valid until the data writer is closed + * writer_state: the state returned by serialize_writer, empty if there is none + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * data_writer: set to the data writer handle + */ + int (*create_data_writer)(const SparkDsTaskContext* task, const uint8_t* writer_state, + int64_t writer_state_length, int64_t epoch_id, void** data_writer); + + /* + * Writes a batch of data. + * + * array: a struct array whose fields are the columns + * schema: the schema of the array + */ + int (*write)(void* data_writer, struct ArrowArray* array, struct ArrowSchema* schema); + + /* + * Commits the data written by the task. If it fails, Spark calls abort_data_writer. If it + * succeeds, Spark does not call abort_data_writer, even if it then fails the task, for example + * because set_message was called more than once. + * + * set_message: called once with the message to pass to commit on the driver, or not at all for + * an empty message + */ + int (*commit_data_writer)(void* data_writer, SparkDsSink set_message, void* sink); + + /* Aborts the data written by the task. */ + int (*abort_data_writer)(void* data_writer); + + /* Releases a data writer handle, after its commit or abort. */ + void (*close_data_writer)(void* data_writer); +} SparkDsLibrary; + +/* + * The function that a library exports. Spark calls it once in each JVM, when it loads the library, + * before any other function. + * + * host: Spark, valid until the JVM exits + * library: set to the function table of the library, for the version of this interface that it + * implements, or the newest one up to host->abi_version if it implements several. Spark + * reads the table after the call, so it is typically the address of a static struct. + * + * Returns 0 on success, or any other value if the library cannot run, after it logs why with + * host->log. + */ +SPARK_DS_EXPORT int spark_ds_init(const SparkDsHost* host, const SparkDsLibrary** library); + +#ifdef __cplusplus +} +#endif + +#endif /* SPARK_DATASOURCE_H */ From e6b2915359f6c362d8b5545b43ec2783fd8d41c1 Mon Sep 17 00:00:00 2001 From: Gengliang Wang Date: Thu, 15 Oct 2026 23:29:39 +0000 Subject: [PATCH 2/6] Align Rust SDK with native planner and capability lifecycle --- .../main/rust/native-data-source/README.md | 24 ++- .../src/main/rust/native-data-source/smoke.py | 27 ++- .../native-data-source/tests/c_abi_smoke.c | 200 +++++++++++++++--- .../native-data-source/tests/c_abi_smoke.rs | 2 +- sql/native/rust/README.md | 83 ++++++-- sql/native/rust/src/bridge.rs | 111 +++++++--- sql/native/rust/src/export.rs | 6 +- sql/native/rust/src/lib.rs | 12 +- .../rust/tests/fixtures/spark_datasource.h | 19 +- 9 files changed, 385 insertions(+), 99 deletions(-) diff --git a/examples/src/main/rust/native-data-source/README.md b/examples/src/main/rust/native-data-source/README.md index 9405136309be7..1275bf92b9905 100644 --- a/examples/src/main/rust/native-data-source/README.md +++ b/examples/src/main/rust/native-data-source/README.md @@ -23,13 +23,19 @@ The connector in [src/lib.rs](src/lib.rs) exports `rust_sdk_range` through the S does not use JNI, raw pointers, or unsafe code. The example targets C ABI version 1 from [Spark PR #59424][native-source-pr], pinned to commit -`51af65c63a9369d793aab1a198ddde4c90c8f85b`. The SDK and example build without Spark or a JVM. +`534dfca15723c7208061b1e9e4d80c5757360a2a`. The SDK and example build without Spark or a JVM. Running the Spark smoke script requires a distribution built with that native data source framework and Java 22 or later on the driver and executors. The JNI SDK prototype in PR #59442 targets a different interface. -The source advertises only batch reads. It does not implement writes or Structured Streaming; -the framework rejects those operations from its declared capabilities. +The source's native capability mask is fixed to `SPARK_DS_BATCH_READ` (`1`). It does not implement +writes or Structured Streaming. During analysis, Spark creates and closes a source with the table +options to inspect its capabilities. It exposes `BATCH_WRITE` and `TRUNCATE` to Data Source V2 only +when declared, so this example rejects batch append and overwrite before evaluating their input. +The framework still reports the Data Source V2 flags `BATCH_READ`, `MICRO_BATCH_READ`, and +`STREAMING_WRITE` to keep operations on the V2 path instead of falling back to V1. These routing +flags do not add native capabilities: each operation checks the native mask again with its own +options. A source created just for capability discovery is closed without creating a scan. ## Data and options @@ -65,6 +71,13 @@ zero-column batches for operations such as `COUNT(*)`. - a serializable `RangeReaderFactory` with output columns and batch size; - disjoint `RangePartition` values with inclusive starts and exclusive ends. +This API matches the framework's separation of driver planning and executor reading. Spark's +`NativeBatchPlanner` owns the native handle; its `NativePartitionReaderFactory` carries the +serialized Rust factory and schema to executors. Java `ScanBuilder.build()` caches a scan but +does not call Rust `plan()`. The first request for batch partitions or a reader factory triggers +native `partitions` and then `serialize_reader`, which share the SDK's cached plan. Spark then +closes the native driver handle. The Rust API and C ABI layout remain unchanged. + The SDK serializes this data for the executors. Each factory opens a lazy Arrow reader using a partition and a lifetime-bound `TaskContext`. The example checks task cancellation before generating each batch and keeps memory bounded by the configured batch size. Neither the factory @@ -117,8 +130,9 @@ df.where("id >= 10 AND id < 20").select("label", "id").show() ``` [smoke.py](smoke.py) checks rows, projection, unfiltered and filtered limits, zero-column counts, -and residual predicates against Spark. Run it against the pinned framework before relying on -end-to-end compatibility. +and residual predicates against Spark. It also checks that batch append and overwrite fail during +analysis before evaluating an input containing `raise_error` and a shuffle. Run it against the +pinned framework before relying on end-to-end compatibility. [sdk]: ../../../../../sql/native/rust/README.md [native-source-pr]: https://github.com/apache/spark/pull/59424 diff --git a/examples/src/main/rust/native-data-source/smoke.py b/examples/src/main/rust/native-data-source/smoke.py index 6f78bd1ec2548..c6b2ddb90749d 100644 --- a/examples/src/main/rust/native-data-source/smoke.py +++ b/examples/src/main/rust/native-data-source/smoke.py @@ -19,6 +19,7 @@ import argparse from pathlib import Path +from pyspark.errors import AnalysisException from pyspark.sql import SparkSession, functions as F @@ -28,7 +29,11 @@ def main(): args = parser.parse_args() package = args.package.resolve(strict=True) - spark = SparkSession.builder.appName("Rust native data source smoke").getOrCreate() + spark = ( + SparkSession.builder.appName("Rust native data source smoke") + .config("spark.sql.adaptive.enabled", "true") + .getOrCreate() + ) try: spark.addArtifact(str(package)) df = ( @@ -68,6 +73,26 @@ def main(): assert [row.id for row in labels] == [3, 13, 23] residual_limit = df.where(F.col("label") == "row-21").limit(1).select("id").collect() assert [row.id for row in residual_limit] == [21] + + # Analysis must reject both modes before AQE can run the input shuffle. + failing_input = ( + spark.range(10) + .selectExpr( + "if(id >= 0, raise_error('input was evaluated'), id) AS id", "'row' AS label" + ) + .repartition(2) + ) + for mode in ("append", "overwrite"): + try: + failing_input.write.format("rust_sdk_range").mode(mode).save() + except AnalysisException as error: + assert error.getCondition() == "UNSUPPORTED_FEATURE.TABLE_OPERATION", str(error) + assert error.getMessageParameters() == { + "tableName": "`rust_sdk_range`", + "operation": "batch write", + }, str(error) + else: + raise AssertionError(f"{mode} should reject this batch-read-only source") print("Rust native data source Spark smoke checks passed") finally: spark.stop() diff --git a/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c b/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c index bdbdae8f039e8..f6f71b3edfc55 100644 --- a/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c +++ b/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -142,16 +143,17 @@ static void* create_source(const char* const* keys, const char* const* values, i return source; } -static Plan plan_scan( +static void check_capability_discovery( + const char* const* keys, const char* const* values, int64_t count) { + /* Analysis may create a source only to inspect its capabilities, then immediately close it. */ + void* source = create_source(keys, values, count); + library->close(source); +} + +static void* create_planner( const char* const* keys, const char* const* values, - int64_t option_count, - const char* const* predicates, - const int32_t* expected_uses, - int64_t predicate_count, - int32_t limit, - const char* const* projection, - int64_t projection_count) { + int64_t option_count) { void* source = create_source(keys, values, option_count); struct ArrowSchema schema = {0}; check_status(library->schema(source, &schema)); @@ -169,6 +171,43 @@ static Plan plan_scan( schema.release(&schema); } library->close(source); + return reader; +} + +static Plan finish_plan(void* reader) { + Plan plan = {0}; + check_status(library->partitions(reader, add_bytes, &plan.partitions)); + ByteList repeated = {0}; + check_status(library->partitions(reader, add_bytes, &repeated)); + CHECK(repeated.length == plan.partitions.length, "Repeated planning changed partition count"); + for (size_t i = 0; i < repeated.length; i++) { + check_bytes(&plan.partitions.values[i], &repeated.values[i]); + } + free_list(&repeated); + + ByteList state = {0}; + check_status(library->serialize_reader(reader, add_bytes, &state)); + check_status(library->serialize_reader(reader, add_bytes, &state)); + CHECK(state.length == 2, "Each serialization must emit exactly one state value"); + check_bytes(&state.values[0], &state.values[1]); + plan.state = state.values[0]; + free(state.values[1].data); + free(state.values); + library->close_reader(reader); + return plan; +} + +static Plan plan_scan( + const char* const* keys, + const char* const* values, + int64_t option_count, + const char* const* predicates, + const int32_t* expected_uses, + int64_t predicate_count, + int32_t limit, + const char* const* projection, + int64_t projection_count) { + void* reader = create_planner(keys, values, option_count); if (predicate_count != 0) { int32_t* uses = calloc((size_t) predicate_count, sizeof(int32_t)); @@ -191,26 +230,7 @@ static Plan plan_scan( CHECK(pruned, "The Range source must apply its projection"); } - Plan plan = {0}; - check_status(library->partitions(reader, add_bytes, &plan.partitions)); - ByteList repeated = {0}; - check_status(library->partitions(reader, add_bytes, &repeated)); - CHECK(repeated.length == plan.partitions.length, "Repeated planning changed partition count"); - for (size_t i = 0; i < repeated.length; i++) { - check_bytes(&plan.partitions.values[i], &repeated.values[i]); - } - free_list(&repeated); - - ByteList state = {0}; - check_status(library->serialize_reader(reader, add_bytes, &state)); - check_status(library->serialize_reader(reader, add_bytes, &state)); - CHECK(state.length == 2, "Each serialization must emit exactly one state value"); - check_bytes(&state.values[0], &state.values[1]); - plan.state = state.values[0]; - free(state.values[1].data); - free(state.values); - library->close_reader(reader); - return plan; + return finish_plan(reader); } static void free_plan(Plan* plan) { @@ -316,6 +336,40 @@ static void consume_plan( CHECK(rows == expected_rows, "Unexpected total row count"); } +static void check_independent_planners( + const char* const* keys, + const char* const* values, + int64_t count, + const char* lower, + const char* upper) { + void* first = create_planner(keys, values, count); + void* second = create_planner(keys, values, count); + int32_t use = SPARK_DS_PREDICATE_UNUSED; + check_status(library->push_predicates(first, &upper, 1, &use)); + CHECK(use == SPARK_DS_PREDICATE_COMPLETE, "The first planner did not accept its filter"); + const char* label_only[] = {"label"}; + bool pruned = false; + check_status(library->prune_columns(first, label_only, 1, &pruned)); + CHECK(pruned, "The first planner did not accept its projection"); + Plan first_plan = finish_plan(first); + + /* The second planner must neither inherit the first filter nor change its serialized plan. */ + use = SPARK_DS_PREDICATE_UNUSED; + check_status(library->push_predicates(second, &lower, 1, &use)); + CHECK(use == SPARK_DS_PREDICATE_COMPLETE, "The second planner did not accept its filter"); + const char* id_only[] = {"id"}; + pruned = false; + check_status(library->prune_columns(second, id_only, 1, &pruned)); + CHECK(pruned, "The second planner did not accept its projection"); + Plan second_plan = finish_plan(second); + + /* Both driver planners are already closed; executor reads only get the copied bytes. */ + consume_plan(&second_plan, 13, 12, id_only, 1, false); + consume_plan(&first_plan, 10, 12, label_only, 1, false); + free_plan(&second_plan); + free_plan(&first_plan); +} + static void check_cancellation(const Plan* plan) { CHECK(plan->partitions.length == 1, "Cancellation test requires one partition"); SparkDsTaskContext task = task_context(false); @@ -347,6 +401,88 @@ static void check_cancellation(const Plan* plan) { interrupted = false; } +enum StreamOperation { + STREAM_SCHEMA, + STREAM_NEXT, + STREAM_ERROR, + STREAM_RELEASE +}; + +typedef struct StreamCall { + struct ArrowArrayStream* stream; + enum StreamOperation operation; + struct ArrowSchema schema; + struct ArrowArray batch; + int status; + char error[256]; +} StreamCall; + +static void* run_stream_call(void* argument) { + StreamCall* call = argument; + switch (call->operation) { + case STREAM_SCHEMA: + call->status = call->stream->get_schema(call->stream, &call->schema); + break; + case STREAM_NEXT: + call->status = call->stream->get_next(call->stream, &call->batch); + break; + case STREAM_ERROR: { + const char* error = call->stream->get_last_error(call->stream); + CHECK(error != NULL, "A stream error must remain available on another thread"); + snprintf(call->error, sizeof(call->error), "%s", error); + break; + } + case STREAM_RELEASE: + call->stream->release(call->stream); + break; + } + return NULL; +} + +static void stream_call_on_thread(StreamCall* call) { + pthread_t thread; + CHECK(pthread_create(&thread, NULL, run_stream_call, call) == 0, + "Cannot create a stream callback thread"); + CHECK(pthread_join(thread, NULL) == 0, "Cannot join a stream callback thread"); +} + +static void check_cross_thread_stream(const Plan* plan) { + SparkDsTaskContext task = task_context(false); + expected_task = &task; + interrupted = false; + interruption_checks = 0; + struct ArrowArrayStream stream = open_stream(plan, 0, &task); + + /* Callbacks do not overlap, but each worker has a different lifetime and thread-local state. */ + StreamCall schema = {.stream = &stream, .operation = STREAM_SCHEMA}; + stream_call_on_thread(&schema); + CHECK(schema.status == 0 && schema.schema.release != NULL, "Cannot get schema on a new thread"); + + StreamCall first = {.stream = &stream, .operation = STREAM_NEXT}; + stream_call_on_thread(&first); + CHECK(first.status == 0 && first.batch.release != NULL, "Cannot read a batch on a new thread"); + CHECK(interruption_checks > 0, "The worker did not call the task cancellation callback"); + + interrupted = true; + StreamCall next = {.stream = &stream, .operation = STREAM_NEXT}; + stream_call_on_thread(&next); + CHECK(next.status != 0, "The worker must observe live task cancellation"); + StreamCall error = {.stream = &stream, .operation = STREAM_ERROR}; + stream_call_on_thread(&error); + CHECK(strstr(error.error, "cancel") != NULL, "The stream error changed across threads"); + if (next.batch.release != NULL) { + next.batch.release(&next.batch); + } + + StreamCall release = {.stream = &stream, .operation = STREAM_RELEASE}; + stream_call_on_thread(&release); + CHECK(stream.release == NULL, "Cross-thread release must clear the stream callback"); + check_batch(&schema.schema, &first.batch, 0, 2); + first.batch.release(&first.batch); + schema.schema.release(&schema.schema); + interrupted = false; +} + static void check_error_message(void) { const char* keys[] = {"batchsize"}; const char* values[] = {"0"}; @@ -359,7 +495,8 @@ static void check_error_message(void) { }; uint64_t capabilities = 0; void* source = NULL; - CHECK(library->create(&info, &capabilities, &source) != 0, "Invalid options must return an error"); + CHECK(library->create(&info, &capabilities, &source) != 0, + "Invalid options must return an error"); const char* error = library->last_error(); CHECK(error != NULL && strstr(error, "batchSize must be positive") != NULL, "The C error channel did not preserve the provider's message"); @@ -428,6 +565,10 @@ int main(int argc, char** argv) { const char* id_only[] = {"id"}; const char* original[] = {"id", "label"}; + check_capability_discovery(keys, values, 4); + check_capability_discovery(NULL, NULL, 0); + check_independent_planners(keys, values, 4, lower, upper); + Plan filtered = plan_scan(keys, values, 4, predicates, uses, 2, -1, reordered, 2); consume_plan(&filtered, 13, 9, reordered, 2, false); free_plan(&filtered); @@ -450,6 +591,7 @@ int main(int argc, char** argv) { Plan cancellable = plan_scan(short_keys, one_partition, 3, NULL, NULL, 0, -1, NULL, -1); consume_plan(&cancellable, 0, 20, original, 2, true); check_cancellation(&cancellable); + check_cross_thread_stream(&cancellable); free_plan(&cancellable); check_error_message(); diff --git a/examples/src/main/rust/native-data-source/tests/c_abi_smoke.rs b/examples/src/main/rust/native-data-source/tests/c_abi_smoke.rs index 78d7b5e1657c8..b5d63b3f7ad35 100644 --- a/examples/src/main/rust/native-data-source/tests/c_abi_smoke.rs +++ b/examples/src/main/rust/native-data-source/tests/c_abi_smoke.rs @@ -43,7 +43,7 @@ fn c_abi_arrow_round_trip() -> Result<(), Box> { 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"]) + .args(["-std=c11", "-Wall", "-Wextra", "-Werror", "-pthread"]) .arg("-I") .arg(header) .arg(manifest.join("tests/c_abi_smoke.c")) diff --git a/sql/native/rust/README.md b/sql/native/rust/README.md index 4749102680f5d..12b5d4bed925a 100644 --- a/sql/native/rust/README.md +++ b/sql/native/rust/README.md @@ -19,7 +19,7 @@ limitations under the License. This experimental SDK implements **batch reads** over the C ABI proposed in [Spark PR #59424](https://github.com/apache/spark/pull/59424), pinned to commit -`51af65c63a9369d793aab1a198ddde4c90c8f85b`. Spark calls the library through the +`534dfca15723c7208061b1e9e4d80c5757360a2a`. Spark calls the library through the Java Foreign Function and Memory (FFM) API. This PR contains the SDK and example, not the Spark framework; running the Spark example requires that framework and Java 22 or newer on the driver and executors. @@ -39,7 +39,8 @@ to include the example. A connector implements three traits from [`src/lib.rs`](src/lib.rs): * `BatchSource`: parse driver options, discover the schema, and create an - independent scan builder. Spark can drop the source immediately afterwards. + independent scan builder. Spark can also create and immediately drop a source + just to discover its capabilities; schema inference or a scan need not follow. * `ScanBuilder`: negotiate predicates, limit, and projection, then consume itself in `plan()` to return a typed `ScanPlan`. * `BatchReaderFactory`: serialize executor configuration and open one lazy Arrow @@ -60,13 +61,25 @@ Executor, per task: v restored factory + partition + TaskContext -> Arrow reader -> Arrow C Stream ``` -`ScanBuilder::plan()` runs on the first native `partitions` or `serialize_reader` -call, after pushdown negotiation. It is not the Java DSv2 `ScanBuilder.build()` -callback. The driver runtime caches the encoded result and shares it across -repeated callbacks without cloning its byte buffers. A failed negotiation or -planning call is terminal and retains its original error, including a panic. -Closing a driver handle releases its local resources; it cannot invalidate -factory or partition bytes already copied to Spark. +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` calls `partitions`, then +`serialize_reader`, and closes the native planner handle in `finally`. The +Rust `ScanBuilder::plan()` consumes its builder at that first native planning +call. The SDK still caches the encoded result so both native callbacks use the +same factory and partitions without planning twice or cloning byte buffers. +A failed negotiation or planning call retains its original error, including a +panic. 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 restores the Rust factory from those bytes without creating or accessing +the driver registry. 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. Task IDs @@ -81,16 +94,33 @@ spark_native_source::export_data_source!(RangeSource); ``` The macro exports only `spark_ds_init` and returns Spark's versioned function -table. `create` declares only `SPARK_DS_BATCH_READ`; streaming, write, and -truncate capabilities are absent, and their function slots are null. Query-level -pushdown results are independent of these source capabilities. The connector -uses safe Rust traits without C pointers. Rust trait objects never cross the -binary boundary between libraries. +table. The connector uses safe Rust traits without C pointers. Rust trait objects +never cross the binary boundary between libraries. The implementation separates public connector traits, driver handle/planning state in `bridge.rs`, Arrow ownership in `stream.rs`, and C bindings/adapters in `abi.rs` and `export.rs`. It does not change Spark's native ABI or JVM classes. +## Source capabilities + +Every `BatchSource` created by this SDK declares exactly `SPARK_DS_BATCH_READ`. +Streaming, write, and truncate flags remain absent, and unsupported function +slots are null. These native flags describe what the source implements; +query-level predicate, limit, and projection results are separate 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. Keep `BatchSource::new` +lightweight and safe to drop immediately; do not assume that `schema` or `scan` +will follow. Executor resources belong in `BatchReaderFactory::open`. + +The pinned framework propagates native `BATCH_WRITE` and `TRUNCATE` flags to +`Table.capabilities()`, so a batch-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 this SDK's batch-read-only native mask. + ## Arrow contract `BatchReaderFactory::open` returns a lifetime-bound `Box`. Return fixed-schema batches lazily with bounded memory. Exported arrays own their buffers independently of the reader: Spark may release a stream before releasing its previously returned batches. +Stream callbacks, including release, may run on different threads, but never +concurrently. Readers must be `Send`; they do not need to be `Sync` or tied to +the thread that opened the stream. Stream errors are owned by the stream and +remain readable after a callback moves to another thread. + The SDK validates the restored negotiated schema before exporting a reader and checks each batch's schema and recursive array offsets. It compares Spark logical types, ignoring nullability and Arrow metadata, accepting the supported @@ -152,7 +187,7 @@ conservative defaults so a connector can add support incrementally. ## Build and test without Spark Use Rust 1.85 or newer (as required by Arrow 58) and a C compiler. The C integration -test uses `dlopen` on Linux or macOS. From this directory: +test uses `dlopen` and pthreads on Linux or macOS. From this directory: ```bash cargo build --workspace --locked @@ -165,22 +200,26 @@ Build before testing so the C harness loads the current example dynamic library. The harness compiles against the unmodified framework header in [`tests/fixtures/spark_datasource.h`](tests/fixtures/spark_datasource.h), copied from the pinned commit above. It calls the actual function table and validates -capabilities and optional slots, driver handle independence, pushdown and -planning, serialized executor reads, cancellation, and Arrow buffer ownership. +capability-only source discovery and optional slots, independent driver planners, +pushdown and planning, serialized executor reads, cancellation, sequential +cross-thread stream callbacks, and Arrow buffers surviving stream release. This checks the C boundary without requiring Java; it does not substitute for Spark planner and task-lifecycle testing. Unit tests additionally cover error and panic containment, terminal planning -failures, schema/layout validation, sized ABI structures, task-context lifetimes, -integer boundaries, projections, and deterministic partition plans. +failures, retries after the driver registry is dropped, schema/layout validation, +sized ABI structures, task-context lifetimes, integer boundaries, projections, +and deterministic partition plans. ## Package and run with Spark See the [Range example guide][range] for packaging and `spark-submit` commands. Its `package.py` creates a single-platform `.sparkpkg`; `smoke.py` checks reading, -filtered reads, projection, and unfiltered limit pushdown against a Spark -installation containing the pinned #59424 framework. No full Spark integration -is implied by the standalone Rust/C tests. +filtered reads, projection, unfiltered limit pushdown, and analysis-time rejection +of unsupported batch writes against the pinned #59424 framework. The write +checks use failing input behind a shuffle to verify Spark rejects the operation +before computing its data. No full Spark integration is implied by the +standalone Rust/C tests. ## Scope and limitations diff --git a/sql/native/rust/src/bridge.rs b/sql/native/rust/src/bridge.rs index a83a2be3c5e9a..0bb0ed6eea038 100644 --- a/sql/native/rust/src/bridge.rs +++ b/sql/native/rust/src/bridge.rs @@ -13,7 +13,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Transport-independent driver handles and executor plan restoration. +//! Driver-only handles and planners, with stateless executor plan restoration. use std::any::Any; use std::collections::HashMap; @@ -42,7 +42,7 @@ type HandleMap = Mutex>>>; pub struct Bridge { next_handle: AtomicI64, sources: HandleMap, - readers: HandleMap>, + planners: HandleMap>, } impl Default for Bridge { @@ -50,7 +50,7 @@ impl Default for Bridge { Self { next_handle: AtomicI64::new(1), sources: Mutex::new(HashMap::new()), - readers: Mutex::new(HashMap::new()), + planners: Mutex::new(HashMap::new()), } } } @@ -80,7 +80,7 @@ impl Bridge { stream::validate_schema(&schema)?; let source = get(&self.sources, handle)?; let builder = lock(&source).scan(schema.clone())?; - self.insert(&self.readers, Reader::new(builder, schema)) + self.insert(&self.planners, Planner::new(builder, schema)) } pub fn push_predicates( @@ -88,7 +88,7 @@ impl Bridge { handle: i64, predicates: &[Predicate], ) -> Result> { - let reader = get(&self.readers, handle)?; + let reader = get(&self.planners, handle)?; let mut reader = lock(&reader); let support = reader.negotiate(|builder| { let support = builder.push_filters(predicates)?; @@ -102,7 +102,7 @@ impl Bridge { } pub fn push_limit(&self, handle: i64, limit: usize) -> Result { - let reader = get(&self.readers, handle)?; + let reader = get(&self.planners, handle)?; let mut reader = lock(&reader); let has_residual = reader.has_residual; reader.negotiate(|builder| { @@ -115,7 +115,7 @@ impl Bridge { } pub fn prune_columns(&self, handle: i64, columns: &[String]) -> Result { - let reader = get(&self.readers, handle)?; + let reader = get(&self.planners, handle)?; let mut reader = lock(&reader); let indices: Vec<_> = columns .iter() @@ -130,7 +130,7 @@ impl Bridge { } pub fn planned(&self, handle: i64) -> Result> { - let reader = get(&self.readers, handle)?; + let reader = get(&self.planners, handle)?; let result = lock(&reader).plan(); result } @@ -141,27 +141,27 @@ impl Bridge { } pub fn close_reader(&self, handle: i64) -> Result<()> { - remove(&self.readers, handle); + remove(&self.planners, handle); Ok(()) } +} - /// Restores only serialized values; no driver-side handle is needed. - /// - /// # 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( - &self, - task: *const crate::abi::SparkDsTaskContext, - state: &[u8], - partition: &[u8], - ) -> Result { - let state: ReaderState<::Factory> = decode(state)?; - let task = unsafe { TaskContext::from_raw(task)? }; - let reader = state.factory.open(decode(partition)?, task)?; - unsafe { stream::export(reader, &Arc::new(state.schema)) } - } +/// Restores an executor reader using only its factory type and serialized values. +/// No source, planner, 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( + task: *const crate::abi::SparkDsTaskContext, + state: &[u8], + partition: &[u8], +) -> Result { + let state: ReaderState = decode(state)?; + let task = unsafe { TaskContext::from_raw(task)? }; + let reader = state.factory.open(decode(partition)?, task)?; + unsafe { stream::export(reader, &Arc::new(state.schema)) } } fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { @@ -184,7 +184,7 @@ fn remove(handles: &HandleMap, handle: i64) { drop(value); } -struct Reader { +struct Planner { phase: Phase, schema: SchemaRef, has_residual: bool, @@ -207,7 +207,7 @@ struct ReaderState { schema: Schema, } -impl Reader { +impl Planner { fn new(builder: B, schema: SchemaRef) -> Self { Self { phase: Phase::Negotiating(builder), @@ -326,6 +326,7 @@ pub fn parse_predicates(predicates: &[String]) -> Result> { mod tests { use std::sync::atomic::AtomicUsize; + use arrow_array::ffi_stream::ArrowArrayStreamReader; use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator}; use arrow_schema::{DataType, Field}; @@ -640,6 +641,54 @@ mod tests { bridge.close_reader(second).unwrap(); } + #[test] + fn executor_retries_use_only_serialized_values_after_the_driver_is_dropped() { + let bridge = Bridge::::default(); + let source = Source::new(Options::new()).unwrap(); + let builds = source.builds.clone(); + let source_drops = source.drops.clone(); + let builder_drops = source.builder_drops.clone(); + let handle = bridge.insert(&bridge.sources, source).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let planner = bridge.create_reader(handle, schema).unwrap(); + bridge.close_source(handle).unwrap(); + let plan = bridge.planned(planner).unwrap(); + // Model Spark copying the factory and partition descriptors out of the plugin. + let state = plan.state.clone(); + let partitions = plan.partitions.clone(); + drop(plan); + bridge.close_reader(planner).unwrap(); + drop(bridge); + assert_eq!(builds.load(Ordering::SeqCst), 1); + assert_eq!(source_drops.load(Ordering::SeqCst), 1); + assert_eq!(builder_drops.load(Ordering::SeqCst), 1); + + for attempt in 0..2 { + for (index, partition) in partitions.iter().enumerate() { + let task = crate::abi::SparkDsTaskContext { + size: std::mem::size_of::() as i64, + stage_id: 1, + stage_attempt_number: 0, + partition_id: index as i32, + attempt_number: attempt, + task_attempt_id: i64::from(attempt) * 2 + index as i64, + is_interrupted: None, + }; + // The executor entry point depends only on Factory, not Source or Builder. + let stream = unsafe { open_reader::(&task, &state, partition) }.unwrap(); + let mut reader = ArrowArrayStreamReader::try_new(stream).unwrap(); + let batch = reader.next().unwrap().unwrap(); + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.values(), &[42 + [2, 7][index]]); + assert!(reader.next().is_none()); + } + } + } + #[test] fn executor_rejects_a_factory_schema_different_from_the_negotiated_schema() { let bridge = Bridge::::default(); @@ -662,8 +711,10 @@ mod tests { }; // The factory returns one id column, but the driver negotiated an empty schema. // The expected schema must survive serialization and the removal of driver handles. - let error = - error_message(unsafe { bridge.open_reader(&task, &plan.state, &plan.partitions[0]) }); + drop(bridge); + let error = error_message(unsafe { + open_reader::(&task, &plan.state, &plan.partitions[0]) + }); assert!(error.contains("schema"), "{error}"); } diff --git a/sql/native/rust/src/export.rs b/sql/native/rust/src/export.rs index 7f0bf2b30b452..5749473298827 100644 --- a/sql/native/rust/src/export.rs +++ b/sql/native/rust/src/export.rs @@ -24,7 +24,9 @@ macro_rules! export_data_source { use super::*; use std::ffi::{c_char, c_int, c_void}; use $crate::abi::{self, ArrowArrayStream, ArrowSchema, SparkDsLibrary, SparkDsSink}; - use $crate::bridge::Bridge; + use $crate::bridge::{open_reader, Bridge}; + + type Factory = <<$source as $crate::BatchSource>::Scan as $crate::ScanBuilder>::Factory; fn bridge() -> &'static Bridge<$source> { static BRIDGE: std::sync::OnceLock> = std::sync::OnceLock::new(); @@ -207,7 +209,7 @@ macro_rules! export_data_source { abi::require_output(output)?; let state = unsafe { abi::slice(state, state_length)? }; let partition = unsafe { abi::slice(partition, partition_length)? }; - let stream = unsafe { bridge().open_reader(task, state, partition)? }; + let stream = unsafe { open_reader::(task, state, partition)? }; unsafe { output.write(stream) }; Ok(()) }) diff --git a/sql/native/rust/src/lib.rs b/sql/native/rust/src/lib.rs index 2fa4343c96b08..5857a784b301e 100644 --- a/sql/native/rust/src/lib.rs +++ b/sql/native/rust/src/lib.rs @@ -79,11 +79,15 @@ pub struct ScanPlan { /// Driver-side source configuration and schema discovery. /// /// Registering this trait declares only the batch-read capability. Writes and -/// streaming remain absent from the native function table. +/// streaming remain absent from the native function table. Spark also creates +/// sources to discover table capabilities during query analysis; such a source +/// can be dropped without calling `schema` or `scan`. pub trait BatchSource: Send + Sized + 'static { type Scan: ScanBuilder; - /// Spark passes option keys in lower case. + /// Spark passes option keys in lower case. Keep construction lightweight: + /// a capability probe uses only table options, while a scan may override + /// them with query options. Open executor resources in the factory. fn new(options: Options) -> Result; fn schema(&self) -> Result; @@ -104,6 +108,8 @@ pub trait BatchReaderFactory: Serialize + DeserializeOwned + Send + 'static { /// Return lazy batches with a fixed schema and bounded memory. The reader /// may retain `task` until the stream is released. Check cancellation during /// long operations and use timeouts for blocking I/O. + /// Spark may call stream callbacks on different threads, never concurrently; + /// the reader must be `Send` but does not need to be `Sync`. /// /// The SDK checks Spark-compatible Arrow types, the negotiated schema, and /// zero array offsets recursively. It reports incompatible layouts instead @@ -122,6 +128,8 @@ pub trait BatchReaderFactory: Serialize + DeserializeOwned + Send + 'static { /// `serialize_reader` call and caches the result, including a planning failure. /// This is later than Java DSv2 `ScanBuilder.build()`. Only serialized factory /// and partition descriptors travel to executors; driver handles stay local. +/// It corresponds to Spark's driver-only `ColumnarBatchPlanner`: the native +/// planner calls `partitions`, then `serialize_reader`, then closes the handle. pub trait ScanBuilder: Send + Sized + 'static { type Factory: BatchReaderFactory; diff --git a/sql/native/rust/tests/fixtures/spark_datasource.h b/sql/native/rust/tests/fixtures/spark_datasource.h index 3ff46d2871560..16e21494dba85 100644 --- a/sql/native/rust/tests/fixtures/spark_datasource.h +++ b/sql/native/rust/tests/fixtures/spark_datasource.h @@ -50,6 +50,10 @@ * creates a data source with create on the driver for each schema inference, scan and write, and * closes it right after. create also returns the capabilities of the data source: Spark reports * the operations that a data source does not declare as unsupported, without calling the library. + * Spark also creates a data source while it analyzes a query, to reject a batch write that the + * data source does not support before it computes the data to write. That data source has the + * options of the table, those of CREATE TABLE for a table in a catalog, which the options of a + * write can override, so Spark checks the capabilities of each operation again. * * - Batch reads. On the driver, Spark creates a reader with create_reader, pushes operations down * by calling, each at most once and in this order, push_predicates, push_limit and @@ -149,8 +153,9 @@ * from a thread other than the threads that used it, such as a JVM Cleaner thread when Spark * releases a handle once it is no longer referenced, for example the last stream writer of a * streaming query. A SparkDsSink must be called on the thread of the call that it was passed - * to, during that call. The functions of SparkDsHost and SparkDsTaskContext can be called from - * any thread. + * to, during that call. The functions of SparkDsHost and SparkDsTaskContext, and the callbacks + * of the stream that read initializes, can be called from any thread. The callbacks of the + * stream never run concurrently. * * Compatibility * ------------- @@ -575,11 +580,11 @@ typedef struct SparkDsLibrary { * there is none * partition: a partition returned by partitions or stream_partitions * out: the stream to initialize. Each array of the stream is a struct array whose fields are the - * columns to read: the columns of the schema to read, or the pruned columns. Spark calls the - * stream from the thread of the task that reads the partition. get_next may block, for - * example on network I/O, but Spark cannot interrupt native code: when the task is killed, - * Spark only stops reading once get_next returns, and then releases the stream. So a - * get_next that blocks forever keeps the task running, unless the task reaper + * columns to read: the columns of the schema to read, or the pruned columns. The + * callbacks of the stream can be called from any thread, but never concurrently. get_next + * may block, for example on network I/O, but Spark cannot interrupt native code: when the + * task is killed, Spark only stops reading once get_next returns, and then releases the + * stream. So a get_next that blocks forever keeps the task running, unless the task reaper * (spark.task.reaper.enabled and spark.task.reaper.killTimeout) kills the executor. Use * timeouts for blocking operations, and check task->is_interrupted. */ From 0bd726541f6c0444eb1b0f3841e52a2ff9050593 Mon Sep 17 00:00:00 2001 From: Gengliang Wang Date: Fri, 16 Oct 2026 20:22:14 +0000 Subject: [PATCH 3/6] Expose per-source capabilities in the Rust native SDK --- .../main/rust/native-data-source/README.md | 9 +- .../main/rust/native-data-source/src/lib.rs | 6 +- sql/native/rust/README.md | 49 +++- sql/native/rust/src/bridge.rs | 62 ++++- sql/native/rust/src/export.rs | 4 +- sql/native/rust/src/lib.rs | 35 ++- sql/native/rust/tests/capabilities.rs | 248 ++++++++++++++++++ 7 files changed, 378 insertions(+), 35 deletions(-) create mode 100644 sql/native/rust/tests/capabilities.rs diff --git a/examples/src/main/rust/native-data-source/README.md b/examples/src/main/rust/native-data-source/README.md index 1275bf92b9905..6811ef60c86a2 100644 --- a/examples/src/main/rust/native-data-source/README.md +++ b/examples/src/main/rust/native-data-source/README.md @@ -28,9 +28,12 @@ Running the Spark smoke script requires a distribution built with that native da framework and Java 22 or later on the driver and executors. The JNI SDK prototype in PR #59442 targets a different interface. -The source's native capability mask is fixed to `SPARK_DS_BATCH_READ` (`1`). It does not implement -writes or Structured Streaming. During analysis, Spark creates and closes a source with the table -options to inspect its capabilities. It exposes `BATCH_WRITE` and `TRUNCATE` to Data Source V2 only +The source explicitly implements `BatchSource::capabilities()` to return +`vec![SourceCapability::BatchRead]`, producing the native mask `SPARK_DS_BATCH_READ` (`1`). +Connectors can override this method based on their options; the SDK validates the declaration +against its implemented callbacks. This example does not implement writes or Structured Streaming. +During analysis, Spark creates and closes a source with the table options to inspect its +capabilities. It exposes `BATCH_WRITE` and `TRUNCATE` to Data Source V2 only when declared, so this example rejects batch append and overwrite before evaluating their input. The framework still reports the Data Source V2 flags `BATCH_READ`, `MICRO_BATCH_READ`, and `STREAMING_WRITE` to keep operations on the V2 path instead of falling back to V1. These routing diff --git a/examples/src/main/rust/native-data-source/src/lib.rs b/examples/src/main/rust/native-data-source/src/lib.rs index 360459dffefa9..dbca41207ab7a 100644 --- a/examples/src/main/rust/native-data-source/src/lib.rs +++ b/examples/src/main/rust/native-data-source/src/lib.rs @@ -25,7 +25,7 @@ use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; use serde::{Deserialize, Serialize}; use spark_native_source::{ BatchReader, BatchReaderFactory, BatchSource, FilterSupport, LimitSupport, Options, Predicate, - Result, ScanBuilder, ScanPlan, TaskContext, + Result, ScanBuilder, ScanPlan, SourceCapability, TaskContext, }; pub struct RangeSource { @@ -86,6 +86,10 @@ impl BatchSource for RangeSource { }) } + fn capabilities(&self) -> Vec { + vec![SourceCapability::BatchRead] + } + fn schema(&self) -> Result { output_schema(&["id".into(), "label".into()]) } diff --git a/sql/native/rust/README.md b/sql/native/rust/README.md index 12b5d4bed925a..471874866a062 100644 --- a/sql/native/rust/README.md +++ b/sql/native/rust/README.md @@ -38,8 +38,8 @@ to include the example. A connector implements three traits from [`src/lib.rs`](src/lib.rs): -* `BatchSource`: parse driver options, discover the schema, and create an - independent scan builder. Spark can also create and immediately drop a source +* `BatchSource`: parse driver options, declare capabilities, discover the schema, + and create an independent scan builder. Spark can also create and drop a source just to discover its capabilities; schema inference or a scan need not follow. * `ScanBuilder`: negotiate predicates, limit, and projection, then consume itself in `plan()` to return a typed `ScanPlan`. @@ -103,23 +103,48 @@ state in `bridge.rs`, Arrow ownership in `stream.rs`, and C bindings/adapters in ## Source capabilities -Every `BatchSource` created by this SDK declares exactly `SPARK_DS_BATCH_READ`. -Streaming, write, and truncate flags remain absent, and unsupported function -slots are null. These native flags describe what the source implements; -query-level predicate, limit, and projection results are separate contracts. +Connector developers declare native capabilities through +`BatchSource::capabilities(&self) -> Vec`. The default is +`vec![SourceCapability::BatchRead]`; override it to describe the source instance +created with its effective options. For example, inside a source implementation: + +```rust +fn capabilities(&self) -> Vec { + if self.batch_reads_enabled { + vec![SourceCapability::BatchRead] + } else { + vec![] + } +} +``` + +The SDK evaluates this method once during `create`, validates the declaration, +and returns the resulting native flag mask to Spark. Duplicate entries are +ignored. An empty set still allows schema discovery and closing the source; +batch scan creation fails before calling the connector's `scan` method. The +declaration stays fixed for that source instance. + +`SourceCapability` names all current native operations: `BatchRead`, +`MicroBatchRead`, `BatchWrite`, `StreamingWrite`, and `Truncate`. This batch-read +prototype accepts only `BatchRead`; declaring other variants fails source +creation with an explicit error because their callbacks are not implemented. +Unsupported function slots remain null. These native flags describe supported +operations; query-level predicate, limit, and projection results are separate +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. Keep `BatchSource::new` -lightweight and safe to drop immediately; do not assume that `schema` or `scan` -will follow. Executor resources belong in `BatchReaderFactory::open`. +them with query options and checks capabilities again on a new source instance. +Keep `BatchSource::new` and `capabilities` lightweight and safe to drop +immediately; do not assume that `schema` or `scan` will follow. Executor resources +belong in `BatchReaderFactory::open`. The pinned framework propagates native `BATCH_WRITE` and `TRUNCATE` flags to `Table.capabilities()`, so a batch-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 this SDK's batch-read-only native mask. +flags are not a literal copy of the native mask declared by the Rust source. ## Arrow contract @@ -206,7 +231,9 @@ cross-thread stream callbacks, and Arrow buffers surviving stream release. This checks the C boundary without requiring Java; it does not substitute for Spark planner and task-lifecycle testing. -Unit tests additionally cover error and panic containment, terminal planning +Rust tests additionally exercise the exported function table with option-dependent +capabilities, empty and duplicate declarations, and unsupported declarations or +panics during capability discovery. They also cover terminal planning failures, retries after the driver registry is dropped, schema/layout validation, sized ABI structures, task-context lifetimes, integer boundaries, projections, and deterministic partition plans. diff --git a/sql/native/rust/src/bridge.rs b/sql/native/rust/src/bridge.rs index 0bb0ed6eea038..4b52c36b87da8 100644 --- a/sql/native/rust/src/bridge.rs +++ b/sql/native/rust/src/bridge.rs @@ -27,7 +27,7 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize}; use crate::{ BatchReaderFactory, BatchSource, FilterSupport, LimitSupport, Options, Predicate, Result, - ScanBuilder, TaskContext, + ScanBuilder, SourceCapability, TaskContext, }; #[path = "stream.rs"] @@ -41,7 +41,7 @@ type HandleMap = Mutex>>>; /// so independent sources and scans do not serialize behind a global lock. pub struct Bridge { next_handle: AtomicI64, - sources: HandleMap, + sources: HandleMap>, planners: HandleMap>, } @@ -65,13 +65,36 @@ impl Bridge { Ok(handle) } - pub fn create(&self, options: Options) -> Result { - self.insert(&self.sources, S::new(options)?) + pub fn create(&self, options: Options) -> Result<(i64, u64)> { + self.register_source(S::new(options)?) + } + + fn register_source(&self, source: S) -> Result<(i64, u64)> { + let mut capabilities = 0; + for capability in source.capabilities() { + match capability { + SourceCapability::BatchRead => capabilities |= crate::abi::BATCH_READ, + _ => { + return Err(format!( + "capability {capability:?} is not supported by this batch-read SDK" + ) + .into()); + } + } + } + let handle = self.insert( + &self.sources, + SourceState { + source, + capabilities, + }, + )?; + Ok((handle, capabilities)) } pub fn schema(&self, handle: i64) -> Result { let source = get(&self.sources, handle)?; - let schema = lock(&source).schema()?; + let schema = lock(&source).source.schema()?; stream::validate_schema(&schema)?; Ok(schema) } @@ -79,7 +102,13 @@ impl Bridge { pub fn create_reader(&self, handle: i64, schema: SchemaRef) -> Result { stream::validate_schema(&schema)?; let source = get(&self.sources, handle)?; - let builder = lock(&source).scan(schema.clone())?; + let builder = { + let source = lock(&source); + if source.capabilities & crate::abi::BATCH_READ == 0 { + return Err("source does not support batch reads".into()); + } + source.source.scan(schema.clone())? + }; self.insert(&self.planners, Planner::new(builder, schema)) } @@ -184,6 +213,11 @@ fn remove(handles: &HandleMap, handle: i64) { drop(value); } +struct SourceState { + source: S, + capabilities: u64, +} + struct Planner { phase: Phase, schema: SchemaRef, @@ -466,7 +500,7 @@ mod tests { let source = Source::new(Options::new()).unwrap(); let builds = source.builds.clone(); let drops = source.drops.clone(); - let handle = bridge.insert(&bridge.sources, source).unwrap(); + let (handle, _) = bridge.register_source(source).unwrap(); let schema = bridge.schema(handle).unwrap(); let reader = bridge.create_reader(handle, schema).unwrap(); bridge.close_source(handle).unwrap(); @@ -499,7 +533,7 @@ mod tests { source.filter_support = support; source.limit_support = LimitSupport::Global; let limits = source.limits.clone(); - let handle = bridge.insert(&bridge.sources, source).unwrap(); + let (handle, _) = bridge.register_source(source).unwrap(); let reader = bridge .create_reader(handle, bridge.schema(handle).unwrap()) .unwrap(); @@ -521,7 +555,7 @@ mod tests { let mut source = Source::new(Options::new()).unwrap(); source.limit_support = support; let limits = source.limits.clone(); - let handle = bridge.insert(&bridge.sources, source).unwrap(); + let (handle, _) = bridge.register_source(source).unwrap(); let reader = bridge .create_reader(handle, bridge.schema(handle).unwrap()) .unwrap(); @@ -547,7 +581,7 @@ mod tests { source.failure = failure; let builds = source.builds.clone(); let drops = source.builder_drops.clone(); - let handle = bridge.insert(&bridge.sources, source).unwrap(); + let (handle, _) = bridge.register_source(source).unwrap(); let reader = bridge .create_reader(handle, bridge.schema(handle).unwrap()) .unwrap(); @@ -578,7 +612,7 @@ mod tests { let builds = source.builds.clone(); let limits = source.limits.clone(); let drops = source.builder_drops.clone(); - let handle = bridge.insert(&bridge.sources, source).unwrap(); + let (handle, _) = bridge.register_source(source).unwrap(); let reader = bridge .create_reader(handle, bridge.schema(handle).unwrap()) .unwrap(); @@ -603,7 +637,7 @@ mod tests { source.limit_support = LimitSupport::Global; let builds = source.builds.clone(); let limits = source.limits.clone(); - let handle = bridge.insert(&bridge.sources, source).unwrap(); + let (handle, _) = bridge.register_source(source).unwrap(); let first = bridge .create_reader(handle, bridge.schema(handle).unwrap()) .unwrap(); @@ -648,7 +682,7 @@ mod tests { let builds = source.builds.clone(); let source_drops = source.drops.clone(); let builder_drops = source.builder_drops.clone(); - let handle = bridge.insert(&bridge.sources, source).unwrap(); + let (handle, _) = bridge.register_source(source).unwrap(); let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); let planner = bridge.create_reader(handle, schema).unwrap(); bridge.close_source(handle).unwrap(); @@ -692,7 +726,7 @@ mod tests { #[test] fn executor_rejects_a_factory_schema_different_from_the_negotiated_schema() { let bridge = Bridge::::default(); - let handle = bridge.create(Options::new()).unwrap(); + let (handle, _) = bridge.create(Options::new()).unwrap(); let reader = bridge .create_reader(handle, bridge.schema(handle).unwrap()) .unwrap(); diff --git a/sql/native/rust/src/export.rs b/sql/native/rust/src/export.rs index 5749473298827..4d5e19001aebc 100644 --- a/sql/native/rust/src/export.rs +++ b/sql/native/rust/src/export.rs @@ -86,9 +86,9 @@ macro_rules! export_data_source { abi::require_output(capabilities)?; abi::require_output(source)?; let options = unsafe { abi::options(info)? }; - let handle = bridge().create(options)?; + let (handle, declared_capabilities) = bridge().create(options)?; unsafe { - capabilities.write(abi::BATCH_READ); + capabilities.write(declared_capabilities); source.write(handle as usize as *mut c_void); } Ok(()) diff --git a/sql/native/rust/src/lib.rs b/sql/native/rust/src/lib.rs index 5857a784b301e..d1a61c7776025 100644 --- a/sql/native/rust/src/lib.rs +++ b/sql/native/rust/src/lib.rs @@ -43,6 +43,23 @@ pub type Result = std::result::Result; pub type Options = BTreeMap; pub type BatchReader<'task> = Box; +/// Native operations supported by a source instance with its effective options. +/// +/// These are native capabilities, not Spark's Data Source V2 routing flags or +/// per-scan pushdown results. This batch-read SDK currently accepts only +/// [`BatchRead`](Self::BatchRead); declaring another capability fails source +/// creation because the corresponding callbacks are not implemented. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum SourceCapability { + BatchRead, + MicroBatchRead, + BatchWrite, + StreamingWrite, + /// Supports writes that replace the table's existing data. + Truncate, +} + /// How completely a connector evaluates a pushed predicate. /// Spark currently rechecks all predicates, including exact ones. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -78,10 +95,11 @@ pub struct ScanPlan { /// Driver-side source configuration and schema discovery. /// -/// Registering this trait declares only the batch-read capability. Writes and -/// streaming remain absent from the native function table. Spark also creates -/// sources to discover table capabilities during query analysis; such a source -/// can be dropped without calling `schema` or `scan`. +/// Override [`capabilities`](Self::capabilities) to declare the operations this +/// instance supports. Writes and streaming remain absent from the native +/// function table. Spark also creates sources to discover table capabilities +/// during query analysis; such a source can be dropped without calling `schema` +/// or `scan`. pub trait BatchSource: Send + Sized + 'static { type Scan: ScanBuilder; @@ -90,6 +108,15 @@ pub trait BatchSource: Send + Sized + 'static { /// them with query options. Open executor resources in the factory. fn new(options: Options) -> Result; + /// Declares this instance's capabilities, which may depend on its options. + /// The SDK evaluates this once during source creation, validates the set + /// against its implemented callbacks, and keeps it fixed for that instance. + /// Duplicate entries are ignored. An empty set permits schema discovery but + /// rejects batch scan creation without calling `scan`. + fn capabilities(&self) -> Vec { + vec![SourceCapability::BatchRead] + } + fn schema(&self) -> Result; /// The schema may be supplied by the user; reject unsupported schemas. diff --git a/sql/native/rust/tests/capabilities.rs b/sql/native/rust/tests/capabilities.rs new file mode 100644 index 0000000000000..c18a7dd603504 --- /dev/null +++ b/sql/native/rust/tests/capabilities.rs @@ -0,0 +1,248 @@ +// 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. + +use std::ffi::{c_int, c_void, CStr, CString}; +use std::mem::size_of; +use std::ptr; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +use arrow_schema::{Schema, SchemaRef}; +use serde::{Deserialize, Serialize}; +use spark_native_source::abi::{ + ArrowSchema, SparkDsCreateInfo, SparkDsHost, SparkDsLibrary, ABI_VERSION, +}; +use spark_native_source::{ + BatchReader, BatchReaderFactory, BatchSource, Options, Result, ScanBuilder, ScanPlan, + SourceCapability, TaskContext, +}; + +static DECLARATIONS: AtomicUsize = AtomicUsize::new(0); +static SCHEMAS: AtomicUsize = AtomicUsize::new(0); +static SCANS: AtomicUsize = AtomicUsize::new(0); +static DROPS: AtomicUsize = AtomicUsize::new(0); + +struct Source { + mode: String, +} + +impl BatchSource for Source { + type Scan = Builder; + + fn new(options: Options) -> Result { + Ok(Self { + mode: options.get("capabilities").cloned().unwrap(), + }) + } + + fn capabilities(&self) -> Vec { + DECLARATIONS.fetch_add(1, Ordering::SeqCst); + match self.mode.as_str() { + "empty" => vec![], + "batch" => vec![SourceCapability::BatchRead], + "duplicates" => vec![SourceCapability::BatchRead, SourceCapability::BatchRead], + "micro_batch" => vec![SourceCapability::MicroBatchRead], + "batch_write" => vec![SourceCapability::BatchWrite], + "read_write" => vec![SourceCapability::BatchRead, SourceCapability::BatchWrite], + "streaming_write" => vec![SourceCapability::StreamingWrite], + "truncate" => vec![SourceCapability::Truncate], + "panic" => panic!("capability discovery failed"), + other => panic!("unexpected test mode: {other}"), + } + } + + fn schema(&self) -> Result { + SCHEMAS.fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(Schema::empty())) + } + + fn scan(&self, _: SchemaRef) -> Result { + SCANS.fetch_add(1, Ordering::SeqCst); + Ok(Builder) + } +} + +impl Drop for Source { + fn drop(&mut self) { + DROPS.fetch_add(1, Ordering::SeqCst); + } +} + +struct Builder; + +impl ScanBuilder for Builder { + type Factory = Factory; + + fn plan(self) -> Result> { + Ok(ScanPlan { + factory: Factory, + partitions: vec![], + }) + } +} + +#[derive(Serialize, Deserialize)] +struct Factory; + +impl BatchReaderFactory for Factory { + type Partition = (); + + fn open<'task>(&self, _: (), _: TaskContext<'task>) -> Result> { + Err("capability tests do not open partition readers".into()) + } +} + +spark_native_source::export_data_source!(Source); + +extern "C" { + fn spark_ds_init(host: *const SparkDsHost, library: *mut *const SparkDsLibrary) -> c_int; +} + +struct SourceHandle<'a> { + library: &'a SparkDsLibrary, + pointer: *mut c_void, +} + +impl Drop for SourceHandle<'_> { + fn drop(&mut self) { + unsafe { self.library.close.unwrap()(self.pointer) }; + } +} + +fn create(library: &SparkDsLibrary, mode: &str) -> (c_int, u64, *mut c_void) { + let name = CString::new("capability_test").unwrap(); + let key = CString::new("capabilities").unwrap(); + let value = CString::new(mode).unwrap(); + let keys = [key.as_ptr()]; + let values = [value.as_ptr()]; + let info = SparkDsCreateInfo { + size: size_of::() as i64, + name: name.as_ptr(), + option_keys: keys.as_ptr(), + option_values: values.as_ptr(), + num_options: 1, + }; + let mut mask = 0; + let mut source = ptr::null_mut(); + let status = unsafe { library.create.unwrap()(&info, &mut mask, &mut source) }; + (status, mask, source) +} + +fn error(library: &SparkDsLibrary) -> String { + let message = unsafe { library.last_error.unwrap()() }; + assert!(!message.is_null(), "failed C call must set last_error"); + unsafe { CStr::from_ptr(message) }.to_str().unwrap().into() +} + +#[test] +fn source_capabilities_control_the_exported_abi_and_source_lifecycle() { + // One lifecycle test owns these counters; no parallel test resets shared state. + let spark_version = CString::new("test").unwrap(); + let host = SparkDsHost { + size: size_of::() as i64, + abi_version: ABI_VERSION, + spark_version: spark_version.as_ptr(), + log: None, + }; + let mut table = ptr::null(); + assert_eq!(unsafe { spark_ds_init(&host, &mut table) }, 0); + assert!(!table.is_null()); + let library = unsafe { &*table }; + + // Analysis can ask only for capabilities, then close without schema or scan discovery. + let (status, mask, pointer) = create(library, "batch"); + assert_eq!(status, 0); + assert_eq!(mask, 1); + assert!(!pointer.is_null()); + drop(SourceHandle { library, pointer }); + assert_eq!(DECLARATIONS.load(Ordering::SeqCst), 1); + assert_eq!(SCHEMAS.load(Ordering::SeqCst), 0); + assert_eq!(SCANS.load(Ordering::SeqCst), 0); + assert_eq!(DROPS.load(Ordering::SeqCst), 1); + + // Keep all instances alive together: the last declaration must not affect earlier sources. + let sources: Vec<_> = [("empty", 0), ("batch", 1), ("duplicates", 1)] + .into_iter() + .map(|(mode, expected)| { + let (status, mask, pointer) = create(library, mode); + assert_eq!(status, 0); + assert_eq!(mask, expected, "{mode}"); + assert!(!pointer.is_null()); + (expected, SourceHandle { library, pointer }) + }) + .collect(); + assert_eq!(DECLARATIONS.load(Ordering::SeqCst), 4); + + for (mask, source) in sources { + let mut schema = ArrowSchema::empty(); + assert_eq!( + unsafe { library.schema.unwrap()(source.pointer, &mut schema) }, + 0 + ); + let scans_before = SCANS.load(Ordering::SeqCst); + let mut reader = ptr::null_mut(); + let status = + unsafe { library.create_reader.unwrap()(source.pointer, &mut schema, &mut reader) }; + if mask == 0 { + assert_ne!(status, 0); + assert_eq!(error(library), "source does not support batch reads"); + assert!(reader.is_null()); + assert_eq!(SCANS.load(Ordering::SeqCst), scans_before); + } else { + assert_eq!(status, 0); + assert!(!reader.is_null()); + assert_eq!(SCANS.load(Ordering::SeqCst), scans_before + 1); + unsafe { library.close_reader.unwrap()(reader) }; + } + assert_eq!(DECLARATIONS.load(Ordering::SeqCst), 4); + } + assert_eq!(SCHEMAS.load(Ordering::SeqCst), 3); + assert_eq!(SCANS.load(Ordering::SeqCst), 2); + assert_eq!(DROPS.load(Ordering::SeqCst), 4); + + for (mode, capability) in [ + ("micro_batch", "MicroBatchRead"), + ("batch_write", "BatchWrite"), + ("read_write", "BatchWrite"), + ("streaming_write", "StreamingWrite"), + ("truncate", "Truncate"), + ] { + let drops_before = DROPS.load(Ordering::SeqCst); + let declarations_before = DECLARATIONS.load(Ordering::SeqCst); + let (status, mask, source) = create(library, mode); + assert_ne!(status, 0, "{mode}"); + assert_eq!( + error(library), + format!("capability {capability} is not supported by this batch-read SDK") + ); + assert_eq!(mask, 0); + assert!(source.is_null()); + assert_eq!(DROPS.load(Ordering::SeqCst), drops_before + 1); + assert_eq!(DECLARATIONS.load(Ordering::SeqCst), declarations_before + 1); + } + + let drops_before = DROPS.load(Ordering::SeqCst); + let (status, mask, source) = create(library, "panic"); + assert_ne!(status, 0); + assert_eq!( + error(library), + "native provider panicked: capability discovery failed" + ); + assert_eq!(mask, 0); + assert!(source.is_null()); + assert_eq!(DROPS.load(Ordering::SeqCst), drops_before + 1); + assert_eq!(SCHEMAS.load(Ordering::SeqCst), 3); + assert_eq!(SCANS.load(Ordering::SeqCst), 2); +} From cfd1255efe5d347830262299d3ee3b4b1944e22b Mon Sep 17 00:00:00 2001 From: Gengliang Wang Date: Fri, 16 Oct 2026 21:06:39 +0000 Subject: [PATCH 4/6] Add typed batch-write APIs to the Rust native SDK --- .../src/main/rust/native-data-sink/Cargo.toml | 34 + .../src/main/rust/native-data-sink/README.md | 137 ++++ .../src/main/rust/native-data-sink/smoke.py | 127 +++ .../src/main/rust/native-data-sink/src/lib.rs | 746 +++++++++++++++++ .../main/rust/native-data-source/README.md | 9 +- .../main/rust/native-data-source/package.py | 3 +- .../main/rust/native-data-source/src/lib.rs | 16 +- sql/native/rust/Cargo.lock | 11 + sql/native/rust/Cargo.toml | 5 +- sql/native/rust/README.md | 149 +++- sql/native/rust/src/abi.rs | 96 ++- sql/native/rust/src/bridge.rs | 122 ++- sql/native/rust/src/export.rs | 188 ++++- sql/native/rust/src/lib.rs | 138 +++- sql/native/rust/src/provider.rs | 151 ++++ sql/native/rust/src/stream.rs | 4 +- sql/native/rust/src/task.rs | 12 +- sql/native/rust/src/write.rs | 330 ++++++++ sql/native/rust/tests/batch_write.rs | 766 ++++++++++++++++++ sql/native/rust/tests/capabilities.rs | 16 +- 20 files changed, 2909 insertions(+), 151 deletions(-) create mode 100644 examples/src/main/rust/native-data-sink/Cargo.toml create mode 100644 examples/src/main/rust/native-data-sink/README.md create mode 100644 examples/src/main/rust/native-data-sink/smoke.py create mode 100644 examples/src/main/rust/native-data-sink/src/lib.rs create mode 100644 sql/native/rust/src/provider.rs create mode 100644 sql/native/rust/src/write.rs create mode 100644 sql/native/rust/tests/batch_write.rs diff --git a/examples/src/main/rust/native-data-sink/Cargo.toml b/examples/src/main/rust/native-data-sink/Cargo.toml new file mode 100644 index 0000000000000..2b244a0eb96c5 --- /dev/null +++ b/examples/src/main/rust/native-data-sink/Cargo.toml @@ -0,0 +1,34 @@ +# 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. + +[package] +name = "spark-native-sink-demo" +workspace = "../../../../../sql/native/rust" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true + +[lib] +name = "spark_datasource_rust_sdk_sink" +path = "src/lib.rs" +crate-type = ["cdylib", "rlib"] + +[dependencies] +spark-native-source = { path = "../../../../../sql/native/rust" } +arrow-array.workspace = true +arrow-schema.workspace = true +serde.workspace = true +serde_json.workspace = true diff --git a/examples/src/main/rust/native-data-sink/README.md b/examples/src/main/rust/native-data-sink/README.md new file mode 100644 index 0000000000000..78384f64799ac --- /dev/null +++ b/examples/src/main/rust/native-data-sink/README.md @@ -0,0 +1,137 @@ + + +# Rust native batch sink + +This write-only example implements `DataSource`, `BatchWriteSource`, `BatchWrite`, +`BatchWriterFactory`, and `BatchWriter` from the [Rust native data source SDK][sdk]. +It registers `export_data_source!(JsonSink, batch_write, truncate)` and implements +no read traits. The SDK supplies the C ABI entry point, typed factory and commit +message serialization, Arrow batch import, capability checks, and error handling. + +The example targets C ABI version 1 from [Spark PR #59424][native-source-pr], pinned +to `534dfca15723c7208061b1e9e4d80c5757360a2a`. Running it in Spark requires that +framework and Java 22 or later on the driver and executors. The standalone Rust +tests do not need Spark or a JVM. + +## Supported data and options + +The sink accepts exactly two columns, in this order: `id BIGINT` and `label STRING`. +Both nullable and non-nullable input fields are accepted. Null values are preserved. +Task writers encode batches incrementally as JSON lines, avoiding an additional +file-format dependency in this small example. + +| Option | Default | Meaning | +| --- | --- | --- | +| `path` | Required | Absolute directory shared by the driver and every executor | +| `readonly` | `false` | Remove both batch-write and truncate capabilities when `true` | + +Unknown options are rejected. Source creation and capability discovery do not create +files. The default instance advertises only native `BatchWrite` and `Truncate`; it +does not advertise batch reads or streaming operations. Spark performs its native +capability check again with the operation's effective options. + +## Commit and abort protocol + +Each write operation holds a `.write-lock` directory under `path`, rejecting another +concurrent writer. It assigns an operation ID and stages tasks in `.staging/`. +Each task attempt exclusively creates its own file, writes Arrow batches, flushes and +syncs the file, then returns a typed message containing its operation, file, partition, +and row count. Task commit alone does not publish data. + +The driver receives the winning attempt messages. It validates those messages and +renames the entire staging directory to `objects/`, preventing new tasks +from opening files for that operation. It removes every file not named by a winning +message, including speculative attempts and attempts whose messages never arrived. +It then writes and syncs a new manifest and atomically renames it to `manifest.json`. +That one manifest replacement is the commit point: + +- Append includes the previous manifest's files plus the winning new files. +- Truncate includes only the winning new files, including an empty result. +- Previously published data files are retained, so an existing manifest snapshot + stays readable and a failed overwrite cannot delete the old data. +- Abort removes the entire operation, including files with no commit message and + a sealed directory whose manifest publication failed. It leaves the current + manifest and its data unchanged. Dropping an unfinished coordinator also performs + best-effort cleanup. Dropping a prepared task writer leaves its file for the driver. + +To inspect output, read only the files listed in `manifest.json`. Recursively reading +the destination directory would also pick up data from previous truncate operations. +For example, after a successful write: + +```python +import json +from pathlib import Path + +root = Path("/absolute/shared/path/native-sink") +manifest = json.loads((root / "manifest.json").read_text()) +files = [str(root / entry["path"]) for entry in manifest["files"]] +if files: + spark.read.schema("id LONG, label STRING").json(files).show() +``` + +## Build, package, and run + +From the repository root: + +```sh +cargo build --manifest-path sql/native/rust/Cargo.toml --workspace --locked +cargo test --manifest-path sql/native/rust/Cargo.toml -p spark-native-sink-demo --locked + +python3 examples/src/main/rust/native-data-source/package.py \ + --source rust_sdk_sink \ + --library sql/native/rust/target/debug/libspark_datasource_rust_sdk_sink.so \ + --output sql/native/rust/target/rust_sdk_sink.sparkpkg + +"$SPARK_HOME/bin/spark-submit" --master 'local[2]' \ + examples/src/main/rust/native-data-sink/smoke.py \ + --package sql/native/rust/target/rust_sdk_sink.sparkpkg +``` + +Use `libspark_datasource_rust_sdk_sink.dylib` for macOS. The smoke script uses a +temporary local directory and requires a local Spark master. It checks append, +overwrite, a failed overwrite that preserves published data and cleans staging, +and write rejection with `readonly=true`. Rust tests additionally exercise losing +attempts, missing commit messages, a driver publication failure, and lock cleanup. + +An application can write with: + +```python +spark.addArtifact("sql/native/rust/target/rust_sdk_sink.sparkpkg") +data = spark.range(10).selectExpr("id", "concat('row-', id) AS label") +data.write.format("rust_sdk_sink").mode("append").save("/absolute/shared/path/native-sink") +data.write.format("rust_sdk_sink").mode("overwrite").save("/absolute/shared/path/native-sink") +``` + +## Limits of this example + +Use a local POSIX filesystem, or a shared filesystem with atomic directory renames, +atomic file replacement, and coherent exclusive creation. The same absolute path must +refer to the same storage on all hosts. Executor-local paths in a distributed cluster, +object-store URLs, and Windows are unsupported. Only one write operation can own a +destination at a time; the lock is not a distributed lease or fencing protocol. + +This example demonstrates atomic visibility during normal execution and abort, not +crash recovery or exactly-once publication across driver retries. It syncs files but +does not implement a crash-durable metadata journal. Process or machine failure can +leave a stale `.write-lock`, staged files, or unreferenced objects. Recover the lock +only after confirming no operation still owns it. Retained files from old manifests +also need garbage collection with an appropriate reader-retention policy; this demo +does not implement that policy. + +[sdk]: ../../../../../sql/native/rust/README.md +[native-source-pr]: https://github.com/apache/spark/pull/59424 diff --git a/examples/src/main/rust/native-data-sink/smoke.py b/examples/src/main/rust/native-data-sink/smoke.py new file mode 100644 index 0000000000000..81ad7416ee5a4 --- /dev/null +++ b/examples/src/main/rust/native-data-sink/smoke.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +# 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. + +"""Run with Java 22+ and local spark-submit from a build that includes PR #59424.""" + +import argparse +import json +from pathlib import Path +import tempfile + +from pyspark.errors import AnalysisException +from pyspark.sql import SparkSession + + +def committed_rows(root): + manifest = json.loads((root / "manifest.json").read_text()) + assert manifest["version"] == 1, manifest + result = [] + for entry in manifest["files"]: + rows = [json.loads(line) for line in (root / entry["path"]).read_text().splitlines()] + assert len(rows) == entry["rows"], entry + result.extend((row["id"], row["label"]) for row in rows) + return sorted(result) + + +def assert_clean(root): + assert not (root / ".write-lock").exists(), "The operation lock was not released" + staging = root / ".staging" + assert not staging.exists() or not list(staging.iterdir()), "Staged files were not cleaned" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--package", type=Path, required=True) + args = parser.parse_args() + package = args.package.resolve(strict=True) + + spark = ( + SparkSession.builder.appName("Rust native batch sink smoke") + .config("spark.sql.adaptive.enabled", "false") + .config("spark.task.maxFailures", "1") + .getOrCreate() + ) + try: + if not spark.sparkContext.master.startswith("local"): + raise ValueError( + "This smoke script uses driver-local temporary files; use --master local" + ) + spark.addArtifact(str(package)) + + def data(start, end): + return spark.range(start, end, numPartitions=2).selectExpr( + "id", "concat('row-', id) AS label" + ) + + with tempfile.TemporaryDirectory(prefix="rust-native-sink-") as directory: + root = Path(directory) + + def write(frame, mode): + frame.write.format("rust_sdk_sink").mode(mode).save(str(root)) + assert_clean(root) + + write(data(0, 5), "append") + assert committed_rows(root) == [(i, f"row-{i}") for i in range(5)] + write(data(5, 9), "append") + assert committed_rows(root) == [(i, f"row-{i}") for i in range(9)] + + old_manifest = json.loads((root / "manifest.json").read_text()) + write(data(20, 23), "overwrite") + assert committed_rows(root) == [(i, f"row-{i}") for i in range(20, 23)] + assert all((root / entry["path"]).exists() for entry in old_manifest["files"]) + + # No input shuffle: the failure happens inside the write job, after + # the driver has created the coordinator and its staging directory. + before = (root / "manifest.json").read_bytes() + previous_operations = sorted(path.name for path in (root / "objects").iterdir()) + failing = spark.range(0, 8, numPartitions=2).selectExpr( + "id", + "if(id >= 4, cast(raise_error('intentional native sink input failure') AS STRING), " + "concat('row-', id)) AS label", + ) + try: + failing.write.format("rust_sdk_sink").mode("overwrite").save(str(root)) + except Exception as error: + assert "intentional native sink input failure" in str(error), str(error) + else: + raise AssertionError("The input error should abort the overwrite") + assert (root / "manifest.json").read_bytes() == before + assert sorted(path.name for path in (root / "objects").iterdir()) == previous_operations + assert committed_rows(root) == [(i, f"row-{i}") for i in range(20, 23)] + assert_clean(root) + + for mode in ("append", "overwrite"): + try: + ( + data(0, 1) + .write.format("rust_sdk_sink") + .mode(mode) + .option("readonly", "true") + .save(str(root)) + ) + except AnalysisException as error: + assert error.getCondition() == "UNSUPPORTED_FEATURE.TABLE_OPERATION", str(error) + else: + raise AssertionError(f"readonly=true should reject {mode}") + assert (root / "manifest.json").read_bytes() == before + assert_clean(root) + print("Rust native batch sink Spark smoke checks passed") + finally: + spark.stop() + + +if __name__ == "__main__": + main() diff --git a/examples/src/main/rust/native-data-sink/src/lib.rs b/examples/src/main/rust/native-data-sink/src/lib.rs new file mode 100644 index 0000000000000..cae8d43ad04b8 --- /dev/null +++ b/examples/src/main/rust/native-data-sink/src/lib.rs @@ -0,0 +1,746 @@ +// 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 write-only JSONL sink demonstrating task staging and atomic manifest commits. +//! +//! This example requires one local or shared POSIX filesystem. It is not an +//! object-store connector, a concurrent writer protocol, or a durable table format. + +use std::collections::BTreeSet; +use std::fs::{self, File, OpenOptions}; +use std::io::{BufWriter, ErrorKind, Write}; +use std::path::{Component, Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use arrow_array::{Array, Int64Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use serde::{Deserialize, Serialize}; +use spark_native_source::{ + BatchWrite, BatchWriteSource, BatchWriter, BatchWriterFactory, DataSource, Options, Result, + SourceCapability, TaskContext, TaskWriter, WriteMode, +}; + +spark_native_source::export_data_source!(JsonSink, batch_write, truncate); + +static NEXT_ID: AtomicU64 = AtomicU64::new(0); + +/// Construction and capability discovery do not access the filesystem. +pub struct JsonSink { + path: PathBuf, + readonly: bool, +} + +impl DataSource for JsonSink { + fn new(mut options: Options) -> Result { + let path = PathBuf::from( + options + .remove("path") + .ok_or("The path option is required")?, + ); + if !path.is_absolute() { + return Err("The path must be absolute and shared by the driver and executors".into()); + } + let readonly = match options.remove("readonly").as_deref() { + None | Some("false") => false, + Some("true") => true, + _ => return Err("The readonly option must be true or false".into()), + }; + if !options.is_empty() { + return Err( + format!("Unknown options: {:?}", options.keys().collect::>()).into(), + ); + } + Ok(Self { path, readonly }) + } + + fn schema(&self) -> Result { + Ok(sink_schema()) + } + + fn capabilities(&self, supported: &[SourceCapability]) -> Vec { + if self.readonly { + Vec::new() + } else { + supported.to_vec() + } + } +} + +impl BatchWriteSource for JsonSink { + type Write = JsonWrite; + + fn write(&self, schema: SchemaRef, mode: WriteMode) -> Result { + if self.readonly { + return Err("This sink is readonly".into()); + } + validate_schema(&schema)?; + JsonWrite::new(self.path.clone(), mode) + } +} + +fn sink_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("label", DataType::Utf8, true), + ])) +} + +fn validate_schema(schema: &Schema) -> Result<()> { + let fields = schema.fields(); + if fields.len() != 2 + || fields[0].name() != "id" + || fields[0].data_type() != &DataType::Int64 + || fields[1].name() != "label" + || fields[1].data_type() != &DataType::Utf8 + { + return Err("Expected columns id BIGINT, label STRING, in that order".into()); + } + Ok(()) +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +struct DataFile { + path: String, + rows: u64, +} + +#[derive(Debug, Deserialize, Serialize)] +struct Manifest { + version: u32, + files: Vec, +} + +impl Default for Manifest { + fn default() -> Self { + Self { + version: 1, + files: Vec::new(), + } + } +} + +fn read_manifest(root: &Path) -> Result { + let file = match File::open(root.join("manifest.json")) { + Ok(file) => file, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(Manifest::default()), + Err(error) => return Err(error.into()), + }; + let manifest: Manifest = serde_json::from_reader(file)?; + if manifest.version != 1 { + return Err("Unsupported sink manifest version".into()); + } + for file in &manifest.files { + let components = Path::new(&file.path).components().collect::>(); + if components.len() != 3 + || components[0] != Component::Normal("objects".as_ref()) + || !components.iter().all(|c| matches!(c, Component::Normal(_))) + { + return Err("Invalid sink manifest data path".into()); + } + } + Ok(manifest) +} + +/// This lock is held for the whole write operation, including all task attempts. +struct WriteLock(PathBuf); + +impl WriteLock { + fn acquire(root: &Path) -> Result { + let path = root.join(".write-lock"); + match fs::create_dir(&path) { + Ok(()) => Ok(Self(path)), + Err(error) if error.kind() == ErrorKind::AlreadyExists => { + Err("Another write owns .write-lock; concurrent writes are unsupported".into()) + } + Err(error) => Err(error.into()), + } + } +} + +impl Drop for WriteLock { + fn drop(&mut self) { + let _ = fs::remove_dir(&self.0); + } +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum WriteState { + Open, + Committed, + Aborted, +} + +pub struct JsonWrite { + root: PathBuf, + operation: String, + staging: PathBuf, + objects: PathBuf, + temporary_manifest: PathBuf, + previous: Manifest, + mode: WriteMode, + state: WriteState, + sealed: bool, + lock: Option, +} + +impl JsonWrite { + fn new(root: PathBuf, mode: WriteMode) -> Result { + if !cfg!(unix) { + return Err("This demo requires a POSIX filesystem".into()); + } + fs::create_dir_all(&root)?; + let lock = WriteLock::acquire(&root)?; + let previous = read_manifest(&root)?; + fs::create_dir_all(root.join(".staging"))?; + fs::create_dir_all(root.join("objects"))?; + let staging = unique_directory(&root.join(".staging"))?; + let operation = staging.file_name().unwrap().to_str().unwrap().to_owned(); + Ok(Self { + objects: root.join("objects").join(&operation), + temporary_manifest: root.join(format!(".manifest-{operation}.json")), + root, + operation, + staging, + previous, + mode, + state: WriteState::Open, + sealed: false, + lock: Some(lock), + }) + } + + fn publish(&mut self, messages: Vec) -> Result<()> { + let mut filenames = BTreeSet::new(); + let mut partitions = BTreeSet::new(); + for message in &messages { + if message.operation != self.operation + || !is_filename(&message.filename) + || !filenames.insert(message.filename.clone()) + { + return Err("Invalid or duplicate task commit message".into()); + } + if let Some(partition) = message.partition { + if !partitions.insert(partition) { + return Err("More than one winning attempt for the same partition".into()); + } + } + if !self.staging.join(&message.filename).is_file() { + return Err("A winning task's staged file is missing".into()); + } + } + + // Seal the entire operation directory. New attempts can no longer open + // a staged file, and existing losing files remain invisible to readers. + fs::rename(&self.staging, &self.objects)?; + self.sealed = true; + for entry in fs::read_dir(&self.objects)? { + let entry = entry?; + if !filenames.contains(&entry.file_name().to_string_lossy().into_owned()) { + fs::remove_file(entry.path())?; + } + } + + let mut files = match self.mode { + WriteMode::Append => self.previous.files.clone(), + WriteMode::Truncate => Vec::new(), + }; + files.extend(messages.into_iter().map(|message| DataFile { + path: format!("objects/{}/{}", self.operation, message.filename), + rows: message.rows, + })); + let manifest = Manifest { version: 1, files }; + let file = OpenOptions::new() + .write(true) + .create_new(true) + .open(&self.temporary_manifest)?; + let mut writer = BufWriter::new(file); + serde_json::to_writer(&mut writer, &manifest)?; + writer.write_all(b"\n")?; + writer.flush()?; + writer.get_ref().sync_all()?; + drop(writer); + + // This single rename is the commit point. No existing data is deleted, + // including during truncate. No fallible work follows publication. + fs::rename(&self.temporary_manifest, self.root.join("manifest.json"))?; + self.state = WriteState::Committed; + self.lock.take(); + Ok(()) + } + + fn cleanup(&self) -> Result<()> { + // Do not rely on commit messages: failed tasks and losing attempts may + // have staged files without returning a message to the driver. + let mut first_error = None; + for result in [ + fs::remove_dir_all(&self.staging), + if self.sealed { + fs::remove_dir_all(&self.objects) + } else { + Ok(()) + }, + fs::remove_file(&self.temporary_manifest), + ] { + if let Err(error) = result { + if error.kind() != ErrorKind::NotFound && first_error.is_none() { + first_error = Some(error); + } + } + } + match first_error { + Some(error) => Err(error.into()), + None => Ok(()), + } + } +} + +impl BatchWrite for JsonWrite { + type Factory = JsonWriterFactory; + + fn factory(&mut self) -> Result { + if self.state != WriteState::Open { + return Err("The write operation is no longer open".into()); + } + Ok(JsonWriterFactory { + staging: self.staging.clone(), + operation: self.operation.clone(), + }) + } + + fn commit(&mut self, messages: Vec) -> Result<()> { + match self.state { + WriteState::Committed => Ok(()), + WriteState::Aborted => Err("The write operation was aborted".into()), + WriteState::Open => self.publish(messages), + } + } + + fn abort(&mut self, _messages: Vec>) -> Result<()> { + if self.state == WriteState::Committed { + return Ok(()); + } + let result = self.cleanup(); + self.state = WriteState::Aborted; + self.lock.take(); + result + } +} + +impl Drop for JsonWrite { + fn drop(&mut self) { + if self.state != WriteState::Committed { + let _ = self.cleanup(); + } + } +} + +/// Serializable values only: the driver lock and open files stay on their hosts. +#[derive(Deserialize, Serialize)] +pub struct JsonWriterFactory { + staging: PathBuf, + operation: String, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct CommitMessage { + operation: String, + filename: String, + partition: Option, + rows: u64, +} + +impl BatchWriterFactory for JsonWriterFactory { + type Message = CommitMessage; + + fn open<'task>(&self, task: TaskContext<'task>) -> Result> { + if task.is_interrupted() { + return Err("Task cancelled before opening the sink".into()); + } + // Do not recreate staging: its absence means the driver sealed or + // aborted this operation. create_new also separates speculative tasks. + for _ in 0..16 { + let filename = format!( + "part-{}-attempt-{}-{}.jsonl", + task.partition_id().unwrap_or(-1), + task.task_attempt_id().unwrap_or(-1), + unique_id()? + ); + let path = self.staging.join(&filename); + match OpenOptions::new().write(true).create_new(true).open(&path) { + Ok(file) => { + return Ok(Box::new(JsonTaskWriter { + task, + path, + message: CommitMessage { + operation: self.operation.clone(), + filename, + partition: task.partition_id(), + rows: 0, + }, + writer: Some(BufWriter::new(file)), + prepared: false, + failed: false, + })); + } + Err(error) if error.kind() == ErrorKind::AlreadyExists => continue, + Err(error) => return Err(error.into()), + } + } + Err("Could not create a unique task attempt file".into()) + } +} + +#[derive(Serialize)] +struct Row<'a> { + id: Option, + label: Option<&'a str>, +} + +struct JsonTaskWriter<'task> { + task: TaskContext<'task>, + path: PathBuf, + message: CommitMessage, + writer: Option>, + prepared: bool, + failed: bool, +} + +impl JsonTaskWriter<'_> { + fn discard_buffer(&mut self) { + if let Some(writer) = self.writer.take() { + // BufWriter normally flushes on drop; abort discards buffered rows. + let (file, _) = writer.into_parts(); + drop(file); + } + } + + fn write_batch(&mut self, batch: RecordBatch) -> Result<()> { + if self.prepared || self.failed { + return Err("The task writer is no longer accepting batches".into()); + } + validate_schema(batch.schema().as_ref())?; + let writer = self.writer.as_mut().ok_or("The task writer is closed")?; + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let labels = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for index in 0..batch.num_rows() { + if index % 1024 == 0 && self.task.is_interrupted() { + return Err("Task cancelled while writing the sink".into()); + } + let row = Row { + id: (!ids.is_null(index)).then(|| ids.value(index)), + label: (!labels.is_null(index)).then(|| labels.value(index)), + }; + serde_json::to_writer(&mut *writer, &row)?; + writer.write_all(b"\n")?; + } + self.message.rows = self + .message + .rows + .checked_add(batch.num_rows() as u64) + .ok_or("Task row count overflow")?; + Ok(()) + } + + fn prepare(&mut self) -> Result { + if self.failed { + return Err("Cannot commit a failed task writer".into()); + } + if self.prepared { + return Ok(self.message.clone()); + } + if self.task.is_interrupted() { + return Err("Task cancelled before committing the sink".into()); + } + let mut writer = self.writer.take().ok_or("The task writer is closed")?; + writer.flush()?; + writer.get_ref().sync_all()?; + drop(writer); + self.prepared = true; + Ok(self.message.clone()) + } +} + +impl BatchWriter for JsonTaskWriter<'_> { + type Message = CommitMessage; + + fn write(&mut self, batch: RecordBatch) -> Result<()> { + let result = self.write_batch(batch); + self.failed |= result.is_err(); + result + } + + fn commit(&mut self) -> Result { + let result = self.prepare(); + self.failed |= result.is_err(); + result + } + + fn abort(&mut self) -> Result<()> { + self.discard_buffer(); + self.prepared = false; + self.failed = true; + match fs::remove_file(&self.path) { + Ok(()) => Ok(()), + Err(error) if error.kind() == ErrorKind::NotFound => Ok(()), + Err(error) => Err(error.into()), + } + } +} + +impl Drop for JsonTaskWriter<'_> { + fn drop(&mut self) { + if !self.prepared { + self.discard_buffer(); + let _ = fs::remove_file(&self.path); + } + } +} + +fn is_filename(name: &str) -> bool { + let mut components = Path::new(name).components(); + matches!(components.next(), Some(Component::Normal(_))) && components.next().is_none() +} + +fn unique_id() -> Result { + Ok(format!( + "{}-{}-{}", + std::process::id(), + SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos(), + NEXT_ID.fetch_add(1, Ordering::Relaxed) + )) +} + +fn unique_directory(parent: &Path) -> Result { + for _ in 0..16 { + let path = parent.join(unique_id()?); + match fs::create_dir(&path) { + Ok(()) => return Ok(path), + Err(error) if error.kind() == ErrorKind::AlreadyExists => continue, + Err(error) => return Err(error.into()), + } + } + Err("Could not create a unique write operation directory".into()) +} + +#[cfg(test)] +mod tests { + use super::*; + + struct TempDirectory(PathBuf); + + impl TempDirectory { + fn new() -> Self { + Self(unique_directory(&std::env::temp_dir()).unwrap()) + } + + fn source(&self) -> JsonSink { + JsonSink::new(Options::from([( + "path".to_owned(), + self.0.to_str().unwrap().to_owned(), + )])) + .unwrap() + } + } + + impl Drop for TempDirectory { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } + } + + fn batch(ids: Vec>) -> RecordBatch { + let labels = ids + .iter() + .map(|id| id.map(|id| format!("row-{id}"))) + .collect::>(); + RecordBatch::try_new( + sink_schema(), + vec![ + Arc::new(Int64Array::from(ids)), + Arc::new(StringArray::from(labels)), + ], + ) + .unwrap() + } + + fn stage(factory: &JsonWriterFactory, ids: Vec>) -> CommitMessage { + // Executors reconstruct the factory without the source or driver write. + let factory: JsonWriterFactory = + serde_json::from_slice(&serde_json::to_vec(factory).unwrap()).unwrap(); + let mut writer = factory.open(TaskContext::detached()).unwrap(); + writer.write(batch(ids)).unwrap(); + let message = writer.commit().unwrap(); + serde_json::from_slice(&serde_json::to_vec(&message).unwrap()).unwrap() + } + + fn committed_rows(path: &Path) -> Vec { + read_manifest(path) + .unwrap() + .files + .iter() + .flat_map(|file| { + let text = fs::read_to_string(path.join(&file.path)).unwrap(); + let rows = text + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect::>(); + assert_eq!(rows.len() as u64, file.rows); + rows + }) + .collect() + } + + fn commit_ids(source: &JsonSink, mode: WriteMode, ids: Vec>) { + let mut write = source.write(sink_schema(), mode).unwrap(); + let message = stage(&write.factory().unwrap(), ids); + write.commit(vec![message]).unwrap(); + } + + #[test] + fn append_publishes_all_committed_rows_and_preserves_nulls() { + let directory = TempDirectory::new(); + let source = directory.source(); + commit_ids(&source, WriteMode::Append, vec![Some(1), None]); + commit_ids(&source, WriteMode::Append, vec![Some(2)]); + assert_eq!( + committed_rows(&directory.0), + vec![ + serde_json::json!({"id": 1, "label": "row-1"}), + serde_json::json!({"id": null, "label": null}), + serde_json::json!({"id": 2, "label": "row-2"}), + ] + ); + assert_eq!(read_manifest(&directory.0).unwrap().files.len(), 2); + assert!(!directory.0.join(".write-lock").exists()); + } + + #[test] + fn truncate_replaces_manifest_but_retains_previously_published_files() { + let directory = TempDirectory::new(); + let source = directory.source(); + commit_ids(&source, WriteMode::Append, vec![Some(1)]); + let previous = read_manifest(&directory.0).unwrap(); + commit_ids(&source, WriteMode::Truncate, vec![Some(2)]); + assert_eq!( + committed_rows(&directory.0), + vec![serde_json::json!({"id": 2, "label": "row-2"})] + ); + assert!(directory.0.join(&previous.files[0].path).is_file()); + commit_ids(&source, WriteMode::Truncate, Vec::new()); + assert!(committed_rows(&directory.0).is_empty()); + } + + #[test] + fn driver_publishes_only_winning_attempts_and_removes_unreported_files() { + let directory = TempDirectory::new(); + let mut write = directory + .source() + .write(sink_schema(), WriteMode::Append) + .unwrap(); + let factory = write.factory().unwrap(); + let winning = stage(&factory, vec![Some(1)]); + let _losing = stage(&factory, vec![Some(2)]); + let _missing_message = stage(&factory, vec![Some(3)]); + assert_eq!(fs::read_dir(&write.staging).unwrap().count(), 3); + write.commit(vec![winning]).unwrap(); + assert!(!write.staging.exists()); + assert_eq!(fs::read_dir(&write.objects).unwrap().count(), 1); + assert_eq!( + committed_rows(&directory.0), + vec![serde_json::json!({"id": 1, "label": "row-1"})] + ); + assert!(factory.open(TaskContext::detached()).is_err()); + } + + #[test] + fn abort_cleans_the_operation_even_with_missing_messages() { + let directory = TempDirectory::new(); + let source = directory.source(); + commit_ids(&source, WriteMode::Append, vec![Some(1)]); + let before = fs::read(directory.0.join("manifest.json")).unwrap(); + let mut write = source.write(sink_schema(), WriteMode::Truncate).unwrap(); + let factory = write.factory().unwrap(); + let message = stage(&factory, vec![Some(2)]); + let _missing = stage(&factory, vec![Some(3)]); + let mut incomplete = factory.open(TaskContext::detached()).unwrap(); + incomplete.write(batch(vec![Some(4)])).unwrap(); + write.abort(vec![Some(message), None]).unwrap(); + assert!(!write.staging.exists()); + assert!(!write.objects.exists()); + assert!(!directory.0.join(".write-lock").exists()); + assert_eq!(fs::read(directory.0.join("manifest.json")).unwrap(), before); + assert_eq!(committed_rows(&directory.0).len(), 1); + assert!(factory.open(TaskContext::detached()).is_err()); + } + + #[test] + fn failed_publication_keeps_previous_manifest_and_abort_removes_sealed_files() { + let directory = TempDirectory::new(); + let source = directory.source(); + commit_ids(&source, WriteMode::Append, vec![Some(1)]); + let before = fs::read(directory.0.join("manifest.json")).unwrap(); + let mut write = source.write(sink_schema(), WriteMode::Truncate).unwrap(); + let message = stage(&write.factory().unwrap(), vec![Some(2)]); + // Force create_new to fail after the operation directory has been sealed. + fs::write(&write.temporary_manifest, b"publication failure").unwrap(); + assert!(write.commit(vec![message]).is_err()); + assert!(write.objects.exists()); + assert_eq!(fs::read(directory.0.join("manifest.json")).unwrap(), before); + write.abort(vec![None]).unwrap(); + assert!(!write.objects.exists()); + assert!(!write.temporary_manifest.exists()); + assert_eq!(committed_rows(&directory.0).len(), 1); + } + + #[test] + fn concurrent_write_is_rejected_and_drop_releases_staging_and_lock() { + let directory = TempDirectory::new(); + let source = directory.source(); + let mut write = source.write(sink_schema(), WriteMode::Append).unwrap(); + let _message = stage(&write.factory().unwrap(), vec![Some(1)]); + let staging = write.staging.clone(); + assert!(source.write(sink_schema(), WriteMode::Append).is_err()); + drop(write); + assert!(!staging.exists()); + assert!(!directory.0.join(".write-lock").exists()); + assert!(source.write(sink_schema(), WriteMode::Append).is_ok()); + } + + #[test] + fn readonly_removes_all_write_capabilities_without_creating_files() { + let directory = TempDirectory::new(); + let path = directory.0.join("not-created"); + let source = JsonSink::new(Options::from([ + ("path".to_owned(), path.to_str().unwrap().to_owned()), + ("readonly".to_owned(), "true".to_owned()), + ])) + .unwrap(); + assert!(source + .capabilities(&[SourceCapability::BatchWrite, SourceCapability::Truncate]) + .is_empty()); + assert!(source.write(sink_schema(), WriteMode::Append).is_err()); + assert!(!path.exists()); + } +} diff --git a/examples/src/main/rust/native-data-source/README.md b/examples/src/main/rust/native-data-source/README.md index 6811ef60c86a2..3b2b789582f3f 100644 --- a/examples/src/main/rust/native-data-source/README.md +++ b/examples/src/main/rust/native-data-source/README.md @@ -28,10 +28,11 @@ Running the Spark smoke script requires a distribution built with that native da framework and Java 22 or later on the driver and executors. The JNI SDK prototype in PR #59442 targets a different interface. -The source explicitly implements `BatchSource::capabilities()` to return -`vec![SourceCapability::BatchRead]`, producing the native mask `SPARK_DS_BATCH_READ` (`1`). -Connectors can override this method based on their options; the SDK validates the declaration -against its implemented callbacks. This example does not implement writes or Structured Streaming. +Registering the source for batch reads derives `SourceCapability::BatchRead` and the native +mask `SPARK_DS_BATCH_READ` (`1`); the connector does not maintain a separate capability list. +Connectors can restrict registered operations through `DataSource::capabilities()` based on +options. This example does not implement writes or Structured Streaming. The companion +[file sink](../native-data-sink/README.md) demonstrates batch append and overwrite. During analysis, Spark creates and closes a source with the table options to inspect its capabilities. It exposes `BATCH_WRITE` and `TRUNCATE` to Data Source V2 only when declared, so this example rejects batch append and overwrite before evaluating their input. diff --git a/examples/src/main/rust/native-data-source/package.py b/examples/src/main/rust/native-data-source/package.py index 09019d3943458..36d440291fa51 100644 --- a/examples/src/main/rust/native-data-source/package.py +++ b/examples/src/main/rust/native-data-source/package.py @@ -36,6 +36,7 @@ def current_platform(): def main(): parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", default="rust_sdk_range", help="Registered data source name") parser.add_argument("--library", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--platform", default=current_platform(), help="Override for cross builds") @@ -56,7 +57,7 @@ def main(): library_path = f"{args.platform}/{args.library.name}" manifest = { "abiVersion": 1, - "dataSources": ["rust_sdk_range"], + "dataSources": [args.source], "libraries": {args.platform: library_path}, } args.output.parent.mkdir(parents=True, exist_ok=True) diff --git a/examples/src/main/rust/native-data-source/src/lib.rs b/examples/src/main/rust/native-data-source/src/lib.rs index dbca41207ab7a..e29910f3fd0e8 100644 --- a/examples/src/main/rust/native-data-source/src/lib.rs +++ b/examples/src/main/rust/native-data-source/src/lib.rs @@ -24,8 +24,8 @@ use arrow_array::{ArrayRef, Int64Array, RecordBatch, RecordBatchIterator, Record use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; use serde::{Deserialize, Serialize}; use spark_native_source::{ - BatchReader, BatchReaderFactory, BatchSource, FilterSupport, LimitSupport, Options, Predicate, - Result, ScanBuilder, ScanPlan, SourceCapability, TaskContext, + BatchReader, BatchReaderFactory, BatchSource, DataSource, FilterSupport, LimitSupport, Options, + Predicate, Result, ScanBuilder, ScanPlan, TaskContext, }; pub struct RangeSource { @@ -56,9 +56,7 @@ pub struct RangePartition { end: i64, } -impl BatchSource for RangeSource { - type Scan = RangeScan; - +impl DataSource for RangeSource { fn new(options: Options) -> Result { for key in options.keys() { if !matches!( @@ -86,13 +84,13 @@ impl BatchSource for RangeSource { }) } - fn capabilities(&self) -> Vec { - vec![SourceCapability::BatchRead] - } - fn schema(&self) -> Result { output_schema(&["id".into(), "label".into()]) } +} + +impl BatchSource for RangeSource { + type Scan = RangeScan; fn scan(&self, schema: SchemaRef) -> Result { let columns: Vec = schema.fields().iter().map(|f| f.name().clone()).collect(); diff --git a/sql/native/rust/Cargo.lock b/sql/native/rust/Cargo.lock index 2e411f09d9e22..889509f9c4d6a 100644 --- a/sql/native/rust/Cargo.lock +++ b/sql/native/rust/Cargo.lock @@ -458,6 +458,17 @@ dependencies = [ "spark-native-source", ] +[[package]] +name = "spark-native-sink-demo" +version = "0.1.0" +dependencies = [ + "arrow-array", + "arrow-schema", + "serde", + "serde_json", + "spark-native-source", +] + [[package]] name = "spark-native-source" version = "0.1.0" diff --git a/sql/native/rust/Cargo.toml b/sql/native/rust/Cargo.toml index 6b325226ba0a0..5568b681d65ee 100644 --- a/sql/native/rust/Cargo.toml +++ b/sql/native/rust/Cargo.toml @@ -22,7 +22,10 @@ publish.workspace = true description = "Experimental Rust batch data source SDK for Apache Spark PR #59424" [workspace] -members = ["../../../examples/src/main/rust/native-data-source"] +members = [ + "../../../examples/src/main/rust/native-data-source", + "../../../examples/src/main/rust/native-data-sink", +] default-members = ["."] resolver = "2" diff --git a/sql/native/rust/README.md b/sql/native/rust/README.md index 471874866a062..bb7eaa2746ad9 100644 --- a/sql/native/rust/README.md +++ b/sql/native/rust/README.md @@ -17,7 +17,7 @@ limitations under the License. # Rust native data source SDK prototype -This experimental SDK implements **batch reads** over the C ABI proposed in +This experimental SDK implements **batch reads and writes** over the C ABI proposed in [Spark PR #59424](https://github.com/apache/spark/pull/59424), pinned to commit `534dfca15723c7208061b1e9e4d80c5757360a2a`. Spark calls the library through the Java Foreign Function and Memory (FFM) API. This PR contains the SDK and example, @@ -30,22 +30,28 @@ 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. Both crates share this directory's lockfile and -`target` directory. Cargo commands select the SDK by default; use `--workspace` +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 example. ## Connector API and lifecycle -A connector implements three traits from [`src/lib.rs`](src/lib.rs): +A connector implements [`DataSource`](src/lib.rs) for construction, schema discovery, +and optional capability restrictions. Spark can create and immediately drop a source +just to discover capabilities; schema inference or an operation need not follow. -* `BatchSource`: parse driver options, declare capabilities, discover the schema, - and create an independent scan builder. Spark can also create and drop a source - just to discover its capabilities; schema inference or a scan need not follow. +Read support uses three traits: + +* `BatchSource`: create an independent scan builder. * `ScanBuilder`: negotiate predicates, limit, and projection, then consume itself in `plan()` to return a typed `ScanPlan`. * `BatchReaderFactory`: serialize executor configuration and open one lazy Arrow reader from a typed partition and a borrowed `TaskContext`. +Write support is independent: a write-only connector implements `BatchWriteSource`, +`BatchWrite`, `BatchWriterFactory`, and `BatchWriter`, without implementing reads. +The write lifecycle is described below. + 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 @@ -103,49 +109,109 @@ state in `bridge.rs`, Arrow ownership in `stream.rs`, and C bindings/adapters in ## Source capabilities -Connector developers declare native capabilities through -`BatchSource::capabilities(&self) -> Vec`. The default is -`vec![SourceCapability::BatchRead]`; override it to describe the source instance -created with its effective options. For example, inside a source implementation: +Registration determines which operation traits the connector implements. The SDK +installs their native callbacks and derives the default capability set: + +```rust +export_data_source!(RangeSource, batch_read); +export_data_source!(FileSink, batch_write, truncate); +export_data_source!(ReadWriteSource, batch_read, batch_write); +``` + +These are alternative registrations; invoke the macro only once per library. +The one-argument form is shorthand for `batch_read`. `truncate` is optional and +requires `batch_write`; without it the source supports only append writes. +Unregistered operation families have null native callbacks. The registration +requires the corresponding Rust traits at compile time. + +Most connectors do not implement `capabilities`. A connector can restrict the +registered set for the current options by overriding this `DataSource` method: ```rust -fn capabilities(&self) -> Vec { - if self.batch_reads_enabled { - vec![SourceCapability::BatchRead] - } else { - vec![] - } +fn capabilities(&self, supported: &[SourceCapability]) -> Vec { + supported.iter().copied().filter(|capability| { + !self.read_only || !matches!(capability, + SourceCapability::BatchWrite | SourceCapability::Truncate) + }).collect() } ``` -The SDK evaluates this method once during `create`, validates the declaration, -and returns the resulting native flag mask to Spark. Duplicate entries are -ignored. An empty set still allows schema discovery and closing the source; -batch scan creation fails before calling the connector's `scan` method. The -declaration stays fixed for that source instance. - -`SourceCapability` names all current native operations: `BatchRead`, -`MicroBatchRead`, `BatchWrite`, `StreamingWrite`, and `Truncate`. This batch-read -prototype accepts only `BatchRead`; declaring other variants fails source -creation with an explicit error because their callbacks are not implemented. -Unsupported function slots remain null. These native flags describe supported -operations; query-level predicate, limit, and projection results are separate +The SDK evaluates this once during `create`, validates the declaration, and returns +its native flag mask. Duplicate entries are ignored. Enabling an unregistered +operation is an error; `Truncate` also requires `BatchWrite`. An empty set permits +schema discovery and close, but rejects operation creation before calling the +connector. The declaration stays fixed for that instance. Streaming capabilities +are not supported yet. Predicate, limit, and projection support 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 `BatchSource::new` and `capabilities` lightweight and safe to drop +Keep `DataSource::new` and `capabilities` lightweight and safe to drop immediately; do not assume that `schema` or `scan` will follow. Executor resources belong in `BatchReaderFactory::open`. The pinned framework propagates native `BATCH_WRITE` and `TRUNCATE` flags to -`Table.capabilities()`, so a batch-only source rejects batch writes during +`Table.capabilities()`, 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. +## Batch write lifecycle + +```text +Driver: + DataSource + BatchWriteSource -> BatchWrite coordinator + | factory() | commit(messages) / abort(messages) + v ^ + serialized factory + schema | typed messages + | | +Executor, per task attempt: v | + BatchWriterFactory::open(TaskContext) + -> BatchWriter::write(RecordBatch)* + -> commit() / abort() -> Drop +``` + +`BatchWriteSource::write(schema, mode)` validates the input and returns an +independent driver coordinator. `BatchWrite::factory()` produces serializable +executor configuration; the SDK caches the encoded factory and input schema, +including a terminal error if creation or serialization fails. Restoring an +executor writer uses only those bytes, without creating a driver registry. + +Each `BatchWriterFactory::open` returns a task-bound writer. `write` takes ownership +of a `RecordBatch`, so its Arrow buffers can be retained after the callback. The +SDK moves the incoming Arrow array and schema, checks the schema against the +negotiated input, 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. Check cancellation during lengthy writes and use I/O timeouts. + +Task `commit()` returns a typed, serde-serializable message. Driver `commit` +receives one winning message per partition; driver `abort` receives +`Vec>`, preserving missing-task positions. A failed task or driver +commit leaves the same object available for abort. The SDK prevents further writes +after a failure or successful commit, caches successful task commit messages, and +avoids repeating successful commit/abort callbacks. 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. + +Spark may close a coordinator without commit/abort if factory serialization fails. +It also creates writers late, after prerequisite AQE shuffle stages, but input +computation can still fail during write tasks. Constructors must therefore avoid +irreversible changes. This API is batch-only: non-null streaming query IDs and +epoch IDs other than `-1` are rejected. + +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. + ## Arrow contract `BatchReaderFactory::open` returns a lifetime-bound `Box; @@ -294,6 +296,59 @@ pub unsafe fn options(info: *const SparkDsCreateInfo) -> Result { .collect()) } +/// # Safety +/// `info` must follow the create_writer contract of the C header. +pub unsafe fn write_mode(info: *const SparkDsWriteInfo) -> Result { + let size = unsafe { struct_size(info)? }; + let pointer = info.cast(); + let mode: i32 = unsafe { field(pointer, size, offset_of!(SparkDsWriteInfo, mode)) } + .ok_or("write info does not provide a mode")?; + let query_id: *const c_char = + unsafe { field(pointer, size, offset_of!(SparkDsWriteInfo, query_id)) } + .ok_or("write info does not provide a query ID")?; + if !query_id.is_null() { + return Err("streaming writes are not supported".into()); + } + match mode { + 0 => Ok(WriteMode::Append), + 1 => Ok(WriteMode::Truncate), + _ => Err(format!("unsupported native write mode: {mode}").into()), + } +} + +pub fn batch_epoch(epoch: i64) -> Result<()> { + if epoch != -1 { + return Err("batch writes require epoch ID -1".into()); + } + Ok(()) +} + +/// # Safety +/// Both arrays must contain count readable values. Non-null message pointers +/// must address their lengths of bytes and remain valid for the returned borrow. +pub unsafe fn messages<'a>( + pointers: *const *const u8, + lengths: *const i64, + count: i64, +) -> Result>> { + let pointers = unsafe { slice(pointers, count)? }; + let lengths = unsafe { slice(lengths, count)? }; + pointers + .iter() + .zip(lengths) + .map(|(&pointer, &length)| { + if pointer.is_null() { + if length != 0 { + return Err("null commit message must have zero length".into()); + } + Ok(None) + } else { + Ok(Some(unsafe { slice(pointer, length)? })) + } + }) + .collect() +} + pub fn predicate_use(support: FilterSupport) -> i32 { match support { FilterSupport::Unsupported => 0, @@ -359,4 +414,43 @@ mod tests { let info = (&prefix as *const i64).cast::(); assert!(unsafe { options(info) }.is_err()); } + + #[test] + fn write_info_rejects_truncated_unknown_and_streaming_requests() { + let prefix = size_of::() as i64; + let truncated = (&prefix as *const i64).cast::(); + assert!(unsafe { write_mode(truncated) }.is_err()); + let mut info = SparkDsWriteInfo { + size: size_of::() as i64, + mode: 0, + query_id: std::ptr::null(), + }; + assert_eq!(unsafe { write_mode(&info) }.unwrap(), WriteMode::Append); + info.mode = 1; + assert_eq!(unsafe { write_mode(&info) }.unwrap(), WriteMode::Truncate); + info.mode = 2; + assert!(unsafe { write_mode(&info) }.is_err()); + info.mode = 0; + info.query_id = c"stream".as_ptr(); + assert!(unsafe { write_mode(&info) }.is_err()); + assert!(batch_epoch(-1).is_ok()); + assert!(batch_epoch(0).is_err()); + } + + #[test] + fn abort_messages_distinguish_missing_and_empty_messages() { + let bytes = [7_u8, 8]; + let pointers = [std::ptr::null(), bytes.as_ptr(), bytes.as_ptr()]; + let lengths = [0, 0, 2]; + let values = unsafe { messages(pointers.as_ptr(), lengths.as_ptr(), 3) }.unwrap(); + assert_eq!(values, vec![None, Some(&[][..]), Some(&bytes[..])]); + assert!(unsafe { messages(std::ptr::null(), std::ptr::null(), 0) } + .unwrap() + .is_empty()); + assert!(unsafe { messages(std::ptr::null(), std::ptr::null(), -1) }.is_err()); + let invalid_lengths = [1, 0, 2]; + assert!(unsafe { messages(pointers.as_ptr(), invalid_lengths.as_ptr(), 3) }.is_err()); + let invalid_lengths = [0, -1, 2]; + assert!(unsafe { messages(pointers.as_ptr(), invalid_lengths.as_ptr(), 3) }.is_err()); + } } diff --git a/sql/native/rust/src/bridge.rs b/sql/native/rust/src/bridge.rs index 4b52c36b87da8..cb807126f8ff7 100644 --- a/sql/native/rust/src/bridge.rs +++ b/sql/native/rust/src/bridge.rs @@ -26,12 +26,15 @@ use arrow_schema::{Schema, SchemaRef}; use serde::{de::DeserializeOwned, Deserialize, Serialize}; use crate::{ - BatchReaderFactory, BatchSource, FilterSupport, LimitSupport, Options, Predicate, Result, - ScanBuilder, SourceCapability, TaskContext, + BatchReaderFactory, DataSource, FilterSupport, LimitSupport, Options, Predicate, Result, + ScanBuilder, SourceCapability, TaskContext, WriteMode, }; +use crate::provider::Provider; +use crate::write::DriverWriter; + #[path = "stream.rs"] -mod stream; +pub(crate) mod stream; const STATE_VERSION: u32 = 1; @@ -39,23 +42,25 @@ type HandleMap = Mutex>>>; /// One registry per loaded plugin. Provider code runs under a per-handle lock, /// so independent sources and scans do not serialize behind a global lock. -pub struct Bridge { +pub struct Bridge { next_handle: AtomicI64, - sources: HandleMap>, - planners: HandleMap>, + sources: HandleMap>, + planners: HandleMap>, + writers: HandleMap>, } -impl Default for Bridge { +impl Default for Bridge

{ fn default() -> Self { Self { next_handle: AtomicI64::new(1), sources: Mutex::new(HashMap::new()), planners: Mutex::new(HashMap::new()), + writers: Mutex::new(HashMap::new()), } } } -impl Bridge { +impl Bridge

{ fn insert(&self, handles: &HandleMap, value: T) -> Result { let handle = self .next_handle @@ -66,21 +71,24 @@ impl Bridge { } pub fn create(&self, options: Options) -> Result<(i64, u64)> { - self.register_source(S::new(options)?) + self.register_source(P::Source::new(options)?) } - fn register_source(&self, source: S) -> Result<(i64, u64)> { + fn register_source(&self, source: P::Source) -> Result<(i64, u64)> { let mut capabilities = 0; - for capability in source.capabilities() { - match capability { - SourceCapability::BatchRead => capabilities |= crate::abi::BATCH_READ, - _ => { - return Err(format!( - "capability {capability:?} is not supported by this batch-read SDK" - ) - .into()); - } + for capability in source.capabilities(P::CAPABILITIES) { + if !P::CAPABILITIES.contains(&capability) { + return Err(format!("capability {capability:?} is not registered").into()); } + capabilities |= match capability { + SourceCapability::BatchRead => crate::abi::BATCH_READ, + SourceCapability::BatchWrite => crate::abi::BATCH_WRITE, + SourceCapability::Truncate => crate::abi::TRUNCATE, + _ => return Err(format!("capability {capability:?} is not implemented").into()), + }; + } + if capabilities & crate::abi::TRUNCATE != 0 && capabilities & crate::abi::BATCH_WRITE == 0 { + return Err("Truncate requires BatchWrite".into()); } let handle = self.insert( &self.sources, @@ -107,7 +115,7 @@ impl Bridge { if source.capabilities & crate::abi::BATCH_READ == 0 { return Err("source does not support batch reads".into()); } - source.source.scan(schema.clone())? + P::scan(&source.source, schema.clone())? }; self.insert(&self.planners, Planner::new(builder, schema)) } @@ -164,6 +172,45 @@ impl Bridge { result } + pub fn create_writer(&self, handle: i64, schema: SchemaRef, mode: WriteMode) -> Result { + stream::validate_schema(&schema)?; + let source = get(&self.sources, handle)?; + let writer = { + let source = lock(&source); + if source.capabilities & crate::abi::BATCH_WRITE == 0 { + return Err("source does not support batch writes".into()); + } + if mode == WriteMode::Truncate && source.capabilities & crate::abi::TRUNCATE == 0 { + return Err("source does not support truncate".into()); + } + P::write(&source.source, schema.clone(), mode)? + }; + self.insert(&self.writers, DriverWriter::new(writer, schema)) + } + + pub fn serialized_writer(&self, handle: i64) -> Result>> { + let writer = get(&self.writers, handle)?; + let result = lock(&writer).serialized(); + result + } + + pub fn commit(&self, handle: i64, messages: &[Option<&[u8]>]) -> Result<()> { + let writer = get(&self.writers, handle)?; + let result = lock(&writer).commit(messages); + result + } + + pub fn abort(&self, handle: i64, messages: &[Option<&[u8]>]) -> Result<()> { + let writer = get(&self.writers, handle)?; + let result = lock(&writer).abort(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(()) @@ -314,14 +361,14 @@ struct Envelope { value: T, } -fn encode(value: &T) -> Result> { +pub(crate) fn encode(value: &T) -> Result> { Ok(serde_json::to_vec(&Envelope { version: STATE_VERSION, value, })?) } -fn decode(bytes: &[u8]) -> Result { +pub(crate) fn decode(bytes: &[u8]) -> Result { let envelope: Envelope = serde_json::from_slice(bytes)?; if envelope.version != STATE_VERSION { return Err(format!("unsupported native state version: {}", envelope.version).into()); @@ -329,7 +376,9 @@ fn decode(bytes: &[u8]) -> Result { Ok(envelope.value) } -fn provider_call(action: impl FnOnce() -> Result) -> std::result::Result { +pub(crate) fn provider_call( + action: impl FnOnce() -> Result, +) -> std::result::Result { catch_unwind(AssertUnwindSafe(|| { action().map_err(|error| error.to_string()) })) @@ -365,7 +414,8 @@ mod tests { use arrow_schema::{DataType, Field}; use super::*; - use crate::{BatchReader, ScanPlan}; + use crate::provider::ReadOnly; + use crate::{BatchReader, BatchSource, ScanPlan}; #[derive(Clone, Copy)] enum FailureMode { @@ -416,9 +466,7 @@ mod tests { } } - impl BatchSource for Source { - type Scan = Builder; - + impl DataSource for Source { fn new(_: Options) -> Result { Ok(Self { builds: Arc::default(), @@ -434,6 +482,10 @@ mod tests { fn schema(&self) -> Result { Ok(Arc::new(Schema::empty())) } + } + + impl BatchSource for Source { + type Scan = Builder; fn scan(&self, _: SchemaRef) -> Result { Ok(Builder { @@ -496,7 +548,7 @@ mod tests { #[test] fn source_can_close_before_planning_and_plan_is_built_once() { - let bridge = Bridge::::default(); + let bridge = Bridge::>::default(); let source = Source::new(Options::new()).unwrap(); let builds = source.builds.clone(); let drops = source.drops.clone(); @@ -528,7 +580,7 @@ mod tests { #[test] fn residual_filters_prevent_calling_the_providers_limit_method() { for support in [FilterSupport::PruningOnly, FilterSupport::Unsupported] { - let bridge = Bridge::::default(); + let bridge = Bridge::>::default(); let mut source = Source::new(Options::new()).unwrap(); source.filter_support = support; source.limit_support = LimitSupport::Global; @@ -551,7 +603,7 @@ mod tests { #[test] fn only_global_limits_return_true_to_the_abi() { for support in [LimitSupport::PerPartition, LimitSupport::Global] { - let bridge = Bridge::::default(); + let bridge = Bridge::>::default(); let mut source = Source::new(Options::new()).unwrap(); source.limit_support = support; let limits = source.limits.clone(); @@ -576,7 +628,7 @@ mod tests { "native provider panicked: planning failed", ), ] { - let bridge = Bridge::::default(); + let bridge = Bridge::>::default(); let mut source = Source::new(Options::new()).unwrap(); source.failure = failure; let builds = source.builds.clone(); @@ -606,7 +658,7 @@ mod tests { "native provider panicked: negotiation failed", ), ] { - let bridge = Bridge::::default(); + let bridge = Bridge::>::default(); let mut source = Source::new(Options::new()).unwrap(); source.failure = failure; let builds = source.builds.clone(); @@ -631,7 +683,7 @@ mod tests { #[test] fn scans_of_one_source_keep_pushdown_and_cached_plans_independent() { - let bridge = Bridge::::default(); + let bridge = Bridge::>::default(); let mut source = Source::new(Options::new()).unwrap(); source.filter_support = FilterSupport::PruningOnly; source.limit_support = LimitSupport::Global; @@ -677,7 +729,7 @@ mod tests { #[test] fn executor_retries_use_only_serialized_values_after_the_driver_is_dropped() { - let bridge = Bridge::::default(); + let bridge = Bridge::>::default(); let source = Source::new(Options::new()).unwrap(); let builds = source.builds.clone(); let source_drops = source.drops.clone(); @@ -725,7 +777,7 @@ mod tests { #[test] fn executor_rejects_a_factory_schema_different_from_the_negotiated_schema() { - let bridge = Bridge::::default(); + let bridge = Bridge::>::default(); let (handle, _) = bridge.create(Options::new()).unwrap(); let reader = bridge .create_reader(handle, bridge.schema(handle).unwrap()) diff --git a/sql/native/rust/src/export.rs b/sql/native/rust/src/export.rs index 4d5e19001aebc..fc3b7c9963330 100644 --- a/sql/native/rust/src/export.rs +++ b/sql/native/rust/src/export.rs @@ -13,23 +13,53 @@ // See the License for the specific language governing permissions and // limitations under the License. -/// Exports one batch-read provider through version 1 of Spark's C data source ABI. +/// Exports a provider through version 1 of Spark's C data source ABI. /// /// Call this macro once in a cdylib. The only exported symbol is spark_ds_init; /// providers implement the safe Rust traits and do not handle C pointers. +/// Register `batch_read`, `batch_write`, or both (in that order). Append `truncate` +/// to opt in to overwrite support. The one-argument form registers batch reads. +/// Registration generates the default capabilities and installs matching callbacks; +/// unsupported operation families have null callbacks. #[macro_export] macro_rules! export_data_source { ($source:ty) => { + $crate::export_data_source!($source, batch_read); + }; + ($source:ty, batch_read) => { + $crate::export_data_source!(@provider $crate::provider::ReadOnly<$source>); + }; + ($source:ty, batch_write) => { + $crate::export_data_source!(@provider $crate::provider::WriteOnly<$source, false>); + }; + ($source:ty, batch_write, truncate) => { + $crate::export_data_source!(@provider $crate::provider::WriteOnly<$source, true>); + }; + ($source:ty, batch_read, batch_write) => { + $crate::export_data_source!(@provider $crate::provider::ReadWrite<$source, false>); + }; + ($source:ty, batch_read, batch_write, truncate) => { + $crate::export_data_source!(@provider $crate::provider::ReadWrite<$source, true>); + }; + (@provider $provider:ty) => { mod __spark_native_datasource_exports { use super::*; use std::ffi::{c_char, c_int, c_void}; use $crate::abi::{self, ArrowArrayStream, ArrowSchema, SparkDsLibrary, SparkDsSink}; use $crate::bridge::{open_reader, Bridge}; - type Factory = <<$source as $crate::BatchSource>::Scan as $crate::ScanBuilder>::Factory; + use $crate::provider::Provider as _; - fn bridge() -> &'static Bridge<$source> { - static BRIDGE: std::sync::OnceLock> = std::sync::OnceLock::new(); + type Registration = $provider; + type Factory = <::Scan + as $crate::ScanBuilder>::Factory; + type WriterFactory = <::Write + as $crate::BatchWrite>::Factory; + type Message = ::Message; + + fn bridge() -> &'static Bridge { + static BRIDGE: std::sync::OnceLock> = + std::sync::OnceLock::new(); BRIDGE.get_or_init(Bridge::default) } @@ -39,32 +69,40 @@ macro_rules! export_data_source { last_error: Some(abi::last_error), create: Some(create), schema: Some(schema), - create_reader: Some(create_reader), + create_reader: if Registration::READ { Some(create_reader) } else { None }, create_stream_reader: None, - create_writer: None, + create_writer: if Registration::WRITE { Some(create_writer) } else { None }, close: Some(close), - push_predicates: Some(push_predicates), - push_limit: Some(push_limit), - prune_columns: Some(prune_columns), - partitions: Some(partitions), - serialize_reader: Some(serialize_reader), - close_reader: Some(close_reader), + push_predicates: if Registration::READ { Some(push_predicates) } else { None }, + push_limit: if Registration::READ { Some(push_limit) } else { None }, + prune_columns: if Registration::READ { Some(prune_columns) } else { None }, + partitions: if Registration::READ { Some(partitions) } else { None }, + serialize_reader: if Registration::READ { Some(serialize_reader) } else { None }, + close_reader: if Registration::READ { Some(close_reader) } else { None }, initial_offset: None, latest_offset: None, stream_partitions: None, serialize_stream_reader: None, commit_offset: None, close_stream_reader: None, - read: Some(read), - serialize_writer: None, - commit: None, - abort: None, - close_writer: None, - create_data_writer: None, - write: None, - commit_data_writer: None, - abort_data_writer: None, - close_data_writer: None, + read: if Registration::READ { Some(read) } else { None }, + serialize_writer: if Registration::WRITE { Some(serialize_writer) } else { None }, + commit: if Registration::WRITE { Some(commit) } else { None }, + abort: if Registration::WRITE { Some(abort) } else { None }, + close_writer: if Registration::WRITE { Some(close_writer) } else { None }, + create_data_writer: if Registration::WRITE { + Some(create_data_writer) + } else { + None + }, + write: if Registration::WRITE { Some(write) } else { None }, + commit_data_writer: if Registration::WRITE { + Some(commit_data_writer) + } else { + None + }, + abort_data_writer: if Registration::WRITE { Some(abort_data_writer) } else { None }, + close_data_writer: if Registration::WRITE { Some(close_data_writer) } else { None }, }; #[export_name = "spark_ds_init"] @@ -214,6 +252,112 @@ macro_rules! export_data_source { Ok(()) }) } + + unsafe extern "C" fn create_writer( + source: *mut c_void, + schema: *mut ArrowSchema, + info: *const abi::SparkDsWriteInfo, + writer: *mut *mut c_void, + ) -> c_int { + abi::call(|| { + abi::require_output(writer)?; + let mode = unsafe { abi::write_mode(info)? }; + let schema = unsafe { abi::import_schema(schema)? }; + let handle = bridge().create_writer(source as usize as i64, schema, mode)?; + unsafe { writer.write(handle as usize as *mut c_void) }; + Ok(()) + }) + } + + unsafe extern "C" fn serialize_writer( + writer: *mut c_void, + sink: SparkDsSink, + data: *mut c_void, + ) -> c_int { + abi::call(|| { + let state = bridge().serialized_writer(writer as usize as i64)?; + unsafe { abi::emit(sink, data, &state) } + }) + } + + unsafe extern "C" fn commit( + writer: *mut c_void, + epoch: i64, + messages: *const *const u8, + lengths: *const i64, + count: i64, + ) -> c_int { + abi::call(|| { + abi::batch_epoch(epoch)?; + let messages = unsafe { abi::messages(messages, lengths, count)? }; + bridge().commit(writer as usize as i64, &messages) + }) + } + + unsafe extern "C" fn abort( + writer: *mut c_void, + epoch: i64, + messages: *const *const u8, + lengths: *const i64, + count: i64, + ) -> c_int { + abi::call(|| { + abi::batch_epoch(epoch)?; + let messages = unsafe { abi::messages(messages, lengths, count)? }; + bridge().abort(writer as usize as i64, &messages) + }) + } + + unsafe extern "C" fn close_writer(writer: *mut c_void) { + let _ = abi::call(|| bridge().close_writer(writer as usize as i64)); + } + + unsafe extern "C" fn create_data_writer( + task: *const abi::SparkDsTaskContext, + state: *const u8, + length: i64, + epoch: i64, + writer: *mut *mut c_void, + ) -> c_int { + abi::call(|| { + abi::require_output(writer)?; + abi::batch_epoch(epoch)?; + let state = unsafe { abi::slice(state, length)? }; + let handle = unsafe { $crate::write::open::(task, state)? }; + unsafe { writer.write(handle) }; + Ok(()) + }) + } + + unsafe extern "C" fn write( + writer: *mut c_void, + array: *mut abi::ArrowArray, + schema: *mut ArrowSchema, + ) -> c_int { + abi::call(|| unsafe { + $crate::write::write_batch::(writer, array, schema) + }) + } + + unsafe extern "C" fn commit_data_writer( + writer: *mut c_void, + sink: SparkDsSink, + data: *mut c_void, + ) -> c_int { + abi::call(|| { + sink.ok_or("null result sink")?; + let message = unsafe { $crate::write::commit_task::(writer)? }; + unsafe { abi::emit(sink, data, &message) } + }) + } + + unsafe extern "C" fn abort_data_writer(writer: *mut c_void) -> c_int { + abi::call(|| unsafe { $crate::write::abort_task::(writer) }) + } + + unsafe extern "C" fn close_data_writer(writer: *mut c_void) { + let _ = abi::call(|| unsafe { $crate::write::close_task::(writer) }); + } } }; } diff --git a/sql/native/rust/src/lib.rs b/sql/native/rust/src/lib.rs index d1a61c7776025..3a73fd9273b87 100644 --- a/sql/native/rust/src/lib.rs +++ b/sql/native/rust/src/lib.rs @@ -13,19 +13,19 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Experimental batch-read SDK for the C ABI proposed in Spark PR #59424. +//! Experimental batch data source SDK for the C ABI proposed in Spark PR #59424. //! -//! Implement [`BatchSource`], [`ScanBuilder`] and [`BatchReaderFactory`], then -//! invoke [`export_data_source!`] once in a `cdylib` crate. The SDK owns handles, -//! plan serialization, C ABI adaptation, and Arrow C Stream export. Rust traits +//! Implement [`DataSource`] and the batch-read or batch-write traits, 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. use std::collections::BTreeMap; -use arrow_array::RecordBatchReader; +use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::SchemaRef; use serde::de::DeserializeOwned; -use serde::Serialize; +use serde::{Deserialize, Serialize}; #[doc(hidden)] pub mod abi; @@ -33,7 +33,11 @@ pub mod abi; pub mod bridge; mod export; mod predicate; +#[doc(hidden)] +pub mod provider; mod task; +#[doc(hidden)] +pub mod write; pub use predicate::Predicate; pub use task::TaskContext; @@ -42,13 +46,14 @@ pub type Error = Box; pub type Result = std::result::Result; pub type Options = BTreeMap; pub type BatchReader<'task> = Box; +pub type TaskWriter<'task, M> = Box + 'task>; /// Native operations supported by a source instance with its effective options. /// /// These are native capabilities, not Spark's Data Source V2 routing flags or -/// per-scan pushdown results. This batch-read SDK currently accepts only -/// [`BatchRead`](Self::BatchRead); declaring another capability fails source -/// creation because the corresponding callbacks are not implemented. +/// per-scan pushdown results. Registration determines the available operations; +/// a source instance may disable them but cannot enable unregistered operations. +/// Streaming operations are not implemented by this SDK yet. #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[non_exhaustive] pub enum SourceCapability { @@ -93,37 +98,114 @@ pub struct ScanPlan { pub partitions: Vec, } -/// Driver-side source configuration and schema discovery. +/// Driver-side source configuration, schema discovery, and capability restrictions. /// -/// Override [`capabilities`](Self::capabilities) to declare the operations this -/// instance supports. Writes and streaming remain absent from the native -/// function table. Spark also creates sources to discover table capabilities -/// during query analysis; such a source can be dropped without calling `schema` -/// or `scan`. -pub trait BatchSource: Send + Sized + 'static { - type Scan: ScanBuilder; - - /// Spark passes option keys in lower case. Keep construction lightweight: - /// a capability probe uses only table options, while a scan may override - /// them with query options. Open executor resources in the factory. +/// Spark may create a source only to inspect its capabilities, then drop it. +/// Keep construction lightweight. Readers and writers must outlive their source. +pub trait DataSource: Send + Sized + 'static { + /// Spark passes option keys in lower case. A capability probe uses table + /// options; an operation may override them with query options. fn new(options: Options) -> Result; - /// Declares this instance's capabilities, which may depend on its options. - /// The SDK evaluates this once during source creation, validates the set - /// against its implemented callbacks, and keeps it fixed for that instance. - /// Duplicate entries are ignored. An empty set permits schema discovery but - /// rejects batch scan creation without calling `scan`. - fn capabilities(&self) -> Vec { - vec![SourceCapability::BatchRead] + /// By default, enable every registered operation. Override this to restrict + /// the current instance, for example for a read-only destination. The SDK + /// evaluates this once, ignores duplicates, and rejects unregistered flags. + /// Truncate requires BatchWrite. An empty set still permits schema discovery. + fn capabilities(&self, supported: &[SourceCapability]) -> Vec { + supported.to_vec() } fn schema(&self) -> Result; +} + +/// Driver-side entry point for batch reads. +/// +/// Register with `export_data_source!(Source, batch_read)` (or its one-argument +/// shorthand). Metadata and construction belong to [`DataSource`]. +pub trait BatchSource: DataSource { + type Scan: ScanBuilder; /// The schema may be supplied by the user; reject unsupported schemas. /// The returned builder must remain valid after this source is dropped. fn scan(&self, schema: SchemaRef) -> Result; } +/// How a batch write changes the destination. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum WriteMode { + Append, + /// Replace existing data at driver commit, never at writer creation. + Truncate, +} + +/// Driver-side entry point for batch writes, independent of read support. +/// +/// Register `batch_write` and, if implemented, `truncate`. The SDK derives native +/// capabilities and callbacks from registration, then applies instance restrictions. +pub trait BatchWriteSource: DataSource { + type Write: BatchWrite; + + /// Validate the input schema and create an independent write coordinator. + /// Construction may be followed only by Drop, without commit or abort. + fn write(&self, schema: SchemaRef, mode: WriteMode) -> Result; +} + +/// A driver coordinator for one batch write. +/// +/// Keep irreversible publication, including removing overwritten data, in commit. +/// Drop releases resources; it must not publish data or perform fallible flushing. +/// Failed commit is followed by abort on the same coordinator. +pub trait BatchWrite: Send + 'static { + type Factory: BatchWriterFactory; + + /// Return serializable executor configuration. The SDK caches the encoded + /// factory and negotiated schema, including serialization failures. + /// Include a unique operation ID when needed; Spark supplies no batch query ID. + fn factory(&mut self) -> Result; + + /// 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<::Message>, + ) -> Result<()>; + + /// Abort after input, task, or driver commit failure. None preserves the + /// position of a partition whose commit message did not reach the driver. + /// Executor loss can leave orphaned staged data; plan for its cleanup. + fn abort( + &mut self, + messages: Vec::Message>>, + ) -> Result<()>; +} + +/// Serializable configuration for independent task-attempt writers. +/// +/// Open connections and files in open, not in the factory. A returned writer +/// cannot borrow this temporary factory, but may retain the task context until +/// close. Callbacks can move between threads sequentially; writers must be Send. +pub trait BatchWriterFactory: Serialize + DeserializeOwned + Send + 'static { + type Message: Serialize + DeserializeOwned + Send + 'static; + + fn open<'task>(&self, task: TaskContext<'task>) -> Result>; +} + +/// One task attempt's writer. The SDK closes it after commit or abort. +pub trait BatchWriter: Send { + type Message; + + /// Takes ownership of an imported Arrow batch, which may be retained after + /// this call. Process batches incrementally and check task cancellation. + fn write(&mut self, batch: RecordBatch) -> Result<()>; + + /// Finish staging and return a typed message for the driver. Do not make the + /// whole write visible here. On failure (including message serialization), + /// the SDK keeps this writer available for abort. + fn commit(&mut self) -> Result; + + fn abort(&mut self) -> Result<()>; +} + /// Serializable configuration for creating independent executor readers. /// /// This type carries the configuration or snapshot descriptor; connections and diff --git a/sql/native/rust/src/provider.rs b/sql/native/rust/src/provider.rs new file mode 100644 index 0000000000000..895dad7f35135 --- /dev/null +++ b/sql/native/rust/src/provider.rs @@ -0,0 +1,151 @@ +// 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. + +//! Compile-time registration of supported callback families. + +use std::marker::PhantomData; + +use arrow_schema::SchemaRef; +use serde::{Deserialize, Serialize}; + +use crate::{ + BatchReader, BatchReaderFactory, BatchSource, BatchWrite, BatchWriteSource, BatchWriterFactory, + DataSource, Result, ScanBuilder, ScanPlan, SourceCapability, TaskContext, TaskWriter, + WriteMode, +}; + +pub trait Provider: Send + 'static { + type Source: DataSource; + type Scan: ScanBuilder; + type Write: BatchWrite; + const READ: bool; + const WRITE: bool; + const CAPABILITIES: &'static [SourceCapability]; + + fn scan(source: &Self::Source, schema: SchemaRef) -> Result; + fn write(source: &Self::Source, schema: SchemaRef, mode: WriteMode) -> Result; +} + +pub struct ReadOnly(PhantomData); +pub struct WriteOnly(PhantomData); +pub struct ReadWrite(PhantomData); + +impl Provider for ReadOnly { + type Source = S; + type Scan = S::Scan; + type Write = Unavailable; + const READ: bool = true; + const WRITE: bool = false; + const CAPABILITIES: &'static [SourceCapability] = &[SourceCapability::BatchRead]; + + fn scan(source: &S, schema: SchemaRef) -> Result { + source.scan(schema) + } + + fn write(_: &S, _: SchemaRef, _: WriteMode) -> Result { + Err("batch writes are not registered".into()) + } +} + +impl Provider for WriteOnly { + type Source = S; + type Scan = Unavailable; + type Write = S::Write; + const READ: bool = false; + const WRITE: bool = true; + const CAPABILITIES: &'static [SourceCapability] = if TRUNCATE { + &[SourceCapability::BatchWrite, SourceCapability::Truncate] + } else { + &[SourceCapability::BatchWrite] + }; + + fn scan(_: &S, _: SchemaRef) -> Result { + Err("batch reads are not registered".into()) + } + + fn write(source: &S, schema: SchemaRef, mode: WriteMode) -> Result { + source.write(schema, mode) + } +} + +impl Provider for ReadWrite { + type Source = S; + type Scan = S::Scan; + type Write = S::Write; + const READ: bool = true; + const WRITE: bool = true; + const CAPABILITIES: &'static [SourceCapability] = if TRUNCATE { + &[ + SourceCapability::BatchRead, + SourceCapability::BatchWrite, + SourceCapability::Truncate, + ] + } else { + &[SourceCapability::BatchRead, SourceCapability::BatchWrite] + }; + + fn scan(source: &S, schema: SchemaRef) -> Result { + source.scan(schema) + } + + fn write(source: &S, schema: SchemaRef, mode: WriteMode) -> Result { + source.write(schema, mode) + } +} + +// Keeps the adapter generic without requiring connectors to implement operations +// they do not support. Corresponding C function slots are always null. +#[derive(Serialize, Deserialize)] +pub struct Unavailable; + +impl ScanBuilder for Unavailable { + type Factory = Self; + + fn plan(self) -> Result> { + Err("batch reads are not registered".into()) + } +} + +impl BatchReaderFactory for Unavailable { + type Partition = (); + + fn open<'task>(&self, _: (), _: TaskContext<'task>) -> Result> { + Err("batch reads are not registered".into()) + } +} + +impl BatchWrite for Unavailable { + type Factory = Self; + + fn factory(&mut self) -> Result { + Err("batch writes are not registered".into()) + } + + fn commit(&mut self, _: Vec<()>) -> Result<()> { + Err("batch writes are not registered".into()) + } + + fn abort(&mut self, _: Vec>) -> Result<()> { + Err("batch writes are not registered".into()) + } +} + +impl BatchWriterFactory for Unavailable { + type Message = (); + + fn open<'task>(&self, _: TaskContext<'task>) -> Result> { + Err("batch writes are not registered".into()) + } +} diff --git a/sql/native/rust/src/stream.rs b/sql/native/rust/src/stream.rs index 8a0958dc307e9..c760f3e8b17ff 100644 --- a/sql/native/rust/src/stream.rs +++ b/sql/native/rust/src/stream.rs @@ -160,12 +160,12 @@ unsafe extern "C" fn release(stream: *mut FFI_ArrowArrayStream) { // Compare Spark logical types, ignoring Arrow nullability/metadata and the // interchangeable physical encodings of strings and binary data. Spark also // distinguishes only the presence of a time zone, not its spelling. -pub(super) fn validate_schema(schema: &SchemaRef) -> Result<()> { +pub(crate) fn validate_schema(schema: &SchemaRef) -> Result<()> { canonical_fields(schema.fields())?; Ok(()) } -fn canonical_fields(fields: &Fields) -> Result { +pub(crate) fn canonical_fields(fields: &Fields) -> Result { fields .iter() .map(|field| { diff --git a/sql/native/rust/src/task.rs b/sql/native/rust/src/task.rs index 2252f82c8e32a..5aab58ce8aa6b 100644 --- a/sql/native/rust/src/task.rs +++ b/sql/native/rust/src/task.rs @@ -19,10 +19,11 @@ use std::mem::offset_of; use crate::abi::{self, Interrupted, SparkDsTaskContext}; use crate::Result; -/// Task identifiers and cooperative cancellation, borrowed until its Arrow stream is released. +/// Task identifiers and cooperative cancellation, borrowed until its reader or writer is closed. /// -/// A reader may retain this value and check it during long operations. It cannot retain it -/// beyond the lifetime of its returned reader. Blocking I/O still needs its own timeouts. +/// A reader or writer may retain this value and check it during long operations. +/// It cannot retain it beyond the lifetime of its returned reader or writer. +/// Blocking I/O still needs its own timeouts. /// /// ```compile_fail /// use spark_native_source::TaskContext; @@ -49,7 +50,7 @@ unsafe impl Send for TaskContext<'_> {} unsafe impl Sync for TaskContext<'_> {} impl TaskContext<'static> { - /// A context without Spark task metadata, for testing a reader directly. + /// A context without Spark task metadata, for testing a reader or writer directly. pub fn detached() -> Self { Self { pointer: std::ptr::null(), @@ -69,7 +70,8 @@ impl<'task> TaskContext<'task> { /// /// # Safety /// The host must keep the context and its callbacks valid for `'task`, including calls from - /// other threads. In the Spark ABI this lasts until the exported Arrow stream is released. + /// other threads. Spark keeps it alive until the Arrow stream is released or + /// data writer closes. #[doc(hidden)] pub unsafe fn from_raw(pointer: *const SparkDsTaskContext) -> Result { let size = unsafe { abi::struct_size(pointer)? }; diff --git a/sql/native/rust/src/write.rs b/sql/native/rust/src/write.rs new file mode 100644 index 0000000000000..486db70b3207b --- /dev/null +++ b/sql/native/rust/src/write.rs @@ -0,0 +1,330 @@ +// 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 write coordination and task-local writers restored from serialized state. + +use std::ffi::c_void; +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::sync::Arc; + +use arrow_array::{ffi, Array, RecordBatch, StructArray}; +use arrow_schema::{DataType, Schema, SchemaRef}; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; + +use crate::abi::{self, ArrowArray, ArrowSchema, SparkDsTaskContext}; +use crate::bridge::stream::{canonical_fields, validate_schema}; +use crate::bridge::{decode, encode, panic_message, provider_call}; +use crate::{BatchWrite, BatchWriterFactory, Result, TaskContext, TaskWriter}; + +type Cached = std::result::Result; + +enum Phase { + Active, + Committed(T), + Failed(String), + Aborted(Cached<()>), +} + +impl Phase { + fn active(&self) -> Result<()> { + match self { + Self::Active => Ok(()), + Self::Failed(message) => Err(message.clone().into()), + Self::Committed(_) => Err("writer was already committed".into()), + Self::Aborted(_) => Err("writer was already aborted".into()), + } + } +} + +fn cached(result: &Cached) -> Result { + result.clone().map_err(Into::into) +} + +#[derive(Serialize, Deserialize)] +struct WriterState { + factory: F, + schema: Schema, +} + +/// The driver keeps this coordinator until close, including after a failed commit. +pub struct DriverWriter { + writer: Option, + schema: SchemaRef, + serialized: Option>>>, + phase: Phase<()>, +} + +impl DriverWriter { + pub fn new(writer: W, schema: SchemaRef) -> Self { + Self { + writer: Some(writer), + schema, + serialized: None, + phase: Phase::Active, + } + } + + pub fn serialized(&mut self) -> Result>> { + if let Some(result) = &self.serialized { + return cached(result); + } + self.phase.active()?; + let result = provider_call(|| { + validate_schema(&self.schema)?; + let factory = self.writer.as_mut().expect("writer is present").factory()?; + Ok(Arc::new(encode(&WriterState { + factory, + schema: self.schema.as_ref().clone(), + })?)) + }); + if let Err(message) = &result { + self.phase = Phase::Failed(message.clone()); + } + self.serialized = Some(result); + cached( + self.serialized + .as_ref() + .expect("serialization result is present"), + ) + } + + pub fn commit(&mut self, messages: &[Option<&[u8]>]) -> Result<()> { + if matches!(self.phase, Phase::Committed(())) { + return Ok(()); + } + self.phase.active()?; + let result = provider_call(|| { + let messages = messages + .iter() + .map(|message| decode(message.ok_or("missing committed task message")?)) + .collect::>>()?; + self.writer + .as_mut() + .expect("writer is present") + .commit(messages) + }); + self.phase = match &result { + Ok(()) => Phase::Committed(()), + Err(message) => Phase::Failed(message.clone()), + }; + cached(&result) + } + + pub fn abort(&mut self, messages: &[Option<&[u8]>]) -> Result<()> { + match &self.phase { + Phase::Committed(()) => return Ok(()), + Phase::Aborted(result) => return cached(result), + Phase::Active | Phase::Failed(_) => {} + } + let result = provider_call(|| { + let messages = messages + .iter() + .map(|message| message.map(decode).transpose()) + .collect::>>()?; + self.writer + .as_mut() + .expect("writer is present") + .abort(messages) + }); + self.phase = Phase::Aborted(result.clone()); + cached(&result) + } +} + +impl Drop for DriverWriter { + fn drop(&mut self) { + if let Some(writer) = self.writer.take() { + drop_provider(writer); + } + } +} + +struct TaskState<'task, M> { + writer: Option>, + schema: SchemaRef, + phase: Phase>>, +} + +impl Drop for TaskState<'_, M> { + fn drop(&mut self) { + if let Some(writer) = self.writer.take() { + drop_provider(writer); + } + } +} + +fn drop_provider(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); + } +} + +/// Restores a task writer without creating or consulting any driver registry. +/// +/// # Safety +/// `task` must satisfy TaskContext::from_raw and remain valid until close_task. +/// The returned handle must only be used with callbacks for F::Message, never +/// concurrently, and must be closed exactly once while the plugin stays loaded. +pub unsafe fn open( + task: *const SparkDsTaskContext, + state: &[u8], +) -> Result<*mut c_void> { + let WriterState { factory, schema }: WriterState = decode(state)?; + let schema = Arc::new(schema); + validate_schema(&schema)?; + let task = unsafe { TaskContext::from_raw(task)? }; + let (writer, open_error) = match provider_call(|| factory.open(task)) { + Ok(writer) => (Some(writer), None), + Err(message) => (None, Some(message)), + }; + let state = Box::new(TaskState { + writer, + schema, + phase: Phase::Active, + }); + // Keep guarded writer destruction in place before dropping the temporary + // factory, whose destructor is connector code too. + let factory_drop = provider_call(|| { + drop(factory); + Ok(()) + }); + if let Some(message) = open_error { + return Err(message.into()); + } + factory_drop.map_err(crate::Error::from)?; + Ok(Box::into_raw(state).cast()) +} + +/// # Safety +/// `handle` must be a live handle from open for M, accessed without concurrency. +unsafe fn task_state<'a, M>(handle: *mut c_void) -> Result<&'a mut TaskState<'a, M>> { + abi::require_output(handle)?; + Ok(unsafe { &mut *handle.cast::>() }) +} + +/// Imports owned Arrow data and passes it to the task writer. +/// +/// # Safety +/// `handle` must follow open's contract. Both input structs must be valid Arrow +/// C data that this call may move; their buffers must satisfy the Arrow schema. +pub unsafe fn write_batch( + handle: *mut c_void, + array: *mut ArrowArray, + schema: *mut ArrowSchema, +) -> Result<()> +where + M: Serialize + DeserializeOwned + Send + 'static, +{ + let state = unsafe { task_state::(handle)? }; + state.phase.active()?; + let result = provider_call(|| { + let batch = unsafe { import_batch(array, schema, &state.schema)? }; + state + .writer + .as_mut() + .expect("writer is present") + .write(batch) + }); + if let Err(message) = &result { + state.phase = Phase::Failed(message.clone()); + } + cached(&result) +} + +unsafe fn import_batch( + array: *mut ArrowArray, + schema: *mut ArrowSchema, + expected: &SchemaRef, +) -> Result { + abi::require_output(array)?; + abi::require_output(schema)?; + let array = unsafe { ArrowArray::from_raw(array) }; + let schema = unsafe { ArrowSchema::from_raw(schema) }; + let data = unsafe { ffi::from_ffi(array, &schema)? }; + if !matches!(data.data_type(), DataType::Struct(_)) { + return Err("write input must be an Arrow struct array".into()); + } + let array = StructArray::from(data); + if array.null_count() != 0 { + return Err("write input cannot contain null rows".into()); + } + let batch = RecordBatch::from(array); + if canonical_fields(batch.schema().fields())? != canonical_fields(expected.fields())? { + return Err("write batch schema does not match the planned schema".into()); + } + Ok(batch) +} + +/// Commits a task once and caches its serialized message, including failures. +/// +/// # Safety +/// `handle` must follow open's contract for M. +pub unsafe fn commit_task(handle: *mut c_void) -> Result>> +where + M: Serialize + DeserializeOwned + Send + 'static, +{ + let state = unsafe { task_state::(handle)? }; + if let Phase::Committed(message) = &state.phase { + return Ok(message.clone()); + } + state.phase.active()?; + let result = provider_call(|| { + let message = state.writer.as_mut().expect("writer is present").commit()?; + Ok(Arc::new(encode(&message)?)) + }); + match result { + Ok(message) => { + state.phase = Phase::Committed(message.clone()); + Ok(message) + } + Err(message) => { + state.phase = Phase::Failed(message.clone()); + Err(message.into()) + } + } +} + +/// Aborts an unfinished task once, including after write or commit failures. +/// +/// # Safety +/// `handle` must follow open's contract for M. +pub unsafe fn abort_task(handle: *mut c_void) -> Result<()> +where + M: Serialize + DeserializeOwned + Send + 'static, +{ + let state = unsafe { task_state::(handle)? }; + match &state.phase { + Phase::Committed(_) => return Ok(()), + Phase::Aborted(result) => return cached(result), + Phase::Active | Phase::Failed(_) => {} + } + let result = provider_call(|| state.writer.as_mut().expect("writer is present").abort()); + state.phase = Phase::Aborted(result.clone()); + cached(&result) +} + +/// Releases a task writer without committing or aborting it implicitly. +/// +/// # Safety +/// `handle` must follow open's contract for M and must not be used again. +pub unsafe fn close_task(handle: *mut c_void) -> Result<()> +where + M: Serialize + DeserializeOwned + Send + 'static, +{ + abi::require_output(handle)?; + drop(unsafe { Box::from_raw(handle.cast::>()) }); + Ok(()) +} diff --git a/sql/native/rust/tests/batch_write.rs b/sql/native/rust/tests/batch_write.rs new file mode 100644 index 0000000000000..eafc4ed68c900 --- /dev/null +++ b/sql/native/rust/tests/batch_write.rs @@ -0,0 +1,766 @@ +// 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. + +use std::ffi::{c_int, c_void, CStr, CString}; +use std::mem::size_of; +use std::ptr; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; + +use arrow_array::ffi::to_ffi; +use arrow_array::{Array, Int32Array, Int64Array, RecordBatch, StructArray}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use serde::{Deserialize, Serialize, Serializer}; +use spark_native_source::abi::{ + ArrowSchema, SparkDsCreateInfo, SparkDsHost, SparkDsLibrary, SparkDsTaskContext, + SparkDsWriteInfo, ABI_VERSION, +}; +use spark_native_source::{ + BatchWrite, BatchWriteSource, BatchWriter, BatchWriterFactory, DataSource, Options, Result, + SourceCapability, TaskContext, TaskWriter, WriteMode, +}; + +// Each test exercises the actual exported table; this lock isolates fixture observations. +static TEST_LOCK: Mutex<()> = Mutex::new(()); +static EVENTS: Mutex> = Mutex::new(Vec::new()); + +#[derive(Clone, Debug, PartialEq)] +enum Event { + WriteCreated(bool), + Factory, + Batch(usize), + TaskCommit, + TaskAbort(bool), + TaskDropped(bool), + DriverCommit(Vec), + DriverAbort(Vec>), +} + +fn record(event: Event) { + EVENTS.lock().unwrap().push(event); +} + +fn events() -> Vec { + EVENTS.lock().unwrap().clone() +} + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])) +} + +struct Source { + failure: String, + capabilities: String, +} + +impl DataSource for Source { + fn new(options: Options) -> Result { + Ok(Self { + failure: options.get("failure").cloned().unwrap_or_default(), + capabilities: options.get("capabilities").cloned().unwrap_or_default(), + }) + } + + fn schema(&self) -> Result { + Ok(schema()) + } + + fn capabilities(&self, supported: &[SourceCapability]) -> Vec { + match self.capabilities.as_str() { + "none" => vec![], + "append" => vec![SourceCapability::BatchWrite], + _ => supported.to_vec(), + } + } +} + +impl BatchWriteSource for Source { + type Write = Driver; + + fn write(&self, input: SchemaRef, mode: WriteMode) -> Result { + if input != schema() { + return Err("unsupported input schema".into()); + } + record(Event::WriteCreated(matches!(mode, WriteMode::Truncate))); + Ok(Driver { + failure: self.failure.clone(), + }) + } +} + +struct Driver { + failure: String, +} + +impl BatchWrite for Driver { + type Factory = Factory; + + fn factory(&mut self) -> Result { + record(Event::Factory); + Ok(Factory { + failure: self.failure.clone(), + }) + } + + fn commit(&mut self, messages: Vec) -> Result<()> { + record(Event::DriverCommit(messages)); + match self.failure.as_str() { + "driver_commit_error" => Err("driver commit failed".into()), + "driver_commit_panic" => panic!("driver commit panicked"), + _ => Ok(()), + } + } + + fn abort(&mut self, messages: Vec>) -> Result<()> { + record(Event::DriverAbort(messages)); + Ok(()) + } +} + +#[derive(Serialize, Deserialize)] +struct Factory { + #[serde(serialize_with = "serialize_factory_mode")] + failure: String, +} + +fn serialize_factory_mode( + mode: &str, + serializer: S, +) -> std::result::Result { + if mode == "factory_serialization" { + return Err(serde::ser::Error::custom("factory serialization failed")); + } + serializer.serialize_str(mode) +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +struct Message { + partition: i32, + attempt: i64, + values: Vec, + batches: usize, + #[serde(serialize_with = "serialize_message_mode")] + failure: String, +} + +fn serialize_message_mode( + mode: &str, + serializer: S, +) -> std::result::Result { + if mode == "message_serialization" { + return Err(serde::ser::Error::custom("message serialization failed")); + } + serializer.serialize_str(mode) +} + +impl BatchWriterFactory for Factory { + type Message = Message; + + fn open<'task>(&self, task: TaskContext<'task>) -> Result> { + Ok(Box::new(Writer { + task, + failure: self.failure.clone(), + batches: vec![], + })) + } +} + +struct Writer<'task> { + task: TaskContext<'task>, + failure: String, + batches: Vec, +} + +impl BatchWriter for Writer<'_> { + type Message = Message; + + fn write(&mut self, batch: RecordBatch) -> Result<()> { + record(Event::Batch(batch.num_rows())); + // Retain the imported buffers after the ABI call returns and its inputs are dropped. + self.batches.push(batch); + match self.failure.as_str() { + "write_error" => Err("writing batch failed".into()), + "write_panic" => panic!("writing batch panicked"), + _ => Ok(()), + } + } + + fn commit(&mut self) -> Result { + record(Event::TaskCommit); + if self.task.is_interrupted() { + return Err("task cancelled during commit".into()); + } + match self.failure.as_str() { + "task_commit_error" => return Err("task commit failed".into()), + "task_commit_panic" => panic!("task commit panicked"), + _ => (), + } + let values = self + .batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + Ok(Message { + partition: self.task.partition_id().unwrap(), + attempt: self.task.task_attempt_id().unwrap(), + values, + batches: self.batches.len(), + failure: self.failure.clone(), + }) + } + + fn abort(&mut self) -> Result<()> { + record(Event::TaskAbort(self.task.is_interrupted())); + self.batches.clear(); + Ok(()) + } +} + +impl Drop for Writer<'_> { + fn drop(&mut self) { + record(Event::TaskDropped(self.task.is_interrupted())); + } +} + +spark_native_source::export_data_source!(Source, batch_write, truncate); + +extern "C" { + fn spark_ds_init(host: *const SparkDsHost, library: *mut *const SparkDsLibrary) -> c_int; +} + +fn library() -> &'static SparkDsLibrary { + let host = SparkDsHost { + size: size_of::() as i64, + abi_version: ABI_VERSION, + spark_version: c"test".as_ptr(), + log: None, + }; + let mut table = ptr::null(); + assert_eq!(unsafe { spark_ds_init(&host, &mut table) }, 0); + assert!(!table.is_null()); + unsafe { &*table } +} + +fn error(library: &SparkDsLibrary) -> String { + let message = unsafe { library.last_error.unwrap()() }; + assert!(!message.is_null()); + unsafe { CStr::from_ptr(message) }.to_str().unwrap().into() +} + +fn success(library: &SparkDsLibrary, status: c_int) { + assert_eq!(status, 0, "{}", error(library)); +} + +struct Handle { + library: &'static SparkDsLibrary, + pointer: *mut c_void, + close: unsafe extern "C" fn(*mut c_void), +} + +impl Drop for Handle { + fn drop(&mut self) { + unsafe { (self.close)(self.pointer) }; + } +} + +fn source(library: &'static SparkDsLibrary, failure: &str, capabilities: &str) -> (Handle, u64) { + let failure = CString::new(failure).unwrap(); + let capabilities = CString::new(capabilities).unwrap(); + let keys = [c"failure".as_ptr(), c"capabilities".as_ptr()]; + let values = [failure.as_ptr(), capabilities.as_ptr()]; + let info = SparkDsCreateInfo { + size: size_of::() as i64, + name: c"write_test".as_ptr(), + option_keys: keys.as_ptr(), + option_values: values.as_ptr(), + num_options: keys.len() as i64, + }; + let mut mask = 0; + let mut pointer = ptr::null_mut(); + let status = unsafe { library.create.unwrap()(&info, &mut mask, &mut pointer) }; + success(library, status); + assert!(!pointer.is_null()); + ( + Handle { + library, + pointer, + close: library.close.unwrap(), + }, + mask, + ) +} + +fn create_driver( + source: &Handle, + mode: i32, + query: *const std::ffi::c_char, +) -> (c_int, *mut c_void) { + let mut input = ArrowSchema::try_from(schema().as_ref()).unwrap(); + let info = SparkDsWriteInfo { + size: size_of::() as i64, + mode, + query_id: query, + }; + let mut pointer = ptr::null_mut(); + let status = unsafe { + source.library.create_writer.unwrap()(source.pointer, &mut input, &info, &mut pointer) + }; + (status, pointer) +} + +fn driver(source: &Handle, mode: i32) -> Handle { + let (status, pointer) = create_driver(source, mode, ptr::null()); + success(source.library, status); + assert!(!pointer.is_null()); + Handle { + library: source.library, + pointer, + close: source.library.close_writer.unwrap(), + } +} + +unsafe extern "C" fn collect(sink: *mut c_void, bytes: *const u8, length: i64) { + let output = unsafe { &mut *sink.cast::>>() }; + output.push(unsafe { std::slice::from_raw_parts(bytes, length as usize) }.to_vec()); +} + +fn serialize(driver: &Handle) -> (c_int, Vec>) { + let mut output = Vec::>::new(); + let status = unsafe { + driver.library.serialize_writer.unwrap()( + driver.pointer, + Some(collect), + (&mut output as *mut Vec>).cast(), + ) + }; + (status, output) +} + +fn state(driver: &Handle) -> Vec { + let (status, mut output) = serialize(driver); + success(driver.library, status); + assert_eq!(output.len(), 1); + output.pop().unwrap() +} + +#[repr(C)] +struct HostTask { + context: SparkDsTaskContext, + cancelled: AtomicBool, +} + +unsafe extern "C" fn interrupted(context: *const SparkDsTaskContext) -> bool { + unsafe { &*context.cast::() } + .cancelled + .load(Ordering::SeqCst) +} + +struct Task { + library: &'static SparkDsLibrary, + pointer: *mut c_void, + host: Box, +} + +impl Drop for Task { + fn drop(&mut self) { + // The host context must outlive the provider's destructor as well as its callbacks. + unsafe { self.library.close_data_writer.unwrap()(self.pointer) }; + } +} + +fn task(library: &'static SparkDsLibrary, state: &[u8], partition: i32) -> Task { + let host = Box::new(HostTask { + context: SparkDsTaskContext { + size: size_of::() as i64, + stage_id: 3, + stage_attempt_number: 1, + partition_id: partition, + attempt_number: 2, + task_attempt_id: 100 + i64::from(partition), + is_interrupted: Some(interrupted), + }, + cancelled: AtomicBool::new(false), + }); + let mut pointer = ptr::null_mut(); + let status = unsafe { + library.create_data_writer.unwrap()( + &host.context, + state.as_ptr(), + state.len() as i64, + -1, + &mut pointer, + ) + }; + success(library, status); + assert!(!pointer.is_null()); + Task { + library, + pointer, + host, + } +} + +fn batch(values: Vec) -> RecordBatch { + RecordBatch::try_new(schema(), vec![Arc::new(Int64Array::from(values))]).unwrap() +} + +fn write(task: &Task, batch: RecordBatch) -> c_int { + let data = StructArray::from(batch).to_data(); + let (mut array, mut schema) = to_ffi(&data).unwrap(); + let status = unsafe { task.library.write.unwrap()(task.pointer, &mut array, &mut schema) }; + if status == 0 { + assert!( + array.is_released(), + "retained input must be moved from the caller" + ); + assert!(schema.release.is_none()); + } + status +} + +fn task_commit(task: &Task) -> (c_int, Vec>) { + let mut output = Vec::>::new(); + let status = unsafe { + task.library.commit_data_writer.unwrap()( + task.pointer, + Some(collect), + (&mut output as *mut Vec>).cast(), + ) + }; + (status, output) +} + +fn task_message(task: &Task) -> Vec { + let (status, mut messages) = task_commit(task); + success(task.library, status); + assert_eq!(messages.len(), 1); + messages.pop().unwrap() +} + +fn driver_finish(driver: &Handle, abort: bool, messages: &[Option<&[u8]>]) -> c_int { + let pointers: Vec<_> = messages + .iter() + .map(|message| message.map_or(ptr::null(), |bytes| bytes.as_ptr())) + .collect(); + let lengths: Vec<_> = messages + .iter() + .map(|message| message.map_or(0, |bytes| bytes.len() as i64)) + .collect(); + let callback = if abort { + driver.library.abort.unwrap() + } else { + driver.library.commit.unwrap() + }; + unsafe { + callback( + driver.pointer, + -1, + pointers.as_ptr(), + lengths.as_ptr(), + messages.len() as i64, + ) + } +} + +#[test] +fn write_registration_and_options_control_capabilities() { + let _guard = TEST_LOCK.lock().unwrap(); + EVENTS.lock().unwrap().clear(); + let library = library(); + assert!(library.create_reader.is_none()); + assert!(library.read.is_none()); + assert!(library.create_stream_reader.is_none()); + assert!(library.create_writer.is_some()); + assert!(library.write.is_some()); + + for (capabilities, expected) in [("", 4 | 16), ("append", 4), ("none", 0)] { + let (source, mask) = source(library, "", capabilities); + assert_eq!(mask, expected); + for mode in [0, 1] { + let (status, pointer) = create_driver(&source, mode, ptr::null()); + let supported = expected & 4 != 0 && (mode == 0 || expected & 16 != 0); + if supported { + success(library, status); + unsafe { library.close_writer.unwrap()(pointer) }; + } else { + assert_ne!(status, 0); + assert!(pointer.is_null()); + assert!(error(library).contains("support")); + } + } + } + assert_eq!( + events(), + vec![ + Event::WriteCreated(false), + Event::WriteCreated(true), + Event::WriteCreated(false) + ] + ); + + let (source, _) = source(library, "", ""); + for (mode, query) in [(99, ptr::null()), (0, c"streaming-query".as_ptr())] { + let (status, pointer) = create_driver(&source, mode, query); + assert_ne!(status, 0); + assert!(pointer.is_null()); + } +} + +#[test] +fn serialized_factory_and_messages_roundtrip_without_driver_source_or_input_buffers() { + let _guard = TEST_LOCK.lock().unwrap(); + EVENTS.lock().unwrap().clear(); + let library = library(); + let (source, _) = source(library, "", ""); + let driver = driver(&source, 1); + drop(source); + let state = state(&driver); + assert_eq!(state, self::state(&driver)); + + let first = task(library, &state, 7); + let input = batch(vec![2, 5, 9]); + let buffer = input.column(0).to_data().buffers()[0].clone(); + success(library, write(&first, input)); + assert!( + buffer.strong_count() > 1, + "writer must own retained Arrow buffers" + ); + success(library, write(&first, batch(vec![]))); + let first_message = task_message(&first); + assert_eq!(first_message, task_message(&first)); + drop(first); + assert_eq!( + buffer.strong_count(), + 1, + "close must release retained Arrow buffers" + ); + + let second = task(library, &state, 8); + let second_message = task_message(&second); + drop(second); + let messages = [ + Some(first_message.as_slice()), + Some(second_message.as_slice()), + ]; + success(library, driver_finish(&driver, false, &messages)); + success(library, driver_finish(&driver, false, &messages)); + assert_eq!( + events(), + vec![ + Event::WriteCreated(true), + Event::Factory, + Event::Batch(3), + Event::Batch(0), + Event::TaskCommit, + Event::TaskDropped(false), + Event::TaskCommit, + Event::TaskDropped(false), + Event::DriverCommit(vec![ + Message { + partition: 7, + attempt: 107, + values: vec![2, 5, 9], + batches: 2, + failure: String::new(), + }, + Message { + partition: 8, + attempt: 108, + values: vec![], + batches: 0, + failure: String::new(), + }, + ]), + ] + ); +} + +#[test] +fn input_schema_and_live_task_cancellation_are_checked() { + let _guard = TEST_LOCK.lock().unwrap(); + let library = library(); + for cancelled in [false, true] { + EVENTS.lock().unwrap().clear(); + let (source, _) = source(library, "", ""); + let driver = driver(&source, 0); + let task = task(library, &state(&driver), 0); + if cancelled { + success(library, write(&task, batch(vec![1]))); + task.host.cancelled.store(true, Ordering::SeqCst); + let (status, messages) = task_commit(&task); + assert_ne!(status, 0); + assert!(messages.is_empty()); + assert!(error(library).contains("cancelled")); + } else { + let wrong_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let input = + RecordBatch::try_new(wrong_schema, vec![Arc::new(Int32Array::from(vec![1]))]) + .unwrap(); + assert_ne!(write(&task, input), 0); + assert!(error(library).contains("schema")); + assert!(!events() + .iter() + .any(|event| matches!(event, Event::Batch(_)))); + } + success(library, unsafe { + library.abort_data_writer.unwrap()(task.pointer) + }); + drop(task); + assert!(events().contains(&Event::TaskAbort(cancelled))); + assert!(events().contains(&Event::TaskDropped(cancelled))); + success(library, driver_finish(&driver, true, &[None])); + } +} + +#[test] +fn task_failures_and_message_serialization_failures_remain_abortable() { + let _guard = TEST_LOCK.lock().unwrap(); + let library = library(); + for failure in [ + "write_error", + "write_panic", + "task_commit_error", + "task_commit_panic", + "message_serialization", + ] { + EVENTS.lock().unwrap().clear(); + let (source, _) = source(library, failure, ""); + let driver = driver(&source, 0); + let task = task(library, &state(&driver), 0); + let status = write(&task, batch(vec![4])); + let write_failed = failure.starts_with("write_"); + assert_eq!(status != 0, write_failed, "{failure}"); + let (status, messages) = task_commit(&task); + assert_ne!(status, 0, "{failure}"); + assert!(messages.is_empty()); + let message = error(library); + assert!( + message.contains("failed") || message.contains("panicked"), + "{message}" + ); + let before_retry = events(); + let (status, messages) = task_commit(&task); + assert_ne!(status, 0); + assert!(messages.is_empty()); + assert_eq!( + events(), + before_retry, + "failed commit cannot run provider commit twice" + ); + success(library, unsafe { + library.abort_data_writer.unwrap()(task.pointer) + }); + success(library, unsafe { + library.abort_data_writer.unwrap()(task.pointer) + }); + drop(task); + assert_eq!( + events() + .iter() + .filter(|event| **event == Event::TaskAbort(false)) + .count(), + 1, + ); + assert_eq!( + events() + .iter() + .filter(|event| **event == Event::TaskCommit) + .count(), + usize::from(!write_failed), + ); + success(library, driver_finish(&driver, true, &[None])); + } +} + +#[test] +fn driver_failures_preserve_messages_and_missing_task_positions_for_abort() { + let _guard = TEST_LOCK.lock().unwrap(); + let library = library(); + for failure in ["driver_commit_error", "driver_commit_panic"] { + EVENTS.lock().unwrap().clear(); + let (source, _) = source(library, failure, ""); + let driver = driver(&source, 0); + let task = task(library, &state(&driver), 1); + success(library, write(&task, batch(vec![42]))); + let message = task_message(&task); + drop(task); + let committed = [Some(message.as_slice())]; + assert_ne!(driver_finish(&driver, false, &committed), 0); + assert!(error(library).contains("driver commit")); + let before_retry = events(); + assert_ne!(driver_finish(&driver, false, &committed), 0); + assert_eq!(events(), before_retry); + let aborted = [None, Some(message.as_slice()), None]; + success(library, driver_finish(&driver, true, &aborted)); + success(library, driver_finish(&driver, true, &aborted)); + let expected = Message { + partition: 1, + attempt: 101, + values: vec![42], + batches: 1, + failure: failure.to_owned(), + }; + assert!(events().contains(&Event::DriverCommit(vec![expected.clone()]))); + assert_eq!( + events() + .iter() + .filter(|event| matches!(event, Event::DriverAbort(_))) + .count(), + 1, + ); + assert!(events().contains(&Event::DriverAbort(vec![None, Some(expected), None]))); + } +} + +#[test] +fn factory_serialization_failure_is_cached_and_driver_can_abort() { + let _guard = TEST_LOCK.lock().unwrap(); + EVENTS.lock().unwrap().clear(); + let library = library(); + let (source, _) = source(library, "factory_serialization", ""); + let driver = driver(&source, 0); + for _ in 0..2 { + let (status, output) = serialize(&driver); + assert_ne!(status, 0); + assert!(output.is_empty()); + assert!(error(library).contains("factory serialization failed")); + } + success(library, driver_finish(&driver, true, &[])); + assert_eq!( + events(), + vec![ + Event::WriteCreated(false), + Event::Factory, + Event::DriverAbort(vec![]) + ], + ); +} diff --git a/sql/native/rust/tests/capabilities.rs b/sql/native/rust/tests/capabilities.rs index c18a7dd603504..80a16a8bda28b 100644 --- a/sql/native/rust/tests/capabilities.rs +++ b/sql/native/rust/tests/capabilities.rs @@ -25,8 +25,8 @@ use spark_native_source::abi::{ ArrowSchema, SparkDsCreateInfo, SparkDsHost, SparkDsLibrary, ABI_VERSION, }; use spark_native_source::{ - BatchReader, BatchReaderFactory, BatchSource, Options, Result, ScanBuilder, ScanPlan, - SourceCapability, TaskContext, + BatchReader, BatchReaderFactory, BatchSource, DataSource, Options, Result, ScanBuilder, + ScanPlan, SourceCapability, TaskContext, }; static DECLARATIONS: AtomicUsize = AtomicUsize::new(0); @@ -38,16 +38,14 @@ struct Source { mode: String, } -impl BatchSource for Source { - type Scan = Builder; - +impl DataSource for Source { fn new(options: Options) -> Result { Ok(Self { mode: options.get("capabilities").cloned().unwrap(), }) } - fn capabilities(&self) -> Vec { + fn capabilities(&self, _: &[SourceCapability]) -> Vec { DECLARATIONS.fetch_add(1, Ordering::SeqCst); match self.mode.as_str() { "empty" => vec![], @@ -67,6 +65,10 @@ impl BatchSource for Source { SCHEMAS.fetch_add(1, Ordering::SeqCst); Ok(Arc::new(Schema::empty())) } +} + +impl BatchSource for Source { + type Scan = Builder; fn scan(&self, _: SchemaRef) -> Result { SCANS.fetch_add(1, Ordering::SeqCst); @@ -225,7 +227,7 @@ fn source_capabilities_control_the_exported_abi_and_source_lifecycle() { assert_ne!(status, 0, "{mode}"); assert_eq!( error(library), - format!("capability {capability} is not supported by this batch-read SDK") + format!("capability {capability} is not registered") ); assert_eq!(mask, 0); assert!(source.is_null()); From 69adf1a46774afdee065fa5100516204f5e13a99 Mon Sep 17 00:00:00 2001 From: Gengliang Wang Date: Mon, 19 Oct 2026 18:46:12 +0000 Subject: [PATCH 5/6] Refine Rust SDK write API naming and lifecycle contract - Rename BatchWriterFactory/BatchWriter to DataWriterFactory/DataWriter, following DSv2 and avoiding a one-letter difference from BatchWrite. - Mark WriteMode, FilterSupport, LimitSupport, and Predicate as #[non_exhaustive] so new variants are not breaking changes. - Document the call-order guarantees the SDK enforces for BatchWrite and DataWriter, and remove the sink example's redundant phase tracking. Co-authored-by: Isaac --- .../src/main/rust/native-data-sink/README.md | 4 +- .../src/main/rust/native-data-sink/src/lib.rs | 97 ++++++------------- sql/native/rust/README.md | 10 +- sql/native/rust/src/export.rs | 2 +- sql/native/rust/src/lib.rs | 49 ++++++++-- sql/native/rust/src/predicate.rs | 1 + sql/native/rust/src/provider.rs | 6 +- sql/native/rust/src/write.rs | 4 +- sql/native/rust/tests/batch_write.rs | 6 +- 9 files changed, 92 insertions(+), 87 deletions(-) diff --git a/examples/src/main/rust/native-data-sink/README.md b/examples/src/main/rust/native-data-sink/README.md index 78384f64799ac..679e9152e483b 100644 --- a/examples/src/main/rust/native-data-sink/README.md +++ b/examples/src/main/rust/native-data-sink/README.md @@ -18,7 +18,7 @@ limitations under the License. # Rust native batch sink This write-only example implements `DataSource`, `BatchWriteSource`, `BatchWrite`, -`BatchWriterFactory`, and `BatchWriter` from the [Rust native data source SDK][sdk]. +`DataWriterFactory`, and `DataWriter` from the [Rust native data source SDK][sdk]. It registers `export_data_source!(JsonSink, batch_write, truncate)` and implements no read traits. The SDK supplies the C ABI entry point, typed factory and commit message serialization, Arrow batch import, capability checks, and error handling. @@ -67,7 +67,7 @@ That one manifest replacement is the commit point: - Abort removes the entire operation, including files with no commit message and a sealed directory whose manifest publication failed. It leaves the current manifest and its data unchanged. Dropping an unfinished coordinator also performs - best-effort cleanup. Dropping a prepared task writer leaves its file for the driver. + best-effort cleanup. Dropping a committed task writer leaves its file for the driver. To inspect output, read only the files listed in `manifest.json`. Recursively reading the destination directory would also pick up data from previous truncate operations. diff --git a/examples/src/main/rust/native-data-sink/src/lib.rs b/examples/src/main/rust/native-data-sink/src/lib.rs index cae8d43ad04b8..570f3fe11e122 100644 --- a/examples/src/main/rust/native-data-sink/src/lib.rs +++ b/examples/src/main/rust/native-data-sink/src/lib.rs @@ -30,7 +30,7 @@ use arrow_array::{Array, Int64Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use serde::{Deserialize, Serialize}; use spark_native_source::{ - BatchWrite, BatchWriteSource, BatchWriter, BatchWriterFactory, DataSource, Options, Result, + BatchWrite, BatchWriteSource, DataSource, DataWriter, DataWriterFactory, Options, Result, SourceCapability, TaskContext, TaskWriter, WriteMode, }; @@ -177,13 +177,6 @@ impl Drop for WriteLock { } } -#[derive(Clone, Copy, PartialEq, Eq)] -enum WriteState { - Open, - Committed, - Aborted, -} - pub struct JsonWrite { root: PathBuf, operation: String, @@ -191,8 +184,8 @@ pub struct JsonWrite { objects: PathBuf, temporary_manifest: PathBuf, previous: Manifest, - mode: WriteMode, - state: WriteState, + truncate: bool, + committed: bool, sealed: bool, lock: Option, } @@ -202,6 +195,11 @@ impl JsonWrite { if !cfg!(unix) { return Err("This demo requires a POSIX filesystem".into()); } + let truncate = match mode { + WriteMode::Append => false, + WriteMode::Truncate => true, + _ => return Err("Unsupported write mode".into()), + }; fs::create_dir_all(&root)?; let lock = WriteLock::acquire(&root)?; let previous = read_manifest(&root)?; @@ -216,8 +214,8 @@ impl JsonWrite { operation, staging, previous, - mode, - state: WriteState::Open, + truncate, + committed: false, sealed: false, lock: Some(lock), }) @@ -254,9 +252,10 @@ impl JsonWrite { } } - let mut files = match self.mode { - WriteMode::Append => self.previous.files.clone(), - WriteMode::Truncate => Vec::new(), + let mut files = if self.truncate { + Vec::new() + } else { + self.previous.files.clone() }; files.extend(messages.into_iter().map(|message| DataFile { path: format!("objects/{}/{}", self.operation, message.filename), @@ -277,7 +276,7 @@ impl JsonWrite { // This single rename is the commit point. No existing data is deleted, // including during truncate. No fallible work follows publication. fs::rename(&self.temporary_manifest, self.root.join("manifest.json"))?; - self.state = WriteState::Committed; + self.committed = true; self.lock.take(); Ok(()) } @@ -308,13 +307,12 @@ impl JsonWrite { } } +// The SDK calls each method at most once, never after a successful commit or an +// abort, so this coordinator does not track its own phase. impl BatchWrite for JsonWrite { type Factory = JsonWriterFactory; fn factory(&mut self) -> Result { - if self.state != WriteState::Open { - return Err("The write operation is no longer open".into()); - } Ok(JsonWriterFactory { staging: self.staging.clone(), operation: self.operation.clone(), @@ -322,19 +320,11 @@ impl BatchWrite for JsonWrite { } fn commit(&mut self, messages: Vec) -> Result<()> { - match self.state { - WriteState::Committed => Ok(()), - WriteState::Aborted => Err("The write operation was aborted".into()), - WriteState::Open => self.publish(messages), - } + self.publish(messages) } fn abort(&mut self, _messages: Vec>) -> Result<()> { - if self.state == WriteState::Committed { - return Ok(()); - } let result = self.cleanup(); - self.state = WriteState::Aborted; self.lock.take(); result } @@ -342,7 +332,7 @@ impl BatchWrite for JsonWrite { impl Drop for JsonWrite { fn drop(&mut self) { - if self.state != WriteState::Committed { + if !self.committed { let _ = self.cleanup(); } } @@ -363,7 +353,7 @@ pub struct CommitMessage { rows: u64, } -impl BatchWriterFactory for JsonWriterFactory { +impl DataWriterFactory for JsonWriterFactory { type Message = CommitMessage; fn open<'task>(&self, task: TaskContext<'task>) -> Result> { @@ -392,8 +382,7 @@ impl BatchWriterFactory for JsonWriterFactory { rows: 0, }, writer: Some(BufWriter::new(file)), - prepared: false, - failed: false, + committed: false, })); } Err(error) if error.kind() == ErrorKind::AlreadyExists => continue, @@ -415,8 +404,7 @@ struct JsonTaskWriter<'task> { path: PathBuf, message: CommitMessage, writer: Option>, - prepared: bool, - failed: bool, + committed: bool, } impl JsonTaskWriter<'_> { @@ -427,11 +415,14 @@ impl JsonTaskWriter<'_> { drop(file); } } +} - fn write_batch(&mut self, batch: RecordBatch) -> Result<()> { - if self.prepared || self.failed { - return Err("The task writer is no longer accepting batches".into()); - } +// The SDK never writes after a failure, commit, or abort, commits at most once, +// and never aborts a committed writer, so this writer does not track its phase. +impl DataWriter for JsonTaskWriter<'_> { + type Message = CommitMessage; + + fn write(&mut self, batch: RecordBatch) -> Result<()> { validate_schema(batch.schema().as_ref())?; let writer = self.writer.as_mut().ok_or("The task writer is closed")?; let ids = batch @@ -463,13 +454,7 @@ impl JsonTaskWriter<'_> { Ok(()) } - fn prepare(&mut self) -> Result { - if self.failed { - return Err("Cannot commit a failed task writer".into()); - } - if self.prepared { - return Ok(self.message.clone()); - } + fn commit(&mut self) -> Result { if self.task.is_interrupted() { return Err("Task cancelled before committing the sink".into()); } @@ -477,30 +462,12 @@ impl JsonTaskWriter<'_> { writer.flush()?; writer.get_ref().sync_all()?; drop(writer); - self.prepared = true; + self.committed = true; Ok(self.message.clone()) } -} - -impl BatchWriter for JsonTaskWriter<'_> { - type Message = CommitMessage; - - fn write(&mut self, batch: RecordBatch) -> Result<()> { - let result = self.write_batch(batch); - self.failed |= result.is_err(); - result - } - - fn commit(&mut self) -> Result { - let result = self.prepare(); - self.failed |= result.is_err(); - result - } fn abort(&mut self) -> Result<()> { self.discard_buffer(); - self.prepared = false; - self.failed = true; match fs::remove_file(&self.path) { Ok(()) => Ok(()), Err(error) if error.kind() == ErrorKind::NotFound => Ok(()), @@ -511,7 +478,7 @@ impl BatchWriter for JsonTaskWriter<'_> { impl Drop for JsonTaskWriter<'_> { fn drop(&mut self) { - if !self.prepared { + if !self.committed { self.discard_buffer(); let _ = fs::remove_file(&self.path); } diff --git a/sql/native/rust/README.md b/sql/native/rust/README.md index bb7eaa2746ad9..f1120d937a6c6 100644 --- a/sql/native/rust/README.md +++ b/sql/native/rust/README.md @@ -49,7 +49,7 @@ Read support uses three traits: reader from a typed partition and a borrowed `TaskContext`. Write support is independent: a write-only connector implements `BatchWriteSource`, -`BatchWrite`, `BatchWriterFactory`, and `BatchWriter`, without implementing reads. +`BatchWrite`, `DataWriterFactory`, and `DataWriter`, without implementing reads. The write lifecycle is described below. The factory and partition types implement `Serialize` and `Deserialize`. @@ -168,8 +168,8 @@ Driver: serialized factory + schema | typed messages | | Executor, per task attempt: v | - BatchWriterFactory::open(TaskContext) - -> BatchWriter::write(RecordBatch)* + DataWriterFactory::open(TaskContext) + -> DataWriter::write(RecordBatch)* -> commit() / abort() -> Drop ``` @@ -179,7 +179,7 @@ executor configuration; the SDK caches the encoded factory and input schema, including a terminal error if creation or serialization fails. Restoring an executor writer uses only those bytes, without creating a driver registry. -Each `BatchWriterFactory::open` returns a task-bound writer. `write` takes ownership +Each `DataWriterFactory::open` returns a task-bound writer. `write` takes ownership of a `RecordBatch`, so its Arrow buffers can be retained after the callback. The SDK moves the incoming Arrow array and schema, checks the schema against the negotiated input, and releases imported resources on failures. Input slices need @@ -192,6 +192,8 @@ receives one winning message per partition; driver `abort` receives commit leaves the same object available for abort. The SDK prevents further writes after a failure or successful commit, caches successful task commit messages, and avoids repeating successful commit/abort callbacks. Drop only releases resources. +Connectors therefore do not need to track their own phase; the `# Lifecycle` +sections of `BatchWrite` and `DataWriter` list the exact call-order guarantees. 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 diff --git a/sql/native/rust/src/export.rs b/sql/native/rust/src/export.rs index fc3b7c9963330..785ee2e074e50 100644 --- a/sql/native/rust/src/export.rs +++ b/sql/native/rust/src/export.rs @@ -55,7 +55,7 @@ macro_rules! export_data_source { as $crate::ScanBuilder>::Factory; type WriterFactory = <::Write as $crate::BatchWrite>::Factory; - type Message = ::Message; + type Message = ::Message; fn bridge() -> &'static Bridge { static BRIDGE: std::sync::OnceLock> = diff --git a/sql/native/rust/src/lib.rs b/sql/native/rust/src/lib.rs index 3a73fd9273b87..18c8935c00d00 100644 --- a/sql/native/rust/src/lib.rs +++ b/sql/native/rust/src/lib.rs @@ -46,7 +46,7 @@ pub type Error = Box; pub type Result = std::result::Result; pub type Options = BTreeMap; pub type BatchReader<'task> = Box; -pub type TaskWriter<'task, M> = Box + 'task>; +pub type TaskWriter<'task, M> = Box + 'task>; /// Native operations supported by a source instance with its effective options. /// @@ -68,6 +68,7 @@ pub enum SourceCapability { /// How completely a connector evaluates a pushed predicate. /// Spark currently rechecks all predicates, including exact ones. #[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] pub enum FilterSupport { /// Implements Spark SQL semantics and returns only matching rows. Exact, @@ -79,6 +80,7 @@ pub enum FilterSupport { /// How a connector limits the rows produced by the whole scan. #[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] pub enum LimitSupport { Unsupported, /// Caps each partition; the ABI boolean is false because it promises a @@ -131,7 +133,12 @@ pub trait BatchSource: DataSource { } /// 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, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[non_exhaustive] pub enum WriteMode { Append, /// Replace existing data at driver commit, never at writer creation. @@ -147,6 +154,8 @@ pub trait BatchWriteSource: DataSource { /// Validate the input schema and create an independent write coordinator. /// Construction may be followed only by Drop, without commit or abort. + /// The SDK calls this only while the instance's capabilities include + /// BatchWrite, and passes `Truncate` only when they include Truncate. fn write(&self, schema: SchemaRef, mode: WriteMode) -> Result; } @@ -154,9 +163,21 @@ pub trait BatchWriteSource: DataSource { /// /// Keep irreversible publication, including removing overwritten data, in commit. /// Drop releases resources; it must not publish data or perform fallible flushing. -/// Failed commit is followed by abort on the same coordinator. +/// +/// # Lifecycle +/// +/// The SDK enforces the following, so implementations need not track their own +/// phase or guard against repeated calls: +/// +/// - `factory` runs at most once, and never after `commit` or `abort`. +/// - `commit` runs at most once, only while no earlier call has failed. It is +/// skipped, and the write fails, when any task message is missing. +/// - `abort` runs at most once, never after a successful `commit`. It may follow +/// a failed `factory` or `commit`, or run without either being called. +/// - Drop runs exactly once, possibly without `commit` or `abort`. +/// - A panic in any method is reported as that method's failure. pub trait BatchWrite: Send + 'static { - type Factory: BatchWriterFactory; + type Factory: DataWriterFactory; /// Return serializable executor configuration. The SDK caches the encoded /// factory and negotiated schema, including serialization failures. @@ -167,7 +188,7 @@ pub trait BatchWrite: Send + 'static { /// and speculative attempts can stage other data that must not be published. fn commit( &mut self, - messages: Vec<::Message>, + messages: Vec<::Message>, ) -> Result<()>; /// Abort after input, task, or driver commit failure. None preserves the @@ -175,7 +196,7 @@ pub trait BatchWrite: Send + 'static { /// Executor loss can leave orphaned staged data; plan for its cleanup. fn abort( &mut self, - messages: Vec::Message>>, + messages: Vec::Message>>, ) -> Result<()>; } @@ -184,14 +205,28 @@ pub trait BatchWrite: Send + 'static { /// Open connections and files in open, not in the factory. A returned writer /// cannot borrow this temporary factory, but may retain the task context until /// close. Callbacks can move between threads sequentially; writers must be Send. -pub trait BatchWriterFactory: Serialize + DeserializeOwned + Send + 'static { +pub trait DataWriterFactory: Serialize + DeserializeOwned + Send + 'static { type Message: Serialize + DeserializeOwned + Send + 'static; fn open<'task>(&self, task: TaskContext<'task>) -> Result>; } /// One task attempt's writer. The SDK closes it after commit or abort. -pub trait BatchWriter: Send { +/// +/// # Lifecycle +/// +/// The SDK enforces the following, so implementations need not track their own +/// phase or guard against repeated calls: +/// +/// - `write` runs zero or more times, and never after any call has failed or +/// after `commit` or `abort`. +/// - `commit` runs at most once, never after a failed `write` or after `abort`. +/// - `abort` runs at most once, never after a successful `commit`. It may follow +/// a failed `write` or `commit`. +/// - Drop runs exactly once when Spark closes the writer, possibly without +/// `commit` or `abort`. +/// - 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 diff --git a/sql/native/rust/src/predicate.rs b/sql/native/rust/src/predicate.rs index f0ff0b8391712..94bcda501f067 100644 --- a/sql/native/rust/src/predicate.rs +++ b/sql/native/rust/src/predicate.rs @@ -24,6 +24,7 @@ use serde_json::Value; /// Function names are extensible: unknown names must remain residual. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "lowercase")] +#[non_exhaustive] pub enum Predicate { Column { name: Vec, diff --git a/sql/native/rust/src/provider.rs b/sql/native/rust/src/provider.rs index 895dad7f35135..29d4c22712c67 100644 --- a/sql/native/rust/src/provider.rs +++ b/sql/native/rust/src/provider.rs @@ -21,8 +21,8 @@ use arrow_schema::SchemaRef; use serde::{Deserialize, Serialize}; use crate::{ - BatchReader, BatchReaderFactory, BatchSource, BatchWrite, BatchWriteSource, BatchWriterFactory, - DataSource, Result, ScanBuilder, ScanPlan, SourceCapability, TaskContext, TaskWriter, + BatchReader, BatchReaderFactory, BatchSource, BatchWrite, BatchWriteSource, DataSource, + DataWriterFactory, Result, ScanBuilder, ScanPlan, SourceCapability, TaskContext, TaskWriter, WriteMode, }; @@ -142,7 +142,7 @@ impl BatchWrite for Unavailable { } } -impl BatchWriterFactory for Unavailable { +impl DataWriterFactory for Unavailable { type Message = (); fn open<'task>(&self, _: TaskContext<'task>) -> Result> { diff --git a/sql/native/rust/src/write.rs b/sql/native/rust/src/write.rs index 486db70b3207b..c9794f1c358e2 100644 --- a/sql/native/rust/src/write.rs +++ b/sql/native/rust/src/write.rs @@ -26,7 +26,7 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize}; use crate::abi::{self, ArrowArray, ArrowSchema, SparkDsTaskContext}; use crate::bridge::stream::{canonical_fields, validate_schema}; use crate::bridge::{decode, encode, panic_message, provider_call}; -use crate::{BatchWrite, BatchWriterFactory, Result, TaskContext, TaskWriter}; +use crate::{BatchWrite, DataWriterFactory, Result, TaskContext, TaskWriter}; type Cached = std::result::Result; @@ -178,7 +178,7 @@ fn drop_provider(value: T) { /// `task` must satisfy TaskContext::from_raw and remain valid until close_task. /// The returned handle must only be used with callbacks for F::Message, never /// concurrently, and must be closed exactly once while the plugin stays loaded. -pub unsafe fn open( +pub unsafe fn open( task: *const SparkDsTaskContext, state: &[u8], ) -> Result<*mut c_void> { diff --git a/sql/native/rust/tests/batch_write.rs b/sql/native/rust/tests/batch_write.rs index eafc4ed68c900..85e446be6f4b7 100644 --- a/sql/native/rust/tests/batch_write.rs +++ b/sql/native/rust/tests/batch_write.rs @@ -28,7 +28,7 @@ use spark_native_source::abi::{ SparkDsWriteInfo, ABI_VERSION, }; use spark_native_source::{ - BatchWrite, BatchWriteSource, BatchWriter, BatchWriterFactory, DataSource, Options, Result, + BatchWrite, BatchWriteSource, DataSource, DataWriter, DataWriterFactory, Options, Result, SourceCapability, TaskContext, TaskWriter, WriteMode, }; @@ -169,7 +169,7 @@ fn serialize_message_mode( serializer.serialize_str(mode) } -impl BatchWriterFactory for Factory { +impl DataWriterFactory for Factory { type Message = Message; fn open<'task>(&self, task: TaskContext<'task>) -> Result> { @@ -187,7 +187,7 @@ struct Writer<'task> { batches: Vec, } -impl BatchWriter for Writer<'_> { +impl DataWriter for Writer<'_> { type Message = Message; fn write(&mut self, batch: RecordBatch) -> Result<()> { From a8df273ed4552adb3e44fea5c740a20b2a0d5562 Mon Sep 17 00:00:00 2001 From: Gengliang Wang Date: Mon, 19 Oct 2026 20:39:47 +0000 Subject: [PATCH 6/6] Address API review of the Rust native data source SDK - Pass the data source name to DataSource::new in a non-exhaustive CreateInfo, and pass writes and executor readers and writers WriteInfo, ReaderContext, and WriterContext, which carry the negotiated schema, so ABI fields can be added. - Add DataSource::init with HostInfo, spark_native_source::log through the host log, and DataSource::INFERS_SCHEMA to leave the schema function unset. - Replace the capability list with a Capabilities set that sources can only narrow, declare truncation with BatchWriteSource::SUPPORTS_TRUNCATE, accept operations in any order, remove the one-argument macro form, and reject panic = "abort" at compile time. - Return a WritePlan from BatchWriteSource::new_write instead of a separate BatchWrite::factory, drop Copy and serde from WriteMode, and add constructors to ScanPlan and WritePlan. - Document that abort can follow a commit whose message fails to serialize, that Drop is the only cleanup when Spark closes a write without abort, and that message i belongs to partition i. - Pass commit messages that cannot be decoded to abort as None, parse each pushed predicate on its own, and include error sources in messages. - Rename BatchSource to BatchReadSource, BatchReaderFactory to PartitionReaderFactory, and the reader and writer aliases to PartitionReader and PartitionWriter, and remove LimitSupport::PerPartition. - Return plain TaskContext IDs, reject task contexts shorter than version 1 of the ABI, and add TaskContext::for_testing. - Add Predicate::literal_value to decode literals by their Spark SQL type. - Hide the macro's support modules in a private __private module. - Pin framework commit 74aa0dff: accept only the types of version 1, without intervals or the null type, and check read batches against the nullability of the negotiated schema. Co-authored-by: Isaac --- .../src/main/rust/native-data-sink/README.md | 12 +- .../src/main/rust/native-data-sink/src/lib.rs | 178 +++--- .../main/rust/native-data-source/README.md | 31 +- .../main/rust/native-data-source/src/lib.rs | 130 ++--- .../native-data-source/tests/c_abi_smoke.c | 44 +- sql/native/rust/README.md | 254 +++++---- sql/native/rust/src/abi.rs | 69 ++- sql/native/rust/src/bridge.rs | 265 ++++++--- sql/native/rust/src/export.rs | 191 ++++--- sql/native/rust/src/host.rs | 87 +++ sql/native/rust/src/lib.rs | 513 +++++++++++++----- sql/native/rust/src/predicate.rs | 290 +++++++++- sql/native/rust/src/provider.rs | 88 ++- sql/native/rust/src/stream.rs | 238 +++++++- sql/native/rust/src/task.rs | 243 ++++++--- sql/native/rust/src/write.rs | 58 +- sql/native/rust/tests/batch_write.rs | 123 +++-- sql/native/rust/tests/capabilities.rs | 190 ++++--- .../rust/tests/fixtures/spark_datasource.h | 36 +- 19 files changed, 2217 insertions(+), 823 deletions(-) create mode 100644 sql/native/rust/src/host.rs diff --git a/examples/src/main/rust/native-data-sink/README.md b/examples/src/main/rust/native-data-sink/README.md index 679e9152e483b..b00bca5d1cf3b 100644 --- a/examples/src/main/rust/native-data-sink/README.md +++ b/examples/src/main/rust/native-data-sink/README.md @@ -19,12 +19,13 @@ limitations under the License. This write-only example implements `DataSource`, `BatchWriteSource`, `BatchWrite`, `DataWriterFactory`, and `DataWriter` from the [Rust native data source SDK][sdk]. -It registers `export_data_source!(JsonSink, batch_write, truncate)` and implements +It registers `export_data_source!(JsonSink, batch_write)`, declares overwrite support +with `const SUPPORTS_TRUNCATE: bool = true` in `BatchWriteSource`, and implements no read traits. The SDK supplies the C ABI entry point, typed factory and commit message serialization, Arrow batch import, capability checks, and error handling. The example targets C ABI version 1 from [Spark PR #59424][native-source-pr], pinned -to `534dfca15723c7208061b1e9e4d80c5757360a2a`. Running it in Spark requires that +to `74aa0dffc53463034fa6549796c4dd599e9705de`. Running it in Spark requires that framework and Java 22 or later on the driver and executors. The standalone Rust tests do not need Spark or a JVM. @@ -53,9 +54,10 @@ Each task attempt exclusively creates its own file, writes Arrow batches, flushe syncs the file, then returns a typed message containing its operation, file, partition, and row count. Task commit alone does not publish data. -The driver receives the winning attempt messages. It validates those messages and -renames the entire staging directory to `objects/`, preventing new tasks -from opening files for that operation. It removes every file not named by a winning +The driver receives the winning attempt messages, the message of partition `i` at +position `i`. It validates those messages and renames the entire staging directory +to `objects/`, preventing new tasks from opening files for that +operation. It removes every file not named by a winning message, including speculative attempts and attempts whose messages never arrived. It then writes and syncs a new manifest and atomically renames it to `manifest.json`. That one manifest replacement is the commit point: diff --git a/examples/src/main/rust/native-data-sink/src/lib.rs b/examples/src/main/rust/native-data-sink/src/lib.rs index 570f3fe11e122..3705bb0a46d01 100644 --- a/examples/src/main/rust/native-data-sink/src/lib.rs +++ b/examples/src/main/rust/native-data-sink/src/lib.rs @@ -30,11 +30,12 @@ use arrow_array::{Array, Int64Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use serde::{Deserialize, Serialize}; use spark_native_source::{ - BatchWrite, BatchWriteSource, DataSource, DataWriter, DataWriterFactory, Options, Result, - SourceCapability, TaskContext, TaskWriter, WriteMode, + BatchWrite, BatchWriteSource, Capabilities, CreateInfo, DataSource, DataWriter, + DataWriterFactory, PartitionWriter, Result, SourceCapability, TaskContext, WriteInfo, + WriteMode, WritePlan, WriterContext, }; -spark_native_source::export_data_source!(JsonSink, batch_write, truncate); +spark_native_source::export_data_source!(JsonSink, batch_write); static NEXT_ID: AtomicU64 = AtomicU64::new(0); @@ -45,7 +46,8 @@ pub struct JsonSink { } impl DataSource for JsonSink { - fn new(mut options: Options) -> Result { + fn new(info: CreateInfo) -> Result { + let mut options = info.options; let path = PathBuf::from( options .remove("path") @@ -71,24 +73,29 @@ impl DataSource for JsonSink { Ok(sink_schema()) } - fn capabilities(&self, supported: &[SourceCapability]) -> Vec { + fn capabilities(&self, registered: Capabilities) -> Capabilities { if self.readonly { - Vec::new() + // Also removes Truncate, which requires batch writes. + registered.without(SourceCapability::BatchWrite) } else { - supported.to_vec() + registered } } } impl BatchWriteSource for JsonSink { - type Write = JsonWrite; + type BatchWrite = JsonWrite; + + const SUPPORTS_TRUNCATE: bool = true; - fn write(&self, schema: SchemaRef, mode: WriteMode) -> Result { + fn new_write(&self, info: WriteInfo) -> Result> { if self.readonly { return Err("This sink is readonly".into()); } - validate_schema(&schema)?; - JsonWrite::new(self.path.clone(), mode) + validate_schema(&info.schema)?; + let write = JsonWrite::new(self.path.clone(), info.mode)?; + let factory = write.writer_factory(); + Ok(WritePlan::new(write, factory)) } } @@ -221,20 +228,28 @@ impl JsonWrite { }) } + fn writer_factory(&self) -> JsonWriterFactory { + JsonWriterFactory { + staging: self.staging.clone(), + operation: self.operation.clone(), + } + } + fn publish(&mut self, messages: Vec) -> Result<()> { let mut filenames = BTreeSet::new(); let mut partitions = BTreeSet::new(); - for message in &messages { + for (index, message) in messages.iter().enumerate() { if message.operation != self.operation || !is_filename(&message.filename) || !filenames.insert(message.filename.clone()) { return Err("Invalid or duplicate task commit message".into()); } - if let Some(partition) = message.partition { - if !partitions.insert(partition) { - return Err("More than one winning attempt for the same partition".into()); - } + // Spark passes the message of partition i at position i. + if usize::try_from(message.partition).ok() != Some(index) + || !partitions.insert(message.partition) + { + return Err("A commit message does not match its partition".into()); } if !self.staging.join(&message.filename).is_file() { return Err("A winning task's staged file is missing".into()); @@ -312,13 +327,6 @@ impl JsonWrite { impl BatchWrite for JsonWrite { type Factory = JsonWriterFactory; - fn factory(&mut self) -> Result { - Ok(JsonWriterFactory { - staging: self.staging.clone(), - operation: self.operation.clone(), - }) - } - fn commit(&mut self, messages: Vec) -> Result<()> { self.publish(messages) } @@ -330,6 +338,8 @@ impl BatchWrite for JsonWrite { } } +// Spark can close a write without abort, for example when serializing its factory +// fails, so an unfinished write also cleans up here. impl Drop for JsonWrite { fn drop(&mut self) { if !self.committed { @@ -349,14 +359,20 @@ pub struct JsonWriterFactory { pub struct CommitMessage { operation: String, filename: String, - partition: Option, + partition: i32, rows: u64, } impl DataWriterFactory for JsonWriterFactory { type Message = CommitMessage; - fn open<'task>(&self, task: TaskContext<'task>) -> Result> { + fn open<'task>( + &self, + context: WriterContext<'task>, + ) -> Result> { + // The SDK checks that each batch has this schema, so the writer checks it once. + validate_schema(&context.schema)?; + let task = context.task; if task.is_interrupted() { return Err("Task cancelled before opening the sink".into()); } @@ -365,8 +381,8 @@ impl DataWriterFactory for JsonWriterFactory { for _ in 0..16 { let filename = format!( "part-{}-attempt-{}-{}.jsonl", - task.partition_id().unwrap_or(-1), - task.task_attempt_id().unwrap_or(-1), + task.partition_id(), + task.task_attempt_id(), unique_id()? ); let path = self.staging.join(&filename); @@ -423,18 +439,17 @@ impl DataWriter for JsonTaskWriter<'_> { type Message = CommitMessage; fn write(&mut self, batch: RecordBatch) -> Result<()> { - validate_schema(batch.schema().as_ref())?; let writer = self.writer.as_mut().ok_or("The task writer is closed")?; let ids = batch .column(0) .as_any() .downcast_ref::() - .unwrap(); + .ok_or("The id column must be an Arrow int64 array")?; let labels = batch .column(1) .as_any() .downcast_ref::() - .unwrap(); + .ok_or("The label column must be an Arrow utf8 array")?; for index in 0..batch.num_rows() { if index % 1024 == 0 && self.task.is_interrupted() { return Err("Task cancelled while writing the sink".into()); @@ -523,11 +538,7 @@ mod tests { } fn source(&self) -> JsonSink { - JsonSink::new(Options::from([( - "path".to_owned(), - self.0.to_str().unwrap().to_owned(), - )])) - .unwrap() + JsonSink::new(info(&[("path", self.0.to_str().unwrap())])).unwrap() } } @@ -537,6 +548,29 @@ mod tests { } } + fn info(options: &[(&str, &str)]) -> CreateInfo { + let options = options + .iter() + .map(|(key, value)| (key.to_string(), value.to_string())) + .collect(); + CreateInfo::new("rust_sdk_sink", options) + } + + fn new_write(source: &JsonSink, mode: WriteMode) -> Result<(JsonWrite, JsonWriterFactory)> { + let plan = source.new_write(WriteInfo::new(sink_schema(), mode))?; + Ok((plan.coordinator, plan.factory)) + } + + fn open( + factory: &JsonWriterFactory, + partition: i32, + ) -> PartitionWriter<'static, CommitMessage> { + let task = TaskContext::for_testing().partition_id(partition).build(); + factory + .open(WriterContext::new(sink_schema(), task)) + .unwrap() + } + fn batch(ids: Vec>) -> RecordBatch { let labels = ids .iter() @@ -552,11 +586,11 @@ mod tests { .unwrap() } - fn stage(factory: &JsonWriterFactory, ids: Vec>) -> CommitMessage { + fn stage(factory: &JsonWriterFactory, partition: i32, ids: Vec>) -> CommitMessage { // Executors reconstruct the factory without the source or driver write. let factory: JsonWriterFactory = serde_json::from_slice(&serde_json::to_vec(factory).unwrap()).unwrap(); - let mut writer = factory.open(TaskContext::detached()).unwrap(); + let mut writer = open(&factory, partition); writer.write(batch(ids)).unwrap(); let message = writer.commit().unwrap(); serde_json::from_slice(&serde_json::to_vec(&message).unwrap()).unwrap() @@ -580,8 +614,8 @@ mod tests { } fn commit_ids(source: &JsonSink, mode: WriteMode, ids: Vec>) { - let mut write = source.write(sink_schema(), mode).unwrap(); - let message = stage(&write.factory().unwrap(), ids); + let (mut write, factory) = new_write(source, mode).unwrap(); + let message = stage(&factory, 0, ids); write.commit(vec![message]).unwrap(); } @@ -622,14 +656,10 @@ mod tests { #[test] fn driver_publishes_only_winning_attempts_and_removes_unreported_files() { let directory = TempDirectory::new(); - let mut write = directory - .source() - .write(sink_schema(), WriteMode::Append) - .unwrap(); - let factory = write.factory().unwrap(); - let winning = stage(&factory, vec![Some(1)]); - let _losing = stage(&factory, vec![Some(2)]); - let _missing_message = stage(&factory, vec![Some(3)]); + let (mut write, factory) = new_write(&directory.source(), WriteMode::Append).unwrap(); + let winning = stage(&factory, 0, vec![Some(1)]); + let _losing = stage(&factory, 0, vec![Some(2)]); + let _missing_message = stage(&factory, 1, vec![Some(3)]); assert_eq!(fs::read_dir(&write.staging).unwrap().count(), 3); write.commit(vec![winning]).unwrap(); assert!(!write.staging.exists()); @@ -638,7 +668,19 @@ mod tests { committed_rows(&directory.0), vec![serde_json::json!({"id": 1, "label": "row-1"})] ); - assert!(factory.open(TaskContext::detached()).is_err()); + let task = TaskContext::for_testing().build(); + assert!(factory + .open(WriterContext::new(sink_schema(), task)) + .is_err()); + } + + #[test] + fn commit_rejects_a_message_at_the_position_of_another_partition() { + let directory = TempDirectory::new(); + let (mut write, factory) = new_write(&directory.source(), WriteMode::Append).unwrap(); + let message = stage(&factory, 1, vec![Some(1)]); + assert!(write.commit(vec![message]).is_err()); + assert!(!directory.0.join("manifest.json").exists()); } #[test] @@ -647,11 +689,10 @@ mod tests { let source = directory.source(); commit_ids(&source, WriteMode::Append, vec![Some(1)]); let before = fs::read(directory.0.join("manifest.json")).unwrap(); - let mut write = source.write(sink_schema(), WriteMode::Truncate).unwrap(); - let factory = write.factory().unwrap(); - let message = stage(&factory, vec![Some(2)]); - let _missing = stage(&factory, vec![Some(3)]); - let mut incomplete = factory.open(TaskContext::detached()).unwrap(); + let (mut write, factory) = new_write(&source, WriteMode::Truncate).unwrap(); + let message = stage(&factory, 0, vec![Some(2)]); + let _missing = stage(&factory, 1, vec![Some(3)]); + let mut incomplete = open(&factory, 2); incomplete.write(batch(vec![Some(4)])).unwrap(); write.abort(vec![Some(message), None]).unwrap(); assert!(!write.staging.exists()); @@ -659,7 +700,10 @@ mod tests { assert!(!directory.0.join(".write-lock").exists()); assert_eq!(fs::read(directory.0.join("manifest.json")).unwrap(), before); assert_eq!(committed_rows(&directory.0).len(), 1); - assert!(factory.open(TaskContext::detached()).is_err()); + let task = TaskContext::for_testing().build(); + assert!(factory + .open(WriterContext::new(sink_schema(), task)) + .is_err()); } #[test] @@ -668,8 +712,8 @@ mod tests { let source = directory.source(); commit_ids(&source, WriteMode::Append, vec![Some(1)]); let before = fs::read(directory.0.join("manifest.json")).unwrap(); - let mut write = source.write(sink_schema(), WriteMode::Truncate).unwrap(); - let message = stage(&write.factory().unwrap(), vec![Some(2)]); + let (mut write, factory) = new_write(&source, WriteMode::Truncate).unwrap(); + let message = stage(&factory, 0, vec![Some(2)]); // Force create_new to fail after the operation directory has been sealed. fs::write(&write.temporary_manifest, b"publication failure").unwrap(); assert!(write.commit(vec![message]).is_err()); @@ -685,29 +729,29 @@ mod tests { fn concurrent_write_is_rejected_and_drop_releases_staging_and_lock() { let directory = TempDirectory::new(); let source = directory.source(); - let mut write = source.write(sink_schema(), WriteMode::Append).unwrap(); - let _message = stage(&write.factory().unwrap(), vec![Some(1)]); + let (write, factory) = new_write(&source, WriteMode::Append).unwrap(); + let _message = stage(&factory, 0, vec![Some(1)]); let staging = write.staging.clone(); - assert!(source.write(sink_schema(), WriteMode::Append).is_err()); + assert!(new_write(&source, WriteMode::Append).is_err()); drop(write); assert!(!staging.exists()); assert!(!directory.0.join(".write-lock").exists()); - assert!(source.write(sink_schema(), WriteMode::Append).is_ok()); + assert!(new_write(&source, WriteMode::Append).is_ok()); } #[test] fn readonly_removes_all_write_capabilities_without_creating_files() { let directory = TempDirectory::new(); let path = directory.0.join("not-created"); - let source = JsonSink::new(Options::from([ - ("path".to_owned(), path.to_str().unwrap().to_owned()), - ("readonly".to_owned(), "true".to_owned()), + let source = JsonSink::new(info(&[ + ("path", path.to_str().unwrap()), + ("readonly", "true"), ])) .unwrap(); - assert!(source - .capabilities(&[SourceCapability::BatchWrite, SourceCapability::Truncate]) - .is_empty()); - assert!(source.write(sink_schema(), WriteMode::Append).is_err()); + let capabilities = source.capabilities(Capabilities::all()); + assert!(!capabilities.contains(SourceCapability::BatchWrite)); + assert!(!capabilities.contains(SourceCapability::Truncate)); + assert!(new_write(&source, WriteMode::Append).is_err()); assert!(!path.exists()); } } diff --git a/examples/src/main/rust/native-data-source/README.md b/examples/src/main/rust/native-data-source/README.md index 3b2b789582f3f..85a62b7be5f6c 100644 --- a/examples/src/main/rust/native-data-source/README.md +++ b/examples/src/main/rust/native-data-source/README.md @@ -19,18 +19,18 @@ limitations under the License. This example implements batch reads using the [Rust native data source SDK][sdk]. The connector in [src/lib.rs](src/lib.rs) exports `rust_sdk_range` through the SDK's -`export_data_source!` macro. The macro exports the C entry point `spark_ds_init`; the connector -does not use JNI, raw pointers, or unsafe code. +`export_data_source!(RangeSource, batch_read)` macro. The macro exports the C entry point +`spark_ds_init`; the connector does not use JNI, raw pointers, or unsafe code. The example targets C ABI version 1 from [Spark PR #59424][native-source-pr], pinned to commit -`534dfca15723c7208061b1e9e4d80c5757360a2a`. The SDK and example build without Spark or a JVM. +`74aa0dffc53463034fa6549796c4dd599e9705de`. The SDK and example build without Spark or a JVM. Running the Spark smoke script requires a distribution built with that native data source framework and Java 22 or later on the driver and executors. The JNI SDK prototype in PR #59442 targets a different interface. -Registering the source for batch reads derives `SourceCapability::BatchRead` and the native +Registering the source for batch reads derives the `BatchRead` capability and the native mask `SPARK_DS_BATCH_READ` (`1`); the connector does not maintain a separate capability list. -Connectors can restrict registered operations through `DataSource::capabilities()` based on +Connectors can narrow registered operations through `DataSource::capabilities()` based on options. This example does not implement writes or Structured Streaming. The companion [file sink](../native-data-sink/README.md) demonstrates batch append and overwrite. During analysis, Spark creates and closes a source with the table options to inspect its @@ -58,9 +58,10 @@ options. A source created just for capability discovery is closed without creati Spark normalizes option names to lower case. Unknown options are rejected. The source evaluates `=`, `<=>`, `<`, `<=`, `>`, and `>=` comparisons between `id` and non-null -integer literals, including conjunctions of supported predicates. It reports these predicates -as `Exact`; Spark still evaluates all predicates again after reading. Unknown expressions and -unsupported literal types remain `Unsupported`. +integer literals, which `Predicate::literal_value` decodes and checks against the range of their +type, including conjunctions of supported predicates. It reports these predicates as `Exact`; +Spark still evaluates all predicates again after reading. Unknown expressions and unsupported +literal types remain `Unsupported`. Spark pushes a limit only when no residual filter remains, which excludes filtered scans in the pinned framework. The connector can apply a global limit by shortening the selected range @@ -69,12 +70,15 @@ zero-column batches for operations such as `COUNT(*)`. ## Planning and reading -`RangeSource` creates a driver-side `RangeScan`. Pushdown methods update that builder, and -`plan(self)` consumes it to produce a `ScanPlan`. The plan contains: +`RangeSource::new_scan` creates a driver-side `RangeScan`. Pushdown methods update that builder, +and `plan(self)` consumes it to produce a `ScanPlan`. The plan contains: -- a serializable `RangeReaderFactory` with output columns and batch size; +- a serializable `RangeReaderFactory` with the batch size; - disjoint `RangePartition` values with inclusive starts and exclusive ends. +The SDK serializes the negotiated schema with the factory, so the reader takes its output +columns, in order, from `ReaderContext::schema` instead of carrying them in the factory. + This API matches the framework's separation of driver planning and executor reading. Spark's `NativeBatchPlanner` owns the native handle; its `NativePartitionReaderFactory` carries the serialized Rust factory and schema to executors. Java `ScanBuilder.build()` caches a scan but @@ -83,8 +87,9 @@ native `partitions` and then `serialize_reader`, which share the SDK's cached pl closes the native driver handle. The Rust API and C ABI layout remain unchanged. The SDK serializes this data for the executors. Each factory opens a lazy Arrow reader using a -partition and a lifetime-bound `TaskContext`. The example checks task cancellation before -generating each batch and keeps memory bounded by the configured batch size. Neither the factory +partition and a `ReaderContext`, which carries the negotiated schema and a lifetime-bound +`TaskContext`. The example checks task cancellation before generating each batch and keeps +memory bounded by the configured batch size. Neither the factory nor a partition contains a driver handle or open connection. A storage connector can also put a snapshot identifier in its factory to make retried reads refer to the same data. diff --git a/examples/src/main/rust/native-data-source/src/lib.rs b/examples/src/main/rust/native-data-source/src/lib.rs index e29910f3fd0e8..50bbdde517541 100644 --- a/examples/src/main/rust/native-data-source/src/lib.rs +++ b/examples/src/main/rust/native-data-source/src/lib.rs @@ -24,8 +24,9 @@ use arrow_array::{ArrayRef, Int64Array, RecordBatch, RecordBatchIterator, Record use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; use serde::{Deserialize, Serialize}; use spark_native_source::{ - BatchReader, BatchReaderFactory, BatchSource, DataSource, FilterSupport, LimitSupport, Options, - Predicate, Result, ScanBuilder, ScanPlan, TaskContext, + BatchReadSource, CreateInfo, DataSource, FilterSupport, LimitSupport, LiteralValue, Options, + PartitionReader, PartitionReaderFactory, Predicate, ReaderContext, Result, ScanBuilder, + ScanPlan, }; pub struct RangeSource { @@ -40,13 +41,14 @@ pub struct RangeScan { upper: i128, num_partitions: usize, limit: Option, + columns: Vec, factory: RangeReaderFactory, } /// Reader configuration travels to each executor; it contains no live objects. +/// The SDK passes the negotiated columns to the reader, so the factory omits them. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct RangeReaderFactory { - columns: Vec, batch_size: usize, } @@ -57,7 +59,8 @@ pub struct RangePartition { } impl DataSource for RangeSource { - fn new(options: Options) -> Result { + fn new(info: CreateInfo) -> Result { + let options = info.options; for key in options.keys() { if !matches!( key.as_str(), @@ -89,10 +92,10 @@ impl DataSource for RangeSource { } } -impl BatchSource for RangeSource { - type Scan = RangeScan; +impl BatchReadSource for RangeSource { + type ScanBuilder = RangeScan; - fn scan(&self, schema: SchemaRef) -> Result { + fn new_scan(&self, schema: SchemaRef) -> Result { let columns: Vec = schema.fields().iter().map(|f| f.name().clone()).collect(); let expected = output_schema(&columns)?; for (field, expected) in schema.fields().iter().zip(expected.fields()) { @@ -110,27 +113,32 @@ impl BatchSource for RangeSource { upper: i128::from(self.end), num_partitions: self.num_partitions, limit: None, + columns, factory: RangeReaderFactory { - columns, batch_size: self.batch_size, }, }) } } -impl BatchReaderFactory for RangeReaderFactory { +impl PartitionReaderFactory for RangeReaderFactory { type Partition = RangePartition; fn open<'task>( &self, partition: RangePartition, - task: TaskContext<'task>, - ) -> Result> { - let factory = self.clone(); - if factory.batch_size == 0 || partition.start > partition.end { + context: ReaderContext<'task>, + ) -> Result> { + let batch_size = self.batch_size; + if batch_size == 0 || partition.start > partition.end { return Err("Invalid serialized range plan".into()); } - let schema = output_schema(&factory.columns)?; + // The negotiated schema has the requested columns in order, with the types + // that the scan validated. + let schema = context.schema; + let columns: Vec = schema.fields().iter().map(|f| f.name().clone()).collect(); + output_schema(&columns)?; + let task = context.task; let batch_schema = Arc::clone(&schema); let mut current = partition.start; let batches = std::iter::from_fn(move || { @@ -143,11 +151,10 @@ impl BatchReaderFactory for RangeReaderFactory { "Range read was cancelled".into(), ))); } - let next = (i128::from(current) + factory.batch_size as i128) - .min(i128::from(partition.end)) as i64; + 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 columns: Vec = factory - .columns + let arrays: Vec = columns .iter() .map(|column| match column.as_str() { "id" => Arc::new(Int64Array::from_iter_values(current..next)) as ArrayRef, @@ -161,7 +168,7 @@ impl BatchReaderFactory for RangeReaderFactory { // Explicit row counts also preserve COUNT(*) with a zero-column projection. Some(RecordBatch::try_new_with_options( Arc::clone(&batch_schema), - columns, + arrays, &RecordBatchOptions::new().with_row_count(Some(num_rows)), )) }); @@ -193,14 +200,11 @@ impl ScanBuilder for RangeScan { } fn prune_columns(&mut self, columns: &[String]) -> Result { - if columns - .iter() - .any(|column| !self.factory.columns.contains(column)) - { + if columns.iter().any(|column| !self.columns.contains(column)) { return Err("Projection contains a column outside the read schema".into()); } output_schema(columns)?; - self.factory.columns = columns.to_vec(); + self.columns = columns.to_vec(); Ok(true) } @@ -225,10 +229,7 @@ impl ScanBuilder for RangeScan { start = end; } } - Ok(ScanPlan { - factory: self.factory, - partitions, - }) + Ok(ScanPlan::new(self.factory, partitions)) } } @@ -308,39 +309,36 @@ fn is_id(predicate: &Predicate) -> bool { matches!(predicate, Predicate::Column { name } if name.len() == 1 && name[0] == "id") } +/// A non-null tinyint, smallint, int, or bigint literal, decoded by the SDK. fn integer_literal(predicate: &Predicate) -> Option { - let Predicate::Literal { data_type, value } = predicate else { - return None; - }; - let value = value.as_i64()?; - match data_type.as_str() { - "tinyint" if i8::try_from(value).is_ok() => Some(value), - "smallint" if i16::try_from(value).is_ok() => Some(value), - "int" if i32::try_from(value).is_ok() => Some(value), - "bigint" => Some(value), + match predicate.literal_value()? { + LiteralValue::Integer(value) => Some(value), _ => None, } } -spark_native_source::export_data_source!(RangeSource); +spark_native_source::export_data_source!(RangeSource, batch_read); #[cfg(test)] mod tests { use super::*; use serde_json::json; + use spark_native_source::TaskContext; + + fn info(values: &[(&str, &str)]) -> CreateInfo { + let options: Options = values + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(); + CreateInfo::new("rust_sdk_range", options) + } fn source(values: &[(&str, &str)]) -> RangeSource { - RangeSource::new( - values - .iter() - .map(|(k, v)| (k.to_string(), v.to_string())) - .collect(), - ) - .unwrap() + RangeSource::new(info(values)).unwrap() } fn scan(source: &RangeSource) -> RangeScan { - source.scan(source.schema().unwrap()).unwrap() + source.new_scan(source.schema().unwrap()).unwrap() } fn comparison(name: &str, value: i64, reversed: bool) -> Predicate { @@ -362,12 +360,19 @@ mod tests { } } - fn read(plan: ScanPlan) -> Vec { + /// Reads every partition with the schema that the SDK negotiated for `columns`. + fn read(plan: ScanPlan, columns: &[&str]) -> Vec { + let columns: Vec = columns.iter().map(|column| column.to_string()).collect(); + let schema = output_schema(&columns).unwrap(); plan.partitions .into_iter() - .flat_map(|partition| { + .enumerate() + .flat_map(|(index, partition)| { + let task = TaskContext::for_testing() + .partition_id(index as i32) + .build(); plan.factory - .open(partition, TaskContext::detached()) + .open(partition, ReaderContext::new(schema.clone(), task)) .unwrap() .map(|batch| batch.unwrap()) .collect::>() @@ -375,6 +380,8 @@ mod tests { .collect() } + const ALL: &[&str] = &["id", "label"]; + fn ids(batches: &[RecordBatch]) -> Vec { batches .iter() @@ -409,7 +416,7 @@ mod tests { assert!(scan.prune_columns(&["label".into(), "id".into()]).unwrap()); let plan = scan.plan().unwrap(); assert_eq!(plan.partitions.len(), 3); - let batches = read(plan); + let batches = read(plan, &["label", "id"]); assert_eq!(ids(&batches), (13..18).collect::>()); for batch in &batches { assert!(batch.num_rows() <= 2); @@ -456,7 +463,7 @@ mod tests { }) .collect(); assert_eq!( - ids(&read(scan.plan().unwrap())), + ids(&read(scan.plan().unwrap(), ALL)), expected, "range={start}..{end}, op={op}, value={literal}, reverse={reversed}" ); @@ -499,7 +506,7 @@ mod tests { vec![FilterSupport::Unsupported; 4] ); assert_eq!( - ids(&read(scan.plan().unwrap())), + ids(&read(scan.plan().unwrap(), ALL)), (0..10).collect::>() ); } @@ -516,7 +523,7 @@ mod tests { first.push_filters(&[and]).unwrap(), vec![FilterSupport::Exact] ); - assert_eq!(ids(&read(first.plan().unwrap())), vec![3, 4, 5, 6]); + assert_eq!(ids(&read(first.plan().unwrap(), ALL)), vec![3, 4, 5, 6]); let mut empty = scan(&source); empty .push_filters(&[comparison(">", 5, false), comparison("<", 3, false)]) @@ -538,10 +545,7 @@ mod tests { serde_json::from_slice(&serde_json::to_vec(partition).unwrap()).unwrap() }) .collect(); - let batches = read(ScanPlan { - factory, - partitions, - }); + let batches = read(ScanPlan::new(factory, partitions), &[]); assert_eq!( batches .iter() @@ -563,7 +567,7 @@ mod tests { limited.push_limit(5).unwrap(); let plan = limited.plan().unwrap(); assert_eq!(plan.partitions.len(), 2); - assert_eq!(ids(&read(plan)), vec![0, 1]); + assert_eq!(ids(&read(plan, ALL)), vec![0, 1]); let empty = source(&[("start", "5"), ("end", "5")]); assert!(scan(&empty).plan().unwrap().partitions.is_empty()); } @@ -578,13 +582,13 @@ mod tests { ("start", "not-an-integer"), ("unknown", "1"), ] { - assert!(RangeSource::new([(key.into(), value.into())].into()).is_err()); + assert!(RangeSource::new(info(&[(key, value)])).is_err()); } let source = source(&[]); let invalid = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); - assert!(source.scan(invalid).is_err()); + assert!(source.new_scan(invalid).is_err()); let labels = output_schema(&["label".into()]).unwrap(); - let mut scan = source.scan(labels).unwrap(); + let mut scan = source.new_scan(labels).unwrap(); assert!(scan.prune_columns(&["id".into()]).is_err()); assert!(output_schema(&["id".into(), "id".into()]).is_err()); } @@ -598,9 +602,11 @@ mod tests { ("batchsize", "2"), ]); let plan = scan(&source).plan().unwrap(); + let context = + ReaderContext::new(source.schema().unwrap(), TaskContext::for_testing().build()); let mut reader = plan .factory - .open(plan.partitions[0].clone(), TaskContext::detached()) + .open(plan.partitions[0].clone(), context) .unwrap(); let first = reader.next().unwrap().unwrap(); let second = reader.next().unwrap().unwrap(); diff --git a/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c b/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c index f6f71b3edfc55..e20497c372e75 100644 --- a/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c +++ b/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c @@ -80,7 +80,7 @@ static bool unavailable_callback(const SparkDsTaskContext* task) { return false; } -static SparkDsTaskContext task_context(bool short_context) { +static SparkDsTaskContext task_context(void) { SparkDsTaskContext task = { .size = sizeof(SparkDsTaskContext), .stage_id = 71, @@ -90,10 +90,6 @@ static SparkDsTaskContext task_context(bool short_context) { .task_attempt_id = INT64_C(12345678901), .is_interrupted = is_interrupted }; - if (short_context) { - task.size = offsetof(SparkDsTaskContext, is_interrupted); - task.is_interrupted = unavailable_callback; - } return task; } @@ -288,11 +284,10 @@ static void consume_plan( int64_t first_id, int64_t expected_rows, const char* const* columns, - int64_t column_count, - bool short_context) { + int64_t column_count) { int64_t rows = 0; for (size_t partition = 0; partition < plan->partitions.length; partition++) { - SparkDsTaskContext task = task_context(short_context); + SparkDsTaskContext task = task_context(); expected_task = &task; interrupted = false; struct ArrowArrayStream stream = open_stream(plan, partition, &task); @@ -364,15 +359,15 @@ static void check_independent_planners( Plan second_plan = finish_plan(second); /* Both driver planners are already closed; executor reads only get the copied bytes. */ - consume_plan(&second_plan, 13, 12, id_only, 1, false); - consume_plan(&first_plan, 10, 12, label_only, 1, false); + consume_plan(&second_plan, 13, 12, id_only, 1); + consume_plan(&first_plan, 10, 12, label_only, 1); free_plan(&second_plan); free_plan(&first_plan); } static void check_cancellation(const Plan* plan) { CHECK(plan->partitions.length == 1, "Cancellation test requires one partition"); - SparkDsTaskContext task = task_context(false); + SparkDsTaskContext task = task_context(); expected_task = &task; interrupted = false; interruption_checks = 0; @@ -446,8 +441,22 @@ static void stream_call_on_thread(StreamCall* call) { CHECK(pthread_join(thread, NULL) == 0, "Cannot join a stream callback thread"); } +/* Every host has the fields of version 1: a shorter context is an error, not partly read. */ +static void check_short_task_context(const Plan* plan) { + SparkDsTaskContext task = task_context(); + task.size = offsetof(SparkDsTaskContext, is_interrupted); + task.is_interrupted = unavailable_callback; + struct ArrowArrayStream stream = {0}; + const Bytes* bytes = &plan->partitions.values[0]; + CHECK(library->read(&task, plan->state.data, plan->state.length, bytes->data, bytes->length, + &stream) != 0, "A task context shorter than version 1 must be rejected"); + const char* error = library->last_error(); + CHECK(error != NULL && strstr(error, "version 1") != NULL, "Missing task context error"); + CHECK(stream.release == NULL, "A failed read must not initialize the stream"); +} + static void check_cross_thread_stream(const Plan* plan) { - SparkDsTaskContext task = task_context(false); + SparkDsTaskContext task = task_context(); expected_task = &task; interrupted = false; interruption_checks = 0; @@ -570,26 +579,27 @@ int main(int argc, char** argv) { check_independent_planners(keys, values, 4, lower, upper); Plan filtered = plan_scan(keys, values, 4, predicates, uses, 2, -1, reordered, 2); - consume_plan(&filtered, 13, 9, reordered, 2, false); + consume_plan(&filtered, 13, 9, reordered, 2); free_plan(&filtered); Plan residual = plan_scan(keys, values, 4, predicates, uses, 3, -1, id_only, 1); - consume_plan(&residual, 13, 9, id_only, 1, false); + consume_plan(&residual, 13, 9, id_only, 1); free_plan(&residual); Plan limited = plan_scan(keys, values, 4, NULL, NULL, 0, 5, id_only, 1); - consume_plan(&limited, 10, 5, id_only, 1, false); + consume_plan(&limited, 10, 5, id_only, 1); free_plan(&limited); const char* short_keys[] = {"end", "numpartitions", "batchsize"}; const char* short_values[] = {"7", "2", "2"}; Plan empty = plan_scan(short_keys, short_values, 3, NULL, NULL, 0, -1, NULL, 0); - consume_plan(&empty, 0, 7, NULL, 0, false); + consume_plan(&empty, 0, 7, NULL, 0); free_plan(&empty); const char* one_partition[] = {"20", "1", "2"}; Plan cancellable = plan_scan(short_keys, one_partition, 3, NULL, NULL, 0, -1, NULL, -1); - consume_plan(&cancellable, 0, 20, original, 2, true); + consume_plan(&cancellable, 0, 20, original, 2); + check_short_task_context(&cancellable); check_cancellation(&cancellable); check_cross_thread_stream(&cancellable); free_plan(&cancellable); diff --git a/sql/native/rust/README.md b/sql/native/rust/README.md index f1120d937a6c6..4e63e876829a9 100644 --- a/sql/native/rust/README.md +++ b/sql/native/rust/README.md @@ -19,7 +19,7 @@ limitations under the License. This experimental SDK implements **batch reads and writes** over the C ABI proposed in [Spark PR #59424](https://github.com/apache/spark/pull/59424), pinned to commit -`534dfca15723c7208061b1e9e4d80c5757360a2a`. Spark calls the library through the +`74aa0dffc53463034fa6549796c4dd599e9705de`. Spark calls the library through the Java Foreign Function and Memory (FFM) API. This PR contains the SDK and example, not the Spark framework; running the Spark example requires that framework and Java 22 or newer on the driver and executors. @@ -37,21 +37,27 @@ to include the example. ## Connector API and lifecycle A connector implements [`DataSource`](src/lib.rs) for construction, schema discovery, -and optional capability restrictions. Spark can create and immediately drop a source -just to discover capabilities; schema inference or an operation need not follow. +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. Read support uses three traits: -* `BatchSource`: create an independent scan builder. +* `BatchReadSource`: create an independent scan builder with `new_scan`. * `ScanBuilder`: negotiate predicates, limit, and projection, then consume itself in `plan()` to return a typed `ScanPlan`. -* `BatchReaderFactory`: serialize executor configuration and open one lazy Arrow - reader from a typed partition and a borrowed `TaskContext`. +* `PartitionReaderFactory`: serialize executor configuration and open one lazy Arrow + reader from a typed partition and a `ReaderContext`, which carries the negotiated + schema and a borrowed `TaskContext`. Write support is independent: a write-only connector implements `BatchWriteSource`, `BatchWrite`, `DataWriterFactory`, and `DataWriter`, without implementing reads. The write lifecycle is described below. +The inputs of these methods are `#[non_exhaustive]` structs (`CreateInfo`, `WriteInfo`, +`ReaderContext`, and `WriterContext`), and the plans (`ScanPlan` and `WritePlan`) have +constructors, so that new ABI fields can reach connectors without breaking them. + 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 @@ -59,12 +65,12 @@ must read the same snapshot. Multiple task attempts open independent readers. ```text Driver: - BatchSource -> ScanBuilder -> ScanPlan { factory, partitions } - | | - mutable negotiation immutable serialized descriptors - | -Executor, per task: v - restored factory + partition + TaskContext -> Arrow reader -> Arrow C Stream + BatchReadSource -> 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 @@ -83,29 +89,49 @@ panic. 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 restores the Rust factory from those bytes without creating or accessing -the driver registry. Closing the driver handles cannot invalidate factory or -partition bytes already copied to Spark, including for retried tasks. +point restores the Rust factory and the negotiated schema 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. Task IDs -are optional for compatibility with shorter ABI structures. `is_interrupted()` -checks Spark's live cancellation callback, not a snapshot of its initial value. -Check it during long operations; blocking I/O still requires timeouts. +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. `is_interrupted()` checks Spark's live cancellation callback, not a snapshot +of its initial value. Check it during long operations; blocking I/O still requires +timeouts. `TaskContext::for_testing()` builds a context with chosen IDs and an +optional interruption flag for tests outside Spark. Register one provider in the connector's `cdylib` crate: ```rust -spark_native_source::export_data_source!(RangeSource); +spark_native_source::export_data_source!(RangeSource, batch_read); ``` 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. `spark_native_source::log` writes to the driver or executor log through +Spark from any thread, and drops messages outside Spark, such as in unit tests. + +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 registered `DataSource` can choose the source to create; the macro registers a +single Rust type per library. + The implementation separates public connector traits, driver handle/planning state in `bridge.rs`, Arrow ownership in `stream.rs`, and C bindings/adapters in -`abi.rs` and `export.rs`. It does not change Spark's native ABI or JVM classes. +`abi.rs` and `export.rs`. The macro reaches them through a hidden `__private` +module, which is not part of the API. It does not change Spark's native ABI or JVM +classes. ## Source capabilities @@ -114,42 +140,47 @@ installs their native callbacks and derives the default capability set: ```rust export_data_source!(RangeSource, batch_read); -export_data_source!(FileSink, batch_write, truncate); +export_data_source!(FileSink, batch_write); export_data_source!(ReadWriteSource, batch_read, batch_write); ``` -These are alternative registrations; invoke the macro only once per library. -The one-argument form is shorthand for `batch_read`. `truncate` is optional and -requires `batch_write`; without it the source supports only append writes. +These are alternative registrations; invoke the macro only once per library, with +its operations in any order. A `BatchWriteSource` opts into overwrite support with +`const SUPPORTS_TRUNCATE: bool = true`, next to the `new_write` that handles +`WriteMode::Truncate`; without it the source supports only append writes. Unregistered operation families have null native callbacks. The registration -requires the corresponding Rust traits at compile time. +requires the corresponding Rust traits at compile time. 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 registered set for the current options by overriding this `DataSource` method: ```rust -fn capabilities(&self, supported: &[SourceCapability]) -> Vec { - supported.iter().copied().filter(|capability| { - !self.read_only || !matches!(capability, - SourceCapability::BatchWrite | SourceCapability::Truncate) - }).collect() +fn capabilities(&self, registered: Capabilities) -> Capabilities { + if self.read_only { + // Also removes Truncate, which requires batch writes. + registered.without(SourceCapability::BatchWrite) + } else { + registered + } } ``` -The SDK evaluates this once during `create`, validates the declaration, and returns -its native flag mask. Duplicate entries are ignored. Enabling an unregistered -operation is an error; `Truncate` also requires `BatchWrite`. An empty set permits -schema discovery and close, but rejects operation creation before calling the -connector. The declaration stays fixed for that instance. Streaming capabilities -are not supported yet. Predicate, limit, and projection support remain per-scan -contracts. +The SDK evaluates this once during `create` and keeps only registered operations, +so a source can narrow its capabilities but never enable an unregistered operation. +An empty set permits schema discovery and close, but rejects operation creation +before calling the connector. The declaration stays fixed for that instance. +`Capabilities::all()` exists to test this method outside Spark. Streaming +capabilities are not supported yet. Predicate, limit, and projection support 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 `scan` will follow. Executor resources -belong in `BatchReaderFactory::open`. +immediately; do not assume that `schema` or `new_scan` will follow. Executor +resources belong in `PartitionReaderFactory::open`. The pinned framework propagates native `BATCH_WRITE` and `TRUNCATE` flags to `Table.capabilities()`, so a read-only source rejects batch writes during @@ -162,36 +193,39 @@ flags are not a literal copy of the native mask declared by the Rust source. ```text Driver: - DataSource + BatchWriteSource -> BatchWrite coordinator - | factory() | commit(messages) / abort(messages) - v ^ - serialized factory + schema | typed messages - | | -Executor, per task attempt: v | - DataWriterFactory::open(TaskContext) + DataSource + BatchWriteSource -> WritePlan { coordinator, factory } + | ^ commit(messages) / abort(messages) + serialized factory + schema | typed messages + | | +Executor, per task attempt: v | + DataWriterFactory::open(WriterContext) -> DataWriter::write(RecordBatch)* -> commit() / abort() -> Drop ``` -`BatchWriteSource::write(schema, mode)` validates the input and returns an -independent driver coordinator. `BatchWrite::factory()` produces serializable -executor configuration; the SDK caches the encoded factory and input schema, -including a terminal error if creation or serialization fails. Restoring an -executor writer uses only those bytes, without creating a driver registry. +`BatchWriteSource::new_write(WriteInfo)` validates the input schema and returns a +`WritePlan`: an independent driver coordinator and the serializable executor +configuration of its task writers. The SDK encodes the factory and input schema +once, including a terminal error if serialization fails. Restoring an executor +writer uses only those bytes, without creating a driver registry. -Each `DataWriterFactory::open` returns a task-bound writer. `write` takes ownership -of a `RecordBatch`, so its Arrow buffers can be retained after the callback. The -SDK moves the incoming Arrow array and schema, checks the schema against the -negotiated input, 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. Check cancellation during lengthy writes and use I/O timeouts. +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. The SDK +moves the incoming Arrow array and schema, checks the schema against the negotiated +input, 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. +Check cancellation during lengthy writes and use I/O timeouts. Task `commit()` returns a typed, serde-serializable message. Driver `commit` -receives one winning message per partition; driver `abort` receives -`Vec>`, preserving missing-task positions. A failed task or driver -commit leaves the same object available for abort. The SDK prevents further writes -after a failure or successful commit, caches successful task commit messages, and -avoids repeating successful commit/abort callbacks. Drop only releases resources. +receives one winning message per partition, at the position of its partition; +driver `abort` receives `Vec>`, 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. The SDK prevents further writes after a +failure or successful commit, caches successful task commit messages, and avoids +repeating successful commit/abort callbacks. Drop only releases resources. Connectors therefore do not need to track their own phase; the `# Lifecycle` sections of `BatchWrite` and `DataWriter` list the exact call-order guarantees. @@ -203,11 +237,12 @@ 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. -Spark may close a coordinator without commit/abort if factory serialization fails. -It also creates writers late, after prerequisite AQE shuffle stages, but input -computation can still fail during write tasks. Constructors must therefore avoid -irreversible changes. This API is batch-only: non-null streaming query IDs and -epoch IDs other than `-1` are rejected. +Spark closes a coordinator without commit or abort if factory serialization fails, +so Drop must also release what only that write staged. It also creates writers +late, after prerequisite AQE shuffle stages, but input computation can still fail +during write tasks. Constructors must therefore avoid irreversible changes. This +API is batch-only: non-null streaming query IDs and epoch IDs other than `-1` are +rejected. 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 @@ -216,10 +251,11 @@ not transaction guarantees provided by the SDK. ## Arrow contract -`BatchReaderFactory::open` returns a lifetime-bound `Box`. Return fixed-schema batches lazily with bounded memory. Exported arrays -own their buffers independently of the reader: Spark may release a stream before -releasing its previously returned batches. +`PartitionReaderFactory::open` returns a lifetime-bound `PartitionReader`, a +`Box`. Return fixed-schema batches lazily with bounded +memory, with the columns of `ReaderContext::schema`. Exported arrays own their +buffers independently of the reader: Spark may release a stream before releasing +its previously returned batches. Stream callbacks, including release, may run on different threads, but never concurrently. Readers must be `Send`; they do not need to be `Sync` or tied to @@ -232,13 +268,18 @@ logical types, ignoring nullability and Arrow metadata, accepting the supported string/binary encodings and equivalent timestamp time zones. It rejects unsupported types, schema drift, and nonzero array offsets with explicit errors. It does not silently copy arrays, cast types, or expand dictionary encoding. - -Supported types include signed integers, float32/64, boolean, strings and binary -(including Large and View variants), decimal128 up to precision 38, date32, -microsecond timestamps, supported Spark intervals, null, structs, lists, and -maps. Dictionary encoding, unsigned integers, decimal256, large/fixed-size -lists, and non-microsecond timestamps are outside this prototype's contract. -Normalize sliced arrays when necessary before returning them. +Batches must also respect the nullability of the negotiated schema: as Spark does, +the SDK rejects a null in a column, field, element, or map value that the schema +declares not nullable, unless its parent value is null. + +The supported types are those of version 1 of the ABI: signed integers, +float32/64, boolean, strings and binary (including Large and View variants), decimal128 +up to precision 38, date32, microsecond timestamps with or without a time zone, +structs, lists, and maps. Spark rejects other Spark types, such as intervals, +variant, time, and the null type, before it passes a schema to the library. +Dictionary encoding, unsigned integers, decimal256, large/fixed-size lists, and +non-microsecond timestamps are outside this prototype's contract. Normalize sliced +arrays when necessary before returning them. An empty projection must retain the number of rows even with no columns. The [Range implementation][range-source] demonstrates this together with bounded @@ -256,15 +297,20 @@ Results correspond to input predicates in order: | `Unsupported` | `SPARK_DS_PREDICATE_UNUSED` | No evaluation | The pinned framework rechecks **all** filters, including `Exact`. Unknown -expression names and unsupported types must remain unsupported. The SDK -preserves literal JSON values, including exact 64-bit integers, rather than -converting all numbers to floating point. It is not a generic Spark expression -evaluator; connectors remain responsible for Spark SQL semantics. +expression names and unsupported types must remain unsupported. The SDK parses +each predicate on its own: one that it cannot parse, for example because a later +Spark encodes it differently, reaches `push_filters` as `Predicate::Unknown` and +stays unsupported instead of failing the scan. The SDK preserves literal JSON +values, including exact 64-bit integers, rather than converting all numbers to +floating point. `Predicate::literal_value` decodes a literal by its Spark SQL type: +integers within the range of their type, exact floats, strings, base64 binary, +plain decimal strings, dates, and timestamps. A literal that it does not decode +returns `None`; keep its predicate unsupported. The SDK is not a generic Spark +expression evaluator; connectors remain responsible for Spark SQL semantics. | Limit result | ABI boolean | Meaning | | --- | --- | --- | -| `Unsupported` | `false` | No limit applied | -| `PerPartition` | `false` | A local cap; the total may exceed the limit | +| `Unsupported` | `false` | No promise; a connector may still cap each partition on its own | | `Global` | `true` | The entire scan returns at most the limit | Spark retains its final limit and currently pushes a limit only when the scan @@ -294,21 +340,23 @@ The harness compiles against the unmodified framework header in [`tests/fixtures/spark_datasource.h`](tests/fixtures/spark_datasource.h), copied from the pinned commit above. It calls the actual function table and validates capability-only source discovery and optional slots, independent driver planners, -pushdown and planning, serialized executor reads, cancellation, sequential -cross-thread stream callbacks, and Arrow buffers surviving stream release. -This checks the C boundary without requiring Java; it does not substitute for -Spark planner and task-lifecycle testing. +pushdown and planning, serialized executor reads, rejection of task contexts shorter +than version 1, cancellation, sequential cross-thread stream callbacks, and Arrow +buffers surviving stream release. This checks the C boundary without requiring +Java; it does not substitute for Spark planner and task-lifecycle testing. Rust tests additionally exercise the exported function table with option-dependent -capabilities, empty and duplicate declarations, and unsupported declarations or -panics during capability discovery. They also cover terminal planning -failures, retries after the driver registry is dropped, schema/layout validation, -sized ABI structures, task-context lifetimes, integer boundaries, projections, -and deterministic partition plans. Write ABI tests cover typed factory/message -round trips, retained Arrow buffers, cancellation, failed writes and commits, -serialization failures, missing abort messages, and repeated callbacks. The sink -tests cover append/truncate publication, losing attempts, and cleanup after failed -publication while preserving the previous manifest. +capabilities, narrowed and unregistered declarations, the source name, library +initialization and logging, and panics during capability discovery. They also +cover terminal planning failures, predicates that do not parse, typed literals, +retries after the driver registry is dropped, schema, layout, and nullability +validation, sized ABI structures, task-context lifetimes, integer boundaries, +projections, and deterministic partition plans. Write ABI tests cover typed +factory/message round trips, retained Arrow buffers, cancellation, failed writes and +commits, serialization failures, missing and undecodable messages, and repeated +callbacks. The sink tests cover append/truncate publication, losing attempts, +messages at the wrong partition position, and cleanup after failed publication while +preserving the previous manifest. ## Package and run with Spark @@ -327,11 +375,11 @@ analysis-time rejection for a source instance with writes disabled. ## Scope and limitations Native sources run trusted code in the driver and executors. The SDK translates -provider errors and unwinding panics into C/Arrow errors. Process aborts, -allocation failure, double panics during unwinding, or panics in custom Arrow -buffer destructors cannot be recovered. A reader's or writer's background work -must stop before its task context expires. Native memory is not reserved through -Spark's memory manager. +provider errors, with the messages of their sources, and unwinding panics into +C/Arrow errors. Process aborts, allocation failure, double panics during unwinding, +or panics in custom Arrow buffer destructors cannot be recovered. A reader's or +writer's background work must stop before its task context expires. Native memory +is not reserved through Spark's memory manager. This prototype does not include Structured Streaming, nested projection, custom metrics, memory accounting, a generic filter evaluator, C++ bindings, or diff --git a/sql/native/rust/src/abi.rs b/sql/native/rust/src/abi.rs index 12cc2211a8aac..9a7489ad588a3 100644 --- a/sql/native/rust/src/abi.rs +++ b/sql/native/rust/src/abi.rs @@ -26,8 +26,9 @@ 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::panic_message; -use crate::{FilterSupport, Options, Result, WriteMode}; +use crate::bridge::{describe, panic_message}; +use crate::host::{self, LogFunction}; +use crate::{CreateInfo, FilterSupport, HostInfo, Options, Result, WriteMode}; pub const ABI_VERSION: i32 = 1; pub const BATCH_READ: u64 = 1; @@ -43,7 +44,7 @@ pub struct SparkDsHost { pub size: i64, pub abi_version: i32, pub spark_version: *const c_char, - pub log: Option, + pub log: Option, } #[repr(C)] @@ -161,7 +162,7 @@ thread_local! { /// Converts provider failures and unwinding panics into the C error convention. pub fn call(action: impl FnOnce() -> Result<()>) -> c_int { LAST_ERROR.with(|error| *error.borrow_mut() = None); - let message = match catch_unwind(AssertUnwindSafe(|| action().map_err(|e| e.to_string()))) { + let message = match catch_unwind(AssertUnwindSafe(|| action().map_err(|e| describe(&*e)))) { Ok(Ok(())) => return 0, Ok(Err(message)) => message, Err(payload) => panic_message(payload), @@ -216,29 +217,45 @@ pub fn require_output(pointer: *mut T) -> Result<()> { Ok(()) } +/// Negotiates the ABI version, installs the host's log, and runs the source's `init`. +/// /// # Safety /// `host` and `output` must follow the `spark_ds_init` contract of the C header. pub unsafe fn initialize( host: *const SparkDsHost, output: *mut *const SparkDsLibrary, library: &'static SparkDsLibrary, + init: fn(&HostInfo) -> Result<()>, ) -> c_int { let result = call(|| { require_output(output)?; unsafe { output.write(std::ptr::null()) }; let size = unsafe { struct_size(host)? }; - let version: i32 = - unsafe { field(host.cast(), size, offset_of!(SparkDsHost, abi_version)) } - .ok_or("host does not provide an ABI version")?; + let pointer = host.cast(); + let version: i32 = unsafe { field(pointer, size, offset_of!(SparkDsHost, abi_version)) } + .ok_or("host does not provide an ABI version")?; if version < ABI_VERSION { return Err(format!("host ABI {version} does not support ABI {ABI_VERSION}").into()); } + let log: Option> = + unsafe { field(pointer, size, offset_of!(SparkDsHost, log)) }; + host::install(log.flatten()); + let spark_version: Option<*const c_char> = + unsafe { field(pointer, size, offset_of!(SparkDsHost, spark_version)) }; + let spark_version = spark_version + .filter(|version| !version.is_null()) + .and_then(|version| unsafe { CStr::from_ptr(version) }.to_str().ok()) + .map(str::to_owned); + init(&HostInfo { + abi_version: version, + spark_version, + })?; unsafe { output.write(library) }; Ok(()) }); if result != 0 { if let Ok(size) = unsafe { struct_size(host) } { - let log: Option> = + let log: Option> = unsafe { field(host.cast(), size, offset_of!(SparkDsHost, log)) }; if let Some(Some(log)) = log { unsafe { log(LOG_ERROR, last_error()) }; @@ -281,19 +298,26 @@ pub unsafe fn strings(data: *const *const c_char, length: i64) -> Result Result { +pub unsafe fn create_info(info: *const SparkDsCreateInfo) -> Result { let size = unsafe { struct_size(info)? }; let pointer = info.cast(); + let name: *const c_char = unsafe { field(pointer, size, offset_of!(SparkDsCreateInfo, name)) } + .ok_or("create info does not provide a name")?; + if name.is_null() { + return Err("create info has a null name".into()); + } + let name = unsafe { CStr::from_ptr(name) }.to_str()?.to_owned(); let keys = unsafe { field(pointer, size, offset_of!(SparkDsCreateInfo, option_keys)) } .ok_or("create info does not provide option keys")?; let values = unsafe { field(pointer, size, offset_of!(SparkDsCreateInfo, option_values)) } .ok_or("create info does not provide option values")?; let count = unsafe { field(pointer, size, offset_of!(SparkDsCreateInfo, num_options)) } .ok_or("create info does not provide an option count")?; - Ok(unsafe { strings(keys, count)? } + let options: Options = unsafe { strings(keys, count)? } .into_iter() .zip(unsafe { strings(values, count)? }) - .collect()) + .collect(); + Ok(CreateInfo::new(name, options)) } /// # Safety @@ -412,7 +436,28 @@ mod tests { assert!(unsafe { strings(std::ptr::null(), 0) }.unwrap().is_empty()); let prefix = size_of::() as i64; let info = (&prefix as *const i64).cast::(); - assert!(unsafe { options(info) }.is_err()); + assert!(unsafe { create_info(info) }.is_err()); + } + + #[test] + fn create_info_carries_the_source_name_and_options() { + let keys = [c"path".as_ptr()]; + let values = [c"/tmp/data".as_ptr()]; + let mut info = SparkDsCreateInfo { + size: size_of::() as i64, + name: c"my_source".as_ptr(), + option_keys: keys.as_ptr(), + option_values: values.as_ptr(), + num_options: 1, + }; + let created = unsafe { create_info(&info) }.unwrap(); + assert_eq!(created.name, "my_source"); + assert_eq!( + created.options, + Options::from([("path".into(), "/tmp/data".into())]) + ); + info.name = std::ptr::null(); + assert!(unsafe { create_info(&info) }.is_err()); } #[test] diff --git a/sql/native/rust/src/bridge.rs b/sql/native/rust/src/bridge.rs index cb807126f8ff7..03cbafbd9779b 100644 --- a/sql/native/rust/src/bridge.rs +++ b/sql/native/rust/src/bridge.rs @@ -26,8 +26,9 @@ use arrow_schema::{Schema, SchemaRef}; use serde::{de::DeserializeOwned, Deserialize, Serialize}; use crate::{ - BatchReaderFactory, DataSource, FilterSupport, LimitSupport, Options, Predicate, Result, - ScanBuilder, SourceCapability, TaskContext, WriteMode, + Capabilities, CreateInfo, DataSource, FilterSupport, LimitSupport, LogLevel, + PartitionReaderFactory, Predicate, ReaderContext, Result, ScanBuilder, SourceCapability, + TaskContext, WriteInfo, WriteMode, }; use crate::provider::Provider; @@ -45,8 +46,8 @@ type HandleMap = Mutex>>>; pub struct Bridge { next_handle: AtomicI64, sources: HandleMap>, - planners: HandleMap>, - writers: HandleMap>, + planners: HandleMap>, + writers: HandleMap>, } impl Default for Bridge

{ @@ -70,26 +71,15 @@ impl Bridge

{ Ok(handle) } - pub fn create(&self, options: Options) -> Result<(i64, u64)> { - self.register_source(P::Source::new(options)?) + pub fn create(&self, info: CreateInfo) -> Result<(i64, u64)> { + self.register_source(P::Source::new(info)?) } fn register_source(&self, source: P::Source) -> Result<(i64, u64)> { - let mut capabilities = 0; - for capability in source.capabilities(P::CAPABILITIES) { - if !P::CAPABILITIES.contains(&capability) { - return Err(format!("capability {capability:?} is not registered").into()); - } - capabilities |= match capability { - SourceCapability::BatchRead => crate::abi::BATCH_READ, - SourceCapability::BatchWrite => crate::abi::BATCH_WRITE, - SourceCapability::Truncate => crate::abi::TRUNCATE, - _ => return Err(format!("capability {capability:?} is not implemented").into()), - }; - } - if capabilities & crate::abi::TRUNCATE != 0 && capabilities & crate::abi::BATCH_WRITE == 0 { - return Err("Truncate requires BatchWrite".into()); - } + // Keep only registered operations: a source can narrow them but never add one. + let capabilities = source + .capabilities(P::CAPABILITIES) + .intersection(P::CAPABILITIES); let handle = self.insert( &self.sources, SourceState { @@ -97,7 +87,7 @@ impl Bridge

{ capabilities, }, )?; - Ok((handle, capabilities)) + Ok((handle, capabilities.bits())) } pub fn schema(&self, handle: i64) -> Result { @@ -112,28 +102,40 @@ impl Bridge

{ let source = get(&self.sources, handle)?; let builder = { let source = lock(&source); - if source.capabilities & crate::abi::BATCH_READ == 0 { + if !source.capabilities.contains(SourceCapability::BatchRead) { return Err("source does not support batch reads".into()); } - P::scan(&source.source, schema.clone())? + P::new_scan(&source.source, schema.clone())? }; self.insert(&self.planners, Planner::new(builder, schema)) } + /// Pushes parsed predicates; None stands for one that did not parse, which the + /// connector receives as Unknown and which stays unsupported whatever it returns. pub fn push_predicates( &self, handle: i64, - predicates: &[Predicate], + predicates: Vec>, ) -> Result> { let reader = get(&self.planners, handle)?; let mut reader = lock(&reader); - let support = reader.negotiate(|builder| { - let support = builder.push_filters(predicates)?; + let parsed: Vec = predicates.iter().map(Option::is_some).collect(); + let predicates: Vec = predicates + .into_iter() + .map(|predicate| predicate.unwrap_or(Predicate::Unknown)) + .collect(); + let mut support = reader.negotiate(|builder| { + let support = builder.push_filters(&predicates)?; if support.len() != predicates.len() { return Err("push_filters must return one result per predicate".into()); } Ok(support) })?; + for (support, parsed) in support.iter_mut().zip(parsed) { + if !parsed { + *support = FilterSupport::Unsupported; + } + } reader.has_residual |= support.iter().any(|s| *s != FilterSupport::Exact); Ok(support) } @@ -172,20 +174,23 @@ impl Bridge

{ result } - pub fn create_writer(&self, handle: i64, schema: SchemaRef, mode: WriteMode) -> Result { - stream::validate_schema(&schema)?; + pub fn create_writer(&self, handle: i64, info: WriteInfo) -> Result { + stream::validate_schema(&info.schema)?; let source = get(&self.sources, handle)?; - let writer = { + let schema = info.schema.clone(); + let plan = { let source = lock(&source); - if source.capabilities & crate::abi::BATCH_WRITE == 0 { + if !source.capabilities.contains(SourceCapability::BatchWrite) { return Err("source does not support batch writes".into()); } - if mode == WriteMode::Truncate && source.capabilities & crate::abi::TRUNCATE == 0 { + if info.mode == WriteMode::Truncate + && !source.capabilities.contains(SourceCapability::Truncate) + { return Err("source does not support truncate".into()); } - P::write(&source.source, schema.clone(), mode)? + P::new_write(&source.source, info)? }; - self.insert(&self.writers, DriverWriter::new(writer, schema)) + self.insert(&self.writers, DriverWriter::new(plan, schema)) } pub fn serialized_writer(&self, handle: i64) -> Result>> { @@ -229,15 +234,16 @@ impl Bridge

{ /// `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( +pub unsafe fn open_reader( task: *const crate::abi::SparkDsTaskContext, state: &[u8], partition: &[u8], ) -> Result { - let state: ReaderState = decode(state)?; + let ReaderState { factory, schema }: ReaderState = decode(state)?; + let schema = Arc::new(schema); let task = unsafe { TaskContext::from_raw(task)? }; - let reader = state.factory.open(decode(partition)?, task)?; - unsafe { stream::export(reader, &Arc::new(state.schema)) } + let reader = factory.open(decode(partition)?, ReaderContext::new(schema.clone(), task))?; + unsafe { stream::export(reader, &schema) } } fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { @@ -262,7 +268,7 @@ fn remove(handles: &HandleMap, handle: i64) { struct SourceState { source: S, - capabilities: u64, + capabilities: Capabilities, } struct Planner { @@ -380,11 +386,27 @@ pub(crate) fn provider_call( action: impl FnOnce() -> Result, ) -> std::result::Result { catch_unwind(AssertUnwindSafe(|| { - action().map_err(|error| error.to_string()) + action().map_err(|error| describe(&*error)) })) .unwrap_or_else(|payload| Err(panic_message(payload))) } +/// The message of an error followed by those of its sources, which `Display` of many +/// errors, such as those with added context, omits. +pub(crate) fn describe(error: &(dyn std::error::Error + 'static)) -> String { + let mut message = error.to_string(); + let mut source = error.source(); + while let Some(cause) = source { + let cause_message = cause.to_string(); + if !message.contains(&cause_message) { + message.push_str(": "); + message.push_str(&cause_message); + } + source = cause.source(); + } + message +} + pub(crate) fn panic_message(payload: Box) -> String { let message = if let Some(message) = payload.downcast_ref::() { message.clone() @@ -398,10 +420,21 @@ pub(crate) fn panic_message(payload: Box) -> String { format!("native provider panicked: {message}") } -pub fn parse_predicates(predicates: &[String]) -> Result> { +/// Parses each predicate on its own. One that does not parse, for example because a +/// later Spark encodes it differently, stays unused instead of failing the scan. +pub fn parse_predicates(predicates: &[String]) -> Vec> { predicates .iter() - .map(|text| Ok(serde_json::from_str(text)?)) + .map(|text| match serde_json::from_str(text) { + Ok(predicate) => Some(predicate), + Err(error) => { + crate::log( + LogLevel::Warn, + &format!("Leaving unused a predicate that cannot be parsed ({error}): {text}"), + ); + None + } + }) .collect() } @@ -415,7 +448,7 @@ mod tests { use super::*; use crate::provider::ReadOnly; - use crate::{BatchReader, BatchSource, ScanPlan}; + use crate::{BatchReadSource, PartitionReader, ScanPlan}; #[derive(Clone, Copy)] enum FailureMode { @@ -467,7 +500,7 @@ mod tests { } impl DataSource for Source { - fn new(_: Options) -> Result { + fn new(_: CreateInfo) -> Result { Ok(Self { builds: Arc::default(), limits: Arc::default(), @@ -484,10 +517,10 @@ mod tests { } } - impl BatchSource for Source { - type Scan = Builder; + impl BatchReadSource for Source { + type ScanBuilder = Builder; - fn scan(&self, _: SchemaRef) -> Result { + fn new_scan(&self, _: SchemaRef) -> Result { Ok(Builder { builds: self.builds.clone(), limits: self.limits.clone(), @@ -503,10 +536,15 @@ mod tests { #[derive(Serialize, Deserialize)] struct Factory(u64); - impl BatchReaderFactory for Factory { + impl PartitionReaderFactory for Factory { type Partition = u64; - fn open<'task>(&self, partition: u64, _: TaskContext<'task>) -> Result> { + // Ignores the negotiated schema of the context, to test that the SDK checks it. + fn open<'task>( + &self, + partition: u64, + _: ReaderContext<'task>, + ) -> Result> { let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); let values = Arc::new(Int64Array::from(vec![(self.0 + partition) as i64])); let batch = RecordBatch::try_new(schema.clone(), vec![values])?; @@ -532,13 +570,14 @@ mod tests { fn plan(self) -> Result> { self.builds.fetch_add(1, Ordering::SeqCst); self.failure.fail("planning failed")?; - Ok(ScanPlan { - factory: Factory(self.value), - partitions: vec![2, 7], - }) + Ok(ScanPlan::new(Factory(self.value), vec![2, 7])) } } + fn info() -> CreateInfo { + CreateInfo::new("bridge_test", Default::default()) + } + fn error_message(result: Result) -> String { result .err() @@ -549,7 +588,7 @@ mod tests { #[test] fn source_can_close_before_planning_and_plan_is_built_once() { let bridge = Bridge::>::default(); - let source = Source::new(Options::new()).unwrap(); + let source = Source::new(info()).unwrap(); let builds = source.builds.clone(); let drops = source.drops.clone(); let (handle, _) = bridge.register_source(source).unwrap(); @@ -581,7 +620,7 @@ mod tests { fn residual_filters_prevent_calling_the_providers_limit_method() { for support in [FilterSupport::PruningOnly, FilterSupport::Unsupported] { let bridge = Bridge::>::default(); - let mut source = Source::new(Options::new()).unwrap(); + let mut source = Source::new(info()).unwrap(); source.filter_support = support; source.limit_support = LimitSupport::Global; let limits = source.limits.clone(); @@ -591,7 +630,7 @@ mod tests { .unwrap(); assert_eq!( bridge - .push_predicates(reader, &[Predicate::Unknown]) + .push_predicates(reader, vec![Some(Predicate::Unknown)]) .unwrap(), vec![support] ); @@ -600,11 +639,43 @@ mod tests { } } + #[test] + fn predicates_that_do_not_parse_stay_unsupported_and_residual() { + let bridge = Bridge::>::default(); + let mut source = Source::new(info()).unwrap(); + source.filter_support = FilterSupport::Exact; + source.limit_support = LimitSupport::Global; + let limits = source.limits.clone(); + let (handle, _) = bridge.register_source(source).unwrap(); + let reader = bridge + .create_reader(handle, bridge.schema(handle).unwrap()) + .unwrap(); + let predicates = parse_predicates(&[ + r#"{"type":"function","name":"IS_NULL","children":[{"type":"column","name":"id"}]}"# + .into(), + r#"{"type":"function","name":"IS_NULL","children":[ + {"type":"column","name":["id"]}]}"# + .into(), + ]); + assert!( + predicates[0].is_none(), + "a column name must be an array of parts" + ); + assert!(predicates[1].is_some()); + // The connector reports every predicate as exact, including the Unknown placeholder. + assert_eq!( + bridge.push_predicates(reader, predicates).unwrap(), + vec![FilterSupport::Unsupported, FilterSupport::Exact] + ); + assert!(!bridge.push_limit(reader, 1).unwrap()); + assert_eq!(limits.load(Ordering::SeqCst), 0); + } + #[test] fn only_global_limits_return_true_to_the_abi() { - for support in [LimitSupport::PerPartition, LimitSupport::Global] { + for support in [LimitSupport::Unsupported, LimitSupport::Global] { let bridge = Bridge::>::default(); - let mut source = Source::new(Options::new()).unwrap(); + let mut source = Source::new(info()).unwrap(); source.limit_support = support; let limits = source.limits.clone(); let (handle, _) = bridge.register_source(source).unwrap(); @@ -629,7 +700,7 @@ mod tests { ), ] { let bridge = Bridge::>::default(); - let mut source = Source::new(Options::new()).unwrap(); + let mut source = Source::new(info()).unwrap(); source.failure = failure; let builds = source.builds.clone(); let drops = source.builder_drops.clone(); @@ -659,7 +730,7 @@ mod tests { ), ] { let bridge = Bridge::>::default(); - let mut source = Source::new(Options::new()).unwrap(); + let mut source = Source::new(info()).unwrap(); source.failure = failure; let builds = source.builds.clone(); let limits = source.limits.clone(); @@ -671,7 +742,10 @@ mod tests { assert_eq!(error_message(bridge.push_limit(reader, 5)), expected); assert_eq!(error_message(bridge.push_limit(reader, 1)), expected); - assert_eq!(error_message(bridge.push_predicates(reader, &[])), expected); + assert_eq!( + error_message(bridge.push_predicates(reader, vec![])), + expected + ); assert_eq!(error_message(bridge.planned(reader)), expected); assert_eq!(limits.load(Ordering::SeqCst), 1); assert_eq!(builds.load(Ordering::SeqCst), 0); @@ -684,7 +758,7 @@ mod tests { #[test] fn scans_of_one_source_keep_pushdown_and_cached_plans_independent() { let bridge = Bridge::>::default(); - let mut source = Source::new(Options::new()).unwrap(); + let mut source = Source::new(info()).unwrap(); source.filter_support = FilterSupport::PruningOnly; source.limit_support = LimitSupport::Global; let builds = source.builds.clone(); @@ -699,7 +773,7 @@ mod tests { bridge.close_source(handle).unwrap(); bridge - .push_predicates(second, &[Predicate::Unknown]) + .push_predicates(second, vec![Some(Predicate::Unknown)]) .unwrap(); assert!(bridge.push_limit(first, 5).unwrap()); assert!(!bridge.push_limit(second, 1).unwrap()); @@ -730,7 +804,7 @@ mod tests { #[test] fn executor_retries_use_only_serialized_values_after_the_driver_is_dropped() { let bridge = Bridge::>::default(); - let source = Source::new(Options::new()).unwrap(); + let source = Source::new(info()).unwrap(); let builds = source.builds.clone(); let source_drops = source.drops.clone(); let builder_drops = source.builder_drops.clone(); @@ -778,7 +852,7 @@ mod tests { #[test] fn executor_rejects_a_factory_schema_different_from_the_negotiated_schema() { let bridge = Bridge::>::default(); - let (handle, _) = bridge.create(Options::new()).unwrap(); + let (handle, _) = bridge.create(info()).unwrap(); let reader = bridge .create_reader(handle, bridge.schema(handle).unwrap()) .unwrap(); @@ -804,6 +878,67 @@ mod tests { assert!(error.contains("schema"), "{error}"); } + #[test] + fn capabilities_keep_only_registered_operations() { + struct Narrowing(Capabilities); + + impl DataSource for Narrowing { + fn new(_: CreateInfo) -> Result { + unreachable!("registered directly") + } + + fn capabilities(&self, _: Capabilities) -> Capabilities { + self.0 + } + } + + impl BatchReadSource for Narrowing { + type ScanBuilder = Builder; + + fn new_scan(&self, _: SchemaRef) -> Result { + Err("not scanned".into()) + } + } + + let bridge = Bridge::>::default(); + for (declared, expected) in [ + (Capabilities::all(), crate::abi::BATCH_READ), + (Capabilities::empty(), 0), + (Capabilities::all().without(SourceCapability::BatchRead), 0), + ] { + let (handle, mask) = bridge.register_source(Narrowing(declared)).unwrap(); + assert_eq!(mask, expected, "{declared:?}"); + bridge.close_source(handle).unwrap(); + } + } + + #[test] + fn error_messages_include_their_sources_once() { + #[derive(Debug)] + struct Context(std::io::Error); + + impl std::fmt::Display for Context { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("cannot open the table") + } + } + + impl std::error::Error for Context { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(&self.0) + } + } + + let cause = std::io::Error::new(std::io::ErrorKind::NotFound, "no such file"); + let error: crate::Error = Box::new(Context(cause)); + assert_eq!(describe(&*error), "cannot open the table: no such file"); + let error = provider_call::<()>(|| Err(Box::new(Context(std::io::Error::other("io"))))); + assert_eq!(error.unwrap_err(), "cannot open the table: io"); + // A message that already contains its source is not repeated. + let arrow = arrow_schema::ArrowError::ExternalError(Box::new(std::io::Error::other("io"))); + assert_eq!(describe(&arrow), arrow.to_string()); + } + #[test] fn rejects_a_different_wire_version() { let error = decode::(br#"{"version":2,"value":42}"#).unwrap_err(); diff --git a/sql/native/rust/src/export.rs b/sql/native/rust/src/export.rs index 785ee2e074e50..307239c4c8526 100644 --- a/sql/native/rust/src/export.rs +++ b/sql/native/rust/src/export.rs @@ -13,96 +13,128 @@ // See the License for the specific language governing permissions and // limitations under the License. -/// Exports a provider through version 1 of Spark's C data source ABI. +/// Exports a data source through version 1 of Spark's C data source ABI. /// -/// Call this macro once in a cdylib. The only exported symbol is spark_ds_init; -/// providers implement the safe Rust traits and do not handle C pointers. -/// Register `batch_read`, `batch_write`, or both (in that order). Append `truncate` -/// to opt in to overwrite support. The one-argument form registers batch reads. -/// Registration generates the default capabilities and installs matching callbacks; -/// unsupported operation families have null callbacks. +/// Call this macro once in a cdylib with the operations that the source implements, +/// in any order: +/// +/// ```rust,ignore +/// export_data_source!(RangeSource, batch_read); +/// export_data_source!(FileSink, batch_write); +/// export_data_source!(ReadWriteSource, batch_read, batch_write); +/// ``` +/// +/// `batch_read` requires [`BatchReadSource`](crate::BatchReadSource), and `batch_write` +/// requires [`BatchWriteSource`](crate::BatchWriteSource), whose `SUPPORTS_TRUNCATE` +/// registers truncation. Registration generates the default capabilities and installs +/// matching callbacks; unsupported operation families have null 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) => { - $crate::export_data_source!($source, batch_read); + (@parse $source:ty; [$read:tt $write:tt]; batch_read $($rest:ident)*) => { + $crate::export_data_source!(@parse $source; [true $write]; $($rest)*); + }; + (@parse $source:ty; [$read:tt $write:tt]; batch_write $($rest:ident)*) => { + $crate::export_data_source!(@parse $source; [$read true]; $($rest)*); }; - ($source:ty, batch_read) => { - $crate::export_data_source!(@provider $crate::provider::ReadOnly<$source>); + (@parse $source:ty; [$read:tt $write:tt]; truncate $($rest:ident)*) => { + compile_error!( + "truncate is not an operation: set `const SUPPORTS_TRUNCATE: bool = true` in the \ + BatchWriteSource implementation" + ); }; - ($source:ty, batch_write) => { - $crate::export_data_source!(@provider $crate::provider::WriteOnly<$source, false>); + (@parse $source:ty; [$read:tt $write:tt]; $other:ident $($rest:ident)*) => { + compile_error!(concat!( + "unknown operation `", + stringify!($other), + "`: expected batch_read or batch_write" + )); }; - ($source:ty, batch_write, truncate) => { - $crate::export_data_source!(@provider $crate::provider::WriteOnly<$source, true>); + (@parse $source:ty; [true false];) => { + $crate::export_data_source!(@provider $crate::__private::provider::ReadOnly<$source>); }; - ($source:ty, batch_read, batch_write) => { - $crate::export_data_source!(@provider $crate::provider::ReadWrite<$source, false>); + (@parse $source:ty; [false true];) => { + $crate::export_data_source!(@provider $crate::__private::provider::WriteOnly<$source>); }; - ($source:ty, batch_read, batch_write, truncate) => { - $crate::export_data_source!(@provider $crate::provider::ReadWrite<$source, true>); + (@parse $source:ty; [true true];) => { + $crate::export_data_source!(@provider $crate::__private::provider::ReadWrite<$source>); }; (@provider $provider: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" + ); + mod __spark_native_datasource_exports { use super::*; use std::ffi::{c_char, c_int, c_void}; - use $crate::abi::{self, ArrowArrayStream, ArrowSchema, SparkDsLibrary, SparkDsSink}; - use $crate::bridge::{open_reader, Bridge}; + use $crate::__private::abi::{ + self, ArrowArrayStream, ArrowSchema, SparkDsLibrary, SparkDsSink, + }; + use $crate::__private::bridge::{open_reader, parse_predicates, Bridge}; - use $crate::provider::Provider as _; + use $crate::__private::provider::Provider as _; - type Registration = $provider; - type Factory = <::Scan + // Prefixed so that they cannot capture the name of the registered source type. + type __SparkRegistration = $provider; + type __SparkSource = + <__SparkRegistration as $crate::__private::provider::Provider>::Source; + type __SparkReaderFactory = <<__SparkRegistration + as $crate::__private::provider::Provider>::ScanBuilder as $crate::ScanBuilder>::Factory; - type WriterFactory = <::Write + type __SparkWriterFactory = <<__SparkRegistration + as $crate::__private::provider::Provider>::BatchWrite as $crate::BatchWrite>::Factory; - type Message = ::Message; + type __SparkMessage = <__SparkWriterFactory as $crate::DataWriterFactory>::Message; - fn bridge() -> &'static Bridge { - static BRIDGE: std::sync::OnceLock> = + fn bridge() -> &'static Bridge<__SparkRegistration> { + static BRIDGE: std::sync::OnceLock> = std::sync::OnceLock::new(); BRIDGE.get_or_init(Bridge::default) } + const READ: bool = __SparkRegistration::READ; + const WRITE: bool = __SparkRegistration::WRITE; + static LIBRARY: SparkDsLibrary = SparkDsLibrary { size: std::mem::size_of::() as i64, abi_version: abi::ABI_VERSION, last_error: Some(abi::last_error), create: Some(create), - schema: Some(schema), - create_reader: if Registration::READ { Some(create_reader) } else { None }, + schema: if <__SparkSource as $crate::DataSource>::INFERS_SCHEMA { + Some(schema) + } else { + None + }, + create_reader: if READ { Some(create_reader) } else { None }, create_stream_reader: None, - create_writer: if Registration::WRITE { Some(create_writer) } else { None }, + create_writer: if WRITE { Some(create_writer) } else { None }, close: Some(close), - push_predicates: if Registration::READ { Some(push_predicates) } else { None }, - push_limit: if Registration::READ { Some(push_limit) } else { None }, - prune_columns: if Registration::READ { Some(prune_columns) } else { None }, - partitions: if Registration::READ { Some(partitions) } else { None }, - serialize_reader: if Registration::READ { Some(serialize_reader) } else { None }, - close_reader: if Registration::READ { Some(close_reader) } else { None }, + push_predicates: if READ { Some(push_predicates) } else { None }, + push_limit: if READ { Some(push_limit) } else { None }, + prune_columns: if READ { Some(prune_columns) } else { None }, + partitions: if READ { Some(partitions) } else { None }, + serialize_reader: if READ { Some(serialize_reader) } else { None }, + close_reader: if READ { Some(close_reader) } else { None }, initial_offset: None, latest_offset: None, stream_partitions: None, serialize_stream_reader: None, commit_offset: None, close_stream_reader: None, - read: if Registration::READ { Some(read) } else { None }, - serialize_writer: if Registration::WRITE { Some(serialize_writer) } else { None }, - commit: if Registration::WRITE { Some(commit) } else { None }, - abort: if Registration::WRITE { Some(abort) } else { None }, - close_writer: if Registration::WRITE { Some(close_writer) } else { None }, - create_data_writer: if Registration::WRITE { - Some(create_data_writer) - } else { - None - }, - write: if Registration::WRITE { Some(write) } else { None }, - commit_data_writer: if Registration::WRITE { - Some(commit_data_writer) - } else { - None - }, - abort_data_writer: if Registration::WRITE { Some(abort_data_writer) } else { None }, - close_data_writer: if Registration::WRITE { Some(close_data_writer) } else { None }, + read: if READ { Some(read) } else { None }, + serialize_writer: if WRITE { Some(serialize_writer) } else { None }, + commit: if WRITE { Some(commit) } else { None }, + abort: if WRITE { Some(abort) } else { None }, + close_writer: if WRITE { Some(close_writer) } else { None }, + create_data_writer: if WRITE { Some(create_data_writer) } else { None }, + write: if WRITE { Some(write) } 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 }, }; #[export_name = "spark_ds_init"] @@ -112,7 +144,14 @@ macro_rules! export_data_source { host: *const abi::SparkDsHost, library: *mut *const SparkDsLibrary, ) -> c_int { - unsafe { abi::initialize(host, library, &LIBRARY) } + unsafe { + abi::initialize( + host, + library, + &LIBRARY, + <__SparkSource as $crate::DataSource>::init, + ) + } } unsafe extern "C" fn create( @@ -123,8 +162,8 @@ macro_rules! export_data_source { abi::call(|| { abi::require_output(capabilities)?; abi::require_output(source)?; - let options = unsafe { abi::options(info)? }; - let (handle, declared_capabilities) = bridge().create(options)?; + let info = unsafe { abi::create_info(info)? }; + let (handle, declared_capabilities) = bridge().create(info)?; unsafe { capabilities.write(declared_capabilities); source.write(handle as usize as *mut c_void); @@ -169,8 +208,8 @@ macro_rules! export_data_source { abi::require_output(uses)?; } let values = unsafe { abi::strings(predicates, count)? }; - let predicates = $crate::bridge::parse_predicates(&values)?; - let support = bridge().push_predicates(reader as usize as i64, &predicates)?; + let predicates = parse_predicates(&values); + let support = bridge().push_predicates(reader as usize as i64, predicates)?; for (index, support) in support.into_iter().enumerate() { unsafe { uses.add(index).write(abi::predicate_use(support)) }; } @@ -247,7 +286,8 @@ macro_rules! export_data_source { abi::require_output(output)?; let state = unsafe { abi::slice(state, state_length)? }; let partition = unsafe { abi::slice(partition, partition_length)? }; - let stream = unsafe { open_reader::(task, state, partition)? }; + let stream = + unsafe { open_reader::<__SparkReaderFactory>(task, state, partition)? }; unsafe { output.write(stream) }; Ok(()) }) @@ -263,7 +303,8 @@ macro_rules! export_data_source { abi::require_output(writer)?; let mode = unsafe { abi::write_mode(info)? }; let schema = unsafe { abi::import_schema(schema)? }; - let handle = bridge().create_writer(source as usize as i64, schema, mode)?; + let info = $crate::WriteInfo::new(schema, mode); + let handle = bridge().create_writer(source as usize as i64, info)?; unsafe { writer.write(handle as usize as *mut c_void) }; Ok(()) }) @@ -323,7 +364,9 @@ macro_rules! export_data_source { abi::require_output(writer)?; abi::batch_epoch(epoch)?; let state = unsafe { abi::slice(state, length)? }; - let handle = unsafe { $crate::write::open::(task, state)? }; + let handle = unsafe { + $crate::__private::write::open::<__SparkWriterFactory>(task, state)? + }; unsafe { writer.write(handle) }; Ok(()) }) @@ -335,7 +378,7 @@ macro_rules! export_data_source { schema: *mut ArrowSchema, ) -> c_int { abi::call(|| unsafe { - $crate::write::write_batch::(writer, array, schema) + $crate::__private::write::write_batch::<__SparkMessage>(writer, array, schema) }) } @@ -346,18 +389,32 @@ macro_rules! export_data_source { ) -> c_int { abi::call(|| { sink.ok_or("null result sink")?; - let message = unsafe { $crate::write::commit_task::(writer)? }; + let message = unsafe { + $crate::__private::write::commit_task::<__SparkMessage>(writer)? + }; unsafe { abi::emit(sink, data, &message) } }) } unsafe extern "C" fn abort_data_writer(writer: *mut c_void) -> c_int { - abi::call(|| unsafe { $crate::write::abort_task::(writer) }) + abi::call(|| unsafe { + $crate::__private::write::abort_task::<__SparkMessage>(writer) + }) } unsafe extern "C" fn close_data_writer(writer: *mut c_void) { - let _ = abi::call(|| unsafe { $crate::write::close_task::(writer) }); + let _ = abi::call(|| unsafe { + $crate::__private::write::close_task::<__SparkMessage>(writer) + }); } } }; + ($source:ty, $($operation:ident),+ $(,)?) => { + $crate::export_data_source!(@parse $source; [false false]; $($operation)+); + }; + ($source:ty $(,)?) => { + compile_error!( + "name the operations of the source, such as export_data_source!(MySource, batch_read)" + ); + }; } diff --git a/sql/native/rust/src/host.rs b/sql/native/rust/src/host.rs new file mode 100644 index 0000000000000..edbbde8bf7784 --- /dev/null +++ b/sql/native/rust/src/host.rs @@ -0,0 +1,87 @@ +// 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 Spark process that loaded the library: its version and its log. + +use std::ffi::{c_char, CString}; +use std::sync::{PoisonError, RwLock}; + +pub(crate) type LogFunction = unsafe extern "C" fn(i32, *const c_char); + +// Copied from the host when Spark initializes the library, so that logging never +// reads the host structure afterwards. +static LOG: RwLock> = RwLock::new(None); + +/// The Spark process that loaded the library, as [`DataSource::init`](crate::DataSource::init) +/// receives it. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct HostInfo { + /// The newest version of the native data source ABI that Spark supports. + pub abi_version: i32, + /// The version of Spark, such as `4.4.0`, when Spark provides it. + pub spark_version: Option, +} + +impl Default for HostInfo { + /// A host that supports the ABI version of this SDK, for testing outside Spark. + fn default() -> Self { + Self { + abi_version: crate::abi::ABI_VERSION, + spark_version: None, + } + } +} + +/// The level of a message logged with [`log`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +#[non_exhaustive] +pub enum LogLevel { + Error, + Warn, + Info, + Debug, + Trace, +} + +impl LogLevel { + /// The SPARK_DS_LOG_* value of this level. + fn code(self) -> i32 { + match self { + Self::Error => 1, + Self::Warn => 2, + Self::Info => 3, + Self::Debug => 4, + Self::Trace => 5, + } + } +} + +/// Logs a message in the log of the Spark driver or executor, with the path of the library. +/// +/// Can be called from any thread. Messages are dropped until Spark initializes the +/// library, for example in unit tests. +pub fn log(level: LogLevel, message: &str) { + let function = *LOG.read().unwrap_or_else(PoisonError::into_inner); + if let Some(function) = function { + let message = CString::new(message.replace('\0', "\\0")).expect("NULs removed"); + unsafe { function(level.code(), message.as_ptr()) }; + } +} + +/// Installs the log of the host that initializes the library. +pub(crate) fn install(function: Option) { + *LOG.write().unwrap_or_else(PoisonError::into_inner) = function; +} diff --git a/sql/native/rust/src/lib.rs b/sql/native/rust/src/lib.rs index 18c8935c00d00..3ebab03e22a0b 100644 --- a/sql/native/rust/src/lib.rs +++ b/sql/native/rust/src/lib.rs @@ -21,50 +21,161 @@ //! stay inside the plugin; they are not an ABI between dynamic libraries. use std::collections::BTreeMap; +use std::fmt; use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::SchemaRef; use serde::de::DeserializeOwned; -use serde::{Deserialize, Serialize}; +use serde::Serialize; -#[doc(hidden)] -pub mod abi; -#[doc(hidden)] -pub mod bridge; +mod abi; +mod bridge; mod export; +mod host; mod predicate; -#[doc(hidden)] -pub mod provider; +mod provider; mod task; +mod write; + +/// Implementation details of [`export_data_source!`] and of the SDK's own tests, not an API. #[doc(hidden)] -pub mod write; +pub mod __private { + pub mod abi { + pub use crate::abi::*; + } + + pub mod bridge { + pub use crate::bridge::{open_reader, parse_predicates, Bridge}; + } + + pub mod provider { + pub use crate::provider::*; + } + + pub mod write { + pub use crate::write::*; + } +} -pub use predicate::Predicate; -pub use task::TaskContext; +pub use host::{log, HostInfo, LogLevel}; +pub use predicate::{LiteralValue, Predicate}; +pub use task::{TaskContext, TaskContextBuilder}; pub type Error = Box; pub type Result = std::result::Result; pub type Options = BTreeMap; -pub type BatchReader<'task> = Box; -pub type TaskWriter<'task, M> = Box + 'task>; +/// The lazy Arrow reader of one partition, which [`PartitionReaderFactory::open`] returns. +pub type PartitionReader<'task> = Box; +/// The writer of one task attempt, which [`DataWriterFactory::open`] returns. +pub type PartitionWriter<'task, M> = Box + 'task>; -/// Native operations supported by a source instance with its effective options. +/// A native operation of a source instance. /// /// These are native capabilities, not Spark's Data Source V2 routing flags or -/// per-scan pushdown results. Registration determines the available operations; -/// a source instance may disable them but cannot enable unregistered operations. -/// Streaming operations are not implemented by this SDK yet. +/// per-scan pushdown results. Streaming operations are not implemented by this SDK yet. #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[non_exhaustive] pub enum SourceCapability { BatchRead, - MicroBatchRead, BatchWrite, - StreamingWrite, - /// Supports writes that replace the table's existing data. + /// Batch writes that replace the existing data. Requires `BatchWrite`. Truncate, } +impl SourceCapability { + const ALL: [Self; 3] = [Self::BatchRead, Self::BatchWrite, Self::Truncate]; + + const fn bit(self) -> u64 { + match self { + Self::BatchRead => abi::BATCH_READ, + Self::BatchWrite => abi::BATCH_WRITE, + Self::Truncate => abi::TRUNCATE, + } + } +} + +/// The native operations that a source instance supports with its effective options. +/// +/// Registration determines the operations of a library. [`DataSource::capabilities`] +/// receives them and can only remove some, so a source cannot enable an operation +/// that it did not register. Removing `BatchWrite` also removes `Truncate`. +#[derive(Clone, Copy, PartialEq, Eq)] +pub struct Capabilities(u64); + +impl Capabilities { + /// No operations. A source without capabilities still supports schema discovery. + pub const fn empty() -> Self { + Self(0) + } + + /// Every operation that this SDK implements, to test [`DataSource::capabilities`]. + /// The SDK keeps only the operations that the library registers. + pub const fn all() -> Self { + Self(abi::BATCH_READ | abi::BATCH_WRITE | abi::TRUNCATE) + } + + pub const fn contains(self, capability: SourceCapability) -> bool { + self.0 & capability.bit() != 0 + } + + pub const fn is_empty(self) -> bool { + self.0 == 0 + } + + /// Removes an operation, and also `Truncate` when it removes `BatchWrite`. + #[must_use] + pub const fn without(self, capability: SourceCapability) -> Self { + let removed = match capability { + SourceCapability::BatchWrite => abi::BATCH_WRITE | abi::TRUNCATE, + other => other.bit(), + }; + Self(self.0 & !removed) + } + + /// The operations that registering batch reads, batch writes, and truncation enables. + pub(crate) const fn registered(read: bool, write: bool, truncate: bool) -> Self { + let mut bits = 0; + if read { + bits |= abi::BATCH_READ; + } + if write { + bits |= abi::BATCH_WRITE; + if truncate { + bits |= abi::TRUNCATE; + } + } + Self(bits) + } + + /// The operations in both sets, without `Truncate` unless `BatchWrite` remains. + pub(crate) const fn intersection(self, other: Self) -> Self { + let bits = self.0 & other.0; + if bits & abi::BATCH_WRITE == 0 { + Self(bits & !abi::TRUNCATE) + } else { + Self(bits) + } + } + + /// The native capability mask, a combination of the SPARK_DS_* flags. + pub(crate) const fn bits(self) -> u64 { + self.0 + } +} + +impl fmt::Debug for Capabilities { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_set() + .entries( + SourceCapability::ALL + .into_iter() + .filter(|capability| self.contains(*capability)), + ) + .finish() + } +} + /// How completely a connector evaluates a pushed predicate. /// Spark currently rechecks all predicates, including exact ones. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -78,58 +189,188 @@ pub enum FilterSupport { Unsupported, } -/// How a connector limits the rows produced by the whole scan. +/// 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, - /// Caps each partition; the ABI boolean is false because it promises a - /// bound on the whole scan. Spark retains its final limit in either case. - PerPartition, - /// Caps the whole scan; the ABI boolean is true. + /// The whole scan returns at most the limit. Spark still applies its limit. Global, } +/// 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, options: Options) -> Self { + Self { + name: name.into(), + options, + } + } +} + +/// Driver-side source configuration, schema discovery, and capability restrictions. +/// +/// Spark creates a source for each schema inference, scan, and write, and may create +/// one only to inspect its capabilities, then drop it. Keep construction lightweight. +/// Scan builders and writes cannot borrow their source, which Spark drops first. +pub trait DataSource: Send + Sized + 'static { + /// 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; + + /// 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(()) + } + + fn new(info: CreateInfo) -> Result; + + /// Returns the operations of this instance, a subset of the registered ones. 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, registered: Capabilities) -> Capabilities { + registered + } + + /// Returns the schema of the data, when the user does not specify one. + fn schema(&self) -> Result { + Err( + "DataSource::schema is not implemented: implement it, or set \ + DataSource::INFERS_SCHEMA to false to require a user-specified schema" + .into(), + ) + } +} + +/// Driver-side entry point for batch reads. +/// +/// Register with `export_data_source!(Source, batch_read)`. Metadata and construction +/// belong to [`DataSource`]. +pub trait BatchReadSource: DataSource { + type ScanBuilder: ScanBuilder; + + /// Creates an independent scan of the schema to read, which the user may have + /// specified; reject unsupported schemas. + fn new_scan(&self, schema: SchemaRef) -> Result; +} + +/// Mutable driver-side negotiation, followed by immutable planning. +/// +/// The SDK consumes this builder on the first native `partitions` or +/// `serialize_reader` call and caches the result, including a planning failure. +/// This is later than Java DSv2 `ScanBuilder.build()`. Only serialized factory +/// and partition descriptors travel to executors; driver handles stay local. +/// It corresponds to Spark's driver-only `ColumnarBatchPlanner`: the native +/// planner calls `partitions`, then `serialize_reader`, then closes the handle. +pub trait ScanBuilder: Send + Sized + 'static { + type Factory: PartitionReaderFactory; + + /// Return one result per input predicate in order. Unknown functions and + /// types must not be reported as exact. Spark currently rechecks all filters. + /// A predicate that the SDK cannot parse arrives as [`Predicate::Unknown`], and + /// the SDK reports it as unsupported whatever this method returns for it. + fn push_filters(&mut self, filters: &[Predicate]) -> Result> { + Ok(vec![FilterSupport::Unsupported; filters.len()]) + } + + /// Count rows after filtering. Spark currently pushes limits only for + /// scans with no filters; the SDK also declines a limit when any predicate + /// it received remains residual. + fn push_limit(&mut self, _limit: usize) -> Result { + Ok(LimitSupport::Unsupported) + } + + /// On true, output exactly these top-level columns in this order. On false, + /// keep the original schema. Empty projection must preserve row counts. + /// Retain internally needed filter columns without returning them. Readers get + /// the resulting schema in [`ReaderContext::schema`]. + fn prune_columns(&mut self, _columns: &[String]) -> Result { + Ok(false) + } + + fn plan(self) -> Result>; +} + /// 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. -pub struct ScanPlan { +#[non_exhaustive] +pub struct ScanPlan { pub factory: F, pub partitions: Vec, } -/// Driver-side source configuration, schema discovery, and capability restrictions. -/// -/// Spark may create a source only to inspect its capabilities, then drop it. -/// Keep construction lightweight. Readers and writers must outlive their source. -pub trait DataSource: Send + Sized + 'static { - /// Spark passes option keys in lower case. A capability probe uses table - /// options; an operation may override them with query options. - fn new(options: Options) -> Result; - - /// By default, enable every registered operation. Override this to restrict - /// the current instance, for example for a read-only destination. The SDK - /// evaluates this once, ignores duplicates, and rejects unregistered flags. - /// Truncate requires BatchWrite. An empty set still permits schema discovery. - fn capabilities(&self, supported: &[SourceCapability]) -> Vec { - supported.to_vec() +impl ScanPlan { + pub fn new(factory: F, partitions: Vec) -> Self { + Self { + factory, + partitions, + } } +} - fn schema(&self) -> Result; +/// What an executor reader receives with its partition. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct ReaderContext<'task> { + /// The negotiated schema: the schema to read, or its pruned columns in order. + /// The reader returns batches with exactly these columns. + pub schema: SchemaRef, + /// The task that reads the partition. + pub task: TaskContext<'task>, } -/// Driver-side entry point for batch reads. +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 } + } +} + +/// Serializable configuration for creating independent executor readers. /// -/// Register with `export_data_source!(Source, batch_read)` (or its one-argument -/// shorthand). Metadata and construction belong to [`DataSource`]. -pub trait BatchSource: DataSource { - type Scan: ScanBuilder; - - /// The schema may be supplied by the user; reject unsupported schemas. - /// The returned builder must remain valid after this source is dropped. - fn scan(&self, schema: SchemaRef) -> Result; +/// This type carries the configuration or snapshot descriptor; connections and +/// other live resources are opened on the executor. The returned reader owns +/// those resources and cannot borrow this short-lived factory. The SDK serializes +/// the negotiated schema with the factory, so the factory need not carry it. +pub trait PartitionReaderFactory: Serialize + DeserializeOwned + Send + 'static { + type Partition: Serialize + DeserializeOwned + Send + 'static; + + /// Return lazy batches with the schema of `context` and bounded memory. The reader + /// may retain the context until the stream is released. Check cancellation during + /// long operations and use timeouts for blocking I/O. + /// Spark may call stream callbacks on different threads, never concurrently; + /// the reader must be `Send` but does not need to be `Sync`. + /// + /// The SDK checks Spark-compatible Arrow types, the negotiated schema and its + /// nullability, and zero array offsets recursively. It reports incompatible layouts + /// instead of copying buffers. Exported batches must own their buffers independently + /// of the reader, because Spark may release the stream before its batches. + fn open<'task>( + &self, + partition: Self::Partition, + context: ReaderContext<'task>, + ) -> Result>; } /// How a batch write changes the destination. @@ -137,7 +378,7 @@ pub trait BatchSource: DataSource { /// 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, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Eq)] #[non_exhaustive] pub enum WriteMode { Append, @@ -145,18 +386,61 @@ pub enum WriteMode { Truncate, } +/// 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 } + } +} + /// Driver-side entry point for batch writes, independent of read support. /// -/// Register `batch_write` and, if implemented, `truncate`. The SDK derives native +/// Register with `export_data_source!(Source, batch_write)`. The SDK derives native /// capabilities and callbacks from registration, then applies instance restrictions. pub trait BatchWriteSource: DataSource { - type Write: BatchWrite; + type BatchWrite: BatchWrite; + + /// Whether [`new_write`](Self::new_write) supports [`WriteMode::Truncate`], which + /// registers the `Truncate` capability. + const SUPPORTS_TRUNCATE: bool = false; - /// Validate the input schema and create an independent write coordinator. - /// Construction may be followed only by Drop, without commit or abort. - /// The SDK calls this only while the instance's capabilities include + /// 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 /// BatchWrite, and passes `Truncate` only when they include Truncate. - fn write(&self, schema: SchemaRef, mode: WriteMode) -> Result; + fn new_write(&self, info: WriteInfo) -> Result>; +} + +/// A driver coordinator and the serializable configuration of its task writers. +/// +/// The SDK serializes the factory once, with the schema of the input, and caches the +/// result, including a serialization failure. Include a unique operation ID in the +/// factory when needed; Spark supplies no batch query ID. +#[non_exhaustive] +pub struct WritePlan { + /// The coordinator that commits or aborts the write on the driver. + pub coordinator: W, + /// The configuration from which each task attempt opens its writer. + pub factory: W::Factory, +} + +impl WritePlan { + pub fn new(coordinator: W, factory: W::Factory) -> Self { + Self { + coordinator, + factory, + } + } } /// A driver coordinator for one batch write. @@ -169,21 +453,20 @@ pub trait BatchWriteSource: DataSource { /// The SDK enforces the following, so implementations need not track their own /// phase or guard against repeated calls: /// -/// - `factory` runs at most once, and never after `commit` or `abort`. -/// - `commit` runs at most once, only while no earlier call has failed. It is -/// skipped, and the write fails, when any task message is missing. +/// - `commit` runs at most once, only while no earlier call has failed, including +/// the serialization of the factory. It is skipped, and the write fails, when +/// any task message is missing or cannot be decoded. /// - `abort` runs at most once, never after a successful `commit`. It may follow -/// a failed `factory` or `commit`, or run without either being called. -/// - Drop runs exactly once, possibly without `commit` or `abort`. +/// a failed `commit`, or run without `commit`. +/// - Drop runs exactly once, possibly without `commit` or `abort`. Spark closes the +/// coordinator without `abort` when serializing the factory fails, so Drop must +/// also release what only this write staged. /// - 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 { type Factory: DataWriterFactory; - /// Return serializable executor configuration. The SDK caches the encoded - /// factory and negotiated schema, including serialization failures. - /// Include a unique operation ID when needed; Spark supplies no batch query ID. - fn factory(&mut self) -> Result; - /// 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( @@ -191,8 +474,8 @@ pub trait BatchWrite: Send + 'static { messages: Vec<::Message>, ) -> Result<()>; - /// Abort after input, task, or driver commit failure. None preserves the - /// position of a partition whose commit message did not reach the driver. + /// 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, @@ -200,15 +483,35 @@ pub trait BatchWrite: Send + 'static { ) -> Result<()>; } +/// 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>, +} + +impl<'task> WriterContext<'task> { + /// The context of a writer opened outside Spark, for example in a test. + pub fn new(schema: SchemaRef, task: TaskContext<'task>) -> Self { + Self { schema, task } + } +} + /// Serializable configuration for independent task-attempt writers. /// /// Open connections and files in open, not in the factory. A returned writer -/// cannot borrow this temporary factory, but may retain the task context until +/// 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, task: TaskContext<'task>) -> Result>; + fn open<'task>( + &self, + context: WriterContext<'task>, + ) -> Result>; } /// One task attempt's writer. The SDK closes it after commit or abort. @@ -221,8 +524,10 @@ pub trait DataWriterFactory: Serialize + DeserializeOwned + Send + 'static { /// - `write` runs zero or more times, and never after any call has failed or /// after `commit` or `abort`. /// - `commit` runs at most once, never after a failed `write` or after `abort`. -/// - `abort` runs at most once, never after a successful `commit`. It may follow -/// a failed `write` or `commit`. +/// - `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`. /// - A panic in any method is reported as that method's failure. @@ -230,7 +535,8 @@ pub trait DataWriter: Send { type Message; /// Takes ownership of an imported Arrow batch, which may be retained after - /// this call. Process batches incrementally and check task cancellation. + /// this call. Its columns match [`WriterContext::schema`] by Spark type. + /// Process batches incrementally and check task cancellation. fn write(&mut self, batch: RecordBatch) -> Result<()>; /// Finish staging and return a typed message for the driver. Do not make the @@ -240,62 +546,3 @@ pub trait DataWriter: Send { fn abort(&mut self) -> Result<()>; } - -/// Serializable configuration for creating independent executor readers. -/// -/// This type carries the configuration or snapshot descriptor; connections and -/// other live resources are opened on the executor. The returned reader owns -/// those resources and cannot borrow this short-lived factory. -pub trait BatchReaderFactory: Serialize + DeserializeOwned + Send + 'static { - type Partition: Serialize + DeserializeOwned + Send + 'static; - - /// Return lazy batches with a fixed schema and bounded memory. The reader - /// may retain `task` until the stream is released. Check cancellation during - /// long operations and use timeouts for blocking I/O. - /// Spark may call stream callbacks on different threads, never concurrently; - /// the reader must be `Send` but does not need to be `Sync`. - /// - /// The SDK checks Spark-compatible Arrow types, the negotiated schema, and - /// zero array offsets recursively. It reports incompatible layouts instead - /// of copying buffers. Exported batches must own their buffers independently - /// of the reader, because Spark may release the stream before its batches. - fn open<'task>( - &self, - partition: Self::Partition, - task: TaskContext<'task>, - ) -> Result>; -} - -/// Mutable driver-side negotiation, followed by immutable planning. -/// -/// The SDK consumes this builder on the first native `partitions` or -/// `serialize_reader` call and caches the result, including a planning failure. -/// This is later than Java DSv2 `ScanBuilder.build()`. Only serialized factory -/// and partition descriptors travel to executors; driver handles stay local. -/// It corresponds to Spark's driver-only `ColumnarBatchPlanner`: the native -/// planner calls `partitions`, then `serialize_reader`, then closes the handle. -pub trait ScanBuilder: Send + Sized + 'static { - type Factory: BatchReaderFactory; - - /// Return one result per input predicate in order. Unknown functions and - /// types must not be reported as exact. Spark currently rechecks all filters. - fn push_filters(&mut self, filters: &[Predicate]) -> Result> { - Ok(vec![FilterSupport::Unsupported; filters.len()]) - } - - /// Count rows after filtering. Spark currently pushes limits only for - /// scans with no filters; the SDK also declines a limit when any predicate - /// it received remains residual. - fn push_limit(&mut self, _limit: usize) -> Result { - Ok(LimitSupport::Unsupported) - } - - /// On true, output exactly these top-level columns in this order. On false, - /// keep the original schema. Empty projection must preserve row counts. - /// Retain internally needed filter columns without returning them. - fn prune_columns(&mut self, _columns: &[String]) -> Result { - Ok(false) - } - - fn plan(self) -> Result>; -} diff --git a/sql/native/rust/src/predicate.rs b/sql/native/rust/src/predicate.rs index 94bcda501f067..538678afe2804 100644 --- a/sql/native/rust/src/predicate.rs +++ b/sql/native/rust/src/predicate.rs @@ -21,6 +21,7 @@ use serde_json::Value; /// Literal values keep their JSON representation so a connector can reject /// unsupported types without lossy numeric conversions. In particular, /// decimal values are strings and timestamps are integer microseconds. +/// [`literal_value`](Self::literal_value) decodes them by their data type. /// Function names are extensible: unknown names must remain residual. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "lowercase")] @@ -42,9 +43,169 @@ pub enum Predicate { Unknown, } +/// The value of a literal, decoded by its Spark SQL data type. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub enum LiteralValue { + /// A null of any data type. + Null, + Boolean(bool), + /// A `tinyint`, `smallint`, `int`, or `bigint`, within the range of its type. + Integer(i64), + /// A `float` or `double`. A `float` is the exact value of the float. + Float(f64), + /// A `string`, which Spark compares by its UTF-8 bytes. + String(String), + Binary(Vec), + /// A `decimal(precision, scale)`, whose value is `unscaled * 10^-scale`. + Decimal { + unscaled: i128, + precision: u8, + scale: i8, + }, + /// A `date`, as days since 1970-01-01. + Date(i32), + /// A `timestamp`, as microseconds since the epoch. + Timestamp(i64), + /// A `timestamp_ntz`, as microseconds since 1970-01-01 00:00:00 without a time zone. + TimestampNtz(i64), +} + +impl Predicate { + /// Decodes a literal by its data type. Returns `None` for other expressions, and + /// for data types or values that this SDK does not decode: keep such predicates + /// unsupported. + pub fn literal_value(&self) -> Option { + let Predicate::Literal { data_type, value } = self else { + return None; + }; + if value.is_null() { + return Some(LiteralValue::Null); + } + Some(match data_type.as_str() { + "boolean" => LiteralValue::Boolean(value.as_bool()?), + "tinyint" => LiteralValue::Integer(i8::try_from(value.as_i64()?).ok()?.into()), + "smallint" => LiteralValue::Integer(i16::try_from(value.as_i64()?).ok()?.into()), + "int" => LiteralValue::Integer(i32::try_from(value.as_i64()?).ok()?.into()), + "bigint" => LiteralValue::Integer(value.as_i64()?), + "float" => { + // Spark writes a float as the double of the same value. + let value = value.as_f64()?; + if f64::from(value as f32) != value { + return None; + } + LiteralValue::Float(value) + } + "double" => LiteralValue::Float(value.as_f64()?), + "string" => LiteralValue::String(value.as_str()?.to_owned()), + "binary" => LiteralValue::Binary(decode_base64(value.as_str()?)?), + "date" => LiteralValue::Date(i32::try_from(value.as_i64()?).ok()?), + "timestamp" => LiteralValue::Timestamp(value.as_i64()?), + "timestamp_ntz" => LiteralValue::TimestampNtz(value.as_i64()?), + other => decode_decimal(other, value.as_str()?)?, + }) + } +} + +/// Decodes a `decimal(p,s)` literal, which Spark writes as a plain decimal string. +fn decode_decimal(data_type: &str, text: &str) -> Option { + let arguments = data_type.strip_prefix("decimal(")?.strip_suffix(')')?; + let (precision, scale) = arguments.split_once(',')?; + let precision: u8 = precision.trim().parse().ok()?; + let scale: i8 = scale.trim().parse().ok()?; + if !(1..=38).contains(&precision) || i16::from(scale) > i16::from(precision) { + return None; + } + let (negative, digits) = match text.strip_prefix('-') { + Some(digits) => (true, digits), + None => (false, text), + }; + let (integer, fraction) = match digits.split_once('.') { + Some((integer, fraction)) if !fraction.is_empty() => (integer, fraction), + Some(_) => return None, + None => (digits, ""), + }; + if integer.is_empty() + || !integer + .bytes() + .chain(fraction.bytes()) + .all(|digit| digit.is_ascii_digit()) + { + return None; + } + let mut unscaled: i128 = 0; + for digit in integer.bytes().chain(fraction.bytes()) { + unscaled = unscaled + .checked_mul(10)? + .checked_add(i128::from(digit - b'0'))?; + } + // The digits have `fraction.len()` decimal places; rescale them to `scale` exactly. + let shift = i32::from(scale) - i32::try_from(fraction.len()).ok()?; + let factor = 10_i128.checked_pow(shift.unsigned_abs())?; + if shift >= 0 { + unscaled = unscaled.checked_mul(factor)?; + } else if unscaled % factor == 0 { + unscaled /= factor; + } else { + return None; + } + if unscaled >= 10_i128.pow(u32::from(precision)) { + return None; + } + Some(LiteralValue::Decimal { + unscaled: if negative { -unscaled } else { unscaled }, + precision, + scale, + }) +} + +/// Decodes standard base64 with padding, as Spark writes binary literals. +fn decode_base64(text: &str) -> Option> { + let text = text.as_bytes(); + if text.len() % 4 != 0 { + return None; + } + let groups = text.len() / 4; + let mut bytes = Vec::with_capacity(groups * 3); + for (index, group) in text.chunks_exact(4).enumerate() { + let padding = group.iter().rev().take_while(|&&byte| byte == b'=').count(); + if padding > 2 || (padding > 0 && index + 1 != groups) { + return None; + } + let mut bits = 0_u32; + for &byte in &group[..4 - padding] { + bits = (bits << 6) | u32::from(base64_value(byte)?); + } + bits <<= 6 * padding; + bytes.extend_from_slice(&bits.to_be_bytes()[1..4 - padding]); + } + Some(bytes) +} + +fn base64_value(byte: u8) -> Option { + match byte { + b'A'..=b'Z' => Some(byte - b'A'), + b'a'..=b'z' => Some(byte - b'a' + 26), + b'0'..=b'9' => Some(byte - b'0' + 52), + b'+' => Some(62), + b'/' => Some(63), + _ => None, + } +} + #[cfg(test)] mod tests { - use super::Predicate; + use serde_json::json; + + use super::{LiteralValue, Predicate}; + + fn literal(data_type: &str, value: serde_json::Value) -> Option { + Predicate::Literal { + data_type: data_type.into(), + value, + } + .literal_value() + } #[test] fn preserves_int64_literals_and_nested_column_paths() { @@ -66,6 +227,11 @@ mod tests { panic!("expected a literal"); }; assert_eq!(value.as_i64(), Some(i64::MAX)); + assert_eq!( + children[1].literal_value(), + Some(LiteralValue::Integer(i64::MAX)) + ); + assert_eq!(children[0].literal_value(), None); } #[test] @@ -81,4 +247,126 @@ mod tests { }; assert_eq!(name, "FUTURE_FUNCTION"); } + + #[test] + fn decodes_literals_by_their_data_type_and_range() { + assert_eq!(literal("int", json!(null)), Some(LiteralValue::Null)); + assert_eq!(literal("void", json!(null)), Some(LiteralValue::Null)); + assert_eq!( + literal("boolean", json!(true)), + Some(LiteralValue::Boolean(true)) + ); + assert_eq!( + literal("tinyint", json!(-128)), + Some(LiteralValue::Integer(-128)) + ); + assert_eq!(literal("tinyint", json!(128)), None); + assert_eq!( + literal("smallint", json!(32767)), + Some(LiteralValue::Integer(32767)) + ); + assert_eq!(literal("smallint", json!(32768)), None); + assert_eq!( + literal("int", json!(i32::MIN)), + Some(LiteralValue::Integer(-1 << 31)) + ); + assert_eq!(literal("int", json!(1_i64 << 31)), None); + assert_eq!(literal("bigint", json!(1.5)), None); + assert_eq!( + literal("double", json!(1.5)), + Some(LiteralValue::Float(1.5)) + ); + assert_eq!( + literal("float", json!(f64::from(0.1_f32))), + Some(LiteralValue::Float(f64::from(0.1_f32))) + ); + assert_eq!( + literal("float", json!(0.1)), + None, + "0.1 is not a float value" + ); + assert_eq!( + literal("string", json!("row-1")), + Some(LiteralValue::String("row-1".into())) + ); + assert_eq!(literal("string", json!(1)), None); + assert_eq!( + literal("date", json!(19000)), + Some(LiteralValue::Date(19000)) + ); + assert_eq!( + literal("timestamp", json!(-1)), + Some(LiteralValue::Timestamp(-1)) + ); + assert_eq!( + literal("timestamp_ntz", json!(i64::MAX)), + Some(LiteralValue::TimestampNtz(i64::MAX)) + ); + assert_eq!(literal("interval day", json!(1)), None); + } + + #[test] + fn decodes_binary_literals_as_padded_base64() { + for (text, bytes) in [ + ("", &b""[..]), + ("Zg==", b"f"), + ("Zm8=", b"fo"), + ("Zm9v", b"foo"), + ("Zm9vYg==", b"foob"), + ("AP8+/w==", &[0x00, 0xff, 0x3e, 0xff]), + ] { + assert_eq!( + literal("binary", json!(text)), + Some(LiteralValue::Binary(bytes.to_vec())), + "{text}" + ); + } + for invalid in ["Zg", "Zg=", "Z===", "Zg==Zg==", "Zm9-", "Zm 9"] { + assert_eq!(literal("binary", json!(invalid)), None, "{invalid}"); + } + } + + #[test] + fn decodes_plain_decimal_strings_at_their_scale() { + let decimal = |unscaled, precision, scale| { + Some(LiteralValue::Decimal { + unscaled, + precision, + scale, + }) + }; + assert_eq!( + literal("decimal(10,2)", json!("123.45")), + decimal(12345, 10, 2) + ); + assert_eq!(literal("decimal(10,2)", json!("-0.05")), decimal(-5, 10, 2)); + assert_eq!(literal("decimal(10,2)", json!("1.00")), decimal(100, 10, 2)); + assert_eq!( + literal("decimal(5,0)", json!("12345")), + decimal(12345, 5, 0) + ); + assert_eq!( + literal("decimal(3,-2)", json!("12300")), + decimal(123, 3, -2) + ); + let max = "9".repeat(38); + assert_eq!( + literal("decimal(38,0)", json!(max)), + decimal(10_i128.pow(38) - 1, 38, 0) + ); + for (data_type, text) in [ + ("decimal(4,2)", "123.45"), + ("decimal(10,2)", "1.234"), + ("decimal(3,-2)", "12345"), + ("decimal(10,2)", "1."), + ("decimal(10,2)", ".5"), + ("decimal(10,2)", "1e3"), + ("decimal(10,2)", "+1.00"), + ("decimal(39,0)", "1"), + ("decimal(2,3)", "0.001"), + ] { + assert_eq!(literal(data_type, json!(text)), None, "{data_type} {text}"); + } + assert_eq!(literal("decimal(10,2)", json!(1.5)), None); + } } diff --git a/sql/native/rust/src/provider.rs b/sql/native/rust/src/provider.rs index 29d4c22712c67..a97398f0e8cb1 100644 --- a/sql/native/rust/src/provider.rs +++ b/sql/native/rust/src/provider.rs @@ -21,87 +21,75 @@ use arrow_schema::SchemaRef; use serde::{Deserialize, Serialize}; use crate::{ - BatchReader, BatchReaderFactory, BatchSource, BatchWrite, BatchWriteSource, DataSource, - DataWriterFactory, Result, ScanBuilder, ScanPlan, SourceCapability, TaskContext, TaskWriter, - WriteMode, + BatchReadSource, BatchWrite, BatchWriteSource, Capabilities, DataSource, DataWriterFactory, + PartitionReader, PartitionReaderFactory, PartitionWriter, ReaderContext, Result, ScanBuilder, + ScanPlan, WriteInfo, WritePlan, WriterContext, }; pub trait Provider: Send + 'static { type Source: DataSource; - type Scan: ScanBuilder; - type Write: BatchWrite; + type ScanBuilder: ScanBuilder; + type BatchWrite: BatchWrite; const READ: bool; const WRITE: bool; - const CAPABILITIES: &'static [SourceCapability]; + const CAPABILITIES: Capabilities; - fn scan(source: &Self::Source, schema: SchemaRef) -> Result; - fn write(source: &Self::Source, schema: SchemaRef, mode: WriteMode) -> Result; + fn new_scan(source: &Self::Source, schema: SchemaRef) -> Result; + fn new_write(source: &Self::Source, info: WriteInfo) -> Result>; } pub struct ReadOnly(PhantomData); -pub struct WriteOnly(PhantomData); -pub struct ReadWrite(PhantomData); +pub struct WriteOnly(PhantomData); +pub struct ReadWrite(PhantomData); -impl Provider for ReadOnly { +impl Provider for ReadOnly { type Source = S; - type Scan = S::Scan; - type Write = Unavailable; + type ScanBuilder = S::ScanBuilder; + type BatchWrite = Unavailable; const READ: bool = true; const WRITE: bool = false; - const CAPABILITIES: &'static [SourceCapability] = &[SourceCapability::BatchRead]; + const CAPABILITIES: Capabilities = Capabilities::registered(true, false, false); - fn scan(source: &S, schema: SchemaRef) -> Result { - source.scan(schema) + fn new_scan(source: &S, schema: SchemaRef) -> Result { + source.new_scan(schema) } - fn write(_: &S, _: SchemaRef, _: WriteMode) -> Result { + fn new_write(_: &S, _: WriteInfo) -> Result> { Err("batch writes are not registered".into()) } } -impl Provider for WriteOnly { +impl Provider for WriteOnly { type Source = S; - type Scan = Unavailable; - type Write = S::Write; + type ScanBuilder = Unavailable; + type BatchWrite = S::BatchWrite; const READ: bool = false; const WRITE: bool = true; - const CAPABILITIES: &'static [SourceCapability] = if TRUNCATE { - &[SourceCapability::BatchWrite, SourceCapability::Truncate] - } else { - &[SourceCapability::BatchWrite] - }; + const CAPABILITIES: Capabilities = Capabilities::registered(false, true, S::SUPPORTS_TRUNCATE); - fn scan(_: &S, _: SchemaRef) -> Result { + fn new_scan(_: &S, _: SchemaRef) -> Result { Err("batch reads are not registered".into()) } - fn write(source: &S, schema: SchemaRef, mode: WriteMode) -> Result { - source.write(schema, mode) + fn new_write(source: &S, info: WriteInfo) -> Result> { + source.new_write(info) } } -impl Provider for ReadWrite { +impl Provider for ReadWrite { type Source = S; - type Scan = S::Scan; - type Write = S::Write; + type ScanBuilder = ::ScanBuilder; + type BatchWrite = ::BatchWrite; const READ: bool = true; const WRITE: bool = true; - const CAPABILITIES: &'static [SourceCapability] = if TRUNCATE { - &[ - SourceCapability::BatchRead, - SourceCapability::BatchWrite, - SourceCapability::Truncate, - ] - } else { - &[SourceCapability::BatchRead, SourceCapability::BatchWrite] - }; - - fn scan(source: &S, schema: SchemaRef) -> Result { - source.scan(schema) + const CAPABILITIES: Capabilities = Capabilities::registered(true, true, S::SUPPORTS_TRUNCATE); + + fn new_scan(source: &S, schema: SchemaRef) -> Result { + source.new_scan(schema) } - fn write(source: &S, schema: SchemaRef, mode: WriteMode) -> Result { - source.write(schema, mode) + fn new_write(source: &S, info: WriteInfo) -> Result> { + source.new_write(info) } } @@ -118,10 +106,10 @@ impl ScanBuilder for Unavailable { } } -impl BatchReaderFactory for Unavailable { +impl PartitionReaderFactory for Unavailable { type Partition = (); - fn open<'task>(&self, _: (), _: TaskContext<'task>) -> Result> { + fn open<'task>(&self, _: (), _: ReaderContext<'task>) -> Result> { Err("batch reads are not registered".into()) } } @@ -129,10 +117,6 @@ impl BatchReaderFactory for Unavailable { impl BatchWrite for Unavailable { type Factory = Self; - fn factory(&mut self) -> Result { - Err("batch writes are not registered".into()) - } - fn commit(&mut self, _: Vec<()>) -> Result<()> { Err("batch writes are not registered".into()) } @@ -145,7 +129,7 @@ impl BatchWrite for Unavailable { impl DataWriterFactory for Unavailable { type Message = (); - fn open<'task>(&self, _: TaskContext<'task>) -> Result> { + fn open<'task>(&self, _: WriterContext<'task>) -> Result> { Err("batch writes are not registered".into()) } } diff --git a/sql/native/rust/src/stream.rs b/sql/native/rust/src/stream.rs index c760f3e8b17ff..e9d46869d4d6c 100644 --- a/sql/native/rust/src/stream.rs +++ b/sql/native/rust/src/stream.rs @@ -21,16 +21,17 @@ use arrow_array::ffi::FFI_ArrowArray; use arrow_array::ffi_stream::FFI_ArrowArrayStream; use arrow_array::{Array, StructArray}; use arrow_data::ArrayData; -use arrow_schema::{ - ffi::FFI_ArrowSchema, DataType, Field, Fields, IntervalUnit, Schema, SchemaRef, TimeUnit, -}; +use arrow_schema::{ffi::FFI_ArrowSchema, DataType, Field, Fields, Schema, SchemaRef, TimeUnit}; -use super::panic_message; -use crate::{BatchReader, Result}; +use super::{describe, panic_message}; +use crate::{PartitionReader, Result}; struct StreamState<'task> { - reader: Option>, + reader: Option>, + // The reader's own schema, which each batch must have. schema: SchemaRef, + // The negotiated schema, whose nullability each batch must respect. + expected: SchemaRef, last_error: Option, finished: bool, } @@ -53,13 +54,14 @@ impl Drop for StreamState<'_> { /// The caller must keep everything borrowed by the reader alive until release, /// and the consumer must obey the Arrow C Stream ownership protocol. pub(super) unsafe fn export<'task>( - reader: BatchReader<'task>, + reader: PartitionReader<'task>, expected_schema: &SchemaRef, ) -> Result { // Establish guarded destruction before calling the provider's schema(). let mut state = Box::new(StreamState { reader: Some(reader), schema: Arc::new(Schema::empty()), + expected: expected_schema.clone(), last_error: None, finished: false, }); @@ -84,7 +86,7 @@ fn callback( action: impl FnOnce(&mut StreamState<'_>) -> Result<()>, ) -> c_int { let result = catch_unwind(AssertUnwindSafe(|| { - action(state).map_err(|e| e.to_string()) + action(state).map_err(|e| describe(&*e)) })); let message = match result { Ok(Ok(())) => return 0, @@ -127,6 +129,9 @@ unsafe extern "C" fn get_next( } let data = StructArray::from(batch).to_data(); validate_layout(&data)?; + for (column, field) in data.child_data().iter().zip(state.expected.fields()) { + check_nulls(column, field, field.name(), None)?; + } let array = FFI_ArrowArray::new(&data); unsafe { output.write(array) }; } @@ -157,9 +162,10 @@ unsafe extern "C" fn release(stream: *mut FFI_ArrowArrayStream) { } } -// Compare Spark logical types, ignoring Arrow nullability/metadata and the -// interchangeable physical encodings of strings and binary data. Spark also -// distinguishes only the presence of a time zone, not its spelling. +// Compare the Spark logical types of version 1 of the ABI, ignoring Arrow +// nullability/metadata and the interchangeable physical encodings of strings and +// binary data. Spark also distinguishes only the presence of a time zone, not its +// spelling. Nullability is checked against the data instead, in check_nulls. pub(crate) fn validate_schema(schema: &SchemaRef) -> Result<()> { canonical_fields(schema.fields())?; Ok(()) @@ -181,17 +187,13 @@ pub(crate) fn canonical_fields(fields: &Fields) -> Result { fn canonical_type(data_type: &DataType) -> Result { use DataType::*; Ok(match data_type { - Null | Boolean | Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Date32 => { - data_type.clone() - } + Boolean | Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Date32 => data_type.clone(), Utf8 | LargeUtf8 | Utf8View => Utf8, Binary | LargeBinary | BinaryView => Binary, Decimal128(precision, _) if *precision > 0 && *precision <= 38 => data_type.clone(), Timestamp(TimeUnit::Microsecond, tz) => { Timestamp(TimeUnit::Microsecond, tz.as_ref().map(|_| "UTC".into())) } - Duration(TimeUnit::Microsecond) - | Interval(IntervalUnit::YearMonth | IntervalUnit::MonthDayNano) => data_type.clone(), Struct(fields) => Struct(canonical_fields(fields)?), List(field) => List(Arc::new(Field::new( "item", @@ -218,6 +220,73 @@ fn canonical_type(data_type: &DataType) -> Result { }) } +/// Rejects a null where the negotiated schema declares a column, a field, an element, +/// or a map value not nullable, unless its parent is null, as Spark does after reading. +/// `referenced` tells which values belong to a parent value that is not null. +fn check_nulls( + data: &ArrayData, + field: &Field, + column: &str, + referenced: Option<&[bool]>, +) -> Result<()> { + let is_referenced = |index: usize| referenced.is_none_or(|mask| mask[index]); + if !field.is_nullable() + && data.null_count() > 0 + && (0..data.len()).any(|index| data.is_null(index) && is_referenced(index)) + { + return Err(format!( + "The column {column} has a null, but the schema to read declares it not nullable" + ) + .into()); + } + let valid = |index: usize| data.is_valid(index) && is_referenced(index); + match field.data_type() { + DataType::Struct(fields) => { + let mask: Vec = (0..data.len()).map(valid).collect(); + for (child, field) in data.child_data().iter().zip(fields) { + let name = format!("{column}.{}", field.name()); + check_nulls(child, field, &name, Some(&mask))?; + } + } + DataType::List(element) => { + let mask = elements(data, valid)?; + let name = format!("{column}.element"); + check_nulls(&data.child_data()[0], element, &name, Some(&mask))?; + } + DataType::Map(entries, _) => { + let DataType::Struct(fields) = entries.data_type() else { + return Err("Arrow map entries must be a struct".into()); + }; + let mask = elements(data, valid)?; + let entries = &data.child_data()[0]; + // Spark does not allow null keys. + let key = fields[0].as_ref().clone().with_nullable(false); + let name = format!("{column}.key"); + check_nulls(&entries.child_data()[0], &key, &name, Some(&mask))?; + let name = format!("{column}.value"); + check_nulls(&entries.child_data()[1], &fields[1], &name, Some(&mask))?; + } + _ => {} + } + Ok(()) +} + +/// Whether each element of a list or map belongs to a value that `valid` accepts. +fn elements(data: &ArrayData, valid: impl Fn(usize) -> bool) -> Result> { + let offsets = data.buffer::(0); + let mut mask = vec![false; data.child_data()[0].len()]; + for index in 0..data.len() { + if valid(index) { + let start = usize::try_from(offsets[index])?; + let end = usize::try_from(offsets[index + 1])?; + mask.get_mut(start..end) + .ok_or("Arrow offsets exceed the length of their values")? + .fill(true); + } + } + Ok(mask) +} + fn validate_layout(data: &ArrayData) -> Result<()> { if data.offset() != 0 { return Err(format!( @@ -237,6 +306,9 @@ fn validate_layout(data: &ArrayData) -> Result<()> { mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; + use arrow_array::builder::{ + BooleanBuilder, Int32Builder, ListBuilder, MapBuilder, StringBuilder, StructBuilder, + }; use arrow_array::ffi_stream::ArrowArrayStreamReader; use arrow_array::{BooleanArray, Int32Array, Int64Array, RecordBatch, RecordBatchReader}; use arrow_schema::{ArrowError, DataType, Field}; @@ -275,7 +347,7 @@ mod tests { } } - fn export(reader: BatchReader<'static>) -> Result { + fn export(reader: PartitionReader<'static>) -> Result { let expected = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); // Test readers own all their state and cannot borrow a host context. unsafe { super::export(reader, &expected) } @@ -491,6 +563,138 @@ mod tests { assert_eq!(drops.load(Ordering::SeqCst), 1); } + fn reader_of(schema: SchemaRef, batch: RecordBatch) -> Reader { + let mut reader = reader(); + reader.schema = schema; + reader.batches = vec![Ok(batch)].into_iter(); + reader + } + + fn first_batch_error(reader: Reader, expected: &SchemaRef) -> Option { + let exported = unsafe { super::export(Box::new(reader), expected) }.unwrap(); + let mut stream = ArrowArrayStreamReader::try_new(exported).unwrap(); + stream.next().unwrap().err().map(|error| error.to_string()) + } + + #[test] + fn intervals_and_the_null_type_are_not_version_1_types() { + for data_type in [ + DataType::Null, + DataType::Duration(TimeUnit::Microsecond), + DataType::Interval(arrow_schema::IntervalUnit::YearMonth), + DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano), + ] { + let schema = Arc::new(Schema::new(vec![Field::new("value", data_type, true)])); + assert!(validate_schema(&schema).is_err()); + } + } + + #[test] + fn nulls_must_respect_the_nullability_of_the_negotiated_schema() { + let values = Arc::new(Int64Array::from(vec![Some(1), None])); + let nullable = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, true)])); + let required = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let batch = RecordBatch::try_new(nullable.clone(), vec![values]).unwrap(); + let error = first_batch_error(reader_of(nullable.clone(), batch.clone()), &required); + assert!(error.unwrap().contains("The column id has a null")); + assert_eq!( + first_batch_error(reader_of(nullable.clone(), batch), &nullable), + None + ); + } + + #[test] + fn nested_nulls_are_allowed_only_under_a_null_parent() { + let required = Fields::from(vec![Field::new("flag", DataType::Boolean, false)]); + let nullable = Fields::from(vec![Field::new("flag", DataType::Boolean, true)]); + let schema = |fields: &Fields| { + Arc::new(Schema::new(vec![Field::new( + "nested", + DataType::Struct(fields.clone()), + true, + )])) + }; + fn flags(builder: &mut StructBuilder) -> &mut BooleanBuilder { + builder.field_builder::(0).unwrap() + } + for (parent_is_valid, rejected) in [(false, false), (true, true)] { + let mut builder = + StructBuilder::new(nullable.clone(), vec![Box::new(BooleanBuilder::new())]); + flags(&mut builder).append_value(true); + builder.append(true); + flags(&mut builder).append_null(); + builder.append(parent_is_valid); + let batch = + RecordBatch::try_new(schema(&nullable), vec![Arc::new(builder.finish())]).unwrap(); + let error = first_batch_error(reader_of(schema(&nullable), batch), &schema(&required)); + assert_eq!(error.is_some(), rejected, "{error:?}"); + if let Some(error) = error { + assert!( + error.contains("The column nested.flag has a null"), + "{error}" + ); + } + } + } + + #[test] + fn list_elements_and_map_values_respect_their_nullability() { + let list_schema = |nullable| { + let element = Arc::new(Field::new("item", DataType::Int32, nullable)); + Arc::new(Schema::new(vec![Field::new( + "values", + DataType::List(element), + true, + )])) + }; + for (first, rejected) in [(Some(1), false), (None, true)] { + let mut builder = ListBuilder::new(Int32Builder::new()); + builder.values().append_option(first); + builder.append(true); + // The second list is null, so its null element does not count. + builder.values().append_null(); + builder.append(false); + let lists = Arc::new(builder.finish()); + let batch = RecordBatch::try_new(list_schema(true), vec![lists]).unwrap(); + let error = first_batch_error(reader_of(list_schema(true), batch), &list_schema(false)); + assert_eq!(error.is_some(), rejected, "{error:?}"); + if let Some(error) = error { + assert!( + error.contains("The column values.element has a null"), + "{error}" + ); + } + } + + let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + builder.keys().append_value("a"); + builder.values().append_null(); + builder.append(true).unwrap(); + let map = builder.finish(); + let actual = Arc::new(Schema::new(vec![Field::new( + "map", + map.data_type().clone(), + true, + )])); + let entries = Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, false), + ]); + let entries = Arc::new(Field::new("entries", DataType::Struct(entries), false)); + let required = Arc::new(Schema::new(vec![Field::new( + "map", + DataType::Map(entries, false), + true, + )])); + let batch = RecordBatch::try_new(actual.clone(), vec![Arc::new(map)]).unwrap(); + let error = first_batch_error(reader_of(actual.clone(), batch.clone()), &required); + assert!(error.unwrap().contains("The column map.value has a null")); + assert_eq!( + first_batch_error(reader_of(actual.clone(), batch), &actual), + None + ); + } + #[test] fn a_batch_cannot_change_the_readers_advertised_schema() { let mut reader = reader(); diff --git a/sql/native/rust/src/task.rs b/sql/native/rust/src/task.rs index 5aab58ce8aa6b..61078ab7eef9b 100644 --- a/sql/native/rust/src/task.rs +++ b/sql/native/rust/src/task.rs @@ -14,7 +14,8 @@ // limitations under the License. use std::marker::PhantomData; -use std::mem::offset_of; +use std::mem::size_of; +use std::sync::atomic::{AtomicBool, Ordering}; use crate::abi::{self, Interrupted, SparkDsTaskContext}; use crate::Result; @@ -32,16 +33,26 @@ use crate::Result; /// task /// } /// ``` -#[derive(Clone, Copy)] +#[derive(Debug, Clone, Copy)] pub struct TaskContext<'task> { - pointer: *const SparkDsTaskContext, - stage_id: Option, - stage_attempt_number: Option, - partition_id: Option, - attempt_number: Option, - task_attempt_id: Option, - interrupted: Option, - lifetime: PhantomData<&'task SparkDsTaskContext>, + stage_id: i32, + stage_attempt_number: i32, + partition_id: i32, + attempt_number: i32, + task_attempt_id: i64, + interruption: Interruption<'task>, +} + +/// Where a task learns that Spark killed it. +#[derive(Debug, Clone, Copy)] +enum Interruption<'task> { + Never, + Host { + context: *const SparkDsTaskContext, + callback: Interrupted, + lifetime: PhantomData<&'task SparkDsTaskContext>, + }, + Flag(&'task AtomicBool), } // The C ABI permits task callbacks from any thread. The lifetime prevents use after the @@ -50,83 +61,156 @@ unsafe impl Send for TaskContext<'_> {} unsafe impl Sync for TaskContext<'_> {} impl TaskContext<'static> { - /// A context without Spark task metadata, for testing a reader or writer directly. - pub fn detached() -> Self { - Self { - pointer: std::ptr::null(), - stage_id: None, - stage_attempt_number: None, - partition_id: None, - attempt_number: None, - task_attempt_id: None, - interrupted: None, - lifetime: PhantomData, + /// Starts a context for testing a reader or writer outside Spark. Its IDs are 0 + /// unless set, and it is never interrupted unless it is given a flag. + pub fn for_testing() -> TaskContextBuilder<'static> { + TaskContextBuilder { + context: TaskContext { + stage_id: 0, + stage_attempt_number: 0, + partition_id: 0, + attempt_number: 0, + task_attempt_id: 0, + interruption: Interruption::Never, + }, } } } impl<'task> TaskContext<'task> { - /// Wraps a host context without reading fields beyond its declared size. + /// Wraps a host context, which has every field of version 1 of the ABI. /// /// # Safety /// The host must keep the context and its callbacks valid for `'task`, including calls from /// other threads. Spark keeps it alive until the Arrow stream is released or /// data writer closes. - #[doc(hidden)] - pub unsafe fn from_raw(pointer: *const SparkDsTaskContext) -> Result { + pub(crate) unsafe fn from_raw(pointer: *const SparkDsTaskContext) -> Result { let size = unsafe { abi::struct_size(pointer)? }; - let bytes = pointer.cast(); + // Fields are only appended to the structure, so every host has those of version 1. + if size < size_of::() as i64 { + return Err(format!( + "the task context has {size} bytes, fewer than in version 1 of the native ABI" + ) + .into()); + } + let context = unsafe { pointer.read_unaligned() }; Ok(Self { - pointer, - stage_id: unsafe { abi::field(bytes, size, offset_of!(SparkDsTaskContext, stage_id)) }, - stage_attempt_number: unsafe { - abi::field( - bytes, - size, - offset_of!(SparkDsTaskContext, stage_attempt_number), - ) + stage_id: context.stage_id, + stage_attempt_number: context.stage_attempt_number, + partition_id: context.partition_id, + attempt_number: context.attempt_number, + task_attempt_id: context.task_attempt_id, + interruption: match context.is_interrupted { + Some(callback) => Interruption::Host { + context: pointer, + callback, + lifetime: PhantomData, + }, + None => Interruption::Never, }, - partition_id: unsafe { - abi::field(bytes, size, offset_of!(SparkDsTaskContext, partition_id)) - }, - attempt_number: unsafe { - abi::field(bytes, size, offset_of!(SparkDsTaskContext, attempt_number)) - }, - task_attempt_id: unsafe { - abi::field(bytes, size, offset_of!(SparkDsTaskContext, task_attempt_id)) - }, - interrupted: unsafe { - abi::field(bytes, size, offset_of!(SparkDsTaskContext, is_interrupted)) - } - .flatten(), - lifetime: PhantomData, }) } - pub fn stage_id(&self) -> Option { + /// The ID of the stage of the task. + pub fn stage_id(&self) -> i32 { self.stage_id } - pub fn stage_attempt_number(&self) -> Option { + /// The attempt number of the stage of the task. + pub fn stage_attempt_number(&self) -> i32 { self.stage_attempt_number } - pub fn partition_id(&self) -> Option { + /// The partition that the task reads or writes. + pub fn partition_id(&self) -> i32 { self.partition_id } - pub fn attempt_number(&self) -> Option { + /// The attempt number of the task. + pub fn attempt_number(&self) -> i32 { self.attempt_number } - pub fn task_attempt_id(&self) -> Option { + /// The ID of the task attempt, unique within the Spark application but not across + /// applications. + pub fn task_attempt_id(&self) -> i64 { self.task_attempt_id } - /// Whether Spark killed this task. Older contexts without the callback return false. + /// Whether Spark killed this task, read from Spark each time. pub fn is_interrupted(&self) -> bool { - self.interrupted - .is_some_and(|callback| unsafe { callback(self.pointer) }) + match self.interruption { + Interruption::Never => false, + Interruption::Host { + context, callback, .. + } => unsafe { callback(context) }, + Interruption::Flag(flag) => flag.load(Ordering::SeqCst), + } + } +} + +/// Builds a [`TaskContext`] for testing a reader or writer outside Spark. +#[derive(Debug, Clone, Copy)] +pub struct TaskContextBuilder<'task> { + context: TaskContext<'task>, +} + +impl<'task> TaskContextBuilder<'task> { + #[must_use] + pub fn stage_id(mut self, stage_id: i32) -> Self { + self.context.stage_id = stage_id; + self + } + + #[must_use] + pub fn stage_attempt_number(mut self, stage_attempt_number: i32) -> Self { + self.context.stage_attempt_number = stage_attempt_number; + self + } + + #[must_use] + pub fn partition_id(mut self, partition_id: i32) -> Self { + self.context.partition_id = partition_id; + self + } + + #[must_use] + pub fn attempt_number(mut self, attempt_number: i32) -> Self { + self.context.attempt_number = attempt_number; + self + } + + #[must_use] + pub fn task_attempt_id(mut self, task_attempt_id: i64) -> Self { + self.context.task_attempt_id = task_attempt_id; + self + } + + /// Reports the task as interrupted whenever `flag` is true. + #[must_use] + pub fn interrupted_by<'flag>(self, flag: &'flag AtomicBool) -> TaskContextBuilder<'flag> { + let TaskContext { + stage_id, + stage_attempt_number, + partition_id, + attempt_number, + task_attempt_id, + .. + } = self.context; + TaskContextBuilder { + context: TaskContext { + stage_id, + stage_attempt_number, + partition_id, + attempt_number, + task_attempt_id, + interruption: Interruption::Flag(flag), + }, + } + } + + pub fn build(self) -> TaskContext<'task> { + self.context } } @@ -164,11 +248,11 @@ mod tests { interrupted: AtomicBool::new(false), }; let context = unsafe { TaskContext::from_raw(&host.task) }.unwrap(); - assert_eq!(context.stage_id(), Some(3)); - assert_eq!(context.stage_attempt_number(), Some(4)); - assert_eq!(context.partition_id(), Some(5)); - assert_eq!(context.attempt_number(), Some(6)); - assert_eq!(context.task_attempt_id(), Some(7)); + assert_eq!(context.stage_id(), 3); + assert_eq!(context.stage_attempt_number(), 4); + assert_eq!(context.partition_id(), 5); + assert_eq!(context.attempt_number(), 6); + assert_eq!(context.task_attempt_id(), 7); assert!(!context.is_interrupted()); host.interrupted.store(true, Ordering::SeqCst); std::thread::scope(|scope| { @@ -177,13 +261,42 @@ mod tests { } #[test] - fn only_reads_fields_present_in_the_declared_prefix() { + fn rejects_a_context_without_every_field_of_version_1() { let prefix = size_of::() as i64; let pointer = (&prefix as *const i64).cast::(); - let context = unsafe { TaskContext::from_raw(pointer) }.unwrap(); - assert_eq!(context.stage_id(), None); - assert_eq!(context.task_attempt_id(), None); + let error = unsafe { TaskContext::from_raw(pointer) }.unwrap_err(); + assert!(error.to_string().contains("version 1"), "{error}"); + assert!(unsafe { TaskContext::from_raw(std::ptr::null()) }.is_err()); + } + + #[test] + fn test_contexts_carry_ids_and_an_optional_interruption_flag() { + let context = TaskContext::for_testing() + .stage_id(1) + .stage_attempt_number(2) + .partition_id(3) + .attempt_number(4) + .task_attempt_id(5) + .build(); + assert_eq!( + ( + context.stage_id(), + context.stage_attempt_number(), + context.partition_id(), + context.attempt_number(), + context.task_attempt_id() + ), + (1, 2, 3, 4, 5) + ); + assert!(!context.is_interrupted()); + let flag = AtomicBool::new(false); + let context = TaskContext::for_testing() + .partition_id(8) + .interrupted_by(&flag) + .build(); + assert_eq!(context.partition_id(), 8); assert!(!context.is_interrupted()); - assert!(!TaskContext::detached().is_interrupted()); + flag.store(true, Ordering::SeqCst); + assert!(context.is_interrupted()); } } diff --git a/sql/native/rust/src/write.rs b/sql/native/rust/src/write.rs index c9794f1c358e2..0e9e53e5e565b 100644 --- a/sql/native/rust/src/write.rs +++ b/sql/native/rust/src/write.rs @@ -25,8 +25,11 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize}; use crate::abi::{self, ArrowArray, ArrowSchema, SparkDsTaskContext}; use crate::bridge::stream::{canonical_fields, validate_schema}; -use crate::bridge::{decode, encode, panic_message, provider_call}; -use crate::{BatchWrite, DataWriterFactory, Result, TaskContext, TaskWriter}; +use crate::bridge::{decode, describe, encode, panic_message, provider_call}; +use crate::{ + BatchWrite, DataWriterFactory, LogLevel, PartitionWriter, Result, TaskContext, WritePlan, + WriterContext, +}; type Cached = std::result::Result; @@ -60,16 +63,18 @@ struct WriterState { /// The driver keeps this coordinator until close, including after a failed commit. pub struct DriverWriter { - writer: Option, + coordinator: Option, + factory: Option, schema: SchemaRef, serialized: Option>>>, phase: Phase<()>, } impl DriverWriter { - pub fn new(writer: W, schema: SchemaRef) -> Self { + pub fn new(plan: WritePlan, schema: SchemaRef) -> Self { Self { - writer: Some(writer), + coordinator: Some(plan.coordinator), + factory: Some(plan.factory), schema, serialized: None, phase: Phase::Active, @@ -83,9 +88,8 @@ impl DriverWriter { self.phase.active()?; let result = provider_call(|| { validate_schema(&self.schema)?; - let factory = self.writer.as_mut().expect("writer is present").factory()?; Ok(Arc::new(encode(&WriterState { - factory, + factory: self.factory.as_ref().expect("factory is present"), schema: self.schema.as_ref().clone(), })?)) }); @@ -110,9 +114,9 @@ impl DriverWriter { .iter() .map(|message| decode(message.ok_or("missing committed task message")?)) .collect::>>()?; - self.writer + self.coordinator .as_mut() - .expect("writer is present") + .expect("coordinator is present") .commit(messages) }); self.phase = match &result { @@ -129,13 +133,29 @@ impl DriverWriter { Phase::Active | Phase::Failed(_) => {} } let result = provider_call(|| { + // Abort is the cleanup path: a message that cannot be decoded is passed as + // None instead of preventing the coordinator from cleaning up. let messages = messages .iter() - .map(|message| message.map(decode).transpose()) - .collect::>>()?; - self.writer + .enumerate() + .map(|(partition, message)| match decode((*message)?) { + Ok(message) => Some(message), + Err(error) => { + crate::log( + LogLevel::Warn, + &format!( + "Aborting a write without the commit message of partition \ + {partition}, which cannot be decoded: {}", + describe(&*error) + ), + ); + None + } + }) + .collect(); + self.coordinator .as_mut() - .expect("writer is present") + .expect("coordinator is present") .abort(messages) }); self.phase = Phase::Aborted(result.clone()); @@ -145,14 +165,17 @@ impl DriverWriter { impl Drop for DriverWriter { fn drop(&mut self) { - if let Some(writer) = self.writer.take() { - drop_provider(writer); + if let Some(coordinator) = self.coordinator.take() { + drop_provider(coordinator); + } + if let Some(factory) = self.factory.take() { + drop_provider(factory); } } } struct TaskState<'task, M> { - writer: Option>, + writer: Option>, schema: SchemaRef, phase: Phase>>, } @@ -186,7 +209,8 @@ pub unsafe fn open( let schema = Arc::new(schema); validate_schema(&schema)?; let task = unsafe { TaskContext::from_raw(task)? }; - let (writer, open_error) = match provider_call(|| factory.open(task)) { + let context = WriterContext::new(schema.clone(), task); + let (writer, open_error) = match provider_call(|| factory.open(context)) { Ok(writer) => (Some(writer), None), Err(message) => (None, Some(message)), }; diff --git a/sql/native/rust/tests/batch_write.rs b/sql/native/rust/tests/batch_write.rs index 85e446be6f4b7..5d474c5c0cf95 100644 --- a/sql/native/rust/tests/batch_write.rs +++ b/sql/native/rust/tests/batch_write.rs @@ -23,13 +23,14 @@ use arrow_array::ffi::to_ffi; use arrow_array::{Array, Int32Array, Int64Array, RecordBatch, StructArray}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use serde::{Deserialize, Serialize, Serializer}; -use spark_native_source::abi::{ +use spark_native_source::__private::abi::{ ArrowSchema, SparkDsCreateInfo, SparkDsHost, SparkDsLibrary, SparkDsTaskContext, SparkDsWriteInfo, ABI_VERSION, }; use spark_native_source::{ - BatchWrite, BatchWriteSource, DataSource, DataWriter, DataWriterFactory, Options, Result, - SourceCapability, TaskContext, TaskWriter, WriteMode, + BatchWrite, BatchWriteSource, Capabilities, CreateInfo, DataSource, DataWriter, + DataWriterFactory, PartitionWriter, Result, SourceCapability, TaskContext, WriteInfo, + WriteMode, WritePlan, WriterContext, }; // Each test exercises the actual exported table; this lock isolates fixture observations. @@ -39,7 +40,6 @@ static EVENTS: Mutex> = Mutex::new(Vec::new()); #[derive(Clone, Debug, PartialEq)] enum Event { WriteCreated(bool), - Factory, Batch(usize), TaskCommit, TaskAbort(bool), @@ -70,37 +70,46 @@ struct Source { } impl DataSource for Source { - fn new(options: Options) -> Result { + // A sink of any input: Spark requires a user-specified schema to read it. + const INFERS_SCHEMA: bool = false; + + fn new(info: CreateInfo) -> Result { Ok(Self { - failure: options.get("failure").cloned().unwrap_or_default(), - capabilities: options.get("capabilities").cloned().unwrap_or_default(), + failure: info.options.get("failure").cloned().unwrap_or_default(), + capabilities: info + .options + .get("capabilities") + .cloned() + .unwrap_or_default(), }) } - fn schema(&self) -> Result { - Ok(schema()) - } - - fn capabilities(&self, supported: &[SourceCapability]) -> Vec { + fn capabilities(&self, registered: Capabilities) -> Capabilities { match self.capabilities.as_str() { - "none" => vec![], - "append" => vec![SourceCapability::BatchWrite], - _ => supported.to_vec(), + "none" => Capabilities::empty(), + "append" => registered.without(SourceCapability::Truncate), + _ => registered, } } } impl BatchWriteSource for Source { - type Write = Driver; + type BatchWrite = Driver; + + const SUPPORTS_TRUNCATE: bool = true; - fn write(&self, input: SchemaRef, mode: WriteMode) -> Result { - if input != schema() { + fn new_write(&self, info: WriteInfo) -> Result> { + if info.schema != schema() { return Err("unsupported input schema".into()); } - record(Event::WriteCreated(matches!(mode, WriteMode::Truncate))); - Ok(Driver { + record(Event::WriteCreated(info.mode == WriteMode::Truncate)); + let driver = Driver { failure: self.failure.clone(), - }) + }; + let factory = Factory { + failure: self.failure.clone(), + }; + Ok(WritePlan::new(driver, factory)) } } @@ -111,13 +120,6 @@ struct Driver { impl BatchWrite for Driver { type Factory = Factory; - fn factory(&mut self) -> Result { - record(Event::Factory); - Ok(Factory { - failure: self.failure.clone(), - }) - } - fn commit(&mut self, messages: Vec) -> Result<()> { record(Event::DriverCommit(messages)); match self.failure.as_str() { @@ -172,9 +174,15 @@ fn serialize_message_mode( impl DataWriterFactory for Factory { type Message = Message; - fn open<'task>(&self, task: TaskContext<'task>) -> Result> { + fn open<'task>( + &self, + context: WriterContext<'task>, + ) -> Result> { + if context.schema != schema() { + return Err("the writer context lost the input schema".into()); + } Ok(Box::new(Writer { - task, + task: context.task, failure: self.failure.clone(), batches: vec![], })) @@ -226,8 +234,8 @@ impl DataWriter for Writer<'_> { }) .collect(); Ok(Message { - partition: self.task.partition_id().unwrap(), - attempt: self.task.task_attempt_id().unwrap(), + partition: self.task.partition_id(), + attempt: self.task.task_attempt_id(), values, batches: self.batches.len(), failure: self.failure.clone(), @@ -247,21 +255,22 @@ impl Drop for Writer<'_> { } } -spark_native_source::export_data_source!(Source, batch_write, truncate); +spark_native_source::export_data_source!(Source, batch_write); extern "C" { fn spark_ds_init(host: *const SparkDsHost, library: *mut *const SparkDsLibrary) -> c_int; } fn library() -> &'static SparkDsLibrary { - let host = SparkDsHost { + // Spark's host structure stays valid until the JVM exits. + let host = Box::leak(Box::new(SparkDsHost { size: size_of::() as i64, abi_version: ABI_VERSION, spark_version: c"test".as_ptr(), log: None, - }; + })); let mut table = ptr::null(); - assert_eq!(unsafe { spark_ds_init(&host, &mut table) }, 0); + assert_eq!(unsafe { spark_ds_init(host, &mut table) }, 0); assert!(!table.is_null()); unsafe { &*table } } @@ -492,6 +501,10 @@ fn write_registration_and_options_control_capabilities() { let _guard = TEST_LOCK.lock().unwrap(); EVENTS.lock().unwrap().clear(); let library = library(); + assert!( + library.schema.is_none(), + "INFERS_SCHEMA = false leaves schema inference to Spark's error" + ); assert!(library.create_reader.is_none()); assert!(library.read.is_none()); assert!(library.create_stream_reader.is_none()); @@ -573,7 +586,6 @@ fn serialized_factory_and_messages_roundtrip_without_driver_source_or_input_buff events(), vec![ Event::WriteCreated(true), - Event::Factory, Event::Batch(3), Event::Batch(0), Event::TaskCommit, @@ -676,6 +688,8 @@ fn task_failures_and_message_serialization_failures_remain_abortable() { before_retry, "failed commit cannot run provider commit twice" ); + // A commit whose message cannot be serialized counts as failed, so abort follows + // a provider commit that returned a message. success(library, unsafe { library.abort_data_writer.unwrap()(task.pointer) }); @@ -757,10 +771,35 @@ fn factory_serialization_failure_is_cached_and_driver_can_abort() { success(library, driver_finish(&driver, true, &[])); assert_eq!( events(), - vec![ - Event::WriteCreated(false), - Event::Factory, - Event::DriverAbort(vec![]) - ], + vec![Event::WriteCreated(false), Event::DriverAbort(vec![])], ); } + +#[test] +fn undecodable_messages_fail_commit_but_abort_still_cleans_up() { + let _guard = TEST_LOCK.lock().unwrap(); + EVENTS.lock().unwrap().clear(); + let library = library(); + let (source, _) = source(library, "", ""); + let driver = driver(&source, 0); + let task = task(library, &state(&driver), 2); + success(library, write(&task, batch(vec![6]))); + let message = task_message(&task); + drop(task); + + let garbage = b"not a commit message"; + assert_ne!(driver_finish(&driver, false, &[Some(garbage)]), 0); + assert!(!events() + .iter() + .any(|event| matches!(event, Event::DriverCommit(_)))); + let aborted = [Some(&garbage[..]), Some(message.as_slice()), None]; + success(library, driver_finish(&driver, true, &aborted)); + let expected = Message { + partition: 2, + attempt: 102, + values: vec![6], + batches: 1, + failure: String::new(), + }; + assert!(events().contains(&Event::DriverAbort(vec![None, Some(expected), None]))); +} diff --git a/sql/native/rust/tests/capabilities.rs b/sql/native/rust/tests/capabilities.rs index 80a16a8bda28b..65ba31dc0a928 100644 --- a/sql/native/rust/tests/capabilities.rs +++ b/sql/native/rust/tests/capabilities.rs @@ -13,49 +13,60 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::ffi::{c_int, c_void, CStr, CString}; +use std::ffi::{c_char, c_int, c_void, CStr, CString}; use std::mem::size_of; use std::ptr; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use arrow_schema::{Schema, SchemaRef}; use serde::{Deserialize, Serialize}; -use spark_native_source::abi::{ +use spark_native_source::__private::abi::{ ArrowSchema, SparkDsCreateInfo, SparkDsHost, SparkDsLibrary, ABI_VERSION, }; use spark_native_source::{ - BatchReader, BatchReaderFactory, BatchSource, DataSource, Options, Result, ScanBuilder, - ScanPlan, SourceCapability, TaskContext, + BatchReadSource, Capabilities, CreateInfo, DataSource, HostInfo, LogLevel, PartitionReader, + PartitionReaderFactory, ReaderContext, Result, ScanBuilder, ScanPlan, SourceCapability, }; static DECLARATIONS: AtomicUsize = AtomicUsize::new(0); static SCHEMAS: AtomicUsize = AtomicUsize::new(0); static SCANS: AtomicUsize = AtomicUsize::new(0); static DROPS: AtomicUsize = AtomicUsize::new(0); +static INITIALIZED: Mutex>> = Mutex::new(Vec::new()); +static LOGS: Mutex> = Mutex::new(Vec::new()); struct Source { mode: String, } impl DataSource for Source { - fn new(options: Options) -> Result { + fn init(host: &HostInfo) -> Result<()> { + INITIALIZED.lock().unwrap().push(host.spark_version.clone()); + if host.spark_version.as_deref() == Some("unsupported") { + return Err("this library does not support the Spark version".into()); + } + spark_native_source::log(LogLevel::Info, "capability test library initialized"); + Ok(()) + } + + fn new(info: CreateInfo) -> Result { + if info.name != "capability_test" { + return Err(format!("unexpected source name {}", info.name).into()); + } Ok(Self { - mode: options.get("capabilities").cloned().unwrap(), + mode: info.options.get("capabilities").cloned().unwrap(), }) } - fn capabilities(&self, _: &[SourceCapability]) -> Vec { + fn capabilities(&self, registered: Capabilities) -> Capabilities { DECLARATIONS.fetch_add(1, Ordering::SeqCst); match self.mode.as_str() { - "empty" => vec![], - "batch" => vec![SourceCapability::BatchRead], - "duplicates" => vec![SourceCapability::BatchRead, SourceCapability::BatchRead], - "micro_batch" => vec![SourceCapability::MicroBatchRead], - "batch_write" => vec![SourceCapability::BatchWrite], - "read_write" => vec![SourceCapability::BatchRead, SourceCapability::BatchWrite], - "streaming_write" => vec![SourceCapability::StreamingWrite], - "truncate" => vec![SourceCapability::Truncate], + "empty" => Capabilities::empty(), + "registered" => registered, + "without_read" => registered.without(SourceCapability::BatchRead), + // Operations that the library did not register are dropped. + "all" => Capabilities::all(), "panic" => panic!("capability discovery failed"), other => panic!("unexpected test mode: {other}"), } @@ -67,10 +78,10 @@ impl DataSource for Source { } } -impl BatchSource for Source { - type Scan = Builder; +impl BatchReadSource for Source { + type ScanBuilder = Builder; - fn scan(&self, _: SchemaRef) -> Result { + fn new_scan(&self, _: SchemaRef) -> Result { SCANS.fetch_add(1, Ordering::SeqCst); Ok(Builder) } @@ -88,30 +99,45 @@ impl ScanBuilder for Builder { type Factory = Factory; fn plan(self) -> Result> { - Ok(ScanPlan { - factory: Factory, - partitions: vec![], - }) + Ok(ScanPlan::new(Factory, vec![])) } } #[derive(Serialize, Deserialize)] struct Factory; -impl BatchReaderFactory for Factory { +impl PartitionReaderFactory for Factory { type Partition = (); - fn open<'task>(&self, _: (), _: TaskContext<'task>) -> Result> { + fn open<'task>(&self, _: (), _: ReaderContext<'task>) -> Result> { Err("capability tests do not open partition readers".into()) } } -spark_native_source::export_data_source!(Source); +spark_native_source::export_data_source!(Source, batch_read); extern "C" { fn spark_ds_init(host: *const SparkDsHost, library: *mut *const SparkDsLibrary) -> c_int; } +unsafe extern "C" fn record_log(level: i32, message: *const c_char) { + let message = unsafe { CStr::from_ptr(message) } + .to_str() + .unwrap() + .to_owned(); + LOGS.lock().unwrap().push((level, message)); +} + +/// A host that, like Spark's, stays valid until the process exits. +fn host(spark_version: &'static CStr) -> &'static SparkDsHost { + Box::leak(Box::new(SparkDsHost { + size: size_of::() as i64, + abi_version: ABI_VERSION, + spark_version: spark_version.as_ptr(), + log: Some(record_log), + })) +} + struct SourceHandle<'a> { library: &'a SparkDsLibrary, pointer: *mut c_void, @@ -123,8 +149,8 @@ impl Drop for SourceHandle<'_> { } } -fn create(library: &SparkDsLibrary, mode: &str) -> (c_int, u64, *mut c_void) { - let name = CString::new("capability_test").unwrap(); +fn create(library: &SparkDsLibrary, name: &str, mode: &str) -> (c_int, u64, *mut c_void) { + let name = CString::new(name).unwrap(); let key = CString::new("capabilities").unwrap(); let value = CString::new(mode).unwrap(); let keys = [key.as_ptr()]; @@ -149,22 +175,45 @@ fn error(library: &SparkDsLibrary) -> String { } #[test] -fn source_capabilities_control_the_exported_abi_and_source_lifecycle() { +fn initialization_logging_and_capabilities_follow_the_source_lifecycle() { // One lifecycle test owns these counters; no parallel test resets shared state. - let spark_version = CString::new("test").unwrap(); - let host = SparkDsHost { - size: size_of::() as i64, - abi_version: ABI_VERSION, - spark_version: spark_version.as_ptr(), - log: None, - }; let mut table = ptr::null(); - assert_eq!(unsafe { spark_ds_init(&host, &mut table) }, 0); + assert_ne!( + unsafe { spark_ds_init(host(c"unsupported"), &mut table) }, + 0 + ); + assert!( + table.is_null(), + "a library whose init fails returns no table" + ); + assert_eq!( + LOGS.lock().unwrap().as_slice(), + [( + 1, + "this library does not support the Spark version".to_owned() + )], + ); + assert_eq!(unsafe { spark_ds_init(host(c"4.4.0-test"), &mut table) }, 0); assert!(!table.is_null()); let library = unsafe { &*table }; + assert_eq!( + INITIALIZED.lock().unwrap().as_slice(), + [ + Some("unsupported".to_owned()), + Some("4.4.0-test".to_owned()) + ], + ); + spark_native_source::log(LogLevel::Warn, "a message with a \0 byte"); + assert_eq!( + LOGS.lock().unwrap()[1..], + [ + (3, "capability test library initialized".to_owned()), + (2, "a message with a \\0 byte".to_owned()), + ], + ); // Analysis can ask only for capabilities, then close without schema or scan discovery. - let (status, mask, pointer) = create(library, "batch"); + let (status, mask, pointer) = create(library, "capability_test", "registered"); assert_eq!(status, 0); assert_eq!(mask, 1); assert!(!pointer.is_null()); @@ -175,17 +224,22 @@ fn source_capabilities_control_the_exported_abi_and_source_lifecycle() { assert_eq!(DROPS.load(Ordering::SeqCst), 1); // Keep all instances alive together: the last declaration must not affect earlier sources. - let sources: Vec<_> = [("empty", 0), ("batch", 1), ("duplicates", 1)] - .into_iter() - .map(|(mode, expected)| { - let (status, mask, pointer) = create(library, mode); - assert_eq!(status, 0); - assert_eq!(mask, expected, "{mode}"); - assert!(!pointer.is_null()); - (expected, SourceHandle { library, pointer }) - }) - .collect(); - assert_eq!(DECLARATIONS.load(Ordering::SeqCst), 4); + let sources: Vec<_> = [ + ("empty", 0), + ("registered", 1), + ("all", 1), + ("without_read", 0), + ] + .into_iter() + .map(|(mode, expected)| { + let (status, mask, pointer) = create(library, "capability_test", mode); + assert_eq!(status, 0); + assert_eq!(mask, expected, "{mode}"); + assert!(!pointer.is_null()); + (expected, SourceHandle { library, pointer }) + }) + .collect(); + assert_eq!(DECLARATIONS.load(Ordering::SeqCst), 5); for (mask, source) in sources { let mut schema = ArrowSchema::empty(); @@ -208,35 +262,21 @@ fn source_capabilities_control_the_exported_abi_and_source_lifecycle() { assert_eq!(SCANS.load(Ordering::SeqCst), scans_before + 1); unsafe { library.close_reader.unwrap()(reader) }; } - assert_eq!(DECLARATIONS.load(Ordering::SeqCst), 4); + assert_eq!(DECLARATIONS.load(Ordering::SeqCst), 5); } - assert_eq!(SCHEMAS.load(Ordering::SeqCst), 3); + assert_eq!(SCHEMAS.load(Ordering::SeqCst), 4); assert_eq!(SCANS.load(Ordering::SeqCst), 2); - assert_eq!(DROPS.load(Ordering::SeqCst), 4); - - for (mode, capability) in [ - ("micro_batch", "MicroBatchRead"), - ("batch_write", "BatchWrite"), - ("read_write", "BatchWrite"), - ("streaming_write", "StreamingWrite"), - ("truncate", "Truncate"), - ] { - let drops_before = DROPS.load(Ordering::SeqCst); - let declarations_before = DECLARATIONS.load(Ordering::SeqCst); - let (status, mask, source) = create(library, mode); - assert_ne!(status, 0, "{mode}"); - assert_eq!( - error(library), - format!("capability {capability} is not registered") - ); - assert_eq!(mask, 0); - assert!(source.is_null()); - assert_eq!(DROPS.load(Ordering::SeqCst), drops_before + 1); - assert_eq!(DECLARATIONS.load(Ordering::SeqCst), declarations_before + 1); - } + assert_eq!(DROPS.load(Ordering::SeqCst), 5); + + // The source receives the name that Spark passes, so it can reject other names. + let (status, mask, source) = create(library, "other_source", "registered"); + assert_ne!(status, 0); + assert_eq!(error(library), "unexpected source name other_source"); + assert_eq!(mask, 0); + assert!(source.is_null()); let drops_before = DROPS.load(Ordering::SeqCst); - let (status, mask, source) = create(library, "panic"); + let (status, mask, source) = create(library, "capability_test", "panic"); assert_ne!(status, 0); assert_eq!( error(library), @@ -245,6 +285,6 @@ fn source_capabilities_control_the_exported_abi_and_source_lifecycle() { assert_eq!(mask, 0); assert!(source.is_null()); assert_eq!(DROPS.load(Ordering::SeqCst), drops_before + 1); - assert_eq!(SCHEMAS.load(Ordering::SeqCst), 3); + assert_eq!(SCHEMAS.load(Ordering::SeqCst), 4); assert_eq!(SCANS.load(Ordering::SeqCst), 2); } diff --git a/sql/native/rust/tests/fixtures/spark_datasource.h b/sql/native/rust/tests/fixtures/spark_datasource.h index 16e21494dba85..db21656252e08 100644 --- a/sql/native/rust/tests/fixtures/spark_datasource.h +++ b/sql/native/rust/tests/fixtures/spark_datasource.h @@ -117,7 +117,8 @@ * source handle right after creating it. * - Memory. What Spark passes to a function, such as strings, byte arrays, SparkDsCreateInfo and * SparkDsWriteInfo, is only valid during the call, except for SparkDsHost and - * SparkDsTaskContext, which document how long they live. Strings are UTF-8 and NUL-terminated. + * SparkDsTaskContext, which document how long they live. Strings are UTF-8 and NUL-terminated: + * Spark fails an operation whose options or column names contain the character U+0000. * The JSON of the predicates is ASCII: it escapes the other characters, which any JSON parser * decodes. The library returns strings and byte arrays by passing them to a SparkDsSink, which * copies them. @@ -134,12 +135,21 @@ * parent refers to, such as a column shorter than its batch. The buffers of an array must hold * all of its values, as the C data interface requires: it does not carry the sizes of the * buffers, so Spark cannot check them, and reads past the end of a buffer that is too short. - * - Schemas. Spark types map to Arrow types as in spark.sql.execution.arrow: for example, bigint - * is int64, string is utf8, timestamp is timestamp[us] with a time zone and timestamp_ntz is - * timestamp[us] without one. Data read from a library may also use large_utf8 and utf8_view for - * strings, large_binary and binary_view for binary data, and any names for the fields of the - * entries of a map. Spark rejects the other Arrow types, such as unsigned integers, decimal256, - * large lists, fixed-size lists and list views, and dictionary encoding. + * - Schemas. These are the types of version 1, with their Arrow types and the formats of the C + * data interface: boolean is bool (b), tinyint, smallint, int and bigint are int8 (c), int16 + * (s), int32 (i) and int64 (l), float and double are float32 (f) and float64 (g), decimal(p, s) + * is decimal128(p, s) (d:p,s), string is utf8 (u), binary is binary (z), date is date32 (tdD), + * timestamp is timestamp[us] with the time zone of the session (tsu:), timestamp_ntz is + * timestamp[us] without a time zone (tsu:), and array, map and struct are list (+l), map (+m) + * and struct (+s). Spark fails an operation whose schema has another type, such as an + * interval, variant or time, before it passes the schema to the library: a later version may + * support other types, for the data sources that declare a capability for them. Data read from + * a library may also use large_utf8 (U) and utf8_view (vu) for strings, large_binary (Z) and + * binary_view (vz) for binary data, and any names for the fields of the entries of a map. + * Spark rejects the other Arrow types, such as unsigned integers, decimal256, large lists, + * fixed-size lists and list views, and dictionary encoding. The data must agree with the + * nullability of the schema to read: Spark fails a read whose data has a null where the schema + * declares a column, a field, an element or a map value not nullable. * - Options. Option keys are passed in lower case, because Spark matches them * case-insensitively. * - Optional functions. A library only sets the functions of the operations it supports in its @@ -443,7 +453,10 @@ typedef struct SparkDsLibrary { * SQL type name such as int, string or decimal(10,2). The value is null, a JSON boolean or * number for boolean and numeric types, a string for string, a base64 string for binary, a * decimal string for decimal, the days since the epoch for date, and the microseconds since - * the epoch for timestamp and timestamp_ntz. + * the epoch for timestamp and timestamp_ntz. Read a literal at the precision of its data type: + * a float literal is the exact value of the float, written as the double of the same value, + * and the values of the integer types, dates and timestamps are exact 64-bit integers, which + * a double does not hold above 2^53. Spark does not pass a string that is not valid UTF-8. * - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in * org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=> * (equal, with null equal to null), <, <=, > and >= with two children, IN with the value to @@ -471,8 +484,11 @@ typedef struct SparkDsLibrary { * example, NOT (i IN (1, NULL)) matches no row. * - Strings are compared by their UTF-8 bytes, which order the characters outside the Basic * Multilingual Plane after all the others, unlike comparisons of UTF-16 strings. - * - To use only a part of a predicate, a reader may take the parts that it does not support as - * true, but not under a NOT, where that would drop rows. + * - To use only a part of a predicate, a reader may take a part that it does not support as + * true, but only when every function between that part and the root of the predicate is AND + * or OR. Under any other function, such as NOT, or = and <=> between two predicates, that + * would drop rows: (id > 3) = (id < 3) is true for the row whose id is 3, but taking id < 3 + * as true makes it false. * The file org/apache/spark/sql/datasource/predicate-cases.json in the jar of Spark SQL lists * predicates as Spark passes them to this function, with the rows of a table for which they are * true, to test a reader with.