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..cca49d0ddbef0 --- /dev/null +++ b/examples/src/main/rust/native-data-source/Cargo.toml @@ -0,0 +1,37 @@ +# 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 + +[dev-dependencies] +jni = { workspace = true, features = ["invocation"] } 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..02c9e4cbd3456 --- /dev/null +++ b/examples/src/main/rust/native-data-source/README.md @@ -0,0 +1,90 @@ + + +# 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 JNI macro. +See the SDK guide for API contracts, ownership rules, and local test instructions. + +The example targets NativeBridge ABI version 1 from [Spark PR #59313][native-source-pr], pinned +to commit `ff7b1ae0150fbcb3b6a7c67bea8accf6aa4ad9bd`. Set `SPARK_HOME` to a Spark distribution +built with that PR before running the Spark commands below. + +## Data and options + +| Column | Spark type | Value | +| --- | --- | --- | +| `id` | `BIGINT`, non-null | Integers in `[start, end)` | +| `label` | `STRING`, non-null | `row-{id}`, such as `row-13` | + +| Option | Default | Meaning | +| --- | --- | --- | +| `start` | `0` | Inclusive start, as a signed 64-bit integer | +| `end` | `100` | Exclusive end; must be at least `start` | +| `numpartitions` | `2` | Maximum number of partitions; must be positive | +| `batchsize` | `1024` | Maximum rows per Arrow batch; must be positive | + +Spark normalizes option names to lower case. Unknown options are rejected. + +The source evaluates `=`, `<=>`, `<`, `<=`, `>`, and `>=` comparisons between `id` and non-null +integer literals, including conjunctions of supported predicates. Other predicates remain in +Spark. A global limit is applied after exact filters and before splitting the range into partitions; +the SDK declines limit pushdown when a residual predicate remains. Column pruning preserves the +requested order and supports zero-column batches for operations such as `COUNT(*)`. + +## Build and run + +Run these commands from the repository root. The workspace, lockfile, and build output belong to +`sql/native/rust`; the example is a member of that workspace. + +```sh +cargo build --manifest-path sql/native/rust/Cargo.toml --workspace --locked + +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 the resulting rows, projection, limits, zero-column counts, and +filters that Spark must evaluate. This Spark smoke script has not yet been run against the +required distribution. The local JVM integration test exercises JNI and Arrow C Stream without +launching Spark; it does not establish end-to-end Spark compatibility. + +[sdk]: ../../../../../sql/native/rust/README.md +[native-source-pr]: https://github.com/apache/spark/pull/59313 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..9f64198df0856 --- /dev/null +++ b/examples/src/main/rust/native-data-source/package.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Package the compiled range connector for the NativeBridge ABI in Spark PR #59313.""" + +import argparse +import json +from pathlib import Path +import platform +import zipfile + + +def current_platform(): + os_name = {"Linux": "linux", "Darwin": "osx", "Windows": "windows"}.get(platform.system()) + arch = { + "x86_64": "x86_64", + "amd64": "x86_64", + "aarch64": "aarch64", + "arm64": "aarch64", + }.get(platform.machine().lower()) + return f"{os_name}-{arch}" if os_name and arch else None + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--library", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--platform", default=current_platform(), help="Override for cross builds") + args = parser.parse_args() + if not args.library.is_file(): + parser.error(f"Library does not exist: {args.library}") + if args.output.suffix != ".sparkpkg": + parser.error("The output file must have the .sparkpkg extension") + if args.platform not in { + f"{os_name}-{arch}" + for os_name in ("linux", "osx", "windows") + for arch in ("x86_64", "aarch64") + }: + parser.error("Specify a supported --platform, such as linux-x86_64") + if args.library.resolve() == args.output.resolve(): + parser.error("The output must differ from the input library") + + library_path = f"{args.platform}/{args.library.name}" + manifest = { + "abiVersion": 1, + "dataSources": ["rust_sdk_range"], + "libraries": {args.platform: library_path}, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(args.output, "w", compression=zipfile.ZIP_DEFLATED) as package: + package.writestr("spark-native-datasource.json", json.dumps(manifest, indent=2) + "\n") + package.write(args.library, library_path) + print(args.output.resolve()) + + +if __name__ == "__main__": + main() diff --git a/examples/src/main/rust/native-data-source/smoke.py b/examples/src/main/rust/native-data-source/smoke.py new file mode 100644 index 0000000000000..db4c81ca2797d --- /dev/null +++ b/examples/src/main/rust/native-data-source/smoke.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. + +"""Run with spark-submit from a Spark distribution that includes PR #59313.""" + +import argparse +from pathlib import Path + +from pyspark.sql import SparkSession, functions as F + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--package", type=Path, required=True) + args = parser.parse_args() + package = args.package.resolve(strict=True) + + spark = SparkSession.builder.appName("Rust native data source smoke").getOrCreate() + try: + spark.addArtifact(str(package)) + df = ( + spark.read.format("rust_sdk_range") + .option("start", 0) + .option("end", 31) + .option("numpartitions", 4) + .option("batchsize", 3) + .load() + ) + assert [tuple(row) for row in df.orderBy("id").collect()] == [ + (i, f"row-{i}") for i in range(31) + ] + + 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 + + # The connector does not evaluate these expressions; Spark must retain the filters. + arithmetic = df.where((F.col("id") % 3) == 1).select("id").orderBy("id").collect() + assert [row.id for row in arithmetic] == list(range(1, 31, 3)) + labels = df.where(F.col("label").endswith("3")).select("id").orderBy("id").collect() + assert [row.id for row in labels] == [3, 13, 23] + residual_limit = df.where(F.col("label") == "row-21").limit(1).select("id").collect() + assert [row.id for row in residual_limit] == [21] + print("Rust native data source Spark smoke checks passed") + finally: + spark.stop() + + +if __name__ == "__main__": + main() diff --git a/examples/src/main/rust/native-data-source/src/lib.rs b/examples/src/main/rust/native-data-source/src/lib.rs new file mode 100644 index 0000000000000..11a52f78b1082 --- /dev/null +++ b/examples/src/main/rust/native-data-source/src/lib.rs @@ -0,0 +1,587 @@ +// 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 JNI 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::{DataType, Field, Schema, SchemaRef}; +use serde::{Deserialize, Serialize}; +use spark_native_source::{ + BatchReader, BatchSource, FilterSupport, LimitSupport, Options, Predicate, 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, + state: RangeState, +} + +/// Reader configuration travels to each executor; it contains no live objects. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RangeState { + columns: Vec, + batch_size: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RangePartition { + start: i64, + end: i64, +} + +impl BatchSource for RangeSource { + type Scan = RangeScan; + + fn new(options: Options) -> Result { + for key in options.keys() { + if !matches!( + key.as_str(), + "start" | "end" | "numpartitions" | "batchsize" + ) { + return Err(format!("Unknown range option: {key}").into()); + } + } + let start = option(&options, "start", 0_i64)?; + let end = option(&options, "end", 100_i64)?; + let num_partitions = option(&options, "numpartitions", 2_usize)?; + let batch_size = option(&options, "batchsize", 1024_usize)?; + if end < start { + return Err("end must be greater than or equal to start".into()); + } + if num_partitions == 0 || batch_size == 0 { + return Err("numPartitions and batchSize must be positive".into()); + } + Ok(Self { + start, + end, + num_partitions, + batch_size, + }) + } + + fn schema(&self) -> Result { + output_schema(&["id".into(), "label".into()]) + } + + fn scan(&self, schema: SchemaRef) -> Result { + let columns: Vec = schema.fields().iter().map(|f| f.name().clone()).collect(); + let expected = output_schema(&columns)?; + for (field, expected) in schema.fields().iter().zip(expected.fields()) { + if field.data_type() != expected.data_type() { + return Err(format!( + "Unsupported type for {}: {}", + field.name(), + field.data_type() + ) + .into()); + } + } + Ok(RangeScan { + lower: i128::from(self.start), + upper: i128::from(self.end), + num_partitions: self.num_partitions, + limit: None, + state: RangeState { + columns, + batch_size: self.batch_size, + }, + }) + } + + fn open(state: RangeState, partition: RangePartition) -> Result { + if state.batch_size == 0 || partition.start > partition.end { + return Err("Invalid serialized range plan".into()); + } + let schema = output_schema(&state.columns)?; + let batch_schema = Arc::clone(&schema); + let mut current = partition.start; + let batches = std::iter::from_fn(move || { + if current == partition.end { + return None; + } + let next = (i128::from(current) + state.batch_size as i128) + .min(i128::from(partition.end)) as i64; + let num_rows = (i128::from(next) - i128::from(current)) as usize; + let columns: Vec = state + .columns + .iter() + .map(|column| match column.as_str() { + "id" => Arc::new(Int64Array::from_iter_values(current..next)) as ArrayRef, + "label" => Arc::new(StringArray::from_iter_values( + (current..next).map(|id| format!("row-{id}")), + )) as ArrayRef, + _ => unreachable!("column names were validated before reading"), + }) + .collect(); + current = next; + // Explicit row counts also preserve COUNT(*) with a zero-column projection. + Some(RecordBatch::try_new_with_options( + Arc::clone(&batch_schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(num_rows)), + )) + }); + Ok(Box::new(RecordBatchIterator::new(batches, schema))) + } +} + +impl ScanBuilder for RangeScan { + type State = RangeState; + type Partition = RangePartition; + + 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.state.columns.contains(column)) + { + return Err("Projection contains a column outside the read schema".into()); + } + output_schema(columns)?; + self.state.columns = columns.to_vec(); + Ok(true) + } + + fn build(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 { + state: self.state, + partitions, + }) + } +} + +fn option(options: &Options, name: &str, default: T) -> Result +where + T: std::str::FromStr, + T::Err: std::fmt::Display, +{ + options.get(name).map_or(Ok(default), |value| { + value + .parse() + .map_err(|error| format!("Invalid {name}: {error}").into()) + }) +} + +fn output_schema(columns: &[String]) -> Result { + let mut seen = BTreeSet::new(); + let fields: Result> = columns + .iter() + .map(|name| { + if !seen.insert(name) { + return Err(format!("Duplicate column: {name}").into()); + } + let data_type = match name.as_str() { + "id" => DataType::Int64, + "label" => DataType::Utf8, + _ => return Err(format!("Unknown column: {name}").into()), + }; + Ok(Field::new(name, data_type, false)) + }) + .collect(); + Ok(Arc::new(Schema::new(fields?))) +} + +/// Closed-open bounds use i128 so strict comparisons never overflow at i64 extrema. +fn predicate_bounds(predicate: &Predicate) -> Option<(i128, i128)> { + let Predicate::Function { name, children } = predicate else { + return None; + }; + if children.len() != 2 { + return None; + } + if name == "AND" { + // Do not mutate the scan when only one side of an AND is supported. + let left = predicate_bounds(&children[0])?; + let right = predicate_bounds(&children[1])?; + return Some((left.0.max(right.0), left.1.min(right.1))); + } + let (value, reversed) = if is_id(&children[0]) { + (integer_literal(&children[1])?, false) + } else if is_id(&children[1]) { + (integer_literal(&children[0])?, true) + } else { + return None; + }; + let name = match (name.as_str(), reversed) { + (">", true) => "<", + (">=", true) => "<=", + ("<", true) => ">", + ("<=", true) => ">=", + (name, _) => name, + }; + let min = i128::from(i64::MIN); + let max = i128::from(i64::MAX) + 1; + let value = i128::from(value); + match name { + "=" | "<=>" => Some((value, value + 1)), + ">" => Some((value + 1, max)), + ">=" => Some((value, max)), + "<" => Some((min, value)), + "<=" => Some((min, value + 1)), + _ => None, + } +} + +fn is_id(predicate: &Predicate) -> bool { + matches!(predicate, Predicate::Column { name } if name.len() == 1 && name[0] == "id") +} + +fn integer_literal(predicate: &Predicate) -> Option { + let Predicate::Literal { data_type, value } = predicate else { + return None; + }; + let value = value.as_i64()?; + match data_type.as_str() { + "tinyint" if i8::try_from(value).is_ok() => Some(value), + "smallint" if i16::try_from(value).is_ok() => Some(value), + "int" if i32::try_from(value).is_ok() => Some(value), + "bigint" => Some(value), + _ => None, + } +} + +spark_native_source::export_data_source!(RangeSource); + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn source(values: &[(&str, &str)]) -> RangeSource { + RangeSource::new( + values + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(), + ) + .unwrap() + } + + fn scan(source: &RangeSource) -> RangeScan { + source.scan(source.schema().unwrap()).unwrap() + } + + fn comparison(name: &str, value: i64, reversed: bool) -> Predicate { + let mut children = vec![ + Predicate::Column { + name: vec!["id".into()], + }, + Predicate::Literal { + data_type: "bigint".into(), + value: json!(value), + }, + ]; + if reversed { + children.reverse(); + } + Predicate::Function { + name: name.into(), + children, + } + } + + fn read(plan: ScanPlan) -> Vec { + plan.partitions + .into_iter() + .flat_map(|partition| { + RangeSource::open(plan.state.clone(), partition) + .unwrap() + .map(|batch| batch.unwrap()) + .collect::>() + }) + .collect() + } + + fn ids(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|batch| { + let index = batch.schema().index_of("id").unwrap(); + batch + .column(index) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect() + } + + #[test] + fn filters_before_global_limit_and_preserves_projected_order() { + let source = source(&[ + ("start", "10"), + ("end", "25"), + ("numpartitions", "3"), + ("batchsize", "2"), + ]); + let mut scan = scan(&source); + assert_eq!( + scan.push_filters(&[comparison(">=", 13, false), comparison("<", 22, false),]) + .unwrap(), + vec![FilterSupport::Exact; 2] + ); + assert_eq!(scan.push_limit(5).unwrap(), LimitSupport::Global); + assert!(scan.prune_columns(&["label".into(), "id".into()]).unwrap()); + let plan = scan.build().unwrap(); + assert_eq!(plan.partitions.len(), 3); + let batches = read(plan); + assert_eq!(ids(&batches), (13..18).collect::>()); + for batch in &batches { + assert!(batch.num_rows() <= 2); + assert_eq!(batch.schema().field(0).name(), "label"); + let labels = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for (row, id) in ids(std::slice::from_ref(batch)).iter().enumerate() { + assert_eq!(labels.value(row), format!("row-{id}")); + } + } + } + + #[test] + fn comparisons_handle_reversed_operands_and_int64_boundaries() { + for (start, end) in [(-2, 3), (i64::MIN, i64::MIN + 3), (i64::MAX - 3, i64::MAX)] { + let source = source(&[("start", &start.to_string()), ("end", &end.to_string())]); + for literal in [i64::MIN, -1, 0, 1, i64::MAX] { + for op in ["=", "<=>", ">", ">=", "<", "<="] { + for reversed in [false, true] { + let mut scan = scan(&source); + assert_eq!( + scan.push_filters(&[comparison(op, literal, reversed)]) + .unwrap(), + vec![FilterSupport::Exact] + ); + let expected: Vec<_> = (start..end) + .filter(|id| { + let (left, right) = if reversed { + (literal, *id) + } else { + (*id, literal) + }; + match op { + "=" | "<=>" => left == right, + ">" => left > right, + ">=" => left >= right, + "<" => left < right, + "<=" => left <= right, + _ => unreachable!(), + } + }) + .collect(); + assert_eq!( + ids(&read(scan.build().unwrap())), + expected, + "range={start}..{end}, op={op}, value={literal}, reverse={reversed}" + ); + } + } + } + } + } + + #[test] + fn unsupported_filters_remain_residual_without_mutating_the_scan() { + let source = source(&[("end", "10")]); + let mut scan = scan(&source); + let unsupported_and = Predicate::Function { + name: "AND".into(), + children: vec![comparison(">", 5, false), Predicate::Unknown], + }; + let null: Predicate = serde_json::from_value(json!({ + "type": "function", "name": "<=>", "children": [ + {"type": "column", "name": ["id"]}, + {"type": "literal", "dataType": "bigint", "value": null} + ] + })) + .unwrap(); + let double: Predicate = serde_json::from_value(json!({ + "type": "function", "name": "=", "children": [ + {"type": "column", "name": ["id"]}, + {"type": "literal", "dataType": "double", "value": 1.5} + ] + })) + .unwrap(); + assert_eq!( + scan.push_filters(&[ + unsupported_and, + null, + double, + comparison("FUTURE_COMPARISON", 0, false), + ]) + .unwrap(), + vec![FilterSupport::Unsupported; 4] + ); + assert_eq!( + ids(&read(scan.build().unwrap())), + (0..10).collect::>() + ); + } + + #[test] + fn conjunctions_and_contradictory_predicates_form_exact_intersections() { + let source = source(&[("end", "10")]); + let mut first = scan(&source); + let and = Predicate::Function { + name: "AND".into(), + children: vec![comparison(">=", 3, false), comparison("<=", 6, false)], + }; + assert_eq!( + first.push_filters(&[and]).unwrap(), + vec![FilterSupport::Exact] + ); + assert_eq!(ids(&read(first.build().unwrap())), vec![3, 4, 5, 6]); + let mut empty = scan(&source); + empty + .push_filters(&[comparison(">", 5, false), comparison("<", 3, false)]) + .unwrap(); + assert!(empty.build().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.build().unwrap(); + let state = serde_json::from_slice(&serde_json::to_vec(&plan.state).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 { state, 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.build().unwrap().partitions.is_empty()); + let mut limited = scan(&provider); + limited.push_limit(2).unwrap(); + limited.push_limit(5).unwrap(); + let plan = limited.build().unwrap(); + assert_eq!(plan.partitions.len(), 2); + assert_eq!(ids(&read(plan)), vec![0, 1]); + let empty = source(&[("start", "5"), ("end", "5")]); + assert!(scan(&empty).build().unwrap().partitions.is_empty()); + } + + #[test] + fn rejects_invalid_options_and_incompatible_user_schemas() { + for (key, value) in [ + ("batchsize", "0"), + ("batchsize", "-1"), + ("numpartitions", "0"), + ("end", "-1"), + ("start", "not-an-integer"), + ("unknown", "1"), + ] { + assert!(RangeSource::new([(key.into(), value.into())].into()).is_err()); + } + let source = source(&[]); + let invalid = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + assert!(source.scan(invalid).is_err()); + let labels = output_schema(&["label".into()]).unwrap(); + let mut scan = source.scan(labels).unwrap(); + assert!(scan.prune_columns(&["id".into()]).is_err()); + assert!(output_schema(&["id".into(), "id".into()]).is_err()); + } + + #[test] + fn opens_a_huge_partition_lazily_without_materializing_it() { + let source = source(&[ + ("start", "-9223372036854775808"), + ("end", "9223372036854775807"), + ("numpartitions", "1"), + ("batchsize", "2"), + ]); + let plan = scan(&source).build().unwrap(); + let mut reader = RangeSource::open(plan.state, plan.partitions[0].clone()).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/java/org/apache/spark/sql/datasource/NativeBridge.java b/examples/src/main/rust/native-data-source/tests/java/org/apache/spark/sql/datasource/NativeBridge.java new file mode 100644 index 0000000000000..cd4914e0e5ca0 --- /dev/null +++ b/examples/src/main/rust/native-data-source/tests/java/org/apache/spark/sql/datasource/NativeBridge.java @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.datasource; + +/** Minimal test fixture with the batch-read signatures of NativeBridge ABI version 1. */ +public final class NativeBridge { + private NativeBridge() {} + + public static void load(String path) { + System.load(path); + } + + public static native int abiVersion(); + + public static native long createDataSource(String name, String[] keys, String[] values); + + public static native void schema(long source, long schemaAddress); + + public static native long createReader(long source, long schemaAddress); + + public static native void closeDataSource(long source); + + public static native boolean[] pushPredicates(long reader, String[] predicates); + + public static native boolean pushLimit(long reader, int limit); + + public static native boolean pruneColumns(long reader, String[] columns); + + public static native byte[][] partitions(long reader); + + public static native byte[] serializeReader(long reader); + + public static native void closeReader(long reader); + + public static native void read(byte[] state, byte[] partition, long streamAddress); +} diff --git a/examples/src/main/rust/native-data-source/tests/jni_smoke.rs b/examples/src/main/rust/native-data-source/tests/jni_smoke.rs new file mode 100644 index 0000000000000..b7f5684d7eb22 --- /dev/null +++ b/examples/src/main/rust/native-data-source/tests/jni_smoke.rs @@ -0,0 +1,407 @@ +// 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::error::Error; +use std::fs; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::time::{SystemTime, UNIX_EPOCH}; + +use arrow_array::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream}; +use arrow_array::{Int64Array, RecordBatch, StringArray}; +use arrow_schema::ffi::FFI_ArrowSchema; +use arrow_schema::{DataType, Schema}; +use jni::objects::{JBooleanArray, JByteArray, JObject, JObjectArray, JString, JValue}; +use jni::sys::{jlong, JNI_TRUE}; +use jni::{InitArgsBuilder, JNIEnv, JNIVersion, JavaVM}; + +type TestResult = Result>; +const BRIDGE: &str = "org/apache/spark/sql/datasource/NativeBridge"; + +struct Classes(PathBuf); + +impl Drop for Classes { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } +} + +struct Plan { + state: Vec, + partitions: Vec>, +} + +fn string_array<'local>( + env: &mut JNIEnv<'local>, + values: &[&str], +) -> jni::errors::Result> { + let array = env.new_object_array(values.len() as i32, "java/lang/String", JObject::null())?; + for (i, value) in values.iter().enumerate() { + let value = env.new_string(value)?; + env.set_object_array_element(&array, i as i32, value)?; + } + Ok(array) +} + +fn source(env: &mut JNIEnv<'_>, options: &[(&str, &str)]) -> jni::errors::Result { + let name = env.new_string("rust_sdk_range")?; + let keys = string_array( + env, + &options.iter().map(|(key, _)| *key).collect::>(), + )?; + let values = string_array( + env, + &options.iter().map(|(_, value)| *value).collect::>(), + )?; + env.call_static_method( + BRIDGE, + "createDataSource", + "(Ljava/lang/String;[Ljava/lang/String;[Ljava/lang/String;)J", + &[ + JValue::Object(name.as_ref()), + JValue::Object(keys.as_ref()), + JValue::Object(values.as_ref()), + ], + )? + .j() +} + +fn plan( + env: &mut JNIEnv<'_>, + options: &[(&str, &str)], + predicates: &[&str], + accepted: &[bool], + limit: Option<(i32, bool)>, + projection: Option<&[&str]>, +) -> TestResult { + let source = source(env, options)?; + assert_ne!(source, 0); + let mut schema = FFI_ArrowSchema::empty(); + let schema_address = &mut schema as *mut FFI_ArrowSchema as jlong; + env.call_static_method( + BRIDGE, + "schema", + "(JJ)V", + &[JValue::Long(source), JValue::Long(schema_address)], + )?; + let imported = Schema::try_from(&schema)?; + assert_eq!(imported.field(0).name(), "id"); + assert_eq!(imported.field(0).data_type(), &DataType::Int64); + assert_eq!(imported.field(1).name(), "label"); + let reader = env + .call_static_method( + BRIDGE, + "createReader", + "(JJ)J", + &[JValue::Long(source), JValue::Long(schema_address)], + )? + .j()?; + assert_ne!(reader, 0); + drop(schema); + // Reader state and imported schema must not borrow the source handle. + env.call_static_method(BRIDGE, "closeDataSource", "(J)V", &[JValue::Long(source)])?; + + let predicates = string_array(env, predicates)?; + let result = env + .call_static_method( + BRIDGE, + "pushPredicates", + "(J[Ljava/lang/String;)[Z", + &[JValue::Long(reader), JValue::Object(predicates.as_ref())], + )? + .l()?; + let result = JBooleanArray::from(result); + assert_eq!(env.get_array_length(&result)? as usize, accepted.len()); + let mut booleans = vec![0; accepted.len()]; + env.get_boolean_array_region(&result, 0, &mut booleans)?; + assert_eq!( + booleans + .iter() + .map(|value| *value == JNI_TRUE) + .collect::>(), + accepted + ); + if let Some((limit, complete)) = limit { + assert_eq!( + env.call_static_method( + BRIDGE, + "pushLimit", + "(JI)Z", + &[JValue::Long(reader), JValue::Int(limit)], + )? + .z()?, + complete + ); + } + if let Some(projection) = projection { + let columns = string_array(env, projection)?; + assert!(env + .call_static_method( + BRIDGE, + "pruneColumns", + "(J[Ljava/lang/String;)Z", + &[JValue::Long(reader), JValue::Object(columns.as_ref())], + )? + .z()?); + } + let partitions = JObjectArray::from( + env.call_static_method(BRIDGE, "partitions", "(J)[[B", &[JValue::Long(reader)])? + .l()?, + ); + let mut partition_bytes = Vec::new(); + for i in 0..env.get_array_length(&partitions)? { + let partition = JByteArray::from(env.get_object_array_element(&partitions, i)?); + partition_bytes.push(env.convert_byte_array(&partition)?); + } + let state = JByteArray::from( + env.call_static_method(BRIDGE, "serializeReader", "(J)[B", &[JValue::Long(reader)])? + .l()?, + ); + let state = env.convert_byte_array(&state)?; + // Executor reads below only have the serialized bytes, just as on another JVM. + env.call_static_method(BRIDGE, "closeReader", "(J)V", &[JValue::Long(reader)])?; + Ok(Plan { + state, + partitions: partition_bytes, + }) +} + +fn open_stream( + env: &mut JNIEnv<'_>, + plan: &Plan, + partition: &[u8], +) -> TestResult { + let state = env.byte_array_from_slice(&plan.state)?; + let partition = env.byte_array_from_slice(partition)?; + let mut stream = FFI_ArrowArrayStream::empty(); + env.call_static_method( + BRIDGE, + "read", + "([B[BJ)V", + &[ + JValue::Object(state.as_ref()), + JValue::Object(partition.as_ref()), + JValue::Long(&mut stream as *mut FFI_ArrowArrayStream as jlong), + ], + )?; + Ok(ArrowArrayStreamReader::try_new(stream)?) +} + +fn exercise_bridge(env: &mut JNIEnv<'_>, library: &Path) -> TestResult<()> { + let library = env.new_string(library.to_str().ok_or("Non-UTF-8 library path")?)?; + env.call_static_method( + BRIDGE, + "load", + "(Ljava/lang/String;)V", + &[JValue::Object(library.as_ref())], + )?; + assert_eq!( + env.call_static_method(BRIDGE, "abiVersion", "()I", &[])? + .i()?, + 1 + ); + + let options = [ + ("start", "10"), + ("end", "25"), + ("numpartitions", "3"), + ("batchsize", "2"), + ]; + let lower = concat!( + r#"{"type":"function","name":">=","children":["#, + r#"{"type":"column","name":["id"]},"#, + r#"{"type":"literal","dataType":"bigint","value":13}]}"# + ); + let upper = concat!( + r#"{"type":"function","name":"<","children":["#, + r#"{"type":"column","name":["id"]},"#, + r#"{"type":"literal","dataType":"bigint","value":22}]}"# + ); + let label = concat!( + r#"{"type":"function","name":"=","children":["#, + r#"{"type":"column","name":["label"]},"#, + r#"{"type":"literal","dataType":"string","value":"row-21"}]}"# + ); + let filtered = plan( + env, + &options, + &[lower, upper], + &[true, true], + Some((5, true)), + Some(&["label", "id"]), + )?; + let mut batches: Vec = Vec::new(); + for partition in &filtered.partitions { + let reader = open_stream(env, &filtered, partition)?; + // Keep all batches alive across later get_next calls and stream destruction. + batches.extend(reader.collect::, _>>()?); + } + assert!(batches.len() >= 3); + let mut ids = Vec::new(); + for batch in &batches { + assert!(batch.num_rows() <= 2); + assert_eq!(batch.schema().field(0).name(), "label"); + assert_eq!(batch.schema().field(1).name(), "id"); + let labels = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + let id = values.value(i); + assert_eq!(labels.value(i), format!("row-{id}")); + ids.push(id); + } + } + ids.sort_unstable(); + assert_eq!(ids, vec![13, 14, 15, 16, 17]); + + // Spark must evaluate the label predicate before LIMIT; keep all candidate rows. + let residual = plan( + env, + &options, + &[lower, upper, label], + &[true, true, false], + Some((1, false)), + Some(&["id"]), + )?; + let mut candidates = Vec::new(); + for partition in &residual.partitions { + for batch in open_stream(env, &residual, partition)? { + let batch = batch?; + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + candidates.extend_from_slice(ids.values()); + } + } + candidates.sort_unstable(); + assert_eq!(candidates, (13..22).collect::>()); + + let unconsumed = plan( + env, + &[("end", "20"), ("numpartitions", "1"), ("batchsize", "2")], + &[], + &[], + None, + None, + )?; + let mut stream = open_stream(env, &unconsumed, &unconsumed.partitions[0])?; + let first = stream.next().ok_or("Missing first batch")??; + drop(stream); + let first_ids = first + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(first_ids.values().as_ref(), &[0, 1]); + + let empty_projection = plan( + env, + &[("end", "7"), ("numpartitions", "2"), ("batchsize", "2")], + &[], + &[], + None, + Some(&[]), + )?; + let mut count = 0; + for partition in &empty_projection.partitions { + for batch in open_stream(env, &empty_projection, partition)? { + let batch = batch?; + assert_eq!(batch.num_columns(), 0); + count += batch.num_rows(); + } + } + assert_eq!(count, 7); + + let error = source(env, &[("batchsize", "0")]).expect_err("Expected a Java exception"); + assert!(matches!(error, jni::errors::Error::JavaException)); + let exception = env.exception_occurred()?; + env.exception_clear()?; + assert!(env.is_instance_of(&exception, "java/lang/RuntimeException")?); + let message = JString::from( + env.call_method(exception, "toString", "()Ljava/lang/String;", &[])? + .l()?, + ); + let message: String = env.get_string(&message)?.into(); + assert!(message.contains("batchSize must be positive"), "{message}"); + Ok(()) +} + +#[test] +fn jni_arrow_round_trip() -> TestResult<()> { + let unique = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos(); + let classes = Classes( + std::env::temp_dir().join(format!("spark-native-jni-{}-{unique}", std::process::id())), + ); + fs::create_dir(&classes.0)?; + let javac = std::env::var_os("JAVA_HOME") + .map(|home| { + PathBuf::from(home) + .join("bin") + .join(format!("javac{}", std::env::consts::EXE_SUFFIX)) + }) + .unwrap_or_else(|| PathBuf::from("javac")); + let fixture = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("tests/java/org/apache/spark/sql/datasource/NativeBridge.java"); + let compiled = Command::new(javac) + .arg("-d") + .arg(&classes.0) + .arg(fixture) + .output()?; + assert!( + compiled.status.success(), + "{}", + String::from_utf8_lossy(&compiled.stderr) + ); + + let executable = std::env::current_exe()?; + let profile_dir = executable + .parent() + .and_then(Path::parent) + .ok_or("Missing profile directory")?; + let library = profile_dir.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 class_path = format!("-Djava.class.path={}", classes.0.display()); + let args = InitArgsBuilder::new() + .version(JNIVersion::V8) + .option(&class_path) + .option("-Xcheck:jni") + .build()?; + let jvm = JavaVM::new(args)?; + let mut env = jvm.attach_current_thread()?; + env.ensure_local_capacity(256)?; + let result = exercise_bridge(&mut env, &library); + if env.exception_check()? { + env.exception_describe()?; + env.exception_clear()?; + } + result +} diff --git a/sql/native/rust/Cargo.lock b/sql/native/rust/Cargo.lock new file mode 100644 index 0000000000000..6c6adb6c2c86a --- /dev/null +++ b/sql/native/rust/Cargo.lock @@ -0,0 +1,896 @@ +# 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. + +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +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", +] + +[[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 = "cesu8" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c" + +[[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 = "combine" +version = "4.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e" +dependencies = [ + "bytes", + "memchr", +] + +[[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 = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[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 = "java-locator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09c46c1fe465c59b1474e665e85e1256c3893dd00927b8d55f63b09044c1e64f" +dependencies = [ + "glob", +] + +[[package]] +name = "jni" +version = "0.21.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a87aa2bb7d2af34197c04845522473242e1aa17c12f4935d5856491a7fb8c97" +dependencies = [ + "cesu8", + "cfg-if", + "combine", + "java-locator", + "jni-sys 0.3.1", + "libloading", + "log", + "thiserror", + "walkdir", + "windows-sys 0.45.0", +] + +[[package]] +name = "jni-sys" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41a652e1f9b6e0275df1f15b32661cf0d4b78d4d87ddec5e0c3c20f097433258" +dependencies = [ + "jni-sys 0.4.1", +] + +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn 2.0.119", +] + +[[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 = "libloading" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f" +dependencies = [ + "cfg-if", + "winapi", +] + +[[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 = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[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", + "jni", + "serde", + "serde_json", + "spark-native-source", +] + +[[package]] +name = "spark-native-source" +version = "0.1.0" +dependencies = [ + "arrow-array", + "arrow-schema", + "jni", + "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 = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[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 = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[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 = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[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 = "windows-sys" +version = "0.45.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" + +[[package]] +name = "windows_i686_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" + +[[package]] +name = "windows_i686_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" + +[[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..967ab4bee4900 --- /dev/null +++ b/sql/native/rust/Cargo.toml @@ -0,0 +1,51 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +[package] +name = "spark-native-source" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +description = "Experimental Rust batch data source SDK for Apache Spark PR #59313" + +[workspace] +members = ["../../../examples/src/main/rust/native-data-source"] +default-members = ["."] +resolver = "2" + +[workspace.package] +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +publish = false + +[workspace.dependencies] +arrow-array = { version = "58", features = ["ffi"] } +arrow-schema = { version = "58", features = ["ffi"] } +jni = "0.21" +serde = { version = "1", features = ["derive"] } +serde_json = "1" + +[dependencies] +arrow-array.workspace = true +arrow-schema.workspace = true +jni.workspace = true +serde.workspace = true +serde_json.workspace = true + +[profile.release] +# Provider panics are translated into Java / 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..4985e0983edf7 --- /dev/null +++ b/sql/native/rust/README.md @@ -0,0 +1,166 @@ + + +# Rust native data source SDK prototype + +This crate demonstrates a Rust API for **batch reads** over +the JNI ABI proposed in [Spark PR #59313](https://github.com/apache/spark/pull/59313), +at commit `ff7b1ae0150fbcb3b6a7c67bea8accf6aa4ad9bd`. It does not include that PR's +Spark implementation. Running the Spark example requires a Spark distribution +with that native data source framework. Rust and JNI tests run independently of +a Spark build. + +This is an API prototype, not a published crate or a compatibility commitment. +It is not part of the Maven/SBT build and does not add a Rust requirement to +building Spark. + +The `spark-native-source` SDK crate and Cargo workspace live in this directory. +The [Range connector][range] is a separate workspace member under +`examples/src/main/rust/native-data-source`, with its JNI integration test, +packaging helper, and Spark smoke script. The example depends on the SDK through +a Cargo path dependency. Both crates share the lockfile and `target` directory +here. By default, Cargo commands in this directory select only the SDK; use +`--workspace` to include the example. + +There are no write or Structured Streaming exports. + +## Connector API + +A connector implements two traits from [`src/lib.rs`](src/lib.rs): + +* `BatchSource`: parse options, return the schema, create a driver-side scan + builder, and open an Arrow reader on an executor. +* `ScanBuilder`: negotiate filters, limit and output columns, then build a plan + containing typed state and typed partitions. + +The state and partition types derive `serde::Serialize` and `serde::Deserialize`. +The SDK encodes and restores them. Connections, pointers and reader objects stay +inside the process that created them. Include snapshot identifiers in the plan +when reads must be repeatable across task retries. + +Register one provider in the connector's `cdylib` crate: + +```rust +spark_native_source::export_data_source!(RangeSource); +``` + +The macro exports the batch-read JNI functions of +`org.apache.spark.sql.datasource.NativeBridge`. It handles Java strings and +arrays, native handles, error conversion, plan serialization, and Arrow C Stream +export. The connector does not use `JNIEnv`, raw pointers, or `unsafe` code. Rust +traits remain inside the plugin and are not used as a cross-library Rust ABI. + +`BatchSource::open` returns `Box`. Existing Arrow +readers can be used when their types and layouts are supported by the pinned +Spark ABI. That implementation requires zero array offsets at every nesting +level and does not support dictionary encoding or every Arrow type. Connectors +must normalize sliced arrays when necessary; the SDK preserves their layout. +Return batches lazily so a partition does not need to fit in memory. The +[Range implementation][range-source] demonstrates +multiple bounded batches, disjoint splits, and an empty output projection that +preserves row counts. + +## Pushdown contracts + +Filter results correspond to the input predicates in order: + +| Result | Connector responsibility | Spark behavior | +| --- | --- | --- | +| `Exact` | Evaluate completely with Spark SQL semantics | May remove the filter | +| `PruningOnly` | Only skip data that cannot match | Keep the residual filter | +| `Unsupported` | No evaluation | Keep the residual filter | + +Unknown expression names and unsupported data types must remain residual. The +SDK preserves literal JSON values, including exact 64-bit integers, instead of +converting every number to floating point. It does not implement a generic +Spark expression evaluator. + +Limit pushdown distinguishes a per-partition cap from a cap on the whole scan: + +| Result | JNI `pushLimit` return | Meaning | +| --- | --- | --- | +| `Unsupported` | `false` | No limit applied | +| `PerPartition` | `false` | A local cap; the total may exceed the limit | +| `Global` | `true` | The entire scan returns at most the limit | + +The boolean mapping follows #59313's existing contract. Spark retains its final +limit in either case. Spark's optimizer decides whether a query allows limit +pushdown; the SDK additionally declines it if any filter it received was not +`Exact`. Capping rows before a residual filter is unsafe and cannot be repaired +by retaining the final Spark limit. + +Projection returns `true` only when the reader will produce exactly the +requested top-level columns, in order. Keep columns needed internally for exact +filters even if they are absent from the output projection. An empty projection +must still produce the correct number of rows. + +The Range example accepts integer comparisons on `id`, including reversed +operands and conjunctions. It intersects the range with the exact predicates, +applies a global limit, and then creates disjoint partitions. Other predicates +remain with Spark. Pushdown methods have conservative defaults in the SDK, so +connectors can add these capabilities incrementally. + +## Build and test without Spark + +Use a Rust toolchain compatible with Arrow 58 and a JDK (17 or newer). Set +`JAVA_HOME` if the installed JDK is not otherwise discoverable. 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 JNI integration test can load the actual connector +dynamic library. The test compiles a small Java fixture containing the batch +ABI declarations, starts a real JVM, invokes the JNI exports, and reads the +resulting Arrow C Stream. This checks the binary boundary independently of +Spark; it does not substitute for testing Spark's planner and task lifecycle. + +The crate unit tests cover filter negotiation, integer boundaries, global +limits across partitions, projections, and serialized plans. The SDK also +tests Arrow stream and JNI lifecycle behavior. + +## Package and run with Spark + +See the [Range example guide][range] for packaging and `spark-submit` commands. +The example's `package.py` creates a `.sparkpkg` artifact for one platform, and +`smoke.py` checks reading and pushdown against a Spark distribution with #59313. +The compiled example library is under this workspace's `target` directory. + +## Scope and remaining work + +The JNI ABI and this demo both run trusted native code inside the driver and +executors. Arrow buffers cross the language boundary by reference. Files still +need decoding, and downstream Spark operators can convert batches to rows. + +The SDK centralizes resource cleanup and translates provider errors and +unwinding panics from JNI calls and reader callbacks. Process aborts and panics +in custom Arrow buffer destructors cannot be recovered. +Exported buffers must remain valid until their consumer releases them; releasing +the stream does not release batches already handed to its consumer. + +This demo does not introduce task cancellation, Spark memory reservations, +custom metrics, nested projection, a generic filter evaluator, C++ bindings, +or a new C ABI. Network readers still need their own I/O timeouts. It does not +fix the implementation issues tracked on #59313 or claim full end-to-end +compatibility with later revisions of that PR. + +[range]: ../../../examples/src/main/rust/native-data-source/README.md +[range-source]: ../../../examples/src/main/rust/native-data-source/src/lib.rs diff --git a/sql/native/rust/src/bridge.rs b/sql/native/rust/src/bridge.rs new file mode 100644 index 0000000000000..27244a0d56c98 --- /dev/null +++ b/sql/native/rust/src/bridge.rs @@ -0,0 +1,504 @@ +// 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. + +//! Implementation details used by the exported JNI functions. + +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::{ffi::FFI_ArrowSchema, Schema, SchemaRef}; +pub use jni; +use jni::objects::{JObject, JObjectArray, JString}; +use jni::sys::{jbooleanArray, jlong, jobjectArray}; +use jni::JNIEnv; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; + +use crate::{BatchSource, FilterSupport, LimitSupport, Options, Predicate, Result, ScanBuilder}; + +#[path = "stream.rs"] +mod stream; + +const STATE_VERSION: u32 = 1; + +type HandleMap = Mutex>>>; + +/// One registry per loaded plugin. Values have independent locks, so calls on +/// unrelated handles do not hold a shared lock while provider code runs. +pub struct Bridge { + next_handle: AtomicI64, + sources: HandleMap, + readers: HandleMap>, +} + +impl Default for Bridge { + fn default() -> Self { + Self { + next_handle: AtomicI64::new(1), + sources: Mutex::new(HashMap::new()), + readers: Mutex::new(HashMap::new()), + } + } +} + +impl Bridge { + fn insert(&self, handles: &HandleMap, value: T) -> Result { + let handle = self + .next_handle + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1)) + .map_err(|_| "native handle IDs exhausted")?; + lock(handles).insert(handle, Arc::new(Mutex::new(value))); + Ok(handle) + } + + pub fn create(&self, options: Options) -> Result { + self.insert(&self.sources, S::new(options)?) + } + + pub fn schema(&self, handle: jlong) -> Result { + let source = get(&self.sources, handle)?; + let result = lock(&source).schema(); + result + } + + pub fn create_reader(&self, handle: jlong, schema: SchemaRef) -> Result { + let source = get(&self.sources, handle)?; + let builder = lock(&source).scan(schema)?; + self.insert(&self.readers, Reader::new(builder)) + } + + pub fn push_predicates(&self, handle: jlong, predicates: &[Predicate]) -> Result> { + let reader = get(&self.readers, handle)?; + let mut reader = lock(&reader); + let support = reader.builder()?.push_filters(predicates)?; + if support.len() != predicates.len() { + // Never permit a later limit after an invalid filter negotiation. + reader.has_residual = true; + return Err("push_filters must return one result per predicate".into()); + } + let exact: Vec = support.iter().map(|s| *s == FilterSupport::Exact).collect(); + reader.has_residual |= exact.iter().any(|s| !s); + Ok(exact) + } + + pub fn push_limit(&self, handle: jlong, limit: usize) -> Result { + let reader = get(&self.readers, handle)?; + let mut reader = lock(&reader); + if reader.has_residual { + return Ok(false); + } + Ok(reader.builder()?.push_limit(limit)? == LimitSupport::Global) + } + + pub fn prune_columns(&self, handle: jlong, columns: &[String]) -> Result { + let reader = get(&self.readers, handle)?; + let result = lock(&reader).builder()?.prune_columns(columns); + result + } + + pub fn partitions(&self, handle: jlong) -> Result>> { + let reader = get(&self.readers, handle)?; + let result = lock(&reader).plan().map(|plan| plan.partitions.clone()); + result + } + + pub fn serialize_reader(&self, handle: jlong) -> Result> { + let reader = get(&self.readers, handle)?; + let result = lock(&reader).plan().map(|plan| plan.state.clone()); + result + } + + pub fn close_source(&self, handle: jlong) -> Result<()> { + remove(&self.sources, handle) + } + + pub fn close_reader(&self, handle: jlong) -> Result<()> { + remove(&self.readers, handle) + } + + pub fn open_reader(&self, state: &[u8], partition: &[u8]) -> Result { + let reader = S::open(decode(state)?, decode(partition)?)?; + stream::export(reader) + } +} + +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + // A provider panic is reported through JNI. Recover the lock so close can + // still remove and release its handle after that failed call. + mutex.lock().unwrap_or_else(|error| error.into_inner()) +} + +fn get(handles: &HandleMap, handle: jlong) -> Result>> { + lock(handles) + .get(&handle) + .cloned() + .ok_or_else(|| format!("unknown or closed native handle: {handle}").into()) +} + +fn remove(handles: &HandleMap, handle: jlong) -> Result<()> { + 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); + Ok(()) +} + +struct Reader { + builder: Option, + plan: Option, + has_residual: bool, +} + +struct EncodedPlan { + state: Vec, + partitions: Vec>, +} + +impl Reader { + fn new(builder: B) -> Self { + Self { + builder: Some(builder), + plan: None, + has_residual: false, + } + } + + fn builder(&mut self) -> Result<&mut B> { + self.builder + .as_mut() + .ok_or_else(|| "scan has already been planned".into()) + } + + fn plan(&mut self) -> Result<&EncodedPlan> { + if self.plan.is_none() { + let builder = self + .builder + .take() + .ok_or("scan planning previously failed")?; + let plan = builder.build()?; + self.plan = Some(EncodedPlan { + state: encode(&plan.state)?, + partitions: plan.partitions.iter().map(encode).collect::>()?, + }); + } + Ok(self.plan.as_ref().expect("plan was initialized")) + } +} + +#[derive(Serialize, Deserialize)] +struct Envelope { + version: u32, + value: T, +} + +fn encode(value: &T) -> Result> { + Ok(serde_json::to_vec(&Envelope { + version: STATE_VERSION, + value, + })?) +} + +fn decode(bytes: &[u8]) -> Result { + let envelope: Envelope = serde_json::from_slice(bytes)?; + if envelope.version != STATE_VERSION { + return Err(format!("unsupported native state version: {}", envelope.version).into()); + } + Ok(envelope.value) +} + +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}") +} + +/// Converts all errors and unwinding panics into Java exceptions. This does +/// not recover from aborting panics, allocation failure, or process faults. +pub fn jni_call<'local, T: Default>( + env: &mut JNIEnv<'local>, + action: impl FnOnce(&mut JNIEnv<'local>) -> Result, +) -> T { + let result = catch_unwind(AssertUnwindSafe(|| action(env).map_err(|e| e.to_string()))); + let message = match result { + Ok(Ok(value)) => return value, + Ok(Err(message)) => message, + Err(payload) => panic_message(payload), + }; + // Keep a pending JVM exception (for example an allocation failure). + if !env.exception_check().unwrap_or(true) { + let _ = env.throw_new("java/lang/RuntimeException", message); + } + T::default() +} + +pub fn strings(env: &mut JNIEnv<'_>, array: &JObjectArray<'_>) -> Result> { + let len = env.get_array_length(array)?; + let mut strings = Vec::with_capacity(len as usize); + for index in 0..len { + let string = JString::from(env.get_object_array_element(array, index)?); + strings.push(env.get_string(&string)?.into()); + env.delete_local_ref(string)?; + } + Ok(strings) +} + +pub fn options( + env: &mut JNIEnv<'_>, + keys: &JObjectArray<'_>, + values: &JObjectArray<'_>, +) -> Result { + let keys = strings(env, keys)?; + let values = strings(env, values)?; + if keys.len() != values.len() { + return Err("option keys and values have different lengths".into()); + } + Ok(keys.into_iter().zip(values).collect()) +} + +pub fn booleans(env: &mut JNIEnv<'_>, values: &[bool]) -> Result { + let array = env.new_boolean_array(values.len().try_into()?)?; + let values: Vec<_> = values.iter().map(|v| u8::from(*v)).collect(); + env.set_boolean_array_region(&array, 0, &values)?; + Ok(array.into_raw()) +} + +pub fn parse_predicates(predicates: &[String]) -> Result> { + predicates + .iter() + .map(|text| Ok(serde_json::from_str(text)?)) + .collect() +} + +pub fn byte_arrays(env: &mut JNIEnv<'_>, values: &[Vec]) -> Result { + let array = env.new_object_array(values.len().try_into()?, "[B", JObject::null())?; + for (index, value) in values.iter().enumerate() { + let bytes = env.byte_array_from_slice(value)?; + env.set_object_array_element(&array, index.try_into()?, &bytes)?; + env.delete_local_ref(bytes)?; + } + Ok(array.into_raw()) +} + +fn check_address(address: jlong) -> Result<*mut T> { + let pointer = address as usize as *mut T; + if pointer.is_null() || !pointer.is_aligned() { + return Err("null or unaligned Arrow structure address".into()); + } + Ok(pointer) +} + +/// Moves a schema from Spark's C structure and releases it after conversion. +/// +/// # Safety +/// The address must point to a live, initialized ArrowSchema that the caller +/// permits this function to move, as required by NativeBridge ABI v1. +pub unsafe fn import_schema(address: jlong) -> Result { + let address = check_address::(address)?; + let schema = unsafe { FFI_ArrowSchema::from_raw(address) }; + Ok(Arc::new(Schema::try_from(&schema)?)) +} + +/// Moves an exported schema into Spark's output ArrowSchema structure. +/// +/// # Safety +/// The address must point to writable, aligned storage for an ArrowSchema +/// whose release callback is null. Spark retains ownership of that storage. +pub unsafe fn export_schema(schema: &SchemaRef, address: jlong) -> Result<()> { + let address = check_address::(address)?; + let schema = FFI_ArrowSchema::try_from(schema.as_ref())?; + unsafe { address.write(schema) }; + Ok(()) +} + +/// Initializes Spark's output stream, transferring ownership of the reader. +/// +/// # Safety +/// The address must point to writable, aligned storage for an ArrowArrayStream +/// whose release callback is null. Spark retains ownership of that storage. +pub unsafe fn export_stream(stream: FFI_ArrowArrayStream, address: jlong) -> Result<()> { + let address = check_address::(address)?; + unsafe { address.write(stream) }; + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicUsize; + + use super::*; + use crate::{BatchReader, ScanPlan}; + + struct Source { + builds: Arc, + limits: Arc, + drops: Arc, + filter_support: FilterSupport, + limit_support: LimitSupport, + } + + impl Drop for Source { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } + } + + struct Builder { + builds: Arc, + limits: Arc, + filter_support: FilterSupport, + limit_support: LimitSupport, + } + + impl BatchSource for Source { + type Scan = Builder; + + fn new(_: Options) -> Result { + Ok(Self { + builds: Arc::default(), + limits: Arc::default(), + drops: Arc::default(), + filter_support: FilterSupport::Unsupported, + limit_support: LimitSupport::Unsupported, + }) + } + + fn schema(&self) -> Result { + Ok(Arc::new(Schema::empty())) + } + + fn scan(&self, _: SchemaRef) -> Result { + Ok(Builder { + builds: self.builds.clone(), + limits: self.limits.clone(), + filter_support: self.filter_support, + limit_support: self.limit_support, + }) + } + + fn open(_: u64, _: u64) -> Result { + unreachable!("these tests inspect driver-side planning") + } + } + + impl ScanBuilder for Builder { + type State = u64; + type Partition = u64; + + fn push_filters(&mut self, filters: &[Predicate]) -> Result> { + Ok(vec![self.filter_support; filters.len()]) + } + + fn push_limit(&mut self, _: usize) -> Result { + self.limits.fetch_add(1, Ordering::SeqCst); + Ok(self.limit_support) + } + + fn build(self) -> Result> { + self.builds.fetch_add(1, Ordering::SeqCst); + Ok(ScanPlan { + state: 42, + partitions: vec![2, 7], + }) + } + } + + #[test] + fn source_can_close_before_planning_and_plan_is_built_once() { + let bridge = Bridge::::default(); + let source = Source::new(Options::new()).unwrap(); + let builds = source.builds.clone(); + let drops = source.drops.clone(); + let handle = bridge.insert(&bridge.sources, source).unwrap(); + let schema = bridge.schema(handle).unwrap(); + let reader = bridge.create_reader(handle, schema).unwrap(); + bridge.close_source(handle).unwrap(); + bridge.close_source(handle).unwrap(); + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert!(bridge.schema(handle).is_err()); + + let partitions = bridge.partitions(reader).unwrap(); + assert_eq!(bridge.partitions(reader).unwrap(), partitions); + assert_eq!(decode::(&partitions[0]).unwrap(), 2); + assert_eq!(decode::(&partitions[1]).unwrap(), 7); + assert_eq!( + decode::(&bridge.serialize_reader(reader).unwrap()).unwrap(), + 42 + ); + assert_eq!(builds.load(Ordering::SeqCst), 1); + assert!(bridge.prune_columns(reader, &[]).is_err()); + bridge.close_reader(reader).unwrap(); + assert!(bridge.partitions(reader).is_err()); + } + + #[test] + fn residual_filters_prevent_calling_the_providers_limit_method() { + for support in [FilterSupport::PruningOnly, FilterSupport::Unsupported] { + let bridge = Bridge::::default(); + let mut source = Source::new(Options::new()).unwrap(); + source.filter_support = support; + source.limit_support = LimitSupport::Global; + let limits = source.limits.clone(); + let handle = bridge.insert(&bridge.sources, source).unwrap(); + let reader = bridge + .create_reader(handle, bridge.schema(handle).unwrap()) + .unwrap(); + assert_eq!( + bridge + .push_predicates(reader, &[Predicate::Unknown]) + .unwrap(), + vec![false] + ); + assert!(!bridge.push_limit(reader, 1).unwrap()); + assert_eq!(limits.load(Ordering::SeqCst), 0); + } + } + + #[test] + fn only_global_limits_return_true_to_jni() { + for support in [LimitSupport::PerPartition, LimitSupport::Global] { + let bridge = Bridge::::default(); + let mut source = Source::new(Options::new()).unwrap(); + source.limit_support = support; + let limits = source.limits.clone(); + let handle = bridge.insert(&bridge.sources, source).unwrap(); + let reader = bridge + .create_reader(handle, bridge.schema(handle).unwrap()) + .unwrap(); + assert_eq!( + bridge.push_limit(reader, 4).unwrap(), + support == LimitSupport::Global + ); + assert_eq!(limits.load(Ordering::SeqCst), 1); + } + } + + #[test] + fn 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..3667dcee8c0fe --- /dev/null +++ b/sql/native/rust/src/export.rs @@ -0,0 +1,187 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/// Exports one batch-read provider as the JNI methods of NativeBridge ABI v1. +/// +/// Call this macro once in a `cdylib`. JNI and Arrow pointer handling stay in +/// the SDK; providers implement the safe [`crate::BatchSource`] trait. +#[macro_export] +macro_rules! export_data_source { + ($source:ty) => { + mod __spark_native_datasource_exports { + use super::*; + use $crate::bridge::jni::objects::{JByteArray, JClass, JObjectArray, JString}; + use $crate::bridge::jni::sys::{ + jboolean, jbooleanArray, jbyteArray, jint, jlong, jobjectArray, + }; + use $crate::bridge::jni::JNIEnv; + use $crate::bridge::{self, Bridge}; + + fn bridge() -> &'static Bridge<$source> { + static BRIDGE: std::sync::OnceLock> = std::sync::OnceLock::new(); + BRIDGE.get_or_init(Bridge::default) + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_abiVersion"] + pub extern "system" fn abi_version(_env: JNIEnv<'_>, _class: JClass<'_>) -> jint { + 1 + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource"] + pub extern "system" fn create_data_source( + mut env: JNIEnv<'_>, + _class: JClass<'_>, + _name: JString<'_>, + keys: JObjectArray<'_>, + values: JObjectArray<'_>, + ) -> jlong { + bridge::jni_call(&mut env, |env| { + bridge().create(bridge::options(env, &keys, &values)?) + }) + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_schema"] + /// # Safety + /// Spark must supply writable ArrowSchema storage as specified by ABI v1. + pub unsafe extern "system" fn schema( + mut env: JNIEnv<'_>, + _class: JClass<'_>, + source: jlong, + address: jlong, + ) { + bridge::jni_call(&mut env, |_| { + let schema = bridge().schema(source)?; + unsafe { bridge::export_schema(&schema, address) } + }) + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_createReader"] + /// # Safety + /// Spark must supply an initialized ArrowSchema that this call may move. + pub unsafe extern "system" fn create_reader( + mut env: JNIEnv<'_>, + _class: JClass<'_>, + source: jlong, + address: jlong, + ) -> jlong { + bridge::jni_call(&mut env, |_| { + let schema = unsafe { bridge::import_schema(address)? }; + bridge().create_reader(source, schema) + }) + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_pushPredicates"] + pub extern "system" fn push_predicates( + mut env: JNIEnv<'_>, + _class: JClass<'_>, + reader: jlong, + predicates: JObjectArray<'_>, + ) -> jbooleanArray { + bridge::jni_call(&mut env, |env| { + let predicates = bridge::strings(env, &predicates)?; + let predicates = bridge::parse_predicates(&predicates)?; + let exact = bridge().push_predicates(reader, &predicates)?; + bridge::booleans(env, &exact) + }) + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_pushLimit"] + pub extern "system" fn push_limit( + mut env: JNIEnv<'_>, + _class: JClass<'_>, + reader: jlong, + limit: jint, + ) -> jboolean { + bridge::jni_call(&mut env, |_| { + Ok(u8::from(bridge().push_limit(reader, limit.try_into()?)?)) + }) + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_pruneColumns"] + pub extern "system" fn prune_columns( + mut env: JNIEnv<'_>, + _class: JClass<'_>, + reader: jlong, + columns: JObjectArray<'_>, + ) -> jboolean { + bridge::jni_call(&mut env, |env| { + Ok(u8::from( + bridge().prune_columns(reader, &bridge::strings(env, &columns)?)?, + )) + }) + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_partitions"] + pub extern "system" fn partitions( + mut env: JNIEnv<'_>, + _class: JClass<'_>, + reader: jlong, + ) -> jobjectArray { + bridge::jni_call(&mut env, |env| { + bridge::byte_arrays(env, &bridge().partitions(reader)?) + }) + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_serializeReader"] + pub extern "system" fn serialize_reader( + mut env: JNIEnv<'_>, + _class: JClass<'_>, + reader: jlong, + ) -> jbyteArray { + bridge::jni_call(&mut env, |env| { + Ok(env + .byte_array_from_slice(&bridge().serialize_reader(reader)?)? + .into_raw()) + }) + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_read"] + /// # Safety + /// Spark must supply writable ArrowArrayStream storage as specified by ABI v1. + pub unsafe extern "system" fn read( + mut env: JNIEnv<'_>, + _class: JClass<'_>, + state: JByteArray<'_>, + partition: JByteArray<'_>, + address: jlong, + ) { + bridge::jni_call(&mut env, |env| { + let state = env.convert_byte_array(&state)?; + let partition = env.convert_byte_array(&partition)?; + let stream = bridge().open_reader(&state, &partition)?; + unsafe { bridge::export_stream(stream, address) } + }) + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_closeReader"] + pub extern "system" fn close_reader( + mut env: JNIEnv<'_>, + _class: JClass<'_>, + reader: jlong, + ) { + bridge::jni_call(&mut env, |_| bridge().close_reader(reader)) + } + + #[export_name = "Java_org_apache_spark_sql_datasource_NativeBridge_closeDataSource"] + pub extern "system" fn close_data_source( + mut env: JNIEnv<'_>, + _class: JClass<'_>, + source: jlong, + ) { + bridge::jni_call(&mut env, |_| bridge().close_source(source)) + } + } + }; +} diff --git a/sql/native/rust/src/lib.rs b/sql/native/rust/src/lib.rs new file mode 100644 index 0000000000000..639c96a3d6bde --- /dev/null +++ b/sql/native/rust/src/lib.rs @@ -0,0 +1,130 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Experimental batch-read SDK for the JNI ABI proposed in Spark PR #59313. +//! +//! Implement [`BatchSource`] and [`ScanBuilder`], then invoke +//! [`export_data_source!`] once in a `cdylib` crate. The SDK owns JNI conversion, +//! handles, plan serialization, and Arrow C Stream export. Rust traits stay +//! inside the plugin; they are not a binary interface between dynamic libraries. + +use std::collections::BTreeMap; + +use arrow_array::RecordBatchReader; +use arrow_schema::SchemaRef; +use serde::de::DeserializeOwned; +use serde::Serialize; + +#[doc(hidden)] +pub mod bridge; +mod export; +mod predicate; + +pub use predicate::Predicate; + +pub type Error = Box; +pub type Result = std::result::Result; +pub type Options = BTreeMap; +pub type BatchReader = Box; + +/// How completely a connector evaluates a pushed predicate. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FilterSupport { + /// Implements Spark SQL semantics and returns only matching rows. + Exact, + /// May skip data, but Spark must evaluate the predicate again. + PruningOnly, + /// Spark must evaluate the predicate. + Unsupported, +} + +/// How a connector limits the rows produced by the whole scan. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LimitSupport { + Unsupported, + /// Caps each partition; Spark retains the global limit. JNI returns false + /// because ABI v1's boolean promises a bound on the *whole* scan. + PerPartition, + /// Caps the whole scan; JNI returns true. Spark still retains its limit. + Global, +} + +/// A driver-created plan containing only serializable values. +/// +/// The SDK serializes the state once and each partition independently. Include +/// snapshot identifiers here when needed for repeatable task attempts. Do not +/// put pointers, open files, connections, or credentials with an assumed +/// process-local lifetime into these values. +pub struct ScanPlan { + pub state: S, + pub partitions: Vec

