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..b00bca5d1cf3b --- /dev/null +++ b/examples/src/main/rust/native-data-sink/README.md @@ -0,0 +1,139 @@ + + +# Rust native batch sink + +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)`, 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 `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. + +## 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, 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: + +- 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 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. +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..3705bb0a46d01 --- /dev/null +++ b/examples/src/main/rust/native-data-sink/src/lib.rs @@ -0,0 +1,757 @@ +// 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, Capabilities, CreateInfo, DataSource, DataWriter, + DataWriterFactory, PartitionWriter, Result, SourceCapability, TaskContext, WriteInfo, + WriteMode, WritePlan, WriterContext, +}; + +spark_native_source::export_data_source!(JsonSink, batch_write); + +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(info: CreateInfo) -> Result { + let mut options = info.options; + 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, registered: Capabilities) -> Capabilities { + if self.readonly { + // Also removes Truncate, which requires batch writes. + registered.without(SourceCapability::BatchWrite) + } else { + registered + } + } +} + +impl BatchWriteSource for JsonSink { + type BatchWrite = JsonWrite; + + const SUPPORTS_TRUNCATE: bool = true; + + fn new_write(&self, info: WriteInfo) -> Result> { + if self.readonly { + return Err("This sink is readonly".into()); + } + validate_schema(&info.schema)?; + let write = JsonWrite::new(self.path.clone(), info.mode)?; + let factory = write.writer_factory(); + Ok(WritePlan::new(write, factory)) + } +} + +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); + } +} + +pub struct JsonWrite { + root: PathBuf, + operation: String, + staging: PathBuf, + objects: PathBuf, + temporary_manifest: PathBuf, + previous: Manifest, + truncate: bool, + committed: bool, + 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()); + } + 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)?; + 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, + truncate, + committed: false, + sealed: false, + lock: Some(lock), + }) + } + + 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 (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()); + } + // 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()); + } + } + + // 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 = 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), + 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.committed = true; + 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(()), + } + } +} + +// 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 commit(&mut self, messages: Vec) -> Result<()> { + self.publish(messages) + } + + fn abort(&mut self, _messages: Vec>) -> Result<()> { + let result = self.cleanup(); + self.lock.take(); + result + } +} + +// 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 { + 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: i32, + rows: u64, +} + +impl DataWriterFactory for JsonWriterFactory { + type Message = CommitMessage; + + 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()); + } + // 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(), + task.task_attempt_id(), + 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)), + committed: 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>, + committed: 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); + } + } +} + +// 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<()> { + let writer = self.writer.as_mut().ok_or("The task writer is closed")?; + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or("The id column must be an Arrow int64 array")?; + let labels = batch + .column(1) + .as_any() + .downcast_ref::() + .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()); + } + 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 commit(&mut self) -> Result { + 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.committed = true; + Ok(self.message.clone()) + } + + fn abort(&mut self) -> Result<()> { + self.discard_buffer(); + 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.committed { + 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(info(&[("path", self.0.to_str().unwrap())])).unwrap() + } + } + + impl Drop for TempDirectory { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } + } + + 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() + .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, 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 = open(&factory, partition); + 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, factory) = new_write(source, mode).unwrap(); + let message = stage(&factory, 0, 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, 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()); + 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"})] + ); + 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] + 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, 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()); + 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); + let task = TaskContext::for_testing().build(); + assert!(factory + .open(WriterContext::new(sink_schema(), task)) + .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, 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()); + 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 (write, factory) = new_write(&source, WriteMode::Append).unwrap(); + let _message = stage(&factory, 0, vec![Some(1)]); + let staging = write.staging.clone(); + assert!(new_write(&source, WriteMode::Append).is_err()); + drop(write); + assert!(!staging.exists()); + assert!(!directory.0.join(".write-lock").exists()); + 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(info(&[ + ("path", path.to_str().unwrap()), + ("readonly", "true"), + ])) + .unwrap(); + 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/.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..85a62b7be5f6c --- /dev/null +++ b/examples/src/main/rust/native-data-source/README.md @@ -0,0 +1,147 @@ + + +# 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!(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 +`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 the `BatchRead` capability and the native +mask `SPARK_DS_BATCH_READ` (`1`); the connector does not maintain a separate capability list. +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 +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 + +| 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, 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 +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::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 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 +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 `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. + +## 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. 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/package.py b/examples/src/main/rust/native-data-source/package.py new file mode 100644 index 0000000000000..36d440291fa51 --- /dev/null +++ b/examples/src/main/rust/native-data-source/package.py @@ -0,0 +1,71 @@ +#!/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("--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") + 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": [args.source], + "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..c6b2ddb90749d --- /dev/null +++ b/examples/src/main/rust/native-data-source/smoke.py @@ -0,0 +1,102 @@ +#!/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.errors import AnalysisException +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") + .config("spark.sql.adaptive.enabled", "true") + .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] + + # 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() + + +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..50bbdde517541 --- /dev/null +++ b/examples/src/main/rust/native-data-source/src/lib.rs @@ -0,0 +1,616 @@ +// 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::{ + BatchReadSource, CreateInfo, DataSource, FilterSupport, LimitSupport, LiteralValue, Options, + PartitionReader, PartitionReaderFactory, Predicate, ReaderContext, Result, ScanBuilder, + ScanPlan, +}; + +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, + 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 { + batch_size: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RangePartition { + start: i64, + end: i64, +} + +impl DataSource for RangeSource { + fn new(info: CreateInfo) -> Result { + let options = info.options; + 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()]) + } +} + +impl BatchReadSource for RangeSource { + type ScanBuilder = RangeScan; + + 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()) { + 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, + columns, + factory: RangeReaderFactory { + batch_size: self.batch_size, + }, + }) + } +} + +impl PartitionReaderFactory for RangeReaderFactory { + type Partition = RangePartition; + + fn open<'task>( + &self, + partition: RangePartition, + 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()); + } + // 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 || { + if current == partition.end { + return None; + } + if task.is_interrupted() { + current = partition.end; + return Some(Err(ArrowError::ComputeError( + "Range read was cancelled".into(), + ))); + } + let next = + (i128::from(current) + batch_size as i128).min(i128::from(partition.end)) as i64; + let num_rows = (i128::from(next) - i128::from(current)) as usize; + let arrays: Vec = 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), + arrays, + &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.columns.contains(column)) { + return Err("Projection contains a column outside the read schema".into()); + } + output_schema(columns)?; + self.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::new(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") +} + +/// A non-null tinyint, smallint, int, or bigint literal, decoded by the SDK. +fn integer_literal(predicate: &Predicate) -> Option { + match predicate.literal_value()? { + LiteralValue::Integer(value) => Some(value), + _ => None, + } +} + +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(info(values)).unwrap() + } + + fn scan(source: &RangeSource) -> RangeScan { + source.new_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, + } + } + + /// 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() + .enumerate() + .flat_map(|(index, partition)| { + let task = TaskContext::for_testing() + .partition_id(index as i32) + .build(); + plan.factory + .open(partition, ReaderContext::new(schema.clone(), task)) + .unwrap() + .map(|batch| batch.unwrap()) + .collect::>() + }) + .collect() + } + + const ALL: &[&str] = &["id", "label"]; + + 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, &["label", "id"]); + 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(), ALL)), + 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(), ALL)), + (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(), ALL)), 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::new(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, ALL)), 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(info(&[(key, value)])).is_err()); + } + let source = source(&[]); + let invalid = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + assert!(source.new_scan(invalid).is_err()); + let labels = output_schema(&["label".into()]).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()); + } + + #[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 context = + ReaderContext::new(source.schema().unwrap(), TaskContext::for_testing().build()); + let mut reader = plan + .factory + .open(plan.partitions[0].clone(), context) + .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..e20497c372e75 --- /dev/null +++ b/examples/src/main/rust/native-data-source/tests/c_abi_smoke.c @@ -0,0 +1,611 @@ +/* + * 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 +#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(void) { + 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 + }; + 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 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) { + 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); + 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)); + 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"); + } + + return finish_plan(reader); +} + +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) { + int64_t rows = 0; + for (size_t partition = 0; partition < plan->partitions.length; partition++) { + SparkDsTaskContext task = task_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_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); + 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(); + 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; +} + +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"); +} + +/* 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(); + 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"}; + 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"}; + + 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); + 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); + 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); + 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); + 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); + check_short_task_context(&cancellable); + check_cancellation(&cancellable); + check_cross_thread_stream(&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..b5d63b3f7ad35 --- /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", "-pthread"]) + .arg("-I") + .arg(header) + .arg(manifest.join("tests/c_abi_smoke.c")) + .arg("-o") + .arg(&harness); + if cfg!(target_os = "linux") { + command.arg("-ldl"); + } + let compiled = command.output()?; + assert!( + compiled.status.success(), + "C ABI harness did not compile:\n{}", + String::from_utf8_lossy(&compiled.stderr) + ); + + let executable = std::env::current_exe()?; + let profile = executable + .parent() + .and_then(Path::parent) + .ok_or("Missing Cargo profile directory")?; + let library = profile.join(format!( + "{}spark_datasource_rust_sdk_range{}", + std::env::consts::DLL_PREFIX, + std::env::consts::DLL_SUFFIX + )); + assert!( + library.is_file(), + "Run cargo build --workspace first: {}", + library.display() + ); + + 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..889509f9c4d6a --- /dev/null +++ b/sql/native/rust/Cargo.lock @@ -0,0 +1,675 @@ +# 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-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" +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..5568b681d65ee --- /dev/null +++ b/sql/native/rust/Cargo.toml @@ -0,0 +1,54 @@ +# 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", + "../../../examples/src/main/rust/native-data-sink", +] +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..4e63e876829a9 --- /dev/null +++ b/sql/native/rust/README.md @@ -0,0 +1,393 @@ + + +# Rust native data source SDK prototype + +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 +`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. + +The Rust and C tests run independently of Spark and Java. The workspace is not +part of Maven/SBT and introduces no Rust requirement for building Spark. This is +an API prototype, not a published crate or a compatibility commitment. + +The [Range connector][range] is a separate workspace member under +`examples/src/main/rust/native-data-source`, with a C integration test, packaging +helper, and Spark smoke script. The [file sink example][sink] is a separate write-only connector. All crates +share this directory's lockfile and `target` directory. Cargo commands select the SDK by default; use `--workspace` +to include the example. + +## Connector API and lifecycle + +A connector implements [`DataSource`](src/lib.rs) for construction, schema discovery, +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: + +* `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`. +* `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 +must read the same snapshot. Multiple task attempts open independent readers. + +```text +Driver: + 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 +`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 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. 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, 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`. 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 + +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); +export_data_source!(ReadWriteSource, batch_read, batch_write); +``` + +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. 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, 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` 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 `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 +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 -> 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::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` 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, 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. + +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 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 +single-writer assumption and orphan cleanup limitations are example constraints, +not transaction guarantees provided by the SDK. + +## Arrow contract + +`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 +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 +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. +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 +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 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 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 +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` and pthreads 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 +capability-only source discovery and optional slots, independent driver planners, +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, 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 + +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, 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. + +The [sink guide][sink] uses the same packager with `--source rust_sdk_sink`. Its +Spark smoke script checks append, overwrite, failed-overwrite cleanup, and +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, 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 +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 + +[sink]: ../../../examples/src/main/rust/native-data-sink/README.md diff --git a/sql/native/rust/src/abi.rs b/sql/native/rust/src/abi.rs new file mode 100644 index 0000000000000..9a7489ad588a3 --- /dev/null +++ b/sql/native/rust/src/abi.rs @@ -0,0 +1,501 @@ +// 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::{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; +pub const BATCH_WRITE: u64 = 4; +pub const TRUNCATE: u64 = 16; +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| describe(&*e)))) { + 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(()) +} + +/// 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 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> = + 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 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")?; + let options: Options = unsafe { strings(keys, count)? } + .into_iter() + .zip(unsafe { strings(values, count)? }) + .collect(); + Ok(CreateInfo::new(name, options)) +} + +/// # 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, + 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 { 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] + 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 new file mode 100644 index 0000000000000..03cbafbd9779b --- /dev/null +++ b/sql/native/rust/src/bridge.rs @@ -0,0 +1,949 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Driver-only handles and planners, with stateless 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::{ + Capabilities, CreateInfo, DataSource, FilterSupport, LimitSupport, LogLevel, + PartitionReaderFactory, Predicate, ReaderContext, Result, ScanBuilder, SourceCapability, + TaskContext, WriteInfo, WriteMode, +}; + +use crate::provider::Provider; +use crate::write::DriverWriter; + +#[path = "stream.rs"] +pub(crate) 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>, + planners: HandleMap>, + writers: HandleMap>, +} + +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

{ + 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, info: CreateInfo) -> Result<(i64, u64)> { + self.register_source(P::Source::new(info)?) + } + + fn register_source(&self, source: P::Source) -> Result<(i64, u64)> { + // 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 { + source, + capabilities, + }, + )?; + Ok((handle, capabilities.bits())) + } + + pub fn schema(&self, handle: i64) -> Result { + let source = get(&self.sources, handle)?; + let schema = lock(&source).source.schema()?; + stream::validate_schema(&schema)?; + Ok(schema) + } + + pub fn create_reader(&self, handle: i64, schema: SchemaRef) -> Result { + stream::validate_schema(&schema)?; + let source = get(&self.sources, handle)?; + let builder = { + let source = lock(&source); + if !source.capabilities.contains(SourceCapability::BatchRead) { + return Err("source does not support batch reads".into()); + } + 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: Vec>, + ) -> Result> { + let reader = get(&self.planners, handle)?; + let mut reader = lock(&reader); + 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) + } + + pub fn push_limit(&self, handle: i64, limit: usize) -> Result { + let reader = get(&self.planners, 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.planners, 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.planners, handle)?; + let result = lock(&reader).plan(); + result + } + + pub fn create_writer(&self, handle: i64, info: WriteInfo) -> Result { + stream::validate_schema(&info.schema)?; + let source = get(&self.sources, handle)?; + let schema = info.schema.clone(); + let plan = { + let source = lock(&source); + if !source.capabilities.contains(SourceCapability::BatchWrite) { + return Err("source does not support batch writes".into()); + } + if info.mode == WriteMode::Truncate + && !source.capabilities.contains(SourceCapability::Truncate) + { + return Err("source does not support truncate".into()); + } + P::new_write(&source.source, info)? + }; + self.insert(&self.writers, DriverWriter::new(plan, 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(()) + } + + pub fn close_reader(&self, handle: i64) -> Result<()> { + remove(&self.planners, handle); + Ok(()) + } +} + +/// 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 ReaderState { factory, schema }: ReaderState = decode(state)?; + let schema = Arc::new(schema); + let task = unsafe { TaskContext::from_raw(task)? }; + let reader = factory.open(decode(partition)?, ReaderContext::new(schema.clone(), task))?; + unsafe { stream::export(reader, &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 SourceState { + source: S, + capabilities: Capabilities, +} + +struct Planner { + 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 Planner { + 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, +} + +pub(crate) fn encode(value: &T) -> Result> { + Ok(serde_json::to_vec(&Envelope { + version: STATE_VERSION, + value, + })?) +} + +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()); + } + Ok(envelope.value) +} + +pub(crate) fn provider_call( + action: impl FnOnce() -> Result, +) -> std::result::Result { + catch_unwind(AssertUnwindSafe(|| { + 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() + } 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}") +} + +/// 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| 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() +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicUsize; + + use arrow_array::ffi_stream::ArrowArrayStreamReader; + use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator}; + use arrow_schema::{DataType, Field}; + + use super::*; + use crate::provider::ReadOnly; + use crate::{BatchReadSource, PartitionReader, 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 DataSource for Source { + fn new(_: CreateInfo) -> 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())) + } + } + + impl BatchReadSource for Source { + type ScanBuilder = Builder; + + fn new_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 PartitionReaderFactory for Factory { + type Partition = u64; + + // 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])?; + 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::new(Factory(self.value), vec![2, 7])) + } + } + + fn info() -> CreateInfo { + CreateInfo::new("bridge_test", Default::default()) + } + + 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(info()).unwrap(); + let builds = source.builds.clone(); + let drops = source.drops.clone(); + 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(); + 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(info()).unwrap(); + source.filter_support = support; + 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(); + assert_eq!( + bridge + .push_predicates(reader, vec![Some(Predicate::Unknown)]) + .unwrap(), + vec![support] + ); + assert!(!bridge.push_limit(reader, 1).unwrap()); + assert_eq!(limits.load(Ordering::SeqCst), 0); + } + } + + #[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::Unsupported, LimitSupport::Global] { + let bridge = Bridge::>::default(); + let mut source = Source::new(info()).unwrap(); + source.limit_support = support; + let limits = source.limits.clone(); + let (handle, _) = bridge.register_source(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(info()).unwrap(); + source.failure = failure; + let builds = source.builds.clone(); + let drops = source.builder_drops.clone(); + let (handle, _) = bridge.register_source(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(info()).unwrap(); + source.failure = failure; + let builds = source.builds.clone(); + let limits = source.limits.clone(); + let drops = source.builder_drops.clone(); + let (handle, _) = bridge.register_source(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, vec![])), + 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(info()).unwrap(); + source.filter_support = FilterSupport::PruningOnly; + source.limit_support = LimitSupport::Global; + let builds = source.builds.clone(); + let limits = source.limits.clone(); + let (handle, _) = bridge.register_source(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, vec![Some(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_retries_use_only_serialized_values_after_the_driver_is_dropped() { + let bridge = Bridge::>::default(); + let source = Source::new(info()).unwrap(); + let builds = source.builds.clone(); + let source_drops = source.drops.clone(); + let builder_drops = source.builder_drops.clone(); + 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(); + 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(); + let (handle, _) = bridge.create(info()).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. + drop(bridge); + let error = error_message(unsafe { + open_reader::(&task, &plan.state, &plan.partitions[0]) + }); + 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(); + 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..307239c4c8526 --- /dev/null +++ b/sql/native/rust/src/export.rs @@ -0,0 +1,420 @@ +// 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 a data source through version 1 of Spark's C data source ABI. +/// +/// 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 { + (@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)*); + }; + (@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" + ); + }; + (@parse $source:ty; [$read:tt $write:tt]; $other:ident $($rest:ident)*) => { + compile_error!(concat!( + "unknown operation `", + stringify!($other), + "`: expected batch_read or batch_write" + )); + }; + (@parse $source:ty; [true false];) => { + $crate::export_data_source!(@provider $crate::__private::provider::ReadOnly<$source>); + }; + (@parse $source:ty; [false true];) => { + $crate::export_data_source!(@provider $crate::__private::provider::WriteOnly<$source>); + }; + (@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::__private::abi::{ + self, ArrowArrayStream, ArrowSchema, SparkDsLibrary, SparkDsSink, + }; + use $crate::__private::bridge::{open_reader, parse_predicates, Bridge}; + + use $crate::__private::provider::Provider as _; + + // 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 __SparkWriterFactory = <<__SparkRegistration + as $crate::__private::provider::Provider>::BatchWrite + as $crate::BatchWrite>::Factory; + type __SparkMessage = <__SparkWriterFactory as $crate::DataWriterFactory>::Message; + + 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: 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 WRITE { Some(create_writer) } else { None }, + close: Some(close), + 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 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"] + /// # 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, + <__SparkSource as $crate::DataSource>::init, + ) + } + } + + 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 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); + } + 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 = 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 { open_reader::<__SparkReaderFactory>(task, state, partition)? }; + unsafe { output.write(stream) }; + 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 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(()) + }) + } + + 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::__private::write::open::<__SparkWriterFactory>(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::__private::write::write_batch::<__SparkMessage>(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::__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::__private::write::abort_task::<__SparkMessage>(writer) + }) + } + + unsafe extern "C" fn close_data_writer(writer: *mut c_void) { + 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 new file mode 100644 index 0000000000000..3ebab03e22a0b --- /dev/null +++ b/sql/native/rust/src/lib.rs @@ -0,0 +1,548 @@ +// 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 data source SDK for the C ABI proposed in Spark PR #59424. +//! +//! 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 std::fmt; + +use arrow_array::{RecordBatch, RecordBatchReader}; +use arrow_schema::SchemaRef; +use serde::de::DeserializeOwned; +use serde::Serialize; + +mod abi; +mod bridge; +mod export; +mod host; +mod predicate; +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 __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 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; +/// 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>; + +/// A native operation of a source instance. +/// +/// These are native capabilities, not Spark's Data Source V2 routing flags or +/// 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, + BatchWrite, + /// 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)] +#[non_exhaustive] +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, +} + +/// Whether a connector limits the rows produced by the whole scan. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum LimitSupport { + /// The scan may return more rows than the limit. A connector may still cap the + /// rows of each partition on its own: Spark applies its limit after reading. + Unsupported, + /// The whole scan returns at most the limit. Spark still applies its limit. + Global, +} + +/// 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. +#[non_exhaustive] +pub struct ScanPlan { + pub factory: F, + pub partitions: Vec, +} + +impl ScanPlan { + pub fn new(factory: F, partitions: Vec) -> Self { + Self { + factory, + partitions, + } + } +} + +/// 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>, +} + +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. +/// +/// 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. +/// +/// New modes may be added. The SDK passes a mode only when the source instance +/// has the matching capability, so a source never receives a mode it did not +/// opt in to; a wildcard arm can return an error. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub enum WriteMode { + Append, + /// Replace existing data at driver commit, never at writer creation. + Truncate, +} + +/// 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 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 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. 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 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. +/// +/// Keep irreversible publication, including removing overwritten data, in commit. +/// Drop releases resources; it must not publish data or perform fallible flushing. +/// +/// # Lifecycle +/// +/// The SDK enforces the following, so implementations need not track their own +/// phase or guard against repeated calls: +/// +/// - `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 `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; + + /// 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 stands for a partition + /// whose commit message did not reach the driver, or cannot be decoded. + /// Executor loss can leave orphaned staged data; plan for its cleanup. + fn abort( + &mut self, + messages: Vec::Message>>, + ) -> 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 context until +/// close. Callbacks can move between threads sequentially; writers must be Send. +pub trait DataWriterFactory: Serialize + DeserializeOwned + Send + 'static { + type Message: Serialize + DeserializeOwned + Send + 'static; + + fn open<'task>( + &self, + context: WriterContext<'task>, + ) -> Result>; +} + +/// One task attempt's writer. The SDK closes it after commit or abort. +/// +/// # 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 `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. +pub trait DataWriter: Send { + type Message; + + /// Takes ownership of an imported Arrow batch, which may be retained after + /// 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 + /// 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<()>; +} diff --git a/sql/native/rust/src/predicate.rs b/sql/native/rust/src/predicate.rs new file mode 100644 index 0000000000000..538678afe2804 --- /dev/null +++ b/sql/native/rust/src/predicate.rs @@ -0,0 +1,372 @@ +// 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. +/// [`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")] +#[non_exhaustive] +pub enum Predicate { + Column { + name: Vec, + }, + Literal { + #[serde(rename = "dataType")] + data_type: String, + value: Value, + }, + Function { + name: String, + children: Vec, + }, + #[serde(other)] + 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 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() { + 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)); + assert_eq!( + children[1].literal_value(), + Some(LiteralValue::Integer(i64::MAX)) + ); + assert_eq!(children[0].literal_value(), None); + } + + #[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"); + } + + #[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 new file mode 100644 index 0000000000000..a97398f0e8cb1 --- /dev/null +++ b/sql/native/rust/src/provider.rs @@ -0,0 +1,135 @@ +// 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::{ + 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 ScanBuilder: ScanBuilder; + type BatchWrite: BatchWrite; + const READ: bool; + const WRITE: bool; + const CAPABILITIES: Capabilities; + + 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); + +impl Provider for ReadOnly { + type Source = S; + type ScanBuilder = S::ScanBuilder; + type BatchWrite = Unavailable; + const READ: bool = true; + const WRITE: bool = false; + const CAPABILITIES: Capabilities = Capabilities::registered(true, false, false); + + fn new_scan(source: &S, schema: SchemaRef) -> Result { + source.new_scan(schema) + } + + fn new_write(_: &S, _: WriteInfo) -> Result> { + Err("batch writes are not registered".into()) + } +} + +impl Provider for WriteOnly { + type Source = S; + type ScanBuilder = Unavailable; + type BatchWrite = S::BatchWrite; + const READ: bool = false; + const WRITE: bool = true; + const CAPABILITIES: Capabilities = Capabilities::registered(false, true, S::SUPPORTS_TRUNCATE); + + fn new_scan(_: &S, _: SchemaRef) -> Result { + Err("batch reads are not registered".into()) + } + + fn new_write(source: &S, info: WriteInfo) -> Result> { + source.new_write(info) + } +} + +impl Provider for ReadWrite { + type Source = S; + type ScanBuilder = ::ScanBuilder; + type BatchWrite = ::BatchWrite; + const READ: bool = true; + const WRITE: bool = true; + const CAPABILITIES: Capabilities = Capabilities::registered(true, true, S::SUPPORTS_TRUNCATE); + + fn new_scan(source: &S, schema: SchemaRef) -> Result { + source.new_scan(schema) + } + + fn new_write(source: &S, info: WriteInfo) -> Result> { + source.new_write(info) + } +} + +// 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 PartitionReaderFactory for Unavailable { + type Partition = (); + + fn open<'task>(&self, _: (), _: ReaderContext<'task>) -> Result> { + Err("batch reads are not registered".into()) + } +} + +impl BatchWrite for Unavailable { + type Factory = Self; + + 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 DataWriterFactory for Unavailable { + type Message = (); + + 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 new file mode 100644 index 0000000000000..e9d46869d4d6c --- /dev/null +++ b/sql/native/rust/src/stream.rs @@ -0,0 +1,721 @@ +// 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, Schema, SchemaRef, TimeUnit}; + +use super::{describe, panic_message}; +use crate::{PartitionReader, Result}; + +struct StreamState<'task> { + 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, +} + +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: 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, + }); + 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| describe(&*e)) + })); + 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)?; + 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) }; + } + 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 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(()) +} + +pub(crate) 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 { + 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())) + } + 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()), + }) +} + +/// 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!( + "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::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}; + + 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: 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) } + } + + 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); + } + + 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(); + 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..61078ab7eef9b --- /dev/null +++ b/sql/native/rust/src/task.rs @@ -0,0 +1,302 @@ +// 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::size_of; +use std::sync::atomic::{AtomicBool, Ordering}; + +use crate::abi::{self, Interrupted, SparkDsTaskContext}; +use crate::Result; + +/// Task identifiers and cooperative cancellation, borrowed until its reader or writer is closed. +/// +/// 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; +/// +/// fn retain_forever(task: TaskContext<'_>) -> TaskContext<'static> { +/// task +/// } +/// ``` +#[derive(Debug, Clone, Copy)] +pub struct TaskContext<'task> { + 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 +// 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> { + /// 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, 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. + pub(crate) unsafe fn from_raw(pointer: *const SparkDsTaskContext) -> Result { + let size = unsafe { abi::struct_size(pointer)? }; + // 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 { + 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, + }, + }) + } + + /// The ID of the stage of the task. + pub fn stage_id(&self) -> i32 { + self.stage_id + } + + /// The attempt number of the stage of the task. + pub fn stage_attempt_number(&self) -> i32 { + self.stage_attempt_number + } + + /// The partition that the task reads or writes. + pub fn partition_id(&self) -> i32 { + self.partition_id + } + + /// The attempt number of the task. + pub fn attempt_number(&self) -> i32 { + self.attempt_number + } + + /// 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, read from Spark each time. + pub fn is_interrupted(&self) -> bool { + 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 + } +} + +#[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(), 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| { + scope.spawn(|| assert!(context.is_interrupted())); + }); + } + + #[test] + fn rejects_a_context_without_every_field_of_version_1() { + let prefix = size_of::() as i64; + let pointer = (&prefix as *const i64).cast::(); + 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()); + 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 new file mode 100644 index 0000000000000..0e9e53e5e565b --- /dev/null +++ b/sql/native/rust/src/write.rs @@ -0,0 +1,354 @@ +// 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, describe, encode, panic_message, provider_call}; +use crate::{ + BatchWrite, DataWriterFactory, LogLevel, PartitionWriter, Result, TaskContext, WritePlan, + WriterContext, +}; + +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 { + coordinator: Option, + factory: Option, + schema: SchemaRef, + serialized: Option>>>, + phase: Phase<()>, +} + +impl DriverWriter { + pub fn new(plan: WritePlan, schema: SchemaRef) -> Self { + Self { + coordinator: Some(plan.coordinator), + factory: Some(plan.factory), + 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)?; + Ok(Arc::new(encode(&WriterState { + factory: self.factory.as_ref().expect("factory is present"), + 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.coordinator + .as_mut() + .expect("coordinator 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(|| { + // 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() + .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("coordinator is present") + .abort(messages) + }); + self.phase = Phase::Aborted(result.clone()); + cached(&result) + } +} + +impl Drop for DriverWriter { + fn drop(&mut self) { + 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>, + 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 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)), + }; + 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..5d474c5c0cf95 --- /dev/null +++ b/sql/native/rust/tests/batch_write.rs @@ -0,0 +1,805 @@ +// 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::__private::abi::{ + ArrowSchema, SparkDsCreateInfo, SparkDsHost, SparkDsLibrary, SparkDsTaskContext, + SparkDsWriteInfo, ABI_VERSION, +}; +use spark_native_source::{ + 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. +static TEST_LOCK: Mutex<()> = Mutex::new(()); +static EVENTS: Mutex> = Mutex::new(Vec::new()); + +#[derive(Clone, Debug, PartialEq)] +enum Event { + WriteCreated(bool), + 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 { + // 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: info.options.get("failure").cloned().unwrap_or_default(), + capabilities: info + .options + .get("capabilities") + .cloned() + .unwrap_or_default(), + }) + } + + fn capabilities(&self, registered: Capabilities) -> Capabilities { + match self.capabilities.as_str() { + "none" => Capabilities::empty(), + "append" => registered.without(SourceCapability::Truncate), + _ => registered, + } + } +} + +impl BatchWriteSource for Source { + type BatchWrite = Driver; + + const SUPPORTS_TRUNCATE: bool = true; + + fn new_write(&self, info: WriteInfo) -> Result> { + if info.schema != schema() { + return Err("unsupported input schema".into()); + } + 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)) + } +} + +struct Driver { + failure: String, +} + +impl BatchWrite for Driver { + type Factory = Factory; + + 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 DataWriterFactory for Factory { + type Message = Message; + + 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: context.task, + failure: self.failure.clone(), + batches: vec![], + })) + } +} + +struct Writer<'task> { + task: TaskContext<'task>, + failure: String, + batches: Vec, +} + +impl DataWriter 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(), + attempt: self.task.task_attempt_id(), + 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); + +extern "C" { + fn spark_ds_init(host: *const SparkDsHost, library: *mut *const SparkDsLibrary) -> c_int; +} + +fn library() -> &'static SparkDsLibrary { + // 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!(!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.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()); + 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::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" + ); + // 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) + }); + 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::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 new file mode 100644 index 0000000000000..65ba31dc0a928 --- /dev/null +++ b/sql/native/rust/tests/capabilities.rs @@ -0,0 +1,290 @@ +// 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, CStr, CString}; +use std::mem::size_of; +use std::ptr; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use arrow_schema::{Schema, SchemaRef}; +use serde::{Deserialize, Serialize}; +use spark_native_source::__private::abi::{ + ArrowSchema, SparkDsCreateInfo, SparkDsHost, SparkDsLibrary, ABI_VERSION, +}; +use spark_native_source::{ + 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 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: info.options.get("capabilities").cloned().unwrap(), + }) + } + + fn capabilities(&self, registered: Capabilities) -> Capabilities { + DECLARATIONS.fetch_add(1, Ordering::SeqCst); + match self.mode.as_str() { + "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}"), + } + } + + fn schema(&self) -> Result { + SCHEMAS.fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(Schema::empty())) + } +} + +impl BatchReadSource for Source { + type ScanBuilder = Builder; + + fn new_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::new(Factory, vec![])) + } +} + +#[derive(Serialize, Deserialize)] +struct Factory; + +impl PartitionReaderFactory for Factory { + type Partition = (); + + fn open<'task>(&self, _: (), _: ReaderContext<'task>) -> Result> { + Err("capability tests do not open partition readers".into()) + } +} + +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, +} + +impl Drop for SourceHandle<'_> { + fn drop(&mut self) { + unsafe { self.library.close.unwrap()(self.pointer) }; + } +} + +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()]; + 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 initialization_logging_and_capabilities_follow_the_source_lifecycle() { + // One lifecycle test owns these counters; no parallel test resets shared state. + let mut table = ptr::null(); + 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, "capability_test", "registered"); + 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), + ("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(); + 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), 5); + } + assert_eq!(SCHEMAS.load(Ordering::SeqCst), 4); + assert_eq!(SCANS.load(Ordering::SeqCst), 2); + 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, "capability_test", "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), 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 new file mode 100644 index 0000000000000..db21656252e08 --- /dev/null +++ b/sql/native/rust/tests/fixtures/spark_datasource.h @@ -0,0 +1,704 @@ +/* + * 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. + * 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 + * 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: + * 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. + * - 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. 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 + * 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, and the callbacks + * of the stream that read initializes, can be called from any thread. The callbacks of the + * stream never run concurrently. + * + * 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. 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 + * 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 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. + * + * 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. 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. + */ + 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 */