, +} + +/// A provider's driver-side configuration and executor-side reader factory. +pub trait BatchSource: Send + Sized + 'static { + type Scan: ScanBuilder; + + /// Spark passes option keys in lower case. + fn new(options: Options) -> Result; + + fn schema(&self) -> Result; + + /// The schema may be supplied by the user; reject unsupported schemas. + /// The returned builder must remain valid after this source is dropped. + fn scan(&self, schema: SchemaRef) -> Result; + + /// Called in an executor with restored values, without a driver object. + /// + /// Return batches lazily with a fixed schema and bounded memory. Buffers + /// already exported to Spark must remain valid until Spark releases them. + /// Use Arrow types supported by the pinned Spark ABI, with zero array + /// offsets at every nesting level. Normalize sliced arrays when necessary; + /// the SDK does not rewrite buffers or expand dictionary encoding. + /// Blocking I/O needs timeouts: ABI v1 does not provide task cancellation. + fn open( + state: ::State, + partition: ::Partition, + ) -> Result; +} + +/// Builds one scan. Spark negotiates filters, limit, and projection before +/// planning partitions; the SDK consumes the builder and caches the plan. +pub trait ScanBuilder: Send + Sized + 'static { + type State: Serialize + DeserializeOwned + Send; + type Partition: Serialize + DeserializeOwned + Send; + + /// Return one result per input predicate in the same order. Only `Exact` + /// allows Spark to remove a residual filter. Unknown functions and types + /// must not be reported as exact. Unsupported filters need not be stored. + fn push_filters(&mut self, filters: &[Predicate]) -> Result> { + Ok(vec![FilterSupport::Unsupported; filters.len()]) + } + + /// Count rows after filtering, never merely scanned rows. Spark decides + /// whether pushing a limit is safe for the query; the SDK also declines + /// limit pushdown 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. An empty projection is valid and + /// must preserve row counts. Keep internally needed filter columns. + fn prune_columns(&mut self, _columns: &[String]) -> Result { + Ok(false) + } + + fn build(self) -> Result>; +} diff --git a/sql/native/rust/src/predicate.rs b/sql/native/rust/src/predicate.rs new file mode 100644 index 0000000000000..f0ff0b8391712 --- /dev/null +++ b/sql/native/rust/src/predicate.rs @@ -0,0 +1,83 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +/// The expression tree documented by `NativeBridge.pushPredicates`. +/// +/// Literal values keep their JSON representation so a connector can reject +/// unsupported types without lossy numeric conversions. In particular, +/// decimal values are strings and timestamps are integer microseconds. +/// Function names are extensible: unknown names must remain residual. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "lowercase")] +pub enum Predicate { + Column { + name: Vec, + }, + Literal { + #[serde(rename = "dataType")] + data_type: String, + value: Value, + }, + Function { + name: String, + children: Vec, + }, + #[serde(other)] + Unknown, +} + +#[cfg(test)] +mod tests { + use super::Predicate; + + #[test] + fn preserves_int64_literals_and_nested_column_paths() { + let json = r#"{"type":"function","name":">=","children":[ + {"type":"column","name":["outer","id"]}, + {"type":"literal","dataType":"bigint","value":9223372036854775807} + ]}"#; + let predicate: Predicate = serde_json::from_str(json).unwrap(); + let Predicate::Function { children, .. } = predicate else { + panic!("expected a function"); + }; + assert_eq!( + children[0], + Predicate::Column { + name: vec!["outer".into(), "id".into()] + } + ); + let Predicate::Literal { value, .. } = &children[1] else { + panic!("expected a literal"); + }; + assert_eq!(value.as_i64(), Some(i64::MAX)); + } + + #[test] + fn preserves_unknown_expressions_for_conservative_fallback() { + let json = r#"{"type":"future_expression","value":1}"#; + assert_eq!( + serde_json::from_str::(json).unwrap(), + Predicate::Unknown + ); + let json = r#"{"type":"function","name":"FUTURE_FUNCTION","children":[]}"#; + let Predicate::Function { name, .. } = serde_json::from_str(json).unwrap() else { + panic!("expected a function"); + }; + assert_eq!(name, "FUTURE_FUNCTION"); + } +} diff --git a/sql/native/rust/src/stream.rs b/sql/native/rust/src/stream.rs new file mode 100644 index 0000000000000..75898592989bc --- /dev/null +++ b/sql/native/rust/src/stream.rs @@ -0,0 +1,270 @@ +// 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_schema::{ffi::FFI_ArrowSchema, Schema, SchemaRef}; + +use super::panic_message; +use crate::{BatchReader, Result}; + +struct StreamState { + reader: Option, + schema: SchemaRef, + last_error: Option, + finished: bool, +} + +impl Drop for StreamState { + fn drop(&mut self) { + // release has no error return channel, but it must never unwind into C. + if let Some(reader) = self.reader.take() { + if let Err(payload) = catch_unwind(AssertUnwindSafe(|| drop(reader))) { + let _ = panic_message(payload); + } + } + } +} + +pub(super) fn export(reader: BatchReader) -> Result { + // Establish guarded destruction before calling the provider's schema(). + let mut state = Box::new(StreamState { + reader: Some(reader), + schema: Arc::new(Schema::empty()), + last_error: None, + finished: false, + }); + state.schema = catch_unwind(AssertUnwindSafe(|| { + state.reader.as_ref().expect("reader is present").schema() + })) + .map_err(panic_message)?; + Ok(FFI_ArrowArrayStream { + get_schema: Some(get_schema), + get_next: Some(get_next), + get_last_error: Some(get_last_error), + release: Some(release), + private_data: Box::into_raw(state).cast::(), + }) +} + +fn callback(state: &mut StreamState, action: impl FnOnce(&mut StreamState) -> Result<()>) -> c_int { + let result = catch_unwind(AssertUnwindSafe(|| { + action(state).map_err(|e| e.to_string()) + })); + let message = match result { + Ok(Ok(())) => return 0, + Ok(Err(message)) => message, + Err(payload) => panic_message(payload), + }; + state.finished = true; + // Error text is a C string. A provider error may contain embedded NULs. + state.last_error = Some(CString::new(message.replace('\0', "\\0")).expect("NULs removed")); + 5 // EIO; get_last_error supplies the provider's actual error message. +} + +unsafe extern "C" fn get_schema( + stream: *mut FFI_ArrowArrayStream, + output: *mut FFI_ArrowSchema, +) -> c_int { + let state = unsafe { &mut *(*stream).private_data.cast::() }; + callback(state, |state| { + let schema = FFI_ArrowSchema::try_from(state.schema.as_ref())?; + unsafe { output.write(schema) }; + Ok(()) + }) +} + +unsafe extern "C" fn get_next( + stream: *mut FFI_ArrowArrayStream, + output: *mut FFI_ArrowArray, +) -> c_int { + let state = unsafe { &mut *(*stream).private_data.cast::() }; + callback(state, |state| { + if state.finished { + unsafe { output.write(FFI_ArrowArray::empty()) }; + return Ok(()); + } + match state.reader.as_mut().expect("reader is present").next() { + Some(batch) => { + let batch = batch?; + if batch.schema() != state.schema { + return Err("reader returned a batch with a different schema".into()); + } + let data = StructArray::from(batch).to_data(); + 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); + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use arrow_array::ffi_stream::ArrowArrayStreamReader; + use arrow_array::{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 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); + } +}