From 56acc8bea7362d8618d1b3939fd2b98122c41401 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 9 Oct 2026 10:09:08 +0900 Subject: [PATCH 01/27] [SPARK-60101][SQL] Add a columnar data source API for JVM and native (Rust, C++) data sources This adds the columnar data source API, a simpler way to plug a data source into Spark than implementing the Data Source V2 interfaces directly. It follows the Python Data Source API, but data is exchanged only as columnar batches: readers return ColumnarBatches, and writers receive Arrow-backed ColumnarBatches. Spark adapts a data source to Data Source V2, with batch and micro-batch reads, batch and streaming writes, and predicate, column and limit pushdown. A data source is implemented either on the JVM, with the interfaces of org.apache.spark.sql.datasource and a DataSourceProvider, or in native code such as Rust or C++, as a shared library that implements the native methods of NativeBridge as JNI functions and exchanges data through the Arrow C data interface. Spark loads each library into its own copy of NativeBridge, so libraries exporting the same functions can be loaded side by side, and Spark itself ships no native code. A native data source is distributed as a native data source package, a zip file with the extension .sparkpkg that contains a manifest and the library for each platform. Packages are added with spark.addArtifact, in Scala and Python, classic and Spark Connect, or found under spark.sql.dataSource.native.paths; executors get them through the artifacts of the session. The change also adds the spark.sql.dataSource.native.enabled and spark.sql.dataSource.native.paths configurations, the native data source error conditions, the arrow-c-data dependency, and a documentation page with examples in Scala, Rust and C++. Co-authored-by: Isaac --- LICENSE-binary | 1 + .../resources/error/error-conditions.json | 46 + dev/deps/spark-deps-hadoop-3-hive-2.3 | 1 + docs/_data/menu-sql.yaml | 11 + docs/sql-columnar-data-sources.md | 573 +++++++++ docs/sql-v2-data-sources.md | 4 +- pom.xml | 19 + python/pyspark/sql/connect/client/artifact.py | 3 +- python/pyspark/sql/session.py | 11 +- .../sql/tests/connect/client/test_artifact.py | 15 + python/pyspark/sql/tests/test_artifact.py | 12 + .../scala/org/apache/spark/sql/Artifact.scala | 12 +- .../sql/errors/QueryCompilationErrors.scala | 6 + .../sql/errors/QueryExecutionErrors.scala | 26 + .../apache/spark/sql/internal/SQLConf.scala | 12 + .../spark/sql/internal/StaticSQLConf.scala | 12 + sql/core/pom.xml | 4 + .../spark/sql/datasource/DataSource.java | 102 ++ .../sql/datasource/DataSourceProvider.java | 65 + .../sql/datasource/DataSourceReader.java | 89 ++ .../datasource/DataSourceStreamReader.java | 73 ++ .../datasource/DataSourceStreamWriter.java | 69 + .../sql/datasource/DataSourceWriter.java | 66 + .../spark/sql/datasource/NativeBridge.java | 321 +++++ .../datasources/v2/ffi/NativeLibrary.java | 292 +++++ .../spark/sql/artifact/ArtifactManager.scala | 13 + .../catalyst/analysis/ResolveDataSource.scala | 5 +- .../spark/sql/classic/DataStreamWriter.scala | 5 +- .../execution/datasources/DataSource.scala | 14 +- .../datasources/v2/NamedTableProvider.scala | 30 + .../v2/columnar/ColumnarTable.scala | 212 ++++ .../v2/columnar/ColumnarWrite.scala | 201 +++ .../datasources/v2/ffi/NativeDataSource.scala | 614 +++++++++ .../v2/ffi/NativeDataSourcePackage.scala | 264 ++++ .../v2/ffi/NativeDataSourceRegistry.scala | 143 +++ .../v2/ffi/NativeDataSourceV2.scala | 66 + .../datasources/v2/ffi/NativePredicates.scala | 92 ++ .../v2/python/PythonDataSourceV2.scala | 5 +- ...pache.spark.sql.sources.DataSourceRegister | 1 + .../test_native_datasource.cc | 1107 +++++++++++++++++ .../datasource/ColumnarDataSourceSuite.scala | 372 ++++++ .../v2/ffi/NativeDataSourceSuite.scala | 604 +++++++++ 42 files changed, 5576 insertions(+), 17 deletions(-) create mode 100644 docs/sql-columnar-data-sources.md create mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSource.java create mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceProvider.java create mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceReader.java create mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamReader.java create mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamWriter.java create mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceWriter.java create mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java create mode 100644 sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/NamedTableProvider.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala create mode 100644 sql/core/src/test/resources/native-datasource/test_native_datasource.cc create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/datasource/ColumnarDataSourceSuite.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala diff --git a/LICENSE-binary b/LICENSE-binary index 7051e09a6383a..53dec42594efa 100644 --- a/LICENSE-binary +++ b/LICENSE-binary @@ -301,6 +301,7 @@ joda-time:joda-time net.sf.opencsv:opencsv net.sf.supercsv:super-csv net.sf.jpam:jpam +org.apache.arrow:arrow-c-data org.apache.arrow:arrow-compression org.apache.arrow:arrow-format org.apache.arrow:arrow-memory-core diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 3a796a031afaf..279a426d02b9f 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -4792,6 +4792,34 @@ ], "sqlState" : "42000" }, + "INVALID_NATIVE_DATA_SOURCE_PACKAGE" : { + "message" : [ + "Invalid native data source package :" + ], + "subClass" : { + "INVALID_LIBRARY" : { + "message" : [ + "Cannot load the library : " + ] + }, + "INVALID_MANIFEST" : { + "message" : [ + "" + ] + }, + "UNSUPPORTED_ABI_VERSION" : { + "message" : [ + "It requires version of the native data source interface, but Spark supports version ." + ] + }, + "UNSUPPORTED_PLATFORM" : { + "message" : [ + "It has no library for the platform . It has libraries for: ." + ] + } + }, + "sqlState" : "22023" + }, "INVALID_NON_DETERMINISTIC_EXPRESSIONS" : { "message" : [ "The operator expects a deterministic expression, but the actual expression is ." @@ -6262,6 +6290,24 @@ ], "sqlState" : "0A000" }, + "NATIVE_DATA_SOURCE_ERROR" : { + "message" : [ + "Failed to native data source : " + ], + "sqlState" : "38000" + }, + "NATIVE_DATA_SOURCE_PACKAGE_CONFLICT" : { + "message" : [ + "More than one native data source package provides the data source : . Remove all but one of them." + ], + "sqlState" : "42710" + }, + "NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND" : { + "message" : [ + "Cannot find the native data source package with the SHA-256 checksum . Add the package with `spark.addArtifact`, or install the same package under `spark.sql.dataSource.native.paths` on every node." + ], + "sqlState" : "42K02" + }, "NEAREST_BY_JOIN" : { "message" : [ "Invalid nearest-by join." diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index 50f978819c3da..b088d623ea7cb 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -21,6 +21,7 @@ antlr4-runtime/4.13.1//antlr4-runtime-4.13.1.jar aopalliance-repackaged/3.0.6//aopalliance-repackaged-3.0.6.jar arpack/3.2.0//arpack-3.2.0.jar arpack_combined_all/0.1//arpack_combined_all-0.1.jar +arrow-c-data/19.0.0//arrow-c-data-19.0.0.jar arrow-compression/19.0.0//arrow-compression-19.0.0.jar arrow-format/19.0.0//arrow-format-19.0.0.jar arrow-memory-core/19.0.0//arrow-memory-core-19.0.0.jar diff --git a/docs/_data/menu-sql.yaml b/docs/_data/menu-sql.yaml index a100f0ccb1ce3..e4b8bb243e47d 100644 --- a/docs/_data/menu-sql.yaml +++ b/docs/_data/menu-sql.yaml @@ -75,6 +75,17 @@ url: sql-v2-data-sources.html#streaming - text: Further Reading url: sql-v2-data-sources.html#further-reading +- text: Columnar Data Source API + url: sql-columnar-data-sources.html + subitems: + - text: Overview + url: sql-columnar-data-sources.html#overview + - text: The API + url: sql-columnar-data-sources.html#the-api + - text: JVM Data Sources + url: sql-columnar-data-sources.html#jvm-data-sources + - text: Native Data Sources + url: sql-columnar-data-sources.html#native-data-sources - text: Performance Tuning url: sql-performance-tuning.html subitems: diff --git a/docs/sql-columnar-data-sources.md b/docs/sql-columnar-data-sources.md new file mode 100644 index 0000000000000..ed1fca5914783 --- /dev/null +++ b/docs/sql-columnar-data-sources.md @@ -0,0 +1,573 @@ +--- +layout: global +title: Columnar Data Source API +displayTitle: Columnar Data Source API +license: | + 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. +--- + +* Table of contents +{:toc} + +## Overview + +The columnar data source API is a simpler way to plug a data source into Spark than implementing +the [Data Source V2](sql-v2-data-sources.html) interfaces directly. It is meant for data sources +that identify their data by options, such as a path or a URL. It follows the +[Python Data Source API](api/python/reference/pyspark.sql/api/pyspark.sql.datasource.DataSource.html), +but data is exchanged only as columnar batches: a data source returns +[`ColumnarBatch`](api/java/org/apache/spark/sql/vectorized/ColumnarBatch.html)es when reading, and +receives them when writing. Spark adapts a data source to Data Source V2, so it supports: + +- batch and micro-batch streaming reads, +- batch and streaming writes, appending to or overwriting the existing data, +- predicate, column and limit pushdown. + +A data source can be implemented in two ways: + +- **On the JVM**, in Java or Scala, with the interfaces of the + [`org.apache.spark.sql.datasource`](api/java/org/apache/spark/sql/datasource/package-summary.html) + package. +- **In native code**, such as Rust or C++, as a shared library that implements the same API + through the [`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) binary + interface. A native data source needs no JVM code: it is distributed as a native data source + package, and Spark loads it when a query uses it. + +## The API + +| Interface | Methods | Description | +|-----------|---------|-------------| +| `DataSource` | `schema`, `reader`, `streamReader`, `writer`, `streamWriter` | Created on the driver with the options of the user, for each schema inference, scan and write. A data source only implements the operations it supports. | +| `DataSourceReader` | `pushPredicates`, `pushLimit`, `pruneColumns`, `partitions`, `read` | Reads the data for a batch scan. Spark pushes operations down and plans the partitions on the driver, then serializes the reader to the executors, where `read` returns the batches of a partition. | +| `DataSourceStreamReader` | `initialOffset`, `latestOffset`, `partitions`, `read`, `commit`, `stop` | Reads the data for a micro-batch streaming scan. Offsets are JSON strings defined by the data source. | +| `DataSourceWriter` | `createWriter`, `commit`, `abort` | Writes the data for a batch write. `createWriter` returns a `DataWriter` for each task, and Spark commits or aborts the write on the driver with the messages of the tasks. | +| `DataSourceStreamWriter` | `createWriter`, `commit`, `abort` | Writes the data for a streaming write, one micro-batch (epoch) at a time. | + +The batches that a writer receives are backed by Apache Arrow: every column is an +[`ArrowColumnVector`](api/java/org/apache/spark/sql/vectorized/ArrowColumnVector.html). A batch has +at most `spark.sql.execution.arrow.maxRecordsPerBatch` rows, and is only valid during the call. + +## JVM Data Sources + +A JVM data source extends +[`DataSourceProvider`](api/java/org/apache/spark/sql/datasource/DataSourceProvider.html), which +creates the `DataSource` with the options given by the user. The following data source reads the +numbers `[0, end)` as a column `id`, in two partitions: + +```scala +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.BigIntVector + +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} +import org.apache.spark.sql.datasource._ +import org.apache.spark.sql.types.{LongType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} + +class RangeProvider extends DataSourceProvider { + override def shortName(): String = "scala_range" + + override def createDataSource(options: CaseInsensitiveStringMap): DataSource = { + new RangeDataSource(options.getLong("end", 10)) + } +} + +class RangeDataSource(end: Long) extends DataSource { + override def schema(): StructType = new StructType().add("id", LongType, nullable = false) + + override def reader(schema: StructType): DataSourceReader = new RangeReader(end) +} + +case class RangePartition(start: Long, end: Long) extends InputPartition + +class RangeReader(end: Long) extends DataSourceReader { + override def partitions(): Array[InputPartition] = { + Array(RangePartition(0, end / 2), RangePartition(end / 2, end)) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + val range = partition.asInstanceOf[RangePartition] + new RangePartitionReader(range.start, range.end) + } +} + +// Returns the ids [start, end) in a single batch, backed by an Arrow vector. +class RangePartitionReader(start: Long, end: Long) extends PartitionReader[ColumnarBatch] { + private val allocator = new RootAllocator() + private val ids = new BigIntVector("id", allocator) + private var done = false + + override def next(): Boolean = { + if (done) return false + val count = (end - start).toInt + ids.allocateNew(count) + (0 until count).foreach(i => ids.set(i, start + i)) + ids.setValueCount(count) + done = true + true + } + + override def get(): ColumnarBatch = { + new ColumnarBatch(Array[ColumnVector](new ArrowColumnVector(ids)), ids.getValueCount) + } + + override def close(): Unit = { + ids.close() + allocator.close() + } +} +``` + +To use the data source by its short name, list the provider in +`META-INF/services/org.apache.spark.sql.sources.DataSourceRegister` in its JAR, like any other data +source. Its fully qualified class name can also be used as the format: + +```scala +spark.read.format("scala_range").option("end", 100).load() +``` + +## Native Data Sources + +### How It Works + +A native data source is a shared library that implements the `native` methods of +[`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) as JNI functions, +named after the class, such as `Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource`. +Spark loads each library into its own copy of the class, so any number of native data sources can +be loaded side by side. The methods follow the API above: Spark implements `DataSource`, +`DataSourceReader` and the other interfaces by calling them. + +- **Handles**: the `create` functions return a pointer to a native object, which Spark releases with + the matching `close` function. +- **Arrow**: schemas and data are exchanged with the + [Arrow C data interface](https://arrow.apache.org/docs/format/CDataInterface.html). A partition + is read by exporting an `ArrowArrayStream`, and a writer receives an `ArrowArray` for each batch. +- **State**: partitions, the state of readers and writers, and commit messages are byte arrays in a + format defined by the library. Spark sends them to the executors and back. +- **Errors**: a library reports an error by throwing a Java exception, for example with the JNI + function `ThrowNew`. Spark reports it as a `NATIVE_DATA_SOURCE_ERROR`. + +A library only exports the functions of the operations it supports. If a function is missing, +Spark treats the operation as unsupported, except for the pushdown and `serialize` functions, which +are optional: + +| Operation | Required functions | Optional functions | +|-----------|--------------------|--------------------| +| All | `abiVersion`, `createDataSource`, `closeDataSource` | `schema` | +| Batch read | `createReader`, `partitions`, `read`, `closeReader` | `pushPredicates`, `pushLimit`, `pruneColumns`, `serializeReader` | +| Streaming read | `createStreamReader`, `initialOffset`, `latestOffset`, `streamPartitions`, `read`, `closeStreamReader` | `serializeStreamReader`, `commitOffset` | +| Batch write | `createWriter`, `commit`, `abort`, `closeWriter`, `createDataWriter`, `write`, `commitDataWriter`, `abortDataWriter` | `serializeWriter` | +| Streaming write | `createStreamWriter`, `commit`, `abort`, `closeWriter`, `createDataWriter`, `write`, `commitDataWriter`, `abortDataWriter` | `serializeWriter` | + +The Javadoc of `NativeBridge` specifies each function, including the JSON format of the pushed +predicates. + +### Packages + +A native data source package is a zip file with the extension `.sparkpkg`. It contains a manifest +named `spark-native-datasource.json`, and the library built for one or more platforms: + +```json +{ + "abiVersion": 1, + "dataSources": ["rust_range"], + "libraries": { + "linux-x86_64": "linux-x86_64/librust_range.so", + "linux-aarch64": "linux-aarch64/librust_range.so", + "osx-aarch64": "osx-aarch64/librust_range.dylib" + } +} +``` + +- `abiVersion` is the version of `NativeBridge` that the library implements, currently 1. +- `dataSources` lists the names of the data sources that the library implements. The names are + case-insensitive, and Spark passes the name to `createDataSource`. +- `libraries` maps each platform, `-` with an `os` of `linux`, `osx` or `windows` and an + `arch` of `x86_64` or `aarch64`, to the path of the library in the package. + +Add a package to a session with `spark.addArtifact`, which works in Scala, Java and Python, and with +Spark Connect: + +```python +spark.addArtifact("/path/to/rust_range.sparkpkg") +spark.read.format("rust_range").option("end", 100).load().show() +``` + +Spark can also find packages under `spark.sql.dataSource.native.paths`, for example in a container +image. In both cases, the executors get the package automatically. Give packages distinct file +names, for example with their version: sessions that are not isolated share the files they add. + +A Java data source with the same name takes precedence over a native one, and so does a Python +data source. + +### Example in Rust + +This data source reads the numbers `[0, end)` as a column `id`, in two partitions. It uses the +[`jni`](https://crates.io/crates/jni) crate, and the `arrow-array` and `arrow-schema` crates of +[arrow-rs](https://github.com/apache/arrow-rs), which export schemas and streams with +`FFI_ArrowSchema` and `FFI_ArrowArrayStream`. `Cargo.toml`: + +```toml +[package] +name = "rust_range" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["cdylib"] + +[dependencies] +arrow-array = { version = "58", features = ["ffi"] } +arrow-schema = { version = "58", features = ["ffi"] } +jni = "0.21" +``` + +`src/lib.rs`: + +```rust +//! A native data source for Apache Spark, written in Rust. It reads the numbers [0, end) as a +//! column `id`, in two partitions. The option `end` defaults to 10. + +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::sync::Arc; + +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator}; +use arrow_schema::ffi::FFI_ArrowSchema; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use jni::objects::{JByteArray, JClass, JObject, JObjectArray, JString}; +use jni::sys::{jint, jlong, jobjectArray}; +use jni::JNIEnv; + +type Result = std::result::Result>; + +struct RangeSource { + end: i64, +} + +struct RangeReader { + end: i64, +} + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])) +} + +/// Runs `f`, and reports its error or panic as a Java exception: neither may cross into the JVM. +fn call(env: &mut JNIEnv, default: T, f: impl FnOnce(&mut JNIEnv) -> Result) -> T { + let message = match catch_unwind(AssertUnwindSafe(|| f(env))) { + Ok(Ok(value)) => return value, + Ok(Err(e)) => e.to_string(), + Err(_) => "the data source panicked".to_string(), + }; + if !env.exception_check().unwrap_or(true) { + let _ = env.throw_new("java/lang/RuntimeException", message); + } + default +} + +fn strings(env: &mut JNIEnv, array: &JObjectArray) -> Result> { + let mut result = Vec::new(); + for i in 0..env.get_array_length(array)? { + let value = JString::from(env.get_object_array_element(array, i)?); + result.push(env.get_string(&value)?.into()); + } + Ok(result) +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_abiVersion( + _env: JNIEnv, + _class: JClass, +) -> jint { + 1 +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource( + mut env: JNIEnv, + _class: JClass, + _name: JString, + keys: JObjectArray, + values: JObjectArray, +) -> jlong { + call(&mut env, 0, |env| { + let mut source = RangeSource { end: 10 }; + for (key, value) in strings(env, &keys)?.into_iter().zip(strings(env, &values)?) { + if key == "end" { + source.end = value.parse()?; + } + } + Ok(Box::into_raw(Box::new(source)) as jlong) + }) +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_schema( + mut env: JNIEnv, + _class: JClass, + _source: jlong, + schema_address: jlong, +) { + call(&mut env, (), |_| { + let schema = FFI_ArrowSchema::try_from(schema().as_ref())?; + unsafe { std::ptr::write(schema_address as *mut FFI_ArrowSchema, schema) }; + Ok(()) + }) +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_createReader( + mut env: JNIEnv, + _class: JClass, + source: jlong, + _schema_address: jlong, +) -> jlong { + call(&mut env, 0, |_| { + let source = unsafe { &*(source as *const RangeSource) }; + Ok(Box::into_raw(Box::new(RangeReader { end: source.end })) as jlong) + }) +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_closeDataSource( + _env: JNIEnv, + _class: JClass, + source: jlong, +) { + drop(unsafe { Box::from_raw(source as *mut RangeSource) }); +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_partitions( + mut env: JNIEnv, + _class: JClass, + reader: jlong, +) -> jobjectArray { + call(&mut env, std::ptr::null_mut(), |env| { + let end = unsafe { &*(reader as *const RangeReader) }.end; + let bounds = [(0, end / 2), (end / 2, end)]; + let partitions = env.new_object_array(bounds.len() as i32, "[B", JObject::null())?; + for (i, (start, end)) in bounds.into_iter().enumerate() { + let bytes = [start.to_le_bytes(), end.to_le_bytes()].concat(); + let partition = env.byte_array_from_slice(&bytes)?; + env.set_object_array_element(&partitions, i as i32, partition)?; + } + Ok(partitions.into_raw()) + }) +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_closeReader( + _env: JNIEnv, + _class: JClass, + reader: jlong, +) { + drop(unsafe { Box::from_raw(reader as *mut RangeReader) }); +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_read( + mut env: JNIEnv, + _class: JClass, + _reader_state: JByteArray, + partition: JByteArray, + stream_address: jlong, +) { + call(&mut env, (), |env| { + let bytes = env.convert_byte_array(&partition)?; + let start = i64::from_le_bytes(bytes[0..8].try_into()?); + let end = i64::from_le_bytes(bytes[8..16].try_into()?); + let ids = Int64Array::from_iter_values(start..end); + let batch = RecordBatch::try_new(schema(), vec![Arc::new(ids)])?; + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema()); + let stream = FFI_ArrowArrayStream::new(Box::new(batches)); + unsafe { std::ptr::write(stream_address as *mut FFI_ArrowArrayStream, stream) }; + Ok(()) + }) +} +``` + +Build the library, and package it: + +```bash +cargo build --release +mkdir -p package/osx-aarch64 +cp target/release/librust_range.dylib package/osx-aarch64/ +cat > package/spark-native-datasource.json < + +#include +#include +#include +#include +#include + +#include +#include + +#define SPARK_JNI(ret, method) \ + extern "C" JNIEXPORT ret JNICALL Java_org_apache_spark_sql_datasource_NativeBridge_##method + +namespace { + +struct RangeSource { + int64_t end = 10; +}; + +struct RangeReader { + int64_t end; +}; + +std::shared_ptr Schema() { + return arrow::schema({arrow::field("id", arrow::int64(), /*nullable=*/false)}); +} + +// Runs `f`, and reports its error as a Java exception: no C++ exception may cross into the JVM. +template +auto Call(JNIEnv* env, F&& f) -> decltype(f()) { + try { + return f(); + } catch (const std::exception& e) { + env->ThrowNew(env->FindClass("java/lang/RuntimeException"), e.what()); + if constexpr (!std::is_void_v) return {}; + } +} + +void Check(const arrow::Status& status) { + if (!status.ok()) throw std::runtime_error(status.ToString()); +} + +std::string ToString(JNIEnv* env, jobject value) { + const char* chars = env->GetStringUTFChars(static_cast(value), nullptr); + std::string result(chars); + env->ReleaseStringUTFChars(static_cast(value), chars); + return result; +} + +} // namespace + +SPARK_JNI(jint, abiVersion)(JNIEnv*, jclass) { return 1; } + +SPARK_JNI(jlong, createDataSource)(JNIEnv* env, jclass, jstring, jobjectArray keys, + jobjectArray values) { + return Call(env, [&] { + auto source = std::make_unique(); + for (jsize i = 0; i < env->GetArrayLength(keys); i++) { + if (ToString(env, env->GetObjectArrayElement(keys, i)) == "end") { + source->end = std::stoll(ToString(env, env->GetObjectArrayElement(values, i))); + } + } + return reinterpret_cast(source.release()); + }); +} + +SPARK_JNI(void, schema)(JNIEnv* env, jclass, jlong, jlong schema_address) { + Call(env, [&] { + Check(arrow::ExportSchema(*Schema(), reinterpret_cast(schema_address))); + }); +} + +SPARK_JNI(jlong, createReader)(JNIEnv* env, jclass, jlong source, jlong) { + return Call(env, [&] { + auto* reader = new RangeReader{reinterpret_cast(source)->end}; + return reinterpret_cast(reader); + }); +} + +SPARK_JNI(void, closeDataSource)(JNIEnv*, jclass, jlong source) { + delete reinterpret_cast(source); +} + +SPARK_JNI(jobjectArray, partitions)(JNIEnv* env, jclass, jlong reader) { + return Call(env, [&] { + int64_t end = reinterpret_cast(reader)->end; + int64_t bounds[2][2] = {{0, end / 2}, {end / 2, end}}; + jobjectArray partitions = env->NewObjectArray(2, env->FindClass("[B"), nullptr); + for (jsize i = 0; i < 2; i++) { + jbyteArray partition = env->NewByteArray(sizeof(bounds[i])); + env->SetByteArrayRegion(partition, 0, sizeof(bounds[i]), + reinterpret_cast(bounds[i])); + env->SetObjectArrayElement(partitions, i, partition); + } + return partitions; + }); +} + +SPARK_JNI(void, closeReader)(JNIEnv*, jclass, jlong reader) { + delete reinterpret_cast(reader); +} + +SPARK_JNI(void, read)(JNIEnv* env, jclass, jbyteArray, jbyteArray partition, + jlong stream_address) { + Call(env, [&] { + int64_t bounds[2]; + env->GetByteArrayRegion(partition, 0, sizeof(bounds), reinterpret_cast(bounds)); + arrow::Int64Builder ids; + for (int64_t id = bounds[0]; id < bounds[1]; id++) Check(ids.Append(id)); + std::shared_ptr array; + Check(ids.Finish(&array)); + auto batch = arrow::RecordBatch::Make(Schema(), array->length(), {array}); + auto batches = arrow::RecordBatchReader::Make({batch}, Schema()).ValueOrDie(); + Check(arrow::ExportRecordBatchReader(batches, + reinterpret_cast(stream_address))); + }); +} +``` + +Build the library, here on macOS with Arrow C++ 24, which requires C++20, and package it as above: + +```bash +c++ -std=c++20 -O2 -shared -fPIC \ + -I"$JAVA_HOME/include" -I"$JAVA_HOME/include/darwin" \ + -I"$ARROW_HOME/include" -L"$ARROW_HOME/lib" -larrow \ + -o libcpp_range.dylib cpp_range.cc +``` + +A library loaded by Spark must find its dependencies, such as `libarrow`, on every node: link them +statically, or install them on the nodes. + +### Configuration + +| Property Name | Default | Meaning | +|---------------|---------|---------| +| `spark.sql.dataSource.native.enabled` | true | Whether Spark loads native data sources. It is a static configuration, so it can only be set when the Spark application starts. | +| `spark.sql.dataSource.native.paths` | (none) | Comma-separated list of native data source packages, and of directories that contain them, on the local file system of the driver. | + +### Security + +A native data source runs native code in the driver and executor processes, with their +privileges, and is not sandboxed. Only add packages that you trust. Set +`spark.sql.dataSource.native.enabled` to false to prevent loading native data sources, for example +on a server shared by several users. diff --git a/docs/sql-v2-data-sources.md b/docs/sql-v2-data-sources.md index 4b4d0f35cb4e4..68ea6298e732a 100644 --- a/docs/sql-v2-data-sources.md +++ b/docs/sql-v2-data-sources.md @@ -40,7 +40,7 @@ Compared to the earlier Data Source V1 API, DSV2 offers: | Feature | Description | |---------|-------------| -| **Java API** | The connector interfaces are pure Java (`org.apache.spark.sql.connector`), removing the Scala dependency that DSV1 required. A [Python Data Source API](api/python/reference/pyspark.sql/api/pyspark.sql.datasource.DataSource.html) (`pyspark.sql.datasource`) is also available as a wrapper for lightweight connectors written entirely in Python. | +| **Java API** | The connector interfaces are pure Java (`org.apache.spark.sql.connector`), removing the Scala dependency that DSV1 required. A [Python Data Source API](api/python/reference/pyspark.sql/api/pyspark.sql.datasource.DataSource.html) (`pyspark.sql.datasource`) is also available as a wrapper for lightweight connectors written entirely in Python, and the [Columnar Data Source API](sql-columnar-data-sources.html) (`org.apache.spark.sql.datasource`) for lightweight columnar connectors written in Java, Scala, or native code such as Rust or C++. | | **Catalog integration** | Connectors can expose namespaces, tables, views, and functions natively through Spark SQL. | | **Operator pushdown** | Connectors can accept pushed-down filters, required columns, aggregates, limits, offsets, and more. | | **Report partitioning and ordering** | Connectors can report the physical layout of data so that Spark can avoid unnecessary shuffles and sorts. | @@ -473,5 +473,7 @@ A `Table` that declares `STREAMING_WRITE` supports streaming writes through - [Data Sources](sql-data-sources.html) for the user-facing guide to built-in data sources (DSV1). - [Python Data Source API](api/python/reference/pyspark.sql/api/pyspark.sql.datasource.DataSource.html) for writing lightweight connectors entirely in Python. +- [Columnar Data Source API](sql-columnar-data-sources.html) for writing lightweight columnar + connectors in Java, Scala, or native code such as Rust or C++. - [Storage Partition Join](sql-performance-tuning.html#storage-partition-join) for how DSV2 partitioning reporting enables join optimizations. diff --git a/pom.xml b/pom.xml index da05308883570..c2eb3de50b8cd 100644 --- a/pom.xml +++ b/pom.xml @@ -2570,6 +2570,25 @@ + + org.apache.arrow + arrow-c-data + ${arrow.version} + + + com.fasterxml.jackson.core + jackson-annotations + + + com.fasterxml.jackson.core + jackson-core + + + io.netty + netty-common + + + org.apache.arrow arrow-compression diff --git a/python/pyspark/sql/connect/client/artifact.py b/python/pyspark/sql/connect/client/artifact.py index 3ab3df4ff5938..f39362a0e153e 100644 --- a/python/pyspark/sql/connect/client/artifact.py +++ b/python/pyspark/sql/connect/client/artifact.py @@ -237,7 +237,8 @@ def _parse_artifacts( name = f"{name}#{parsed.fragment}" artifact = new_archive_artifact(name, LocalFile(local_path)) - elif file: + elif file or name.endswith(".sparkpkg"): + # Native data source packages are added as files. artifact = new_file_artifact(name, LocalFile(local_path)) elif name.endswith(".jar"): artifact = new_jar_artifact(name, LocalFile(local_path)) diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index 984b02ef6e562..3ba7fc1c0977b 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -2267,10 +2267,14 @@ def addArtifacts( .. versionchanged:: 4.0.0 Supports Spark Classic. + .. versionchanged:: 4.4.0 + Supports native data source packages (``.sparkpkg`` files). + Parameters ---------- *path : tuple of str - Artifact's URIs to add. + Artifact's URIs to add. A native data source package (a ``.sparkpkg`` file) makes + the native data sources it contains available to the session. pyfile : bool Whether to add them as Python dependencies such as .py, .egg, .zip or .jar files. The pyfiles are directly inserted into the path when executing Python functions @@ -2292,6 +2296,11 @@ def addArtifacts( "condition": "True together", }, ) + if not (file or pyfile or archive) and all(p.endswith(".sparkpkg") for p in path): + # Native data source packages are artifacts of the session in the JVM. + for p in path: + self._jsparkSession.addArtifact(os.path.abspath(p)) + return for p in path: normalized_path = os.path.abspath(p) target_dir = os.path.join( diff --git a/python/pyspark/sql/tests/connect/client/test_artifact.py b/python/pyspark/sql/tests/connect/client/test_artifact.py index 6e5ed71745950..286f35d539a63 100644 --- a/python/pyspark/sql/tests/connect/client/test_artifact.py +++ b/python/pyspark/sql/tests/connect/client/test_artifact.py @@ -447,6 +447,21 @@ def test_cache_artifact(self): self.assertEqual(actualHash, expected_hash) self.assertEqual(self.artifact_manager.is_cached_artifact(expected_hash), True) + def test_native_data_source_package(self): + with tempfile.TemporaryDirectory(prefix="test_native_data_source_package") as d: + package_path = os.path.join(d, "my_source.sparkpkg") + with open(package_path, "wb") as f: + f.write(b"package") + requests = list( + self.artifact_manager._create_requests( + package_path, pyfile=False, archive=False, file=False + ) + ) + self.assertEqual(len(requests), 1) + artifacts = requests[0].batch.artifacts + self.assertEqual(len(artifacts), 1) + self.assertEqual(artifacts[0].name, os.path.join("files", "my_source.sparkpkg")) + def test_add_not_existing_artifact(self): with tempfile.TemporaryDirectory(prefix="test_add_not_existing_artifact") as d: with self.assertRaises(FileNotFoundError): diff --git a/python/pyspark/sql/tests/test_artifact.py b/python/pyspark/sql/tests/test_artifact.py index 408b33f1ed743..b6a40e3947042 100644 --- a/python/pyspark/sql/tests/test_artifact.py +++ b/python/pyspark/sql/tests/test_artifact.py @@ -105,6 +105,18 @@ def test_artifacts_cannot_be_overwritten(self): messageParameters={"normalized_path": pyfile_path}, ) + def test_add_native_data_source_package(self): + # A new session, so that the package does not affect the other tests. + spark = self.spark.newSession() + with tempfile.TemporaryDirectory(prefix="test_add_native_data_source_package") as d: + package_path = os.path.join(d, "my_source.sparkpkg") + with open(package_path, "wb") as f: + f.write(b"package") + spark.addArtifact(package_path) + packages = spark._jsparkSession.artifactManager().getNativeDataSourcePackages()._1() + self.assertEqual(packages.size(), 1) + self.assertEqual(packages.head().getName(), "my_source.sparkpkg") + def test_add_zipped_package(self): self.check_add_zipped_package(self.spark) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/Artifact.scala b/sql/api/src/main/scala/org/apache/spark/sql/Artifact.scala index 40c7e532c27db..e4fb52e928391 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/Artifact.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/Artifact.scala @@ -38,6 +38,10 @@ private[sql] object Artifact { val CLASS_PREFIX: Path = Paths.get("classes") val JAR_PREFIX: Path = Paths.get("jars") val CACHE_PREFIX: Path = Paths.get("cache") + val FILE_PREFIX: Path = Paths.get("files") + + /** The file extension of native data source packages, which are added as file artifacts. */ + val NATIVE_DATA_SOURCE_PACKAGE_EXTENSION = ".sparkpkg" def newArtifactFromExtension( fileName: String, @@ -48,13 +52,15 @@ private[sql] object Artifact { newJarArtifact(targetFilePath, storage) case cf if cf.endsWith(".class") => newClassArtifact(targetFilePath, storage) + case pkg if pkg.endsWith(NATIVE_DATA_SOURCE_PACKAGE_EXTENSION) => + newFileArtifact(targetFilePath, storage) case other => throw new UnsupportedOperationException(s"Unsupported file format: $other") } } def parseArtifacts(uri: URI): Seq[Artifact] = { - // Currently only local files with extensions .jar and .class are supported. + // Currently only local files with extensions .jar, .class and .sparkpkg are supported. uri.getScheme match { case "file" => val path = Paths.get(uri) @@ -80,6 +86,10 @@ private[sql] object Artifact { newArtifact(CLASS_PREFIX, ".class", targetFilePath, storage) } + def newFileArtifact(targetFilePath: Path, storage: LocalData): Artifact = { + newArtifact(FILE_PREFIX, "", targetFilePath, storage) + } + def newCacheArtifact(id: String, storage: LocalData): Artifact = { newArtifact(CACHE_PREFIX, "", Paths.get(id), storage) } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala index c16a7d20dfacf..7c82867a2c3f1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala @@ -2020,6 +2020,12 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat cause = Some(cause)) } + def nativeDataSourcePackageConflictError(name: String, packages: Seq[String]): Throwable = { + new AnalysisException( + errorClass = "NATIVE_DATA_SOURCE_PACKAGE_CONFLICT", + messageParameters = Map("name" -> name, "packages" -> packages.mkString(", "))) + } + def findMultipleDataSourceError(provider: String, sourceNames: Seq[String]): Throwable = { new AnalysisException( errorClass = "MULTIPLE_DATA_SOURCES", diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala index fd44b835a0ace..95e4757fe8c60 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala @@ -802,6 +802,32 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE cause = error) } + def nativeDataSourceError( + action: String, name: String, msg: String, cause: Throwable): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "NATIVE_DATA_SOURCE_ERROR", + messageParameters = Map("action" -> action, "name" -> name, "msg" -> msg), + cause = cause) + } + + def invalidNativeDataSourcePackageError( + path: String, + subClass: String, + messageParameters: Map[String, String], + cause: Throwable = null): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = s"INVALID_NATIVE_DATA_SOURCE_PACKAGE.$subClass", + messageParameters = Map("path" -> path) ++ messageParameters, + cause = cause) + } + + def nativeDataSourcePackageNotFoundError( + fileName: String, checksum: String): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND", + messageParameters = Map("fileName" -> fileName, "checksum" -> checksum)) + } + def removedClassInSpark2Error(className: String, e: Throwable): SparkClassNotFoundException = { new SparkClassNotFoundException( errorClass = "_LEGACY_ERROR_TEMP_2052", diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index d3211aba4986f..b068dfce5a1d4 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -7207,6 +7207,18 @@ object SQLConf { .booleanConf .createWithDefault(false) + val NATIVE_DATA_SOURCE_PATHS = buildConf("spark.sql.dataSource.native.paths") + .doc("Comma-separated list of native data source packages (.sparkpkg files), and of " + + "directories that contain them, on the local file system of the driver. Spark adds a " + + "package to the artifacts of the session when one of its data sources is used, so that " + + "the executors get it too. Packages can also be added with spark.addArtifact. See " + + "spark.sql.dataSource.native.enabled.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .stringConf + .toSequence + .createWithDefault(Nil) + val CSV_FILTER_PUSHDOWN_ENABLED = buildConf("spark.sql.csv.filterPushdown.enabled") .doc("When true, enable filter pushdown to CSV datasource.") .version("3.0.0") diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala index e02c7296f9378..8b701d386ced5 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala @@ -464,6 +464,18 @@ object StaticSQLConf { .checkValue(_ >= 0, "The maximum total size must not be negative.") .createWithDefault(128 * 1024) // 128 KiB + val NATIVE_DATA_SOURCE_ENABLED = + buildStaticConf("spark.sql.dataSource.native.enabled") + .doc("When true, Spark finds and loads native data sources: data sources implemented in " + + "native code such as Rust or C++, and distributed as native data source packages " + + "(.sparkpkg files). A package is found when it is added with spark.addArtifact, or when " + + "it is under spark.sql.dataSource.native.paths. Loading a native data source runs its " + + "native code in the driver and executor processes.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(true) + val KAFKA_DISALLOWED_OPTIONS = buildStaticConf("spark.sql.kafka.disallowedOptions") .internal() diff --git a/sql/core/pom.xml b/sql/core/pom.xml index 281546b901a7a..7f3db15f70fdf 100644 --- a/sql/core/pom.xml +++ b/sql/core/pom.xml @@ -298,6 +298,10 @@ org.apache.arrow arrow-compression + + org.apache.arrow + arrow-c-data + org.apache.spark spark-udf-worker-proto_${scala.binary.version} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSource.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSource.java new file mode 100644 index 0000000000000..005f645543903 --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSource.java @@ -0,0 +1,102 @@ +/* + * 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; + +import java.util.Map; + +import org.apache.spark.SparkUnsupportedOperationException; +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.types.StructType; + +/** + * A data source implemented with the columnar data source API. + *

+ * The columnar data source API is a simpler way to plug a data source into Spark than + * implementing the Data Source V2 interfaces directly. It follows the Python Data Source API + * ({@code pyspark.sql.datasource}), but data is exchanged only as columnar batches: a source + * returns {@link org.apache.spark.sql.vectorized.ColumnarBatch}es when reading and receives them + * when writing. Spark adapts the source to Data Source V2, so it supports batch and streaming + * reads and writes, and filter, column and limit pushdown. + *

+ * A data source is implemented either on the JVM, registered through a + * {@link DataSourceProvider}, or in native code such as Rust or C++, through the + * {@link NativeBridge} binary interface. + *

+ * Spark creates a new instance on the driver for each schema inference, scan and write. A data + * source only needs to implement the operations it supports; the default implementations throw an + * error saying that the operation is not supported. + * + * @since 4.4.0 + */ +@Evolving +public interface DataSource { + + /** + * Returns the schema of the data. Spark calls it only when the user does not specify a schema. + */ + default StructType schema() { + throw new SparkUnsupportedOperationException( + "UNABLE_TO_INFER_SCHEMA", Map.of("format", getClass().getName())); + } + + /** + * Returns a reader for a batch scan. + * + * @param schema the schema to read: the one the user specified, or the one returned by + * {@link #schema()} + */ + default DataSourceReader reader(StructType schema) { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map.of("description", getClass().getName())); + } + + /** + * Returns a reader for a micro-batch streaming scan. + * + * @param schema the schema to read: the one the user specified, or the one returned by + * {@link #schema()} + */ + default DataSourceStreamReader streamReader(StructType schema) { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", Map.of("description", getClass().getName())); + } + + /** + * Returns a writer for a batch write. + * + * @param schema the schema of the data to write + * @param overwrite whether to replace the existing data (save mode "overwrite") instead of + * appending to it + */ + default DataSourceWriter writer(StructType schema, boolean overwrite) { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", Map.of("description", getClass().getName())); + } + + /** + * Returns a writer for a streaming write. + * + * @param schema the schema of the data to write + * @param overwrite whether to replace the existing data in each micro-batch (output mode + * "complete") instead of appending to it + */ + default DataSourceStreamWriter streamWriter(StructType schema, boolean overwrite) { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", Map.of("description", getClass().getName())); + } +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceProvider.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceProvider.java new file mode 100644 index 0000000000000..4cfd25e351b4a --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceProvider.java @@ -0,0 +1,65 @@ +/* + * 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; + +import java.util.Map; + +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.connector.catalog.Table; +import org.apache.spark.sql.connector.catalog.TableProvider; +import org.apache.spark.sql.connector.expressions.Transform; +import org.apache.spark.sql.execution.datasources.v2.columnar.ColumnarTable; +import org.apache.spark.sql.sources.DataSourceRegister; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; + +/** + * The entry point of a {@link DataSource} implemented on the JVM. + *

+ * Register a provider like any other data source: list the class in + * {@code META-INF/services/org.apache.spark.sql.sources.DataSourceRegister} to use it by its + * {@link #shortName()}, for example {@code spark.read.format("my_source")}, or use its fully + * qualified class name as the format. + * + * @since 4.4.0 + */ +@Evolving +public abstract class DataSourceProvider implements TableProvider, DataSourceRegister { + + /** + * Creates the data source with the options given by the user, such as + * {@code spark.read.option("key", "value")}. + */ + public abstract DataSource createDataSource(CaseInsensitiveStringMap options); + + @Override + public final StructType inferSchema(CaseInsensitiveStringMap options) { + return createDataSource(options).schema(); + } + + @Override + public final Table getTable( + StructType schema, Transform[] partitioning, Map properties) { + return new ColumnarTable(shortName(), schema, this::createDataSource); + } + + @Override + public final boolean supportsExternalMetadata() { + return true; + } +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceReader.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceReader.java new file mode 100644 index 0000000000000..6a04e003e8047 --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceReader.java @@ -0,0 +1,89 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.datasource; + +import java.io.Serializable; + +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.connector.expressions.filter.Predicate; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.PartitionReader; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.vectorized.ColumnarBatch; + +/** + * Reads the data of a {@link DataSource} for a batch scan. + *

+ * Spark creates the reader on the driver with {@link DataSource#reader}. It then pushes + * operations down by calling, each at most once and in this order, {@link #pushPredicates}, + * {@link #pushLimit} and {@link #pruneColumns}, and plans the scan with {@link #partitions()}. + * After that, the reader is serialized and sent to the executors, where {@link #read} is called + * once for each partition. Any state that {@link #read} needs must be set by the time + * {@link #partitions()} returns. + * + * @since 4.4.0 + */ +@Evolving +public interface DataSourceReader extends Serializable { + + /** + * Pushes predicates down to the data source. The predicates are combined with AND. + * + * @return the predicates that Spark still has to evaluate after reading. By default, all of + * them. The data source can still use a returned predicate to skip data. + */ + default Predicate[] pushPredicates(Predicate[] predicates) { + return predicates; + } + + /** + * Pushes a LIMIT down to the data source. + * + * @return whether the reader returns at most {@code limit} rows in total. Spark applies the + * limit again after reading either way. + */ + default boolean pushLimit(int limit) { + return false; + } + + /** + * Prunes the columns to read. + * + * @param requiredSchema the top-level fields of the read schema that the query needs, with + * their full types + * @return whether {@link #read} returns exactly the columns of {@code requiredSchema}, in that + * order. If false, {@link #read} returns all the columns of the read schema. + */ + default boolean pruneColumns(StructType requiredSchema) { + return false; + } + + /** + * Plans the partitions of the scan. Spark reads each partition in a separate task. + */ + InputPartition[] partitions(); + + /** + * Reads a partition. Called on an executor. + *

+ * The columns of the returned batches must match the read schema: the schema passed to + * {@link DataSource#reader}, or the required schema if {@link #pruneColumns} returned true. A + * batch only needs to stay valid until the next call to {@link PartitionReader#next()}. + */ + PartitionReader read(InputPartition partition); +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamReader.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamReader.java new file mode 100644 index 0000000000000..ba043134a2e6a --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamReader.java @@ -0,0 +1,73 @@ +/* + * 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; + +import java.io.Serializable; + +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.PartitionReader; +import org.apache.spark.sql.vectorized.ColumnarBatch; + +/** + * Reads the data of a {@link DataSource} for a micro-batch streaming scan. + *

+ * Spark creates the reader on the driver with {@link DataSource#streamReader} and keeps it for + * the lifetime of the streaming query. Offsets are JSON strings defined by the data source; Spark + * stores them in the checkpoint and passes them back unchanged. For each micro-batch, Spark plans + * the partitions between two offsets with {@link #partitions}, then serializes the reader and + * sends it to the executors, where {@link #read} is called once for each partition. + * + * @since 4.4.0 + */ +@Evolving +public interface DataSourceStreamReader extends Serializable { + + /** + * Returns the offset to start from when the query has no checkpointed offset. + */ + String initialOffset(); + + /** + * Returns the most recent offset available. + */ + String latestOffset(); + + /** + * Plans the partitions that read the data after {@code start} up to and including {@code end}. + */ + InputPartition[] partitions(String start, String end); + + /** + * Reads a partition. Called on an executor. The columns of the returned batches must match the + * schema passed to {@link DataSource#streamReader}. A batch only needs to stay valid until the + * next call to {@link PartitionReader#next()}. + */ + PartitionReader read(InputPartition partition); + + /** + * Tells the data source that Spark has processed all the data up to and including {@code end} + * and will not ask for it again. + */ + default void commit(String end) {} + + /** + * Stops the reader when the streaming query terminates. + */ + default void stop() {} +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamWriter.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamWriter.java new file mode 100644 index 0000000000000..f18ba7bae2e50 --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamWriter.java @@ -0,0 +1,69 @@ +/* + * 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; + +import java.io.Serializable; + +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.connector.write.DataWriter; +import org.apache.spark.sql.connector.write.WriterCommitMessage; +import org.apache.spark.sql.vectorized.ColumnarBatch; + +/** + * Writes data to a {@link DataSource} for a streaming write. + *

+ * Spark creates the writer on the driver with {@link DataSource#streamWriter} and keeps it for + * the lifetime of the streaming query. For each micro-batch (epoch), it serializes the writer and + * sends it to the executors, where {@link #createWriter} is called once for each task. When all + * the tasks of the micro-batch succeed, Spark calls {@link #commit} on the driver with the + * messages returned by the task writers. Otherwise, it calls {@link #abort}. + *

+ * As with {@link DataSourceWriter}, the batches passed to {@link DataWriter#write} are backed by + * Apache Arrow and are only valid during the call. + * + * @since 4.4.0 + */ +@Evolving +public interface DataSourceStreamWriter extends Serializable { + + /** + * Creates the writer for one task of a micro-batch. Called on an executor. + * + * @param partitionId the partition the task writes + * @param taskId the ID of the task attempt + * @param epochId the ID of the micro-batch + */ + DataWriter createWriter(int partitionId, long taskId, long epochId); + + /** + * Commits a micro-batch. Called on the driver when all the tasks of the micro-batch succeeded. + * + * @param epochId the ID of the micro-batch + * @param messages the messages returned by {@link DataWriter#commit()} of each task + */ + void commit(long epochId, WriterCommitMessage[] messages); + + /** + * Aborts a micro-batch. Called on the driver when the micro-batch failed. + * + * @param epochId the ID of the micro-batch + * @param messages the messages returned by the tasks that committed. An element is null if its + * task did not commit. + */ + void abort(long epochId, WriterCommitMessage[] messages); +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceWriter.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceWriter.java new file mode 100644 index 0000000000000..ce1cc320a22d0 --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceWriter.java @@ -0,0 +1,66 @@ +/* + * 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; + +import java.io.Serializable; + +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.connector.write.DataWriter; +import org.apache.spark.sql.connector.write.WriterCommitMessage; +import org.apache.spark.sql.vectorized.ColumnarBatch; + +/** + * Writes data to a {@link DataSource} for a batch write. + *

+ * Spark creates the writer on the driver with {@link DataSource#writer}, then serializes it and + * sends it to the executors, where {@link #createWriter} is called once for each task. When all + * the tasks succeed, Spark calls {@link #commit} on the driver with the messages returned by the + * task writers. Otherwise, it calls {@link #abort}. + *

+ * The batches passed to {@link DataWriter#write} are backed by Apache Arrow: every column is an + * {@link org.apache.spark.sql.vectorized.ArrowColumnVector}. A batch is only valid during the + * call, so a writer that keeps data after the call returns must copy it. + * + * @since 4.4.0 + */ +@Evolving +public interface DataSourceWriter extends Serializable { + + /** + * Creates the writer for one task. Called on an executor. + * + * @param partitionId the partition the task writes + * @param taskId the ID of the task attempt + */ + DataWriter createWriter(int partitionId, long taskId); + + /** + * Commits the write. Called on the driver when all the tasks succeeded. + * + * @param messages the messages returned by {@link DataWriter#commit()} of each task + */ + void commit(WriterCommitMessage[] messages); + + /** + * Aborts the write. Called on the driver when the write failed. + * + * @param messages the messages returned by the tasks that committed. An element is null if its + * task did not commit. + */ + void abort(WriterCommitMessage[] messages); +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java new file mode 100644 index 0000000000000..5570f6593f309 --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java @@ -0,0 +1,321 @@ +/* + * 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; + +import org.apache.spark.annotation.Evolving; + +/** + * The binary interface between Spark and data sources implemented in native code, such as Rust + * or C++. + *

+ * A native data source is a shared library that implements the {@code native} methods of this + * class as JNI functions, for example + * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each + * library into its own copy of this class, so any number of libraries that export the same + * functions can be loaded side by side. The methods follow the {@link DataSource} API: Spark + * implements {@link DataSource}, {@link DataSourceReader} and the other interfaces by calling + * them. + *

+ * The library is distributed in a native data source package: a zip file with the extension + * {@code .sparkpkg} that contains a manifest named {@code spark-native-datasource.json} and the + * library built for one or more platforms. Spark finds the packages added to a session with + * {@code spark.addArtifact}, and the ones under the paths in + * {@code spark.sql.dataSource.native.paths}. See the "Columnar Data Source API" page of the Spark + * SQL guide for details. + * + *

Conventions

+ *
    + *
  • Handles. The {@code create} methods return an opaque, non-zero handle, typically a + * pointer. Spark releases each handle exactly once with the matching {@code close} method, or + * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers. A reader or + * writer handle must not depend on the data source handle it was created from, because Spark + * closes the data source handle right after creating it.
    • + *
  • Errors. To report an error, throw a Java exception, for example with the JNI + * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and return any value. Spark + * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.
    • + *
  • Arrow data. Schemas and data are exchanged with the + * Arrow C data + * interface. A {@code long} argument named {@code ...Address} is the address of a C struct + * allocated by Spark. For an output struct, the library initializes it and Spark takes + * ownership. For an input struct, the library may move it to take ownership; otherwise, Spark + * releases it after the call returns, so the library must not use it afterwards.
    • + *
  • Schemas. Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: + * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, + * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is + * {@code timestamp[us]} without one. Data read from a library may also use + * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, but not + * unsigned integers or dictionary encoding.
    • + *
  • Options. Option keys are passed in lower case, because Spark matches them + * case-insensitively.
    • + *
  • Optional methods. A library only needs to export the functions for the operations + * it supports. If a function is missing, Spark treats the operation as unsupported, except for + * the pushdown methods and the {@code serialize} methods, which then push nothing down and + * return no state.
    • + *
  • Threads. Spark calls the library from multiple threads, and calls on different + * handles may run concurrently. Calls on the same handle never run concurrently.
    • + *
+ * + * @since 4.4.0 + */ +@Evolving +public final class NativeBridge { + + /** The version of this interface. It changes when the interface changes incompatibly. */ + public static final int ABI_VERSION = 1; + + private NativeBridge() {} + + /** + * Loads the library that implements the native methods of this copy of the class. Spark calls + * it on the copy of the class it creates for the library. + */ + static void load(String path) { + System.load(path); + } + + /** Returns the version of this interface that the library implements, {@link #ABI_VERSION}. */ + public static native int abiVersion(); + + // --------------------------------------------------------------------------------------------- + // Data source (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Creates a data source. See {@link DataSourceProvider#createDataSource}. + * + * @param name the name of the data source, as listed in the package manifest. A library can + * implement several data sources. + * @param optionKeys the option keys, in lower case + * @param optionValues the option values, in the same order as the keys + * @return the data source handle + */ + public static native long createDataSource( + String name, String[] optionKeys, String[] optionValues); + + /** + * Exports the schema of the data source. See {@link DataSource#schema()}. + * + * @param schemaAddress the address of the {@code ArrowSchema} to initialize with a struct type + * whose fields are the columns + */ + public static native void schema(long dataSource, long schemaAddress); + + /** + * Creates a reader for a batch scan. See {@link DataSource#reader}. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read + * @return the reader handle + */ + public static native long createReader(long dataSource, long schemaAddress); + + /** + * Creates a reader for a micro-batch streaming scan. See {@link DataSource#streamReader}. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read + * @return the stream reader handle + */ + public static native long createStreamReader(long dataSource, long schemaAddress); + + /** + * Creates a writer for a batch write. See {@link DataSource#writer}. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @return the writer handle + */ + public static native long createWriter(long dataSource, long schemaAddress, boolean overwrite); + + /** + * Creates a writer for a streaming write. See {@link DataSource#streamWriter}. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @return the writer handle + */ + public static native long createStreamWriter( + long dataSource, long schemaAddress, boolean overwrite); + + /** Releases a data source handle. */ + public static native void closeDataSource(long dataSource); + + // --------------------------------------------------------------------------------------------- + // Batch reader (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Pushes predicates down to the reader. See {@link DataSourceReader#pushPredicates}. + *

+ * Each predicate is a JSON expression tree. An expression is one of: + *

    + *
  • {@code {"type": "column", "name": ["a", "b"]}}: a column. A nested field has a name + * part for each level.
    • + *
  • {@code {"type": "literal", "dataType": "bigint", "value": 1}}: a literal. The data type + * is a Spark SQL type name such as {@code int}, {@code string} or {@code decimal(10,2)}. The + * value is {@code null}, a JSON boolean or number for boolean and numeric types, a string + * for {@code string}, a base64 string for {@code binary}, a decimal string for + * {@code decimal}, the days since the epoch for {@code date}, and the microseconds since the + * epoch for {@code timestamp} and {@code timestamp_ntz}.
    • + *
  • {@code {"type": "function", "name": ">", "children": [...]}}: a predicate or function, + * with the names of {@code org.apache.spark.sql.connector.expressions.filter.Predicate} and + * {@code org.apache.spark.sql.connector.expressions.GeneralScalarExpression}, such as + * {@code =}, {@code <}, {@code IN}, {@code IS_NULL}, {@code STARTS_WITH}, {@code AND}, + * {@code OR} and {@code NOT}.
    • + *
+ * Spark does not pass predicates that it cannot express this way. + * + * @return for each predicate, whether the reader evaluates it completely, so that Spark does + * not have to evaluate it again + */ + public static native boolean[] pushPredicates(long reader, String[] predicates); + + /** + * Pushes a LIMIT down to the reader. See {@link DataSourceReader#pushLimit}. + */ + public static native boolean pushLimit(long reader, int limit); + + /** + * Prunes the columns to read. See {@link DataSourceReader#pruneColumns}. + * + * @param columnNames the names of the top-level columns to read + * @return whether {@link #read} returns exactly these columns, in this order + */ + public static native boolean pruneColumns(long reader, String[] columnNames); + + /** + * Plans the partitions of the scan. See {@link DataSourceReader#partitions()}. + * + * @return the partitions, each serialized in a format defined by the library + */ + public static native byte[][] partitions(long reader); + + /** + * Serializes the state of the reader that {@link #read} needs on the executors, such as the + * pushed predicates. Spark calls it after {@link #partitions}. + */ + public static native byte[] serializeReader(long reader); + + /** Releases a batch reader handle. */ + public static native void closeReader(long reader); + + // --------------------------------------------------------------------------------------------- + // Stream reader (driver) + // --------------------------------------------------------------------------------------------- + + /** Returns the initial offset as JSON. See {@link DataSourceStreamReader#initialOffset()}. */ + public static native String initialOffset(long streamReader); + + /** Returns the latest offset as JSON. See {@link DataSourceStreamReader#latestOffset()}. */ + public static native String latestOffset(long streamReader); + + /** + * Plans the partitions between two offsets. See {@link DataSourceStreamReader#partitions}. + * + * @return the partitions, each serialized in a format defined by the library + */ + public static native byte[][] streamPartitions(long streamReader, String start, String end); + + /** + * Serializes the state of the stream reader that {@link #read} needs on the executors. Spark + * calls it after each call to {@link #streamPartitions}. + */ + public static native byte[] serializeStreamReader(long streamReader); + + /** Commits an offset. See {@link DataSourceStreamReader#commit}. */ + public static native void commitOffset(long streamReader, String end); + + /** Releases a stream reader handle. See {@link DataSourceStreamReader#stop()}. */ + public static native void closeStreamReader(long streamReader); + + // --------------------------------------------------------------------------------------------- + // Partition reader (executor) + // --------------------------------------------------------------------------------------------- + + /** + * Reads a partition of a batch or streaming scan. See {@link DataSourceReader#read}. + * + * @param readerState the state returned by {@link #serializeReader} or + * {@link #serializeStreamReader}, or an empty array if there is none + * @param partition a partition returned by {@link #partitions} or {@link #streamPartitions} + * @param streamAddress the address of the {@code ArrowArrayStream} to initialize. Each array + * of the stream is a struct array whose fields are the columns. Spark + * calls the stream from the thread of the task that reads the partition. + */ + public static native void read(byte[] readerState, byte[] partition, long streamAddress); + + // --------------------------------------------------------------------------------------------- + // Writer (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Serializes the state of the writer that {@link #createDataWriter} needs on the executors. + * Spark calls it once, after creating the writer. + */ + public static native byte[] serializeWriter(long writer); + + /** + * Commits a write. See {@link DataSourceWriter#commit} and {@link DataSourceStreamWriter#commit}. + * + * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write + * @param messages the messages returned by {@link #commitDataWriter} of each task + */ + public static native void commit(long writer, long epochId, byte[][] messages); + + /** + * Aborts a write. See {@link DataSourceWriter#abort} and {@link DataSourceStreamWriter#abort}. + * + * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write + * @param messages the messages returned by {@link #commitDataWriter} of each task. An element + * is null if its task did not commit. + */ + public static native void abort(long writer, long epochId, byte[][] messages); + + /** Releases a writer handle. */ + public static native void closeWriter(long writer); + + // --------------------------------------------------------------------------------------------- + // Data writer (executor) + // --------------------------------------------------------------------------------------------- + + /** + * Creates the writer for one task. See {@link DataSourceWriter#createWriter}. + * + * @param writerState the state returned by {@link #serializeWriter}, or an empty array if there + * is none + * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write + * @return the data writer handle + */ + public static native long createDataWriter( + byte[] writerState, int partitionId, long taskId, long epochId); + + /** + * Writes a batch of data. + * + * @param arrayAddress the address of the {@code ArrowArray}: a struct array whose fields are + * the columns + * @param schemaAddress the address of the {@code ArrowSchema} of the array + */ + public static native void write(long dataWriter, long arrayAddress, long schemaAddress); + + /** + * Commits the data written by the task and releases the data writer handle. + * + * @return the message to pass to {@link #commit} on the driver + */ + public static native byte[] commitDataWriter(long dataWriter); + + /** Aborts the data written by the task and releases the data writer handle. */ + public static native void abortDataWriter(long dataWriter); +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java new file mode 100644 index 0000000000000..3b2343ee142d8 --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -0,0 +1,292 @@ +/* + * 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.execution.datasources.v2.ffi; + +import java.io.IOException; +import java.io.InputStream; +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.MethodType; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; + +import org.apache.spark.sql.datasource.NativeBridge; + +/** + * A native data source library, loaded into its own copy of {@link NativeBridge}. + *

+ * JNI links the native methods of a class to the libraries loaded by the class loader of that + * class, and a library belongs to the class loader of the class that loads it. Defining a copy of + * {@link NativeBridge} in a separate class loader for each library therefore lets all the + * libraries export the same JNI functions. Spark calls the methods of a copy through method + * handles, because the copy is a different class than the {@link NativeBridge} it is compiled + * against. Calling a method that the library does not export throws + * {@link UnsatisfiedLinkError}. + */ +final class NativeLibrary { + + private static final String BRIDGE_CLASS = NativeBridge.class.getName(); + + private final MethodHandle abiVersion; + private final MethodHandle createDataSource; + private final MethodHandle schema; + private final MethodHandle createReader; + private final MethodHandle createStreamReader; + private final MethodHandle createWriter; + private final MethodHandle createStreamWriter; + private final MethodHandle closeDataSource; + private final MethodHandle pushPredicates; + private final MethodHandle pushLimit; + private final MethodHandle pruneColumns; + private final MethodHandle partitions; + private final MethodHandle serializeReader; + private final MethodHandle closeReader; + private final MethodHandle initialOffset; + private final MethodHandle latestOffset; + private final MethodHandle streamPartitions; + private final MethodHandle serializeStreamReader; + private final MethodHandle commitOffset; + private final MethodHandle closeStreamReader; + private final MethodHandle read; + private final MethodHandle serializeWriter; + private final MethodHandle commit; + private final MethodHandle abort; + private final MethodHandle closeWriter; + private final MethodHandle createDataWriter; + private final MethodHandle write; + private final MethodHandle commitDataWriter; + private final MethodHandle abortDataWriter; + + private NativeLibrary(Class bridge) throws ReflectiveOperationException { + Class jlong = long.class; + Class bytes = byte[].class; + Class strings = String[].class; + abiVersion = find(bridge, "abiVersion", int.class); + createDataSource = find(bridge, "createDataSource", jlong, String.class, strings, strings); + schema = find(bridge, "schema", void.class, jlong, jlong); + createReader = find(bridge, "createReader", jlong, jlong, jlong); + createStreamReader = find(bridge, "createStreamReader", jlong, jlong, jlong); + createWriter = find(bridge, "createWriter", jlong, jlong, jlong, boolean.class); + createStreamWriter = find(bridge, "createStreamWriter", jlong, jlong, jlong, boolean.class); + closeDataSource = find(bridge, "closeDataSource", void.class, jlong); + pushPredicates = find(bridge, "pushPredicates", boolean[].class, jlong, strings); + pushLimit = find(bridge, "pushLimit", boolean.class, jlong, int.class); + pruneColumns = find(bridge, "pruneColumns", boolean.class, jlong, strings); + partitions = find(bridge, "partitions", byte[][].class, jlong); + serializeReader = find(bridge, "serializeReader", bytes, jlong); + closeReader = find(bridge, "closeReader", void.class, jlong); + initialOffset = find(bridge, "initialOffset", String.class, jlong); + latestOffset = find(bridge, "latestOffset", String.class, jlong); + streamPartitions = + find(bridge, "streamPartitions", byte[][].class, jlong, String.class, String.class); + serializeStreamReader = find(bridge, "serializeStreamReader", bytes, jlong); + commitOffset = find(bridge, "commitOffset", void.class, jlong, String.class); + closeStreamReader = find(bridge, "closeStreamReader", void.class, jlong); + read = find(bridge, "read", void.class, bytes, bytes, jlong); + serializeWriter = find(bridge, "serializeWriter", bytes, jlong); + commit = find(bridge, "commit", void.class, jlong, jlong, byte[][].class); + abort = find(bridge, "abort", void.class, jlong, jlong, byte[][].class); + closeWriter = find(bridge, "closeWriter", void.class, jlong); + createDataWriter = find(bridge, "createDataWriter", jlong, bytes, int.class, jlong, jlong); + write = find(bridge, "write", void.class, jlong, jlong, jlong); + commitDataWriter = find(bridge, "commitDataWriter", bytes, jlong); + abortDataWriter = find(bridge, "abortDataWriter", void.class, jlong); + } + + private static MethodHandle find( + Class bridge, String name, Class returnType, Class... parameterTypes) + throws ReflectiveOperationException { + return MethodHandles.publicLookup() + .findStatic(bridge, name, MethodType.methodType(returnType, parameterTypes)); + } + + /** + * Loads the library at the given path into a new copy of {@link NativeBridge}. + */ + static NativeLibrary load(String path) throws Throwable { + Class bridge = new BridgeClassLoader(NativeBridge.class.getClassLoader()) + .loadClass(BRIDGE_CLASS); + Method load = bridge.getDeclaredMethod("load", String.class); + load.setAccessible(true); + try { + load.invoke(null, path); + } catch (InvocationTargetException e) { + throw e.getCause(); + } + return new NativeLibrary(bridge); + } + + /** + * A class loader that defines its own copy of {@link NativeBridge}, and delegates all the other + * classes to its parent. + */ + private static final class BridgeClassLoader extends ClassLoader { + static { + ClassLoader.registerAsParallelCapable(); + } + + BridgeClassLoader(ClassLoader parent) { + super(parent); + } + + @Override + protected Class loadClass(String name, boolean resolve) throws ClassNotFoundException { + if (!BRIDGE_CLASS.equals(name)) { + return super.loadClass(name, resolve); + } + synchronized (getClassLoadingLock(name)) { + Class loaded = findLoadedClass(name); + if (loaded == null) { + String resource = name.replace('.', '/') + ".class"; + try (InputStream in = getParent().getResourceAsStream(resource)) { + if (in == null) { + throw new ClassNotFoundException(name); + } + byte[] bytes = in.readAllBytes(); + loaded = defineClass(name, bytes, 0, bytes.length); + } catch (IOException e) { + throw new ClassNotFoundException(name, e); + } + } + if (resolve) { + resolveClass(loaded); + } + return loaded; + } + } + } + + int abiVersion() throws Throwable { + return (int) abiVersion.invokeExact(); + } + + long createDataSource(String name, String[] optionKeys, String[] optionValues) + throws Throwable { + return (long) createDataSource.invokeExact(name, optionKeys, optionValues); + } + + void schema(long dataSource, long schemaAddress) throws Throwable { + schema.invokeExact(dataSource, schemaAddress); + } + + long createReader(long dataSource, long schemaAddress) throws Throwable { + return (long) createReader.invokeExact(dataSource, schemaAddress); + } + + long createStreamReader(long dataSource, long schemaAddress) throws Throwable { + return (long) createStreamReader.invokeExact(dataSource, schemaAddress); + } + + long createWriter(long dataSource, long schemaAddress, boolean overwrite) throws Throwable { + return (long) createWriter.invokeExact(dataSource, schemaAddress, overwrite); + } + + long createStreamWriter(long dataSource, long schemaAddress, boolean overwrite) + throws Throwable { + return (long) createStreamWriter.invokeExact(dataSource, schemaAddress, overwrite); + } + + void closeDataSource(long dataSource) throws Throwable { + closeDataSource.invokeExact(dataSource); + } + + boolean[] pushPredicates(long reader, String[] predicates) throws Throwable { + return (boolean[]) pushPredicates.invokeExact(reader, predicates); + } + + boolean pushLimit(long reader, int limit) throws Throwable { + return (boolean) pushLimit.invokeExact(reader, limit); + } + + boolean pruneColumns(long reader, String[] columnNames) throws Throwable { + return (boolean) pruneColumns.invokeExact(reader, columnNames); + } + + byte[][] partitions(long reader) throws Throwable { + return (byte[][]) partitions.invokeExact(reader); + } + + byte[] serializeReader(long reader) throws Throwable { + return (byte[]) serializeReader.invokeExact(reader); + } + + void closeReader(long reader) throws Throwable { + closeReader.invokeExact(reader); + } + + String initialOffset(long streamReader) throws Throwable { + return (String) initialOffset.invokeExact(streamReader); + } + + String latestOffset(long streamReader) throws Throwable { + return (String) latestOffset.invokeExact(streamReader); + } + + byte[][] streamPartitions(long streamReader, String start, String end) throws Throwable { + return (byte[][]) streamPartitions.invokeExact(streamReader, start, end); + } + + byte[] serializeStreamReader(long streamReader) throws Throwable { + return (byte[]) serializeStreamReader.invokeExact(streamReader); + } + + void commitOffset(long streamReader, String end) throws Throwable { + commitOffset.invokeExact(streamReader, end); + } + + void closeStreamReader(long streamReader) throws Throwable { + closeStreamReader.invokeExact(streamReader); + } + + void read(byte[] readerState, byte[] partition, long streamAddress) throws Throwable { + read.invokeExact(readerState, partition, streamAddress); + } + + byte[] serializeWriter(long writer) throws Throwable { + return (byte[]) serializeWriter.invokeExact(writer); + } + + void commit(long writer, long epochId, byte[][] messages) throws Throwable { + commit.invokeExact(writer, epochId, messages); + } + + void abort(long writer, long epochId, byte[][] messages) throws Throwable { + abort.invokeExact(writer, epochId, messages); + } + + void closeWriter(long writer) throws Throwable { + closeWriter.invokeExact(writer); + } + + long createDataWriter(byte[] writerState, int partitionId, long taskId, long epochId) + throws Throwable { + return (long) createDataWriter.invokeExact(writerState, partitionId, taskId, epochId); + } + + void write(long dataWriter, long arrayAddress, long schemaAddress) throws Throwable { + write.invokeExact(dataWriter, arrayAddress, schemaAddress); + } + + byte[] commitDataWriter(long dataWriter) throws Throwable { + return (byte[]) commitDataWriter.invokeExact(dataWriter); + } + + void abortDataWriter(long dataWriter) throws Throwable { + abortDataWriter.invokeExact(dataWriter); + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala index 2a03a58502c60..281b2a079cc3d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala @@ -153,6 +153,19 @@ class ArtifactManager(session: SparkSession) extends AutoCloseable with Logging */ def getPythonIncludes: Seq[String] = pythonIncludeList.asScala.toSeq + /** + * Returns the native data source packages added to this session, as the local files of the + * driver, and the UUID under which the executors store them, if the session is isolated. + */ + private[sql] def getNativeDataSourcePackages: (Seq[File], Option[String]) = { + val packages = sparkContextRelativePaths.asScala.collect { + case (SparkContextResourceType.FILE, path, _) + if path.toString.endsWith(Artifact.NATIVE_DATA_SOURCE_PACKAGE_EXTENSION) => + artifactPath.resolve(path).toFile + }.toSeq + (packages, Option(state).map(_.uuid)) + } + protected[sql] def getCachedBlockId(hash: String): Option[CacheId] = { Option(hashToCachedIdMap.get(hash)).map(_.id) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala index f22008d3b2321..3183cc8c8c12a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala @@ -32,8 +32,7 @@ import org.apache.spark.sql.connector.catalog.TableCapability._ import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.command.DDLUtils import org.apache.spark.sql.execution.datasources.{DataSource, LogicalRelation} -import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Utils, FileDataSourceV2} -import org.apache.spark.sql.execution.datasources.v2.python.PythonDataSourceV2 +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Utils, FileDataSourceV2, NamedTableProvider} import org.apache.spark.sql.execution.streaming.runtime.StreamingRelation import org.apache.spark.sql.sources.StreamSourceProvider import org.apache.spark.sql.types.StructType @@ -159,7 +158,7 @@ class ResolveDataSource(sparkSession: SparkSession) extends Rule[LogicalPlan] { optionsWithPath.originalMap val dsOptions = new CaseInsensitiveStringMap(finalOptions.asJava) provider match { - case p: PythonDataSourceV2 => p.setShortName(source) + case p: NamedTableProvider => p.setShortName(source) case _ => } val table = diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala index 1301985e30158..1e7989e4cc6c5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala @@ -40,8 +40,7 @@ import org.apache.spark.sql.connector.expressions.{ClusterByTransform, FieldRefe import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.command.DDLUtils import org.apache.spark.sql.execution.datasources.{DataSource, DataSourceUtils} -import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Utils, FileDataSourceV2} -import org.apache.spark.sql.execution.datasources.v2.python.PythonDataSourceV2 +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Utils, FileDataSourceV2, NamedTableProvider} import org.apache.spark.sql.execution.streaming._ import org.apache.spark.sql.execution.streaming.runtime.RealTimeModeAllowlist import org.apache.spark.sql.execution.streaming.sources._ @@ -283,7 +282,7 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) extends streaming.D None } provider match { - case p: PythonDataSourceV2 => p.setShortName(source) + case p: NamedTableProvider => p.setShortName(source) case _ => } val table = DataSourceV2Utils.getTableFromProvider( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala index 52b9e803a64ed..49c85d949374a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala @@ -48,7 +48,8 @@ import org.apache.spark.sql.execution.datasources.jdbc.JdbcRelationProvider import org.apache.spark.sql.execution.datasources.json.JsonFileFormat import org.apache.spark.sql.execution.datasources.orc.OrcFileFormat import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat -import org.apache.spark.sql.execution.datasources.v2.FileDataSourceV2 +import org.apache.spark.sql.execution.datasources.v2.{FileDataSourceV2, NamedTableProvider} +import org.apache.spark.sql.execution.datasources.v2.ffi.{NativeDataSourceRegistry, NativeDataSourceV2} import org.apache.spark.sql.execution.datasources.v2.orc.OrcDataSourceV2 import org.apache.spark.sql.execution.datasources.v2.python.PythonDataSourceV2 import org.apache.spark.sql.execution.datasources.xml.XmlFileFormat @@ -682,6 +683,7 @@ object DataSource extends Logging { val serviceLoader = ServiceLoader.load(classOf[DataSourceRegister], loader) lazy val isUserDefinedDataSource = SparkSession.getActiveSession.exists( _.sessionState.dataSourceManager.dataSourceExists(provider)) + lazy val isNativeDataSource = NativeDataSourceRegistry.exists(provider, conf) try { serviceLoader.asScala.filter(_.shortName().equalsIgnoreCase(provider1)).toList match { @@ -703,6 +705,8 @@ object DataSource extends Logging { throw QueryCompilationErrors.failedToFindKafkaDataSourceError(provider1) } else if (isUserDefinedDataSource) { classOf[PythonDataSourceV2] + } else if (isNativeDataSource) { + classOf[NativeDataSourceV2] } else { throw QueryExecutionErrors.dataSourceNotFoundError(provider1, error) } @@ -718,9 +722,9 @@ object DataSource extends Logging { } } case head :: Nil => - // We do not check whether the provider is a Python data source - // (isUserDefinedDataSource) to avoid the lookup cost. Java data sources - // always take precedence over Python user-defined data sources. + // We do not check whether the provider is a Python or native data source + // (isUserDefinedDataSource, isNativeDataSource) to avoid the lookup cost. Java data + // sources always take precedence over Python and native data sources. head.getClass case sources => // There are multiple registered aliases for the input. If there is single datasource @@ -783,7 +787,7 @@ object DataSource extends Logging { case t: TableProvider if !useV1Sources.contains(cls.getCanonicalName.toLowerCase(Locale.ROOT)) => t match { - case p: PythonDataSourceV2 => p.setShortName(provider) + case p: NamedTableProvider => p.setShortName(provider) case _ => } Some(t) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/NamedTableProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/NamedTableProvider.scala new file mode 100644 index 0000000000000..badff781627c9 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/NamedTableProvider.scala @@ -0,0 +1,30 @@ +/* + * 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.execution.datasources.v2 + +import org.apache.spark.sql.connector.catalog.TableProvider + +/** + * A [[TableProvider]] that serves the data sources registered under their own short names, such + * as Python and native data sources. `DataSource.lookupDataSource` resolves those names to the + * provider class, so the caller has to tell each new instance which name it was created for with + * [[setShortName]] before using it. + */ +trait NamedTableProvider extends TableProvider { + def setShortName(name: String): Unit +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala new file mode 100644 index 0000000000000..5ecd1c8fc806b --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala @@ -0,0 +1,212 @@ +/* + * 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.execution.datasources.v2.columnar + +import java.util +import java.util.function.{Function => JFunction} + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.connector.catalog.{SupportsRead, SupportsWrite, Table, TableCapability} +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read._ +import org.apache.spark.sql.connector.read.streaming.{MicroBatchStream, Offset} +import org.apache.spark.sql.connector.write.{LogicalWriteInfo, WriteBuilder} +import org.apache.spark.sql.datasource.{DataSource, DataSourceReader, DataSourceStreamReader} +import org.apache.spark.sql.internal.connector.SupportsMetadata +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** + * The Data Source V2 [[Table]] of a data source implemented with the columnar data source API. + * It creates a new [[DataSource]] for each scan and each write, with their options. + */ +class ColumnarTable( + shortName: String, + tableSchema: StructType, + createDataSource: JFunction[CaseInsensitiveStringMap, DataSource]) + extends Table with SupportsRead with SupportsWrite { + + override def name(): String = shortName + + override def schema(): StructType = tableSchema + + override def capabilities(): util.Set[TableCapability] = + util.EnumSet.of(BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE) + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new ColumnarScanBuilder(tableSchema, () => createDataSource.apply(options)) + } + + override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = { + new ColumnarWriteBuilder(shortName, info, () => createDataSource.apply(info.options())) + } +} + +class ColumnarScanBuilder(schema: StructType, createDataSource: () => DataSource) + extends ScanBuilder + with SupportsPushDownV2Filters + with SupportsPushDownLimit + with SupportsPushDownRequiredColumns { + + private lazy val dataSource = createDataSource() + + // Created on first use, because a streaming scan neither pushes anything down nor reads in + // batch, and its data source may not support batch scans at all. + private var reader: DataSourceReader = _ + + private var pushed: Array[Predicate] = Array.empty + private var pushedLimit: Option[Int] = None + private var readSchema: StructType = schema + + private def batchReader: DataSourceReader = { + if (reader == null) { + reader = dataSource.reader(schema) + } + reader + } + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + val remaining = batchReader.pushPredicates(predicates) + pushed = predicates.filterNot(remaining.contains) + remaining + } + + override def pushedPredicates(): Array[Predicate] = pushed + + override def pushLimit(limit: Int): Boolean = { + val isPushed = batchReader.pushLimit(limit) + if (isPushed) { + pushedLimit = Some(limit) + } + isPushed + } + + // Spark always applies the limit again: a data source is free to return more rows. + override def isPartiallyPushed(): Boolean = true + + override def pruneColumns(requiredSchema: StructType): Unit = { + // Only top-level columns are pruned by the data source. Spark prunes the nested fields of the + // columns it reads. + val required = + StructType(requiredSchema.fieldNames.flatMap(name => schema.find(_.name == name))) + if (required.length < schema.length && batchReader.pruneColumns(required)) { + readSchema = required + } + } + + override def build(): Scan = { + new ColumnarScan(schema, readSchema, () => batchReader, () => dataSource, pushed, pushedLimit) + } +} + +class ColumnarScan( + fullSchema: StructType, + prunedSchema: StructType, + batchReader: () => DataSourceReader, + dataSource: () => DataSource, + pushedPredicates: Array[Predicate], + pushedLimit: Option[Int]) + extends Scan with SupportsMetadata { + + // Spark may call `toBatch` more than once for the same scan; the partitions are planned once. + private lazy val batch = new ColumnarBatchScan(batchReader()) + + override def readSchema(): StructType = prunedSchema + + // The class that implements the data source, such as NativeDataSource for a native one. + override def description(): String = dataSource().getClass.getSimpleName + + override def columnarSupportMode(): Scan.ColumnarSupportMode = + Scan.ColumnarSupportMode.SUPPORTED + + override def toBatch: Batch = batch + + override def toMicroBatchStream(checkpointLocation: String): MicroBatchStream = { + new ColumnarMicroBatchStream(dataSource().streamReader(fullSchema)) + } + + override def getMetaData(): Map[String, String] = { + Map( + "PushedPredicates" -> pushedPredicates.mkString("[", ", ", "]"), + "ReadSchema" -> prunedSchema.simpleString + ) ++ pushedLimit.map(limit => "PushedLimit" -> s"LIMIT $limit") + } +} + +class ColumnarBatchScan(reader: DataSourceReader) extends Batch { + private lazy val partitions: Array[InputPartition] = reader.partitions() + + override def planInputPartitions(): Array[InputPartition] = partitions + + override def createReaderFactory(): PartitionReaderFactory = { + // The reader is serialized to the executors, so it has to be planned first. + partitions + new ColumnarPartitionReaderFactory(reader) + } +} + +case class ColumnarStreamOffset(json: String) extends Offset + +class ColumnarMicroBatchStream(reader: DataSourceStreamReader) extends MicroBatchStream { + + override def initialOffset(): Offset = ColumnarStreamOffset(reader.initialOffset()) + + override def latestOffset(): Offset = ColumnarStreamOffset(reader.latestOffset()) + + override def planInputPartitions(start: Offset, end: Offset): Array[InputPartition] = { + reader.partitions(start.json(), end.json()) + } + + override def createReaderFactory(): PartitionReaderFactory = { + new ColumnarStreamPartitionReaderFactory(reader) + } + + override def deserializeOffset(json: String): Offset = ColumnarStreamOffset(json) + + override def commit(end: Offset): Unit = reader.commit(end.json()) + + override def stop(): Unit = reader.stop() +} + +/** + * Columnar data sources only produce columnar batches, so Spark never asks for rows: the scan + * reports [[Scan.ColumnarSupportMode.SUPPORTED]]. + */ +abstract class ColumnarOnlyPartitionReaderFactory extends PartitionReaderFactory { + override def createReader(partition: InputPartition): PartitionReader[InternalRow] = { + throw SparkException.internalError("Columnar data sources only support columnar reads.") + } + + override def supportColumnarReads(partition: InputPartition): Boolean = true +} + +class ColumnarPartitionReaderFactory(reader: DataSourceReader) + extends ColumnarOnlyPartitionReaderFactory { + override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = { + reader.read(partition) + } +} + +class ColumnarStreamPartitionReaderFactory(reader: DataSourceStreamReader) + extends ColumnarOnlyPartitionReaderFactory { + override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = { + reader.read(partition) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala new file mode 100644 index 0000000000000..5a7386d6e9145 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala @@ -0,0 +1,201 @@ +/* + * 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.execution.datasources.v2.columnar + +import scala.jdk.CollectionConverters._ + +import org.apache.arrow.vector.VectorSchemaRoot + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.connector.metric.CustomTaskMetric +import org.apache.spark.sql.connector.write._ +import org.apache.spark.sql.connector.write.streaming.{StreamingDataWriterFactory, StreamingWrite} +import org.apache.spark.sql.datasource.{DataSource, DataSourceStreamWriter, DataSourceWriter} +import org.apache.spark.sql.execution.arrow.ArrowWriter +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +class ColumnarWriteBuilder( + shortName: String, + info: LogicalWriteInfo, + createDataSource: () => DataSource) + extends WriteBuilder with SupportsTruncate { + + private var overwrite = false + + override def truncate(): WriteBuilder = { + overwrite = true + this + } + + override def build(): Write = new ColumnarWrite(shortName, info, overwrite, createDataSource) +} + +class ColumnarWrite( + shortName: String, + info: LogicalWriteInfo, + overwrite: Boolean, + createDataSource: () => DataSource) + extends Write { + + override def description(): String = shortName + + override def toBatch: BatchWrite = { + new ColumnarBatchWrite(createDataSource().writer(info.schema(), overwrite), info.schema()) + } + + override def toStreaming: StreamingWrite = { + new ColumnarStreamingWrite( + createDataSource().streamWriter(info.schema(), overwrite), info.schema()) + } +} + +class ColumnarBatchWrite(writer: DataSourceWriter, schema: StructType) extends BatchWrite { + + override def createBatchWriterFactory(info: PhysicalWriteInfo): DataWriterFactory = { + new ColumnarDataWriterFactory(writer, ArrowBatchingOptions(schema)) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = writer.commit(messages) + + override def abort(messages: Array[WriterCommitMessage]): Unit = writer.abort(messages) +} + +class ColumnarStreamingWrite(writer: DataSourceStreamWriter, schema: StructType) + extends StreamingWrite { + + override def createStreamingWriterFactory( + info: PhysicalWriteInfo): StreamingDataWriterFactory = { + new ColumnarStreamingDataWriterFactory(writer, ArrowBatchingOptions(schema)) + } + + override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + writer.commit(epochId, messages) + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + writer.abort(epochId, messages) + } +} + +class ColumnarDataWriterFactory(writer: DataSourceWriter, options: ArrowBatchingOptions) + extends DataWriterFactory { + override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = { + new ArrowBatchingDataWriter(writer.createWriter(partitionId, taskId), options) + } +} + +class ColumnarStreamingDataWriterFactory( + writer: DataSourceStreamWriter, + options: ArrowBatchingOptions) + extends StreamingDataWriterFactory { + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[InternalRow] = { + new ArrowBatchingDataWriter(writer.createWriter(partitionId, taskId, epochId), options) + } +} + +/** + * How [[ArrowBatchingDataWriter]] converts rows to Arrow, captured on the driver from the + * session configuration. + */ +case class ArrowBatchingOptions(schema: StructType, timeZoneId: String, maxRecordsPerBatch: Int) + +object ArrowBatchingOptions { + def apply(schema: StructType): ArrowBatchingOptions = { + val conf = SQLConf.get + val maxRecords = conf.arrowMaxRecordsPerBatch + ArrowBatchingOptions( + schema, conf.sessionLocalTimeZone, if (maxRecords > 0) maxRecords else Int.MaxValue) + } +} + +/** + * Converts the rows of a write task to Arrow-backed [[ColumnarBatch]]es for the [[DataWriter]] + * of a columnar data source. A new Arrow vector schema root is created for each batch, so a + * native writer that took ownership of the data of a batch can keep it after the call. + */ +class ArrowBatchingDataWriter(delegate: DataWriter[ColumnarBatch], options: ArrowBatchingOptions) + extends DataWriter[InternalRow] { + + private val arrowSchema = ArrowUtils.toArrowSchema( + options.schema, options.timeZoneId, errorOnDuplicatedFieldNames = true, largeVarTypes = false) + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator("columnar data source writer", 0, Long.MaxValue) + private var root: VectorSchemaRoot = _ + private var arrowWriter: ArrowWriter = _ + private var numRows = 0 + + override def write(record: InternalRow): Unit = { + if (root == null) { + root = VectorSchemaRoot.create(arrowSchema, allocator) + arrowWriter = ArrowWriter.create(root) + } + arrowWriter.write(record) + numRows += 1 + if (numRows >= options.maxRecordsPerBatch) { + flush() + } + } + + private def flush(): Unit = { + if (root != null) { + arrowWriter.finish() + val columns = root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector) + try { + delegate.write(new ColumnarBatch(columns.toArray, root.getRowCount)) + } finally { + releaseBatch() + } + } + } + + private def releaseBatch(): Unit = { + if (root != null) { + root.close() + root = null + arrowWriter = null + numRows = 0 + } + } + + override def commit(): WriterCommitMessage = { + flush() + delegate.commit() + } + + override def abort(): Unit = { + releaseBatch() + delegate.abort() + } + + override def close(): Unit = { + try { + delegate.close() + } finally { + releaseBatch() + Utils.tryLogNonFatalError(allocator.close()) + } + } + + override def currentMetricsValues(): Array[CustomTaskMetric] = delegate.currentMetricsValues() +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala new file mode 100644 index 0000000000000..dce0d4b16f01a --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -0,0 +1,614 @@ +/* + * 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.execution.datasources.v2.ffi + +import java.lang.ref.Cleaner +import java.util.concurrent.atomic.AtomicLong + +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, Data} +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} +import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.arrow.vector.types.pojo.Field + +import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.internal.Logging +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} +import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} +import org.apache.spark.sql.datasource.{DataSource, DataSourceReader, DataSourceStreamReader, DataSourceStreamWriter, DataSourceWriter} +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +/** + * A [[DataSource]] implemented by a native library, through + * `org.apache.spark.sql.datasource.NativeBridge`. + * + * @param name the name of the data source, as listed in the manifest of the package + */ +class NativeDataSource( + pkg: NativeDataSourcePackage, + name: String, + options: CaseInsensitiveStringMap) + extends DataSource { + import NativeCalls._ + + override def schema(): StructType = { + val library = NativeLibraries.get(pkg) + withDataSource(library, "infer the schema of") { dataSource => + call(name, "infer the schema of") { + optional(NativeArrow.importSchema(address => library.schema(dataSource, address))) + }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) + } + } + + override def reader(schema: StructType): DataSourceReader = { + val library = NativeLibraries.get(pkg) + val handle = create(library, "plan a scan of", "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema)(address => library.createReader(dataSource, address)) + } + new NativeDataSourceReader(pkg, name, schema, new NativeHandle(handle, library.closeReader)) + } + + override def streamReader(schema: StructType): DataSourceStreamReader = { + val library = NativeLibraries.get(pkg) + val handle = create(library, "plan a scan of", "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema)(address => library.createStreamReader(dataSource, address)) + } + new NativeDataSourceStreamReader( + pkg, name, schema, new NativeHandle(handle, library.closeStreamReader)) + } + + override def writer(schema: StructType, overwrite: Boolean): DataSourceWriter = { + val library = NativeLibraries.get(pkg) + val handle = new NativeHandle( + create(library, "plan a write to", "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource => + NativeArrow.exportSchema(schema) { address => + library.createWriter(dataSource, address, overwrite) + } + }, + library.closeWriter) + new NativeDataSourceWriter(pkg, name, handle, serializeWriter(library, handle)) + } + + override def streamWriter(schema: StructType, overwrite: Boolean): DataSourceStreamWriter = { + val library = NativeLibraries.get(pkg) + val handle = new NativeHandle( + create(library, "plan a write to", "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createStreamWriter(dataSource, address, overwrite) + } + }, + library.closeWriter) + new NativeDataSourceStreamWriter(pkg, name, handle, serializeWriter(library, handle)) + } + + /** Runs `f` with a new native data source, which is closed afterwards. */ + private def withDataSource[T](library: NativeLibrary, action: String)(f: Long => T): T = { + // The keys of a case-insensitive map are in lower case. + val entries = options.entrySet().asScala.toArray + val dataSource = call(name, "create") { + library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue)) + } + Utils.tryWithSafeFinally(f(dataSource)) { + call(name, action)(library.closeDataSource(dataSource)) + } + } + + /** Creates a reader or a writer with `f`, which fails if the library does not implement it. */ + private def create(library: NativeLibrary, action: String, unsupportedError: String)( + f: Long => Long): Long = { + withDataSource(library, action) { dataSource => + call(name, action)(optional(f(dataSource))).getOrElse { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) + } + } + } + + private def serializeWriter(library: NativeLibrary, handle: NativeHandle): Array[Byte] = { + try { + bytesOrEmpty(call(name, "plan a write to")(optional(library.serializeWriter(handle.get)))) + } catch { + case e: Throwable => + handle.close() + throw e + } + } +} + +/** Reads a native data source for a batch scan. */ +class NativeDataSourceReader( + pkg: NativeDataSourcePackage, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends DataSourceReader { + import NativeCalls._ + + // Set on the driver when the partitions are planned, and serialized to the executors. + private var readerState: Array[Byte] = Array.emptyByteArray + private var readSchema: StructType = schema + @transient private var plannedPartitions: Array[InputPartition] = _ + + @transient private lazy val library = NativeLibraries.get(pkg) + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + // Only the predicates that can be expressed in JSON are passed to the library. + val encoded = predicates.zipWithIndex.flatMap { case (predicate, i) => + NativePredicates.toJson(predicate).map(_ -> i) + } + val pushed = if (encoded.isEmpty) { + Set.empty[Int] + } else { + call(name, "plan a scan of") { + optional(library.pushPredicates(handle.get, encoded.map(_._1))) + } match { + case Some(flags) if flags != null && flags.length == encoded.length => + encoded.map(_._2).zip(flags).collect { case (i, true) => i }.toSet + case Some(flags) => + throw error(name, "plan a scan of", s"pushPredicates returned " + + s"${Option(flags).map(_.length).getOrElse(0)} values for ${encoded.length} predicates.") + case None => Set.empty[Int] + } + } + predicates.zipWithIndex.collect { case (predicate, i) if !pushed.contains(i) => predicate } + } + + override def pushLimit(limit: Int): Boolean = { + call(name, "plan a scan of")(optional(library.pushLimit(handle.get, limit))).getOrElse(false) + } + + override def pruneColumns(requiredSchema: StructType): Boolean = { + val pruned = call(name, "plan a scan of") { + optional(library.pruneColumns(handle.get, requiredSchema.fieldNames)) + }.getOrElse(false) + if (pruned) { + readSchema = requiredSchema + } + pruned + } + + override def partitions(): Array[InputPartition] = synchronized { + if (plannedPartitions == null) { + try { + val partitions = call(name, "plan a scan of")(library.partitions(handle.get)) + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeReader(handle.get)))) + plannedPartitions = NativeInputPartition.from(partitions) + } finally { + // The executors only need the serialized state. + handle.close() + } + } + plannedPartitions + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + pkg, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, readSchema) + } +} + +/** Reads a native data source for a micro-batch streaming scan. */ +class NativeDataSourceStreamReader( + pkg: NativeDataSourcePackage, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends DataSourceStreamReader { + import NativeCalls._ + + // Set on the driver whenever the partitions of a micro-batch are planned, and serialized to + // the executors with the partitions of that micro-batch. + private var readerState: Array[Byte] = Array.emptyByteArray + + @transient private lazy val library = NativeLibraries.get(pkg) + + override def initialOffset(): String = { + call(name, "plan a scan of")(library.initialOffset(handle.get)) + } + + override def latestOffset(): String = { + call(name, "plan a scan of")(library.latestOffset(handle.get)) + } + + override def partitions(start: String, end: String): Array[InputPartition] = { + val partitions = call(name, "plan a scan of") { + library.streamPartitions(handle.get, start, end) + } + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeStreamReader(handle.get)))) + NativeInputPartition.from(partitions) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + pkg, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, schema) + } + + override def commit(end: String): Unit = { + call(name, "commit an offset of")(optional(library.commitOffset(handle.get, end))) + } + + override def stop(): Unit = handle.close() +} + +/** Writes to a native data source for a batch write. */ +class NativeDataSourceWriter( + pkg: NativeDataSourcePackage, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends DataSourceWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(pkg) + + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(pkg, name, writerState, partitionId, taskId, epochId = -1L) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + handle.close() + } + + override def abort(messages: Array[WriterCommitMessage]): Unit = { + Utils.tryWithSafeFinally { + call(name, "abort a write to") { + library.abort(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + }(handle.close()) + } +} + +/** Writes to a native data source for a streaming write. */ +class NativeDataSourceStreamWriter( + pkg: NativeDataSourcePackage, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends DataSourceStreamWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(pkg) + + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(pkg, name, writerState, partitionId, taskId, epochId) + } + + override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "abort a write to") { + library.abort(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } +} + +/** A partition planned by a native library. */ +case class NativeInputPartition(index: Int, bytes: Array[Byte]) extends InputPartition + +object NativeInputPartition { + def from(partitions: Array[Array[Byte]]): Array[InputPartition] = { + Option(partitions).getOrElse(Array.empty).zipWithIndex.map { case (bytes, i) => + NativeInputPartition(i, Option(bytes).getOrElse(Array.emptyByteArray)) + } + } +} + +/** The message returned by the data writer of a task of a native library. */ +case class NativeWriterCommitMessage(bytes: Array[Byte]) extends WriterCommitMessage + +object NativeWriterCommitMessage { + def toBytes(messages: Array[WriterCommitMessage]): Array[Array[Byte]] = messages.map { + case message: NativeWriterCommitMessage => message.bytes + case null => null + case other => + throw SparkException.internalError(s"Unexpected commit message: ${other.getClass.getName}") + } +} + +/** + * Reads a partition of a native data source: the Arrow stream exported by the library is + * imported into Arrow vectors that back the returned batch. + */ +class NativePartitionReader( + pkg: NativeDataSourcePackage, + name: String, + readerState: Array[Byte], + partition: Array[Byte], + schema: StructType) + extends PartitionReader[ColumnarBatch] { + import NativeCalls._ + + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) + private var reader: ArrowReader = _ + private var root: VectorSchemaRoot = _ + private var batch: ColumnarBatch = _ + + try { + reader = openStream() + root = call(name, "read")(reader.getVectorSchemaRoot) + checkSchema() + batch = new ColumnarBatch( + root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector).toArray) + } catch { + case e: Throwable => + close() + throw e + } + + private def openStream(): ArrowReader = { + val library = NativeLibraries.get(pkg) + val stream = ArrowArrayStream.allocateNew(allocator) + try { + call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) + .getOrElse { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) + } + Data.importArrayStream(allocator, stream) + } catch { + case e: Throwable => + stream.release() + throw e + } finally { + stream.close() + } + } + + private def checkSchema(): Unit = { + def isDictionaryEncoded(field: Field): Boolean = { + field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) + } + root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + + "is not supported.") + } + val actual = ArrowUtils.fromArrowSchema(root.getSchema) + val matches = actual.length == schema.length && actual.zip(schema).forall { + case (actualField, field) => + actualField.name == field.name && + DataType.equalsIgnoreNullability(actualField.dataType, field.dataType) + } + if (!matches) { + throw error(name, "read", s"The data has the schema ${actual.toDDL}, but the expected " + + s"schema is ${schema.toDDL}.") + } + } + + override def next(): Boolean = { + val hasNext = call(name, "read")(reader.loadNextBatch()) + if (hasNext) { + batch.setNumRows(root.getRowCount) + } + hasNext + } + + override def get(): ColumnarBatch = batch + + override def close(): Unit = { + Utils.tryWithSafeFinally { + if (reader != null) { + reader.close() + } + }(Utils.tryLogNonFatalError(allocator.close())) + } +} + +/** Writes the batches of a task to the data writer of a native library. */ +class NativeDataWriter( + pkg: NativeDataSourcePackage, + name: String, + writerState: Array[Byte], + partitionId: Int, + taskId: Long, + epochId: Long) + extends DataWriter[ColumnarBatch] { + import NativeCalls._ + + private val library = NativeLibraries.get(pkg) + // Zero once the data writer was committed or aborted, which releases it. + private var handle: Long = call(name, "write to") { + library.createDataWriter(writerState, partitionId, taskId, epochId) + } + + override def write(batch: ColumnarBatch): Unit = { + if (handle == 0) { + throw SparkException.internalError("The data writer was already committed or aborted.") + } + val vectors = (0 until batch.numCols()).map { i => + batch.column(i) match { + case vector: ArrowColumnVector => vector.getValueVector.asInstanceOf[FieldVector] + case vector => throw SparkException.internalError( + s"Expected an Arrow column vector, but got ${vector.getClass.getName}.") + } + } + val root = new VectorSchemaRoot( + vectors.map(_.getField).asJava, vectors.asJava, batch.numRows()) + val array = ArrowArray.allocateNew(ArrowUtils.rootAllocator) + val arrowSchema = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) + try { + Data.exportVectorSchemaRoot(ArrowUtils.rootAllocator, root, null, array, arrowSchema) + call(name, "write to") { + library.write(handle, array.memoryAddress(), arrowSchema.memoryAddress()) + } + } finally { + // Releases the batch unless the library moved it to take ownership. + array.release() + array.close() + arrowSchema.release() + arrowSchema.close() + } + } + + override def commit(): WriterCommitMessage = { + val dataWriter = handle + handle = 0 + NativeWriterCommitMessage( + bytesOrEmpty(call(name, "write to")(Some(library.commitDataWriter(dataWriter))))) + } + + override def abort(): Unit = { + if (handle != 0) { + val dataWriter = handle + handle = 0 + call(name, "write to")(library.abortDataWriter(dataWriter)) + } + } + + override def close(): Unit = abort() +} + +/** + * A handle to a native object. It is released with `close()`, or once it becomes unreachable if + * it was never closed. + */ +final class NativeHandle(handle: Long, release: Long => Unit) { + private val state = new NativeHandle.State(handle, release) + private val cleanable = NativeHandle.cleaner.register(this, state) + + def get: Long = { + val value = state.handle.get() + if (value == 0) { + throw SparkException.internalError("The native handle was already released.") + } + value + } + + def close(): Unit = cleanable.clean() +} + +object NativeHandle extends Logging { + private val cleaner = Cleaner.create() + + // Must not refer to the handle, which would then never become unreachable. + private class State(value: Long, release: Long => Unit) extends Runnable { + val handle = new AtomicLong(value) + + override def run(): Unit = { + val value = handle.getAndSet(0) + if (value != 0) { + try { + release(value) + } catch { + case e @ (NonFatal(_) | _: LinkageError) => + logWarning("Failed to release a native data source handle.", e) + } + } + } + } +} + +/** Helpers to call native data source libraries. */ +object NativeCalls { + /** Calls a native library, and rethrows the errors it reports as NATIVE_DATA_SOURCE_ERROR. */ + def call[T](name: String, action: String)(f: => T): T = { + try { + f + } catch { + case e: Throwable with SparkThrowable => throw e + case e: UnsatisfiedLinkError => + throw QueryExecutionErrors.nativeDataSourceError( + action, name, s"The library does not implement ${e.getMessage}.", e) + case NonFatal(e) => + throw QueryExecutionErrors.nativeDataSourceError(action, name, message(e), e) + } + } + + // How Arrow reports the error message of an Arrow C stream. + private val StreamError = """(?s)\[errno \d+\] CDataJniException\{errno=\d+, message=(.*)\}""".r + + private def message(e: Throwable): String = Option(e.getMessage).getOrElse(e.toString) match { + case StreamError(message) => message + case message => message + } + + /** Calls a function that the library may not implement. Returns None if it does not. */ + def optional[T](f: => T): Option[T] = { + try { + Some(f) + } catch { + case _: UnsatisfiedLinkError => None + } + } + + /** Returns a NATIVE_DATA_SOURCE_ERROR for an error detected by Spark. */ + def error(name: String, action: String, msg: String): Throwable = { + QueryExecutionErrors.nativeDataSourceError(action, name, msg, null) + } + + def bytesOrEmpty(bytes: Option[Array[Byte]]): Array[Byte] = { + bytes.flatMap(Option(_)).getOrElse(Array.emptyByteArray) + } +} + +/** Exchanges schemas with native libraries through the Arrow C data interface. */ +object NativeArrow { + /** Exports the schema as an `ArrowSchema`, and calls `f` with its address. */ + def exportSchema[T](schema: StructType)(f: Long => T): T = { + val arrowSchema = ArrowUtils.toArrowSchema( + schema, SQLConf.get.sessionLocalTimeZone, errorOnDuplicatedFieldNames = true, + largeVarTypes = false) + val struct = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) + try { + Data.exportSchema(ArrowUtils.rootAllocator, arrowSchema, null, struct) + f(struct.memoryAddress()) + } finally { + // Releases the schema unless the library moved it to take ownership. + struct.release() + struct.close() + } + } + + /** Calls `f` with the address of an `ArrowSchema` to initialize, and imports it. */ + def importSchema(f: Long => Unit): StructType = { + val struct = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) + var imported = false + try { + f(struct.memoryAddress()) + imported = true + // Releases and closes the struct. + ArrowUtils.fromArrowSchema(Data.importSchema(ArrowUtils.rootAllocator, struct, null)) + } finally { + if (!imported) { + struct.release() + struct.close() + } + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala new file mode 100644 index 0000000000000..6fe7a9b8e1c9b --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -0,0 +1,264 @@ +/* + * 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.execution.datasources.v2.ffi + +import java.io.{File, FileInputStream, IOException} +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, StandardCopyOption} +import java.security.MessageDigest +import java.util.{HexFormat, Locale} +import java.util.concurrent.ConcurrentHashMap +import java.util.zip.ZipFile + +import scala.jdk.CollectionConverters._ + +import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} + +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.PATH +import org.apache.spark.sql.Artifact +import org.apache.spark.sql.datasource.NativeBridge +import org.apache.spark.sql.errors.QueryExecutionErrors +import org.apache.spark.util.Utils + +/** + * The manifest of a native data source package, `spark-native-datasource.json`: + * {{{ + * { + * "abiVersion": 1, + * "dataSources": ["my_source"], + * "libraries": { + * "linux-x86_64": "linux-x86_64/libmy_source.so", + * "osx-aarch64": "osx-aarch64/libmy_source.dylib" + * } + * } + * }}} + * + * @param abiVersion the version of [[NativeBridge]] that the libraries implement + * @param dataSources the names of the data sources that the libraries implement + * @param libraries the path of the library in the package for each platform, as returned by + * [[NativePlatform.current]] + */ +case class NativeDataSourceManifest( + abiVersion: Int, + dataSources: Seq[String], + libraries: Map[String, String]) + +/** + * A native data source package: a zip file that contains a [[NativeDataSourceManifest]] and the + * libraries it lists. The SHA-256 checksum of the file identifies the package, so that an + * executor can check that it loads the same package as the driver. + * + * @param path the path of the package file on the driver + * @param artifactUUID the UUID of the session whose artifacts contain the package, if the session + * is isolated. Executors store the artifacts of such a session separately. + */ +case class NativeDataSourcePackage( + path: String, + sha256: String, + manifest: NativeDataSourceManifest, + artifactUUID: Option[String] = None) { + def fileName: String = new File(path).getName +} + +object NativeDataSourcePackage { + val FILE_EXTENSION: String = Artifact.NATIVE_DATA_SOURCE_PACKAGE_EXTENSION + val MANIFEST_NAME = "spark-native-datasource.json" + + private case class CacheKey(path: String, length: Long, lastModified: Long) + + private val packages = new ConcurrentHashMap[CacheKey, NativeDataSourcePackage]() + private val mapper = new ObjectMapper() + + /** Reads the package in the given file. The result is cached until the file changes. */ + def read(file: File): NativeDataSourcePackage = { + val canonical = file.getCanonicalFile + val key = CacheKey(canonical.getPath, canonical.length(), canonical.lastModified()) + packages.computeIfAbsent(key, _ => parse(canonical)) + } + + private def parse(file: File): NativeDataSourcePackage = { + def invalid(reason: String, cause: Throwable = null): Throwable = { + QueryExecutionErrors.invalidNativeDataSourcePackageError( + file.getPath, "INVALID_MANIFEST", Map("reason" -> reason), cause) + } + val manifest = try { + Utils.tryWithResource(new ZipFile(file)) { zip => + val entry = Option(zip.getEntry(MANIFEST_NAME)).getOrElse { + throw invalid(s"The package does not contain the manifest $MANIFEST_NAME.") + } + val json = Utils.tryWithResource(zip.getInputStream(entry)) { in => + new String(in.readAllBytes(), StandardCharsets.UTF_8) + } + val manifest = parseManifest(json).fold(reason => throw invalid(reason), identity) + manifest.libraries.values.find(zip.getEntry(_) == null).foreach { library => + throw invalid(s"The package does not contain the library $library.") + } + manifest + } + } catch { + case e: IOException => throw invalid(s"The package is not a valid zip file: $e", e) + } + NativeDataSourcePackage(file.getPath, sha256(file), manifest) + } + + private[ffi] def parseManifest(json: String): Either[String, NativeDataSourceManifest] = { + val root: JsonNode = try { + mapper.readTree(json) + } catch { + case e: IOException => return Left(s"The manifest is not valid JSON: ${e.getMessage}") + } + if (root == null || !root.isObject) { + return Left("The manifest must be a JSON object.") + } + val abiVersion = root.get("abiVersion") + if (abiVersion == null || !abiVersion.isInt) { + return Left("The manifest must have an integer field 'abiVersion'.") + } + val dataSources = root.get("dataSources") + if (dataSources == null || !dataSources.isArray || dataSources.isEmpty || + !dataSources.elements().asScala.forall(n => n.isTextual && n.asText().nonEmpty)) { + return Left("The manifest must have a field 'dataSources' with a non-empty array of " + + "data source names.") + } + val libraries = root.get("libraries") + if (libraries == null || !libraries.isObject || libraries.isEmpty || + !libraries.properties().asScala.forall(e => e.getValue.isTextual)) { + return Left("The manifest must have a field 'libraries' with an object that maps " + + "platforms to the paths of the libraries in the package.") + } + Right(NativeDataSourceManifest( + abiVersion.asInt(), + dataSources.elements().asScala.map(_.asText()).toSeq, + libraries.properties().asScala.map(e => e.getKey -> e.getValue.asText()).toMap)) + } + + private def sha256(file: File): String = { + val digest = MessageDigest.getInstance("SHA-256") + Utils.tryWithResource(new FileInputStream(file)) { in => + val buffer = new Array[Byte](64 * 1024) + var read = in.read(buffer) + while (read != -1) { + digest.update(buffer, 0, read) + read = in.read(buffer) + } + } + HexFormat.of().formatHex(digest.digest()) + } +} + +/** The platform names used in [[NativeDataSourceManifest.libraries]]. */ +object NativePlatform { + /** The platform of this JVM, such as `linux-x86_64`, `linux-aarch64` or `osx-aarch64`. */ + lazy val current: String = name(System.getProperty("os.name"), System.getProperty("os.arch")) + + private[ffi] def name(osName: String, osArch: String): String = { + val os = osName.toLowerCase(Locale.ROOT) match { + case n if n.startsWith("linux") => "linux" + case n if n.startsWith("mac") || n.startsWith("darwin") => "osx" + case n if n.startsWith("windows") => "windows" + case n => n.replaceAll("[^a-z0-9]", "") + } + val arch = osArch.toLowerCase(Locale.ROOT) match { + case "amd64" | "x86_64" | "x64" => "x86_64" + case "aarch64" | "arm64" => "aarch64" + case a => a + } + s"$os-$arch" + } +} + +/** + * The native data source libraries loaded in this JVM. A library is extracted from its package + * and loaded once, on first use, and stays loaded until the JVM exits. + */ +object NativeLibraries extends Logging { + private val libraries = new ConcurrentHashMap[String, NativeLibrary]() + + private lazy val extractionDir: File = Utils.createTempDir(namePrefix = "native-datasources") + + /** Returns the library of the given package, and loads it if needed. */ + private[ffi] def get(pkg: NativeDataSourcePackage): NativeLibrary = { + val library = libraries.get(pkg.sha256) + if (library != null) library else libraries.computeIfAbsent(pkg.sha256, _ => load(pkg)) + } + + private def load(pkg: NativeDataSourcePackage): NativeLibrary = { + val platform = NativePlatform.current + val libraryPath = pkg.manifest.libraries.getOrElse(platform, { + throw QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, + "UNSUPPORTED_PLATFORM", + Map( + "platform" -> platform, + "platforms" -> pkg.manifest.libraries.keys.toSeq.sorted.mkString(", "))) + }) + val packageFile = locate(pkg) + val libraryFile = new File(new File(extractionDir, pkg.sha256), new File(libraryPath).getName) + if (!libraryFile.isFile) { + extract(packageFile, libraryPath, libraryFile) + } + logInfo(log"Loading the native data source library ${MDC(PATH, libraryFile)}.") + + def invalidLibrary(reason: String, cause: Throwable): Throwable = { + QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, "INVALID_LIBRARY", Map("library" -> libraryPath, "reason" -> reason), cause) + } + val library = try { + NativeLibrary.load(libraryFile.getPath) + } catch { + case e: UnsatisfiedLinkError => throw invalidLibrary(e.getMessage, e) + } + val abiVersion = try { + library.abiVersion() + } catch { + case e: UnsatisfiedLinkError => + throw invalidLibrary("It does not implement NativeBridge.abiVersion.", e) + } + if (abiVersion != NativeBridge.ABI_VERSION) { + throw QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, + "UNSUPPORTED_ABI_VERSION", + Map("version" -> abiVersion.toString, "supported" -> NativeBridge.ABI_VERSION.toString)) + } + library + } + + /** Finds a copy of the package on this node: the one on the driver, or a distributed one. */ + private def locate(pkg: NativeDataSourcePackage): File = { + (new File(pkg.path) +: NativeDataSourceRegistry.distributedCopies(pkg)) + .find(file => file.isFile && NativeDataSourcePackage.read(file).sha256 == pkg.sha256) + .getOrElse { + throw QueryExecutionErrors.nativeDataSourcePackageNotFoundError(pkg.fileName, pkg.sha256) + } + } + + private def extract(packageFile: File, libraryPath: String, target: File): Unit = { + target.getParentFile.mkdirs() + val temp = File.createTempFile(target.getName, ".tmp", target.getParentFile) + try { + Utils.tryWithResource(new ZipFile(packageFile)) { zip => + Utils.tryWithResource(zip.getInputStream(zip.getEntry(libraryPath))) { in => + Files.copy(in, temp.toPath, StandardCopyOption.REPLACE_EXISTING) + } + } + Files.move(temp.toPath, target.toPath, StandardCopyOption.REPLACE_EXISTING) + } finally { + temp.delete() + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala new file mode 100644 index 0000000000000..94d76493e630d --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -0,0 +1,143 @@ +/* + * 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.execution.datasources.v2.ffi + +import java.io.File +import java.nio.file.Paths + +import org.apache.spark.{SparkEnv, SparkFiles} +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.PATH +import org.apache.spark.sql.Artifact +import org.apache.spark.sql.classic.SparkSession +import org.apache.spark.sql.datasource.NativeBridge +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.util.Utils + +/** + * Finds the native data source packages of the active session on the driver: the packages added + * with `spark.addArtifact`, and the packages listed in, or in a directory listed in, + * `spark.sql.dataSource.native.paths`. + */ +object NativeDataSourceRegistry extends Logging { + + /** Returns whether a package provides the native data source with the given name. */ + def exists(name: String, conf: SQLConf): Boolean = lookup(name, conf).isDefined + + /** + * Finds the package that provides the native data source with the given name, which is + * case-insensitive. Returns the name of the data source as listed in the manifest, and the + * package. + */ + def lookup(name: String, conf: SQLConf): Option[(String, NativeDataSourcePackage)] = { + if (!conf.getConf(StaticSQLConf.NATIVE_DATA_SOURCE_ENABLED)) { + return None + } + val found = packages(conf) + .flatMap(pkg => pkg.manifest.dataSources.find(_.equalsIgnoreCase(name)).map(_ -> pkg)) + // The same package can be found more than once, for example when it is in a configured + // directory and was also added as an artifact. + .groupBy(_._2.sha256).values.map(_.head).toSeq + found match { + case Seq() => None + case Seq(result @ (_, pkg)) => + if (pkg.manifest.abiVersion != NativeBridge.ABI_VERSION) { + throw QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, + "UNSUPPORTED_ABI_VERSION", + Map( + "version" -> pkg.manifest.abiVersion.toString, + "supported" -> NativeBridge.ABI_VERSION.toString)) + } + Some(result) + case _ => + throw QueryCompilationErrors.nativeDataSourcePackageConflictError( + name, found.map(_._2.path).sorted) + } + } + + /** + * Makes the package available to the executors of the active session. Returns the package with + * the location of the copies of the executors. + */ + def distribute(pkg: NativeDataSourcePackage): NativeDataSourcePackage = { + SparkSession.getActiveSession match { + // In local mode, the executors run in the driver and read the package where it is. + case Some(session) if !session.sparkContext.isLocal => addToArtifacts(session, pkg) + case _ => pkg + } + } + + /** Adds the package as an artifact of the session, unless it already is one. */ + private[ffi] def addToArtifacts( + session: SparkSession, + pkg: NativeDataSourcePackage): NativeDataSourcePackage = { + val artifactManager = session.artifactManager + val (added, artifactUUID) = artifactManager.getNativeDataSourcePackages + if (!added.exists(_.getName == pkg.fileName)) { + artifactManager.addLocalArtifacts(Artifact.newFileArtifact( + Paths.get(pkg.fileName), new Artifact.LocalFile(Paths.get(pkg.path))) :: Nil) + } + pkg.copy(artifactUUID = artifactUUID) + } + + /** + * The local copies, on an executor, of a package distributed with [[distribute]]. The files of + * an isolated session are stored in a directory named after the UUID of the session. + */ + def distributedCopies(pkg: NativeDataSourcePackage): Seq[File] = { + if (SparkEnv.get == null) { + Nil + } else { + val root = new File(SparkFiles.getRootDirectory()) + pkg.artifactUUID.map(uuid => new File(new File(root, uuid), pkg.fileName)).toSeq :+ + new File(root, pkg.fileName) + } + } + + private def packages(conf: SQLConf): Seq[NativeDataSourcePackage] = { + (configuredPackageFiles(conf) ++ artifactPackageFiles()).map(_.getCanonicalFile).distinct + .map(NativeDataSourcePackage.read) + } + + private def configuredPackageFiles(conf: SQLConf): Seq[File] = { + conf.getConf(SQLConf.NATIVE_DATA_SOURCE_PATHS).flatMap { path => + val uri = Utils.resolveURI(path) + val file = new File(uri.getPath) + if (uri.getScheme != null && uri.getScheme != "file") { + logWarning(log"Ignoring the native data source path ${MDC(PATH, path)}: only local " + + log"paths are supported. Add remote packages with spark.addArtifact instead.") + Nil + } else if (file.isDirectory) { + file.listFiles() + .filter(f => f.isFile && f.getName.endsWith(NativeDataSourcePackage.FILE_EXTENSION)) + .sortBy(_.getName) + .toSeq + } else if (file.isFile) { + Seq(file) + } else { + logWarning(log"The native data source path ${MDC(PATH, path)} does not exist.") + Nil + } + } + } + + private def artifactPackageFiles(): Seq[File] = { + SparkSession.getActiveSession.toSeq.flatMap(_.artifactManager.getNativeDataSourcePackages._1) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala new file mode 100644 index 0000000000000..d518af1062e26 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala @@ -0,0 +1,66 @@ +/* + * 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.execution.datasources.v2.ffi + +import java.util + +import org.apache.spark.sql.connector.catalog.Table +import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.execution.datasources.v2.NamedTableProvider +import org.apache.spark.sql.execution.datasources.v2.columnar.ColumnarTable +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * The Data Source V2 provider of the native data sources: data sources implemented by a native + * library through `org.apache.spark.sql.datasource.NativeBridge`. + */ +class NativeDataSourceV2 extends NamedTableProvider { + private var shortName: String = _ + + override def setShortName(name: String): Unit = { + assert(shortName == null) + shortName = name + } + + // The name of the data source as listed in the manifest, and its package. Resolved on first + // use, which is while the query is analyzed, with the configuration of its session. + private lazy val resolved: (String, NativeDataSourcePackage) = { + assert(shortName != null) + val (name, pkg) = NativeDataSourceRegistry.lookup(shortName, SQLConf.get).getOrElse { + throw QueryCompilationErrors.dataSourceDoesNotExist(shortName) + } + (name, NativeDataSourceRegistry.distribute(pkg)) + } + + override def inferSchema(options: CaseInsensitiveStringMap): StructType = { + val (name, pkg) = resolved + new NativeDataSource(pkg, name, options).schema() + } + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: util.Map[String, String]): Table = { + val (name, pkg) = resolved + new ColumnarTable(shortName, schema, options => new NativeDataSource(pkg, name, options)) + } + + override def supportsExternalMetadata(): Boolean = true +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala new file mode 100644 index 0000000000000..7c6c6d2bbffd1 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala @@ -0,0 +1,92 @@ +/* + * 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.execution.datasources.v2.ffi + +import java.util.Base64 + +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.node.ObjectNode + +import org.apache.spark.sql.connector.expressions.{Expression, GeneralScalarExpression, Literal, NamedReference} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.types._ + +/** + * Encodes Data Source V2 predicates as the JSON expression trees passed to + * `NativeBridge.pushPredicates`. + */ +object NativePredicates { + private val mapper = new ObjectMapper() + + /** Returns the JSON of the predicate, or None if it cannot be expressed in JSON. */ + def toJson(predicate: Predicate): Option[String] = { + encode(predicate).map(mapper.writeValueAsString) + } + + private def encode(expression: Expression): Option[ObjectNode] = expression match { + case reference: NamedReference => + val node = mapper.createObjectNode().put("type", "column") + val name = node.putArray("name") + reference.fieldNames().foreach(name.add) + Some(node) + + case literal: Literal[_] => + encodeLiteral(literal.value(), literal.dataType()) + + case function: GeneralScalarExpression => + val children = function.children().map(encode) + if (children.forall(_.isDefined)) { + val node = mapper.createObjectNode() + .put("type", "function") + .put("name", function.name()) + val array = node.putArray("children") + children.foreach(child => array.add(child.get)) + Some(node) + } else { + None + } + + case _ => None + } + + private def encodeLiteral(value: Any, dataType: DataType): Option[ObjectNode] = { + val node = mapper.createObjectNode() + .put("type", "literal") + .put("dataType", dataType.catalogString) + if (value == null) { + Some(node.putNull("value")) + } else { + val encoded = (value, dataType) match { + case (b: Boolean, BooleanType) => Some(node.put("value", b)) + case (n: Byte, ByteType) => Some(node.put("value", n.toInt)) + case (n: Short, ShortType) => Some(node.put("value", n)) + case (n: Int, IntegerType | DateType) => Some(node.put("value", n)) + case (n: Long, LongType | TimestampType | TimestampNTZType) => Some(node.put("value", n)) + case (n: Float, FloatType) if !n.isNaN && !n.isInfinite => Some(node.put("value", n)) + case (n: Double, DoubleType) if !n.isNaN && !n.isInfinite => Some(node.put("value", n)) + // A source compares strings by their bytes, so only push down binary collations. + case (s, st: StringType) if st.isUTF8BinaryCollation => Some(node.put("value", s.toString)) + case (bytes: Array[Byte], BinaryType) => + Some(node.put("value", Base64.getEncoder.encodeToString(bytes))) + case (d: Decimal, _: DecimalType) => + Some(node.put("value", d.toJavaBigDecimal.toPlainString)) + case _ => None + } + encoded + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala index 9fcf81746fe30..cd04bba0c288e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.execution.datasources.v2.python import org.apache.spark.sql.SparkSession import org.apache.spark.sql.connector.catalog._ import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.execution.datasources.v2.NamedTableProvider import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap @@ -26,10 +27,10 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap /** * Data Source V2 wrapper for Python Data Source. */ -class PythonDataSourceV2 extends TableProvider { +class PythonDataSourceV2 extends NamedTableProvider { private var name: String = _ - def setShortName(str: String): Unit = { + override def setShortName(str: String): Unit = { assert(name == null) name = str } diff --git a/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister b/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister index c1fc7234d7c19..7c0052d49ec15 100644 --- a/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister +++ b/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister @@ -30,3 +30,4 @@ org.apache.spark.sql.streaming.sources.FakeWriteOnly org.apache.spark.sql.streaming.sources.FakeNoWrite org.apache.spark.sql.streaming.sources.FakeWriteSupportingExternalMetadata org.apache.spark.sql.streaming.sources.FakeWriteSupportProviderV1Fallback +org.apache.spark.sql.datasource.TestRangeProvider diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc new file mode 100644 index 0000000000000..6f55494d383e5 --- /dev/null +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -0,0 +1,1107 @@ +/* + * 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 native data source library used by NativeDataSourceSuite. It implements the JNI functions of +// org.apache.spark.sql.datasource.NativeBridge without any dependency, building the Arrow C data +// interface structs by hand. It provides three data sources: +// +// native_range Reads rows [0, end) with columns of several types. Supports predicate, +// column and limit pushdown. +// native_sink Writes the rows as CSV files under the directory in the option "path". +// native_counter Streams the rows [0, max_offset) with a single column "id", making "step" +// more rows available each time Spark asks for the latest offset. +// +// The suite compiles it with these macros to get several distinct libraries: +// +// NAME_PREFIX A prefix for the names of the data sources, "" by default. +// ID_SIGN Multiplies the ids produced by native_range, 1 by default. +// READ_ONLY Leaves out the functions for writes, streaming reads and pushdown. +// ABI_VERSION The version returned by abiVersion(), 1 by default. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#ifndef NAME_PREFIX +#define NAME_PREFIX "" +#endif +#ifndef ID_SIGN +#define ID_SIGN 1 +#endif +#ifndef ABI_VERSION +#define ABI_VERSION 1 +#endif + +#define SPARK_JNI(ret, method) \ + extern "C" JNIEXPORT ret JNICALL Java_org_apache_spark_sql_datasource_NativeBridge_##method + +// --------------------------------------------------------------------------------------------- +// The Arrow C data interface: https://arrow.apache.org/docs/format/CDataInterface.html +// --------------------------------------------------------------------------------------------- + +#ifndef ARROW_C_DATA_INTERFACE +#define ARROW_C_DATA_INTERFACE + +#define ARROW_FLAG_NULLABLE 2 + +struct ArrowSchema { + const char* format; + const char* name; + const char* metadata; + int64_t flags; + int64_t n_children; + struct ArrowSchema** children; + struct ArrowSchema* dictionary; + void (*release)(struct ArrowSchema*); + void* private_data; +}; + +struct ArrowArray { + int64_t length; + int64_t null_count; + int64_t offset; + int64_t n_buffers; + int64_t n_children; + const void** buffers; + struct ArrowArray** children; + struct ArrowArray* dictionary; + void (*release)(struct ArrowArray*); + void* private_data; +}; + +#endif // ARROW_C_DATA_INTERFACE + +#ifndef ARROW_C_STREAM_INTERFACE +#define ARROW_C_STREAM_INTERFACE + +struct ArrowArrayStream { + int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out); + int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out); + const char* (*get_last_error)(struct ArrowArrayStream*); + void (*release)(struct ArrowArrayStream*); + void* private_data; +}; + +#endif // ARROW_C_STREAM_INTERFACE + +namespace { + +// --------------------------------------------------------------------------------------------- +// Columns +// --------------------------------------------------------------------------------------------- + +struct ColumnType { + const char* name; + const char* format; +}; + +// The columns of native_range, with their Arrow formats. native_counter only has "id". +const ColumnType kColumns[] = { + {"id", "l"}, // bigint + {"mod", "i"}, // int + {"name", "u"}, // string, null when id is a multiple of 5 + {"value", "g"}, // double + {"flag", "b"}, // boolean + {"day", "tdD"}, // date + {"ts", "tsu:UTC"} // timestamp +}; + +const ColumnType* findColumn(const std::string& name) { + for (const auto& column : kColumns) { + if (name == column.name) return &column; + } + throw std::runtime_error("unknown column " + name); +} + +std::vector allColumns(bool counter) { + if (counter) return {"id"}; + std::vector names; + for (const auto& column : kColumns) names.push_back(column.name); + return names; +} + +// --------------------------------------------------------------------------------------------- +// Exporting schemas +// --------------------------------------------------------------------------------------------- + +struct SchemaData { + std::string format; + std::string name; + std::vector> children; + std::vector child_pointers; +}; + +void releaseSchema(ArrowSchema* schema) { + auto* data = static_cast(schema->private_data); + for (auto& child : data->children) { + if (child->release != nullptr) child->release(child.get()); + } + delete data; + schema->release = nullptr; +} + +void initSchema(ArrowSchema* out, const std::string& format, const std::string& name, + int64_t flags) { + auto* data = new SchemaData{format, name, {}, {}}; + out->format = data->format.c_str(); + out->name = data->name.c_str(); + out->metadata = nullptr; + out->flags = flags; + out->n_children = 0; + out->children = nullptr; + out->dictionary = nullptr; + out->release = releaseSchema; + out->private_data = data; +} + +// Exports a struct type with the given columns. +void exportSchema(const std::vector& columns, ArrowSchema* out) { + initSchema(out, "+s", "", 0); + auto* data = static_cast(out->private_data); + for (const auto& name : columns) { + auto child = std::make_unique(); + initSchema(child.get(), findColumn(name)->format, name, ARROW_FLAG_NULLABLE); + data->child_pointers.push_back(child.get()); + data->children.push_back(std::move(child)); + } + out->n_children = static_cast(data->children.size()); + out->children = data->child_pointers.data(); +} + +// Reads the column names of a struct schema exported by Spark, and releases it. +std::vector importColumns(ArrowSchema* schema) { + std::vector columns; + for (int64_t i = 0; i < schema->n_children; i++) { + columns.push_back(schema->children[i]->name); + } + schema->release(schema); + return columns; +} + +// --------------------------------------------------------------------------------------------- +// Exporting arrays +// --------------------------------------------------------------------------------------------- + +struct ArrayData { + std::vector> buffers; + std::vector buffer_pointers; + std::vector> children; + std::vector child_pointers; +}; + +void releaseArray(ArrowArray* array) { + auto* data = static_cast(array->private_data); + for (auto& child : data->children) { + if (child->release != nullptr) child->release(child.get()); + } + delete data; + array->release = nullptr; +} + +// Initializes an array that owns `buffers`. An empty buffer is exported as null. +void initArray(ArrowArray* out, int64_t length, int64_t null_count, + std::vector> buffers) { + auto* data = new ArrayData{std::move(buffers), {}, {}, {}}; + for (auto& buffer : data->buffers) { + data->buffer_pointers.push_back(buffer.empty() ? nullptr : buffer.data()); + } + out->length = length; + out->null_count = null_count; + out->offset = 0; + out->n_buffers = static_cast(data->buffers.size()); + out->n_children = 0; + out->buffers = data->buffer_pointers.data(); + out->children = nullptr; + out->dictionary = nullptr; + out->release = releaseArray; + out->private_data = data; +} + +template +std::vector toBytes(const std::vector& values) { + std::vector bytes(values.size() * sizeof(T)); + if (!bytes.empty()) std::memcpy(bytes.data(), values.data(), bytes.size()); + return bytes; +} + +void setBit(std::vector& bitmap, int64_t i) { + bitmap[i / 8] |= static_cast(1 << (i % 8)); +} + +// Exports the values of a column for the ids [lo, hi). +void exportColumn(const std::string& name, int64_t lo, int64_t hi, ArrowArray* out) { + int64_t n = hi - lo; + std::vector no_validity; + if (name == "id") { + std::vector values; + for (int64_t id = lo; id < hi; id++) values.push_back(id * ID_SIGN); + initArray(out, n, 0, {no_validity, toBytes(values)}); + } else if (name == "mod") { + std::vector values; + for (int64_t id = lo; id < hi; id++) values.push_back(static_cast(id % 3)); + initArray(out, n, 0, {no_validity, toBytes(values)}); + } else if (name == "name") { + std::vector validity((n + 7) / 8, 0); + std::vector offsets{0}; + std::string chars; + int64_t null_count = 0; + for (int64_t id = lo; id < hi; id++) { + if (id % 5 == 0) { + null_count++; + } else { + setBit(validity, id - lo); + chars += "n" + std::to_string(id); + } + offsets.push_back(static_cast(chars.size())); + } + std::vector chars_bytes(chars.begin(), chars.end()); + initArray(out, n, null_count, {validity, toBytes(offsets), chars_bytes}); + } else if (name == "value") { + std::vector values; + for (int64_t id = lo; id < hi; id++) values.push_back(static_cast(id) * 0.5); + initArray(out, n, 0, {no_validity, toBytes(values)}); + } else if (name == "flag") { + std::vector bits((n + 7) / 8, 0); + for (int64_t id = lo; id < hi; id++) { + if (id % 2 == 0) setBit(bits, id - lo); + } + initArray(out, n, 0, {no_validity, bits}); + } else if (name == "day") { + std::vector values; + for (int64_t id = lo; id < hi; id++) values.push_back(static_cast(id)); + initArray(out, n, 0, {no_validity, toBytes(values)}); + } else if (name == "ts") { + std::vector values; + for (int64_t id = lo; id < hi; id++) values.push_back(id * 1000000); + initArray(out, n, 0, {no_validity, toBytes(values)}); + } else { + throw std::runtime_error("unknown column " + name); + } +} + +// Exports a struct array with the given columns for the ids [lo, hi). +void exportBatch(const std::vector& columns, int64_t lo, int64_t hi, + ArrowArray* out) { + initArray(out, hi - lo, 0, {std::vector()}); + auto* data = static_cast(out->private_data); + for (const auto& name : columns) { + auto child = std::make_unique(); + exportColumn(name, lo, hi, child.get()); + data->child_pointers.push_back(child.get()); + data->children.push_back(std::move(child)); + } + out->n_children = static_cast(data->children.size()); + out->children = data->child_pointers.data(); +} + +// --------------------------------------------------------------------------------------------- +// A minimal JSON parser, for the predicates and the offsets +// --------------------------------------------------------------------------------------------- + +struct Json { + enum Kind { kNull, kBool, kNumber, kString, kArray, kObject } kind = kNull; + bool boolean = false; + double number = 0; + std::string text; + std::vector items; + std::map fields; + + const Json& operator[](const std::string& key) const { + static const Json missing; + auto it = fields.find(key); + return it == fields.end() ? missing : it->second; + } +}; + +class JsonParser { + public: + explicit JsonParser(const std::string& text) : text_(text) {} + + Json parse() { + Json value = parseValue(); + skipSpaces(); + if (pos_ != text_.size()) fail(); + return value; + } + + private: + [[noreturn]] void fail() { throw std::runtime_error("invalid JSON: " + text_); } + + void skipSpaces() { + while (pos_ < text_.size() && std::isspace(static_cast(text_[pos_]))) pos_++; + } + + bool consume(char c) { + skipSpaces(); + if (pos_ < text_.size() && text_[pos_] == c) { + pos_++; + return true; + } + return false; + } + + void expect(char c) { + if (!consume(c)) fail(); + } + + bool consumeWord(const char* word) { + size_t length = std::strlen(word); + if (text_.compare(pos_, length, word) == 0) { + pos_ += length; + return true; + } + return false; + } + + std::string parseString() { + expect('"'); + std::string result; + while (pos_ < text_.size() && text_[pos_] != '"') { + char c = text_[pos_++]; + if (c == '\\') { + if (pos_ >= text_.size()) fail(); + char escaped = text_[pos_++]; + switch (escaped) { + case 'n': result += '\n'; break; + case 't': result += '\t'; break; + case 'u': result += '?'; pos_ += 4; break; + default: result += escaped; + } + } else { + result += c; + } + } + expect('"'); + return result; + } + + Json parseValue() { + skipSpaces(); + Json value; + if (pos_ >= text_.size()) fail(); + char c = text_[pos_]; + if (c == '{') { + value.kind = Json::kObject; + pos_++; + if (!consume('}')) { + do { + skipSpaces(); + std::string key = parseString(); + expect(':'); + value.fields[key] = parseValue(); + } while (consume(',')); + expect('}'); + } + } else if (c == '[') { + value.kind = Json::kArray; + pos_++; + if (!consume(']')) { + do { + value.items.push_back(parseValue()); + } while (consume(',')); + expect(']'); + } + } else if (c == '"') { + value.kind = Json::kString; + value.text = parseString(); + } else if (consumeWord("true")) { + value.kind = Json::kBool; + value.boolean = true; + } else if (consumeWord("false")) { + value.kind = Json::kBool; + } else if (consumeWord("null")) { + value.kind = Json::kNull; + } else { + value.kind = Json::kNumber; + char* end = nullptr; + value.number = std::strtod(text_.c_str() + pos_, &end); + if (end == text_.c_str() + pos_) fail(); + pos_ = end - text_.c_str(); + } + return value; + } + + const std::string& text_; + size_t pos_ = 0; +}; + +// --------------------------------------------------------------------------------------------- +// State, serialized as "key=value" lines +// --------------------------------------------------------------------------------------------- + +using Options = std::map; + +std::string serialize(const Options& options) { + std::string result; + for (const auto& entry : options) result += entry.first + "=" + entry.second + "\n"; + return result; +} + +Options deserialize(const std::string& text) { + Options options; + std::istringstream lines(text); + std::string line; + while (std::getline(lines, line)) { + size_t eq = line.find('='); + if (eq != std::string::npos) options[line.substr(0, eq)] = line.substr(eq + 1); + } + return options; +} + +std::string get(const Options& options, const std::string& key, const std::string& fallback) { + auto it = options.find(key); + return it == options.end() ? fallback : it->second; +} + +int64_t getLong(const Options& options, const std::string& key, int64_t fallback) { + auto it = options.find(key); + return it == options.end() ? fallback : std::stoll(it->second); +} + +std::string join(const std::vector& values) { + std::string result; + for (size_t i = 0; i < values.size(); i++) result += (i > 0 ? "," : "") + values[i]; + return result; +} + +std::vector split(const std::string& text) { + std::vector values; + std::string value; + std::istringstream stream(text); + while (std::getline(stream, value, ',')) values.push_back(value); + return values; +} + +// --------------------------------------------------------------------------------------------- +// JNI helpers +// --------------------------------------------------------------------------------------------- + +void throwJava(JNIEnv* env, const std::string& message) { + jclass cls = env->FindClass("java/lang/RuntimeException"); + if (cls != nullptr) env->ThrowNew(cls, message.c_str()); +} + +// Runs `f`, and turns a C++ exception into a Java exception: none may cross the boundary. +template +T guarded(JNIEnv* env, T fallback, F&& f) { + try { + return f(); + } catch (const std::exception& e) { + throwJava(env, e.what()); + return fallback; + } +} + +template +void guardedVoid(JNIEnv* env, F&& f) { + try { + f(); + } catch (const std::exception& e) { + throwJava(env, e.what()); + } +} + +std::string toString(JNIEnv* env, jstring value) { + if (value == nullptr) return ""; + const char* chars = env->GetStringUTFChars(value, nullptr); + std::string result(chars); + env->ReleaseStringUTFChars(value, chars); + return result; +} + +std::vector toStrings(JNIEnv* env, jobjectArray values) { + std::vector result; + jsize n = env->GetArrayLength(values); + for (jsize i = 0; i < n; i++) { + auto value = static_cast(env->GetObjectArrayElement(values, i)); + result.push_back(toString(env, value)); + env->DeleteLocalRef(value); + } + return result; +} + +std::string toBytes(JNIEnv* env, jbyteArray bytes) { + if (bytes == nullptr) return ""; + jsize n = env->GetArrayLength(bytes); + std::string result(static_cast(n), '\0'); + if (n > 0) env->GetByteArrayRegion(bytes, 0, n, reinterpret_cast(&result[0])); + return result; +} + +jbyteArray newBytes(JNIEnv* env, const std::string& bytes) { + jbyteArray result = env->NewByteArray(static_cast(bytes.size())); + env->SetByteArrayRegion(result, 0, static_cast(bytes.size()), + reinterpret_cast(bytes.data())); + return result; +} + +jobjectArray newByteArrays(JNIEnv* env, const std::vector& values) { + jobjectArray result = + env->NewObjectArray(static_cast(values.size()), env->FindClass("[B"), nullptr); + for (size_t i = 0; i < values.size(); i++) { + jbyteArray value = newBytes(env, values[i]); + env->SetObjectArrayElement(result, static_cast(i), value); + env->DeleteLocalRef(value); + } + return result; +} + +[[maybe_unused]] +std::vector toByteStrings(JNIEnv* env, jobjectArray values, bool* has_null) { + std::vector result; + jsize n = env->GetArrayLength(values); + for (jsize i = 0; i < n; i++) { + auto value = static_cast(env->GetObjectArrayElement(values, i)); + if (value == nullptr) { + *has_null = true; + } else { + result.push_back(toBytes(env, value)); + env->DeleteLocalRef(value); + } + } + return result; +} + +template +jlong toHandle(T* object) { + return reinterpret_cast(object); +} + +template +T* fromHandle(jlong handle) { + return reinterpret_cast(handle); +} + +// --------------------------------------------------------------------------------------------- +// Data sources +// --------------------------------------------------------------------------------------------- + +std::string name(const std::string& base) { return std::string(NAME_PREFIX) + base; } + +struct DataSource { + std::string name; + Options options; +}; + +void maybeFail(const Options& options, const std::string& step) { + if (get(options, "fail", "") == step) throw std::runtime_error("injected failure in " + step); +} + +// The driver-side reader of native_range and native_counter. +struct Reader { + Options state; // Serialized to the executors. +}; + +// Reads the ids [lo, hi) of the columns in the state, in batches. +struct Stream { + std::vector columns; + int64_t next; + int64_t end; + int64_t batch_size; + bool fail; + std::string last_error; +}; + +int streamGetSchema(ArrowArrayStream* stream, ArrowSchema* out) { + auto* state = static_cast(stream->private_data); + exportSchema(state->columns, out); + return 0; +} + +int streamGetNext(ArrowArrayStream* stream, ArrowArray* out) { + auto* state = static_cast(stream->private_data); + if (state->fail) { + state->last_error = "injected failure in get_next"; + return EIO; + } + if (state->next >= state->end) { + out->release = nullptr; // The end of the stream. + return 0; + } + int64_t hi = std::min(state->end, state->next + state->batch_size); + exportBatch(state->columns, state->next, hi, out); + state->next = hi; + return 0; +} + +const char* streamGetLastError(ArrowArrayStream* stream) { + return static_cast(stream->private_data)->last_error.c_str(); +} + +void streamRelease(ArrowArrayStream* stream) { + delete static_cast(stream->private_data); + stream->release = nullptr; +} + +// The driver-side writer of native_sink. +struct Writer { + Options state; // path, overwrite, fail. +}; + +// Writes the rows of a task as CSV lines to a temporary file, renamed on commit. +struct DataWriter { + std::string temp_path; + FILE* file; + Options state; +}; + +[[maybe_unused]] void makeDirs(const std::string& path) { + for (size_t pos = path.find('/', 1); pos != std::string::npos; pos = path.find('/', pos + 1)) { + mkdir(path.substr(0, pos).c_str(), 0755); + } + mkdir(path.c_str(), 0755); +} + +[[maybe_unused]] bool isValid(const ArrowArray* array, int64_t i) { + if (array->null_count == 0 || array->buffers[0] == nullptr) return true; + int64_t bit = array->offset + i; + return (static_cast(array->buffers[0])[bit / 8] >> (bit % 8)) & 1; +} + +// Formats the value at row `i` of a column, or "null". The helpers below are only used by +// writes, which READ_ONLY leaves out. +[[maybe_unused]] +std::string formatValue(const ArrowSchema* schema, const ArrowArray* array, int64_t i) { + if (!isValid(array, i)) return "null"; + std::string format = schema->format; + int64_t row = array->offset + i; + if (format == "l" || format.rfind("ts", 0) == 0) { + return std::to_string(static_cast(array->buffers[1])[row]); + } else if (format == "i" || format == "tdD") { + return std::to_string(static_cast(array->buffers[1])[row]); + } else if (format == "g") { + std::ostringstream out; + out << static_cast(array->buffers[1])[row]; + return out.str(); + } else if (format == "b") { + return ((static_cast(array->buffers[1])[row / 8] >> (row % 8)) & 1) + ? "true" : "false"; + } else if (format == "u") { + const int32_t* offsets = static_cast(array->buffers[1]); + const char* chars = static_cast(array->buffers[2]); + return std::string(chars + offsets[row], static_cast(offsets[row + 1] - offsets[row])); + } + throw std::runtime_error("unsupported Arrow format " + format); +} + +[[maybe_unused]] +std::vector listFiles(const std::string& dir, const std::string& prefix) { + std::vector files; + DIR* handle = opendir(dir.c_str()); + if (handle == nullptr) return files; + while (dirent* entry = readdir(handle)) { + std::string file = entry->d_name; + if (file.rfind(prefix, 0) == 0) files.push_back(file); + } + closedir(handle); + return files; +} + +[[maybe_unused]] void writeFile(const std::string& path, const std::string& content) { + FILE* file = std::fopen(path.c_str(), "w"); + if (file == nullptr) throw std::runtime_error("cannot write " + path); + std::fputs(content.c_str(), file); + std::fclose(file); +} + +} // namespace + +// --------------------------------------------------------------------------------------------- +// NativeBridge +// --------------------------------------------------------------------------------------------- + +SPARK_JNI(jint, abiVersion)(JNIEnv*, jclass) { return ABI_VERSION; } + +SPARK_JNI(jlong, createDataSource)(JNIEnv* env, jclass, jstring name_value, jobjectArray keys, + jobjectArray values) { + return guarded(env, 0, [&] { + auto* source = new DataSource{toString(env, name_value), {}}; + std::unique_ptr owner(source); + auto key_strings = toStrings(env, keys); + auto value_strings = toStrings(env, values); + for (size_t i = 0; i < key_strings.size(); i++) { + source->options[key_strings[i]] = value_strings[i]; + } + maybeFail(source->options, "create"); + if (source->name != name("native_range") && source->name != name("native_sink") && + source->name != name("native_counter")) { + throw std::runtime_error("unknown data source " + source->name); + } + return toHandle(owner.release()); + }); +} + +SPARK_JNI(void, schema)(JNIEnv* env, jclass, jlong handle, jlong schema_address) { + guardedVoid(env, [&] { + auto* source = fromHandle(handle); + maybeFail(source->options, "schema"); + if (source->name == name("native_sink")) { + throw std::runtime_error("native_sink has no schema; it can only be written"); + } + exportSchema(allColumns(source->name == name("native_counter")), + reinterpret_cast(schema_address)); + }); +} + +SPARK_JNI(jlong, createReader)(JNIEnv* env, jclass, jlong handle, jlong schema_address) { + return guarded(env, 0, [&] { + auto* source = fromHandle(handle); + if (source->name != name("native_range")) { + throw std::runtime_error(source->name + " does not support batch reads"); + } + auto columns = importColumns(reinterpret_cast(schema_address)); + for (const auto& column : columns) findColumn(column); + auto* reader = new Reader(); + reader->state["columns"] = join(columns); + reader->state["lo"] = "0"; + reader->state["hi"] = get(source->options, "end", "10"); + reader->state["partitions"] = get(source->options, "partitions", "2"); + reader->state["batch_size"] = get(source->options, "batch_size", "1000"); + reader->state["fail"] = get(source->options, "fail", ""); + return toHandle(reader); + }); +} + +#ifndef READ_ONLY + +SPARK_JNI(jlong, createStreamReader)(JNIEnv* env, jclass, jlong handle, jlong schema_address) { + return guarded(env, 0, [&] { + auto* source = fromHandle(handle); + if (source->name != name("native_counter")) { + throw std::runtime_error(source->name + " does not support streaming reads"); + } + auto* reader = new Reader(); + reader->state["columns"] = join(importColumns(reinterpret_cast(schema_address))); + reader->state["max_offset"] = get(source->options, "max_offset", "10"); + reader->state["step"] = get(source->options, "step", reader->state["max_offset"]); + reader->state["current"] = "0"; + reader->state["batch_size"] = get(source->options, "batch_size", "1000"); + reader->state["commit_path"] = get(source->options, "commit_path", ""); + return toHandle(reader); + }); +} + +#endif // READ_ONLY + +SPARK_JNI(void, closeDataSource)(JNIEnv*, jclass, jlong handle) { + delete fromHandle(handle); +} + +#ifndef READ_ONLY + +SPARK_JNI(jbooleanArray, pushPredicates)(JNIEnv* env, jclass, jlong handle, + jobjectArray predicates) { + return guarded(env, nullptr, [&] { + auto* reader = fromHandle(handle); + auto json = toStrings(env, predicates); + std::vector accepted; + int64_t lo = getLong(reader->state, "lo", 0); + int64_t hi = getLong(reader->state, "hi", 0); + for (const auto& text : json) { + // Accepts "id " for the comparison operators. + Json predicate = JsonParser(text).parse(); + const Json& children = predicate["children"]; + std::string op = predicate["name"].text; + bool comparison = op == ">" || op == ">=" || op == "<" || op == "<=" || op == "="; + bool ok = ID_SIGN == 1 && predicate["type"].text == "function" && comparison && + children.items.size() == 2 && children.items[0]["type"].text == "column" && + children.items[0]["name"].items.size() == 1 && + children.items[0]["name"].items[0].text == "id" && + children.items[1]["type"].text == "literal" && + children.items[1]["value"].kind == Json::kNumber; + if (ok) { + auto value = static_cast(children.items[1]["value"].number); + if (op == ">") lo = std::max(lo, value + 1); + if (op == ">=") lo = std::max(lo, value); + if (op == "<") hi = std::min(hi, value); + if (op == "<=") hi = std::min(hi, value + 1); + if (op == "=") { + lo = std::max(lo, value); + hi = std::min(hi, value + 1); + } + } + accepted.push_back(ok ? JNI_TRUE : JNI_FALSE); + } + reader->state["lo"] = std::to_string(lo); + reader->state["hi"] = std::to_string(std::max(lo, hi)); + reader->state["predicates"] = std::to_string(json.size()); + jbooleanArray result = env->NewBooleanArray(static_cast(accepted.size())); + env->SetBooleanArrayRegion(result, 0, static_cast(accepted.size()), accepted.data()); + return result; + }); +} + +SPARK_JNI(jboolean, pushLimit)(JNIEnv* env, jclass, jlong handle, jint limit) { + return guarded(env, JNI_FALSE, [&] { + fromHandle(handle)->state["limit"] = std::to_string(limit); + return JNI_TRUE; + }); +} + +SPARK_JNI(jboolean, pruneColumns)(JNIEnv* env, jclass, jlong handle, jobjectArray columns) { + return guarded(env, JNI_FALSE, [&] { + auto names = toStrings(env, columns); + for (const auto& column : names) findColumn(column); + fromHandle(handle)->state["columns"] = join(names); + return JNI_TRUE; + }); +} + +#endif // READ_ONLY + +SPARK_JNI(jobjectArray, partitions)(JNIEnv* env, jclass, jlong handle) { + return guarded(env, nullptr, [&] { + auto* reader = fromHandle(handle); + maybeFail(reader->state, "plan"); + int64_t lo = getLong(reader->state, "lo", 0); + int64_t hi = getLong(reader->state, "hi", 0); + int64_t count = getLong(reader->state, "partitions", 1); + if (reader->state.count("limit") > 0) { + // A pushed limit is applied by reading a single partition with at most `limit` rows. + hi = std::min(hi, lo + getLong(reader->state, "limit", 0)); + count = 1; + } + std::vector partitions; + int64_t step = std::max(1, (hi - lo + count - 1) / count); + for (int64_t start = lo; start < hi; start += step) { + int64_t stop = std::min(hi, start + step); + partitions.push_back(std::to_string(start) + "," + std::to_string(stop)); + } + return newByteArrays(env, partitions); + }); +} + +SPARK_JNI(jbyteArray, serializeReader)(JNIEnv* env, jclass, jlong handle) { + return guarded(env, nullptr, [&] { + return newBytes(env, serialize(fromHandle(handle)->state)); + }); +} + +SPARK_JNI(void, closeReader)(JNIEnv*, jclass, jlong handle) { delete fromHandle(handle); } + +#ifndef READ_ONLY + +SPARK_JNI(jstring, initialOffset)(JNIEnv* env, jclass, jlong) { + return env->NewStringUTF("{\"offset\":0}"); +} + +// Each call makes `step` more rows available, up to `max_offset`. +SPARK_JNI(jstring, latestOffset)(JNIEnv* env, jclass, jlong handle) { + return guarded(env, nullptr, [&] { + Options& state = fromHandle(handle)->state; + int64_t latest = std::min(getLong(state, "max_offset", 10), + getLong(state, "current", 0) + getLong(state, "step", 10)); + state["current"] = std::to_string(latest); + return env->NewStringUTF(("{\"offset\":" + std::to_string(latest) + "}").c_str()); + }); +} + +SPARK_JNI(jobjectArray, streamPartitions)(JNIEnv* env, jclass, jlong, jstring start, + jstring end) { + return guarded(env, nullptr, [&] { + auto lo = static_cast(JsonParser(toString(env, start)).parse()["offset"].number); + auto hi = static_cast(JsonParser(toString(env, end)).parse()["offset"].number); + std::vector partitions; + if (hi > lo) partitions.push_back(std::to_string(lo) + "," + std::to_string(hi)); + return newByteArrays(env, partitions); + }); +} + +SPARK_JNI(jbyteArray, serializeStreamReader)(JNIEnv* env, jclass, jlong handle) { + return guarded(env, nullptr, [&] { + return newBytes(env, serialize(fromHandle(handle)->state)); + }); +} + +SPARK_JNI(void, commitOffset)(JNIEnv* env, jclass, jlong handle, jstring end) { + guardedVoid(env, [&] { + std::string path = get(fromHandle(handle)->state, "commit_path", ""); + if (!path.empty()) writeFile(path, toString(env, end)); + }); +} + +SPARK_JNI(void, closeStreamReader)(JNIEnv*, jclass, jlong handle) { + delete fromHandle(handle); +} + +#endif // READ_ONLY + +SPARK_JNI(void, read)(JNIEnv* env, jclass, jbyteArray reader_state, jbyteArray partition, + jlong stream_address) { + guardedVoid(env, [&] { + Options state = deserialize(toBytes(env, reader_state)); + maybeFail(state, "open"); + auto bounds = split(toBytes(env, partition)); + if (bounds.size() != 2) throw std::runtime_error("invalid partition"); + auto* stream = reinterpret_cast(stream_address); + stream->get_schema = streamGetSchema; + stream->get_next = streamGetNext; + stream->get_last_error = streamGetLastError; + stream->release = streamRelease; + stream->private_data = new Stream{split(get(state, "columns", "")), std::stoll(bounds[0]), + std::stoll(bounds[1]), getLong(state, "batch_size", 1000), + get(state, "fail", "") == "read", ""}; + }); +} + +#ifndef READ_ONLY + +namespace { + +jlong createAnyWriter(JNIEnv* env, jlong handle, jlong schema_address, jboolean overwrite) { + return guarded(env, 0, [&] { + auto* source = fromHandle(handle); + if (source->name != name("native_sink")) { + throw std::runtime_error(source->name + " does not support writes"); + } + // Spark releases the schema after the call, since it is not moved here. + auto* schema = reinterpret_cast(schema_address); + if (schema->n_children == 0) throw std::runtime_error("no columns to write"); + auto* writer = new Writer(); + writer->state["path"] = get(source->options, "path", ""); + writer->state["overwrite"] = overwrite ? "true" : "false"; + writer->state["fail"] = get(source->options, "fail", ""); + if (writer->state["path"].empty()) { + delete writer; + throw std::runtime_error("the option 'path' is required"); + } + return toHandle(writer); + }); +} + +} // namespace + +SPARK_JNI(jlong, createWriter)(JNIEnv* env, jclass, jlong handle, jlong schema_address, + jboolean overwrite) { + return createAnyWriter(env, handle, schema_address, overwrite); +} + +SPARK_JNI(jlong, createStreamWriter)(JNIEnv* env, jclass, jlong handle, jlong schema_address, + jboolean overwrite) { + return createAnyWriter(env, handle, schema_address, overwrite); +} + +SPARK_JNI(jbyteArray, serializeWriter)(JNIEnv* env, jclass, jlong handle) { + return guarded(env, nullptr, [&] { + return newBytes(env, serialize(fromHandle(handle)->state)); + }); +} + +SPARK_JNI(void, commit)(JNIEnv* env, jclass, jlong handle, jlong epoch_id, + jobjectArray messages) { + guardedVoid(env, [&] { + auto* writer = fromHandle(handle); + std::string path = writer->state["path"]; + maybeFail(writer->state, "commit"); + bool has_null = false; + auto temp_files = toByteStrings(env, messages, &has_null); + if (has_null) throw std::runtime_error("commit got a null message"); + if (writer->state["overwrite"] == "true") { + for (const auto& file : listFiles(path, "part-")) unlink((path + "/" + file).c_str()); + } + for (const auto& temp_file : temp_files) { + std::string file_name = temp_file.substr(temp_file.rfind('/') + 1); + std::rename(temp_file.c_str(), (path + "/part-" + file_name).c_str()); + } + writeFile(path + (epoch_id < 0 ? "/_SUCCESS" : "/_EPOCH_" + std::to_string(epoch_id)), + std::to_string(temp_files.size())); + }); +} + +SPARK_JNI(void, abort)(JNIEnv* env, jclass, jlong handle, jlong, jobjectArray messages) { + guardedVoid(env, [&] { + auto* writer = fromHandle(handle); + bool has_null = false; + for (const auto& temp_file : toByteStrings(env, messages, &has_null)) { + unlink(temp_file.c_str()); + } + writeFile(writer->state["path"] + "/_ABORTED", ""); + }); +} + +SPARK_JNI(void, closeWriter)(JNIEnv*, jclass, jlong handle) { delete fromHandle(handle); } + +SPARK_JNI(jlong, createDataWriter)(JNIEnv* env, jclass, jbyteArray writer_state, + jint partition_id, jlong task_id, jlong epoch_id) { + return guarded(env, 0, [&] { + Options state = deserialize(toBytes(env, writer_state)); + std::string dir = state["path"] + "/_temporary"; + makeDirs(dir); + std::string temp_path = dir + "/" + std::to_string(epoch_id) + "-" + + std::to_string(partition_id) + "-" + std::to_string(task_id) + ".csv"; + FILE* file = std::fopen(temp_path.c_str(), "w"); + if (file == nullptr) throw std::runtime_error("cannot write " + temp_path); + return toHandle(new DataWriter{temp_path, file, state}); + }); +} + +SPARK_JNI(void, write)(JNIEnv* env, jclass, jlong handle, jlong array_address, + jlong schema_address) { + guardedVoid(env, [&] { + auto* writer = fromHandle(handle); + auto* array = reinterpret_cast(array_address); + auto* schema = reinterpret_cast(schema_address); + // The batch is released here, so Spark does not release it after the call. + std::unique_ptr array_owner( + array, [](ArrowArray* a) { a->release(a); }); + std::unique_ptr schema_owner( + schema, [](ArrowSchema* s) { s->release(s); }); + maybeFail(writer->state, "write"); + for (int64_t i = 0; i < array->length; i++) { + std::string line; + for (int64_t c = 0; c < array->n_children; c++) { + if (c > 0) line += ","; + line += formatValue(schema->children[c], array->children[c], array->offset + i); + } + line += "\n"; + std::fputs(line.c_str(), writer->file); + } + }); +} + +SPARK_JNI(jbyteArray, commitDataWriter)(JNIEnv* env, jclass, jlong handle) { + return guarded(env, nullptr, [&] { + std::unique_ptr writer(fromHandle(handle)); + std::fclose(writer->file); + return newBytes(env, writer->temp_path); + }); +} + +SPARK_JNI(void, abortDataWriter)(JNIEnv*, jclass, jlong handle) { + std::unique_ptr writer(fromHandle(handle)); + std::fclose(writer->file); + unlink(writer->temp_path.c_str()); +} + +#endif // READ_ONLY diff --git a/sql/core/src/test/scala/org/apache/spark/sql/datasource/ColumnarDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/datasource/ColumnarDataSourceSuite.scala new file mode 100644 index 0000000000000..a674188f63c18 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/datasource/ColumnarDataSourceSuite.scala @@ -0,0 +1,372 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.datasource + +import java.nio.charset.StandardCharsets.UTF_8 +import java.util.concurrent.ConcurrentHashMap + +import scala.collection.mutable.ArrayBuffer + +import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.connector.expressions.{Literal, NamedReference} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} +import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation +import org.apache.spark.sql.execution.datasources.v2.columnar.ColumnarScan +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.streaming.StreamingQuery +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +class ColumnarDataSourceSuite extends QueryTest with SharedSparkSession { + import testImplicits._ + + private val format = "test_columnar_range" + + private def scan(df: org.apache.spark.sql.classic.DataFrame): ColumnarScan = { + df.queryExecution.optimizedPlan.collectFirst { + case relation: DataSourceV2ScanRelation => relation.scan.asInstanceOf[ColumnarScan] + }.get + } + + private def findError(e: Throwable, condition: String): SparkThrowable = { + Iterator.iterate(e)(_.getCause).takeWhile(_ != null) + .collectFirst { case t: SparkThrowable if t.getCondition == condition => t } + .getOrElse(throw e) + } + + test("batch read") { + val df = spark.read.format(format).option("end", "10").load() + assert(df.schema == StructType.fromDDL("id BIGINT, name STRING")) + checkAnswer(df, (0L until 10L).map(id => Row(id, s"name$id"))) + // A data source can also be used by the class name of its provider. + checkAnswer( + spark.read.format(classOf[TestRangeProvider].getName).load(), + (0L until 10L).map(id => Row(id, s"name$id"))) + checkAnswer( + spark.read.schema("id BIGINT").format(format).option("end", "3").load(), + Seq(Row(0L), Row(1L), Row(2L))) + } + + test("push down predicates, a limit and the required columns") { + val filtered = spark.read.format(format).option("end", "10").load().filter("id > 6") + checkAnswer(filtered, (7L until 10L).map(id => Row(id, s"name$id"))) + assert(scan(filtered).getMetaData()("PushedPredicates").contains("id > 6")) + + val pruned = spark.read.format(format).option("end", "4").load().select("name") + checkAnswer(pruned, (0L until 4L).map(id => Row(s"name$id"))) + assert(scan(pruned).readSchema() == StructType.fromDDL("name STRING")) + + val limited = spark.read.format(format).option("end", "10").load().limit(2) + assert(limited.collect().length == 2) + assert(scan(limited).getMetaData()("PushedLimit") == "LIMIT 2") + } + + test("batch write") { + withSQLConf(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "3") { + val table = "batch_write" + (0L until 7L).map(id => (id, s"v$id")).toDF("id", "name").coalesce(1) + .write.format(format).option("table", table).mode("append").save() + assert(TestColumnarSink.rows(table) == (0L until 7L).map(id => (id, s"v$id"))) + // The rows are written in Arrow-backed batches of at most 3 rows. + assert(TestColumnarSink.batchSizes(table) == Seq(3, 3, 1)) + + Seq((10L, "a")).toDF("id", "name") + .write.format(format).option("table", table).mode("append").save() + assert(TestColumnarSink.rows(table).length == 8) + Seq((20L, "b")).toDF("id", "name") + .write.format(format).option("table", table).mode("overwrite").save() + assert(TestColumnarSink.rows(table) == Seq((20L, "b"))) + } + } + + test("abort a failed write") { + val table = "failed_write" + val e = intercept[SparkException] { + Seq((1L, "a")).toDF("id", "name").write.format(format).option("table", table) + .option("fail", "true").mode("append").save() + } + assert(Utils.exceptionString(e).contains("injected failure")) + assert(TestColumnarSink.aborted.contains(table)) + assert(!TestColumnarSink.committed.containsKey(table)) + } + + test("streaming read and write") { + val table = "streaming_write" + val query = spark.readStream.format(format).option("end", "5").load() + .writeStream + .format(format) + .option("table", table) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try { + query.processAllAvailable() + } finally { + query.stop() + } + assert(TestColumnarSink.rows(table).sorted == (0L until 5L).map(id => (id, s"name$id"))) + assert(TestColumnarSink.epochs.get(table) == Seq(0L)) + } + + test("operations that the data source does not implement") { + val provider = classOf[TestEmptyProvider].getName + val description = classOf[TestEmptyDataSource].getName + checkError( + exception = intercept[SparkUnsupportedOperationException] { + spark.read.format(provider).load() + }, + condition = "UNABLE_TO_INFER_SCHEMA", + parameters = Map("format" -> description)) + checkError( + exception = intercept[SparkUnsupportedOperationException] { + spark.read.schema("id BIGINT").format(provider).load().collect() + }, + condition = "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", + parameters = Map("description" -> description)) + checkError( + exception = intercept[SparkUnsupportedOperationException] { + spark.range(1).write.format(provider).mode("append").save() + }, + condition = "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", + parameters = Map("description" -> description)) + + def checkStreaming(condition: String)(start: => StreamingQuery): Unit = { + val e = intercept[Exception] { + val query = start + try query.processAllAvailable() finally query.stop() + } + checkError( + exception = findError(e, condition), + condition = condition, + parameters = Map("description" -> description)) + } + checkStreaming("DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { + spark.readStream.schema("id BIGINT").format(provider).load() + .writeStream.format("noop").start() + } + checkStreaming("DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { + spark.readStream.format("rate").load().writeStream.format(provider) + .option("checkpointLocation", Utils.createTempDir().getPath).start() + } + } +} + +/** Reads the ids [0, end) and their names, and writes to [[TestColumnarSink]]. */ +class TestRangeProvider extends DataSourceProvider { + override def shortName(): String = "test_columnar_range" + + override def createDataSource(options: CaseInsensitiveStringMap): DataSource = { + new TestRangeDataSource(options) + } +} + +class TestRangeDataSource(options: CaseInsensitiveStringMap) extends DataSource { + private val end = options.getLong("end", 10) + private val table = options.getOrDefault("table", "default") + + override def schema(): StructType = StructType.fromDDL("id BIGINT, name STRING") + + override def reader(schema: StructType): DataSourceReader = new TestRangeReader(end, schema) + + override def streamReader(schema: StructType): DataSourceStreamReader = { + new TestRangeStreamReader(end, schema) + } + + override def writer(schema: StructType, overwrite: Boolean): DataSourceWriter = { + new TestSinkWriter(table, overwrite, options.getBoolean("fail", false)) + } + + override def streamWriter(schema: StructType, overwrite: Boolean): DataSourceStreamWriter = { + new TestSinkStreamWriter(table) + } +} + +case class TestRangePartition(lo: Long, hi: Long) extends InputPartition + +/** Accepts the predicates `id > `. */ +class TestRangeReader(end: Long, schema: StructType) extends DataSourceReader { + private var lo = 0L + private var limit = Long.MaxValue + private var readSchema = schema + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + predicates.filterNot { predicate => + (predicate.name(), predicate.children()) match { + case (">", Array(column: NamedReference, literal: Literal[_])) + if column.fieldNames().sameElements(Array("id")) => + lo = math.max(lo, literal.value().asInstanceOf[Long] + 1) + true + case _ => false + } + } + } + + override def pushLimit(limit: Int): Boolean = { + this.limit = limit + true + } + + override def pruneColumns(requiredSchema: StructType): Boolean = { + readSchema = requiredSchema + true + } + + override def partitions(): Array[InputPartition] = { + val hi = if (limit == Long.MaxValue) end else math.min(end, lo + limit) + val mid = (lo + hi) / 2 + Array(TestRangePartition(lo, mid), TestRangePartition(mid, hi)) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + val range = partition.asInstanceOf[TestRangePartition] + new TestBatchReader(range.lo, range.hi, readSchema) + } +} + +class TestRangeStreamReader(end: Long, schema: StructType) extends DataSourceStreamReader { + private def offset(json: String): Long = json.stripPrefix("{\"offset\":").stripSuffix("}").toLong + + override def initialOffset(): String = "{\"offset\":0}" + + override def latestOffset(): String = s"""{"offset":$end}""" + + override def partitions(start: String, end: String): Array[InputPartition] = { + Array(TestRangePartition(offset(start), offset(end))) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + val range = partition.asInstanceOf[TestRangePartition] + new TestBatchReader(range.lo, range.hi, schema) + } +} + +/** Returns the ids [lo, hi) and their names in a single batch. */ +class TestBatchReader(lo: Long, hi: Long, schema: StructType) + extends PartitionReader[ColumnarBatch] { + private var consumed = false + private lazy val batch = { + val numRows = (hi - lo).toInt + val vectors = OnHeapColumnVector.allocateColumns(math.max(numRows, 1), schema) + for (i <- 0 until numRows; (name, column) <- schema.fieldNames.zipWithIndex) { + name match { + case "id" => vectors(column).putLong(i, lo + i) + case "name" => vectors(column).putByteArray(i, s"name${lo + i}".getBytes(UTF_8)) + } + } + new ColumnarBatch(vectors.map(v => v: ColumnVector), numRows) + } + + override def next(): Boolean = { + val hasNext = !consumed + consumed = true + hasNext + } + + override def get(): ColumnarBatch = batch + + override def close(): Unit = if (consumed) batch.close() +} + +/** The data written by the test writers, by table. */ +object TestColumnarSink { + val committed = new ConcurrentHashMap[String, Seq[TestCommitMessage]]() + val aborted = ConcurrentHashMap.newKeySet[String]() + val epochs = new ConcurrentHashMap[String, Seq[Long]]() + + def rows(table: String): Seq[(Long, String)] = committed.get(table).flatMap(_.rows) + + def batchSizes(table: String): Seq[Int] = committed.get(table).flatMap(_.batchSizes) + + def commit(table: String, messages: Array[WriterCommitMessage], overwrite: Boolean): Unit = { + val newMessages = messages.map(_.asInstanceOf[TestCommitMessage]).toSeq + committed.compute(table, (_, old) => { + if (old == null || overwrite) newMessages else old ++ newMessages + }) + } +} + +case class TestCommitMessage(rows: Seq[(Long, String)], batchSizes: Seq[Int]) + extends WriterCommitMessage + +class TestDataWriter(fail: Boolean) extends DataWriter[ColumnarBatch] { + private val rows = ArrayBuffer.empty[(Long, String)] + private val batchSizes = ArrayBuffer.empty[Int] + + override def write(batch: ColumnarBatch): Unit = { + if (fail) { + throw new IllegalStateException("injected failure") + } + assert((0 until batch.numCols()).forall(batch.column(_).isInstanceOf[ArrowColumnVector])) + batchSizes += batch.numRows() + (0 until batch.numRows()).foreach { i => + rows += ((batch.column(0).getLong(i), batch.column(1).getUTF8String(i).toString)) + } + } + + override def commit(): WriterCommitMessage = TestCommitMessage(rows.toSeq, batchSizes.toSeq) + + override def abort(): Unit = {} + + override def close(): Unit = {} +} + +class TestSinkWriter(table: String, overwrite: Boolean, fail: Boolean) extends DataSourceWriter { + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new TestDataWriter(fail) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = { + TestColumnarSink.commit(table, messages, overwrite) + } + + override def abort(messages: Array[WriterCommitMessage]): Unit = { + TestColumnarSink.aborted.add(table) + } +} + +class TestSinkStreamWriter(table: String) extends DataSourceStreamWriter { + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[ColumnarBatch] = new TestDataWriter(fail = false) + + override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + TestColumnarSink.commit(table, messages, overwrite = false) + TestColumnarSink.epochs.compute(table, (_, old) => Option(old).getOrElse(Nil) :+ epochId) + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + TestColumnarSink.aborted.add(table) + } +} + +/** A data source that implements nothing. */ +class TestEmptyProvider extends DataSourceProvider { + override def shortName(): String = "test_columnar_empty" + + override def createDataSource(options: CaseInsensitiveStringMap): DataSource = { + new TestEmptyDataSource + } +} + +class TestEmptyDataSource extends DataSource diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala new file mode 100644 index 0000000000000..9edbf1d2357ac --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -0,0 +1,604 @@ +/* + * 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.execution.datasources.v2.ffi + +import java.io.{File, FileOutputStream} +import java.nio.charset.StandardCharsets.UTF_8 +import java.nio.file.Files +import java.sql.{Date, Timestamp} +import java.time.{Instant, LocalDate} +import java.util.concurrent.atomic.AtomicInteger +import java.util.zip.{ZipEntry, ZipOutputStream} + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.spark.{SparkRuntimeException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.sql.{AnalysisException, QueryTest, Row} +import org.apache.spark.sql.classic.{DataFrame, SparkSession} +import org.apache.spark.sql.connector.expressions.{Expression, FieldReference, GeneralScalarExpression, LiteralValue} +import org.apache.spark.sql.connector.expressions.filter.{And, Predicate} +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation +import org.apache.spark.sql.execution.datasources.v2.columnar.ColumnarScan +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.sql.streaming.StreamingQuery +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.types.UTF8String +import org.apache.spark.util.Utils + +/** + * Tests native data sources with the library in `native-datasource/test_native_datasource.cc`, + * compiled with the C++ compiler of the machine. The tests that load a library are skipped when + * there is no C++ compiler or no JNI headers. + */ +class NativeDataSourceSuite extends QueryTest with SharedSparkSession { + + private lazy val buildDir = Utils.createTempDir(namePrefix = "native-datasource-test") + + // The executors store the files of the sessions that are not isolated in the same directory, + // and jobs outside of SQL executions get the files of all the sessions. Each package has its + // own name, and the sessions are kept until the end of the suite, so that none of their + // artifacts is removed while another test runs. + private val packageCount = new AtomicInteger() + private val sessions = ArrayBuffer.empty[SparkSession] + + // The C++ compiler and the JNI include directories, if available. + private lazy val toolchain: Option[(String, Seq[String])] = { + val include = new File(System.getProperty("java.home"), "include") + val platformInclude = Option(include.listFiles()).toSeq.flatten + .find(dir => new File(dir, "jni_md.h").isFile) + val compiler = Seq("c++", "g++", "clang++").find(Utils.checkCommandAvailable) + for { + compiler <- compiler + platformInclude <- platformInclude + if new File(include, "jni.h").isFile + } yield (compiler, Seq(include.getPath, platformInclude.getPath)) + } + + private def compile(name: String, defines: String*): File = { + val (compiler, includes) = toolchain.get + val source = new File(buildDir, "test_native_datasource.cc") + if (!source.exists()) { + val in = getClass.getClassLoader + .getResourceAsStream("native-datasource/test_native_datasource.cc") + Utils.tryWithResource(in)(Files.copy(_, source.toPath)) + } + val library = new File(buildDir, System.mapLibraryName(name)) + val command = Seq(compiler, "-std=c++17", "-O1", "-shared", "-fPIC") ++ + includes.map("-I" + _) ++ defines ++ Seq("-o", library.getPath, source.getPath) + val process = new ProcessBuilder(command: _*).redirectErrorStream(true).start() + val output = new String(process.getInputStream.readAllBytes(), UTF_8) + assert(process.waitFor() == 0, s"Failed to compile the test library: $output") + library + } + + // The test library and its variants, each implementing differently named data sources. + private lazy val library = compile("test_native_datasource") + private lazy val readOnlyLibrary = + compile("test_native_datasource_ro", "-DREAD_ONLY", "-DNAME_PREFIX=\"ro_\"") + private lazy val negatingLibrary = + compile("test_native_datasource_b", "-DNAME_PREFIX=\"b_\"", "-DID_SIGN=-1") + private lazy val abiV2Library = + compile("test_native_datasource_v2", "-DABI_VERSION=2", "-DNAME_PREFIX=\"v2_\"") + + private def dataSources(prefix: String): Seq[String] = + Seq("native_range", "native_sink", "native_counter").map(prefix + _) + + private def manifestJson( + dataSources: Seq[String], + libraries: Map[String, String], + abiVersion: Int = 1): String = { + val names = dataSources.map(n => "\"" + n + "\"").mkString(", ") + val paths = libraries.map { case (p, path) => "\"" + p + "\": \"" + path + "\"" } + .mkString(", ") + s"""{"abiVersion": $abiVersion, "dataSources": [$names], "libraries": {$paths}}""" + } + + private def createZip(dir: File, fileName: String, entries: (String, Array[Byte])*): File = { + dir.mkdirs() + val file = new File(dir, fileName) + Utils.tryWithResource(new ZipOutputStream(new FileOutputStream(file))) { zip => + entries.foreach { case (path, bytes) => + zip.putNextEntry(new ZipEntry(path)) + zip.write(bytes) + zip.closeEntry() + } + } + file + } + + /** Creates a package with the library built for this platform. */ + private def createPackage( + library: File, + dataSources: Seq[String], + dir: File = Utils.createTempDir(), + fileName: String = s"test-${packageCount.incrementAndGet()}.sparkpkg", + abiVersion: Int = 1, + platform: String = NativePlatform.current): File = { + val path = s"$platform/${library.getName}" + createZip( + dir, + fileName, + NativeDataSourcePackage.MANIFEST_NAME -> + manifestJson(dataSources, Map(platform -> path), abiVersion).getBytes(UTF_8), + path -> Files.readAllBytes(library.toPath)) + } + + private lazy val defaultPackage = createPackage(library, dataSources("")) + + /** A new session that finds the native data source packages at the given paths. */ + // Like Python data sources, native data sources are found in the active session. + private def newSession(paths: File*): SparkSession = { + val session = spark.newSession() + SparkSession.setActiveSession(session) + sessions += session + session.conf.set(SQLConf.NATIVE_DATA_SOURCE_PATHS.key, paths.map(_.getPath).mkString(",")) + session + } + + private def nativeTest(name: String)(f: => Unit): Unit = test(name) { + assume(toolchain.isDefined, "A C++ compiler and the JNI headers are required.") + try f finally SparkSession.setActiveSession(spark) + } + + // Checks the answer without converting the DataFrame to an RDD: a job outside of a SQL + // execution gets the files of all the sessions, and the streaming queries of the tests create + // sessions whose files are removed when they are garbage collected. + private def checkRows(df: => DataFrame, expected: Seq[Row]): Unit = { + QueryTest.checkAnswer(df, expected, checkToRDD = false) + } + + private def range(session: SparkSession, options: (String, String)*): DataFrame = { + session.read.format("native_range").options(options.toMap).load() + } + + private def scan(df: DataFrame): ColumnarScan = { + df.queryExecution.optimizedPlan.collectFirst { + case relation: DataSourceV2ScanRelation => relation.scan.asInstanceOf[ColumnarScan] + }.get + } + + private def expectedRow(id: Long): Row = Row( + id, + (id % 3).toInt, + if (id % 5 == 0) null else s"n$id", + id * 0.5, + id % 2 == 0, + Date.valueOf(LocalDate.ofEpochDay(id)), + Timestamp.from(Instant.ofEpochSecond(id))) + + /** Finds the error with the given condition among the causes of a failure. */ + private def findError(e: Throwable, condition: String): SparkThrowable = { + Iterator.iterate(e)(_.getCause).takeWhile(_ != null) + .collectFirst { case t: SparkThrowable if t.getCondition == condition => t } + .getOrElse(throw e) + } + + // Formats a double as the test library does. + private def format(d: Double): String = if (d == d.toLong) d.toLong.toString else d.toString + + test("parse manifests") { + assert(NativeDataSourcePackage.parseManifest( + """{"abiVersion": 1, "dataSources": ["a", "b"], "libraries": {"linux-x86_64": "x.so"}}""") == + Right(NativeDataSourceManifest(1, Seq("a", "b"), Map("linux-x86_64" -> "x.so")))) + Seq( + "not json" -> "The manifest is not valid JSON", + "[1]" -> "The manifest must be a JSON object.", + """{"dataSources": ["a"], "libraries": {"p": "x"}}""" -> "'abiVersion'", + """{"abiVersion": 1, "dataSources": [], "libraries": {"p": "x"}}""" -> "'dataSources'", + """{"abiVersion": 1, "dataSources": [1], "libraries": {"p": "x"}}""" -> "'dataSources'", + """{"abiVersion": 1, "dataSources": ["a"]}""" -> "'libraries'", + """{"abiVersion": 1, "dataSources": ["a"], "libraries": {"p": 1}}""" -> "'libraries'" + ).foreach { case (json, error) => + val result = NativeDataSourcePackage.parseManifest(json) + assert(result.isLeft && result.swap.toOption.get.contains(error), s"$json: $result") + } + } + + test("platform names") { + assert(NativePlatform.name("Linux", "amd64") == "linux-x86_64") + assert(NativePlatform.name("Linux", "aarch64") == "linux-aarch64") + assert(NativePlatform.name("Mac OS X", "aarch64") == "osx-aarch64") + assert(NativePlatform.name("Mac OS X", "x86_64") == "osx-x86_64") + assert(NativePlatform.name("Windows 11", "amd64") == "windows-x86_64") + assert(NativePlatform.current == NativePlatform.name( + System.getProperty("os.name"), System.getProperty("os.arch"))) + } + + test("encode predicates as JSON") { + def predicate(name: String, children: Expression*): Predicate = { + new Predicate(name, children.toArray) + } + val id = FieldReference(Seq("id")) + assert(NativePredicates.toJson(predicate(">", id, LiteralValue(5L, LongType))).contains( + """{"type":"function","name":">","children":[{"type":"column","name":["id"]},""" + + """{"type":"literal","dataType":"bigint","value":5}]}""")) + assert(NativePredicates.toJson(new And( + predicate("IS_NULL", FieldReference(Seq("a", "b"))), + predicate("=", id, new GeneralScalarExpression("+", Array(id, LiteralValue(1, IntegerType)))) + )).contains( + """{"type":"function","name":"AND","children":[{"type":"function","name":"IS_NULL",""" + + """"children":[{"type":"column","name":["a","b"]}]},{"type":"function","name":"=",""" + + """"children":[{"type":"column","name":["id"]},{"type":"function","name":"+",""" + + """"children":[{"type":"column","name":["id"]},""" + + """{"type":"literal","dataType":"int","value":1}]}]}]}""")) + + def literal(value: Any, dataType: DataType): Option[String] = { + NativePredicates.toJson(predicate("=", id, LiteralValue(value, dataType))) + .map(json => json.substring(json.indexOf("\"dataType\""), json.length - 3)) + } + assert(literal(null, LongType).contains(""""dataType":"bigint","value":null""")) + assert(literal(true, BooleanType).contains(""""dataType":"boolean","value":true""")) + assert(literal(1.toByte, ByteType).contains(""""dataType":"tinyint","value":1""")) + assert(literal(2.toShort, ShortType).contains(""""dataType":"smallint","value":2""")) + assert(literal(1.5f, FloatType).contains(""""dataType":"float","value":1.5""")) + assert(literal(2.5, DoubleType).contains(""""dataType":"double","value":2.5""")) + assert(literal(UTF8String.fromString("a\"b"), StringType) + .contains(""""dataType":"string","value":"a\"b"""")) + assert(literal(Array[Byte](1, 2), BinaryType) + .contains(""""dataType":"binary","value":"AQI="""")) + assert(literal(Decimal("12.30"), DecimalType(10, 2)) + .contains(""""dataType":"decimal(10,2)","value":"12.30"""")) + assert(literal(10, DateType).contains(""""dataType":"date","value":10""")) + assert(literal(10L, TimestampType).contains(""""dataType":"timestamp","value":10""")) + assert(literal(10L, TimestampNTZType).contains(""""dataType":"timestamp_ntz","value":10""")) + // Values that JSON cannot represent, and types without an encoding, are not pushed down. + assert(literal(Double.NaN, DoubleType).isEmpty) + assert(literal(Float.PositiveInfinity, FloatType).isEmpty) + assert(literal(UTF8String.fromString("a"), StringType("UTF8_LCASE")).isEmpty) + assert(literal(null, ArrayType(IntegerType)).nonEmpty) + assert(literal(Array(1), ArrayType(IntegerType)).isEmpty) + } + + test("disable native data sources") { + val conf = new SQLConf + conf.setConf(SQLConf.NATIVE_DATA_SOURCE_PATHS, Seq(defaultPackage.getPath)) + conf.setConf(StaticSQLConf.NATIVE_DATA_SOURCE_ENABLED, false) + assert(NativeDataSourceRegistry.lookup("native_range", conf).isEmpty) + } + + nativeTest("read with an inferred schema") { + val session = newSession(defaultPackage) + val df = range(session, "end" -> "20") + assert(df.schema == StructType.fromDDL( + "id BIGINT, mod INT, name STRING, value DOUBLE, flag BOOLEAN, day DATE, ts TIMESTAMP")) + checkRows(df, (0L until 20L).map(expectedRow)) + // The data source names are case-insensitive, and so are the option keys. + checkRows( + session.read.format("NATIVE_RANGE").option("END", "3").load(), + (0L until 3L).map(expectedRow)) + } + + nativeTest("read in partitions and batches") { + val df = range(newSession(defaultPackage), "end" -> "50", "partitions" -> "4", + "batch_size" -> "3") + assert(df.rdd.getNumPartitions == 4) + checkRows(df, (0L until 50L).map(expectedRow)) + } + + nativeTest("read with a user-specified schema") { + val session = newSession(defaultPackage) + checkRows( + session.read.schema("id BIGINT, name STRING").format("native_range").option("end", "6") + .load(), + (0L until 6L).map(id => Row(id, if (id % 5 == 0) null else s"n$id"))) + + val e = intercept[Exception] { + session.read.schema("id INT").format("native_range").load().collect() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "read", + "name" -> "native_range", + "msg" -> "The data has the schema id BIGINT, but the expected schema is id INT.")) + } + + nativeTest("push down predicates") { + val df = range(newSession(defaultPackage), "end" -> "20").filter("id > 5 AND id <= 12") + checkRows(df, (6L to 12L).map(expectedRow)) + val pushed = scan(df).getMetaData()("PushedPredicates") + assert(pushed.contains("id > 5") && pushed.contains("id <= 12"), pushed) + + // A predicate that the library does not accept is evaluated by Spark. + val partly = range(newSession(defaultPackage), "end" -> "20") + .filter("id < 10 AND name LIKE 'n1%'") + checkRows(partly, Seq(1L).map(expectedRow)) + assert(scan(partly).getMetaData()("PushedPredicates").contains("id < 10")) + } + + nativeTest("prune columns") { + val df = range(newSession(defaultPackage), "end" -> "10").select("name", "id") + checkRows(df, (0L until 10L).map(id => Row(if (id % 5 == 0) null else s"n$id", id))) + assert(scan(df).readSchema() == StructType.fromDDL("id BIGINT, name STRING")) + + // A count reads no column at all. + assert(range(newSession(defaultPackage), "end" -> "37", "batch_size" -> "10").count() == 37) + } + + nativeTest("push down a limit") { + val df = range(newSession(defaultPackage), "end" -> "100", "partitions" -> "5").limit(3) + assert(df.collect().length == 3) + assert(scan(df).getMetaData()("PushedLimit") == "LIMIT 3") + } + + nativeTest("batch write") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + def write(mode: String, ids: Range): Unit = { + import session.implicits._ + ids.map(i => (i.toLong, s"v$i")).toDF("id", "name").repartition(2) + .write.format("native_sink").option("path", dir.getPath).mode(mode).save() + } + def written(): Set[String] = dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala).toSet + + write("append", 0 until 4) + assert(written() == (0 until 4).map(i => s"$i,v$i").toSet) + assert(new File(dir, "_SUCCESS").isFile) + write("append", 4 until 6) + assert(written() == (0 until 6).map(i => s"$i,v$i").toSet) + write("overwrite", 10 until 12) + assert(written() == (10 until 12).map(i => s"$i,v$i").toSet) + + // Nulls and the other types of the test library. + range(session, "end" -> "6").write.format("native_sink").option("path", dir.getPath) + .mode("overwrite").save() + assert(written() == (0 until 6).map { id => + val name = if (id % 5 == 0) "null" else s"n$id" + s"$id,${id % 3},$name,${format(id * 0.5)},${id % 2 == 0},$id,${id * 1000000L}" + }.toSet) + } + + nativeTest("abort a failed write") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val e = intercept[Exception] { + session.range(10).write.format("native_sink").option("path", dir.getPath) + .option("fail", "write").mode("append").save() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "write to", + "name" -> "native_sink", + "msg" -> "injected failure in write")) + assert(new File(dir, "_ABORTED").isFile) + assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) + } + + nativeTest("streaming read and write") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val commitPath = new File(Utils.createTempDir(), "committed") + val query = session.readStream.format("native_counter") + .option("max_offset", "12").option("step", "5").option("batch_size", "2") + .option("commit_path", commitPath.getPath) + .load() + .writeStream + .format("native_sink") + .option("path", dir.getPath) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try { + query.processAllAvailable() + } finally { + query.stop() + } + // Each call to latestOffset makes 5 more rows available: [0, 5), [5, 10) and [10, 12). + assert((0 to 2).forall(epoch => new File(dir, s"_EPOCH_$epoch").isFile)) + val written = dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala) + assert(written.sorted.toSeq == (0 until 12).map(_.toString).sorted) + // Spark commits the offsets of a micro-batch when it runs the next one. + assert(Files.readString(commitPath.toPath) == "{\"offset\":10}") + + // The streaming source can also be read in micro-batches into a memory sink. + val memoryQuery = session.readStream.format("native_counter").option("max_offset", "4") + .load().writeStream.format("memory").queryName("native_counter_memory").start() + try { + memoryQuery.processAllAvailable() + checkRows(session.table("native_counter_memory"), (0L until 4L).map(Row(_))) + } finally { + memoryQuery.stop() + } + } + + nativeTest("report the errors of native code") { + def check(action: String, msg: String, name: String = "native_range")( + f: SparkSession => Unit): Unit = { + val e = intercept[Exception](f(newSession(defaultPackage))) + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map("action" -> action, "name" -> name, "msg" -> msg)) + } + check("create", "injected failure in create") { session => + range(session, "fail" -> "create") + } + check("infer the schema of", "injected failure in schema") { session => + range(session, "fail" -> "schema") + } + check("plan a scan of", "injected failure in plan") { session => + range(session, "fail" -> "plan").collect() + } + check("read", "injected failure in open") { session => + range(session, "fail" -> "open").collect() + } + check("read", "injected failure in get_next") { session => + range(session, "fail" -> "read").collect() + } + // An error of the library is not an unsupported operation. + check("plan a scan of", "native_sink does not support batch reads", "native_sink") { + session => session.read.schema("id BIGINT").format("native_sink").load().collect() + } + } + + nativeTest("operations that the library does not implement") { + val session = newSession(createPackage(readOnlyLibrary, dataSources("ro_"))) + // The library does not push anything down, so Spark evaluates the filter. + val df = session.read.format("ro_native_range").option("end", "10").load().filter("id > 6") + checkRows(df, (7L until 10L).map(expectedRow)) + assert(scan(df).getMetaData()("PushedPredicates") == "[]") + + checkError( + exception = intercept[SparkUnsupportedOperationException] { + session.range(1).write.format("ro_native_sink").mode("append").save() + }, + condition = "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", + parameters = Map("description" -> "ro_native_sink")) + def checkStreaming(condition: String, name: String)(start: => StreamingQuery): Unit = { + val e = intercept[Exception] { + val query = start + try query.processAllAvailable() finally query.stop() + } + checkError( + exception = findError(e, condition), + condition = condition, + parameters = Map("description" -> name)) + } + checkStreaming("DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", "ro_native_counter") { + session.readStream.format("ro_native_counter").load().writeStream.format("noop").start() + } + checkStreaming("DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", "ro_native_sink") { + session.readStream.format("rate").load().writeStream.format("ro_native_sink") + .option("checkpointLocation", Utils.createTempDir().getPath).start() + } + } + + nativeTest("load libraries that export the same functions side by side") { + val session = newSession( + defaultPackage, createPackage(negatingLibrary, dataSources("b_"))) + checkRows(range(session, "end" -> "3").select("id"), Seq(Row(0L), Row(1L), Row(2L))) + checkRows( + session.read.format("b_native_range").option("end", "3").load().select("id"), + Seq(Row(0L), Row(-1L), Row(-2L))) + } + + nativeTest("add packages with spark.addArtifact") { + val session = newSession() + session.addArtifact(defaultPackage.getPath) + val (added, artifactUUID) = session.artifactManager.getNativeDataSourcePackages + assert(added.map(_.getName) == Seq(defaultPackage.getName)) + checkRows(range(session, "end" -> "4"), (0L until 4L).map(expectedRow)) + + // The executors find the package in the files of the session. + val pkg = NativeDataSourcePackage.read(added.head).copy(artifactUUID = artifactUUID) + assert(NativeDataSourceRegistry.distributedCopies(pkg).exists(_.isFile)) + + // A package under spark.sql.dataSource.native.paths is added as an artifact when it is used, + // except in local mode, where the executors run in the driver. + val other = newSession(defaultPackage) + range(other, "end" -> "1").collect() + assert(other.artifactManager.getNativeDataSourcePackages._1.isEmpty) + val distributed = + NativeDataSourceRegistry.addToArtifacts(other, NativeDataSourcePackage.read(defaultPackage)) + val (otherPackages, otherArtifactUUID) = other.artifactManager.getNativeDataSourcePackages + assert(otherPackages.map(_.getName) == Seq(defaultPackage.getName)) + assert(distributed.artifactUUID == otherArtifactUUID) + } + + nativeTest("find packages in directories") { + val dir = Utils.createTempDir() + createPackage(library, dataSources(""), dir = dir, fileName = "a.sparkpkg") + // Files without the extension are ignored. + createZip(dir, "ignored.zip", "x" -> Array.emptyByteArray) + checkRows(range(newSession(dir), "end" -> "2"), (0L until 2L).map(expectedRow)) + } + + nativeTest("conflicting packages") { + val dir = Utils.createTempDir() + val first = createPackage(library, dataSources(""), dir = dir, fileName = "a.sparkpkg") + val second = createPackage(library, Seq("native_range"), dir = dir, fileName = "b.sparkpkg") + checkError( + exception = intercept[AnalysisException](range(newSession(dir))), + condition = "NATIVE_DATA_SOURCE_PACKAGE_CONFLICT", + parameters = Map( + "name" -> "native_range", + "packages" -> Seq(first, second).map(_.getCanonicalPath).mkString(", "))) + // The same package found twice is not a conflict: here, in the configured paths and in the + // artifacts of the session. + val session = newSession(first) + session.addArtifact(first.getPath) + checkRows(range(session, "end" -> "2"), Seq(expectedRow(0L), expectedRow(1L))) + } + + nativeTest("Java data sources take precedence") { + val session = newSession(createPackage(library, Seq("json"))) + val dir = Utils.createTempDir().getPath + "/json" + session.range(2).write.json(dir) + checkRows(session.read.format("json").load(dir), Seq(Row(0L), Row(1L))) + } + + nativeTest("invalid packages") { + val dir = Utils.createTempDir() + def check(pkg: File, subClass: String, params: Map[String, String]): Unit = { + checkError( + exception = intercept[SparkRuntimeException](range(newSession(pkg)).collect()), + condition = s"INVALID_NATIVE_DATA_SOURCE_PACKAGE.$subClass", + parameters = params) + } + + val noManifest = createZip(dir, "no_manifest.sparkpkg", "x" -> Array.emptyByteArray) + check(noManifest, "INVALID_MANIFEST", Map( + "path" -> noManifest.getCanonicalPath, + "reason" -> "The package does not contain the manifest spark-native-datasource.json.")) + + val missingLibrary = createZip(dir, "missing_library.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + Seq("native_range"), Map(NativePlatform.current -> "lib.so")).getBytes(UTF_8)) + check(missingLibrary, "INVALID_MANIFEST", Map( + "path" -> missingLibrary.getCanonicalPath, + "reason" -> "The package does not contain the library lib.so.")) + + val newerAbi = createPackage(library, Seq("native_range"), dir = dir, + fileName = "newer_abi.sparkpkg", abiVersion = 2) + check(newerAbi, "UNSUPPORTED_ABI_VERSION", Map( + "path" -> newerAbi.getCanonicalPath, "version" -> "2", "supported" -> "1")) + + val otherPlatform = createPackage(library, Seq("native_range"), dir = dir, + fileName = "other_platform.sparkpkg", platform = "other-platform") + check(otherPlatform, "UNSUPPORTED_PLATFORM", Map( + "path" -> otherPlatform.getCanonicalPath, + "platform" -> NativePlatform.current, + "platforms" -> "other-platform")) + + val path = s"${NativePlatform.current}/${System.mapLibraryName("not_a_library")}" + val notALibrary = createZip(dir, "not_a_library.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + Seq("native_range"), Map(NativePlatform.current -> path)).getBytes(UTF_8), + path -> "not a library".getBytes(UTF_8)) + val e = intercept[SparkRuntimeException](range(newSession(notALibrary))) + assert(e.getCondition == "INVALID_NATIVE_DATA_SOURCE_PACKAGE.INVALID_LIBRARY") + assert(e.getMessageParameters.get("path") == notALibrary.getCanonicalPath) + assert(e.getMessageParameters.get("library") == path) + + // The library implements another version of the interface than its manifest says. + val mismatchedAbi = createPackage(abiV2Library, dataSources("v2_"), dir = dir, + fileName = "mismatched_abi.sparkpkg") + checkError( + exception = intercept[SparkRuntimeException] { + newSession(mismatchedAbi).read.format("v2_native_range").load() + }, + condition = "INVALID_NATIVE_DATA_SOURCE_PACKAGE.UNSUPPORTED_ABI_VERSION", + parameters = Map( + "path" -> mismatchedAbi.getCanonicalPath, "version" -> "2", "supported" -> "1")) + } +} From b032a0a34d504de79d41f4751b99a6dd7450036a Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 9 Oct 2026 10:28:25 +0900 Subject: [PATCH 02/27] Make NativeBridge the only public API of native data sources Downstream data sources are only native, so the JVM interfaces of the columnar data source API (DataSource, DataSourceReader, DataSourceWriter, DataSourceStreamReader, DataSourceStreamWriter and DataSourceProvider) are not public anymore. They become the internal ColumnarDataSource traits that the Data Source V2 adapter and the native data sources share, and the Javadoc of NativeBridge now specifies the whole interface on its own. The documentation page is renamed to "Native Data Sources". Co-authored-by: Isaac --- docs/_data/menu-sql.yaml | 20 +- ...-sources.md => sql-native-data-sources.md} | 179 +++++------------- docs/sql-v2-data-sources.md | 6 +- .../spark/sql/datasource/DataSource.java | 102 ---------- .../sql/datasource/DataSourceProvider.java | 65 ------- .../sql/datasource/DataSourceReader.java | 89 --------- .../datasource/DataSourceStreamReader.java | 73 ------- .../datasource/DataSourceStreamWriter.java | 69 ------- .../sql/datasource/DataSourceWriter.java | 66 ------- .../spark/sql/datasource/NativeBridge.java | 106 ++++++++--- .../v2/columnar/ColumnarDataSource.scala | 128 +++++++++++++ .../v2/columnar/ColumnarTable.scala | 32 ++-- .../v2/columnar/ColumnarWrite.scala | 15 +- .../datasources/v2/ffi/NativeDataSource.scala | 22 +-- ...pache.spark.sql.sources.DataSourceRegister | 2 +- .../columnar}/ColumnarDataSourceSuite.scala | 107 +++-------- 16 files changed, 326 insertions(+), 755 deletions(-) rename docs/{sql-columnar-data-sources.md => sql-native-data-sources.md} (69%) delete mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSource.java delete mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceProvider.java delete mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceReader.java delete mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamReader.java delete mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamWriter.java delete mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceWriter.java create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala rename sql/core/src/test/scala/org/apache/spark/sql/{datasource => execution/datasources/v2/columnar}/ColumnarDataSourceSuite.scala (77%) diff --git a/docs/_data/menu-sql.yaml b/docs/_data/menu-sql.yaml index e4b8bb243e47d..3112614546113 100644 --- a/docs/_data/menu-sql.yaml +++ b/docs/_data/menu-sql.yaml @@ -75,17 +75,19 @@ url: sql-v2-data-sources.html#streaming - text: Further Reading url: sql-v2-data-sources.html#further-reading -- text: Columnar Data Source API - url: sql-columnar-data-sources.html +- text: Native Data Sources + url: sql-native-data-sources.html subitems: - text: Overview - url: sql-columnar-data-sources.html#overview - - text: The API - url: sql-columnar-data-sources.html#the-api - - text: JVM Data Sources - url: sql-columnar-data-sources.html#jvm-data-sources - - text: Native Data Sources - url: sql-columnar-data-sources.html#native-data-sources + url: sql-native-data-sources.html#overview + - text: How It Works + url: sql-native-data-sources.html#how-it-works + - text: Packages + url: sql-native-data-sources.html#packages + - text: Example in Rust + url: sql-native-data-sources.html#example-in-rust + - text: Example in C++ + url: sql-native-data-sources.html#example-in-c - text: Performance Tuning url: sql-performance-tuning.html subitems: diff --git a/docs/sql-columnar-data-sources.md b/docs/sql-native-data-sources.md similarity index 69% rename from docs/sql-columnar-data-sources.md rename to docs/sql-native-data-sources.md index ed1fca5914783..40ae89ca6baa2 100644 --- a/docs/sql-columnar-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -1,7 +1,7 @@ --- layout: global -title: Columnar Data Source API -displayTitle: Columnar Data Source API +title: Native Data Sources +displayTitle: Native Data Sources license: | Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with @@ -24,141 +24,49 @@ license: | ## Overview -The columnar data source API is a simpler way to plug a data source into Spark than implementing -the [Data Source V2](sql-v2-data-sources.html) interfaces directly. It is meant for data sources -that identify their data by options, such as a path or a URL. It follows the -[Python Data Source API](api/python/reference/pyspark.sql/api/pyspark.sql.datasource.DataSource.html), -but data is exchanged only as columnar batches: a data source returns -[`ColumnarBatch`](api/java/org/apache/spark/sql/vectorized/ColumnarBatch.html)es when reading, and -receives them when writing. Spark adapts a data source to Data Source V2, so it supports: +A native data source is a data source implemented in native code, such as Rust or C++, without +any JVM code. It is a shared library that implements the +[`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) binary interface, and +it is distributed as a native data source package: a single file that Spark loads when a query +uses the data source. Spark adapts it to [Data Source V2](sql-v2-data-sources.html), so it +supports: - batch and micro-batch streaming reads, - batch and streaming writes, appending to or overwriting the existing data, - predicate, column and limit pushdown. -A data source can be implemented in two ways: +Data is exchanged only as columnar batches, through the +[Arrow C data interface](https://arrow.apache.org/docs/format/CDataInterface.html), so Spark reads +the data that a library returns without copying it. -- **On the JVM**, in Java or Scala, with the interfaces of the - [`org.apache.spark.sql.datasource`](api/java/org/apache/spark/sql/datasource/package-summary.html) - package. -- **In native code**, such as Rust or C++, as a shared library that implements the same API - through the [`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) binary - interface. A native data source needs no JVM code: it is distributed as a native data source - package, and Spark loads it when a query uses it. +## How It Works -## The API - -| Interface | Methods | Description | -|-----------|---------|-------------| -| `DataSource` | `schema`, `reader`, `streamReader`, `writer`, `streamWriter` | Created on the driver with the options of the user, for each schema inference, scan and write. A data source only implements the operations it supports. | -| `DataSourceReader` | `pushPredicates`, `pushLimit`, `pruneColumns`, `partitions`, `read` | Reads the data for a batch scan. Spark pushes operations down and plans the partitions on the driver, then serializes the reader to the executors, where `read` returns the batches of a partition. | -| `DataSourceStreamReader` | `initialOffset`, `latestOffset`, `partitions`, `read`, `commit`, `stop` | Reads the data for a micro-batch streaming scan. Offsets are JSON strings defined by the data source. | -| `DataSourceWriter` | `createWriter`, `commit`, `abort` | Writes the data for a batch write. `createWriter` returns a `DataWriter` for each task, and Spark commits or aborts the write on the driver with the messages of the tasks. | -| `DataSourceStreamWriter` | `createWriter`, `commit`, `abort` | Writes the data for a streaming write, one micro-batch (epoch) at a time. | - -The batches that a writer receives are backed by Apache Arrow: every column is an -[`ArrowColumnVector`](api/java/org/apache/spark/sql/vectorized/ArrowColumnVector.html). A batch has -at most `spark.sql.execution.arrow.maxRecordsPerBatch` rows, and is only valid during the call. - -## JVM Data Sources - -A JVM data source extends -[`DataSourceProvider`](api/java/org/apache/spark/sql/datasource/DataSourceProvider.html), which -creates the `DataSource` with the options given by the user. The following data source reads the -numbers `[0, end)` as a column `id`, in two partitions: - -```scala -import org.apache.arrow.memory.RootAllocator -import org.apache.arrow.vector.BigIntVector - -import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} -import org.apache.spark.sql.datasource._ -import org.apache.spark.sql.types.{LongType, StructType} -import org.apache.spark.sql.util.CaseInsensitiveStringMap -import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} - -class RangeProvider extends DataSourceProvider { - override def shortName(): String = "scala_range" - - override def createDataSource(options: CaseInsensitiveStringMap): DataSource = { - new RangeDataSource(options.getLong("end", 10)) - } -} - -class RangeDataSource(end: Long) extends DataSource { - override def schema(): StructType = new StructType().add("id", LongType, nullable = false) - - override def reader(schema: StructType): DataSourceReader = new RangeReader(end) -} - -case class RangePartition(start: Long, end: Long) extends InputPartition - -class RangeReader(end: Long) extends DataSourceReader { - override def partitions(): Array[InputPartition] = { - Array(RangePartition(0, end / 2), RangePartition(end / 2, end)) - } - - override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { - val range = partition.asInstanceOf[RangePartition] - new RangePartitionReader(range.start, range.end) - } -} - -// Returns the ids [start, end) in a single batch, backed by an Arrow vector. -class RangePartitionReader(start: Long, end: Long) extends PartitionReader[ColumnarBatch] { - private val allocator = new RootAllocator() - private val ids = new BigIntVector("id", allocator) - private var done = false - - override def next(): Boolean = { - if (done) return false - val count = (end - start).toInt - ids.allocateNew(count) - (0 until count).foreach(i => ids.set(i, start + i)) - ids.setValueCount(count) - done = true - true - } - - override def get(): ColumnarBatch = { - new ColumnarBatch(Array[ColumnVector](new ArrowColumnVector(ids)), ids.getValueCount) - } - - override def close(): Unit = { - ids.close() - allocator.close() - } -} -``` - -To use the data source by its short name, list the provider in -`META-INF/services/org.apache.spark.sql.sources.DataSourceRegister` in its JAR, like any other data -source. Its fully qualified class name can also be used as the format: - -```scala -spark.read.format("scala_range").option("end", 100).load() -``` - -## Native Data Sources - -### How It Works - -A native data source is a shared library that implements the `native` methods of -[`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) as JNI functions, +A native data source library implements the `native` methods of `NativeBridge` as JNI functions, named after the class, such as `Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource`. Spark loads each library into its own copy of the class, so any number of native data sources can -be loaded side by side. The methods follow the API above: Spark implements `DataSource`, -`DataSourceReader` and the other interfaces by calling them. - -- **Handles**: the `create` functions return a pointer to a native object, which Spark releases with - the matching `close` function. -- **Arrow**: schemas and data are exchanged with the - [Arrow C data interface](https://arrow.apache.org/docs/format/CDataInterface.html). A partition - is read by exporting an `ArrowArrayStream`, and a writer receives an `ArrowArray` for each batch. -- **State**: partitions, the state of readers and writers, and commit messages are byte arrays in a - format defined by the library. Spark sends them to the executors and back. -- **Errors**: a library reports an error by throwing a Java exception, for example with the JNI - function `ThrowNew`. Spark reports it as a `NATIVE_DATA_SOURCE_ERROR`. +be loaded side by side. + +Spark calls the functions of the library as follows: + +- **Data sources**: for each schema inference, scan and write, Spark creates a data source on the + driver with `createDataSource`, which receives the options of the user, and closes it right + after. When the user does not specify a schema, Spark asks for it with `schema`. +- **Batch reads**: on the driver, Spark creates a reader with `createReader`, pushes operations down + with `pushPredicates`, `pushLimit` and `pruneColumns`, plans the partitions with `partitions`, + and serializes the state of the reader with `serializeReader`. Then it calls `read` on the + executors, in one task for each partition, and `read` exports an `ArrowArrayStream`. +- **Streaming reads**: on the driver, Spark creates a stream reader with `createStreamReader`. For + each micro-batch, it asks for the `latestOffset`, plans the partitions with `streamPartitions`, + and calls `read` on the executors. Offsets are JSON strings defined by the library. +- **Writes**: on the driver, Spark creates a writer with `createWriter` or `createStreamWriter`. On + the executors, each task creates a data writer with `createDataWriter`, writes the batches of its + data with `write`, and calls `commitDataWriter`. Then the driver calls `commit` with the messages + of the tasks, or `abort` if the write failed. + +The library defines the format of the partitions, of the state of its readers and writers, and of +the commit messages, which Spark sends to the executors and back as byte arrays. A library reports +an error by throwing a Java exception, for example with the JNI function `ThrowNew`, and Spark +reports it as a `NATIVE_DATA_SOURCE_ERROR`. A library only exports the functions of the operations it supports. If a function is missing, Spark treats the operation as unsupported, except for the pushdown and `serialize` functions, which @@ -172,10 +80,11 @@ are optional: | Batch write | `createWriter`, `commit`, `abort`, `closeWriter`, `createDataWriter`, `write`, `commitDataWriter`, `abortDataWriter` | `serializeWriter` | | Streaming write | `createStreamWriter`, `commit`, `abort`, `closeWriter`, `createDataWriter`, `write`, `commitDataWriter`, `abortDataWriter` | `serializeWriter` | -The Javadoc of `NativeBridge` specifies each function, including the JSON format of the pushed -predicates. +The [Javadoc of `NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) +specifies each function, including the ownership of the Arrow structs and the JSON format of the +pushed predicates. -### Packages +## Packages A native data source package is a zip file with the extension `.sparkpkg`. It contains a manifest named `spark-native-datasource.json`, and the library built for one or more platforms: @@ -213,7 +122,7 @@ names, for example with their version: sessions that are not isolated share the A Java data source with the same name takes precedence over a native one, and so does a Python data source. -### Example in Rust +## Example in Rust This data source reads the numbers `[0, end)` as a column `id`, in two partitions. It uses the [`jni`](https://crates.io/crates/jni) crate, and the `arrow-array` and `arrow-schema` crates of @@ -417,7 +326,7 @@ EOF (cd package && zip -r ../rust_range.sparkpkg .) ``` -### Example in C++ +## Example in C++ The same data source in C++, with [Arrow C++](https://arrow.apache.org/docs/cpp/), which exports schemas and streams with `arrow::ExportSchema` and `arrow::ExportRecordBatchReader`: @@ -558,14 +467,14 @@ c++ -std=c++20 -O2 -shared -fPIC \ A library loaded by Spark must find its dependencies, such as `libarrow`, on every node: link them statically, or install them on the nodes. -### Configuration +## Configuration | Property Name | Default | Meaning | |---------------|---------|---------| | `spark.sql.dataSource.native.enabled` | true | Whether Spark loads native data sources. It is a static configuration, so it can only be set when the Spark application starts. | | `spark.sql.dataSource.native.paths` | (none) | Comma-separated list of native data source packages, and of directories that contain them, on the local file system of the driver. | -### Security +## Security A native data source runs native code in the driver and executor processes, with their privileges, and is not sandboxed. Only add packages that you trust. Set diff --git a/docs/sql-v2-data-sources.md b/docs/sql-v2-data-sources.md index 68ea6298e732a..df1674e98ea1c 100644 --- a/docs/sql-v2-data-sources.md +++ b/docs/sql-v2-data-sources.md @@ -40,7 +40,7 @@ Compared to the earlier Data Source V1 API, DSV2 offers: | Feature | Description | |---------|-------------| -| **Java API** | The connector interfaces are pure Java (`org.apache.spark.sql.connector`), removing the Scala dependency that DSV1 required. A [Python Data Source API](api/python/reference/pyspark.sql/api/pyspark.sql.datasource.DataSource.html) (`pyspark.sql.datasource`) is also available as a wrapper for lightweight connectors written entirely in Python, and the [Columnar Data Source API](sql-columnar-data-sources.html) (`org.apache.spark.sql.datasource`) for lightweight columnar connectors written in Java, Scala, or native code such as Rust or C++. | +| **Java API** | The connector interfaces are pure Java (`org.apache.spark.sql.connector`), removing the Scala dependency that DSV1 required. A [Python Data Source API](api/python/reference/pyspark.sql/api/pyspark.sql.datasource.DataSource.html) (`pyspark.sql.datasource`) is also available as a wrapper for lightweight connectors written entirely in Python. [Native data sources](sql-native-data-sources.html) can be written entirely in native code, such as Rust or C++. | | **Catalog integration** | Connectors can expose namespaces, tables, views, and functions natively through Spark SQL. | | **Operator pushdown** | Connectors can accept pushed-down filters, required columns, aggregates, limits, offsets, and more. | | **Report partitioning and ordering** | Connectors can report the physical layout of data so that Spark can avoid unnecessary shuffles and sorts. | @@ -473,7 +473,7 @@ A `Table` that declares `STREAMING_WRITE` supports streaming writes through - [Data Sources](sql-data-sources.html) for the user-facing guide to built-in data sources (DSV1). - [Python Data Source API](api/python/reference/pyspark.sql/api/pyspark.sql.datasource.DataSource.html) for writing lightweight connectors entirely in Python. -- [Columnar Data Source API](sql-columnar-data-sources.html) for writing lightweight columnar - connectors in Java, Scala, or native code such as Rust or C++. +- [Native Data Sources](sql-native-data-sources.html) for writing connectors entirely in native + code, such as Rust or C++. - [Storage Partition Join](sql-performance-tuning.html#storage-partition-join) for how DSV2 partitioning reporting enables join optimizations. diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSource.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSource.java deleted file mode 100644 index 005f645543903..0000000000000 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSource.java +++ /dev/null @@ -1,102 +0,0 @@ -/* - * 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; - -import java.util.Map; - -import org.apache.spark.SparkUnsupportedOperationException; -import org.apache.spark.annotation.Evolving; -import org.apache.spark.sql.types.StructType; - -/** - * A data source implemented with the columnar data source API. - *

- * The columnar data source API is a simpler way to plug a data source into Spark than - * implementing the Data Source V2 interfaces directly. It follows the Python Data Source API - * ({@code pyspark.sql.datasource}), but data is exchanged only as columnar batches: a source - * returns {@link org.apache.spark.sql.vectorized.ColumnarBatch}es when reading and receives them - * when writing. Spark adapts the source to Data Source V2, so it supports batch and streaming - * reads and writes, and filter, column and limit pushdown. - *

- * A data source is implemented either on the JVM, registered through a - * {@link DataSourceProvider}, or in native code such as Rust or C++, through the - * {@link NativeBridge} binary interface. - *

- * Spark creates a new instance on the driver for each schema inference, scan and write. A data - * source only needs to implement the operations it supports; the default implementations throw an - * error saying that the operation is not supported. - * - * @since 4.4.0 - */ -@Evolving -public interface DataSource { - - /** - * Returns the schema of the data. Spark calls it only when the user does not specify a schema. - */ - default StructType schema() { - throw new SparkUnsupportedOperationException( - "UNABLE_TO_INFER_SCHEMA", Map.of("format", getClass().getName())); - } - - /** - * Returns a reader for a batch scan. - * - * @param schema the schema to read: the one the user specified, or the one returned by - * {@link #schema()} - */ - default DataSourceReader reader(StructType schema) { - throw new SparkUnsupportedOperationException( - "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map.of("description", getClass().getName())); - } - - /** - * Returns a reader for a micro-batch streaming scan. - * - * @param schema the schema to read: the one the user specified, or the one returned by - * {@link #schema()} - */ - default DataSourceStreamReader streamReader(StructType schema) { - throw new SparkUnsupportedOperationException( - "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", Map.of("description", getClass().getName())); - } - - /** - * Returns a writer for a batch write. - * - * @param schema the schema of the data to write - * @param overwrite whether to replace the existing data (save mode "overwrite") instead of - * appending to it - */ - default DataSourceWriter writer(StructType schema, boolean overwrite) { - throw new SparkUnsupportedOperationException( - "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", Map.of("description", getClass().getName())); - } - - /** - * Returns a writer for a streaming write. - * - * @param schema the schema of the data to write - * @param overwrite whether to replace the existing data in each micro-batch (output mode - * "complete") instead of appending to it - */ - default DataSourceStreamWriter streamWriter(StructType schema, boolean overwrite) { - throw new SparkUnsupportedOperationException( - "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", Map.of("description", getClass().getName())); - } -} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceProvider.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceProvider.java deleted file mode 100644 index 4cfd25e351b4a..0000000000000 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceProvider.java +++ /dev/null @@ -1,65 +0,0 @@ -/* - * 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; - -import java.util.Map; - -import org.apache.spark.annotation.Evolving; -import org.apache.spark.sql.connector.catalog.Table; -import org.apache.spark.sql.connector.catalog.TableProvider; -import org.apache.spark.sql.connector.expressions.Transform; -import org.apache.spark.sql.execution.datasources.v2.columnar.ColumnarTable; -import org.apache.spark.sql.sources.DataSourceRegister; -import org.apache.spark.sql.types.StructType; -import org.apache.spark.sql.util.CaseInsensitiveStringMap; - -/** - * The entry point of a {@link DataSource} implemented on the JVM. - *

- * Register a provider like any other data source: list the class in - * {@code META-INF/services/org.apache.spark.sql.sources.DataSourceRegister} to use it by its - * {@link #shortName()}, for example {@code spark.read.format("my_source")}, or use its fully - * qualified class name as the format. - * - * @since 4.4.0 - */ -@Evolving -public abstract class DataSourceProvider implements TableProvider, DataSourceRegister { - - /** - * Creates the data source with the options given by the user, such as - * {@code spark.read.option("key", "value")}. - */ - public abstract DataSource createDataSource(CaseInsensitiveStringMap options); - - @Override - public final StructType inferSchema(CaseInsensitiveStringMap options) { - return createDataSource(options).schema(); - } - - @Override - public final Table getTable( - StructType schema, Transform[] partitioning, Map properties) { - return new ColumnarTable(shortName(), schema, this::createDataSource); - } - - @Override - public final boolean supportsExternalMetadata() { - return true; - } -} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceReader.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceReader.java deleted file mode 100644 index 6a04e003e8047..0000000000000 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceReader.java +++ /dev/null @@ -1,89 +0,0 @@ -/* - * 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; - -import java.io.Serializable; - -import org.apache.spark.annotation.Evolving; -import org.apache.spark.sql.connector.expressions.filter.Predicate; -import org.apache.spark.sql.connector.read.InputPartition; -import org.apache.spark.sql.connector.read.PartitionReader; -import org.apache.spark.sql.types.StructType; -import org.apache.spark.sql.vectorized.ColumnarBatch; - -/** - * Reads the data of a {@link DataSource} for a batch scan. - *

- * Spark creates the reader on the driver with {@link DataSource#reader}. It then pushes - * operations down by calling, each at most once and in this order, {@link #pushPredicates}, - * {@link #pushLimit} and {@link #pruneColumns}, and plans the scan with {@link #partitions()}. - * After that, the reader is serialized and sent to the executors, where {@link #read} is called - * once for each partition. Any state that {@link #read} needs must be set by the time - * {@link #partitions()} returns. - * - * @since 4.4.0 - */ -@Evolving -public interface DataSourceReader extends Serializable { - - /** - * Pushes predicates down to the data source. The predicates are combined with AND. - * - * @return the predicates that Spark still has to evaluate after reading. By default, all of - * them. The data source can still use a returned predicate to skip data. - */ - default Predicate[] pushPredicates(Predicate[] predicates) { - return predicates; - } - - /** - * Pushes a LIMIT down to the data source. - * - * @return whether the reader returns at most {@code limit} rows in total. Spark applies the - * limit again after reading either way. - */ - default boolean pushLimit(int limit) { - return false; - } - - /** - * Prunes the columns to read. - * - * @param requiredSchema the top-level fields of the read schema that the query needs, with - * their full types - * @return whether {@link #read} returns exactly the columns of {@code requiredSchema}, in that - * order. If false, {@link #read} returns all the columns of the read schema. - */ - default boolean pruneColumns(StructType requiredSchema) { - return false; - } - - /** - * Plans the partitions of the scan. Spark reads each partition in a separate task. - */ - InputPartition[] partitions(); - - /** - * Reads a partition. Called on an executor. - *

- * The columns of the returned batches must match the read schema: the schema passed to - * {@link DataSource#reader}, or the required schema if {@link #pruneColumns} returned true. A - * batch only needs to stay valid until the next call to {@link PartitionReader#next()}. - */ - PartitionReader read(InputPartition partition); -} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamReader.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamReader.java deleted file mode 100644 index ba043134a2e6a..0000000000000 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamReader.java +++ /dev/null @@ -1,73 +0,0 @@ -/* - * 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; - -import java.io.Serializable; - -import org.apache.spark.annotation.Evolving; -import org.apache.spark.sql.connector.read.InputPartition; -import org.apache.spark.sql.connector.read.PartitionReader; -import org.apache.spark.sql.vectorized.ColumnarBatch; - -/** - * Reads the data of a {@link DataSource} for a micro-batch streaming scan. - *

- * Spark creates the reader on the driver with {@link DataSource#streamReader} and keeps it for - * the lifetime of the streaming query. Offsets are JSON strings defined by the data source; Spark - * stores them in the checkpoint and passes them back unchanged. For each micro-batch, Spark plans - * the partitions between two offsets with {@link #partitions}, then serializes the reader and - * sends it to the executors, where {@link #read} is called once for each partition. - * - * @since 4.4.0 - */ -@Evolving -public interface DataSourceStreamReader extends Serializable { - - /** - * Returns the offset to start from when the query has no checkpointed offset. - */ - String initialOffset(); - - /** - * Returns the most recent offset available. - */ - String latestOffset(); - - /** - * Plans the partitions that read the data after {@code start} up to and including {@code end}. - */ - InputPartition[] partitions(String start, String end); - - /** - * Reads a partition. Called on an executor. The columns of the returned batches must match the - * schema passed to {@link DataSource#streamReader}. A batch only needs to stay valid until the - * next call to {@link PartitionReader#next()}. - */ - PartitionReader read(InputPartition partition); - - /** - * Tells the data source that Spark has processed all the data up to and including {@code end} - * and will not ask for it again. - */ - default void commit(String end) {} - - /** - * Stops the reader when the streaming query terminates. - */ - default void stop() {} -} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamWriter.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamWriter.java deleted file mode 100644 index f18ba7bae2e50..0000000000000 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceStreamWriter.java +++ /dev/null @@ -1,69 +0,0 @@ -/* - * 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; - -import java.io.Serializable; - -import org.apache.spark.annotation.Evolving; -import org.apache.spark.sql.connector.write.DataWriter; -import org.apache.spark.sql.connector.write.WriterCommitMessage; -import org.apache.spark.sql.vectorized.ColumnarBatch; - -/** - * Writes data to a {@link DataSource} for a streaming write. - *

- * Spark creates the writer on the driver with {@link DataSource#streamWriter} and keeps it for - * the lifetime of the streaming query. For each micro-batch (epoch), it serializes the writer and - * sends it to the executors, where {@link #createWriter} is called once for each task. When all - * the tasks of the micro-batch succeed, Spark calls {@link #commit} on the driver with the - * messages returned by the task writers. Otherwise, it calls {@link #abort}. - *

- * As with {@link DataSourceWriter}, the batches passed to {@link DataWriter#write} are backed by - * Apache Arrow and are only valid during the call. - * - * @since 4.4.0 - */ -@Evolving -public interface DataSourceStreamWriter extends Serializable { - - /** - * Creates the writer for one task of a micro-batch. Called on an executor. - * - * @param partitionId the partition the task writes - * @param taskId the ID of the task attempt - * @param epochId the ID of the micro-batch - */ - DataWriter createWriter(int partitionId, long taskId, long epochId); - - /** - * Commits a micro-batch. Called on the driver when all the tasks of the micro-batch succeeded. - * - * @param epochId the ID of the micro-batch - * @param messages the messages returned by {@link DataWriter#commit()} of each task - */ - void commit(long epochId, WriterCommitMessage[] messages); - - /** - * Aborts a micro-batch. Called on the driver when the micro-batch failed. - * - * @param epochId the ID of the micro-batch - * @param messages the messages returned by the tasks that committed. An element is null if its - * task did not commit. - */ - void abort(long epochId, WriterCommitMessage[] messages); -} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceWriter.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceWriter.java deleted file mode 100644 index ce1cc320a22d0..0000000000000 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/DataSourceWriter.java +++ /dev/null @@ -1,66 +0,0 @@ -/* - * 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; - -import java.io.Serializable; - -import org.apache.spark.annotation.Evolving; -import org.apache.spark.sql.connector.write.DataWriter; -import org.apache.spark.sql.connector.write.WriterCommitMessage; -import org.apache.spark.sql.vectorized.ColumnarBatch; - -/** - * Writes data to a {@link DataSource} for a batch write. - *

- * Spark creates the writer on the driver with {@link DataSource#writer}, then serializes it and - * sends it to the executors, where {@link #createWriter} is called once for each task. When all - * the tasks succeed, Spark calls {@link #commit} on the driver with the messages returned by the - * task writers. Otherwise, it calls {@link #abort}. - *

- * The batches passed to {@link DataWriter#write} are backed by Apache Arrow: every column is an - * {@link org.apache.spark.sql.vectorized.ArrowColumnVector}. A batch is only valid during the - * call, so a writer that keeps data after the call returns must copy it. - * - * @since 4.4.0 - */ -@Evolving -public interface DataSourceWriter extends Serializable { - - /** - * Creates the writer for one task. Called on an executor. - * - * @param partitionId the partition the task writes - * @param taskId the ID of the task attempt - */ - DataWriter createWriter(int partitionId, long taskId); - - /** - * Commits the write. Called on the driver when all the tasks succeeded. - * - * @param messages the messages returned by {@link DataWriter#commit()} of each task - */ - void commit(WriterCommitMessage[] messages); - - /** - * Aborts the write. Called on the driver when the write failed. - * - * @param messages the messages returned by the tasks that committed. An element is null if its - * task did not commit. - */ - void abort(WriterCommitMessage[] messages); -} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java index 5570f6593f309..bee42c03847d4 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java @@ -27,16 +27,40 @@ * class as JNI functions, for example * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each * library into its own copy of this class, so any number of libraries that export the same - * functions can be loaded side by side. The methods follow the {@link DataSource} API: Spark - * implements {@link DataSource}, {@link DataSourceReader} and the other interfaces by calling - * them. + * functions can be loaded side by side. Spark adapts a native data source to Data Source V2, with + * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and + * limit pushdown. *

* The library is distributed in a native data source package: a zip file with the extension * {@code .sparkpkg} that contains a manifest named {@code spark-native-datasource.json} and the * library built for one or more platforms. Spark finds the packages added to a session with * {@code spark.addArtifact}, and the ones under the paths in - * {@code spark.sql.dataSource.native.paths}. See the "Columnar Data Source API" page of the Spark - * SQL guide for details. + * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL + * guide for details and examples. + * + *

Lifecycle

+ * Spark creates a data source with {@link #createDataSource} on the driver for each schema + * inference, scan and write, and closes it right after. + *
    + *
  • Batch reads. On the driver, Spark creates a reader with {@link #createReader}, + * pushes operations down by calling, each at most once and in this order, + * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions + * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, + * and closes it. Then it calls {@link #read} on the executors, in one task for each + * partition.
    • + *
  • Streaming reads. On the driver, Spark creates a stream reader with + * {@link #createStreamReader} and keeps it until the streaming query stops. For each + * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with + * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then + * calls {@link #read} on the executors. Offsets are JSON strings defined by the library.
    • + *
  • Writes. On the driver, Spark creates a writer with {@link #createWriter} or + * {@link #createStreamWriter}, and serializes its state with {@link #serializeWriter}. On the + * executors, each task creates a data writer with {@link #createDataWriter}, calls + * {@link #write} for each batch of at most {@code spark.sql.execution.arrow.maxRecordsPerBatch} + * rows, and calls {@link #commitDataWriter}, or {@link #abortDataWriter} if it fails. Then the + * driver calls {@link #commit} with the messages of the tasks if they all succeeded, or + * {@link #abort} otherwise. A streaming write commits or aborts each micro-batch (epoch).
    • + *
* *

Conventions

*
    @@ -45,6 +69,8 @@ * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers. A reader or * writer handle must not depend on the data source handle it was created from, because Spark * closes the data source handle right after creating it. + *
  • State. Partitions, the state of readers and writers, and commit messages are byte + * arrays in a format defined by the library. Spark sends them to the executors and back.
    • *
  • Errors. To report an error, throw a Java exception, for example with the JNI * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and return any value. Spark * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.
    • @@ -96,7 +122,8 @@ static void load(String path) { // --------------------------------------------------------------------------------------------- /** - * Creates a data source. See {@link DataSourceProvider#createDataSource}. + * Creates a data source with the options given by the user, such as + * {@code spark.read.option("key", "value")}. * * @param name the name of the data source, as listed in the package manifest. A library can * implement several data sources. @@ -108,7 +135,7 @@ public static native long createDataSource( String name, String[] optionKeys, String[] optionValues); /** - * Exports the schema of the data source. See {@link DataSource#schema()}. + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. * * @param schemaAddress the address of the {@code ArrowSchema} to initialize with a struct type * whose fields are the columns @@ -116,33 +143,39 @@ public static native long createDataSource( public static native void schema(long dataSource, long schemaAddress); /** - * Creates a reader for a batch scan. See {@link DataSource#reader}. + * Creates a reader for a batch scan. * - * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} * @return the reader handle */ public static native long createReader(long dataSource, long schemaAddress); /** - * Creates a reader for a micro-batch streaming scan. See {@link DataSource#streamReader}. + * Creates a reader for a micro-batch streaming scan. * - * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} * @return the stream reader handle */ public static native long createStreamReader(long dataSource, long schemaAddress); /** - * Creates a writer for a batch write. See {@link DataSource#writer}. + * Creates a writer for a batch write. * * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data (save mode "overwrite") instead of + * appending to it * @return the writer handle */ public static native long createWriter(long dataSource, long schemaAddress, boolean overwrite); /** - * Creates a writer for a streaming write. See {@link DataSource#streamWriter}. + * Creates a writer for a streaming write. * * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data in each micro-batch (output mode + * "complete") instead of appending to it * @return the writer handle */ public static native long createStreamWriter( @@ -156,7 +189,7 @@ public static native long createStreamWriter( // --------------------------------------------------------------------------------------------- /** - * Pushes predicates down to the reader. See {@link DataSourceReader#pushPredicates}. + * Pushes predicates down to the reader. The predicates are combined with AND. *

    * Each predicate is a JSON expression tree. An expression is one of: *

      @@ -177,25 +210,30 @@ public static native long createStreamWriter( * Spark does not pass predicates that it cannot express this way. * * @return for each predicate, whether the reader evaluates it completely, so that Spark does - * not have to evaluate it again + * not have to evaluate it again. The reader can still use the other predicates to skip + * data. */ public static native boolean[] pushPredicates(long reader, String[] predicates); /** - * Pushes a LIMIT down to the reader. See {@link DataSourceReader#pushLimit}. + * Pushes a LIMIT down to the reader. + * + * @return whether the reader returns at most {@code limit} rows in total. Spark applies the + * limit again after reading either way. */ public static native boolean pushLimit(long reader, int limit); /** - * Prunes the columns to read. See {@link DataSourceReader#pruneColumns}. + * Prunes the columns to read. * - * @param columnNames the names of the top-level columns to read - * @return whether {@link #read} returns exactly these columns, in this order + * @param columnNames the names of the top-level columns that the query needs + * @return whether {@link #read} returns exactly these columns, in this order. If not, it + * returns all the columns of the schema to read. */ public static native boolean pruneColumns(long reader, String[] columnNames); /** - * Plans the partitions of the scan. See {@link DataSourceReader#partitions()}. + * Plans the partitions of the scan. Spark reads each partition in a separate task. * * @return the partitions, each serialized in a format defined by the library */ @@ -214,14 +252,14 @@ public static native long createStreamWriter( // Stream reader (driver) // --------------------------------------------------------------------------------------------- - /** Returns the initial offset as JSON. See {@link DataSourceStreamReader#initialOffset()}. */ + /** Returns the offset to start from when the streaming query has no checkpoint, as JSON. */ public static native String initialOffset(long streamReader); - /** Returns the latest offset as JSON. See {@link DataSourceStreamReader#latestOffset()}. */ + /** Returns the most recent offset available, as JSON. */ public static native String latestOffset(long streamReader); /** - * Plans the partitions between two offsets. See {@link DataSourceStreamReader#partitions}. + * Plans the partitions that read the data after {@code start} up to and including {@code end}. * * @return the partitions, each serialized in a format defined by the library */ @@ -233,10 +271,13 @@ public static native long createStreamWriter( */ public static native byte[] serializeStreamReader(long streamReader); - /** Commits an offset. See {@link DataSourceStreamReader#commit}. */ + /** + * Tells the library that Spark has processed all the data up to and including {@code end}, + * and will not ask for it again. + */ public static native void commitOffset(long streamReader, String end); - /** Releases a stream reader handle. See {@link DataSourceStreamReader#stop()}. */ + /** Releases a stream reader handle, when the streaming query stops. */ public static native void closeStreamReader(long streamReader); // --------------------------------------------------------------------------------------------- @@ -244,14 +285,15 @@ public static native long createStreamWriter( // --------------------------------------------------------------------------------------------- /** - * Reads a partition of a batch or streaming scan. See {@link DataSourceReader#read}. + * Reads a partition of a batch or streaming scan. * * @param readerState the state returned by {@link #serializeReader} or * {@link #serializeStreamReader}, or an empty array if there is none * @param partition a partition returned by {@link #partitions} or {@link #streamPartitions} * @param streamAddress the address of the {@code ArrowArrayStream} to initialize. Each array - * of the stream is a struct array whose fields are the columns. Spark - * calls the stream from the thread of the task that reads the partition. + * of the stream is a struct array whose fields are the columns to read: + * the columns of the schema to read, or the pruned columns. Spark calls + * the stream from the thread of the task that reads the partition. */ public static native void read(byte[] readerState, byte[] partition, long streamAddress); @@ -266,7 +308,7 @@ public static native long createStreamWriter( public static native byte[] serializeWriter(long writer); /** - * Commits a write. See {@link DataSourceWriter#commit} and {@link DataSourceStreamWriter#commit}. + * Commits a write, when all the tasks succeeded. * * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write * @param messages the messages returned by {@link #commitDataWriter} of each task @@ -274,7 +316,7 @@ public static native long createStreamWriter( public static native void commit(long writer, long epochId, byte[][] messages); /** - * Aborts a write. See {@link DataSourceWriter#abort} and {@link DataSourceStreamWriter#abort}. + * Aborts a write, when it failed. * * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write * @param messages the messages returned by {@link #commitDataWriter} of each task. An element @@ -290,10 +332,12 @@ public static native long createStreamWriter( // --------------------------------------------------------------------------------------------- /** - * Creates the writer for one task. See {@link DataSourceWriter#createWriter}. + * Creates the writer for one task. * * @param writerState the state returned by {@link #serializeWriter}, or an empty array if there * is none + * @param partitionId the partition that the task writes + * @param taskId the ID of the task attempt * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write * @return the data writer handle */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala new file mode 100644 index 0000000000000..edd1b464121b8 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala @@ -0,0 +1,128 @@ +/* + * 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.execution.datasources.v2.columnar + +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} +import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** + * A data source that exchanges its data as columnar batches. [[ColumnarTable]] adapts it to Data + * Source V2. It is implemented by the native data sources, see + * `org.apache.spark.sql.execution.datasources.v2.ffi.NativeDataSource`. + * + * A new instance is created on the driver for each schema inference, scan and write. + */ +trait ColumnarDataSource { + + /** Returns the schema of the data, when the user does not specify it. */ + def schema(): StructType + + /** Returns a reader for a batch scan of the given schema. */ + def reader(schema: StructType): ColumnarReader + + /** Returns a reader for a micro-batch streaming scan of the given schema. */ + def streamReader(schema: StructType): ColumnarStreamReader + + /** + * Returns a writer for a batch write. + * + * @param overwrite whether to replace the existing data (save mode "overwrite") instead of + * appending to it + */ + def writer(schema: StructType, overwrite: Boolean): ColumnarWriter + + /** + * Returns a writer for a streaming write. + * + * @param overwrite whether to replace the existing data in each micro-batch (output mode + * "complete") instead of appending to it + */ + def streamWriter(schema: StructType, overwrite: Boolean): ColumnarStreamWriter +} + +/** + * Reads the data of a [[ColumnarDataSource]] for a batch scan. Spark pushes operations down by + * calling, each at most once and in this order, `pushPredicates`, `pushLimit` and `pruneColumns`, + * and plans the scan with `partitions`, on the driver. Then it serializes the reader to the + * executors, where `read` is called for each partition. + */ +trait ColumnarReader extends Serializable { + + /** Returns the predicates that Spark still has to evaluate after reading. */ + def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = predicates + + /** Returns whether the reader returns at most `limit` rows. Spark applies it again anyway. */ + def pushLimit(limit: Int): Boolean = false + + /** + * Returns whether `read` returns exactly the columns of `requiredSchema`, the top-level columns + * that the query needs, instead of all the columns of the read schema. + */ + def pruneColumns(requiredSchema: StructType): Boolean = false + + def partitions(): Array[InputPartition] + + /** Reads a partition on an executor. A batch only needs to be valid until the next one. */ + def read(partition: InputPartition): PartitionReader[ColumnarBatch] +} + +/** + * Reads the data of a [[ColumnarDataSource]] for a micro-batch streaming scan, with offsets in + * JSON. Spark plans the partitions of each micro-batch on the driver, then serializes the reader + * to the executors, where `read` is called for each partition. + */ +trait ColumnarStreamReader extends Serializable { + def initialOffset(): String + + def latestOffset(): String + + /** Plans the partitions that read the data after `start` up to and including `end`. */ + def partitions(start: String, end: String): Array[InputPartition] + + def read(partition: InputPartition): PartitionReader[ColumnarBatch] + + /** Tells the data source that Spark has processed all the data up to and including `end`. */ + def commit(end: String): Unit = {} + + def stop(): Unit = {} +} + +/** + * Writes data to a [[ColumnarDataSource]] for a batch write. Spark serializes the writer to the + * executors, where `createWriter` is called for each task, and then commits or aborts the write + * on the driver. The batches passed to the data writers are backed by Arrow, and are only valid + * during the call. + */ +trait ColumnarWriter extends Serializable { + def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] + + def commit(messages: Array[WriterCommitMessage]): Unit + + def abort(messages: Array[WriterCommitMessage]): Unit +} + +/** Writes data to a [[ColumnarDataSource]] for a streaming write, one micro-batch at a time. */ +trait ColumnarStreamWriter extends Serializable { + def createWriter(partitionId: Int, taskId: Long, epochId: Long): DataWriter[ColumnarBatch] + + def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit + + def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala index 5ecd1c8fc806b..6b78219147647 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala @@ -17,7 +17,6 @@ package org.apache.spark.sql.execution.datasources.v2.columnar import java.util -import java.util.function.{Function => JFunction} import org.apache.spark.SparkException import org.apache.spark.sql.catalyst.InternalRow @@ -27,20 +26,19 @@ import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.read._ import org.apache.spark.sql.connector.read.streaming.{MicroBatchStream, Offset} import org.apache.spark.sql.connector.write.{LogicalWriteInfo, WriteBuilder} -import org.apache.spark.sql.datasource.{DataSource, DataSourceReader, DataSourceStreamReader} import org.apache.spark.sql.internal.connector.SupportsMetadata import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.sql.vectorized.ColumnarBatch /** - * The Data Source V2 [[Table]] of a data source implemented with the columnar data source API. - * It creates a new [[DataSource]] for each scan and each write, with their options. + * The Data Source V2 [[Table]] of a [[ColumnarDataSource]]. It creates a new data source for each + * scan and each write, with their options. */ class ColumnarTable( shortName: String, tableSchema: StructType, - createDataSource: JFunction[CaseInsensitiveStringMap, DataSource]) + createDataSource: CaseInsensitiveStringMap => ColumnarDataSource) extends Table with SupportsRead with SupportsWrite { override def name(): String = shortName @@ -51,15 +49,15 @@ class ColumnarTable( util.EnumSet.of(BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE) override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new ColumnarScanBuilder(tableSchema, () => createDataSource.apply(options)) + new ColumnarScanBuilder(tableSchema, () => createDataSource(options)) } override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = { - new ColumnarWriteBuilder(shortName, info, () => createDataSource.apply(info.options())) + new ColumnarWriteBuilder(shortName, info, () => createDataSource(info.options())) } } -class ColumnarScanBuilder(schema: StructType, createDataSource: () => DataSource) +class ColumnarScanBuilder(schema: StructType, createDataSource: () => ColumnarDataSource) extends ScanBuilder with SupportsPushDownV2Filters with SupportsPushDownLimit @@ -69,13 +67,13 @@ class ColumnarScanBuilder(schema: StructType, createDataSource: () => DataSource // Created on first use, because a streaming scan neither pushes anything down nor reads in // batch, and its data source may not support batch scans at all. - private var reader: DataSourceReader = _ + private var reader: ColumnarReader = _ private var pushed: Array[Predicate] = Array.empty private var pushedLimit: Option[Int] = None private var readSchema: StructType = schema - private def batchReader: DataSourceReader = { + private def batchReader: ColumnarReader = { if (reader == null) { reader = dataSource.reader(schema) } @@ -119,8 +117,8 @@ class ColumnarScanBuilder(schema: StructType, createDataSource: () => DataSource class ColumnarScan( fullSchema: StructType, prunedSchema: StructType, - batchReader: () => DataSourceReader, - dataSource: () => DataSource, + batchReader: () => ColumnarReader, + dataSource: () => ColumnarDataSource, pushedPredicates: Array[Predicate], pushedLimit: Option[Int]) extends Scan with SupportsMetadata { @@ -150,7 +148,7 @@ class ColumnarScan( } } -class ColumnarBatchScan(reader: DataSourceReader) extends Batch { +class ColumnarBatchScan(reader: ColumnarReader) extends Batch { private lazy val partitions: Array[InputPartition] = reader.partitions() override def planInputPartitions(): Array[InputPartition] = partitions @@ -164,7 +162,7 @@ class ColumnarBatchScan(reader: DataSourceReader) extends Batch { case class ColumnarStreamOffset(json: String) extends Offset -class ColumnarMicroBatchStream(reader: DataSourceStreamReader) extends MicroBatchStream { +class ColumnarMicroBatchStream(reader: ColumnarStreamReader) extends MicroBatchStream { override def initialOffset(): Offset = ColumnarStreamOffset(reader.initialOffset()) @@ -186,7 +184,7 @@ class ColumnarMicroBatchStream(reader: DataSourceStreamReader) extends MicroBatc } /** - * Columnar data sources only produce columnar batches, so Spark never asks for rows: the scan + * A [[ColumnarDataSource]] only produces columnar batches, so Spark never asks for rows: the scan * reports [[Scan.ColumnarSupportMode.SUPPORTED]]. */ abstract class ColumnarOnlyPartitionReaderFactory extends PartitionReaderFactory { @@ -197,14 +195,14 @@ abstract class ColumnarOnlyPartitionReaderFactory extends PartitionReaderFactory override def supportColumnarReads(partition: InputPartition): Boolean = true } -class ColumnarPartitionReaderFactory(reader: DataSourceReader) +class ColumnarPartitionReaderFactory(reader: ColumnarReader) extends ColumnarOnlyPartitionReaderFactory { override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = { reader.read(partition) } } -class ColumnarStreamPartitionReaderFactory(reader: DataSourceStreamReader) +class ColumnarStreamPartitionReaderFactory(reader: ColumnarStreamReader) extends ColumnarOnlyPartitionReaderFactory { override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = { reader.read(partition) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala index 5a7386d6e9145..af5e530dfda29 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala @@ -24,7 +24,6 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.connector.metric.CustomTaskMetric import org.apache.spark.sql.connector.write._ import org.apache.spark.sql.connector.write.streaming.{StreamingDataWriterFactory, StreamingWrite} -import org.apache.spark.sql.datasource.{DataSource, DataSourceStreamWriter, DataSourceWriter} import org.apache.spark.sql.execution.arrow.ArrowWriter import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType @@ -35,7 +34,7 @@ import org.apache.spark.util.Utils class ColumnarWriteBuilder( shortName: String, info: LogicalWriteInfo, - createDataSource: () => DataSource) + createDataSource: () => ColumnarDataSource) extends WriteBuilder with SupportsTruncate { private var overwrite = false @@ -52,7 +51,7 @@ class ColumnarWrite( shortName: String, info: LogicalWriteInfo, overwrite: Boolean, - createDataSource: () => DataSource) + createDataSource: () => ColumnarDataSource) extends Write { override def description(): String = shortName @@ -67,7 +66,7 @@ class ColumnarWrite( } } -class ColumnarBatchWrite(writer: DataSourceWriter, schema: StructType) extends BatchWrite { +class ColumnarBatchWrite(writer: ColumnarWriter, schema: StructType) extends BatchWrite { override def createBatchWriterFactory(info: PhysicalWriteInfo): DataWriterFactory = { new ColumnarDataWriterFactory(writer, ArrowBatchingOptions(schema)) @@ -78,7 +77,7 @@ class ColumnarBatchWrite(writer: DataSourceWriter, schema: StructType) extends B override def abort(messages: Array[WriterCommitMessage]): Unit = writer.abort(messages) } -class ColumnarStreamingWrite(writer: DataSourceStreamWriter, schema: StructType) +class ColumnarStreamingWrite(writer: ColumnarStreamWriter, schema: StructType) extends StreamingWrite { override def createStreamingWriterFactory( @@ -95,7 +94,7 @@ class ColumnarStreamingWrite(writer: DataSourceStreamWriter, schema: StructType) } } -class ColumnarDataWriterFactory(writer: DataSourceWriter, options: ArrowBatchingOptions) +class ColumnarDataWriterFactory(writer: ColumnarWriter, options: ArrowBatchingOptions) extends DataWriterFactory { override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = { new ArrowBatchingDataWriter(writer.createWriter(partitionId, taskId), options) @@ -103,7 +102,7 @@ class ColumnarDataWriterFactory(writer: DataSourceWriter, options: ArrowBatching } class ColumnarStreamingDataWriterFactory( - writer: DataSourceStreamWriter, + writer: ColumnarStreamWriter, options: ArrowBatchingOptions) extends StreamingDataWriterFactory { override def createWriter( @@ -131,7 +130,7 @@ object ArrowBatchingOptions { /** * Converts the rows of a write task to Arrow-backed [[ColumnarBatch]]es for the [[DataWriter]] - * of a columnar data source. A new Arrow vector schema root is created for each batch, so a + * of a [[ColumnarDataSource]]. A new Arrow vector schema root is created for each batch, so a * native writer that took ownership of the data of a batch can keep it after the call. */ class ArrowBatchingDataWriter(delegate: DataWriter[ColumnarBatch], options: ArrowBatchingOptions) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index dce0d4b16f01a..d52194b5d70b1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -32,8 +32,8 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} -import org.apache.spark.sql.datasource.{DataSource, DataSourceReader, DataSourceStreamReader, DataSourceStreamWriter, DataSourceWriter} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarDataSource, ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataType, StructType} import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} @@ -41,7 +41,7 @@ import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, Column import org.apache.spark.util.Utils /** - * A [[DataSource]] implemented by a native library, through + * A [[ColumnarDataSource]] implemented by a native library, through * `org.apache.spark.sql.datasource.NativeBridge`. * * @param name the name of the data source, as listed in the manifest of the package @@ -50,7 +50,7 @@ class NativeDataSource( pkg: NativeDataSourcePackage, name: String, options: CaseInsensitiveStringMap) - extends DataSource { + extends ColumnarDataSource { import NativeCalls._ override def schema(): StructType = { @@ -62,7 +62,7 @@ class NativeDataSource( } } - override def reader(schema: StructType): DataSourceReader = { + override def reader(schema: StructType): ColumnarReader = { val library = NativeLibraries.get(pkg) val handle = create(library, "plan a scan of", "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { dataSource => @@ -71,7 +71,7 @@ class NativeDataSource( new NativeDataSourceReader(pkg, name, schema, new NativeHandle(handle, library.closeReader)) } - override def streamReader(schema: StructType): DataSourceStreamReader = { + override def streamReader(schema: StructType): ColumnarStreamReader = { val library = NativeLibraries.get(pkg) val handle = create(library, "plan a scan of", "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { dataSource => @@ -81,7 +81,7 @@ class NativeDataSource( pkg, name, schema, new NativeHandle(handle, library.closeStreamReader)) } - override def writer(schema: StructType, overwrite: Boolean): DataSourceWriter = { + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { val library = NativeLibraries.get(pkg) val handle = new NativeHandle( create(library, "plan a write to", "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource => @@ -93,7 +93,7 @@ class NativeDataSource( new NativeDataSourceWriter(pkg, name, handle, serializeWriter(library, handle)) } - override def streamWriter(schema: StructType, overwrite: Boolean): DataSourceStreamWriter = { + override def streamWriter(schema: StructType, overwrite: Boolean): ColumnarStreamWriter = { val library = NativeLibraries.get(pkg) val handle = new NativeHandle( create(library, "plan a write to", "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { @@ -145,7 +145,7 @@ class NativeDataSourceReader( name: String, schema: StructType, @transient private val handle: NativeHandle) - extends DataSourceReader { + extends ColumnarReader { import NativeCalls._ // Set on the driver when the partitions are planned, and serialized to the executors. @@ -218,7 +218,7 @@ class NativeDataSourceStreamReader( name: String, schema: StructType, @transient private val handle: NativeHandle) - extends DataSourceStreamReader { + extends ColumnarStreamReader { import NativeCalls._ // Set on the driver whenever the partitions of a micro-batch are planned, and serialized to @@ -262,7 +262,7 @@ class NativeDataSourceWriter( name: String, @transient private val handle: NativeHandle, writerState: Array[Byte]) - extends DataSourceWriter { + extends ColumnarWriter { import NativeCalls._ @transient private lazy val library = NativeLibraries.get(pkg) @@ -293,7 +293,7 @@ class NativeDataSourceStreamWriter( name: String, @transient private val handle: NativeHandle, writerState: Array[Byte]) - extends DataSourceStreamWriter { + extends ColumnarStreamWriter { import NativeCalls._ @transient private lazy val library = NativeLibraries.get(pkg) diff --git a/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister b/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister index 7c0052d49ec15..8fcf427ec8a75 100644 --- a/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister +++ b/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister @@ -30,4 +30,4 @@ org.apache.spark.sql.streaming.sources.FakeWriteOnly org.apache.spark.sql.streaming.sources.FakeNoWrite org.apache.spark.sql.streaming.sources.FakeWriteSupportingExternalMetadata org.apache.spark.sql.streaming.sources.FakeWriteSupportProviderV1Fallback -org.apache.spark.sql.datasource.TestRangeProvider +org.apache.spark.sql.execution.datasources.v2.columnar.TestRangeProvider diff --git a/sql/core/src/test/scala/org/apache/spark/sql/datasource/ColumnarDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala similarity index 77% rename from sql/core/src/test/scala/org/apache/spark/sql/datasource/ColumnarDataSourceSuite.scala rename to sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala index a674188f63c18..0211f81acda61 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/datasource/ColumnarDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala @@ -14,30 +14,35 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.spark.sql.datasource +package org.apache.spark.sql.execution.datasources.v2.columnar import java.nio.charset.StandardCharsets.UTF_8 +import java.util import java.util.concurrent.ConcurrentHashMap import scala.collection.mutable.ArrayBuffer -import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.SparkException import org.apache.spark.sql.{QueryTest, Row} -import org.apache.spark.sql.connector.expressions.{Literal, NamedReference} +import org.apache.spark.sql.connector.catalog.{Table, TableProvider} +import org.apache.spark.sql.connector.expressions.{Literal, NamedReference, Transform} import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation -import org.apache.spark.sql.execution.datasources.v2.columnar.ColumnarScan import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.streaming.StreamingQuery +import org.apache.spark.sql.sources.DataSourceRegister import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} import org.apache.spark.util.Utils +/** + * Tests the Data Source V2 adapter of [[ColumnarDataSource]] with a data source implemented on the + * JVM. `NativeDataSourceSuite` tests the native data sources, which implement it. + */ class ColumnarDataSourceSuite extends QueryTest with SharedSparkSession { import testImplicits._ @@ -49,12 +54,6 @@ class ColumnarDataSourceSuite extends QueryTest with SharedSparkSession { }.get } - private def findError(e: Throwable, condition: String): SparkThrowable = { - Iterator.iterate(e)(_.getCause).takeWhile(_ != null) - .collectFirst { case t: SparkThrowable if t.getCondition == condition => t } - .getOrElse(throw e) - } - test("batch read") { val df = spark.read.format(format).option("end", "10").load() assert(df.schema == StructType.fromDDL("id BIGINT, name STRING")) @@ -127,76 +126,43 @@ class ColumnarDataSourceSuite extends QueryTest with SharedSparkSession { assert(TestColumnarSink.rows(table).sorted == (0L until 5L).map(id => (id, s"name$id"))) assert(TestColumnarSink.epochs.get(table) == Seq(0L)) } - - test("operations that the data source does not implement") { - val provider = classOf[TestEmptyProvider].getName - val description = classOf[TestEmptyDataSource].getName - checkError( - exception = intercept[SparkUnsupportedOperationException] { - spark.read.format(provider).load() - }, - condition = "UNABLE_TO_INFER_SCHEMA", - parameters = Map("format" -> description)) - checkError( - exception = intercept[SparkUnsupportedOperationException] { - spark.read.schema("id BIGINT").format(provider).load().collect() - }, - condition = "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", - parameters = Map("description" -> description)) - checkError( - exception = intercept[SparkUnsupportedOperationException] { - spark.range(1).write.format(provider).mode("append").save() - }, - condition = "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", - parameters = Map("description" -> description)) - - def checkStreaming(condition: String)(start: => StreamingQuery): Unit = { - val e = intercept[Exception] { - val query = start - try query.processAllAvailable() finally query.stop() - } - checkError( - exception = findError(e, condition), - condition = condition, - parameters = Map("description" -> description)) - } - checkStreaming("DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { - spark.readStream.schema("id BIGINT").format(provider).load() - .writeStream.format("noop").start() - } - checkStreaming("DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { - spark.readStream.format("rate").load().writeStream.format(provider) - .option("checkpointLocation", Utils.createTempDir().getPath).start() - } - } } /** Reads the ids [0, end) and their names, and writes to [[TestColumnarSink]]. */ -class TestRangeProvider extends DataSourceProvider { +class TestRangeProvider extends TableProvider with DataSourceRegister { override def shortName(): String = "test_columnar_range" - override def createDataSource(options: CaseInsensitiveStringMap): DataSource = { - new TestRangeDataSource(options) + override def inferSchema(options: CaseInsensitiveStringMap): StructType = { + new TestRangeDataSource(options).schema() } + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: util.Map[String, String]): Table = { + new ColumnarTable(shortName(), schema, new TestRangeDataSource(_)) + } + + override def supportsExternalMetadata(): Boolean = true } -class TestRangeDataSource(options: CaseInsensitiveStringMap) extends DataSource { +class TestRangeDataSource(options: CaseInsensitiveStringMap) extends ColumnarDataSource { private val end = options.getLong("end", 10) private val table = options.getOrDefault("table", "default") override def schema(): StructType = StructType.fromDDL("id BIGINT, name STRING") - override def reader(schema: StructType): DataSourceReader = new TestRangeReader(end, schema) + override def reader(schema: StructType): ColumnarReader = new TestRangeReader(end, schema) - override def streamReader(schema: StructType): DataSourceStreamReader = { + override def streamReader(schema: StructType): ColumnarStreamReader = { new TestRangeStreamReader(end, schema) } - override def writer(schema: StructType, overwrite: Boolean): DataSourceWriter = { + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { new TestSinkWriter(table, overwrite, options.getBoolean("fail", false)) } - override def streamWriter(schema: StructType, overwrite: Boolean): DataSourceStreamWriter = { + override def streamWriter(schema: StructType, overwrite: Boolean): ColumnarStreamWriter = { new TestSinkStreamWriter(table) } } @@ -204,7 +170,7 @@ class TestRangeDataSource(options: CaseInsensitiveStringMap) extends DataSource case class TestRangePartition(lo: Long, hi: Long) extends InputPartition /** Accepts the predicates `id > `. */ -class TestRangeReader(end: Long, schema: StructType) extends DataSourceReader { +class TestRangeReader(end: Long, schema: StructType) extends ColumnarReader { private var lo = 0L private var limit = Long.MaxValue private var readSchema = schema @@ -243,7 +209,7 @@ class TestRangeReader(end: Long, schema: StructType) extends DataSourceReader { } } -class TestRangeStreamReader(end: Long, schema: StructType) extends DataSourceStreamReader { +class TestRangeStreamReader(end: Long, schema: StructType) extends ColumnarStreamReader { private def offset(json: String): Long = json.stripPrefix("{\"offset\":").stripSuffix("}").toLong override def initialOffset(): String = "{\"offset\":0}" @@ -330,7 +296,7 @@ class TestDataWriter(fail: Boolean) extends DataWriter[ColumnarBatch] { override def close(): Unit = {} } -class TestSinkWriter(table: String, overwrite: Boolean, fail: Boolean) extends DataSourceWriter { +class TestSinkWriter(table: String, overwrite: Boolean, fail: Boolean) extends ColumnarWriter { override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { new TestDataWriter(fail) } @@ -344,7 +310,7 @@ class TestSinkWriter(table: String, overwrite: Boolean, fail: Boolean) extends D } } -class TestSinkStreamWriter(table: String) extends DataSourceStreamWriter { +class TestSinkStreamWriter(table: String) extends ColumnarStreamWriter { override def createWriter( partitionId: Int, taskId: Long, @@ -359,14 +325,3 @@ class TestSinkStreamWriter(table: String) extends DataSourceStreamWriter { TestColumnarSink.aborted.add(table) } } - -/** A data source that implements nothing. */ -class TestEmptyProvider extends DataSourceProvider { - override def shortName(): String = "test_columnar_empty" - - override def createDataSource(options: CaseInsensitiveStringMap): DataSource = { - new TestEmptyDataSource - } -} - -class TestEmptyDataSource extends DataSource From 1fc3fc748b19ec93e67d19585c9933577e1326c9 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 9 Oct 2026 10:41:55 +0900 Subject: [PATCH 03/27] Pass the options of catalog tables to native data sources A table created with CREATE TABLE ... USING ... OPTIONS (...) passes its options to TableProvider.getTable, while its scans and writes only get the options of their query. The columnar table now merges the options of the table under the options of each scan and write, so that SELECT and INSERT INTO on such a table use them. Co-authored-by: Isaac --- .../v2/columnar/ColumnarTable.scala | 22 ++++++++++++++++--- .../v2/ffi/NativeDataSourceV2.scala | 3 ++- .../test_native_datasource.cc | 9 +++++++- .../v2/columnar/ColumnarDataSourceSuite.scala | 16 +++++++++++++- .../v2/ffi/NativeDataSourceSuite.scala | 22 +++++++++++++++++++ 5 files changed, 66 insertions(+), 6 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala index 6b78219147647..a59ae7d8e261c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala @@ -17,6 +17,7 @@ package org.apache.spark.sql.execution.datasources.v2.columnar import java.util +import java.util.Locale import org.apache.spark.SparkException import org.apache.spark.sql.catalyst.InternalRow @@ -33,11 +34,17 @@ import org.apache.spark.sql.vectorized.ColumnarBatch /** * The Data Source V2 [[Table]] of a [[ColumnarDataSource]]. It creates a new data source for each - * scan and each write, with their options. + * scan and each write, with their options on top of the options of the table. + * + * @param tableOptions the options of the table, given to `TableProvider.getTable`. They are the + * options of `CREATE TABLE ... USING ... OPTIONS (...)` for a table in a + * catalog, while the scans and the writes of the table only get the options + * of their query. */ class ColumnarTable( shortName: String, tableSchema: StructType, + tableOptions: util.Map[String, String], createDataSource: CaseInsensitiveStringMap => ColumnarDataSource) extends Table with SupportsRead with SupportsWrite { @@ -49,11 +56,20 @@ class ColumnarTable( util.EnumSet.of(BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE) override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new ColumnarScanBuilder(tableSchema, () => createDataSource(options)) + new ColumnarScanBuilder(tableSchema, () => createDataSource(withTableOptions(options))) } override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = { - new ColumnarWriteBuilder(shortName, info, () => createDataSource(info.options())) + new ColumnarWriteBuilder( + shortName, info, () => createDataSource(withTableOptions(info.options()))) + } + + private def withTableOptions(options: CaseInsensitiveStringMap): CaseInsensitiveStringMap = { + val merged = new util.HashMap[String, String]() + tableOptions.forEach((key, value) => merged.put(key.toLowerCase(Locale.ROOT), value)) + // The keys of a case-insensitive map are in lower case. + options.forEach((key, value) => merged.put(key, value)) + new CaseInsensitiveStringMap(merged) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala index d518af1062e26..e02a3e5fc5407 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala @@ -59,7 +59,8 @@ class NativeDataSourceV2 extends NamedTableProvider { partitioning: Array[Transform], properties: util.Map[String, String]): Table = { val (name, pkg) = resolved - new ColumnarTable(shortName, schema, options => new NativeDataSource(pkg, name, options)) + new ColumnarTable( + shortName, schema, properties, options => new NativeDataSource(pkg, name, options)) } override def supportsExternalMetadata(): Boolean = true diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index 6f55494d383e5..eaa52ed534c8d 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -989,7 +989,14 @@ jlong createAnyWriter(JNIEnv* env, jlong handle, jlong schema_address, jboolean auto* schema = reinterpret_cast(schema_address); if (schema->n_children == 0) throw std::runtime_error("no columns to write"); auto* writer = new Writer(); - writer->state["path"] = get(source->options, "path", ""); + // The location of a table in a catalog is passed as a URI. + std::string path = get(source->options, "path", ""); + if (path.rfind("file://", 0) == 0) { + path = path.substr(7); + } else if (path.rfind("file:", 0) == 0) { + path = path.substr(5); + } + writer->state["path"] = path; writer->state["overwrite"] = overwrite ? "true" : "false"; writer->state["fail"] = get(source->options, "fail", ""); if (writer->state["path"].empty()) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala index 0211f81acda61..af8c09c22bc5f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala @@ -81,6 +81,20 @@ class ColumnarDataSourceSuite extends QueryTest with SharedSparkSession { assert(scan(limited).getMetaData()("PushedLimit") == "LIMIT 2") } + test("tables in a catalog") { + withTable("columnar_range") { + sql(s"CREATE TABLE columnar_range USING $format OPTIONS (end 3, table 'catalog_write')") + // The options of the table are used by the scans of the table, unless a scan sets them. + checkAnswer(spark.table("columnar_range"), (0L until 3L).map(id => Row(id, s"name$id"))) + checkAnswer( + spark.read.option("END", "2").table("columnar_range"), + (0L until 2L).map(id => Row(id, s"name$id"))) + + sql("INSERT INTO columnar_range VALUES (1, 'a')") + assert(TestColumnarSink.rows("catalog_write") == Seq((1L, "a"))) + } + } + test("batch write") { withSQLConf(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "3") { val table = "batch_write" @@ -140,7 +154,7 @@ class TestRangeProvider extends TableProvider with DataSourceRegister { schema: StructType, partitioning: Array[Transform], properties: util.Map[String, String]): Table = { - new ColumnarTable(shortName(), schema, new TestRangeDataSource(_)) + new ColumnarTable(shortName(), schema, properties, new TestRangeDataSource(_)) } override def supportsExternalMetadata(): Boolean = true diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 9edbf1d2357ac..b3af8de6fd3f4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -366,6 +366,28 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { }.toSet) } + nativeTest("tables in a catalog") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + try { + session.sql("CREATE TABLE native_range_table USING native_range OPTIONS (end 3)") + checkRows(session.table("native_range_table"), (0L until 3L).map(expectedRow)) + checkRows( + session.read.option("end", "2").table("native_range_table"), + (0L until 2L).map(expectedRow)) + + session.sql("CREATE TABLE native_sink_table (id BIGINT, name STRING) USING native_sink " + + s"OPTIONS (path '${dir.getPath}')") + session.sql("INSERT INTO native_sink_table VALUES (1, 'a'), (2, 'b')") + val written = dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala) + assert(written.toSet == Set("1,a", "2,b")) + } finally { + session.sql("DROP TABLE IF EXISTS native_range_table") + session.sql("DROP TABLE IF EXISTS native_sink_table") + } + } + nativeTest("abort a failed write") { val session = newSession(defaultPackage) val dir = Utils.createTempDir() From aae6a4f43eab825303af3d67b170682b99e6efe2 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 9 Oct 2026 10:49:44 +0900 Subject: [PATCH 04/27] Find the native data source packages installed in SPARK_HOME automatically Like the Python data sources installed in the Python path, the native data source packages in the native-datasources directory of SPARK_HOME are now found without any configuration. The packages of a session, added with spark.addArtifact or under spark.sql.dataSource.native.paths, take precedence over the installed ones, like the Python data sources registered at runtime take precedence over the installed ones. Co-authored-by: Isaac --- docs/sql-native-data-sources.md | 29 +++++++++----- .../apache/spark/sql/internal/SQLConf.scala | 5 ++- .../spark/sql/internal/StaticSQLConf.scala | 7 ++-- .../spark/sql/datasource/NativeBridge.java | 5 ++- .../v2/ffi/NativeDataSourceRegistry.scala | 40 +++++++++++++------ .../v2/ffi/NativeDataSourceSuite.scala | 34 +++++++++++++++- 6 files changed, 90 insertions(+), 30 deletions(-) diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 40ae89ca6baa2..e3bffd445712f 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -107,17 +107,26 @@ named `spark-native-datasource.json`, and the library built for one or more plat - `libraries` maps each platform, `-` with an `os` of `linux`, `osx` or `windows` and an `arch` of `x86_64` or `aarch64`, to the path of the library in the package. -Add a package to a session with `spark.addArtifact`, which works in Scala, Java and Python, and with -Spark Connect: +Spark finds packages in three places, and the executors get them automatically: -```python -spark.addArtifact("/path/to/rust_range.sparkpkg") -spark.read.format("rust_range").option("end", 100).load().show() -``` +- **Installed packages**: like the Python data sources installed in the Python path, the packages in + the `native-datasources` directory of `SPARK_HOME` are found automatically, without any + configuration. With PySpark installed by pip, `SPARK_HOME` is the directory of the `pyspark` + package. +- **Artifacts of a session**: `spark.addArtifact` adds a package to a session, in Scala, Java and + Python, and with Spark Connect: + + ```python + spark.addArtifact("/path/to/rust_range.sparkpkg") + spark.read.format("rust_range").option("end", 100).load().show() + ``` + +- **Configured paths**: the packages listed in `spark.sql.dataSource.native.paths`, or in the + directories it lists, for example in a container image. -Spark can also find packages under `spark.sql.dataSource.native.paths`, for example in a container -image. In both cases, the executors get the package automatically. Give packages distinct file -names, for example with their version: sessions that are not isolated share the files they add. +The packages of a session, added as artifacts or configured, take precedence over the installed +ones. Give packages distinct file names, for example with their version: sessions that are not +isolated share the files they add. A Java data source with the same name takes precedence over a native one, and so does a Python data source. @@ -472,7 +481,7 @@ statically, or install them on the nodes. | Property Name | Default | Meaning | |---------------|---------|---------| | `spark.sql.dataSource.native.enabled` | true | Whether Spark loads native data sources. It is a static configuration, so it can only be set when the Spark application starts. | -| `spark.sql.dataSource.native.paths` | (none) | Comma-separated list of native data source packages, and of directories that contain them, on the local file system of the driver. | +| `spark.sql.dataSource.native.paths` | (none) | Comma-separated list of native data source packages, and of directories that contain them, on the local file system of the driver. They take precedence over the packages installed in `SPARK_HOME`. | ## Security diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index b068dfce5a1d4..df5ed6e95c5c8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -7211,8 +7211,9 @@ object SQLConf { .doc("Comma-separated list of native data source packages (.sparkpkg files), and of " + "directories that contain them, on the local file system of the driver. Spark adds a " + "package to the artifacts of the session when one of its data sources is used, so that " + - "the executors get it too. Packages can also be added with spark.addArtifact. See " + - "spark.sql.dataSource.native.enabled.") + "the executors get it too. Packages can also be added with spark.addArtifact. These " + + "packages take precedence over the ones installed in the native-datasources directory " + + "of SPARK_HOME. See spark.sql.dataSource.native.enabled.") .version("4.4.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .stringConf diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala index 8b701d386ced5..87d4c89c9a0f9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala @@ -468,9 +468,10 @@ object StaticSQLConf { buildStaticConf("spark.sql.dataSource.native.enabled") .doc("When true, Spark finds and loads native data sources: data sources implemented in " + "native code such as Rust or C++, and distributed as native data source packages " + - "(.sparkpkg files). A package is found when it is added with spark.addArtifact, or when " + - "it is under spark.sql.dataSource.native.paths. Loading a native data source runs its " + - "native code in the driver and executor processes.") + "(.sparkpkg files). A package is found when it is installed in the native-datasources " + + "directory of SPARK_HOME, when it is added with spark.addArtifact, or when it is under " + + "spark.sql.dataSource.native.paths. Loading a native data source runs its native code " + + "in the driver and executor processes.") .version("4.4.0") .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) .booleanConf diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java index bee42c03847d4..d6e5dc27ef882 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java @@ -33,8 +33,9 @@ *

      * The library is distributed in a native data source package: a zip file with the extension * {@code .sparkpkg} that contains a manifest named {@code spark-native-datasource.json} and the - * library built for one or more platforms. Spark finds the packages added to a session with - * {@code spark.addArtifact}, and the ones under the paths in + * library built for one or more platforms. Spark finds the packages installed in the + * {@code native-datasources} directory of {@code SPARK_HOME} automatically, as well as the + * packages added to a session with {@code spark.addArtifact} and the ones under the paths in * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL * guide for details and examples. * diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala index 94d76493e630d..3b9ccca628775 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.datasources.v2.ffi import java.io.File import java.nio.file.Paths -import org.apache.spark.{SparkEnv, SparkFiles} +import org.apache.spark.{SparkContext, SparkEnv, SparkFiles} import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.PATH import org.apache.spark.sql.Artifact @@ -30,12 +30,16 @@ import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.util.Utils /** - * Finds the native data source packages of the active session on the driver: the packages added - * with `spark.addArtifact`, and the packages listed in, or in a directory listed in, - * `spark.sql.dataSource.native.paths`. + * Finds the native data source packages on the driver. Like the Python data sources installed in + * the Python path, the packages installed in `$SPARK_HOME/native-datasources` are found + * automatically. The active session can also add packages with `spark.addArtifact`, or with + * `spark.sql.dataSource.native.paths`, which take precedence over the installed ones. */ object NativeDataSourceRegistry extends Logging { + /** The directory of `SPARK_HOME` that contains the installed native data source packages. */ + val INSTALLED_PACKAGES_DIR = "native-datasources" + /** Returns whether a package provides the native data source with the given name. */ def exists(name: String, conf: SQLConf): Boolean = lookup(name, conf).isDefined @@ -48,7 +52,12 @@ object NativeDataSourceRegistry extends Logging { if (!conf.getConf(StaticSQLConf.NATIVE_DATA_SOURCE_ENABLED)) { return None } - val found = packages(conf) + lookup(name, configuredPackageFiles(conf) ++ artifactPackageFiles()) + .orElse(lookup(name, installedPackageFiles())) + } + + private def lookup(name: String, files: Seq[File]): Option[(String, NativeDataSourcePackage)] = { + val found = files.map(_.getCanonicalFile).distinct.map(NativeDataSourcePackage.read) .flatMap(pkg => pkg.manifest.dataSources.find(_.equalsIgnoreCase(name)).map(_ -> pkg)) // The same package can be found more than once, for example when it is in a configured // directory and was also added as an artifact. @@ -110,9 +119,12 @@ object NativeDataSourceRegistry extends Logging { } } - private def packages(conf: SQLConf): Seq[NativeDataSourcePackage] = { - (configuredPackageFiles(conf) ++ artifactPackageFiles()).map(_.getCanonicalFile).distinct - .map(NativeDataSourcePackage.read) + /** The packages installed in `$SPARK_HOME/native-datasources`, if the directory exists. */ + private[ffi] def installedPackageFiles(): Seq[File] = { + SparkContext.getActive.flatMap(_.getSparkHome()).toSeq + .map(home => new File(home, INSTALLED_PACKAGES_DIR)) + .filter(_.isDirectory) + .flatMap(packageFilesIn) } private def configuredPackageFiles(conf: SQLConf): Seq[File] = { @@ -124,10 +136,7 @@ object NativeDataSourceRegistry extends Logging { log"paths are supported. Add remote packages with spark.addArtifact instead.") Nil } else if (file.isDirectory) { - file.listFiles() - .filter(f => f.isFile && f.getName.endsWith(NativeDataSourcePackage.FILE_EXTENSION)) - .sortBy(_.getName) - .toSeq + packageFilesIn(file) } else if (file.isFile) { Seq(file) } else { @@ -137,6 +146,13 @@ object NativeDataSourceRegistry extends Logging { } } + private def packageFilesIn(dir: File): Seq[File] = { + dir.listFiles() + .filter(f => f.isFile && f.getName.endsWith(NativeDataSourcePackage.FILE_EXTENSION)) + .sortBy(_.getName) + .toSeq + } + private def artifactPackageFiles(): Seq[File] = { SparkSession.getActiveSession.toSeq.flatMap(_.artifactManager.getNativeDataSourcePackages._1) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index b3af8de6fd3f4..794017b227d59 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -27,7 +27,7 @@ import java.util.zip.{ZipEntry, ZipOutputStream} import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ -import org.apache.spark.{SparkRuntimeException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.{SparkConf, SparkRuntimeException, SparkThrowable, SparkUnsupportedOperationException} import org.apache.spark.sql.{AnalysisException, QueryTest, Row} import org.apache.spark.sql.classic.{DataFrame, SparkSession} import org.apache.spark.sql.connector.expressions.{Expression, FieldReference, GeneralScalarExpression, LiteralValue} @@ -50,6 +50,13 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { private lazy val buildDir = Utils.createTempDir(namePrefix = "native-datasource-test") + // The SPARK_HOME of the suite, where packages are installed in INSTALLED_PACKAGES_DIR. + private lazy val installedSparkHome = Utils.createTempDir(namePrefix = "spark-home") + + override protected def sparkConf: SparkConf = { + super.sparkConf.set("spark.home", installedSparkHome.getPath) + } + // The executors store the files of the sessions that are not isolated in the same directory, // and jobs outside of SQL executions get the files of all the sessions. Each package has its // own name, and the sessions are kept until the end of the suite, so that none of their @@ -95,6 +102,10 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { compile("test_native_datasource_b", "-DNAME_PREFIX=\"b_\"", "-DID_SIGN=-1") private lazy val abiV2Library = compile("test_native_datasource_v2", "-DABI_VERSION=2", "-DNAME_PREFIX=\"v2_\"") + private lazy val installedLibrary = + compile("test_native_datasource_installed", "-DNAME_PREFIX=\"installed_\"") + private lazy val installedNegatingLibrary = compile( + "test_native_datasource_installed_b", "-DNAME_PREFIX=\"installed_\"", "-DID_SIGN=-1") private def dataSources(prefix: String): Seq[String] = Seq("native_range", "native_sink", "native_counter").map(prefix + _) @@ -537,6 +548,27 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(distributed.artifactUUID == otherArtifactUUID) } + nativeTest("find installed packages automatically") { + val installDir = + new File(installedSparkHome, NativeDataSourceRegistry.INSTALLED_PACKAGES_DIR) + val installed = createPackage(installedLibrary, dataSources("installed_"), dir = installDir) + try { + assert(NativeDataSourceRegistry.installedPackageFiles() == Seq(installed)) + // The session neither adds the package nor configures a path. + checkRows( + newSession().read.format("installed_native_range").option("end", "3").load().select("id"), + Seq(Row(0L), Row(1L), Row(2L))) + + // A package of the session takes precedence over an installed one. + val session = newSession(createPackage(installedNegatingLibrary, dataSources("installed_"))) + checkRows( + session.read.format("installed_native_range").option("end", "3").load().select("id"), + Seq(Row(0L), Row(-1L), Row(-2L))) + } finally { + installed.delete() + } + } + nativeTest("find packages in directories") { val dir = Utils.createTempDir() createPackage(library, dataSources(""), dir = dir, fileName = "a.sparkpkg") From a5d6708b75c84011ff546e7089eef2f93e507eea Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 9 Oct 2026 12:41:28 +0900 Subject: [PATCH 05/27] Fix the Connect client compatibility check and the documentation build - Exclude the `org.apache.spark.sql.datasource` package, which has no counterpart in the Spark Connect client, from the client compatibility check, like the other packages specific to classic Spark. - Escape the Rust and C++ examples of the native data sources page from Liquid, which failed to parse `{{0, end / 2}, ...}` in the C++ example. Co-authored-by: Isaac --- docs/sql-native-data-sources.md | 4 ++++ .../connect/client/CheckConnectJvmClientCompatibility.scala | 1 + 2 files changed, 5 insertions(+) diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index e3bffd445712f..32fc23d4df2bc 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -155,6 +155,7 @@ jni = "0.21" `src/lib.rs`: +{% raw %} ```rust //! A native data source for Apache Spark, written in Rust. It reads the numbers [0, end) as a //! column `id`, in two partitions. The option `end` defaults to 10. @@ -318,6 +319,7 @@ pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_read( }) } ``` +{% endraw %} Build the library, and package it: @@ -340,6 +342,7 @@ EOF The same data source in C++, with [Arrow C++](https://arrow.apache.org/docs/cpp/), which exports schemas and streams with `arrow::ExportSchema` and `arrow::ExportRecordBatchReader`: +{% raw %} ```cpp // A native data source for Apache Spark, written in C++ with Arrow C++. It reads the numbers // [0, end) as a column `id`, in two partitions. The option `end` defaults to 10. @@ -463,6 +466,7 @@ SPARK_JNI(void, read)(JNIEnv* env, jclass, jbyteArray, jbyteArray partition, }); } ``` +{% endraw %} Build the library, here on macOS with Arrow C++ 24, which requires C++20, and package it as above: diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala index 2ce4cb0166a19..4081e73ccf757 100644 --- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala +++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala @@ -162,6 +162,7 @@ object CheckConnectJvmClientCompatibility { ProblemFilters.exclude[Problem]("org.apache.spark.sql.columnar.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.connector.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.classic.*"), + ProblemFilters.exclude[Problem]("org.apache.spark.sql.datasource.*"), // Native data sources ProblemFilters.exclude[Problem]("org.apache.spark.sql.execution.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.internal.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.jdbc.*"), From 4a2026e6edc53cb63b8167cbdce4d7d5399d835d Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 9 Oct 2026 13:34:45 +0900 Subject: [PATCH 06/27] Find installed native data source libraries automatically Instead of the native-datasources directory of SPARK_HOME, find the library of a native data source where native libraries are usually installed, by its file name: the library spark_datasource_, such as libspark_datasource_.so on Linux, implements the data source . Spark looks for it in the native library path: the directories of java.library.path, which include LD_LIBRARY_PATH on Linux, and then the lib directory of each installation prefix in PATH, such as /usr/local/lib for /usr/local/bin. The executors load the library from the same path as the driver, or else find it in their own native library path. The packages of a session still take precedence. Also: - Keep the error of a library that is loaded but rejected, because it cannot be loaded again in another class loader: using it again failed with "already loaded in another classloader". - Add NativeDataSourceConnectSuite, which tests native data sources with a Spark Connect client and server, and move the helpers that build the test library to NativeDataSourceTestUtils. Co-authored-by: Isaac --- .../resources/error/error-conditions.json | 24 +++ docs/sql-native-data-sources.md | 114 ++++++++--- .../sql/errors/QueryExecutionErrors.scala | 18 ++ .../apache/spark/sql/internal/SQLConf.scala | 4 +- .../spark/sql/internal/StaticSQLConf.scala | 9 +- .../NativeDataSourceConnectSuite.scala | 90 ++++++++ .../spark/sql/datasource/NativeBridge.java | 14 +- .../datasources/v2/ffi/NativeDataSource.scala | 57 +++--- .../v2/ffi/NativeDataSourcePackage.scala | 123 ++++++++--- .../v2/ffi/NativeDataSourceRegistry.scala | 103 +++++++--- .../v2/ffi/NativeDataSourceV2.scala | 18 +- .../v2/ffi/NativeDataSourceSuite.scala | 192 ++++++++---------- .../v2/ffi/NativeDataSourceTestUtils.scala | 148 ++++++++++++++ 13 files changed, 673 insertions(+), 241 deletions(-) create mode 100644 sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 279a426d02b9f..a20ee51b340d6 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -4792,6 +4792,24 @@ ], "sqlState" : "42000" }, + "INVALID_NATIVE_DATA_SOURCE_LIBRARY" : { + "message" : [ + "Invalid native data source library :" + ], + "subClass" : { + "CANNOT_LOAD" : { + "message" : [ + "" + ] + }, + "UNSUPPORTED_ABI_VERSION" : { + "message" : [ + "It requires version of the native data source interface, but Spark supports version ." + ] + } + }, + "sqlState" : "22023" + }, "INVALID_NATIVE_DATA_SOURCE_PACKAGE" : { "message" : [ "Invalid native data source package :" @@ -6296,6 +6314,12 @@ ], "sqlState" : "38000" }, + "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND" : { + "message" : [ + "Cannot find the native data source library , which is installed at on the driver. Install the same library on every node, at the same path or in the native library path." + ], + "sqlState" : "42K02" + }, "NATIVE_DATA_SOURCE_PACKAGE_CONFLICT" : { "message" : [ "More than one native data source package provides the data source : . Remove all but one of them." diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 32fc23d4df2bc..dc79c7cc8a61a 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -26,10 +26,10 @@ license: | A native data source is a data source implemented in native code, such as Rust or C++, without any JVM code. It is a shared library that implements the -[`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) binary interface, and -it is distributed as a native data source package: a single file that Spark loads when a query -uses the data source. Spark adapts it to [Data Source V2](sql-v2-data-sources.html), so it -supports: +[`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) binary interface. +Spark finds the library automatically when it is installed like other native libraries, or when a +session adds it in a package, and loads it when a query uses the data source. Spark adapts it to +[Data Source V2](sql-v2-data-sources.html), so it supports: - batch and micro-batch streaming reads, - batch and streaming writes, appending to or overwriting the existing data, @@ -84,19 +84,63 @@ The [Javadoc of `NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeB specifies each function, including the ownership of the Arrow structs and the JSON format of the pushed predicates. -## Packages +## Finding Native Data Sources -A native data source package is a zip file with the extension `.sparkpkg`. It contains a manifest -named `spark-native-datasource.json`, and the library built for one or more platforms: +When no Java or Python data source has the name that a query uses, Spark looks for a native data +source with that name: in the packages of the session, and then in the installed libraries. A Java +data source with the same name takes precedence over a native one, and so does a Python data +source. + +### Installed Libraries + +Like a Python data source installed in the Python path, a native data source library installed +like other native libraries is found automatically, without any configuration. The library is +named after the data source, in lower case, with the prefix `spark_datasource_`: the library of +`rust_range` is `libspark_datasource_rust_range.so` on Linux, +`libspark_datasource_rust_range.dylib` on macOS and `spark_datasource_rust_range.dll` on Windows. +Spark passes the name of the data source to `createDataSource`. A library that implements several +data sources is installed under each of their names, for example with symbolic links. + +Spark looks for the library in the native library path, which consists of, in order: + +1. the directories of `java.library.path`, where the JVM finds native libraries. They include + `LD_LIBRARY_PATH` on Linux, `DYLD_LIBRARY_PATH` on macOS and `PATH` on Windows, which also + contain the directories set with `spark.driver.extraLibraryPath` and + `spark.executor.extraLibraryPath`. +2. the `lib` directory of each installation prefix in `PATH`: `/lib` for each directory + `/bin` in `PATH`. + +So Spark finds the libraries where native libraries are usually installed: + +- `/usr/local/lib`, the default directory of `make install`, CMake (`cmake --install`) and + [cargo-c](https://github.com/lu-zero/cargo-c) (`cargo cinstall`), and of Homebrew on Intel macs; +- `/opt/homebrew/lib`, for Homebrew on Apple silicon; +- `$CONDA_PREFIX/lib` and `$VIRTUAL_ENV/lib`, in an activated conda environment or Python virtual + environment; +- `/usr/lib`, for the packages of a Linux distribution. + +Cargo does not install libraries: copy the library that `cargo build --release` builds into one of +these directories. A library in any other directory is found with `LD_LIBRARY_PATH`, or with +`spark.driver.extraLibraryPath` and `spark.executor.extraLibraryPath`. + +The executors load the library from the same path as the driver, or else find it in their own +native library path, so install it on every node, for example in the container image of the +cluster. + +### Packages + +A native data source package distributes a library with a session, for one or more platforms. It +is a zip file with the extension `.sparkpkg`, which contains a manifest named +`spark-native-datasource.json` and the libraries: ```json { "abiVersion": 1, "dataSources": ["rust_range"], "libraries": { - "linux-x86_64": "linux-x86_64/librust_range.so", - "linux-aarch64": "linux-aarch64/librust_range.so", - "osx-aarch64": "osx-aarch64/librust_range.dylib" + "linux-x86_64": "linux-x86_64/libspark_datasource_rust_range.so", + "linux-aarch64": "linux-aarch64/libspark_datasource_rust_range.so", + "osx-aarch64": "osx-aarch64/libspark_datasource_rust_range.dylib" } } ``` @@ -107,12 +151,8 @@ named `spark-native-datasource.json`, and the library built for one or more plat - `libraries` maps each platform, `-` with an `os` of `linux`, `osx` or `windows` and an `arch` of `x86_64` or `aarch64`, to the path of the library in the package. -Spark finds packages in three places, and the executors get them automatically: +A session finds the packages in two places, and the executors get them automatically: -- **Installed packages**: like the Python data sources installed in the Python path, the packages in - the `native-datasources` directory of `SPARK_HOME` are found automatically, without any - configuration. With PySpark installed by pip, `SPARK_HOME` is the directory of the `pyspark` - package. - **Artifacts of a session**: `spark.addArtifact` adds a package to a session, in Scala, Java and Python, and with Spark Connect: @@ -122,21 +162,19 @@ Spark finds packages in three places, and the executors get them automatically: ``` - **Configured paths**: the packages listed in `spark.sql.dataSource.native.paths`, or in the - directories it lists, for example in a container image. - -The packages of a session, added as artifacts or configured, take precedence over the installed -ones. Give packages distinct file names, for example with their version: sessions that are not -isolated share the files they add. + directories it lists, on the driver. -A Java data source with the same name takes precedence over a native one, and so does a Python -data source. +The packages of a session take precedence over the installed libraries. Give packages distinct +file names, for example with their version: sessions that are not isolated share the files they +add. ## Example in Rust This data source reads the numbers `[0, end)` as a column `id`, in two partitions. It uses the [`jni`](https://crates.io/crates/jni) crate, and the `arrow-array` and `arrow-schema` crates of [arrow-rs](https://github.com/apache/arrow-rs), which export schemas and streams with -`FFI_ArrowSchema` and `FFI_ArrowArrayStream`. `Cargo.toml`: +`FFI_ArrowSchema` and `FFI_ArrowArrayStream`. `Cargo.toml`, where the name of the library is the +name that Spark finds when it is installed: ```toml [package] @@ -145,6 +183,7 @@ version = "0.1.0" edition = "2021" [lib] +name = "spark_datasource_rust_range" crate-type = ["cdylib"] [dependencies] @@ -321,17 +360,29 @@ pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_read( ``` {% endraw %} -Build the library, and package it: +Build the library, here on macOS, and install it: ```bash cargo build --release +cp target/release/libspark_datasource_rust_range.dylib /usr/local/lib/ +``` + +Spark then finds the data source `rust_range` automatically: + +```python +spark.read.format("rust_range").option("end", 100).load().show() +``` + +Or package it, to add it to a session with `spark.addArtifact`: + +```bash mkdir -p package/osx-aarch64 -cp target/release/librust_range.dylib package/osx-aarch64/ +cp target/release/libspark_datasource_rust_range.dylib package/osx-aarch64/ cat > package/spark-native-datasource.json < fileName, "checksum" -> checksum)) } + def invalidNativeDataSourceLibraryError( + path: String, + subClass: String, + messageParameters: Map[String, String], + cause: Throwable = null): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = s"INVALID_NATIVE_DATA_SOURCE_LIBRARY.$subClass", + messageParameters = Map("path" -> path) ++ messageParameters, + cause = cause) + } + + def nativeDataSourceLibraryNotFoundError( + fileName: String, path: String): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND", + messageParameters = Map("fileName" -> fileName, "path" -> path)) + } + def removedClassInSpark2Error(className: String, e: Throwable): SparkClassNotFoundException = { new SparkClassNotFoundException( errorClass = "_LEGACY_ERROR_TEMP_2052", diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index df5ed6e95c5c8..787b58857116c 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -7212,8 +7212,8 @@ object SQLConf { "directories that contain them, on the local file system of the driver. Spark adds a " + "package to the artifacts of the session when one of its data sources is used, so that " + "the executors get it too. Packages can also be added with spark.addArtifact. These " + - "packages take precedence over the ones installed in the native-datasources directory " + - "of SPARK_HOME. See spark.sql.dataSource.native.enabled.") + "packages take precedence over the native data source libraries installed in the native " + + "library path. See spark.sql.dataSource.native.enabled.") .version("4.4.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .stringConf diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala index 87d4c89c9a0f9..6da375360c4e5 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala @@ -467,9 +467,12 @@ object StaticSQLConf { val NATIVE_DATA_SOURCE_ENABLED = buildStaticConf("spark.sql.dataSource.native.enabled") .doc("When true, Spark finds and loads native data sources: data sources implemented in " + - "native code such as Rust or C++, and distributed as native data source packages " + - "(.sparkpkg files). A package is found when it is installed in the native-datasources " + - "directory of SPARK_HOME, when it is added with spark.addArtifact, or when it is under " + + "native code such as Rust or C++. The library of a data source, such as my_source, is " + + "found when it is installed in the native library path as spark_datasource_my_source, " + + "for example libspark_datasource_my_source.so in /usr/local/lib on Linux. The native " + + "library path consists of the directories of java.library.path, and the lib directory " + + "of each installation prefix in PATH. Libraries can also be distributed in native data " + + "source packages (.sparkpkg files), added with spark.addArtifact or under " + "spark.sql.dataSource.native.paths. Loading a native data source runs its native code " + "in the driver and executor processes.") .version("4.4.0") diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala new file mode 100644 index 0000000000000..8cfe0e7e7c717 --- /dev/null +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala @@ -0,0 +1,90 @@ +/* + * 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.connect + +import java.nio.file.Files +import java.util.UUID + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.SparkException +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.execution.datasources.v2.ffi.NativeDataSourceTestUtils._ +import org.apache.spark.util.Utils + +/** + * Tests native data sources with Spark Connect: a client adds packages to its session with + * `spark.addArtifact`, and the server finds the data sources in the packages of the session or in + * the installed libraries. The tests are skipped when there is no C++ compiler or no JNI headers. + */ +class NativeDataSourceConnectSuite extends SparkConnectServerTest { + + private lazy val library = compile("test_native_datasource") + private lazy val installedLibrary = + compile("test_native_datasource_installed", "-DNAME_PREFIX=\"installed_\"") + + private def nativeTest(name: String)(f: => Unit): Unit = test(name) { + assume(toolchain.isDefined, "A C++ compiler and the JNI headers are required.") + f + } + + private def withNewSession(f: SparkSession => Unit): Unit = { + withSession(sessionId = UUID.randomUUID().toString)(f) + } + + nativeTest("add packages with spark.addArtifact") { + val pkg = createPackage(library, dataSources("")) + withNewSession { session => + session.addArtifact(pkg.getPath) + val df = session.read.format("native_range").option("end", "3").load() + assert(df.select("id").collect().map(_.getLong(0)).sorted.toSeq == Seq(0L, 1L, 2L)) + + val dir = Utils.createTempDir() + session.range(2).selectExpr("id", "concat('v', id) AS name") + .write.format("native_sink").option("path", dir.getPath).mode("append").save() + assert(dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala).toSet == Set("0,v0", "1,v1")) + + try { + session.sql("CREATE TABLE native_connect_table USING native_range OPTIONS (end 2)") + assert(session.sql("SELECT id FROM native_connect_table").collect() + .map(_.getLong(0)).sorted.toSeq == Seq(0L, 1L)) + } finally { + session.sql("DROP TABLE IF EXISTS native_connect_table") + } + } + + // The packages of a session are not visible to the other sessions. + withNewSession { other => + val e = intercept[SparkException] { + other.read.format("native_range").load().collect() + } + assert(e.getCondition == "DATA_SOURCE_NOT_FOUND") + } + } + + nativeTest("find installed libraries automatically") { + val dir = Utils.createTempDir() + installLibrary(installedLibrary, "installed_native_range", dir) + withLibraryPath(dir) { + withNewSession { session => + val df = session.read.format("installed_native_range").option("end", "3").load() + assert(df.select("id").collect().map(_.getLong(0)).sorted.toSeq == Seq(0L, 1L, 2L)) + } + } + } +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java index d6e5dc27ef882..b49d3027ce13c 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java @@ -31,11 +31,15 @@ * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and * limit pushdown. *

      - * The library is distributed in a native data source package: a zip file with the extension - * {@code .sparkpkg} that contains a manifest named {@code spark-native-datasource.json} and the - * library built for one or more platforms. Spark finds the packages installed in the - * {@code native-datasources} directory of {@code SPARK_HOME} automatically, as well as the - * packages added to a session with {@code spark.addArtifact} and the ones under the paths in + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of {@code my_source} is + * {@code spark_datasource_my_source}, such as {@code libspark_datasource_my_source.so} in + * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of + * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux, and in the + * {@code lib} directory of each installation prefix in {@code PATH}. A library can also be + * distributed in a native data source package: a zip file with the extension {@code .sparkpkg} + * that contains a manifest named {@code spark-native-datasource.json} and the library built for + * one or more platforms, added to a session with {@code spark.addArtifact} or under the paths in * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL * guide for details and examples. * diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index d52194b5d70b1..633561f7a6c67 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -44,17 +44,19 @@ import org.apache.spark.util.Utils * A [[ColumnarDataSource]] implemented by a native library, through * `org.apache.spark.sql.datasource.NativeBridge`. * - * @param name the name of the data source, as listed in the manifest of the package + * @param location where the library is + * @param name the name of the data source, as listed in the manifest of its package or in lower + * case for an installed library */ class NativeDataSource( - pkg: NativeDataSourcePackage, + location: NativeLibraryLocation, name: String, options: CaseInsensitiveStringMap) extends ColumnarDataSource { import NativeCalls._ override def schema(): StructType = { - val library = NativeLibraries.get(pkg) + val library = NativeLibraries.get(location) withDataSource(library, "infer the schema of") { dataSource => call(name, "infer the schema of") { optional(NativeArrow.importSchema(address => library.schema(dataSource, address))) @@ -63,26 +65,27 @@ class NativeDataSource( } override def reader(schema: StructType): ColumnarReader = { - val library = NativeLibraries.get(pkg) + val library = NativeLibraries.get(location) val handle = create(library, "plan a scan of", "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { dataSource => NativeArrow.exportSchema(schema)(address => library.createReader(dataSource, address)) } - new NativeDataSourceReader(pkg, name, schema, new NativeHandle(handle, library.closeReader)) + new NativeDataSourceReader( + location, name, schema, new NativeHandle(handle, library.closeReader)) } override def streamReader(schema: StructType): ColumnarStreamReader = { - val library = NativeLibraries.get(pkg) + val library = NativeLibraries.get(location) val handle = create(library, "plan a scan of", "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { dataSource => NativeArrow.exportSchema(schema)(address => library.createStreamReader(dataSource, address)) } new NativeDataSourceStreamReader( - pkg, name, schema, new NativeHandle(handle, library.closeStreamReader)) + location, name, schema, new NativeHandle(handle, library.closeStreamReader)) } override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { - val library = NativeLibraries.get(pkg) + val library = NativeLibraries.get(location) val handle = new NativeHandle( create(library, "plan a write to", "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource => NativeArrow.exportSchema(schema) { address => @@ -90,11 +93,11 @@ class NativeDataSource( } }, library.closeWriter) - new NativeDataSourceWriter(pkg, name, handle, serializeWriter(library, handle)) + new NativeDataSourceWriter(location, name, handle, serializeWriter(library, handle)) } override def streamWriter(schema: StructType, overwrite: Boolean): ColumnarStreamWriter = { - val library = NativeLibraries.get(pkg) + val library = NativeLibraries.get(location) val handle = new NativeHandle( create(library, "plan a write to", "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { dataSource => @@ -103,7 +106,7 @@ class NativeDataSource( } }, library.closeWriter) - new NativeDataSourceStreamWriter(pkg, name, handle, serializeWriter(library, handle)) + new NativeDataSourceStreamWriter(location, name, handle, serializeWriter(library, handle)) } /** Runs `f` with a new native data source, which is closed afterwards. */ @@ -141,7 +144,7 @@ class NativeDataSource( /** Reads a native data source for a batch scan. */ class NativeDataSourceReader( - pkg: NativeDataSourcePackage, + location: NativeLibraryLocation, name: String, schema: StructType, @transient private val handle: NativeHandle) @@ -153,7 +156,7 @@ class NativeDataSourceReader( private var readSchema: StructType = schema @transient private var plannedPartitions: Array[InputPartition] = _ - @transient private lazy val library = NativeLibraries.get(pkg) + @transient private lazy val library = NativeLibraries.get(location) override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { // Only the predicates that can be expressed in JSON are passed to the library. @@ -208,13 +211,13 @@ class NativeDataSourceReader( override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { new NativePartitionReader( - pkg, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, readSchema) + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, readSchema) } } /** Reads a native data source for a micro-batch streaming scan. */ class NativeDataSourceStreamReader( - pkg: NativeDataSourcePackage, + location: NativeLibraryLocation, name: String, schema: StructType, @transient private val handle: NativeHandle) @@ -225,7 +228,7 @@ class NativeDataSourceStreamReader( // the executors with the partitions of that micro-batch. private var readerState: Array[Byte] = Array.emptyByteArray - @transient private lazy val library = NativeLibraries.get(pkg) + @transient private lazy val library = NativeLibraries.get(location) override def initialOffset(): String = { call(name, "plan a scan of")(library.initialOffset(handle.get)) @@ -246,7 +249,7 @@ class NativeDataSourceStreamReader( override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { new NativePartitionReader( - pkg, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, schema) + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, schema) } override def commit(end: String): Unit = { @@ -258,17 +261,17 @@ class NativeDataSourceStreamReader( /** Writes to a native data source for a batch write. */ class NativeDataSourceWriter( - pkg: NativeDataSourcePackage, + location: NativeLibraryLocation, name: String, @transient private val handle: NativeHandle, writerState: Array[Byte]) extends ColumnarWriter { import NativeCalls._ - @transient private lazy val library = NativeLibraries.get(pkg) + @transient private lazy val library = NativeLibraries.get(location) override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { - new NativeDataWriter(pkg, name, writerState, partitionId, taskId, epochId = -1L) + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId = -1L) } override def commit(messages: Array[WriterCommitMessage]): Unit = { @@ -289,20 +292,20 @@ class NativeDataSourceWriter( /** Writes to a native data source for a streaming write. */ class NativeDataSourceStreamWriter( - pkg: NativeDataSourcePackage, + location: NativeLibraryLocation, name: String, @transient private val handle: NativeHandle, writerState: Array[Byte]) extends ColumnarStreamWriter { import NativeCalls._ - @transient private lazy val library = NativeLibraries.get(pkg) + @transient private lazy val library = NativeLibraries.get(location) override def createWriter( partitionId: Int, taskId: Long, epochId: Long): DataWriter[ColumnarBatch] = { - new NativeDataWriter(pkg, name, writerState, partitionId, taskId, epochId) + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId) } override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { @@ -346,7 +349,7 @@ object NativeWriterCommitMessage { * imported into Arrow vectors that back the returned batch. */ class NativePartitionReader( - pkg: NativeDataSourcePackage, + location: NativeLibraryLocation, name: String, readerState: Array[Byte], partition: Array[Byte], @@ -373,7 +376,7 @@ class NativePartitionReader( } private def openStream(): ArrowReader = { - val library = NativeLibraries.get(pkg) + val library = NativeLibraries.get(location) val stream = ArrowArrayStream.allocateNew(allocator) try { call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) @@ -432,7 +435,7 @@ class NativePartitionReader( /** Writes the batches of a task to the data writer of a native library. */ class NativeDataWriter( - pkg: NativeDataSourcePackage, + location: NativeLibraryLocation, name: String, writerState: Array[Byte], partitionId: Int, @@ -441,7 +444,7 @@ class NativeDataWriter( extends DataWriter[ColumnarBatch] { import NativeCalls._ - private val library = NativeLibraries.get(pkg) + private val library = NativeLibraries.get(location) // Zero once the data writer was committed or aborted, which releases it. private var handle: Long = call(name, "write to") { library.createDataWriter(writerState, partitionId, taskId, epochId) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala index 6fe7a9b8e1c9b..b8e393f816492 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -25,6 +25,7 @@ import java.util.concurrent.ConcurrentHashMap import java.util.zip.ZipFile import scala.jdk.CollectionConverters._ +import scala.util.Try import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} @@ -58,6 +59,13 @@ case class NativeDataSourceManifest( dataSources: Seq[String], libraries: Map[String, String]) +/** + * Where the library of a native data source is: in a [[NativeDataSourcePackage]], or installed in + * the native library path, as an [[InstalledNativeLibrary]]. It is sent to the executors, which + * load the library with [[NativeLibraries.get]]. + */ +sealed trait NativeLibraryLocation extends Serializable + /** * A native data source package: a zip file that contains a [[NativeDataSourceManifest]] and the * libraries it lists. The SHA-256 checksum of the file identifies the package, so that an @@ -71,7 +79,7 @@ case class NativeDataSourcePackage( path: String, sha256: String, manifest: NativeDataSourceManifest, - artifactUUID: Option[String] = None) { + artifactUUID: Option[String] = None) extends NativeLibraryLocation { def fileName: String = new File(path).getName } @@ -161,6 +169,14 @@ object NativeDataSourcePackage { } } +/** + * A native data source library installed in the native library path, where + * [[NativeDataSourceRegistry]] finds it by its file name. It must be installed on every node. + * + * @param path the path of the library on the driver + */ +case class InstalledNativeLibrary(path: String) extends NativeLibraryLocation + /** The platform names used in [[NativeDataSourceManifest.libraries]]. */ object NativePlatform { /** The platform of this JVM, such as `linux-x86_64`, `linux-aarch64` or `osx-aarch64`. */ @@ -183,21 +199,34 @@ object NativePlatform { } /** - * The native data source libraries loaded in this JVM. A library is extracted from its package - * and loaded once, on first use, and stays loaded until the JVM exits. + * The native data source libraries loaded in this JVM. A library is loaded once, on first use, + * after it is extracted from its package if needed, and stays loaded until the JVM exits. */ object NativeLibraries extends Logging { - private val libraries = new ConcurrentHashMap[String, NativeLibrary]() + // The libraries of the packages, by the checksum of the package, and the installed libraries, + // by their canonical path: the JVM loads a library file in a single class loader. A library + // that is loaded but rejected stays loaded, so its error is kept: it cannot be loaded again. + private val packagedLibraries = new ConcurrentHashMap[String, Try[NativeLibrary]]() + private val installedLibraries = new ConcurrentHashMap[String, Try[NativeLibrary]]() private lazy val extractionDir: File = Utils.createTempDir(namePrefix = "native-datasources") - /** Returns the library of the given package, and loads it if needed. */ - private[ffi] def get(pkg: NativeDataSourcePackage): NativeLibrary = { - val library = libraries.get(pkg.sha256) - if (library != null) library else libraries.computeIfAbsent(pkg.sha256, _ => load(pkg)) + /** Returns the library at the given location, and loads it if needed. */ + private[ffi] def get(location: NativeLibraryLocation): NativeLibrary = location match { + case pkg: NativeDataSourcePackage => + getOrLoad(packagedLibraries, pkg.sha256)(loadPackaged(pkg)) + case installed: InstalledNativeLibrary => + val file = locate(installed).getCanonicalFile + getOrLoad(installedLibraries, file.getPath)(loadInstalled(file)) + } + + private def getOrLoad(libraries: ConcurrentHashMap[String, Try[NativeLibrary]], key: String)( + load: => Try[NativeLibrary]): NativeLibrary = { + val library = libraries.get(key) + (if (library != null) library else libraries.computeIfAbsent(key, _ => load)).get } - private def load(pkg: NativeDataSourcePackage): NativeLibrary = { + private def loadPackaged(pkg: NativeDataSourcePackage): Try[NativeLibrary] = { val platform = NativePlatform.current val libraryPath = pkg.manifest.libraries.getOrElse(platform, { throw QueryExecutionErrors.invalidNativeDataSourcePackageError( @@ -212,30 +241,54 @@ object NativeLibraries extends Logging { if (!libraryFile.isFile) { extract(packageFile, libraryPath, libraryFile) } - logInfo(log"Loading the native data source library ${MDC(PATH, libraryFile)}.") + load( + libraryFile, + (reason, cause) => QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, "INVALID_LIBRARY", Map("library" -> libraryPath, "reason" -> reason), cause), + parameters => QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, "UNSUPPORTED_ABI_VERSION", parameters)) + } - def invalidLibrary(reason: String, cause: Throwable): Throwable = { - QueryExecutionErrors.invalidNativeDataSourcePackageError( - pkg.path, "INVALID_LIBRARY", Map("library" -> libraryPath, "reason" -> reason), cause) - } + private def loadInstalled(file: File): Try[NativeLibrary] = { + load( + file, + (reason, cause) => QueryExecutionErrors.invalidNativeDataSourceLibraryError( + file.getPath, "CANNOT_LOAD", Map("reason" -> reason), cause), + parameters => QueryExecutionErrors.invalidNativeDataSourceLibraryError( + file.getPath, "UNSUPPORTED_ABI_VERSION", parameters)) + } + + /** + * Loads a library, and checks the version of the interface that it implements. Throws the error + * if the library cannot be loaded, and returns it if the library is loaded but rejected. + * + * @param invalidLibrary returns the error for an invalid library, with the reason + * @param unsupportedAbiVersion returns the error for an unsupported version of the interface, + * with the parameters `version` and `supported` + */ + private def load( + file: File, + invalidLibrary: (String, Throwable) => Throwable, + unsupportedAbiVersion: Map[String, String] => Throwable): Try[NativeLibrary] = { + logInfo(log"Loading the native data source library ${MDC(PATH, file)}.") val library = try { - NativeLibrary.load(libraryFile.getPath) + NativeLibrary.load(file.getPath) } catch { case e: UnsatisfiedLinkError => throw invalidLibrary(e.getMessage, e) } - val abiVersion = try { - library.abiVersion() - } catch { - case e: UnsatisfiedLinkError => - throw invalidLibrary("It does not implement NativeBridge.abiVersion.", e) - } - if (abiVersion != NativeBridge.ABI_VERSION) { - throw QueryExecutionErrors.invalidNativeDataSourcePackageError( - pkg.path, - "UNSUPPORTED_ABI_VERSION", - Map("version" -> abiVersion.toString, "supported" -> NativeBridge.ABI_VERSION.toString)) + Try { + val abiVersion = try { + library.abiVersion() + } catch { + case e: UnsatisfiedLinkError => + throw invalidLibrary("It does not implement NativeBridge.abiVersion.", e) + } + if (abiVersion != NativeBridge.ABI_VERSION) { + throw unsupportedAbiVersion( + Map("version" -> abiVersion.toString, "supported" -> NativeBridge.ABI_VERSION.toString)) + } + library } - library } /** Finds a copy of the package on this node: the one on the driver, or a distributed one. */ @@ -247,6 +300,22 @@ object NativeLibraries extends Logging { } } + /** + * Finds an installed library on this node: where the driver found it, or else in the native + * library path of this node. + */ + private def locate(installed: InstalledNativeLibrary): File = { + val file = new File(installed.path) + if (file.isFile) { + file + } else { + NativeDataSourceRegistry.findInstalledLibrary(file.getName).getOrElse { + throw QueryExecutionErrors.nativeDataSourceLibraryNotFoundError( + file.getName, installed.path) + } + } + } + private def extract(packageFile: File, libraryPath: String, target: File): Unit = { target.getParentFile.mkdirs() val temp = File.createTempFile(target.getName, ".tmp", target.getParentFile) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala index 3b9ccca628775..2aef77d05b572 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -18,8 +18,9 @@ package org.apache.spark.sql.execution.datasources.v2.ffi import java.io.File import java.nio.file.Paths +import java.util.Locale -import org.apache.spark.{SparkContext, SparkEnv, SparkFiles} +import org.apache.spark.{SparkEnv, SparkFiles} import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.PATH import org.apache.spark.sql.Artifact @@ -30,33 +31,37 @@ import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.util.Utils /** - * Finds the native data source packages on the driver. Like the Python data sources installed in - * the Python path, the packages installed in `$SPARK_HOME/native-datasources` are found - * automatically. The active session can also add packages with `spark.addArtifact`, or with - * `spark.sql.dataSource.native.paths`, which take precedence over the installed ones. + * Finds the native data sources on the driver. Like the Python data sources installed in the + * Python path, the native data source libraries installed in the native library path are found + * automatically, by their file name: the library `spark_datasource_`, such as + * `libspark_datasource_my_source.so` on Linux, implements the data source ``. The active + * session can also add packages with `spark.addArtifact`, or with + * `spark.sql.dataSource.native.paths`, which take precedence over the installed libraries. */ object NativeDataSourceRegistry extends Logging { - /** The directory of `SPARK_HOME` that contains the installed native data source packages. */ - val INSTALLED_PACKAGES_DIR = "native-datasources" + /** The prefix of the names of the installed native data source libraries. */ + val LIBRARY_NAME_PREFIX = "spark_datasource_" - /** Returns whether a package provides the native data source with the given name. */ + /** Returns whether the native data source with the given name exists. */ def exists(name: String, conf: SQLConf): Boolean = lookup(name, conf).isDefined /** - * Finds the package that provides the native data source with the given name, which is - * case-insensitive. Returns the name of the data source as listed in the manifest, and the - * package. + * Finds the package or the installed library that provides the native data source with the + * given name, which is case-insensitive. Returns the name of the data source, as listed in the + * manifest of the package or in lower case for an installed library, and where its library is. */ - def lookup(name: String, conf: SQLConf): Option[(String, NativeDataSourcePackage)] = { + def lookup(name: String, conf: SQLConf): Option[(String, NativeLibraryLocation)] = { if (!conf.getConf(StaticSQLConf.NATIVE_DATA_SOURCE_ENABLED)) { return None } - lookup(name, configuredPackageFiles(conf) ++ artifactPackageFiles()) - .orElse(lookup(name, installedPackageFiles())) + lookupPackage(name, configuredPackageFiles(conf) ++ artifactPackageFiles()) + .orElse(lookupInstalledLibrary(name)) } - private def lookup(name: String, files: Seq[File]): Option[(String, NativeDataSourcePackage)] = { + private def lookupPackage( + name: String, + files: Seq[File]): Option[(String, NativeDataSourcePackage)] = { val found = files.map(_.getCanonicalFile).distinct.map(NativeDataSourcePackage.read) .flatMap(pkg => pkg.manifest.dataSources.find(_.equalsIgnoreCase(name)).map(_ -> pkg)) // The same package can be found more than once, for example when it is in a configured @@ -80,16 +85,62 @@ object NativeDataSourceRegistry extends Logging { } } + private def lookupInstalledLibrary(name: String): Option[(String, InstalledNativeLibrary)] = { + // Other names, such as class names, are not the names of libraries. + if (!name.matches("[A-Za-z0-9_]+")) { + return None + } + val dataSource = name.toLowerCase(Locale.ROOT) + findInstalledLibrary(libraryFileName(dataSource)) + .map(file => dataSource -> InstalledNativeLibrary(file.getPath)) + } + + /** Finds the library with the given file name in the native library path of this node. */ + private[ffi] def findInstalledLibrary(fileName: String): Option[File] = { + librarySearchPath(System.getProperty("java.library.path"), System.getenv("PATH")) + .map(new File(_, fileName)) + .find(_.isFile) + } + /** - * Makes the package available to the executors of the active session. Returns the package with - * the location of the copies of the executors. + * The file name of the installed library of the native data source with the given name, such + * as `libspark_datasource_my_source.so` on Linux, `libspark_datasource_my_source.dylib` on + * macOS and `spark_datasource_my_source.dll` on Windows. */ - def distribute(pkg: NativeDataSourcePackage): NativeDataSourcePackage = { - SparkSession.getActiveSession match { - // In local mode, the executors run in the driver and read the package where it is. - case Some(session) if !session.sparkContext.isLocal => addToArtifacts(session, pkg) - case _ => pkg + def libraryFileName(name: String): String = { + System.mapLibraryName(LIBRARY_NAME_PREFIX + name.toLowerCase(Locale.ROOT)) + } + + /** + * The native library path: the directories where native data source libraries are installed, + * in the order they are searched. They are the directories of `java.library.path`, where the + * JVM finds native libraries, which include `LD_LIBRARY_PATH` on Linux, `DYLD_LIBRARY_PATH` on + * macOS and `PATH` on Windows, and then the `lib` directory of each installation prefix in + * `PATH`, such as `/usr/local/lib` for `/usr/local/bin`. Relative directories are ignored. + */ + private[ffi] def librarySearchPath(javaLibraryPath: String, path: String): Seq[File] = { + def directories(value: String): Seq[File] = { + Option(value).toSeq.flatMap(_.split(File.pathSeparator)).map(new File(_)).filter(_.isAbsolute) } + val prefixLibraries = directories(path) + .filter(_.getName == "bin") + .flatMap(bin => Option(bin.getParentFile)) + .map(new File(_, "lib")) + (directories(javaLibraryPath) ++ prefixLibraries).distinct + } + + /** + * Makes the library available to the executors of the active session. Returns where the + * executors find it. An installed library is already installed on every node. + */ + def distribute(location: NativeLibraryLocation): NativeLibraryLocation = location match { + case pkg: NativeDataSourcePackage => + SparkSession.getActiveSession match { + // In local mode, the executors run in the driver and read the package where it is. + case Some(session) if !session.sparkContext.isLocal => addToArtifacts(session, pkg) + case _ => pkg + } + case installed: InstalledNativeLibrary => installed } /** Adds the package as an artifact of the session, unless it already is one. */ @@ -119,14 +170,6 @@ object NativeDataSourceRegistry extends Logging { } } - /** The packages installed in `$SPARK_HOME/native-datasources`, if the directory exists. */ - private[ffi] def installedPackageFiles(): Seq[File] = { - SparkContext.getActive.flatMap(_.getSparkHome()).toSeq - .map(home => new File(home, INSTALLED_PACKAGES_DIR)) - .filter(_.isDirectory) - .flatMap(packageFilesIn) - } - private def configuredPackageFiles(conf: SQLConf): Seq[File] = { conf.getConf(SQLConf.NATIVE_DATA_SOURCE_PATHS).flatMap { path => val uri = Utils.resolveURI(path) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala index e02a3e5fc5407..e9a3fcf7d17e9 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala @@ -39,28 +39,28 @@ class NativeDataSourceV2 extends NamedTableProvider { shortName = name } - // The name of the data source as listed in the manifest, and its package. Resolved on first - // use, which is while the query is analyzed, with the configuration of its session. - private lazy val resolved: (String, NativeDataSourcePackage) = { + // The name of the data source, and where its library is. Resolved on first use, which is while + // the query is analyzed, with the configuration of its session. + private lazy val resolved: (String, NativeLibraryLocation) = { assert(shortName != null) - val (name, pkg) = NativeDataSourceRegistry.lookup(shortName, SQLConf.get).getOrElse { + val (name, location) = NativeDataSourceRegistry.lookup(shortName, SQLConf.get).getOrElse { throw QueryCompilationErrors.dataSourceDoesNotExist(shortName) } - (name, NativeDataSourceRegistry.distribute(pkg)) + (name, NativeDataSourceRegistry.distribute(location)) } override def inferSchema(options: CaseInsensitiveStringMap): StructType = { - val (name, pkg) = resolved - new NativeDataSource(pkg, name, options).schema() + val (name, location) = resolved + new NativeDataSource(location, name, options).schema() } override def getTable( schema: StructType, partitioning: Array[Transform], properties: util.Map[String, String]): Table = { - val (name, pkg) = resolved + val (name, location) = resolved new ColumnarTable( - shortName, schema, properties, options => new NativeDataSource(pkg, name, options)) + shortName, schema, properties, options => new NativeDataSource(location, name, options)) } override def supportsExternalMetadata(): Boolean = true diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 794017b227d59..5534361f3fd65 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -16,18 +16,16 @@ */ package org.apache.spark.sql.execution.datasources.v2.ffi -import java.io.{File, FileOutputStream} +import java.io.File import java.nio.charset.StandardCharsets.UTF_8 import java.nio.file.Files import java.sql.{Date, Timestamp} import java.time.{Instant, LocalDate} -import java.util.concurrent.atomic.AtomicInteger -import java.util.zip.{ZipEntry, ZipOutputStream} import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ -import org.apache.spark.{SparkConf, SparkRuntimeException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.{SparkClassNotFoundException, SparkRuntimeException, SparkThrowable, SparkUnsupportedOperationException} import org.apache.spark.sql.{AnalysisException, QueryTest, Row} import org.apache.spark.sql.classic.{DataFrame, SparkSession} import org.apache.spark.sql.connector.expressions.{Expression, FieldReference, GeneralScalarExpression, LiteralValue} @@ -43,57 +41,16 @@ import org.apache.spark.util.Utils /** * Tests native data sources with the library in `native-datasource/test_native_datasource.cc`, - * compiled with the C++ compiler of the machine. The tests that load a library are skipped when - * there is no C++ compiler or no JNI headers. + * compiled with the C++ compiler of the machine by [[NativeDataSourceTestUtils]]. The tests that + * load a library are skipped when there is no C++ compiler or no JNI headers. */ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { + import NativeDataSourceTestUtils._ - private lazy val buildDir = Utils.createTempDir(namePrefix = "native-datasource-test") - - // The SPARK_HOME of the suite, where packages are installed in INSTALLED_PACKAGES_DIR. - private lazy val installedSparkHome = Utils.createTempDir(namePrefix = "spark-home") - - override protected def sparkConf: SparkConf = { - super.sparkConf.set("spark.home", installedSparkHome.getPath) - } - - // The executors store the files of the sessions that are not isolated in the same directory, - // and jobs outside of SQL executions get the files of all the sessions. Each package has its - // own name, and the sessions are kept until the end of the suite, so that none of their - // artifacts is removed while another test runs. - private val packageCount = new AtomicInteger() + // Jobs outside of SQL executions get the files of all the sessions, so the sessions are kept + // until the end of the suite: none of their artifacts is removed while another test runs. private val sessions = ArrayBuffer.empty[SparkSession] - // The C++ compiler and the JNI include directories, if available. - private lazy val toolchain: Option[(String, Seq[String])] = { - val include = new File(System.getProperty("java.home"), "include") - val platformInclude = Option(include.listFiles()).toSeq.flatten - .find(dir => new File(dir, "jni_md.h").isFile) - val compiler = Seq("c++", "g++", "clang++").find(Utils.checkCommandAvailable) - for { - compiler <- compiler - platformInclude <- platformInclude - if new File(include, "jni.h").isFile - } yield (compiler, Seq(include.getPath, platformInclude.getPath)) - } - - private def compile(name: String, defines: String*): File = { - val (compiler, includes) = toolchain.get - val source = new File(buildDir, "test_native_datasource.cc") - if (!source.exists()) { - val in = getClass.getClassLoader - .getResourceAsStream("native-datasource/test_native_datasource.cc") - Utils.tryWithResource(in)(Files.copy(_, source.toPath)) - } - val library = new File(buildDir, System.mapLibraryName(name)) - val command = Seq(compiler, "-std=c++17", "-O1", "-shared", "-fPIC") ++ - includes.map("-I" + _) ++ defines ++ Seq("-o", library.getPath, source.getPath) - val process = new ProcessBuilder(command: _*).redirectErrorStream(true).start() - val output = new String(process.getInputStream.readAllBytes(), UTF_8) - assert(process.waitFor() == 0, s"Failed to compile the test library: $output") - library - } - // The test library and its variants, each implementing differently named data sources. private lazy val library = compile("test_native_datasource") private lazy val readOnlyLibrary = @@ -107,49 +64,6 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { private lazy val installedNegatingLibrary = compile( "test_native_datasource_installed_b", "-DNAME_PREFIX=\"installed_\"", "-DID_SIGN=-1") - private def dataSources(prefix: String): Seq[String] = - Seq("native_range", "native_sink", "native_counter").map(prefix + _) - - private def manifestJson( - dataSources: Seq[String], - libraries: Map[String, String], - abiVersion: Int = 1): String = { - val names = dataSources.map(n => "\"" + n + "\"").mkString(", ") - val paths = libraries.map { case (p, path) => "\"" + p + "\": \"" + path + "\"" } - .mkString(", ") - s"""{"abiVersion": $abiVersion, "dataSources": [$names], "libraries": {$paths}}""" - } - - private def createZip(dir: File, fileName: String, entries: (String, Array[Byte])*): File = { - dir.mkdirs() - val file = new File(dir, fileName) - Utils.tryWithResource(new ZipOutputStream(new FileOutputStream(file))) { zip => - entries.foreach { case (path, bytes) => - zip.putNextEntry(new ZipEntry(path)) - zip.write(bytes) - zip.closeEntry() - } - } - file - } - - /** Creates a package with the library built for this platform. */ - private def createPackage( - library: File, - dataSources: Seq[String], - dir: File = Utils.createTempDir(), - fileName: String = s"test-${packageCount.incrementAndGet()}.sparkpkg", - abiVersion: Int = 1, - platform: String = NativePlatform.current): File = { - val path = s"$platform/${library.getName}" - createZip( - dir, - fileName, - NativeDataSourcePackage.MANIFEST_NAME -> - manifestJson(dataSources, Map(platform -> path), abiVersion).getBytes(UTF_8), - path -> Files.readAllBytes(library.toPath)) - } - private lazy val defaultPackage = createPackage(library, dataSources("")) /** A new session that finds the native data source packages at the given paths. */ @@ -231,6 +145,19 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { System.getProperty("os.name"), System.getProperty("os.arch"))) } + test("native library path") { + def path(directories: String*): String = directories.mkString(File.pathSeparator) + // The directories of java.library.path, and then the lib directory of each installation + // prefix in PATH. Relative and duplicate directories are ignored. + assert(NativeDataSourceRegistry.librarySearchPath( + path("/opt/native", "", "relative", "/usr/lib"), + path("/usr/local/bin", "/home/user/.venv/bin", "/usr/sbin", "relative/bin", "/usr/local/bin") + ) == Seq("/opt/native", "/usr/lib", "/usr/local/lib", "/home/user/.venv/lib").map(new File(_))) + assert(NativeDataSourceRegistry.librarySearchPath(null, null).isEmpty) + assert(NativeDataSourceRegistry.libraryFileName("My_Source") == + System.mapLibraryName("spark_datasource_my_source")) + } + test("encode predicates as JSON") { def predicate(name: String, children: Expression*): Predicate = { new Predicate(name, children.toArray) @@ -276,7 +203,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(literal(Array(1), ArrayType(IntegerType)).isEmpty) } - test("disable native data sources") { + nativeTest("disable native data sources") { val conf = new SQLConf conf.setConf(SQLConf.NATIVE_DATA_SOURCE_PATHS, Seq(defaultPackage.getPath)) conf.setConf(StaticSQLConf.NATIVE_DATA_SOURCE_ENABLED, false) @@ -548,25 +475,51 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(distributed.artifactUUID == otherArtifactUUID) } - nativeTest("find installed packages automatically") { - val installDir = - new File(installedSparkHome, NativeDataSourceRegistry.INSTALLED_PACKAGES_DIR) - val installed = createPackage(installedLibrary, dataSources("installed_"), dir = installDir) - try { - assert(NativeDataSourceRegistry.installedPackageFiles() == Seq(installed)) - // The session neither adds the package nor configures a path. + nativeTest("find installed libraries automatically") { + val dir = Utils.createTempDir() + val installed = installLibrary(installedLibrary, "installed_native_range", dir) + // A library that implements several data sources is installed under each of their names. + Files.createSymbolicLink( + new File(dir, NativeDataSourceRegistry.libraryFileName("installed_native_sink")).toPath, + installed.toPath) + withLibraryPath(dir) { + // The session neither adds a package nor configures a path. The name is case-insensitive. + val session = newSession() checkRows( - newSession().read.format("installed_native_range").option("end", "3").load().select("id"), + session.read.format("Installed_Native_Range").option("end", "3").load().select("id"), Seq(Row(0L), Row(1L), Row(2L))) - - // A package of the session takes precedence over an installed one. - val session = newSession(createPackage(installedNegatingLibrary, dataSources("installed_"))) + val output = Utils.createTempDir() + session.range(2).selectExpr("id", "concat('v', id) AS name") + .write.format("installed_native_sink").option("path", output.getPath).mode("append").save() + assert(output.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala).toSet == Set("0,v0", "1,v1")) + + // A package of the session takes precedence over an installed library. + val other = newSession(createPackage(installedNegatingLibrary, dataSources("installed_"))) checkRows( - session.read.format("installed_native_range").option("end", "3").load().select("id"), + other.read.format("installed_native_range").option("end", "3").load().select("id"), Seq(Row(0L), Row(-1L), Row(-2L))) - } finally { - installed.delete() + + // An executor finds the library in its native library path, if it is not where the driver + // found it. + val elsewhere = new File(Utils.createTempDir(), installed.getName).getPath + assert(NativeLibraries.get(InstalledNativeLibrary(elsewhere)) eq + NativeLibraries.get(InstalledNativeLibrary(installed.getPath))) } + + val missing = new File(Utils.createTempDir(), installed.getName).getPath + checkError( + exception = intercept[SparkRuntimeException] { + NativeLibraries.get(InstalledNativeLibrary(missing)) + }, + condition = "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND", + parameters = Map("fileName" -> installed.getName, "path" -> missing)) + checkError( + exception = intercept[SparkClassNotFoundException] { + newSession().read.format("installed_native_range").load() + }, + condition = "DATA_SOURCE_NOT_FOUND", + parameters = Map("provider" -> "installed_native_range")) } nativeTest("find packages in directories") { @@ -655,4 +608,27 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { parameters = Map( "path" -> mismatchedAbi.getCanonicalPath, "version" -> "2", "supported" -> "1")) } + + nativeTest("invalid installed libraries") { + val dir = Utils.createTempDir() + val notALibrary = new File(dir, NativeDataSourceRegistry.libraryFileName("not_a_library")) + Files.writeString(notALibrary.toPath, "not a library") + val newerAbi = installLibrary(abiV2Library, "v2_native_range", dir) + withLibraryPath(dir) { + val e = intercept[SparkRuntimeException](newSession().read.format("not_a_library").load()) + assert(e.getCondition == "INVALID_NATIVE_DATA_SOURCE_LIBRARY.CANNOT_LOAD") + assert(e.getMessageParameters.get("path") == notALibrary.getCanonicalPath) + + // The library stays loaded, and the error is the same when it is used again. + (1 to 2).foreach { _ => + checkError( + exception = intercept[SparkRuntimeException] { + newSession().read.format("v2_native_range").load() + }, + condition = "INVALID_NATIVE_DATA_SOURCE_LIBRARY.UNSUPPORTED_ABI_VERSION", + parameters = Map( + "path" -> newerAbi.getCanonicalPath, "version" -> "2", "supported" -> "1")) + } + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala new file mode 100644 index 0000000000000..d638f782bc6cf --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala @@ -0,0 +1,148 @@ +/* + * 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.execution.datasources.v2.ffi + +import java.io.{File, FileOutputStream} +import java.nio.charset.StandardCharsets.UTF_8 +import java.nio.file.Files +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicInteger +import java.util.zip.{ZipEntry, ZipOutputStream} + +import org.apache.spark.util.Utils + +/** + * Builds the native data source library of the tests, + * `native-datasource/test_native_datasource.cc`, with the C++ compiler of the machine, and + * packages or installs it. The tests that load the library are skipped when [[toolchain]] is not + * defined. + */ +object NativeDataSourceTestUtils { + + private lazy val buildDir = Utils.createTempDir(namePrefix = "native-datasource-test") + + private val libraries = new ConcurrentHashMap[String, File]() + + // The executors store the files of the sessions that are not isolated in the same directory, so + // each package has its own name. + private val packageCount = new AtomicInteger() + + /** The C++ compiler and the JNI include directories, if available. */ + lazy val toolchain: Option[(String, Seq[String])] = { + val include = new File(System.getProperty("java.home"), "include") + val platformInclude = Option(include.listFiles()).toSeq.flatten + .find(dir => new File(dir, "jni_md.h").isFile) + val compiler = Seq("c++", "g++", "clang++").find(Utils.checkCommandAvailable) + for { + compiler <- compiler + platformInclude <- platformInclude + if new File(include, "jni.h").isFile + } yield (compiler, Seq(include.getPath, platformInclude.getPath)) + } + + /** + * Compiles the test library with the given preprocessor definitions, once for each name. The + * library implements the data sources [[dataSources]], with the prefix `NAME_PREFIX`. + */ + def compile(name: String, defines: String*): File = { + libraries.computeIfAbsent(name, _ => build(name, defines)) + } + + private def build(name: String, defines: Seq[String]): File = { + val (compiler, includes) = toolchain.get + val source = new File(buildDir, "test_native_datasource.cc") + if (!source.exists()) { + val in = getClass.getClassLoader + .getResourceAsStream("native-datasource/test_native_datasource.cc") + Utils.tryWithResource(in)(Files.copy(_, source.toPath)) + } + val library = new File(buildDir, System.mapLibraryName(name)) + val command = Seq(compiler, "-std=c++17", "-O1", "-shared", "-fPIC") ++ + includes.map("-I" + _) ++ defines ++ Seq("-o", library.getPath, source.getPath) + val process = new ProcessBuilder(command: _*).redirectErrorStream(true).start() + val output = new String(process.getInputStream.readAllBytes(), UTF_8) + assert(process.waitFor() == 0, s"Failed to compile the test library: $output") + library + } + + /** The data sources of the test library compiled with the given `NAME_PREFIX`. */ + def dataSources(prefix: String): Seq[String] = + Seq("native_range", "native_sink", "native_counter").map(prefix + _) + + def manifestJson( + dataSources: Seq[String], + libraries: Map[String, String], + abiVersion: Int = 1): String = { + val names = dataSources.map(n => "\"" + n + "\"").mkString(", ") + val paths = libraries.map { case (p, path) => "\"" + p + "\": \"" + path + "\"" } + .mkString(", ") + s"""{"abiVersion": $abiVersion, "dataSources": [$names], "libraries": {$paths}}""" + } + + def createZip(dir: File, fileName: String, entries: (String, Array[Byte])*): File = { + dir.mkdirs() + val file = new File(dir, fileName) + Utils.tryWithResource(new ZipOutputStream(new FileOutputStream(file))) { zip => + entries.foreach { case (path, bytes) => + zip.putNextEntry(new ZipEntry(path)) + zip.write(bytes) + zip.closeEntry() + } + } + file + } + + /** Creates a package with the library built for this platform. */ + def createPackage( + library: File, + dataSources: Seq[String], + dir: File = Utils.createTempDir(), + fileName: String = s"test-${packageCount.incrementAndGet()}.sparkpkg", + abiVersion: Int = 1, + platform: String = NativePlatform.current): File = { + val path = s"$platform/${library.getName}" + createZip( + dir, + fileName, + NativeDataSourcePackage.MANIFEST_NAME -> + manifestJson(dataSources, Map(platform -> path), abiVersion).getBytes(UTF_8), + path -> Files.readAllBytes(library.toPath)) + } + + /** Installs a library as the library of the given native data source, in a directory. */ + def installLibrary(library: File, dataSource: String, dir: File): File = { + val file = new File(dir, NativeDataSourceRegistry.libraryFileName(dataSource)) + Files.copy(library.toPath, file.toPath) + file + } + + /** Runs `f` with a directory added to `java.library.path`, in the native library path. */ + def withLibraryPath[T](dir: File)(f: => T): T = { + val original = System.getProperty("java.library.path") + System.setProperty( + "java.library.path", (dir.getPath +: Option(original).toSeq).mkString(File.pathSeparator)) + try { + f + } finally { + if (original == null) { + System.clearProperty("java.library.path") + } else { + System.setProperty("java.library.path", original) + } + } + } +} From 822d33c534c260ad8db48e2cf22071d30be3e2cc Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 9 Oct 2026 14:59:15 +0900 Subject: [PATCH 07/27] Format NativeDataSourceConnectSuite with scalafmt The Scala linter checks the format of the Spark Connect modules with scalafmt. Co-authored-by: Isaac --- .../NativeDataSourceConnectSuite.scala | 27 ++++++++++++++----- 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala index 8cfe0e7e7c717..efef2b08240a7 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala @@ -54,15 +54,30 @@ class NativeDataSourceConnectSuite extends SparkConnectServerTest { assert(df.select("id").collect().map(_.getLong(0)).sorted.toSeq == Seq(0L, 1L, 2L)) val dir = Utils.createTempDir() - session.range(2).selectExpr("id", "concat('v', id) AS name") - .write.format("native_sink").option("path", dir.getPath).mode("append").save() - assert(dir.listFiles().filter(_.getName.startsWith("part-")) - .flatMap(file => Files.readAllLines(file.toPath).asScala).toSet == Set("0,v0", "1,v1")) + session + .range(2) + .selectExpr("id", "concat('v', id) AS name") + .write + .format("native_sink") + .option("path", dir.getPath) + .mode("append") + .save() + assert( + dir + .listFiles() + .filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala) + .toSet == Set("0,v0", "1,v1")) try { session.sql("CREATE TABLE native_connect_table USING native_range OPTIONS (end 2)") - assert(session.sql("SELECT id FROM native_connect_table").collect() - .map(_.getLong(0)).sorted.toSeq == Seq(0L, 1L)) + assert( + session + .sql("SELECT id FROM native_connect_table") + .collect() + .map(_.getLong(0)) + .sorted + .toSeq == Seq(0L, 1L)) } finally { session.sql("DROP TABLE IF EXISTS native_connect_table") } From 9dd4f88915052a0d47d01c5e16e2052c42cbb06f Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Sat, 10 Oct 2026 10:57:29 +0900 Subject: [PATCH 08/27] Address review comments on native data source lookup, distribution and lifecycle - Skip and log invalid packages during the lookup, unless the manifest of the invalid package lists the requested data source. - Distribute a package when its reader or writer is created, in the active session, which is the cloned session of a streaming query; executors also check the copies in the files of the other sessions by checksum. - Route .sparkpkg files to the JVM in classic addArtifacts regardless of the flags, also in mixed lists, as Spark Connect does. - Do not push predicates on columns with non-binary collated strings, and document the Spark SQL semantics (NaN, -0.0, nulls) of fully evaluated predicates in NativeBridge. - Report a clear conflict when a different package with the same file name is already an artifact of the session. - Close the previous stream writer of a streaming query when it creates a new one, and document the stream writer lifecycle. - Keep the data writer handle when commitDataWriter fails, so that Spark aborts it with abortDataWriter. - Reuse the lookup of DataSource.lookupDataSource in NativeDataSourceV2. - Skip native data sources when registering a Python data source, which takes precedence over them. - Describe the scan by the name of the data source. Co-authored-by: Isaac --- .../resources/error/error-conditions.json | 6 + docs/sql-native-data-sources.md | 15 +- python/pyspark/sql/session.py | 9 +- python/pyspark/sql/tests/test_artifact.py | 24 ++- .../sql/errors/QueryExecutionErrors.scala | 7 + .../spark/sql/datasource/NativeBridge.java | 28 ++- .../spark/sql/DataSourceRegistration.scala | 4 +- .../execution/datasources/DataSource.scala | 17 +- .../v2/columnar/ColumnarDataSource.scala | 6 +- .../v2/columnar/ColumnarTable.scala | 15 +- .../v2/columnar/ColumnarWrite.scala | 2 +- .../datasources/v2/ffi/NativeDataSource.scala | 64 +++++-- .../v2/ffi/NativeDataSourcePackage.scala | 84 ++++++--- .../v2/ffi/NativeDataSourceRegistry.scala | 65 +++++-- .../v2/ffi/NativeDataSourceV2.scala | 10 +- .../datasources/v2/ffi/NativePredicates.scala | 76 +++++--- .../test_native_datasource.cc | 8 +- .../v2/columnar/ColumnarDataSourceSuite.scala | 5 +- .../v2/ffi/NativeDataSourceSuite.scala | 170 +++++++++++++++++- 19 files changed, 499 insertions(+), 116 deletions(-) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index a20ee51b340d6..f10e7253a3908 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -6326,6 +6326,12 @@ ], "sqlState" : "42710" }, + "NATIVE_DATA_SOURCE_PACKAGE_FILE_NAME_CONFLICT" : { + "message" : [ + "Cannot add the native data source package to the session, because the session already has a different file named . Rename one of the packages." + ], + "sqlState" : "42710" + }, "NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND" : { "message" : [ "Cannot find the native data source package with the SHA-256 checksum . Add the package with `spark.addArtifact`, or install the same package under `spark.sql.dataSource.native.paths` on every node." diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index dc79c7cc8a61a..08b1fb715b732 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -60,8 +60,12 @@ Spark calls the functions of the library as follows: and calls `read` on the executors. Offsets are JSON strings defined by the library. - **Writes**: on the driver, Spark creates a writer with `createWriter` or `createStreamWriter`. On the executors, each task creates a data writer with `createDataWriter`, writes the batches of its - data with `write`, and calls `commitDataWriter`. Then the driver calls `commit` with the messages - of the tasks, or `abort` if the write failed. + data with `write`, and calls `commitDataWriter`, or `abortDataWriter` if the task or the commit + fails. Then the driver calls `commit` with the messages of the tasks, or `abort` if the write + failed. A micro-batch streaming query creates a new stream writer for each micro-batch, and Spark + closes the previous one when it creates the next one. Because Spark is not told when a streaming + query stops, the last stream writer of a query is closed once it is no longer referenced, or + when the query is restarted. The library defines the format of the partitions, of the state of its readers and writers, and of the commit messages, which Spark sends to the executors and back as byte arrays. A library reports @@ -82,7 +86,9 @@ are optional: The [Javadoc of `NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) specifies each function, including the ownership of the Arrow structs and the JSON format of the -pushed predicates. +pushed predicates. A predicate that a library reports as evaluated completely must follow the +semantics of Spark SQL, for example for NaN values, which are equal to each other and greater +than any other value. ## Finding Native Data Sources @@ -166,7 +172,8 @@ A session finds the packages in two places, and the executors get them automatic The packages of a session take precedence over the installed libraries. Give packages distinct file names, for example with their version: sessions that are not isolated share the files they -add. +add, and a session cannot use two different packages with the same file name. Spark ignores, and +logs, a package that it cannot read, unless its manifest lists the data source that a query uses. ## Example in Rust diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index 3ba7fc1c0977b..0be056e8ca960 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -2296,11 +2296,12 @@ def addArtifacts( "condition": "True together", }, ) - if not (file or pyfile or archive) and all(p.endswith(".sparkpkg") for p in path): - # Native data source packages are artifacts of the session in the JVM. - for p in path: + # Native data source packages are artifacts of the session in the JVM, whatever the + # flags, as with Spark Connect. + for p in path: + if p.endswith(".sparkpkg"): self._jsparkSession.addArtifact(os.path.abspath(p)) - return + path = tuple(p for p in path if not p.endswith(".sparkpkg")) for p in path: normalized_path = os.path.abspath(p) target_dir = os.path.join( diff --git a/python/pyspark/sql/tests/test_artifact.py b/python/pyspark/sql/tests/test_artifact.py index b6a40e3947042..f681e60552e4f 100644 --- a/python/pyspark/sql/tests/test_artifact.py +++ b/python/pyspark/sql/tests/test_artifact.py @@ -106,16 +106,26 @@ def test_artifacts_cannot_be_overwritten(self): ) def test_add_native_data_source_package(self): - # A new session, so that the package does not affect the other tests. + # A new session, so that the packages do not affect the other tests. spark = self.spark.newSession() with tempfile.TemporaryDirectory(prefix="test_add_native_data_source_package") as d: - package_path = os.path.join(d, "my_source.sparkpkg") - with open(package_path, "wb") as f: - f.write(b"package") - spark.addArtifact(package_path) + + def new_file(name): + file_path = os.path.join(d, name) + with open(file_path, "wb") as f: + f.write(name.encode("utf-8")) + return file_path + + spark.addArtifact(new_file("my_source.sparkpkg")) + # Packages are added to the session whatever the flags, as with Spark Connect, also + # together with other files. + spark.addArtifact(new_file("my_file_source.sparkpkg"), file=True) + spark.addArtifacts(new_file("my_other_source.sparkpkg"), new_file("other.txt")) packages = spark._jsparkSession.artifactManager().getNativeDataSourcePackages()._1() - self.assertEqual(packages.size(), 1) - self.assertEqual(packages.head().getName(), "my_source.sparkpkg") + self.assertEqual( + sorted(packages.apply(i).getName() for i in range(packages.size())), + ["my_file_source.sparkpkg", "my_other_source.sparkpkg", "my_source.sparkpkg"], + ) def test_add_zipped_package(self): self.check_add_zipped_package(self.spark) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala index 862d3c180981a..ccaf02e452839 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala @@ -828,6 +828,13 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE messageParameters = Map("fileName" -> fileName, "checksum" -> checksum)) } + def nativeDataSourcePackageFileNameConflictError( + path: String, fileName: String): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "NATIVE_DATA_SOURCE_PACKAGE_FILE_NAME_CONFLICT", + messageParameters = Map("path" -> path, "fileName" -> fileName)) + } + def invalidNativeDataSourceLibraryError( path: String, subClass: String, diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java index b49d3027ce13c..f7adf1e14dcc8 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java @@ -62,16 +62,22 @@ * {@link #createStreamWriter}, and serializes its state with {@link #serializeWriter}. On the * executors, each task creates a data writer with {@link #createDataWriter}, calls * {@link #write} for each batch of at most {@code spark.sql.execution.arrow.maxRecordsPerBatch} - * rows, and calls {@link #commitDataWriter}, or {@link #abortDataWriter} if it fails. Then the - * driver calls {@link #commit} with the messages of the tasks if they all succeeded, or - * {@link #abort} otherwise. A streaming write commits or aborts each micro-batch (epoch). + * rows, and calls {@link #commitDataWriter}, or {@link #abortDataWriter} if the task or the + * commit fails. Then the driver calls {@link #commit} with the messages of the tasks if they + * all succeeded, or {@link #abort} otherwise, and closes the writer. A streaming write commits + * or aborts each micro-batch (epoch). A micro-batch streaming query creates a new stream writer + * for each micro-batch, and Spark closes the previous stream writer of the query when the query + * creates a new one. Spark is not told when a streaming query stops, so it closes the last + * stream writer of the query once the writer is no longer referenced, or when the query is + * restarted. *

    * *

    Conventions

    *
      *
    • Handles. The {@code create} methods return an opaque, non-zero handle, typically a * pointer. Spark releases each handle exactly once with the matching {@code close} method, or - * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers. A reader or + * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers, though not + * always right after the last use of a stream writer (see Writes above). A reader or * writer handle must not depend on the data source handle it was created from, because Spark * closes the data source handle right after creating it.
      • *
    • State. Partitions, the state of readers and writers, and commit messages are byte @@ -212,7 +218,15 @@ public static native long createStreamWriter( * {@code =}, {@code <}, {@code IN}, {@code IS_NULL}, {@code STARTS_WITH}, {@code AND}, * {@code OR} and {@code NOT}.
      • *
    - * Spark does not pass predicates that it cannot express this way. + * Spark does not pass predicates that it cannot express this way, nor predicates on columns + * that contain strings with a non-binary collation, such as {@code UTF8_LCASE}: strings are + * compared by their UTF-8 bytes. + *

    + * A predicate that the reader evaluates completely must follow the semantics of Spark SQL. In + * particular, for floating-point columns, NaN is equal to NaN and greater than any other value, + * including positive infinity, and {@code -0.0} is equal to {@code 0.0}, unlike IEEE 754 + * comparisons. For example, {@code d > 1.0} is true for a NaN value of {@code d}. Comparisons + * with null are null, so a row only matches if the predicate is true. * * @return for each predicate, whether the reader evaluates it completely, so that Spark does * not have to evaluate it again. The reader can still use the other predicates to skip @@ -359,7 +373,9 @@ public static native long createDataWriter( public static native void write(long dataWriter, long arrayAddress, long schemaAddress); /** - * Commits the data written by the task and releases the data writer handle. + * Commits the data written by the task and releases the data writer handle. If it fails, the + * handle is not released: Spark then calls {@link #abortDataWriter} with it, to abort the data + * written by the task and release the handle. * * @return the message to pass to {@link #commit} on the driver */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataSourceRegistration.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataSourceRegistration.scala index 9d763edb079e7..d725922d95e83 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/DataSourceRegistration.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/DataSourceRegistration.scala @@ -65,7 +65,9 @@ private[sql] class DataSourceRegistration private[sql] (dataSourceManager: DataS if (dataSourceManager.dataSourceExists(name)) return try { - DataSource.lookupDataSource(name, SQLConf.get) + // A Python data source takes precedence over a native data source with the same name, so + // registering it does not depend on the native libraries that happen to be installed. + DataSource.lookupDataSource(name, SQLConf.get, includeNativeDataSources = false) throw QueryCompilationErrors.dataSourceAlreadyExists(name) } catch { case e: SparkClassNotFoundException if e.getCondition == "DATA_SOURCE_NOT_FOUND" => // OK diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala index 49c85d949374a..b23ab97836f8a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala @@ -667,6 +667,19 @@ object DataSource extends Logging { /** Given a provider name, look up the data source class definition. */ def lookupDataSource(provider: String, conf: SQLConf): Class[_] = { + lookupDataSource(provider, conf, includeNativeDataSources = true) + } + + /** + * Given a provider name, look up the data source class definition. + * + * @param includeNativeDataSources whether to look the name up among the native data sources, + * which come after the Java and Python data sources + */ + private[sql] def lookupDataSource( + provider: String, + conf: SQLConf, + includeNativeDataSources: Boolean): Class[_] = { val provider1 = backwardCompatibilityMap.getOrElse(provider, provider) match { case name if name.equalsIgnoreCase("orc") && conf.getConf(SQLConf.ORC_IMPLEMENTATION) == "native" => @@ -683,7 +696,9 @@ object DataSource extends Logging { val serviceLoader = ServiceLoader.load(classOf[DataSourceRegister], loader) lazy val isUserDefinedDataSource = SparkSession.getActiveSession.exists( _.sessionState.dataSourceManager.dataSourceExists(provider)) - lazy val isNativeDataSource = NativeDataSourceRegistry.exists(provider, conf) + // The provider created for the native data source takes the result of this lookup. + lazy val isNativeDataSource = includeNativeDataSources && + NativeDataSourceRegistry.lookupForProvider(provider, conf).isDefined try { serviceLoader.asScala.filter(_.shortName().equalsIgnoreCase(provider1)).toList match { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala index edd1b464121b8..1a8ac1a968e87 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala @@ -49,12 +49,14 @@ trait ColumnarDataSource { def writer(schema: StructType, overwrite: Boolean): ColumnarWriter /** - * Returns a writer for a streaming write. + * Returns a writer for a streaming write. A micro-batch streaming query creates a new one for + * each micro-batch. * * @param overwrite whether to replace the existing data in each micro-batch (output mode * "complete") instead of appending to it + * @param queryId the ID of the streaming query */ - def streamWriter(schema: StructType, overwrite: Boolean): ColumnarStreamWriter + def streamWriter(schema: StructType, overwrite: Boolean, queryId: String): ColumnarStreamWriter } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala index a59ae7d8e261c..17d253d0c9a94 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala @@ -56,7 +56,8 @@ class ColumnarTable( util.EnumSet.of(BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE) override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new ColumnarScanBuilder(tableSchema, () => createDataSource(withTableOptions(options))) + new ColumnarScanBuilder( + shortName, tableSchema, () => createDataSource(withTableOptions(options))) } override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = { @@ -73,7 +74,10 @@ class ColumnarTable( } } -class ColumnarScanBuilder(schema: StructType, createDataSource: () => ColumnarDataSource) +class ColumnarScanBuilder( + shortName: String, + schema: StructType, + createDataSource: () => ColumnarDataSource) extends ScanBuilder with SupportsPushDownV2Filters with SupportsPushDownLimit @@ -126,11 +130,13 @@ class ColumnarScanBuilder(schema: StructType, createDataSource: () => ColumnarDa } override def build(): Scan = { - new ColumnarScan(schema, readSchema, () => batchReader, () => dataSource, pushed, pushedLimit) + new ColumnarScan( + shortName, schema, readSchema, () => batchReader, () => dataSource, pushed, pushedLimit) } } class ColumnarScan( + shortName: String, fullSchema: StructType, prunedSchema: StructType, batchReader: () => ColumnarReader, @@ -144,8 +150,7 @@ class ColumnarScan( override def readSchema(): StructType = prunedSchema - // The class that implements the data source, such as NativeDataSource for a native one. - override def description(): String = dataSource().getClass.getSimpleName + override def description(): String = shortName override def columnarSupportMode(): Scan.ColumnarSupportMode = Scan.ColumnarSupportMode.SUPPORTED diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala index af5e530dfda29..9f167f97ea211 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala @@ -62,7 +62,7 @@ class ColumnarWrite( override def toStreaming: StreamingWrite = { new ColumnarStreamingWrite( - createDataSource().streamWriter(info.schema(), overwrite), info.schema()) + createDataSource().streamWriter(info.schema(), overwrite, info.queryId()), info.schema()) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index 633561f7a6c67..d3794d3279b54 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -16,7 +16,8 @@ */ package org.apache.spark.sql.execution.datasources.v2.ffi -import java.lang.ref.Cleaner +import java.lang.ref.{Cleaner, WeakReference} +import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.atomic.AtomicLong import scala.jdk.CollectionConverters._ @@ -44,7 +45,8 @@ import org.apache.spark.util.Utils * A [[ColumnarDataSource]] implemented by a native library, through * `org.apache.spark.sql.datasource.NativeBridge`. * - * @param location where the library is + * @param location where the library is on the driver. The readers and the writers distribute it + * to the executors of the active session when they are created. * @param name the name of the data source, as listed in the manifest of its package or in lower * case for an installed library */ @@ -65,26 +67,29 @@ class NativeDataSource( } override def reader(schema: StructType): ColumnarReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) val handle = create(library, "plan a scan of", "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { dataSource => NativeArrow.exportSchema(schema)(address => library.createReader(dataSource, address)) } new NativeDataSourceReader( - location, name, schema, new NativeHandle(handle, library.closeReader)) + executorLocation, name, schema, new NativeHandle(handle, library.closeReader)) } override def streamReader(schema: StructType): ColumnarStreamReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) val handle = create(library, "plan a scan of", "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { dataSource => NativeArrow.exportSchema(schema)(address => library.createStreamReader(dataSource, address)) } new NativeDataSourceStreamReader( - location, name, schema, new NativeHandle(handle, library.closeStreamReader)) + executorLocation, name, schema, new NativeHandle(handle, library.closeStreamReader)) } override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) val handle = new NativeHandle( create(library, "plan a write to", "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource => @@ -93,10 +98,14 @@ class NativeDataSource( } }, library.closeWriter) - new NativeDataSourceWriter(location, name, handle, serializeWriter(library, handle)) + new NativeDataSourceWriter(executorLocation, name, handle, serializeWriter(library, handle)) } - override def streamWriter(schema: StructType, overwrite: Boolean): ColumnarStreamWriter = { + override def streamWriter( + schema: StructType, + overwrite: Boolean, + queryId: String): ColumnarStreamWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) val handle = new NativeHandle( create(library, "plan a write to", "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { @@ -106,7 +115,10 @@ class NativeDataSource( } }, library.closeWriter) - new NativeDataSourceStreamWriter(location, name, handle, serializeWriter(library, handle)) + val writer = new NativeDataSourceStreamWriter( + executorLocation, name, handle, serializeWriter(library, handle)) + NativeStreamWriters.register(queryId, handle) + writer } /** Runs `f` with a new native data source, which is closed afterwards. */ @@ -161,7 +173,7 @@ class NativeDataSourceReader( override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { // Only the predicates that can be expressed in JSON are passed to the library. val encoded = predicates.zipWithIndex.flatMap { case (predicate, i) => - NativePredicates.toJson(predicate).map(_ -> i) + NativePredicates.toJson(predicate, schema).map(_ -> i) } val pushed = if (encoded.isEmpty) { Set.empty[Int] @@ -290,7 +302,10 @@ class NativeDataSourceWriter( } } -/** Writes to a native data source for a streaming write. */ +/** + * Writes to a native data source for a streaming write. Its handle is closed when the streaming + * query creates its next stream writer, see [[NativeStreamWriters]]. + */ class NativeDataSourceStreamWriter( location: NativeLibraryLocation, name: String, @@ -445,7 +460,8 @@ class NativeDataWriter( import NativeCalls._ private val library = NativeLibraries.get(location) - // Zero once the data writer was committed or aborted, which releases it. + // Zero once the data writer was committed or aborted, which releases it. It stays set if the + // commit fails, so that Spark aborts the data writer. private var handle: Long = call(name, "write to") { library.createDataWriter(writerState, partitionId, taskId, epochId) } @@ -480,10 +496,12 @@ class NativeDataWriter( } override def commit(): WriterCommitMessage = { - val dataWriter = handle + if (handle == 0) { + throw SparkException.internalError("The data writer was already committed or aborted.") + } + val message = call(name, "write to")(library.commitDataWriter(handle)) handle = 0 - NativeWriterCommitMessage( - bytesOrEmpty(call(name, "write to")(Some(library.commitDataWriter(dataWriter))))) + NativeWriterCommitMessage(bytesOrEmpty(Some(message))) } override def abort(): Unit = { @@ -537,6 +555,26 @@ object NativeHandle extends Logging { } } +/** + * The stream writers of the streaming queries, by query ID. The Data Source V2 streaming write + * has no hook for when the query stops, but a micro-batch streaming query creates a new stream + * writer for each micro-batch, after the previous micro-batch committed or aborted its write, and + * a restarted query creates a new one. So the previous stream writer of a query is closed when + * the query creates a new one, and the last one when it becomes unreachable. + */ +object NativeStreamWriters { + // Weak, so that the last stream writer of a query can still become unreachable. + private val writers = new ConcurrentHashMap[String, WeakReference[NativeHandle]]() + + /** Registers the new stream writer of the query, and closes its previous one. */ + def register(queryId: String, handle: NativeHandle): Unit = { + // Forgets the queries whose last stream writer was released. + writers.values().removeIf(_.get() == null) + Option(writers.put(queryId, new WeakReference(handle))).flatMap(r => Option(r.get())) + .foreach(_.close()) + } +} + /** Helpers to call native data source libraries. */ object NativeCalls { /** Calls a native library, and rethrows the errors it reports as NATIVE_DATA_SOURCE_ERROR. */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala index b8e393f816492..c5f89c1bacc29 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -83,45 +83,83 @@ case class NativeDataSourcePackage( def fileName: String = new File(path).getName } -object NativeDataSourcePackage { +/** + * A package file that is not a valid native data source package. + * + * @param error the error that explains why the package is invalid + * @param dataSources the data sources listed in the manifest of the package, if it can be read + */ +case class InvalidNativeDataSourcePackage(error: Throwable, dataSources: Seq[String]) + +object NativeDataSourcePackage extends Logging { val FILE_EXTENSION: String = Artifact.NATIVE_DATA_SOURCE_PACKAGE_EXTENSION val MANIFEST_NAME = "spark-native-datasource.json" private case class CacheKey(path: String, length: Long, lastModified: Long) - private val packages = new ConcurrentHashMap[CacheKey, NativeDataSourcePackage]() + private[ffi] type ReadResult = Either[InvalidNativeDataSourcePackage, NativeDataSourcePackage] + + private val packages = new ConcurrentHashMap[CacheKey, ReadResult]() private val mapper = new ObjectMapper() /** Reads the package in the given file. The result is cached until the file changes. */ - def read(file: File): NativeDataSourcePackage = { - val canonical = file.getCanonicalFile + def read(file: File): NativeDataSourcePackage = tryRead(file).fold(p => throw p.error, identity) + + /** + * Reads the package in the given file, or returns why it is invalid. The result is cached until + * the file changes, and an invalid package is logged once. + */ + private[ffi] def tryRead(file: File): ReadResult = { + val canonical = try { + file.getCanonicalFile + } catch { + case e: IOException => + return Left(InvalidNativeDataSourcePackage( + invalidManifestError(file, s"The package cannot be read: $e", e), Nil)) + } val key = CacheKey(canonical.getPath, canonical.length(), canonical.lastModified()) - packages.computeIfAbsent(key, _ => parse(canonical)) + packages.computeIfAbsent(key, _ => { + val result = parse(canonical) + result.swap.foreach { invalid => + logWarning(log"Invalid native data source package ${MDC(PATH, canonical)}.", + invalid.error) + } + result + }) + } + + private def invalidManifestError(file: File, reason: String, cause: Throwable = null) = { + QueryExecutionErrors.invalidNativeDataSourcePackageError( + file.getPath, "INVALID_MANIFEST", Map("reason" -> reason), cause) } - private def parse(file: File): NativeDataSourcePackage = { - def invalid(reason: String, cause: Throwable = null): Throwable = { - QueryExecutionErrors.invalidNativeDataSourcePackageError( - file.getPath, "INVALID_MANIFEST", Map("reason" -> reason), cause) + private def parse(file: File): ReadResult = { + def invalid(reason: String, dataSources: Seq[String] = Nil, cause: Throwable = null) = { + Left(InvalidNativeDataSourcePackage(invalidManifestError(file, reason, cause), dataSources)) } val manifest = try { Utils.tryWithResource(new ZipFile(file)) { zip => - val entry = Option(zip.getEntry(MANIFEST_NAME)).getOrElse { - throw invalid(s"The package does not contain the manifest $MANIFEST_NAME.") - } - val json = Utils.tryWithResource(zip.getInputStream(entry)) { in => - new String(in.readAllBytes(), StandardCharsets.UTF_8) + Option(zip.getEntry(MANIFEST_NAME)) match { + case None => invalid(s"The package does not contain the manifest $MANIFEST_NAME.") + case Some(entry) => + val json = Utils.tryWithResource(zip.getInputStream(entry)) { in => + new String(in.readAllBytes(), StandardCharsets.UTF_8) + } + parseManifest(json) match { + case Left(reason) => invalid(reason) + case Right(manifest) => + manifest.libraries.values.find(zip.getEntry(_) == null) match { + case Some(library) => invalid( + s"The package does not contain the library $library.", manifest.dataSources) + case None => Right(manifest) + } + } } - val manifest = parseManifest(json).fold(reason => throw invalid(reason), identity) - manifest.libraries.values.find(zip.getEntry(_) == null).foreach { library => - throw invalid(s"The package does not contain the library $library.") - } - manifest } } catch { - case e: IOException => throw invalid(s"The package is not a valid zip file: $e", e) + case e: IOException => invalid(s"The package is not a valid zip file: $e", cause = e) } - NativeDataSourcePackage(file.getPath, sha256(file), manifest) + manifest.map(NativeDataSourcePackage(file.getPath, sha256(file), _)) } private[ffi] def parseManifest(json: String): Either[String, NativeDataSourceManifest] = { @@ -294,7 +332,9 @@ object NativeLibraries extends Logging { /** Finds a copy of the package on this node: the one on the driver, or a distributed one. */ private def locate(pkg: NativeDataSourcePackage): File = { (new File(pkg.path) +: NativeDataSourceRegistry.distributedCopies(pkg)) - .find(file => file.isFile && NativeDataSourcePackage.read(file).sha256 == pkg.sha256) + .find { file => + file.isFile && NativeDataSourcePackage.tryRead(file).exists(_.sha256 == pkg.sha256) + } .getOrElse { throw QueryExecutionErrors.nativeDataSourcePackageNotFoundError(pkg.fileName, pkg.sha256) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala index 2aef77d05b572..6b66df0e8e889 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -43,8 +43,10 @@ object NativeDataSourceRegistry extends Logging { /** The prefix of the names of the installed native data source libraries. */ val LIBRARY_NAME_PREFIX = "spark_datasource_" - /** Returns whether the native data source with the given name exists. */ - def exists(name: String, conf: SQLConf): Boolean = lookup(name, conf).isDefined + // The result of the last lookup by `DataSource.lookupDataSource` on this thread that found a + // native data source, by the name it looked up. The provider that the caller creates next for + // that name takes it, so that it does not look the name up again: see [[takeLookup]]. + private val lastLookup = new ThreadLocal[(String, (String, NativeLibraryLocation))] /** * Finds the package or the installed library that provides the native data source with the @@ -59,10 +61,32 @@ object NativeDataSourceRegistry extends Logging { .orElse(lookupInstalledLibrary(name)) } + /** + * Like [[lookup]], and keeps the result for the provider that the caller creates next for the + * given name, which takes it with [[takeLookup]]. + */ + def lookupForProvider(name: String, conf: SQLConf): Option[(String, NativeLibraryLocation)] = { + val result = lookup(name, conf) + result.foreach(found => lastLookup.set(name -> found)) + result + } + + /** Takes the result of [[lookupForProvider]] for the given name on this thread, if any. */ + def takeLookup(name: String): Option[(String, NativeLibraryLocation)] = { + Option(lastLookup.get()).filter(_._1.equalsIgnoreCase(name)).map { case (_, found) => + lastLookup.remove() + found + } + } + private def lookupPackage( name: String, files: Seq[File]): Option[(String, NativeDataSourcePackage)] = { - val found = files.map(_.getCanonicalFile).distinct.map(NativeDataSourcePackage.read) + val (invalid, valid) = files.map(NativeDataSourcePackage.tryRead).partitionMap(identity) + // An invalid package does not prevent finding the data sources of the other packages, unless + // its manifest lists the data source. An invalid package is logged when it is read. + invalid.find(_.dataSources.exists(_.equalsIgnoreCase(name))).foreach(p => throw p.error) + val found = valid .flatMap(pkg => pkg.manifest.dataSources.find(_.equalsIgnoreCase(name)).map(_ -> pkg)) // The same package can be found more than once, for example when it is in a configured // directory and was also added as an artifact. @@ -130,8 +154,10 @@ object NativeDataSourceRegistry extends Logging { } /** - * Makes the library available to the executors of the active session. Returns where the - * executors find it. An installed library is already installed on every node. + * Makes the library available to the executors of the active session, which runs the jobs of + * the scan or the write: for a streaming query, it is the session of the query, a clone of the + * session that started it. Returns where the executors find it. An installed library is already + * installed on every node. */ def distribute(location: NativeLibraryLocation): NativeLibraryLocation = location match { case pkg: NativeDataSourcePackage => @@ -149,24 +175,39 @@ object NativeDataSourceRegistry extends Logging { pkg: NativeDataSourcePackage): NativeDataSourcePackage = { val artifactManager = session.artifactManager val (added, artifactUUID) = artifactManager.getNativeDataSourcePackages - if (!added.exists(_.getName == pkg.fileName)) { - artifactManager.addLocalArtifacts(Artifact.newFileArtifact( - Paths.get(pkg.fileName), new Artifact.LocalFile(Paths.get(pkg.path))) :: Nil) + // The executors find the package by its file name, so the session cannot have another + // package with the same file name. + added.find(_.getName == pkg.fileName) match { + case Some(file) => + val sha256 = NativeDataSourcePackage.tryRead(file).map(_.sha256).getOrElse("") + if (sha256 != pkg.sha256) { + throw QueryExecutionErrors.nativeDataSourcePackageFileNameConflictError( + pkg.path, pkg.fileName) + } + case None => + artifactManager.addLocalArtifacts(Artifact.newFileArtifact( + Paths.get(pkg.fileName), new Artifact.LocalFile(Paths.get(pkg.path))) :: Nil) } pkg.copy(artifactUUID = artifactUUID) } /** - * The local copies, on an executor, of a package distributed with [[distribute]]. The files of - * an isolated session are stored in a directory named after the UUID of the session. + * The local copies, on an executor, of a package distributed with [[distribute]], in the order + * to check them. The files of an isolated session are stored in a directory named after the + * UUID of the session; the copies of other sessions are checked last, because the executors + * check the checksum of the package anyway. */ def distributedCopies(pkg: NativeDataSourcePackage): Seq[File] = { if (SparkEnv.get == null) { Nil } else { val root = new File(SparkFiles.getRootDirectory()) - pkg.artifactUUID.map(uuid => new File(new File(root, uuid), pkg.fileName)).toSeq :+ - new File(root, pkg.fileName) + val sessionCopies = Option(root.listFiles()).toSeq.flatten + .filter(_.isDirectory) + .sortBy(_.getName) + .map(new File(_, pkg.fileName)) + (pkg.artifactUUID.map(uuid => new File(new File(root, uuid), pkg.fileName)).toSeq ++ + Seq(new File(root, pkg.fileName)) ++ sessionCopies).distinct } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala index e9a3fcf7d17e9..93876788182f2 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala @@ -33,20 +33,22 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap */ class NativeDataSourceV2 extends NamedTableProvider { private var shortName: String = _ + // The result of the lookup that `DataSource.lookupDataSource` made to find this provider class. + private var lookup: Option[(String, NativeLibraryLocation)] = None override def setShortName(name: String): Unit = { assert(shortName == null) shortName = name + lookup = NativeDataSourceRegistry.takeLookup(name) } - // The name of the data source, and where its library is. Resolved on first use, which is while - // the query is analyzed, with the configuration of its session. + // The name of the data source, and where its library is on the driver. Resolved on first use, + // which is while the query is analyzed, with the configuration of its session. private lazy val resolved: (String, NativeLibraryLocation) = { assert(shortName != null) - val (name, location) = NativeDataSourceRegistry.lookup(shortName, SQLConf.get).getOrElse { + lookup.orElse(NativeDataSourceRegistry.lookup(shortName, SQLConf.get)).getOrElse { throw QueryCompilationErrors.dataSourceDoesNotExist(shortName) } - (name, NativeDataSourceRegistry.distribute(location)) } override def inferSchema(options: CaseInsensitiveStringMap): StructType = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala index 7c6c6d2bbffd1..b7a6a43518ca2 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala @@ -32,35 +32,67 @@ import org.apache.spark.sql.types._ object NativePredicates { private val mapper = new ObjectMapper() - /** Returns the JSON of the predicate, or None if it cannot be expressed in JSON. */ - def toJson(predicate: Predicate): Option[String] = { - encode(predicate).map(mapper.writeValueAsString) + /** + * Returns the JSON of the predicate on the columns of the given schema, or None if it cannot + * be expressed in JSON. + */ + def toJson(predicate: Predicate, schema: StructType): Option[String] = { + encode(predicate, schema).map(mapper.writeValueAsString) } - private def encode(expression: Expression): Option[ObjectNode] = expression match { - case reference: NamedReference => + private def encode(expression: Expression, schema: StructType): Option[ObjectNode] = { + expression match { + case reference: NamedReference => + encodeColumn(reference, schema) + + case literal: Literal[_] => + encodeLiteral(literal.value(), literal.dataType()) + + case function: GeneralScalarExpression => + val children = function.children().map(encode(_, schema)) + if (children.forall(_.isDefined)) { + val node = mapper.createObjectNode() + .put("type", "function") + .put("name", function.name()) + val array = node.putArray("children") + children.foreach(child => array.add(child.get)) + Some(node) + } else { + None + } + + case _ => None + } + } + + private def encodeColumn(reference: NamedReference, schema: StructType): Option[ObjectNode] = { + // A source compares strings by their bytes, so a column that contains strings with a + // non-binary collation, such as UTF8_LCASE, is not pushed down. Neither is a column that is + // not in the schema. + val hasNonBinaryCollation = columnType(reference.fieldNames().toSeq, schema).forall { + _.existsRecursively { + case st: StringType => !st.isUTF8BinaryCollation + case _ => false + } + } + if (hasNonBinaryCollation) { + None + } else { val node = mapper.createObjectNode().put("type", "column") val name = node.putArray("name") reference.fieldNames().foreach(name.add) Some(node) + } + } - case literal: Literal[_] => - encodeLiteral(literal.value(), literal.dataType()) - - case function: GeneralScalarExpression => - val children = function.children().map(encode) - if (children.forall(_.isDefined)) { - val node = mapper.createObjectNode() - .put("type", "function") - .put("name", function.name()) - val array = node.putArray("children") - children.foreach(child => array.add(child.get)) - Some(node) - } else { - None - } - - case _ => None + /** The type of the column or the nested field with the given name parts. */ + private def columnType(names: Seq[String], dataType: DataType): Option[DataType] = { + (names, dataType) match { + case (Seq(), _) => Some(dataType) + case (name +: rest, struct: StructType) => + struct.fields.find(_.name == name).flatMap(field => columnType(rest, field.dataType)) + case _ => None + } } private def encodeLiteral(value: Any, dataType: DataType): Option[ObjectNode] = { diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index eaa52ed534c8d..eca83f7365b1f 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -1099,9 +1099,13 @@ SPARK_JNI(void, write)(JNIEnv* env, jclass, jlong handle, jlong array_address, SPARK_JNI(jbyteArray, commitDataWriter)(JNIEnv* env, jclass, jlong handle) { return guarded(env, nullptr, [&] { - std::unique_ptr writer(fromHandle(handle)); + auto* writer = fromHandle(handle); + // If the commit fails, the data writer stays valid: Spark aborts it with abortDataWriter. + maybeFail(writer->state, "commitDataWriter"); + jbyteArray message = newBytes(env, writer->temp_path); std::fclose(writer->file); - return newBytes(env, writer->temp_path); + delete writer; + return message; }); } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala index af8c09c22bc5f..894bb2b9c8520 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala @@ -176,7 +176,10 @@ class TestRangeDataSource(options: CaseInsensitiveStringMap) extends ColumnarDat new TestSinkWriter(table, overwrite, options.getBoolean("fail", false)) } - override def streamWriter(schema: StructType, overwrite: Boolean): ColumnarStreamWriter = { + override def streamWriter( + schema: StructType, + overwrite: Boolean, + queryId: String): ColumnarStreamWriter = { new TestSinkStreamWriter(table) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 5534361f3fd65..865e1e7d54e4e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -21,21 +21,27 @@ import java.nio.charset.StandardCharsets.UTF_8 import java.nio.file.Files import java.sql.{Date, Timestamp} import java.time.{Instant, LocalDate} +import java.util.UUID import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ -import org.apache.spark.{SparkClassNotFoundException, SparkRuntimeException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.{SparkClassNotFoundException, SparkFiles, SparkRuntimeException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.api.python.{PythonBroadcast, SimplePythonFunction} +import org.apache.spark.broadcast.Broadcast import org.apache.spark.sql.{AnalysisException, QueryTest, Row} import org.apache.spark.sql.classic.{DataFrame, SparkSession} import org.apache.spark.sql.connector.expressions.{Expression, FieldReference, GeneralScalarExpression, LiteralValue} import org.apache.spark.sql.connector.expressions.filter.{And, Predicate} +import org.apache.spark.sql.execution.datasources.DataSource import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation import org.apache.spark.sql.execution.datasources.v2.columnar.ColumnarScan +import org.apache.spark.sql.execution.datasources.v2.python.{PythonDataSourceV2, UserDefinedPythonDataSource} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.streaming.StreamingQuery import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ +import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.Utils @@ -163,10 +169,13 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { new Predicate(name, children.toArray) } val id = FieldReference(Seq("id")) - assert(NativePredicates.toJson(predicate(">", id, LiteralValue(5L, LongType))).contains( + val schema = StructType.fromDDL("id BIGINT, a STRUCT, s STRING, " + + "c STRING COLLATE UTF8_LCASE, d STRUCT>") + def toJson(predicate: Predicate): Option[String] = NativePredicates.toJson(predicate, schema) + assert(toJson(predicate(">", id, LiteralValue(5L, LongType))).contains( """{"type":"function","name":">","children":[{"type":"column","name":["id"]},""" + """{"type":"literal","dataType":"bigint","value":5}]}""")) - assert(NativePredicates.toJson(new And( + assert(toJson(new And( predicate("IS_NULL", FieldReference(Seq("a", "b"))), predicate("=", id, new GeneralScalarExpression("+", Array(id, LiteralValue(1, IntegerType)))) )).contains( @@ -177,7 +186,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { """{"type":"literal","dataType":"int","value":1}]}]}]}""")) def literal(value: Any, dataType: DataType): Option[String] = { - NativePredicates.toJson(predicate("=", id, LiteralValue(value, dataType))) + toJson(predicate("=", id, LiteralValue(value, dataType))) .map(json => json.substring(json.indexOf("\"dataType\""), json.length - 3)) } assert(literal(null, LongType).contains(""""dataType":"bigint","value":null""")) @@ -201,6 +210,36 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(literal(UTF8String.fromString("a"), StringType("UTF8_LCASE")).isEmpty) assert(literal(null, ArrayType(IntegerType)).nonEmpty) assert(literal(Array(1), ArrayType(IntegerType)).isEmpty) + + // A source compares strings by their bytes, so the columns with strings of a non-binary + // collation are not pushed down, and neither are the columns that are not in the schema. + def compare(left: String, right: String): Option[String] = { + toJson(predicate("=", FieldReference(left.split('.').toSeq), + FieldReference(right.split('.').toSeq))) + } + assert(compare("s", "s").nonEmpty) + assert(compare("a.b", "id").nonEmpty) + assert(compare("c", "c").isEmpty) + assert(compare("s", "c").isEmpty) + assert(compare("d", "d").isEmpty) + assert(compare("d.e", "d.e").isEmpty) + assert(compare("missing", "id").isEmpty) + assert(compare("a.missing", "id").isEmpty) + assert(compare("id.b", "id").isEmpty) + } + + test("close the previous stream writer of a streaming query") { + val released = ArrayBuffer.empty[Long] + def handle(value: Long): NativeHandle = new NativeHandle(value, released += _) + val (first, other, second) = (handle(1), handle(2), handle(3)) + NativeStreamWriters.register("query", first) + NativeStreamWriters.register("other query", other) + assert(released.isEmpty) + NativeStreamWriters.register("query", second) + assert(released == Seq(1L)) + second.close() + other.close() + assert(released == Seq(1L, 3L, 2L)) } nativeTest("disable native data sources") { @@ -253,6 +292,9 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { checkRows(df, (6L to 12L).map(expectedRow)) val pushed = scan(df).getMetaData()("PushedPredicates") assert(pushed.contains("id > 5") && pushed.contains("id <= 12"), pushed) + // The scan is described by the name of the data source. + assert(scan(df).description() == "native_range") + assert(!df.queryExecution.executedPlan.toString.contains("NativeDataSource")) // A predicate that the library does not accept is evaluated by Spark. val partly = range(newSession(defaultPackage), "end" -> "20") @@ -344,6 +386,26 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) } + nativeTest("abort a data writer whose commit failed") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val e = intercept[Exception] { + session.range(10).repartition(2).write.format("native_sink").option("path", dir.getPath) + .option("fail", "commitDataWriter").mode("append").save() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "write to", + "name" -> "native_sink", + "msg" -> "injected failure in commitDataWriter")) + // abortDataWriter removed the temporary files of the tasks. + assert(new File(dir, "_ABORTED").isFile) + assert(new File(dir, "_temporary").listFiles().isEmpty) + assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) + } + nativeTest("streaming read and write") { val session = newSession(defaultPackage) val dir = Utils.createTempDir() @@ -473,6 +535,72 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { val (otherPackages, otherArtifactUUID) = other.artifactManager.getNativeDataSourcePackages assert(otherPackages.map(_.getName) == Seq(defaultPackage.getName)) assert(distributed.artifactUUID == otherArtifactUUID) + // Adding the same package again does nothing. + NativeDataSourceRegistry.addToArtifacts(other, NativeDataSourcePackage.read(defaultPackage)) + assert(other.artifactManager.getNativeDataSourcePackages._1.size == 1) + + // A different package with the same file name cannot be added: the executors would not find + // the package by its file name. + val sameFileName = createPackage(negatingLibrary, dataSources("b_"), + dir = Utils.createTempDir(), fileName = defaultPackage.getName) + checkError( + exception = intercept[SparkRuntimeException] { + NativeDataSourceRegistry.addToArtifacts(other, NativeDataSourcePackage.read(sameFileName)) + }, + condition = "NATIVE_DATA_SOURCE_PACKAGE_FILE_NAME_CONFLICT", + parameters = Map( + "path" -> sameFileName.getCanonicalPath, "fileName" -> defaultPackage.getName)) + + // The executors also check the copies in the files of the other sessions, such as the session + // of a streaming query, which is a clone of the session that started it. + val sessionDir = new File(SparkFiles.getRootDirectory(), s"test-${UUID.randomUUID()}") + try { + val copy = new File(sessionDir, defaultPackage.getName) + sessionDir.mkdirs() + Files.copy(defaultPackage.toPath, copy.toPath) + val copies = NativeDataSourceRegistry.distributedCopies(pkg.copy(artifactUUID = Some("x"))) + assert(copies.take(2) == Seq( + new File(new File(SparkFiles.getRootDirectory(), "x"), defaultPackage.getName), + new File(SparkFiles.getRootDirectory(), defaultPackage.getName))) + assert(copies.contains(copy)) + } finally { + Utils.deleteRecursively(sessionDir) + } + } + + nativeTest("look up a native data source once") { + val session = newSession(defaultPackage) + val provider = DataSource.lookupDataSourceV2("native_range", session.sessionState.conf).get + assert(provider.isInstanceOf[NativeDataSourceV2]) + assert(NativeDataSourceRegistry.takeLookup("native_range").isEmpty) + // The provider uses the result of the lookup that found its class, even if the package is + // no longer configured. + session.conf.unset(SQLConf.NATIVE_DATA_SOURCE_PATHS.key) + assert(provider.inferSchema(CaseInsensitiveStringMap.empty()).fieldNames.head == "id") + } + + nativeTest("Python data sources take precedence") { + val dir = Utils.createTempDir() + installLibrary(installedLibrary, "installed_native_range", dir) + withLibraryPath(dir) { + val session = newSession() + val conf = session.sessionState.conf + assert(DataSource.lookupDataSource("installed_native_range", conf) == + classOf[NativeDataSourceV2]) + // A Python data source can be registered with the name of an installed native data + // source, and takes precedence over it. + val pythonDataSource = UserDefinedPythonDataSource(SimplePythonFunction( + command = Seq.empty, + envVars = new java.util.HashMap[String, String](), + pythonIncludes = java.util.Collections.emptyList[String](), + pythonExec = "python3", + pythonVer = "3", + broadcastVars = java.util.Collections.emptyList[Broadcast[PythonBroadcast]](), + accumulator = null)) + session.dataSource.registerPython("installed_native_range", pythonDataSource) + assert(DataSource.lookupDataSource("installed_native_range", conf) == + classOf[PythonDataSourceV2]) + } } nativeTest("find installed libraries automatically") { @@ -563,11 +691,6 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { parameters = params) } - val noManifest = createZip(dir, "no_manifest.sparkpkg", "x" -> Array.emptyByteArray) - check(noManifest, "INVALID_MANIFEST", Map( - "path" -> noManifest.getCanonicalPath, - "reason" -> "The package does not contain the manifest spark-native-datasource.json.")) - val missingLibrary = createZip(dir, "missing_library.sparkpkg", NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( Seq("native_range"), Map(NativePlatform.current -> "lib.so")).getBytes(UTF_8)) @@ -609,6 +732,35 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { "path" -> mismatchedAbi.getCanonicalPath, "version" -> "2", "supported" -> "1")) } + nativeTest("ignore invalid packages that do not provide the data source") { + val dir = Utils.createTempDir() + createPackage(library, dataSources(""), dir = dir, fileName = "a.sparkpkg") + createZip(dir, "no_manifest.sparkpkg", "x" -> Array.emptyByteArray) + Files.writeString(new File(dir, "not_a_zip.sparkpkg").toPath, "not a zip") + val otherSource = createZip(dir, "other_source.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + Seq("other_source"), Map(NativePlatform.current -> "lib.so")).getBytes(UTF_8)) + val session = newSession(dir) + checkRows(range(session, "end" -> "2"), (0L until 2L).map(expectedRow)) + checkError( + exception = intercept[SparkClassNotFoundException] { + session.read.format("com.typo.Missing").load() + }, + condition = "DATA_SOURCE_NOT_FOUND", + parameters = Map("provider" -> "com.typo.Missing")) + val e = intercept[AnalysisException](session.sql("SELECT * FROM missing_db.missing_table")) + assert(e.getCondition == "TABLE_OR_VIEW_NOT_FOUND") + // The invalid package whose manifest lists the data source reports why it is invalid. + checkError( + exception = intercept[SparkRuntimeException] { + session.read.format("other_source").load() + }, + condition = "INVALID_NATIVE_DATA_SOURCE_PACKAGE.INVALID_MANIFEST", + parameters = Map( + "path" -> otherSource.getCanonicalPath, + "reason" -> "The package does not contain the library lib.so.")) + } + nativeTest("invalid installed libraries") { val dir = Utils.createTempDir() val notALibrary = new File(dir, NativeDataSourceRegistry.libraryFileName("not_a_library")) From 38141475b64603f5755d3db705495d4cc1231cc5 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Sat, 10 Oct 2026 17:13:34 +0900 Subject: [PATCH 09/27] Address more review comments on native data source packages and stream writers - Prefer a usable package or an installed library over an invalid package or a package of an unsupported interface version that lists the same data source, and only report those when nothing else provides it. - Clear the lookup kept for the provider at the start of every DataSource.lookupDataSource, so a provider never takes an earlier lookup. - Bound the cache of read packages, and forget the packages of a session when its artifacts are cleaned up. - Record writer closes in the test library and check end to end that each micro-batch closes the stream writer of the previous one, and that each micro-batch creates its stream writer with the ID of the query. Co-authored-by: Isaac --- docs/sql-native-data-sources.md | 3 +- .../spark/sql/artifact/ArtifactManager.scala | 4 ++ .../execution/datasources/DataSource.scala | 2 + .../v2/ffi/NativeDataSourcePackage.scala | 35 +++++++-- .../v2/ffi/NativeDataSourceRegistry.scala | 62 +++++++++------- .../test_native_datasource.cc | 21 +++++- .../v2/columnar/ColumnarDataSourceSuite.scala | 28 ++++++++ .../v2/ffi/NativeDataSourceSuite.scala | 71 ++++++++++++++++++- 8 files changed, 191 insertions(+), 35 deletions(-) diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 08b1fb715b732..717a0cba854cf 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -173,7 +173,8 @@ A session finds the packages in two places, and the executors get them automatic The packages of a session take precedence over the installed libraries. Give packages distinct file names, for example with their version: sessions that are not isolated share the files they add, and a session cannot use two different packages with the same file name. Spark ignores, and -logs, a package that it cannot read, unless its manifest lists the data source that a query uses. +logs, a package that it cannot read or whose version it does not support, unless its manifest +lists the data source that a query uses and no other package or installed library provides it. ## Example in Rust diff --git a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala index 281b2a079cc3d..80a5b2e539bc6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala @@ -36,6 +36,7 @@ import org.apache.spark.internal.{Logging, LogKeys} import org.apache.spark.internal.config.{CONNECT_SCALA_UDF_STUB_PREFIXES, EXECUTOR_USER_CLASS_PATH_FIRST, JAR_IVY_CONNECT_TIMEOUT, JAR_IVY_READ_TIMEOUT, JAR_IVY_REPO_PATH, JAR_IVY_SETTING_PATH} import org.apache.spark.sql.Artifact import org.apache.spark.sql.classic.SparkSession +import org.apache.spark.sql.execution.datasources.v2.ffi.NativeDataSourcePackage import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.util.ArtifactUtils import org.apache.spark.storage.{BlockManager, CacheId, StorageLevel} @@ -631,6 +632,9 @@ object ArtifactManager extends Logging { } } + // Forget the native data source packages read from the artifacts folder + Utils.tryLogNonFatalError(NativeDataSourcePackage.forgetPackagesIn(artifactPath.toFile)) + // Clean up artifacts folder try { Utils.deleteRecursively(artifactPath.toFile) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala index b23ab97836f8a..77087a65a99d2 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala @@ -680,6 +680,8 @@ object DataSource extends Logging { provider: String, conf: SQLConf, includeNativeDataSources: Boolean): Class[_] = { + // The provider created for a native data source only takes the result of this lookup. + NativeDataSourceRegistry.clearLookup() val provider1 = backwardCompatibilityMap.getOrElse(provider, provider) match { case name if name.equalsIgnoreCase("orc") && conf.getConf(SQLConf.ORC_IMPLEMENTATION) == "native" => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala index c5f89c1bacc29..c9bb312bea6b9 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -28,6 +28,7 @@ import scala.jdk.CollectionConverters._ import scala.util.Try import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} +import com.google.common.cache.{Cache, CacheBuilder} import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.PATH @@ -99,7 +100,12 @@ object NativeDataSourcePackage extends Logging { private[ffi] type ReadResult = Either[InvalidNativeDataSourcePackage, NativeDataSourcePackage] - private val packages = new ConcurrentHashMap[CacheKey, ReadResult]() + // Bounded, because each session stores its artifacts in a directory of its own, so a + // long-running server where many sessions add packages reads many different files. + private val MAX_CACHED_PACKAGES = 1000 + + private val packages: Cache[CacheKey, ReadResult] = + CacheBuilder.newBuilder().maximumSize(MAX_CACHED_PACKAGES).build[CacheKey, ReadResult]() private val mapper = new ObjectMapper() /** Reads the package in the given file. The result is cached until the file changes. */ @@ -107,7 +113,7 @@ object NativeDataSourcePackage extends Logging { /** * Reads the package in the given file, or returns why it is invalid. The result is cached until - * the file changes, and an invalid package is logged once. + * the file changes, and an invalid package is logged when it is read. */ private[ffi] def tryRead(file: File): ReadResult = { val canonical = try { @@ -118,14 +124,29 @@ object NativeDataSourcePackage extends Logging { invalidManifestError(file, s"The package cannot be read: $e", e), Nil)) } val key = CacheKey(canonical.getPath, canonical.length(), canonical.lastModified()) - packages.computeIfAbsent(key, _ => { + Option(packages.getIfPresent(key)).getOrElse { val result = parse(canonical) result.swap.foreach { invalid => logWarning(log"Invalid native data source package ${MDC(PATH, canonical)}.", invalid.error) } + packages.put(key, result) result - }) + } + } + + /** + * Forgets the packages read from the files under the given directory, such as the artifacts of + * a session that is cleaned up. + */ + def forgetPackagesIn(dir: File): Unit = { + val prefix = dir.getCanonicalPath + File.separator + packages.asMap().keySet().removeIf(_.path.startsWith(prefix)) + } + + /** Returns whether a package read from the file at the given canonical path is cached. */ + private[ffi] def isCached(canonicalPath: String): Boolean = { + packages.asMap().keySet().asScala.exists(_.path == canonicalPath) } private def invalidManifestError(file: File, reason: String, cause: Throwable = null) = { @@ -137,8 +158,8 @@ object NativeDataSourcePackage extends Logging { def invalid(reason: String, dataSources: Seq[String] = Nil, cause: Throwable = null) = { Left(InvalidNativeDataSourcePackage(invalidManifestError(file, reason, cause), dataSources)) } - val manifest = try { - Utils.tryWithResource(new ZipFile(file)) { zip => + try { + val manifest = Utils.tryWithResource(new ZipFile(file)) { zip => Option(zip.getEntry(MANIFEST_NAME)) match { case None => invalid(s"The package does not contain the manifest $MANIFEST_NAME.") case Some(entry) => @@ -156,10 +177,10 @@ object NativeDataSourcePackage extends Logging { } } } + manifest.map(NativeDataSourcePackage(file.getPath, sha256(file), _)) } catch { case e: IOException => invalid(s"The package is not a valid zip file: $e", cause = e) } - manifest.map(NativeDataSourcePackage(file.getPath, sha256(file), _)) } private[ffi] def parseManifest(json: String): Either[String, NativeDataSourceManifest] = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala index 6b66df0e8e889..ba07c41369662 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -43,9 +43,10 @@ object NativeDataSourceRegistry extends Logging { /** The prefix of the names of the installed native data source libraries. */ val LIBRARY_NAME_PREFIX = "spark_datasource_" - // The result of the last lookup by `DataSource.lookupDataSource` on this thread that found a + // The result of the last lookup by `DataSource.lookupDataSource` on this thread, if it found a // native data source, by the name it looked up. The provider that the caller creates next for - // that name takes it, so that it does not look the name up again: see [[takeLookup]]. + // that name takes it, so that it does not look the name up again: see [[takeLookup]]. Each + // lookup clears it first, so a provider never takes the result of an earlier lookup. private val lastLookup = new ThreadLocal[(String, (String, NativeLibraryLocation))] /** @@ -57,8 +58,29 @@ object NativeDataSourceRegistry extends Logging { if (!conf.getConf(StaticSQLConf.NATIVE_DATA_SOURCE_ENABLED)) { return None } - lookupPackage(name, configuredPackageFiles(conf) ++ artifactPackageFiles()) + val (invalid, valid) = (configuredPackageFiles(conf) ++ artifactPackageFiles()) + .map(NativeDataSourcePackage.tryRead).partitionMap(identity) + // A package of a version of the interface that Spark does not support is unusable too. + val (unsupported, usable) = valid.partitionMap { pkg => + if (pkg.manifest.abiVersion == NativeBridge.ABI_VERSION) Right(pkg) else Left(pkg) + } + lookupPackage(name, usable) .orElse(lookupInstalledLibrary(name)) + .orElse { + // An unusable package is only reported when no other package or installed library + // provides the data source: a stale copy does not hide a valid one. An invalid package + // whose manifest cannot be read is only logged, when it is read. + unsupported.find(_.manifest.dataSources.exists(_.equalsIgnoreCase(name))).foreach { pkg => + throw QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, + "UNSUPPORTED_ABI_VERSION", + Map( + "version" -> pkg.manifest.abiVersion.toString, + "supported" -> NativeBridge.ABI_VERSION.toString)) + } + invalid.find(_.dataSources.exists(_.equalsIgnoreCase(name))).foreach(p => throw p.error) + None + } } /** @@ -71,38 +93,30 @@ object NativeDataSourceRegistry extends Logging { result } - /** Takes the result of [[lookupForProvider]] for the given name on this thread, if any. */ + /** + * Takes the result of the last [[lookupForProvider]] on this thread, if it was for the given + * name, and forgets it. + */ def takeLookup(name: String): Option[(String, NativeLibraryLocation)] = { - Option(lastLookup.get()).filter(_._1.equalsIgnoreCase(name)).map { case (_, found) => - lastLookup.remove() - found - } + val last = Option(lastLookup.get()) + lastLookup.remove() + last.collect { case (lookedUp, found) if lookedUp.equalsIgnoreCase(name) => found } } + /** Forgets the result of the last [[lookupForProvider]] on this thread. */ + def clearLookup(): Unit = lastLookup.remove() + private def lookupPackage( name: String, - files: Seq[File]): Option[(String, NativeDataSourcePackage)] = { - val (invalid, valid) = files.map(NativeDataSourcePackage.tryRead).partitionMap(identity) - // An invalid package does not prevent finding the data sources of the other packages, unless - // its manifest lists the data source. An invalid package is logged when it is read. - invalid.find(_.dataSources.exists(_.equalsIgnoreCase(name))).foreach(p => throw p.error) - val found = valid + packages: Seq[NativeDataSourcePackage]): Option[(String, NativeDataSourcePackage)] = { + val found = packages .flatMap(pkg => pkg.manifest.dataSources.find(_.equalsIgnoreCase(name)).map(_ -> pkg)) // The same package can be found more than once, for example when it is in a configured // directory and was also added as an artifact. .groupBy(_._2.sha256).values.map(_.head).toSeq found match { case Seq() => None - case Seq(result @ (_, pkg)) => - if (pkg.manifest.abiVersion != NativeBridge.ABI_VERSION) { - throw QueryExecutionErrors.invalidNativeDataSourcePackageError( - pkg.path, - "UNSUPPORTED_ABI_VERSION", - Map( - "version" -> pkg.manifest.abiVersion.toString, - "supported" -> NativeBridge.ABI_VERSION.toString)) - } - Some(result) + case Seq(result) => Some(result) case _ => throw QueryCompilationErrors.nativeDataSourcePackageConflictError( name, found.map(_._2.path).sorted) diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index eca83f7365b1f..b5934a2a31b97 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -667,8 +667,15 @@ void streamRelease(ArrowArrayStream* stream) { // The driver-side writer of native_sink. struct Writer { Options state; // path, overwrite, fail. + // The file that closeWriter creates, so that the tests can check that the writer of a committed + // or aborted write is closed: _CLOSED for a batch write, _CLOSED_ for a streaming one. + std::string closed_marker; }; +[[maybe_unused]] std::string closedMarker(jlong epoch_id) { + return epoch_id < 0 ? "/_CLOSED" : "/_CLOSED_" + std::to_string(epoch_id); +} + // Writes the rows of a task as CSV lines to a temporary file, renamed on commit. struct DataWriter { std::string temp_path; @@ -1043,12 +1050,15 @@ SPARK_JNI(void, commit)(JNIEnv* env, jclass, jlong handle, jlong epoch_id, } writeFile(path + (epoch_id < 0 ? "/_SUCCESS" : "/_EPOCH_" + std::to_string(epoch_id)), std::to_string(temp_files.size())); + writer->closed_marker = closedMarker(epoch_id); }); } -SPARK_JNI(void, abort)(JNIEnv* env, jclass, jlong handle, jlong, jobjectArray messages) { +SPARK_JNI(void, abort)(JNIEnv* env, jclass, jlong handle, jlong epoch_id, + jobjectArray messages) { guardedVoid(env, [&] { auto* writer = fromHandle(handle); + writer->closed_marker = closedMarker(epoch_id); bool has_null = false; for (const auto& temp_file : toByteStrings(env, messages, &has_null)) { unlink(temp_file.c_str()); @@ -1057,7 +1067,14 @@ SPARK_JNI(void, abort)(JNIEnv* env, jclass, jlong handle, jlong, jobjectArray me }); } -SPARK_JNI(void, closeWriter)(JNIEnv*, jclass, jlong handle) { delete fromHandle(handle); } +SPARK_JNI(void, closeWriter)(JNIEnv* env, jclass, jlong handle) { + std::unique_ptr writer(fromHandle(handle)); + guardedVoid(env, [&] { + if (!writer->closed_marker.empty()) { + writeFile(writer->state["path"] + writer->closed_marker, ""); + } + }); +} SPARK_JNI(jlong, createDataWriter)(JNIEnv* env, jclass, jbyteArray writer_state, jint partition_id, jlong task_id, jlong epoch_id) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala index 894bb2b9c8520..c41c906b554a4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala @@ -30,6 +30,7 @@ import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.sources.DataSourceRegister @@ -140,6 +141,30 @@ class ColumnarDataSourceSuite extends QueryTest with SharedSparkSession { assert(TestColumnarSink.rows(table).sorted == (0L until 5L).map(id => (id, s"name$id"))) assert(TestColumnarSink.epochs.get(table) == Seq(0L)) } + + test("create a stream writer for each micro-batch") { + val table = "micro_batch_write" + val input = MemoryStream[(Long, String)] + val query = input.toDF().toDF("id", "name") + .writeStream + .format(format) + .option("table", table) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try { + (0L until 3L).foreach { id => + input.addData((id, s"name$id")) + query.processAllAvailable() + } + } finally { + query.stop() + } + assert(TestColumnarSink.rows(table).sorted == (0L until 3L).map(id => (id, s"name$id"))) + assert(TestColumnarSink.epochs.get(table) == Seq(0L, 1L, 2L)) + // Each micro-batch creates its stream writer with the ID of the query, so that the native + // data sources close the stream writer of the previous micro-batch of the query. + assert(TestColumnarSink.streamWriterQueryIds.get(table) == Seq.fill(3)(query.id.toString)) + } } /** Reads the ids [0, end) and their names, and writes to [[TestColumnarSink]]. */ @@ -180,6 +205,8 @@ class TestRangeDataSource(options: CaseInsensitiveStringMap) extends ColumnarDat schema: StructType, overwrite: Boolean, queryId: String): ColumnarStreamWriter = { + TestColumnarSink.streamWriterQueryIds.compute( + table, (_, old) => Option(old).getOrElse(Nil) :+ queryId) new TestSinkStreamWriter(table) } } @@ -275,6 +302,7 @@ object TestColumnarSink { val committed = new ConcurrentHashMap[String, Seq[TestCommitMessage]]() val aborted = ConcurrentHashMap.newKeySet[String]() val epochs = new ConcurrentHashMap[String, Seq[Long]]() + val streamWriterQueryIds = new ConcurrentHashMap[String, Seq[String]]() def rows(table: String): Seq[(Long, String)] = committed.get(table).flatMap(_.rows) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 865e1e7d54e4e..93634772c8dd3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -332,6 +332,8 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { write("append", 0 until 4) assert(written() == (0 until 4).map(i => s"$i,v$i").toSet) assert(new File(dir, "_SUCCESS").isFile) + // The writer is closed after the commit. + assert(new File(dir, "_CLOSED").isFile) write("append", 4 until 6) assert(written() == (0 until 6).map(i => s"$i,v$i").toSet) write("overwrite", 10 until 12) @@ -383,6 +385,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { "name" -> "native_sink", "msg" -> "injected failure in write")) assert(new File(dir, "_ABORTED").isFile) + assert(new File(dir, "_CLOSED").isFile) assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) } @@ -426,6 +429,9 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } // Each call to latestOffset makes 5 more rows available: [0, 5), [5, 10) and [10, 12). assert((0 to 2).forall(epoch => new File(dir, s"_EPOCH_$epoch").isFile)) + // Each micro-batch creates a new stream writer, which closes the one of the previous + // micro-batch. The last one is closed once it is no longer referenced. + assert((0 to 1).forall(epoch => new File(dir, s"_CLOSED_$epoch").isFile)) val written = dir.listFiles().filter(_.getName.startsWith("part-")) .flatMap(file => Files.readAllLines(file.toPath).asScala) assert(written.sorted.toSeq == (0 until 12).map(_.toString).sorted) @@ -577,6 +583,14 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { // no longer configured. session.conf.unset(SQLConf.NATIVE_DATA_SOURCE_PATHS.key) assert(provider.inferSchema(CaseInsensitiveStringMap.empty()).fieldNames.head == "id") + + // The result is only kept until the next lookup on the thread, whatever it finds, so a + // provider never takes the result of an earlier lookup. + session.conf.set(SQLConf.NATIVE_DATA_SOURCE_PATHS.key, defaultPackage.getPath) + assert(DataSource.lookupDataSource("native_range", session.sessionState.conf) == + classOf[NativeDataSourceV2]) + assert(DataSource.lookupDataSource("json", session.sessionState.conf) != null) + assert(NativeDataSourceRegistry.takeLookup("native_range").isEmpty) } nativeTest("Python data sources take precedence") { @@ -750,7 +764,8 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { parameters = Map("provider" -> "com.typo.Missing")) val e = intercept[AnalysisException](session.sql("SELECT * FROM missing_db.missing_table")) assert(e.getCondition == "TABLE_OR_VIEW_NOT_FOUND") - // The invalid package whose manifest lists the data source reports why it is invalid. + // The invalid package whose manifest lists the data source reports why it is invalid, when + // no other package or installed library provides the data source. checkError( exception = intercept[SparkRuntimeException] { session.read.format("other_source").load() @@ -761,6 +776,60 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { "reason" -> "The package does not contain the library lib.so.")) } + nativeTest("stale packages do not hide a valid one") { + // A shared directory still has broken and outdated copies of the package of native_range. + val dir = Utils.createTempDir() + createZip(dir, "native_range-1.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + dataSources(""), Map(NativePlatform.current -> "lib.so")).getBytes(UTF_8)) + createPackage(library, dataSources(""), dir = dir, fileName = "native_range-2.sparkpkg", + abiVersion = 2) + val session = newSession(dir) + checkError( + exception = intercept[SparkRuntimeException](range(session)), + condition = "INVALID_NATIVE_DATA_SOURCE_PACKAGE.UNSUPPORTED_ABI_VERSION", + parameters = Map( + "path" -> new File(dir, "native_range-2.sparkpkg").getCanonicalPath, + "version" -> "2", + "supported" -> "1")) + + // A valid package added to the session takes precedence over them, without a conflict. + session.addArtifact(defaultPackage.getPath) + checkRows(range(session, "end" -> "2"), (0L until 2L).map(expectedRow)) + + // So does an installed library. + val libraries = Utils.createTempDir() + installLibrary(installedLibrary, "installed_native_range", libraries) + createZip(dir, "installed-1.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + Seq("installed_native_range"), Map(NativePlatform.current -> "lib.so")).getBytes(UTF_8)) + withLibraryPath(libraries) { + checkRows( + newSession(dir).read.format("installed_native_range").option("end", "2").load() + .select("id"), + Seq(Row(0L), Row(1L))) + } + } + + nativeTest("forget the packages of a session that is cleaned up") { + val session = spark.newSession() + session.addArtifact(defaultPackage.getPath) + val added = session.artifactManager.getNativeDataSourcePackages._1.head.getCanonicalPath + NativeDataSourcePackage.read(new File(added)) + assert(NativeDataSourcePackage.isCached(added)) + session.artifactManager.cleanUpResourcesForTesting() + assert(!NativeDataSourcePackage.isCached(added)) + + // Only the packages under the directory are forgotten. + val dir = Utils.createTempDir() + val pkg = createPackage(library, dataSources(""), dir = dir).getCanonicalPath + NativeDataSourcePackage.read(new File(pkg)) + NativeDataSourcePackage.forgetPackagesIn(new File(dir.getPath + "x")) + assert(NativeDataSourcePackage.isCached(pkg)) + NativeDataSourcePackage.forgetPackagesIn(dir) + assert(!NativeDataSourcePackage.isCached(pkg)) + } + nativeTest("invalid installed libraries") { val dir = Utils.createTempDir() val notALibrary = new File(dir, NativeDataSourceRegistry.libraryFileName("not_a_library")) From 9ac5a496e6050ed5c253511179adda747aa0a110 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Sat, 10 Oct 2026 18:45:45 +0900 Subject: [PATCH 10/27] Address review comments on sessions, distributed copies, schemas and batching - Resolve the provider of classic DataFrameWriter commands and of DataStreamWriter sinks with the session of the Dataset active, so that the lookups of Python and native data sources use that session. - Distribute a package that is not an artifact of the session as a copy named after its file name and checksum, which the lookups do not count as a package of the session: a configured package can be replaced in place, and packages with the same file name no longer collide. Remove the then unreachable NATIVE_DATA_SOURCE_PACKAGE_FILE_NAME_CONFLICT error. - Log a warning when a package is distributed without an active session. - Skip, and log, a configured directory that cannot be read. - Compare the data with the read schema as it round-trips through Arrow, which does not carry collations, interval fields and user-defined types. - Also end the write batches at spark.sql.execution.arrow.maxBytesPerBatch. - Document that a failed batch commit is aborted, which closes the writer, why NativeBridge.load is called reflectively, and the security of discovering libraries by file name. Co-authored-by: Isaac --- .../resources/error/error-conditions.json | 6 - docs/sql-native-data-sources.md | 18 +- .../sql/errors/QueryExecutionErrors.scala | 7 - .../spark/sql/datasource/NativeBridge.java | 6 +- .../datasources/v2/ffi/NativeLibrary.java | 4 + .../spark/sql/artifact/ArtifactManager.scala | 9 + .../spark/sql/classic/DataFrameWriter.scala | 6 +- .../spark/sql/classic/DataStreamWriter.scala | 78 +++++---- .../v2/columnar/ColumnarWrite.scala | 19 ++- .../datasources/v2/ffi/NativeDataSource.scala | 10 +- .../v2/ffi/NativeDataSourcePackage.scala | 5 +- .../v2/ffi/NativeDataSourceRegistry.scala | 60 ++++--- .../v2/columnar/ColumnarDataSourceSuite.scala | 13 ++ .../v2/ffi/NativeDataSourceSuite.scala | 159 +++++++++++++++--- 14 files changed, 292 insertions(+), 108 deletions(-) diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index f10e7253a3908..a20ee51b340d6 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -6326,12 +6326,6 @@ ], "sqlState" : "42710" }, - "NATIVE_DATA_SOURCE_PACKAGE_FILE_NAME_CONFLICT" : { - "message" : [ - "Cannot add the native data source package to the session, because the session already has a different file named . Rename one of the packages." - ], - "sqlState" : "42710" - }, "NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND" : { "message" : [ "Cannot find the native data source package with the SHA-256 checksum . Add the package with `spark.addArtifact`, or install the same package under `spark.sql.dataSource.native.paths` on every node." diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 717a0cba854cf..b75d24ffdff09 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -168,13 +168,14 @@ A session finds the packages in two places, and the executors get them automatic ``` - **Configured paths**: the packages listed in `spark.sql.dataSource.native.paths`, or in the - directories it lists, on the driver. + directories it lists, on the driver. A package can be replaced in place: the next queries use + the new version. The packages of a session take precedence over the installed libraries. Give packages distinct file names, for example with their version: sessions that are not isolated share the files they -add, and a session cannot use two different packages with the same file name. Spark ignores, and -logs, a package that it cannot read or whose version it does not support, unless its manifest -lists the data source that a query uses and no other package or installed library provides it. +add. Spark ignores, and logs, a package that it cannot read or whose version it does not +support, unless its manifest lists the data source that a query uses and no other package or +installed library provides it. ## Example in Rust @@ -551,7 +552,12 @@ statically, or install them on the nodes. ## Security A native data source runs native code in the driver and executor processes, with their -privileges, and is not sandboxed. Only install libraries and add packages that you trust, and do -not let untrusted users write to the directories of the native library path. Set +privileges, and is not sandboxed. Spark finds installed libraries by their file name, in the +directories of the native library path, and loads one when a query uses its name and no Java or +Python data source has that name, much like it imports the Python data sources installed as +`pyspark_*` packages. So only install libraries and add packages that you trust, and do not let +untrusted users write to the directories of the native library path, or to the directories in +`spark.sql.dataSource.native.paths`. A user who can add artifacts to a session, for example a Spark +Connect client, can also add a package, as they can add JAR files. Set `spark.sql.dataSource.native.enabled` to false to prevent loading native data sources, for example on a server shared by several users. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala index ccaf02e452839..862d3c180981a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala @@ -828,13 +828,6 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE messageParameters = Map("fileName" -> fileName, "checksum" -> checksum)) } - def nativeDataSourcePackageFileNameConflictError( - path: String, fileName: String): SparkRuntimeException = { - new SparkRuntimeException( - errorClass = "NATIVE_DATA_SOURCE_PACKAGE_FILE_NAME_CONFLICT", - messageParameters = Map("path" -> path, "fileName" -> fileName)) - } - def invalidNativeDataSourceLibraryError( path: String, subClass: String, diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java index f7adf1e14dcc8..fc431af1811f1 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java @@ -62,9 +62,11 @@ * {@link #createStreamWriter}, and serializes its state with {@link #serializeWriter}. On the * executors, each task creates a data writer with {@link #createDataWriter}, calls * {@link #write} for each batch of at most {@code spark.sql.execution.arrow.maxRecordsPerBatch} - * rows, and calls {@link #commitDataWriter}, or {@link #abortDataWriter} if the task or the + * rows, which also ends once it holds {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, + * and calls {@link #commitDataWriter}, or {@link #abortDataWriter} if the task or the * commit fails. Then the driver calls {@link #commit} with the messages of the tasks if they - * all succeeded, or {@link #abort} otherwise, and closes the writer. A streaming write commits + * all succeeded, or {@link #abort} otherwise, including when {@link #commit} fails, and closes + * the writer. A streaming write commits * or aborts each micro-batch (epoch). A micro-batch streaming query creates a new stream writer * for each micro-batch, and Spark closes the previous stream writer of the query when the query * creates a new one. Spark is not told when a streaming query stops, so it closes the last diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java index 3b2343ee142d8..61145e186879e 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -121,6 +121,10 @@ private static MethodHandle find( static NativeLibrary load(String path) throws Throwable { Class bridge = new BridgeClassLoader(NativeBridge.class.getClassLoader()) .loadClass(BRIDGE_CLASS); + // The copy itself has to call System.load, because a library belongs to the class loader of + // the class that loads it. NativeBridge.load is package-private, so that it is not part of + // the public interface, so it is called reflectively. Java always allows this for the classes + // on the class path: they are in unnamed modules, whose packages are all open. Method load = bridge.getDeclaredMethod("load", String.class); load.setAccessible(true); try { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala index 80a5b2e539bc6..e0fb397d36cdc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala @@ -167,6 +167,15 @@ class ArtifactManager(session: SparkSession) extends AutoCloseable with Logging (packages, Option(state).map(_.uuid)) } + /** Returns whether a file with the given name was added to this session as a file artifact. */ + private[sql] def hasFileArtifact(fileName: String): Boolean = { + val path = Paths.get("files", fileName) + sparkContextRelativePaths.asScala.exists { + case (SparkContextResourceType.FILE, relativePath, _) => relativePath == path + case _ => false + } + } + protected[sql] def getCachedBlockId(hash: String): Option[CacheId] = { Option(hashToCachedIdMap.get(hash)).map(_.id) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala index 3d0036f865225..eab75bc3e5bbd 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala @@ -629,8 +629,10 @@ final class DataFrameWriter[T] private[sql](ds: Dataset[T]) extends sql.DataFram * Wrap a DataFrameWriter action to track the QueryExecution and time cost, then report to the * user-registered callback functions. */ - private def runCommand(session: SparkSession)(command: LogicalPlan): Unit = { - val qe = new QueryExecution(session, command, df.queryExecution.tracker, + private def runCommand(session: SparkSession)(command: => LogicalPlan): Unit = { + // The command is built with the session active, like a query is analyzed: it looks up the + // data source, and the lookups of the Python and native data sources use the active session. + val qe = new QueryExecution(session, session.withActive(command), df.queryExecution.tracker, shuffleCleanupModeOpt = Some(QueryExecution.determineShuffleCleanupMode(session.sessionState.conf))) qe.assertCommandExecuted() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala index 1e7989e4cc6c5..0369bcfacfbf8 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala @@ -253,51 +253,55 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) extends streaming.D val sink = new ForeachBatchSink[T](foreachBatchWriter, ds.exprEnc) startQuery(sink, extraOptions, catalogTable = catalogTable) } else { - val cls = DataSource.lookupDataSource(source, ds.sparkSession.sessionState.conf) - val disabledSources = - Utils.stringToSeq(ds.sparkSession.sessionState.conf.disabledV2StreamingWriters) - val useV1Source = disabledSources.contains(cls.getCanonicalName) || - // file source v2 does not support streaming yet. - classOf[FileDataSourceV2].isAssignableFrom(cls) - val optionsWithPath = if (path.isEmpty) { extraOptions } else { extraOptions + ("path" -> path.get) } + // The sink is resolved with the session active, like a query is analyzed: the lookups of + // the Python and native data sources use the active session. + val sink = ds.sparkSession.withActive(resolveSink(optionsWithPath)) + startQuery(sink, optionsWithPath, catalogTable = catalogTable) + } + } - val sink = if (classOf[TableProvider].isAssignableFrom(cls) && !useV1Source) { - val provider = cls.getConstructor().newInstance().asInstanceOf[TableProvider] - val sessionOptions = DataSourceV2Utils.extractSessionConfigs( - source = provider, conf = ds.sparkSession.sessionState.conf) - val finalOptions = sessionOptions.filter { case (k, _) => !optionsWithPath.contains(k) } ++ - optionsWithPath.originalMap - val dsOptions = new CaseInsensitiveStringMap(finalOptions.asJava) - // If the source accepts external table metadata, here we pass the schema of input query - // to `getTable`. This is for avoiding schema inference, which can be very expensive. - // If the query schema is not compatible with the existing data, the behavior is undefined. - val outputSchema = if (provider.supportsExternalMetadata()) { - Some(ds.schema) - } else { - None - } - provider match { - case p: NamedTableProvider => p.setShortName(source) - case _ => - } - val table = DataSourceV2Utils.getTableFromProvider( - provider, dsOptions, userSpecifiedSchema = outputSchema) - import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._ - table match { - case table: SupportsWrite if table.supports(STREAMING_WRITE) => - table - case _ => createV1Sink(optionsWithPath) - } + private def resolveSink(optionsWithPath: CaseInsensitiveMap[String]): Table = { + val cls = DataSource.lookupDataSource(source, ds.sparkSession.sessionState.conf) + val disabledSources = + Utils.stringToSeq(ds.sparkSession.sessionState.conf.disabledV2StreamingWriters) + val useV1Source = disabledSources.contains(cls.getCanonicalName) || + // file source v2 does not support streaming yet. + classOf[FileDataSourceV2].isAssignableFrom(cls) + + if (classOf[TableProvider].isAssignableFrom(cls) && !useV1Source) { + val provider = cls.getConstructor().newInstance().asInstanceOf[TableProvider] + val sessionOptions = DataSourceV2Utils.extractSessionConfigs( + source = provider, conf = ds.sparkSession.sessionState.conf) + val finalOptions = sessionOptions.filter { case (k, _) => !optionsWithPath.contains(k) } ++ + optionsWithPath.originalMap + val dsOptions = new CaseInsensitiveStringMap(finalOptions.asJava) + // If the source accepts external table metadata, here we pass the schema of input query + // to `getTable`. This is for avoiding schema inference, which can be very expensive. + // If the query schema is not compatible with the existing data, the behavior is undefined. + val outputSchema = if (provider.supportsExternalMetadata()) { + Some(ds.schema) } else { - createV1Sink(optionsWithPath) + None } - - startQuery(sink, optionsWithPath, catalogTable = catalogTable) + provider match { + case p: NamedTableProvider => p.setShortName(source) + case _ => + } + val table = DataSourceV2Utils.getTableFromProvider( + provider, dsOptions, userSpecifiedSchema = outputSchema) + import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._ + table match { + case table: SupportsWrite if table.supports(STREAMING_WRITE) => + table + case _ => createV1Sink(optionsWithPath) + } + } else { + createV1Sink(optionsWithPath) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala index 9f167f97ea211..f1f051abae4cc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala @@ -115,16 +115,25 @@ class ColumnarStreamingDataWriterFactory( /** * How [[ArrowBatchingDataWriter]] converts rows to Arrow, captured on the driver from the - * session configuration. + * session configuration. Like the other Arrow batches of Spark, a batch ends when it reaches + * `spark.sql.execution.arrow.maxRecordsPerBatch` rows or, after a row, holds at least + * `spark.sql.execution.arrow.maxBytesPerBatch` bytes. */ -case class ArrowBatchingOptions(schema: StructType, timeZoneId: String, maxRecordsPerBatch: Int) +case class ArrowBatchingOptions( + schema: StructType, + timeZoneId: String, + maxRecordsPerBatch: Int, + maxBytesPerBatch: Long) object ArrowBatchingOptions { def apply(schema: StructType): ArrowBatchingOptions = { val conf = SQLConf.get val maxRecords = conf.arrowMaxRecordsPerBatch ArrowBatchingOptions( - schema, conf.sessionLocalTimeZone, if (maxRecords > 0) maxRecords else Int.MaxValue) + schema, + conf.sessionLocalTimeZone, + if (maxRecords > 0) maxRecords else Int.MaxValue, + conf.arrowMaxBytesPerBatch) } } @@ -151,7 +160,9 @@ class ArrowBatchingDataWriter(delegate: DataWriter[ColumnarBatch], options: Arro } arrowWriter.write(record) numRows += 1 - if (numRows >= options.maxRecordsPerBatch) { + // The size is not computed when the limit is the largest possible one, as in Python UDFs. + if (numRows >= options.maxRecordsPerBatch || (options.maxBytesPerBatch < Int.MaxValue && + arrowWriter.sizeInBytes() >= options.maxBytesPerBatch)) { flush() } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index d3794d3279b54..77074530079dc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -290,6 +290,7 @@ class NativeDataSourceWriter( call(name, "commit a write to") { library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) } + // Not closed if the commit fails: Spark then aborts the write, which closes it. handle.close() } @@ -418,14 +419,19 @@ class NativePartitionReader( "is not supported.") } val actual = ArrowUtils.fromArrowSchema(root.getSchema) - val matches = actual.length == schema.length && actual.zip(schema).forall { + // The schema to read as the library got it: Arrow does not carry everything, such as the + // collations of strings, the fields of intervals and user-defined types, which still apply to + // the data that Spark reads. + val expected = ArrowUtils.fromArrowSchema(ArrowUtils.toArrowSchema( + schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = false)) + val matches = actual.length == expected.length && actual.zip(expected).forall { case (actualField, field) => actualField.name == field.name && DataType.equalsIgnoreNullability(actualField.dataType, field.dataType) } if (!matches) { throw error(name, "read", s"The data has the schema ${actual.toDDL}, but the expected " + - s"schema is ${schema.toDDL}.") + s"schema is ${expected.toDDL}.") } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala index c9bb312bea6b9..ab613e92f873e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -75,12 +75,15 @@ sealed trait NativeLibraryLocation extends Serializable * @param path the path of the package file on the driver * @param artifactUUID the UUID of the session whose artifacts contain the package, if the session * is isolated. Executors store the artifacts of such a session separately. + * @param distributedFileName the name of the copy of the package in the artifacts of the session, + * if the package is not itself an artifact of the session */ case class NativeDataSourcePackage( path: String, sha256: String, manifest: NativeDataSourceManifest, - artifactUUID: Option[String] = None) extends NativeLibraryLocation { + artifactUUID: Option[String] = None, + distributedFileName: Option[String] = None) extends NativeLibraryLocation { def fileName: String = new File(path).getName } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala index ba07c41369662..b83366384affc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -177,32 +177,40 @@ object NativeDataSourceRegistry extends Logging { case pkg: NativeDataSourcePackage => SparkSession.getActiveSession match { // In local mode, the executors run in the driver and read the package where it is. - case Some(session) if !session.sparkContext.isLocal => addToArtifacts(session, pkg) - case _ => pkg + case Some(session) if session.sparkContext.isLocal => pkg + case Some(session) => addToArtifacts(session, pkg) + case None => + logWarning(log"Cannot distribute the native data source package " + + log"${MDC(PATH, pkg.path)} to the executors: there is no active session. The " + + log"executors only find it if it is at the same path on every node.") + pkg } case installed: InstalledNativeLibrary => installed } - /** Adds the package as an artifact of the session, unless it already is one. */ + /** + * Makes the package available to the executors of the session, as one of its artifacts. A + * package that the session added with `spark.addArtifact` already is one. Any other package, + * such as a package under `spark.sql.dataSource.native.paths`, is copied to the artifacts under + * a name made of its file name and its checksum: the copy does not have the extension of a + * package, so it is not a package of the session that the lookups find, and two different + * packages with the same file name do not collide. + */ private[ffi] def addToArtifacts( session: SparkSession, pkg: NativeDataSourcePackage): NativeDataSourcePackage = { val artifactManager = session.artifactManager val (added, artifactUUID) = artifactManager.getNativeDataSourcePackages - // The executors find the package by its file name, so the session cannot have another - // package with the same file name. - added.find(_.getName == pkg.fileName) match { - case Some(file) => - val sha256 = NativeDataSourcePackage.tryRead(file).map(_.sha256).getOrElse("") - if (sha256 != pkg.sha256) { - throw QueryExecutionErrors.nativeDataSourcePackageFileNameConflictError( - pkg.path, pkg.fileName) - } - case None => + if (added.exists(_.getCanonicalPath == pkg.path)) { + pkg.copy(artifactUUID = artifactUUID) + } else { + val copyName = s"${pkg.fileName}.${pkg.sha256}" + if (!artifactManager.hasFileArtifact(copyName)) { artifactManager.addLocalArtifacts(Artifact.newFileArtifact( - Paths.get(pkg.fileName), new Artifact.LocalFile(Paths.get(pkg.path))) :: Nil) + Paths.get(copyName), new Artifact.LocalFile(Paths.get(pkg.path))) :: Nil) + } + pkg.copy(artifactUUID = artifactUUID, distributedFileName = Some(copyName)) } - pkg.copy(artifactUUID = artifactUUID) } /** @@ -215,13 +223,14 @@ object NativeDataSourceRegistry extends Logging { if (SparkEnv.get == null) { Nil } else { + val fileName = pkg.distributedFileName.getOrElse(pkg.fileName) val root = new File(SparkFiles.getRootDirectory()) val sessionCopies = Option(root.listFiles()).toSeq.flatten .filter(_.isDirectory) .sortBy(_.getName) - .map(new File(_, pkg.fileName)) - (pkg.artifactUUID.map(uuid => new File(new File(root, uuid), pkg.fileName)).toSeq ++ - Seq(new File(root, pkg.fileName)) ++ sessionCopies).distinct + .map(new File(_, fileName)) + (pkg.artifactUUID.map(uuid => new File(new File(root, uuid), fileName)).toSeq ++ + Seq(new File(root, fileName)) ++ sessionCopies).distinct } } @@ -245,10 +254,17 @@ object NativeDataSourceRegistry extends Logging { } private def packageFilesIn(dir: File): Seq[File] = { - dir.listFiles() - .filter(f => f.isFile && f.getName.endsWith(NativeDataSourcePackage.FILE_EXTENSION)) - .sortBy(_.getName) - .toSeq + Option(dir.listFiles()) match { + case Some(files) => + files + .filter(f => f.isFile && f.getName.endsWith(NativeDataSourcePackage.FILE_EXTENSION)) + .sortBy(_.getName) + .toSeq + case None => + logWarning(log"Ignoring the native data source path ${MDC(PATH, dir)}: the directory " + + log"cannot be read.") + Nil + } } private def artifactPackageFiles(): Seq[File] = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala index c41c906b554a4..3fb43fb937139 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala @@ -114,6 +114,19 @@ class ColumnarDataSourceSuite extends QueryTest with SharedSparkSession { } } + test("batch write in batches of a limited size") { + withSQLConf(SQLConf.ARROW_EXECUTION_MAX_BYTES_PER_BATCH.key -> "4096") { + val table = "batch_write_bytes" + val rows = (0L until 10L).map(id => (id, "x" * 1000 + id)) + rows.toDF("id", "name").coalesce(1) + .write.format(format).option("table", table).mode("append").save() + assert(TestColumnarSink.rows(table) == rows) + // A batch ends once it holds at least 4096 bytes, long before it has 10,000 rows: each row + // takes about 1 KB, so a batch has 5 rows. + assert(TestColumnarSink.batchSizes(table) == Seq(5, 5)) + } + } + test("abort a failed write") { val table = "failed_write" val e = intercept[SparkException] { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 93634772c8dd3..41c3cd6d825ef 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -275,6 +275,12 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { .load(), (0L until 6L).map(id => Row(id, if (id % 5 == 0) null else s"n$id"))) + // The Arrow schema does not have the collation, which applies to the data that Spark reads. + val collated = session.read.schema("id BIGINT, name STRING COLLATE UTF8_LCASE") + .format("native_range").option("end", "6").load() + checkRows(collated.filter("name = 'N1'"), Seq(Row(1L, "n1"))) + assert(scan(collated.filter("name = 'N1'")).getMetaData()("PushedPredicates") == "[]") + val e = intercept[Exception] { session.read.schema("id INT").format("native_range").load().collect() } @@ -389,6 +395,26 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) } + nativeTest("abort a write whose commit failed") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val e = intercept[Exception] { + session.range(10).write.format("native_sink").option("path", dir.getPath) + .option("fail", "commit").mode("append").save() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "commit a write to", + "name" -> "native_sink", + "msg" -> "injected failure in commit")) + assert(e.getSuppressed.isEmpty) + // Spark aborted the write with the writer, which it closed afterwards. + assert(new File(dir, "_ABORTED").isFile) + assert(new File(dir, "_CLOSED").isFile) + } + nativeTest("abort a data writer whose commit failed") { val session = newSession(defaultPackage) val dir = Utils.createTempDir() @@ -531,31 +557,37 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { val pkg = NativeDataSourcePackage.read(added.head).copy(artifactUUID = artifactUUID) assert(NativeDataSourceRegistry.distributedCopies(pkg).exists(_.isFile)) - // A package under spark.sql.dataSource.native.paths is added as an artifact when it is used, - // except in local mode, where the executors run in the driver. - val other = newSession(defaultPackage) + // A package that the session added is distributed as it is. + assert(NativeDataSourceRegistry.addToArtifacts(session, pkg) == pkg) + + // A package under spark.sql.dataSource.native.paths is copied to the artifacts of the session + // when it is used, except in local mode, where the executors run in the driver. The copy is + // named after the checksum of the package, and it is not a package of the session. + val configuredDir = Utils.createTempDir() + val configured = createPackage(library, dataSources(""), dir = configuredDir, + fileName = defaultPackage.getName) + val other = newSession(configured) range(other, "end" -> "1").collect() + assert(!other.artifactManager.hasFileArtifact(defaultPackage.getName)) + val configuredPackage = NativeDataSourcePackage.read(configured) + val distributed = NativeDataSourceRegistry.addToArtifacts(other, configuredPackage) + val copyName = s"${defaultPackage.getName}.${configuredPackage.sha256}" + assert(distributed == configuredPackage.copy( + artifactUUID = other.artifactManager.getNativeDataSourcePackages._2, + distributedFileName = Some(copyName))) + assert(other.artifactManager.hasFileArtifact(copyName)) assert(other.artifactManager.getNativeDataSourcePackages._1.isEmpty) - val distributed = - NativeDataSourceRegistry.addToArtifacts(other, NativeDataSourcePackage.read(defaultPackage)) - val (otherPackages, otherArtifactUUID) = other.artifactManager.getNativeDataSourcePackages - assert(otherPackages.map(_.getName) == Seq(defaultPackage.getName)) - assert(distributed.artifactUUID == otherArtifactUUID) + assert(NativeDataSourceRegistry.distributedCopies(distributed).exists(_.isFile)) // Adding the same package again does nothing. - NativeDataSourceRegistry.addToArtifacts(other, NativeDataSourcePackage.read(defaultPackage)) - assert(other.artifactManager.getNativeDataSourcePackages._1.size == 1) + assert(NativeDataSourceRegistry.addToArtifacts(other, configuredPackage) == distributed) - // A different package with the same file name cannot be added: the executors would not find - // the package by its file name. + // A package with the same file name as a package of the session does not collide with it. + other.addArtifact(defaultPackage.getPath) val sameFileName = createPackage(negatingLibrary, dataSources("b_"), dir = Utils.createTempDir(), fileName = defaultPackage.getName) - checkError( - exception = intercept[SparkRuntimeException] { - NativeDataSourceRegistry.addToArtifacts(other, NativeDataSourcePackage.read(sameFileName)) - }, - condition = "NATIVE_DATA_SOURCE_PACKAGE_FILE_NAME_CONFLICT", - parameters = Map( - "path" -> sameFileName.getCanonicalPath, "fileName" -> defaultPackage.getName)) + val sameFileNameCopy = + NativeDataSourceRegistry.addToArtifacts(other, NativeDataSourcePackage.read(sameFileName)) + assert(sameFileNameCopy.distributedFileName.exists(_.startsWith(defaultPackage.getName + "."))) // The executors also check the copies in the files of the other sessions, such as the session // of a streaming query, which is a clone of the session that started it. @@ -574,6 +606,95 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } } + test("distribute a package without an active session") { + val pkg = NativeDataSourcePackage( + "/path/to/x.sparkpkg", "checksum", NativeDataSourceManifest(1, Seq("x"), Map.empty)) + val logAppender = new LogAppender("distribute without an active session") + SparkSession.clearActiveSession() + try { + withLogAppender(logAppender, level = Some(org.apache.logging.log4j.Level.WARN)) { + assert(NativeDataSourceRegistry.distribute(pkg) == pkg) + } + } finally { + SparkSession.setActiveSession(spark) + } + assert(logAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains("there is no active session"))) + } + + nativeTest("write with a session that is not the active one") { + val session = newSession() + session.addArtifact(defaultPackage.getPath) + SparkSession.setActiveSession(spark) + val dir = Utils.createTempDir() + session.range(3).selectExpr("id", "concat('v', id) AS name") + .write.format("native_sink").option("path", dir.getPath).mode("append").save() + val written = dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala) + assert(written.toSet == Set("0,v0", "1,v1", "2,v2")) + + val streamDir = Utils.createTempDir() + val query = session.readStream.format("native_counter").option("max_offset", "3").load() + .writeStream + .format("native_sink") + .option("path", streamDir.getPath) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try { + query.processAllAvailable() + } finally { + query.stop() + } + assert(streamDir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala).sorted.toSeq == Seq("0", "1", "2")) + } + + nativeTest("ignore configured directories that cannot be read") { + val dir = Utils.createTempDir() + createPackage(library, dataSources(""), dir = dir, fileName = "a.sparkpkg") + assert(dir.setReadable(false)) + try { + assume(dir.listFiles() == null, "The directory can still be read, for example by root.") + val session = newSession(dir) + checkError( + exception = intercept[SparkClassNotFoundException] { + session.read.format("no_such_source").load() + }, + condition = "DATA_SOURCE_NOT_FOUND", + parameters = Map("provider" -> "no_such_source")) + val e = intercept[AnalysisException](session.sql("SELECT * FROM missing_db.missing_table")) + assert(e.getCondition == "TABLE_OR_VIEW_NOT_FOUND") + } finally { + dir.setReadable(true) + } + } + + nativeTest("replace a configured package in place") { + val dir = Utils.createTempDir() + val configured = createPackage(library, dataSources(""), dir = dir, fileName = "x.sparkpkg") + val session = newSession(configured) + checkRows(range(session, "end" -> "2"), (0L until 2L).map(expectedRow)) + // On a cluster, the package is copied to the artifacts of the session when it is used. + val first = + NativeDataSourceRegistry.addToArtifacts(session, NativeDataSourcePackage.read(configured)) + + // An operator replaces the package with a new version: here, with one more file. + val path = s"${NativePlatform.current}/${library.getName}" + createZip(dir, "x.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + dataSources(""), Map(NativePlatform.current -> path)).getBytes(UTF_8), + path -> Files.readAllBytes(library.toPath), + "README" -> "version 2".getBytes(UTF_8)) + // The copy of the previous version is not a package of the session, so the session uses the + // new version, which is copied under another name. + checkRows(range(session, "end" -> "3"), (0L until 3L).map(expectedRow)) + val second = + NativeDataSourceRegistry.addToArtifacts(session, NativeDataSourcePackage.read(configured)) + assert(second.sha256 != first.sha256) + assert(second.distributedFileName != first.distributedFileName) + assert(session.artifactManager.getNativeDataSourcePackages._1.isEmpty) + } + nativeTest("look up a native data source once") { val session = newSession(defaultPackage) val provider = DataSource.lookupDataSourceV2("native_range", session.sessionState.conf).get From acdc7eb9def74e454b3db40cc30750441f4d661c Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Sun, 11 Oct 2026 07:18:17 +0900 Subject: [PATCH 11/27] Check the Arrow data of native data sources, create writers late, push limits once - Read the Arrow stream of a native library array by array, and reject an array with a non-zero offset at any level, which Arrow Java would read from the start of its buffers, with a NATIVE_DATA_SOURCE_ERROR that names the column. - Name the fields of the entries of maps `key` and `value` before creating the vectors, as ArrowColumnVector looks them up by name, so that maps built by arrow-rs can be read; reject decimal256, and report any other column that ArrowColumnVector cannot read with its name. - Create the writer of a batch or streaming write with its writer factory, after the stages that the write depends on ran, so that Spark commits or aborts each writer it creates, also when such a stage fails with adaptive execution. - Ask the reader for a limit only once, as NativeBridge documents, although Spark pushes the limit of LIMIT ... OFFSET twice. - Document these in NativeBridge and the native data sources page. Co-authored-by: Isaac --- docs/sql-native-data-sources.md | 3 +- .../spark/sql/datasource/NativeBridge.java | 16 +- .../v2/columnar/ColumnarTable.scala | 20 ++- .../v2/columnar/ColumnarWrite.scala | 39 ++++- .../datasources/v2/ffi/NativeDataSource.scala | 158 ++++++++++++++---- .../test_native_datasource.cc | 121 +++++++++++--- .../v2/columnar/ColumnarDataSourceSuite.scala | 17 ++ .../v2/ffi/NativeDataSourceSuite.scala | 47 ++++++ 8 files changed, 349 insertions(+), 72 deletions(-) diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index b75d24ffdff09..ab9a2e11ad40a 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -88,7 +88,8 @@ The [Javadoc of `NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeB specifies each function, including the ownership of the Arrow structs and the JSON format of the pushed predicates. A predicate that a library reports as evaluated completely must follow the semantics of Spark SQL, for example for NaN values, which are equal to each other and greater -than any other value. +than any other value. The arrays that a library exports must have an offset of 0: copy a sliced +array, such as the result of `RecordBatch::slice` in arrow-rs, before exporting it. ## Finding Native Data Sources diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java index fc431af1811f1..554741c042f58 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java @@ -58,8 +58,10 @@ * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then * calls {@link #read} on the executors. Offsets are JSON strings defined by the library. - *

  • Writes. On the driver, Spark creates a writer with {@link #createWriter} or - * {@link #createStreamWriter}, and serializes its state with {@link #serializeWriter}. On the + *
  • Writes. On the driver, once it computed the data to write, right before it starts + * the tasks of the write (or of a micro-batch), Spark creates a writer with + * {@link #createWriter} or {@link #createStreamWriter}, and serializes its state with + * {@link #serializeWriter}. If computing the data fails, Spark creates no writer. On the * executors, each task creates a data writer with {@link #createDataWriter}, calls * {@link #write} for each batch of at most {@code spark.sql.execution.arrow.maxRecordsPerBatch} * rows, which also ends once it holds {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, @@ -92,13 +94,17 @@ * interface. A {@code long} argument named {@code ...Address} is the address of a C struct * allocated by Spark. For an output struct, the library initializes it and Spark takes * ownership. For an input struct, the library may move it to take ownership; otherwise, Spark - * releases it after the call returns, so the library must not use it afterwards.
    • + * releases it after the call returns, so the library must not use it afterwards. The arrays + * that the library exports must have an offset of 0, at every level: Arrow Java ignores the + * offsets of the arrays it imports, so Spark rejects an array with another offset, such as a + * zero-copy slice. Copy a sliced array before exporting it. *
  • Schemas. Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is * {@code timestamp[us]} without one. Data read from a library may also use - * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, but not - * unsigned integers or dictionary encoding.
    • + * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, and any names + * for the fields of the entries of a map, but not unsigned integers, {@code decimal256} or + * dictionary encoding. *
  • Options. Option keys are passed in lower case, because Spark matches them * case-insensitively.
    • *
  • Optional methods. A library only needs to export the functions for the operations diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala index 17d253d0c9a94..eb973fe12a8f6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala @@ -91,6 +91,8 @@ class ColumnarScanBuilder( private var pushed: Array[Predicate] = Array.empty private var pushedLimit: Option[Int] = None + // The limit passed to the reader, and whether the reader applies it. + private var limitResult: Option[(Int, Boolean)] = None private var readSchema: StructType = schema private def batchReader: ColumnarReader = { @@ -108,12 +110,18 @@ class ColumnarScanBuilder( override def pushedPredicates(): Array[Predicate] = pushed - override def pushLimit(limit: Int): Boolean = { - val isPushed = batchReader.pushLimit(limit) - if (isPushed) { - pushedLimit = Some(limit) - } - isPushed + override def pushLimit(limit: Int): Boolean = limitResult match { + // Spark can push a limit more than once, for example the limit of LIMIT ... OFFSET, but the + // reader is only asked once. A limit at least as large as the first one is already applied. + case Some((firstLimit, isPushed)) => + isPushed && limit >= firstLimit + case None => + val isPushed = batchReader.pushLimit(limit) + limitResult = Some(limit -> isPushed) + if (isPushed) { + pushedLimit = Some(limit) + } + isPushed } // Spark always applies the limit again: a data source is free to return more rows. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala index f1f051abae4cc..4c9b814cbadc5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala @@ -57,31 +57,56 @@ class ColumnarWrite( override def description(): String = shortName override def toBatch: BatchWrite = { - new ColumnarBatchWrite(createDataSource().writer(info.schema(), overwrite), info.schema()) + new ColumnarBatchWrite( + () => createDataSource().writer(info.schema(), overwrite), info.schema()) } override def toStreaming: StreamingWrite = { new ColumnarStreamingWrite( - createDataSource().streamWriter(info.schema(), overwrite, info.queryId()), info.schema()) + () => createDataSource().streamWriter(info.schema(), overwrite, info.queryId()), + info.schema()) } } -class ColumnarBatchWrite(writer: ColumnarWriter, schema: StructType) extends BatchWrite { +/** + * A batch write. The writer of the data source is created with the writer factory, once Spark + * computed the stages that the write depends on: Spark only aborts the write if it fails later. + */ +class ColumnarBatchWrite(createWriter: () => ColumnarWriter, schema: StructType) + extends BatchWrite { + + private var writer: ColumnarWriter = _ override def createBatchWriterFactory(info: PhysicalWriteInfo): DataWriterFactory = { + if (writer == null) { + writer = createWriter() + } new ColumnarDataWriterFactory(writer, ArrowBatchingOptions(schema)) } override def commit(messages: Array[WriterCommitMessage]): Unit = writer.commit(messages) - override def abort(messages: Array[WriterCommitMessage]): Unit = writer.abort(messages) + override def abort(messages: Array[WriterCommitMessage]): Unit = { + if (writer != null) { + writer.abort(messages) + } + } } -class ColumnarStreamingWrite(writer: ColumnarStreamWriter, schema: StructType) +/** + * A streaming write. As for a [[ColumnarBatchWrite]], the writer of the data source is created + * with the writer factory, once Spark computed the stages of the micro-batch. + */ +class ColumnarStreamingWrite(createWriter: () => ColumnarStreamWriter, schema: StructType) extends StreamingWrite { + private var writer: ColumnarStreamWriter = _ + override def createStreamingWriterFactory( info: PhysicalWriteInfo): StreamingDataWriterFactory = { + if (writer == null) { + writer = createWriter() + } new ColumnarStreamingDataWriterFactory(writer, ArrowBatchingOptions(schema)) } @@ -90,7 +115,9 @@ class ColumnarStreamingWrite(writer: ColumnarStreamWriter, schema: StructType) } override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { - writer.abort(epochId, messages) + if (writer != null) { + writer.abort(epochId, messages) + } } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index 77074530079dc..52f2978c27806 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -23,10 +23,10 @@ import java.util.concurrent.atomic.AtomicLong import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal -import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, Data} +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, CDataDictionaryProvider, Data} import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} -import org.apache.arrow.vector.ipc.ArrowReader -import org.apache.arrow.vector.types.pojo.Field +import org.apache.arrow.vector.complex.MapVector +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} import org.apache.spark.internal.Logging @@ -39,6 +39,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataType, StructType} import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.unsafe.Platform import org.apache.spark.util.Utils /** @@ -361,8 +362,8 @@ object NativeWriterCommitMessage { } /** - * Reads a partition of a native data source: the Arrow stream exported by the library is - * imported into Arrow vectors that back the returned batch. + * Reads a partition of a native data source: each array of the Arrow stream exported by the + * library is checked, and imported into the Arrow vectors that back the returned batch. */ class NativePartitionReader( location: NativeLibraryLocation, @@ -375,41 +376,62 @@ class NativePartitionReader( private val allocator = ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) - private var reader: ArrowReader = _ + // Read here rather than with `Data.importArrayStream`, to check each array before importing it. + private val stream = ArrowArrayStream.allocateNew(allocator) + private val dictionaries = new CDataDictionaryProvider() private var root: VectorSchemaRoot = _ private var batch: ColumnarBatch = _ try { - reader = openStream() - root = call(name, "read")(reader.getVectorSchemaRoot) + openStream() + root = VectorSchemaRoot.create(importSchema(), allocator) checkSchema() - batch = new ColumnarBatch( - root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector).toArray) + batch = new ColumnarBatch(root.getFieldVectors.asScala.map(columnVector).toArray) } catch { case e: Throwable => close() throw e } - private def openStream(): ArrowReader = { + private def openStream(): Unit = { val library = NativeLibraries.get(location) - val stream = ArrowArrayStream.allocateNew(allocator) + call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) + .getOrElse { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) + } + } + + private def importSchema(): Schema = { + val arrowSchema = ArrowSchema.allocateNew(allocator) try { - call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) - .getOrElse { - throw new SparkUnsupportedOperationException( - "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) - } - Data.importArrayStream(allocator, stream) - } catch { - case e: Throwable => - stream.release() - throw e + call(name, "read")(stream.getSchema(arrowSchema)) + val imported = call(name, "read")(Data.importSchema(allocator, arrowSchema, dictionaries)) + new Schema( + imported.getFields.asScala.map(withMapEntryNames).asJava, imported.getCustomMetadata) } finally { - stream.close() + arrowSchema.close() } } + /** + * Names the fields of the entries of the maps `key` and `value`, which ArrowColumnVector looks + * up by name. Arrow does not fix these names: arrow-rs names them `keys` and `values`, for + * example. The arrays are imported by position, so the names do not change how they are read. + */ + private def withMapEntryNames(field: Field): Field = { + val children = field.getChildren.asScala.map(withMapEntryNames).toSeq + val renamed = (field.getType, children) match { + case (_: ArrowType.Map, Seq(entries)) if entries.getChildren.size == 2 => + val Seq(key, value) = entries.getChildren.asScala.toSeq + Seq(new Field(entries.getName, entries.getFieldType, Seq( + new Field(MapVector.KEY_NAME, key.getFieldType, key.getChildren), + new Field(MapVector.VALUE_NAME, value.getFieldType, value.getChildren)).asJava)) + case _ => children + } + new Field(field.getName, field.getFieldType, renamed.asJava) + } + private def checkSchema(): Unit = { def isDictionaryEncoded(field: Field): Boolean = { field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) @@ -418,6 +440,17 @@ class NativePartitionReader( throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + "is not supported.") } + // The Arrow types that convert to Spark types, but that ArrowColumnVector cannot read. + def unsupportedType(field: Field): Option[ArrowType] = field.getType match { + case decimal: ArrowType.Decimal if decimal.getBitWidth != 128 => Some(decimal) + case _ => field.getChildren.asScala.view.flatMap(unsupportedType).headOption + } + root.getSchema.getFields.asScala.foreach { field => + unsupportedType(field).foreach { arrowType => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"$arrowType, which is not supported.") + } + } val actual = ArrowUtils.fromArrowSchema(root.getSchema) // The schema to read as the library got it: Arrow does not carry everything, such as the // collations of strings, the fields of intervals and user-defined types, which still apply to @@ -435,21 +468,80 @@ class NativePartitionReader( } } + private def columnVector(vector: FieldVector): ColumnVector = { + try { + new ArrowColumnVector(vector) + } catch { + case NonFatal(e) => + throw QueryExecutionErrors.nativeDataSourceError("read", name, + s"Spark cannot read the column ${vector.getName} of the Arrow type " + + s"${vector.getField.getType}: ${message(e)}", e) + } + } + override def next(): Boolean = { - val hasNext = call(name, "read")(reader.loadNextBatch()) - if (hasNext) { - batch.setNumRows(root.getRowCount) + val array = ArrowArray.allocateNew(allocator) + try { + call(name, "read")(stream.getNext(array)) + val snapshot = array.snapshot() + if (snapshot.release == 0) { + // The end of the stream. + false + } else { + try { + checkOffsets(snapshot) + } catch { + case e: Throwable => + array.release() + throw e + } + call(name, "read")(Data.importIntoVectorSchemaRoot(allocator, array, root, dictionaries)) + batch.setNumRows(root.getRowCount) + true + } + } finally { + array.close() + } + } + + /** + * Checks that the array of a batch, and the arrays of its columns at every level, have an + * offset of 0: Arrow Java ignores the offsets of the arrays it imports, so it would read a + * sliced array from the start of its buffers. + */ + private def checkOffsets(array: ArrowArray.Snapshot): Unit = { + def children(array: ArrowArray.Snapshot): Seq[ArrowArray.Snapshot] = { + (0 until array.n_children.toInt).map { i => + ArrowArray.wrap(Platform.getLong(null, array.children + 8L * i)).snapshot() + } + } + def checkColumn(array: ArrowArray.Snapshot, column: String): Unit = { + if (array.offset != 0) { + throw error(name, "read", s"The column $column has an offset of ${array.offset}, but " + + "Spark only reads arrays with an offset of 0. Copy sliced arrays before exporting them.") + } + children(array).foreach(checkColumn(_, column)) + } + if (array.offset != 0) { + throw error(name, "read", s"The batch has an offset of ${array.offset}, but Spark only " + + "reads arrays with an offset of 0. Copy sliced arrays before exporting them.") + } + children(array).zip(root.getSchema.getFields.asScala).foreach { case (column, field) => + checkColumn(column, field.getName) } - hasNext } override def get(): ColumnarBatch = batch override def close(): Unit = { Utils.tryWithSafeFinally { - if (reader != null) { - reader.close() + if (root != null) { + root.close() } + dictionaries.close() + // Releases the stream of the library, unless it was not exported. + stream.release() + stream.close() }(Utils.tryLogNonFatalError(allocator.close())) } } @@ -600,9 +692,11 @@ object NativeCalls { // How Arrow reports the error message of an Arrow C stream. private val StreamError = """(?s)\[errno \d+\] CDataJniException\{errno=\d+, message=(.*)\}""".r - private def message(e: Throwable): String = Option(e.getMessage).getOrElse(e.toString) match { - case StreamError(message) => message - case message => message + private[ffi] def message(e: Throwable): String = { + Option(e.getMessage).getOrElse(e.toString) match { + case StreamError(message) => message + case message => message + } } /** Calls a function that the library may not implement. Returns None if it does not. */ diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index b5934a2a31b97..5b05dc30105f1 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -20,7 +20,10 @@ // interface structs by hand. It provides three data sources: // // native_range Reads rows [0, end) with columns of several types. Supports predicate, -// column and limit pushdown. +// column and limit pushdown. With the option "extra", the schema has one more +// column: "map" adds a map whose entry fields are named "keys" and "values", +// as arrow-rs names them, and "decimal256" adds a decimal256. With the option +// "slice", the columns of the batches have an offset of 1. // native_sink Writes the rows as CSV files under the directory in the option "path". // native_counter Streams the rows [0, max_offset) with a single column "id", making "step" // more rows available each time Spark asks for the latest offset. @@ -137,17 +140,31 @@ const ColumnType kColumns[] = { {"ts", "tsu:UTC"} // timestamp }; +// The columns that native_range adds with the option "extra", by its value. +const std::map kExtraColumns = { + {"map", {"m", "+m"}}, // map: {"k": id} + {"decimal256", {"dec", "d:10,2,256"}} // decimal256(10, 2): id +}; + const ColumnType* findColumn(const std::string& name) { for (const auto& column : kColumns) { if (name == column.name) return &column; } + for (const auto& extra : kExtraColumns) { + if (name == extra.second.name) return &extra.second; + } throw std::runtime_error("unknown column " + name); } -std::vector allColumns(bool counter) { +std::vector allColumns(bool counter, const std::string& extra = "") { if (counter) return {"id"}; std::vector names; for (const auto& column : kColumns) names.push_back(column.name); + if (!extra.empty()) { + auto it = kExtraColumns.find(extra); + if (it == kExtraColumns.end()) throw std::runtime_error("unknown extra column " + extra); + names.push_back(it->second.name); + } return names; } @@ -185,18 +202,30 @@ void initSchema(ArrowSchema* out, const std::string& format, const std::string& out->private_data = data; } +// Adds a child to a schema initialized by initSchema, and returns it. +ArrowSchema* addChild(ArrowSchema* parent, const std::string& format, const std::string& name, + int64_t flags) { + auto* data = static_cast(parent->private_data); + auto child = std::make_unique(); + initSchema(child.get(), format, name, flags); + data->child_pointers.push_back(child.get()); + data->children.push_back(std::move(child)); + parent->n_children = static_cast(data->children.size()); + parent->children = data->child_pointers.data(); + return data->children.back().get(); +} + // Exports a struct type with the given columns. void exportSchema(const std::vector& columns, ArrowSchema* out) { initSchema(out, "+s", "", 0); - auto* data = static_cast(out->private_data); for (const auto& name : columns) { - auto child = std::make_unique(); - initSchema(child.get(), findColumn(name)->format, name, ARROW_FLAG_NULLABLE); - data->child_pointers.push_back(child.get()); - data->children.push_back(std::move(child)); + ArrowSchema* column = addChild(out, findColumn(name)->format, name, ARROW_FLAG_NULLABLE); + if (name == "m") { + ArrowSchema* entries = addChild(column, "+s", "entries", 0); + addChild(entries, "u", "keys", 0); + addChild(entries, "l", "values", ARROW_FLAG_NULLABLE); + } } - out->n_children = static_cast(data->children.size()); - out->children = data->child_pointers.data(); } // Reads the column names of a struct schema exported by Spark, and releases it. @@ -259,6 +288,16 @@ void setBit(std::vector& bitmap, int64_t i) { bitmap[i / 8] |= static_cast(1 << (i % 8)); } +// Adds a child to an array initialized by initArray, and returns it to be initialized. +ArrowArray* addChild(ArrowArray* parent) { + auto* data = static_cast(parent->private_data); + data->children.push_back(std::make_unique()); + data->child_pointers.push_back(data->children.back().get()); + parent->n_children = static_cast(data->children.size()); + parent->children = data->child_pointers.data(); + return data->children.back().get(); +} + // Exports the values of a column for the ids [lo, hi). void exportColumn(const std::string& name, int64_t lo, int64_t hi, ArrowArray* out) { int64_t n = hi - lo; @@ -305,24 +344,52 @@ void exportColumn(const std::string& name, int64_t lo, int64_t hi, ArrowArray* o std::vector values; for (int64_t id = lo; id < hi; id++) values.push_back(id * 1000000); initArray(out, n, 0, {no_validity, toBytes(values)}); + } else if (name == "m") { + // A single entry for each row. + std::vector offsets; + for (int64_t i = 0; i <= n; i++) offsets.push_back(static_cast(i)); + initArray(out, n, 0, {no_validity, toBytes(offsets)}); + ArrowArray* entries = addChild(out); + initArray(entries, n, 0, {no_validity}); + std::vector key_offsets{0}; + std::string keys; + std::vector values; + for (int64_t id = lo; id < hi; id++) { + keys += "k" + std::to_string(id); + key_offsets.push_back(static_cast(keys.size())); + values.push_back(id); + } + initArray(addChild(entries), n, 0, + {no_validity, toBytes(key_offsets), std::vector(keys.begin(), keys.end())}); + initArray(addChild(entries), n, 0, {no_validity, toBytes(values)}); + } else if (name == "dec") { + // 256-bit little-endian unscaled values of the scale 2. + std::vector words; + for (int64_t id = lo; id < hi; id++) { + int64_t unscaled = id * 100; + int64_t sign = unscaled < 0 ? -1 : 0; + words.insert(words.end(), {unscaled, sign, sign, sign}); + } + initArray(out, n, 0, {no_validity, toBytes(words)}); } else { throw std::runtime_error("unknown column " + name); } } -// Exports a struct array with the given columns for the ids [lo, hi). -void exportBatch(const std::vector& columns, int64_t lo, int64_t hi, +// Exports a struct array with the given columns for the ids [lo, hi). A sliced batch exports +// the columns for [lo - 1, hi), with an offset of 1, like a zero-copy slice. +void exportBatch(const std::vector& columns, int64_t lo, int64_t hi, bool slice, ArrowArray* out) { initArray(out, hi - lo, 0, {std::vector()}); - auto* data = static_cast(out->private_data); for (const auto& name : columns) { - auto child = std::make_unique(); - exportColumn(name, lo, hi, child.get()); - data->child_pointers.push_back(child.get()); - data->children.push_back(std::move(child)); + ArrowArray* column = addChild(out); + exportColumn(name, slice ? lo - 1 : lo, hi, column); + if (slice) { + column->offset = 1; + column->length = hi - lo; + column->null_count = -1; // Unknown. + } } - out->n_children = static_cast(data->children.size()); - out->children = data->child_pointers.data(); } // --------------------------------------------------------------------------------------------- @@ -630,6 +697,7 @@ struct Stream { int64_t end; int64_t batch_size; bool fail; + bool slice; std::string last_error; }; @@ -650,7 +718,7 @@ int streamGetNext(ArrowArrayStream* stream, ArrowArray* out) { return 0; } int64_t hi = std::min(state->end, state->next + state->batch_size); - exportBatch(state->columns, state->next, hi, out); + exportBatch(state->columns, state->next, hi, state->slice, out); state->next = hi; return 0; } @@ -776,7 +844,8 @@ SPARK_JNI(void, schema)(JNIEnv* env, jclass, jlong handle, jlong schema_address) if (source->name == name("native_sink")) { throw std::runtime_error("native_sink has no schema; it can only be written"); } - exportSchema(allColumns(source->name == name("native_counter")), + bool counter = source->name == name("native_counter"); + exportSchema(allColumns(counter, get(source->options, "extra", "")), reinterpret_cast(schema_address)); }); } @@ -796,6 +865,7 @@ SPARK_JNI(jlong, createReader)(JNIEnv* env, jclass, jlong handle, jlong schema_a reader->state["partitions"] = get(source->options, "partitions", "2"); reader->state["batch_size"] = get(source->options, "batch_size", "1000"); reader->state["fail"] = get(source->options, "fail", ""); + reader->state["slice"] = get(source->options, "slice", ""); return toHandle(reader); }); } @@ -871,7 +941,10 @@ SPARK_JNI(jbooleanArray, pushPredicates)(JNIEnv* env, jclass, jlong handle, SPARK_JNI(jboolean, pushLimit)(JNIEnv* env, jclass, jlong handle, jint limit) { return guarded(env, JNI_FALSE, [&] { - fromHandle(handle)->state["limit"] = std::to_string(limit); + auto* reader = fromHandle(handle); + // NativeBridge calls pushLimit at most once. + if (reader->state.count("limit") > 0) throw std::runtime_error("pushLimit was called twice"); + reader->state["limit"] = std::to_string(limit); return JNI_TRUE; }); } @@ -978,7 +1051,8 @@ SPARK_JNI(void, read)(JNIEnv* env, jclass, jbyteArray reader_state, jbyteArray p stream->release = streamRelease; stream->private_data = new Stream{split(get(state, "columns", "")), std::stoll(bounds[0]), std::stoll(bounds[1]), getLong(state, "batch_size", 1000), - get(state, "fail", "") == "read", ""}; + get(state, "fail", "") == "read", + get(state, "slice", "") == "true", ""}; }); } @@ -1010,6 +1084,9 @@ jlong createAnyWriter(JNIEnv* env, jlong handle, jlong schema_address, jboolean delete writer; throw std::runtime_error("the option 'path' is required"); } + // Records that a writer was created, which Spark then commits or aborts. + makeDirs(path); + writeFile(path + "/_CREATED", ""); return toHandle(writer); }); } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala index 3fb43fb937139..bc94db65cf0ab 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala @@ -127,6 +127,20 @@ class ColumnarDataSourceSuite extends QueryTest with SharedSparkSession { } } + test("create the writer once the data to write is computed") { + val table = "failed_upstream" + // The error happens in a shuffle stage, which adaptive query execution runs before the write, + // so no writer is created. Spark would not abort it. + val e = intercept[Exception] { + spark.range(10).selectExpr("id", "if(id >= 0, raise_error('boom'), 'x') AS name") + .repartition(2) + .write.format(format).option("table", table).mode("append").save() + } + assert(e.getMessage.contains("boom")) + assert(!TestColumnarSink.writers.contains(table)) + assert(!TestColumnarSink.aborted.contains(table)) + } + test("abort a failed write") { val table = "failed_write" val e = intercept[SparkException] { @@ -211,6 +225,7 @@ class TestRangeDataSource(options: CaseInsensitiveStringMap) extends ColumnarDat } override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + TestColumnarSink.writers.add(table) new TestSinkWriter(table, overwrite, options.getBoolean("fail", false)) } @@ -314,6 +329,8 @@ class TestBatchReader(lo: Long, hi: Long, schema: StructType) object TestColumnarSink { val committed = new ConcurrentHashMap[String, Seq[TestCommitMessage]]() val aborted = ConcurrentHashMap.newKeySet[String]() + // The tables for which a writer was created. + val writers = ConcurrentHashMap.newKeySet[String]() val epochs = new ConcurrentHashMap[String, Seq[Long]]() val streamWriterQueryIds = new ConcurrentHashMap[String, Seq[String]]() diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 41c3cd6d825ef..9a4be5b8d344e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -322,6 +322,37 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { val df = range(newSession(defaultPackage), "end" -> "100", "partitions" -> "5").limit(3) assert(df.collect().length == 3) assert(scan(df).getMetaData()("PushedLimit") == "LIMIT 3") + + // Spark pushes the limit of LIMIT ... OFFSET twice, but the library, which fails if + // pushLimit is called again, is only asked once. + val withOffset = range(newSession(defaultPackage), "end" -> "100", "partitions" -> "5") + .limit(10).offset(5) + checkRows(withOffset, (5L until 10L).map(expectedRow)) + assert(scan(withOffset).getMetaData()("PushedLimit") == "LIMIT 10") + } + + nativeTest("read maps whose entry fields have other names") { + // The library names the fields of the entries `keys` and `values`, as arrow-rs does. + val df = range(newSession(defaultPackage), "end" -> "3", "extra" -> "map").select("id", "m") + assert(df.schema("m").dataType == MapType(StringType, LongType, valueContainsNull = true)) + checkRows(df, (0L until 3L).map(id => Row(id, Map(s"k$id" -> id)))) + } + + nativeTest("reject the Arrow data that Spark cannot read") { + def check(msg: String, options: (String, String)*): Unit = { + val e = intercept[Exception] { + range(newSession(defaultPackage), options :+ ("end" -> "4"): _*).collect() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map("action" -> "read", "name" -> "native_range", "msg" -> msg)) + } + // Arrow Java would read a slice from the start of its buffers. + check("The column id has an offset of 1, but Spark only reads arrays with an offset of 0. " + + "Copy sliced arrays before exporting them.", "slice" -> "true") + check("The column dec contains the Arrow type Decimal(10, 2, 256), which is not supported.", + "extra" -> "decimal256") } nativeTest("batch write") { @@ -395,6 +426,22 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) } + nativeTest("create no writer when computing the data to write fails") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + // The error happens in a shuffle stage, which adaptive query execution runs before the write. + val e = intercept[Exception] { + session.range(10).selectExpr("id", "if(id >= 0, raise_error('boom'), 'x') AS name") + .repartition(2) + .write.format("native_sink").option("path", dir.getPath).mode("append").save() + } + assert(e.getMessage.contains("boom")) + assert(!new File(dir, "_CREATED").exists()) + + session.range(2).write.format("native_sink").option("path", dir.getPath).mode("append").save() + assert(new File(dir, "_CREATED").isFile) + } + nativeTest("abort a write whose commit failed") { val session = newSession(defaultPackage) val dir = Utils.createTempDir() From a8f178b3e2474777d1c506913788342e2867a179 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Sun, 11 Oct 2026 19:02:14 +0900 Subject: [PATCH 12/27] Check the lengths of the arrays that native data sources export - Before importing a batch from a native library, check that each column has at least as many values as the batch has rows, each struct field as many as its struct, and the elements of each list or map as many as its last offset says. Arrow Java's import sets the length of every array to what its parent refers to and extends a shorter array with nulls, so such a column read as nulls, and only the lengths that the library exported show the mismatch. - Document in NativeBridge and the native data sources page that the buffers must hold all the values of their arrays, which Spark cannot check because the C data interface does not carry the sizes of buffers. Co-authored-by: Isaac --- docs/sql-native-data-sources.md | 5 +- .../spark/sql/datasource/NativeBridge.java | 6 +- .../datasources/v2/ffi/NativeDataSource.scala | 58 +++++++++++++++---- .../test_native_datasource.cc | 28 +++++++-- .../v2/ffi/NativeDataSourceSuite.scala | 7 +++ 5 files changed, 86 insertions(+), 18 deletions(-) diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index ab9a2e11ad40a..ff3aa2b5bd01e 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -89,7 +89,10 @@ specifies each function, including the ownership of the Arrow structs and the JS pushed predicates. A predicate that a library reports as evaluated completely must follow the semantics of Spark SQL, for example for NaN values, which are equal to each other and greater than any other value. The arrays that a library exports must have an offset of 0: copy a sliced -array, such as the result of `RecordBatch::slice` in arrow-rs, before exporting it. +array, such as the result of `RecordBatch::slice` in arrow-rs, before exporting it. Spark also +checks that each array has at least as many values as its parent refers to, but it cannot check +that the buffers of the arrays hold those values, because the C data interface does not carry the +sizes of buffers: a library that exports a buffer that is too short makes Spark read past its end. ## Finding Native Data Sources diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java index 554741c042f58..ed455937b5108 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java @@ -97,7 +97,11 @@ * releases it after the call returns, so the library must not use it afterwards. The arrays * that the library exports must have an offset of 0, at every level: Arrow Java ignores the * offsets of the arrays it imports, so Spark rejects an array with another offset, such as a - * zero-copy slice. Copy a sliced array before exporting it.
    • + * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with + * fewer values than its parent refers to, such as a column shorter than its batch. The buffers + * of an array must hold all of its values, as the C data interface requires: it does not carry + * the sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that + * is too short. *
  • Schemas. Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index 52f2978c27806..a9f3dd06c6ada 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -489,7 +489,7 @@ class NativePartitionReader( false } else { try { - checkOffsets(snapshot) + checkArrays(snapshot) } catch { case e: Throwable => array.release() @@ -505,29 +505,65 @@ class NativePartitionReader( } /** - * Checks that the array of a batch, and the arrays of its columns at every level, have an - * offset of 0: Arrow Java ignores the offsets of the arrays it imports, so it would read a - * sliced array from the start of its buffers. + * Checks the array of a batch, and the arrays of its columns at every level, before importing + * them. Arrow Java trusts the arrays, and would read them wrongly otherwise: + * - An array must have an offset of 0. Arrow Java ignores the offsets of the arrays it imports, + * so it would read a sliced array from the start of its buffers. + * - An array must have at least as many values as its parent refers to: a column as many as + * the batch has rows, a field as many as its struct has values, and the elements of a list or + * a map as many as its last offset says. Arrow Java would extend a shorter array with nulls. + * Only the lengths that the library exported show this: the import sets the length of each + * array to what its parent refers to. + * The buffers are not checked, since the C data interface does not carry their sizes: Arrow + * Java takes them from the lengths, so the library must export buffers that hold every value. */ - private def checkOffsets(array: ArrowArray.Snapshot): Unit = { + private def checkArrays(batch: ArrowArray.Snapshot): Unit = { def children(array: ArrowArray.Snapshot): Seq[ArrowArray.Snapshot] = { (0 until array.n_children.toInt).map { i => ArrowArray.wrap(Platform.getLong(null, array.children + 8L * i)).snapshot() } } - def checkColumn(array: ArrowArray.Snapshot, column: String): Unit = { + // The last offset of a list or a map, in its second buffer. Arrow Java rejects the array if + // the buffer is missing. + def lastOffset(array: ArrowArray.Snapshot): Long = { + val offsets = if (array.length > 0 && array.n_buffers >= 2) { + Platform.getLong(null, array.buffers + 8L) + } else { + 0L + } + if (offsets == 0) 0 else Platform.getInt(null, offsets + 4L * array.length) + } + // The number of values that an array refers to in each of its children. + def referenced(array: ArrowArray.Snapshot, arrowType: ArrowType): Long = arrowType match { + case _: ArrowType.Struct => array.length + case _: ArrowType.List | _: ArrowType.Map => lastOffset(array) + // A list view has no last offset: its lists refer to ranges of its elements in any order. + case _ => 0 + } + def checkColumn(array: ArrowArray.Snapshot, field: Field, column: String): Unit = { if (array.offset != 0) { throw error(name, "read", s"The column $column has an offset of ${array.offset}, but " + "Spark only reads arrays with an offset of 0. Copy sliced arrays before exporting them.") } - children(array).foreach(checkColumn(_, column)) + val length = referenced(array, field.getType) + children(array).zip(field.getChildren.asScala).foreach { case (child, childField) => + if (child.length < length) { + throw error(name, "read", s"The column $column contains an array with a length of " + + s"${child.length}, but its parent refers to $length values.") + } + checkColumn(child, childField, column) + } } - if (array.offset != 0) { - throw error(name, "read", s"The batch has an offset of ${array.offset}, but Spark only " + + if (batch.offset != 0) { + throw error(name, "read", s"The batch has an offset of ${batch.offset}, but Spark only " + "reads arrays with an offset of 0. Copy sliced arrays before exporting them.") } - children(array).zip(root.getSchema.getFields.asScala).foreach { case (column, field) => - checkColumn(column, field.getName) + children(batch).zip(root.getSchema.getFields.asScala).foreach { case (column, field) => + if (column.length < batch.length) { + throw error(name, "read", s"The column ${field.getName} has a length of " + + s"${column.length}, but the batch has a length of ${batch.length}.") + } + checkColumn(column, field, field.getName) } } diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index 5b05dc30105f1..7217ab8e6523b 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -23,7 +23,11 @@ // column and limit pushdown. With the option "extra", the schema has one more // column: "map" adds a map whose entry fields are named "keys" and "values", // as arrow-rs names them, and "decimal256" adds a decimal256. With the option -// "slice", the columns of the batches have an offset of 1. +// "slice", the columns of the batches have an offset of 1. With the option +// "short", an array of the batches has one value fewer than its parent refers +// to, although its buffers hold them all: "id" shortens the column id, +// "entries" the entries of the map and their fields, and "values" the values +// of the entries. // native_sink Writes the rows as CSV files under the directory in the option "path". // native_counter Streams the rows [0, max_offset) with a single column "id", making "step" // more rows available each time Spark asks for the latest offset. @@ -377,9 +381,10 @@ void exportColumn(const std::string& name, int64_t lo, int64_t hi, ArrowArray* o } // Exports a struct array with the given columns for the ids [lo, hi). A sliced batch exports -// the columns for [lo - 1, hi), with an offset of 1, like a zero-copy slice. +// the columns for [lo - 1, hi), with an offset of 1, like a zero-copy slice. `short_array` names +// the array to export with one value fewer than its parent refers to, as the option "short". void exportBatch(const std::vector& columns, int64_t lo, int64_t hi, bool slice, - ArrowArray* out) { + const std::string& short_array, ArrowArray* out) { initArray(out, hi - lo, 0, {std::vector()}); for (const auto& name : columns) { ArrowArray* column = addChild(out); @@ -389,6 +394,16 @@ void exportBatch(const std::vector& columns, int64_t lo, int64_t hi column->length = hi - lo; column->null_count = -1; // Unknown. } + if (short_array == "id" && name == "id") { + column->length--; + } else if (short_array == "entries" && name == "m") { + // The last offset of the map refers to one entry more than there are. + ArrowArray* entries = column->children[0]; + entries->length--; + for (int64_t i = 0; i < entries->n_children; i++) entries->children[i]->length--; + } else if (short_array == "values" && name == "m") { + column->children[0]->children[1]->length--; + } } } @@ -698,6 +713,7 @@ struct Stream { int64_t batch_size; bool fail; bool slice; + std::string short_array; std::string last_error; }; @@ -718,7 +734,7 @@ int streamGetNext(ArrowArrayStream* stream, ArrowArray* out) { return 0; } int64_t hi = std::min(state->end, state->next + state->batch_size); - exportBatch(state->columns, state->next, hi, state->slice, out); + exportBatch(state->columns, state->next, hi, state->slice, state->short_array, out); state->next = hi; return 0; } @@ -866,6 +882,7 @@ SPARK_JNI(jlong, createReader)(JNIEnv* env, jclass, jlong handle, jlong schema_a reader->state["batch_size"] = get(source->options, "batch_size", "1000"); reader->state["fail"] = get(source->options, "fail", ""); reader->state["slice"] = get(source->options, "slice", ""); + reader->state["short"] = get(source->options, "short", ""); return toHandle(reader); }); } @@ -1052,7 +1069,8 @@ SPARK_JNI(void, read)(JNIEnv* env, jclass, jbyteArray reader_state, jbyteArray p stream->private_data = new Stream{split(get(state, "columns", "")), std::stoll(bounds[0]), std::stoll(bounds[1]), getLong(state, "batch_size", 1000), get(state, "fail", "") == "read", - get(state, "slice", "") == "true", ""}; + get(state, "slice", "") == "true", get(state, "short", ""), + ""}; }); } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 9a4be5b8d344e..07db1bb8c47cf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -351,6 +351,13 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { // Arrow Java would read a slice from the start of its buffers. check("The column id has an offset of 1, but Spark only reads arrays with an offset of 0. " + "Copy sliced arrays before exporting them.", "slice" -> "true") + // Arrow Java would extend an array with fewer values than its parent refers to with nulls. + // Each batch has 2 rows. + check("The column id has a length of 1, but the batch has a length of 2.", "short" -> "id") + Seq("entries", "values").foreach { array => + check("The column m contains an array with a length of 1, but its parent refers to 2 " + + "values.", "extra" -> "map", "short" -> array) + } check("The column dec contains the Arrow type Decimal(10, 2, 256), which is not supported.", "extra" -> "decimal256") } From ff7b1ae0150fbcb3b6a7c67bea8accf6aa4ad9bd Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Tue, 13 Oct 2026 15:03:34 +0900 Subject: [PATCH 13/27] Address review comments on the compatibility, predicates and checks of native data sources - NativeBridge states its compatibility rules, the guarantees of task attempts, epoch commits and replayed stream ranges, blocking reads, and the threads that release handles. - Only a fixed, documented set of predicates is pushed down. - List views are rejected, since their lengths are not checked. - Executors only load an installed library with the checksum of the driver's file. - Classic PySpark addArtifact passes URIs of packages to the JVM as is. - New NativeDataSourceClusterSuite, in local-cluster mode. Co-authored-by: Isaac --- .../resources/error/error-conditions.json | 2 +- docs/sql-native-data-sources.md | 39 ++++- python/pyspark/sql/session.py | 11 +- python/pyspark/sql/tests/test_artifact.py | 19 ++- .../sql/errors/QueryExecutionErrors.scala | 4 +- .../spark/sql/datasource/NativeBridge.java | 100 ++++++++---- .../datasources/v2/ffi/NativeDataSource.scala | 7 +- .../v2/ffi/NativeDataSourcePackage.scala | 46 ++++-- .../v2/ffi/NativeDataSourceRegistry.scala | 13 +- .../datasources/v2/ffi/NativePredicates.scala | 14 +- .../test_native_datasource.cc | 29 +++- .../v2/ffi/NativeDataSourceClusterSuite.scala | 147 ++++++++++++++++++ .../v2/ffi/NativeDataSourceSuite.scala | 44 +++++- 13 files changed, 404 insertions(+), 71 deletions(-) create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index a20ee51b340d6..f08fcf3748be4 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -6316,7 +6316,7 @@ }, "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND" : { "message" : [ - "Cannot find the native data source library , which is installed at on the driver. Install the same library on every node, at the same path or in the native library path." + "Cannot find the native data source library with the SHA-256 checksum , which is installed at on the driver. Install the same build of the library on every node, at the same path or in the native library path." ], "sqlState" : "42K02" }, diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index ff3aa2b5bd01e..8dcfa59694fd1 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -94,6 +94,37 @@ checks that each array has at least as many values as its parent refers to, but that the buffers of the arrays hold those values, because the C data interface does not carry the sizes of buffers: a library that exports a buffer that is too short makes Spark read past its end. +The `get_next` function of the stream that `read` exports may block, for example on network I/O, +but Spark cannot interrupt native code: when a task is killed, Spark stops reading only once +`get_next` returns. Use timeouts for blocking operations, or enable the task reaper with +`spark.task.reaper.enabled` and `spark.task.reaper.killTimeout`, which kills an executor whose +killed task does not end. + +### Compatibility + +The JVM binds the functions of a library by their names only, so the names and the signatures of +the methods of `NativeBridge` are the interface. Its version, 1, is the `abiVersion` that a +library returns and that the manifest of a package declares. Within a version, the signature and +the semantics of an existing function never change, and new functions are only added as optional +functions, so a library keeps working with later releases of Spark. New values may appear where +the Javadoc says so: for example, later releases may push down predicates with other names, which +a library must not report as evaluated completely. An incompatible change comes with a new +version, with new names or a new class. Spark supports version 1 at least in all the releases of +Spark 4 from 4.4.0. + +### Limitations + +Version 1 of the interface does not cover these features of Data Source V2 yet. They can be added +later as optional functions: + +- statistics: Spark estimates the size of a native scan with `spark.sql.defaultSizeInBytes`, so a + native scan is never the broadcast side of a join, unless the query has a broadcast hint; +- custom metrics, shown in the Spark UI; +- admission control for streaming reads, such as a maximum number of rows for each micro-batch, + and continuous streaming; +- the pushdown of aggregates, of the sort order and of the partitioning of the data; +- a function to cancel a read in progress. + ## Finding Native Data Sources When no Java or Python data source has the name that a query uses, Spark looks for a native data @@ -134,8 +165,12 @@ these directories. A library in any other directory is found with `LD_LIBRARY_PA `spark.driver.extraLibraryPath` and `spark.executor.extraLibraryPath`. The executors load the library from the same path as the driver, or else find it in their own -native library path, so install it on every node, for example in the container image of the -cluster. +native library path, which they get from their own environment, such as the `PATH` of the +executor processes. So install the library on every node, for example in the container image of +the cluster. An executor only loads the same build of the library as the driver, which it checks +with the SHA-256 checksum of the file, because the partitions and the state that the driver +serializes are only meant for that build: a query fails with +`NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND` on a node where a different build is installed. ### Packages diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index 0be056e8ca960..ed1133bdcb04a 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -20,6 +20,7 @@ import warnings from collections.abc import Sized from functools import cached_property, reduce +from pathlib import PureWindowsPath from threading import RLock from types import TracebackType from typing import ( @@ -41,6 +42,7 @@ no_type_check, overload, ) +from urllib.parse import urlparse from pyspark.conf import SparkConf from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError @@ -2297,10 +2299,15 @@ def addArtifacts( }, ) # Native data source packages are artifacts of the session in the JVM, whatever the - # flags, as with Spark Connect. + # flags, as with Spark Connect. The JVM resolves URIs, such as file:///tmp/a.sparkpkg, + # but resolves a relative path against its own working directory, so it gets an absolute + # path instead. A Windows path, such as C:/a.sparkpkg, is a path too. for p in path: if p.endswith(".sparkpkg"): - self._jsparkSession.addArtifact(os.path.abspath(p)) + if urlparse(p).scheme and not PureWindowsPath(p).drive: + self._jsparkSession.addArtifact(p) + else: + self._jsparkSession.addArtifact(os.path.abspath(p)) path = tuple(p for p in path if not p.endswith(".sparkpkg")) for p in path: normalized_path = os.path.abspath(p) diff --git a/python/pyspark/sql/tests/test_artifact.py b/python/pyspark/sql/tests/test_artifact.py index f681e60552e4f..5e65cf7f9866f 100644 --- a/python/pyspark/sql/tests/test_artifact.py +++ b/python/pyspark/sql/tests/test_artifact.py @@ -16,6 +16,7 @@ # import os import tempfile +from pathlib import Path from pyspark.errors import PySparkRuntimeError from pyspark.sql.functions import assert_true, lit, udf @@ -121,10 +122,26 @@ def new_file(name): # together with other files. spark.addArtifact(new_file("my_file_source.sparkpkg"), file=True) spark.addArtifacts(new_file("my_other_source.sparkpkg"), new_file("other.txt")) + # A URI, and a path relative to the working directory of Python, which is not the + # one of the JVM. + spark.addArtifact(Path(new_file("my_uri_source.sparkpkg")).as_uri()) + new_file("my_relative_source.sparkpkg") + cwd = os.getcwd() + os.chdir(d) + try: + spark.addArtifact("my_relative_source.sparkpkg") + finally: + os.chdir(cwd) packages = spark._jsparkSession.artifactManager().getNativeDataSourcePackages()._1() self.assertEqual( sorted(packages.apply(i).getName() for i in range(packages.size())), - ["my_file_source.sparkpkg", "my_other_source.sparkpkg", "my_source.sparkpkg"], + [ + "my_file_source.sparkpkg", + "my_other_source.sparkpkg", + "my_relative_source.sparkpkg", + "my_source.sparkpkg", + "my_uri_source.sparkpkg", + ], ) def test_add_zipped_package(self): diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala index 862d3c180981a..30dcebcd85474 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala @@ -840,10 +840,10 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE } def nativeDataSourceLibraryNotFoundError( - fileName: String, path: String): SparkRuntimeException = { + fileName: String, checksum: String, path: String): SparkRuntimeException = { new SparkRuntimeException( errorClass = "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND", - messageParameters = Map("fileName" -> fileName, "path" -> path)) + messageParameters = Map("fileName" -> fileName, "checksum" -> checksum, "path" -> path)) } def removedClassInSpark2Error(className: String, e: Throwable): SparkClassNotFoundException = { diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java index ed455937b5108..ad8b7fdb06e8f 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java @@ -57,23 +57,33 @@ * {@link #createStreamReader} and keeps it until the streaming query stops. For each * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then - * calls {@link #read} on the executors. Offsets are JSON strings defined by the library.
    • + * calls {@link #read} on the executors. Offsets are JSON strings defined by the library. When + * a streaming query restarts, from its checkpoint, Spark creates a new stream reader and may + * call {@link #streamPartitions} again for the last range that it planned before, so the + * partitions of a range must read the same data each time. *
  • Writes. On the driver, once it computed the data to write, right before it starts - * the tasks of the write (or of a micro-batch), Spark creates a writer with - * {@link #createWriter} or {@link #createStreamWriter}, and serializes its state with - * {@link #serializeWriter}. If computing the data fails, Spark creates no writer. On the - * executors, each task creates a data writer with {@link #createDataWriter}, calls - * {@link #write} for each batch of at most {@code spark.sql.execution.arrow.maxRecordsPerBatch} - * rows, which also ends once it holds {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, - * and calls {@link #commitDataWriter}, or {@link #abortDataWriter} if the task or the - * commit fails. Then the driver calls {@link #commit} with the messages of the tasks if they - * all succeeded, or {@link #abort} otherwise, including when {@link #commit} fails, and closes - * the writer. A streaming write commits - * or aborts each micro-batch (epoch). A micro-batch streaming query creates a new stream writer - * for each micro-batch, and Spark closes the previous stream writer of the query when the query - * creates a new one. Spark is not told when a streaming query stops, so it closes the last - * stream writer of the query once the writer is no longer referenced, or when the query is - * restarted.
    • + * the tasks of the write (or of a micro-batch), Spark creates a writer with {@link #createWriter} + * or {@link #createStreamWriter}, and serializes its state with {@link #serializeWriter}. If + * computing the data fails, Spark creates no writer. On the executors, each task creates a data + * writer with {@link #createDataWriter}, calls {@link #write} for each batch of at most + * {@code spark.sql.execution.arrow.maxRecordsPerBatch} rows, which also ends once it holds + * {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, and calls {@link #commitDataWriter}, + * or {@link #abortDataWriter} if the task or the commit fails. Then the driver calls + * {@link #commit} with the messages of the tasks if they all succeeded, or {@link #abort} + * otherwise, including when {@link #commit} fails, and closes the writer. Several attempts of a + * task may write the same partition, after a failure or with speculative execution, each with a + * data writer of its own. Spark lets the attempts of a partition commit one at a time, and passes + * to {@link #commit} the message of one committed attempt for each partition, so the write + * consists of the data of these attempts only. It aborts the data writers of the other attempts, + * except on an executor that is lost, where a data writer is neither committed nor aborted, and + * may even have committed before its message reached the driver. A streaming write commits or + * aborts each micro-batch (epoch). Spark may commit the same epoch more than once, for example + * when a streaming query restarts after a failure, so committing an epoch that is already + * committed must have no effect, for exactly-once output. A micro-batch streaming query creates a + * new stream writer for each micro-batch, and Spark closes the previous stream writer of the + * query when the query creates a new one. Spark is not told when a streaming query stops, so it + * closes the last stream writer of the query once the writer is no longer referenced, or when the + * query is restarted. *
* *

Conventions

@@ -107,8 +117,8 @@ * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is * {@code timestamp[us]} without one. Data read from a library may also use * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, and any names - * for the fields of the entries of a map, but not unsigned integers, {@code decimal256} or - * dictionary encoding. + * for the fields of the entries of a map, but not unsigned integers, {@code decimal256}, list + * views or dictionary encoding. *
  • Options. Option keys are passed in lower case, because Spark matches them * case-insensitively.
    • *
  • Optional methods. A library only needs to export the functions for the operations @@ -116,15 +126,36 @@ * the pushdown methods and the {@code serialize} methods, which then push nothing down and * return no state.
    • *
  • Threads. Spark calls the library from multiple threads, and calls on different - * handles may run concurrently. Calls on the same handle never run concurrently.
    • + * handles may run concurrently. Calls on the same handle never run concurrently. A handle may + * be released from a thread other than the threads that used it, such as a JVM + * {@code Cleaner} thread when Spark releases a handle once it is no longer referenced, for + * example the last stream writer of a streaming query. * * + *

    Compatibility

    + * The JVM binds the functions of a library by their names only, so the name of this class, the + * names of its methods and their signatures are the interface. Within an {@link #ABI_VERSION}: + *
      + *
    • the signature and the semantics of an existing method never change;
      • + *
    • new methods may be added, only as optional methods, so that a library that does not + * export them keeps working as before;
      • + *
    • new values may appear where this documentation says so, such as the names of pushed + * predicates.
      • + *
    + * An incompatible change uses a new {@link #ABI_VERSION}, with new names or a new class. A + * library reports the version it implements with {@link #abiVersion}, and Spark rejects a library + * of a version it does not support. Spark supports version 1 at least in all the releases of + * Spark 4 from 4.4.0. + * * @since 4.4.0 */ @Evolving public final class NativeBridge { - /** The version of this interface. It changes when the interface changes incompatibly. */ + /** + * The version of this interface. It changes when the interface changes incompatibly: see + * Compatibility above. + */ public static final int ABI_VERSION = 1; private NativeBridge() {} @@ -224,15 +255,21 @@ public static native long createStreamWriter( * for {@code string}, a base64 string for {@code binary}, a decimal string for * {@code decimal}, the days since the epoch for {@code date}, and the microseconds since the * epoch for {@code timestamp} and {@code timestamp_ntz}. - *
  • {@code {"type": "function", "name": ">", "children": [...]}}: a predicate or function, - * with the names of {@code org.apache.spark.sql.connector.expressions.filter.Predicate} and - * {@code org.apache.spark.sql.connector.expressions.GeneralScalarExpression}, such as - * {@code =}, {@code <}, {@code IN}, {@code IS_NULL}, {@code STARTS_WITH}, {@code AND}, - * {@code OR} and {@code NOT}.
    • + *
  • {@code {"type": "function", "name": ">", "children": [...]}}: a predicate, with the + * name it has in {@code org.apache.spark.sql.connector.expressions.filter.Predicate}. The + * names are {@code =}, {@code <>}, {@code <=>} (equal, with null equal to null), {@code <}, + * {@code <=}, {@code >} and {@code >=} with two children, {@code IN} with the value to look + * for and then the values to look in, {@code IS_NULL} and {@code IS_NOT_NULL} with one + * child, {@code STARTS_WITH}, {@code ENDS_WITH} and {@code CONTAINS} with the string and the + * string to look for, {@code AND} and {@code OR} with two predicates, and {@code NOT} with + * one. The children of the other functions are columns, literals or predicates.
    • * - * Spark does not pass predicates that it cannot express this way, nor predicates on columns - * that contain strings with a non-binary collation, such as {@code UTF8_LCASE}: strings are - * compared by their UTF-8 bytes. + * Spark does not pass predicates that it cannot express this way, such as predicates that use + * arithmetic or other functions, nor predicates on columns that contain strings with a + * non-binary collation, such as {@code UTF8_LCASE}: strings are compared by their UTF-8 bytes. + * A later version of Spark may pass predicates with other names without changing + * {@link #ABI_VERSION}, so a reader must not report a predicate as evaluated completely if it + * contains a name that the reader does not know, or a data type that it does not support. *

    * A predicate that the reader evaluates completely must follow the semantics of Spark SQL. In * particular, for floating-point columns, NaN is equal to NaN and greater than any other value, @@ -325,6 +362,13 @@ public static native long createStreamWriter( * of the stream is a struct array whose fields are the columns to read: * the columns of the schema to read, or the pruned columns. Spark calls * the stream from the thread of the task that reads the partition. + * {@code get_next} may block, for example on network I/O, but Spark + * cannot interrupt native code: when the task is killed, Spark only + * stops reading once {@code get_next} returns, and then releases the + * stream. So a {@code get_next} that blocks forever keeps the task + * running, unless the task reaper ({@code spark.task.reaper.enabled} + * and {@code spark.task.reaper.killTimeout}) kills the executor. Use + * timeouts for blocking operations. */ public static native void read(byte[] readerState, byte[] partition, long streamAddress); diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index a9f3dd06c6ada..e2307c70c3a42 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -440,9 +440,12 @@ class NativePartitionReader( throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + "is not supported.") } - // The Arrow types that convert to Spark types, but that ArrowColumnVector cannot read. + // The Arrow types that convert to Spark types, but that Spark does not read: ArrowColumnVector + // cannot read decimals of other widths than 128 bits, and the lengths of the elements of a + // list view cannot be checked without a pass over all its lists (see checkArrays). def unsupportedType(field: Field): Option[ArrowType] = field.getType match { case decimal: ArrowType.Decimal if decimal.getBitWidth != 128 => Some(decimal) + case listView: ArrowType.ListView => Some(listView) case _ => field.getChildren.asScala.view.flatMap(unsupportedType).headOption } root.getSchema.getFields.asScala.foreach { field => @@ -537,7 +540,7 @@ class NativePartitionReader( def referenced(array: ArrowArray.Snapshot, arrowType: ArrowType): Long = arrowType match { case _: ArrowType.Struct => array.length case _: ArrowType.List | _: ArrowType.Map => lastOffset(array) - // A list view has no last offset: its lists refer to ranges of its elements in any order. + // The other types that Spark reads have no children. List views are rejected before. case _ => 0 } def checkColumn(array: ArrowArray.Snapshot, field: Field, column: String): Unit = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala index ab613e92f873e..8db977cb1fc8e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -217,7 +217,7 @@ object NativeDataSourcePackage extends Logging { libraries.properties().asScala.map(e => e.getKey -> e.getValue.asText()).toMap)) } - private def sha256(file: File): String = { + private[ffi] def sha256(file: File): String = { val digest = MessageDigest.getInstance("SHA-256") Utils.tryWithResource(new FileInputStream(file)) { in => val buffer = new Array[Byte](64 * 1024) @@ -233,11 +233,37 @@ object NativeDataSourcePackage extends Logging { /** * A native data source library installed in the native library path, where - * [[NativeDataSourceRegistry]] finds it by its file name. It must be installed on every node. + * [[NativeDataSourceRegistry]] finds it by its file name. It must be installed on every node. The + * SHA-256 checksum of the file identifies the build of the library, so that an executor can check + * that it loads the same build as the driver: the state that the driver serializes is only + * meant for the same build. * * @param path the path of the library on the driver + * @param sha256 the SHA-256 checksum of the library on the driver */ -case class InstalledNativeLibrary(path: String) extends NativeLibraryLocation +case class InstalledNativeLibrary(path: String, sha256: String) extends NativeLibraryLocation + +object InstalledNativeLibrary { + private case class CacheKey(path: String, length: Long, lastModified: Long) + + // Bounded, like the cache of the packages, although few libraries are installed. + private val checksums: Cache[CacheKey, String] = + CacheBuilder.newBuilder().maximumSize(1000).build[CacheKey, String]() + + /** The installed library in the given file. */ + def apply(file: File): InstalledNativeLibrary = InstalledNativeLibrary(file.getPath, sha256(file)) + + /** The SHA-256 checksum of the given file. The result is cached until the file changes. */ + private[ffi] def sha256(file: File): String = { + val canonical = file.getCanonicalFile + val key = CacheKey(canonical.getPath, canonical.length(), canonical.lastModified()) + Option(checksums.getIfPresent(key)).getOrElse { + val checksum = NativeDataSourcePackage.sha256(canonical) + checksums.put(key, checksum) + checksum + } + } +} /** The platform names used in [[NativeDataSourceManifest.libraries]]. */ object NativePlatform { @@ -365,19 +391,17 @@ object NativeLibraries extends Logging { } /** - * Finds an installed library on this node: where the driver found it, or else in the native - * library path of this node. + * Finds an installed library on this node with the same checksum as on the driver: where the + * driver found it, or else in the native library path of this node. */ private def locate(installed: InstalledNativeLibrary): File = { val file = new File(installed.path) - if (file.isFile) { - file - } else { - NativeDataSourceRegistry.findInstalledLibrary(file.getName).getOrElse { + (file +: NativeDataSourceRegistry.findInstalledLibraries(file.getName)) + .find(f => f.isFile && InstalledNativeLibrary.sha256(f) == installed.sha256) + .getOrElse { throw QueryExecutionErrors.nativeDataSourceLibraryNotFoundError( - file.getName, installed.path) + file.getName, installed.sha256, installed.path) } - } } private def extract(packageFile: File, libraryPath: String, target: File): Unit = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala index b83366384affc..d5b09ce578a5b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -129,15 +129,18 @@ object NativeDataSourceRegistry extends Logging { return None } val dataSource = name.toLowerCase(Locale.ROOT) - findInstalledLibrary(libraryFileName(dataSource)) - .map(file => dataSource -> InstalledNativeLibrary(file.getPath)) + findInstalledLibraries(libraryFileName(dataSource)).headOption + .map(file => dataSource -> InstalledNativeLibrary(file)) } - /** Finds the library with the given file name in the native library path of this node. */ - private[ffi] def findInstalledLibrary(fileName: String): Option[File] = { + /** + * Finds the libraries with the given file name in the native library path of this node, in the + * order of the path. + */ + private[ffi] def findInstalledLibraries(fileName: String): Seq[File] = { librarySearchPath(System.getProperty("java.library.path"), System.getenv("PATH")) .map(new File(_, fileName)) - .find(_.isFile) + .filter(_.isFile) } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala index b7a6a43518ca2..5a1c5f5892a70 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala @@ -32,9 +32,19 @@ import org.apache.spark.sql.types._ object NativePredicates { private val mapper = new ObjectMapper() + /** + * The names of the predicates that are passed to the libraries, as listed in + * `NativeBridge.pushPredicates`. They are part of the interface, so a predicate with any other + * name, such as an arithmetic function, is not passed: its semantics, for example on overflow, + * would be part of the interface too, and the names of Data Source V2 grow with new releases. + */ + val FUNCTION_NAMES: Set[String] = Set( + "=", "<>", "<=>", "<", "<=", ">", ">=", "IN", "IS_NULL", "IS_NOT_NULL", + "STARTS_WITH", "ENDS_WITH", "CONTAINS", "AND", "OR", "NOT") + /** * Returns the JSON of the predicate on the columns of the given schema, or None if it cannot - * be expressed in JSON. + * be expressed in JSON with the functions in [[FUNCTION_NAMES]]. */ def toJson(predicate: Predicate, schema: StructType): Option[String] = { encode(predicate, schema).map(mapper.writeValueAsString) @@ -48,7 +58,7 @@ object NativePredicates { case literal: Literal[_] => encodeLiteral(literal.value(), literal.dataType()) - case function: GeneralScalarExpression => + case function: GeneralScalarExpression if FUNCTION_NAMES.contains(function.name()) => val children = function.children().map(encode(_, schema)) if (children.forall(_.isDefined)) { val node = mapper.createObjectNode() diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index 7217ab8e6523b..4aba24753fca4 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -22,12 +22,12 @@ // native_range Reads rows [0, end) with columns of several types. Supports predicate, // column and limit pushdown. With the option "extra", the schema has one more // column: "map" adds a map whose entry fields are named "keys" and "values", -// as arrow-rs names them, and "decimal256" adds a decimal256. With the option -// "slice", the columns of the batches have an offset of 1. With the option -// "short", an array of the batches has one value fewer than its parent refers -// to, although its buffers hold them all: "id" shortens the column id, -// "entries" the entries of the map and their fields, and "values" the values -// of the entries. +// as arrow-rs names them, "decimal256" adds a decimal256, and "listview" adds +// a list view. With the option "slice", the columns of the batches have an +// offset of 1. With the option "short", an array of the batches has one value +// fewer than its parent refers to, although its buffers hold them all: "id" +// shortens the column id, "entries" the entries of the map and their fields, +// and "values" the values of the entries. // native_sink Writes the rows as CSV files under the directory in the option "path". // native_counter Streams the rows [0, max_offset) with a single column "id", making "step" // more rows available each time Spark asks for the latest offset. @@ -147,7 +147,8 @@ const ColumnType kColumns[] = { // The columns that native_range adds with the option "extra", by its value. const std::map kExtraColumns = { {"map", {"m", "+m"}}, // map: {"k": id} - {"decimal256", {"dec", "d:10,2,256"}} // decimal256(10, 2): id + {"decimal256", {"dec", "d:10,2,256"}}, // decimal256(10, 2): id + {"listview", {"lv", "+vl"}} // list_view: [id] }; const ColumnType* findColumn(const std::string& name) { @@ -228,6 +229,8 @@ void exportSchema(const std::vector& columns, ArrowSchema* out) { ArrowSchema* entries = addChild(column, "+s", "entries", 0); addChild(entries, "u", "keys", 0); addChild(entries, "l", "values", ARROW_FLAG_NULLABLE); + } else if (name == "lv") { + addChild(column, "l", "item", ARROW_FLAG_NULLABLE); } } } @@ -375,6 +378,18 @@ void exportColumn(const std::string& name, int64_t lo, int64_t hi, ArrowArray* o words.insert(words.end(), {unscaled, sign, sign, sign}); } initArray(out, n, 0, {no_validity, toBytes(words)}); + } else if (name == "lv") { + // A list of a single element for each row. + std::vector offsets; + std::vector sizes; + std::vector values; + for (int64_t id = lo; id < hi; id++) { + offsets.push_back(static_cast(id - lo)); + sizes.push_back(1); + values.push_back(id); + } + initArray(out, n, 0, {no_validity, toBytes(offsets), toBytes(sizes)}); + initArray(addChild(out), n, 0, {no_validity, toBytes(values)}); } else { throw std::runtime_error("unknown column " + name); } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala new file mode 100644 index 0000000000000..d261654f7cede --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala @@ -0,0 +1,147 @@ +/* + * 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.execution.datasources.v2.ffi + +import java.io.File +import java.nio.file.Files + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.classic.{DataFrame, SparkSession} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.tags.ExtendedSQLTest +import org.apache.spark.util.Utils + +/** + * Tests native data sources in local-cluster mode, where the executors run in JVMs of their own: + * Spark distributes the packages to them, and they find, check and load the libraries + * themselves. [[NativeDataSourceSuite]] tests the rest in local mode, where the executors use the + * packages of the driver. + */ +@ExtendedSQLTest +class NativeDataSourceClusterSuite extends QueryTest { + import NativeDataSourceTestUtils._ + + protected var spark: SparkSession = _ + + // Jobs outside of SQL executions get the files of all the sessions, so the sessions are kept + // until the end of the suite, as in NativeDataSourceSuite. + private val sessions = ArrayBuffer.empty[SparkSession] + + private lazy val library = compile("test_native_datasource") + // Another build, so that its package has another checksum, which the executors did not load. + private lazy val otherLibrary = + compile("test_native_datasource_cluster", "-DNAME_PREFIX=\"cluster_\"") + + override def beforeAll(): Unit = { + super.beforeAll() + spark = SparkSession.builder() + .master("local-cluster[2,1,1024]") + .appName("NativeDataSourceClusterSuite") + .getOrCreate() + } + + override def afterAll(): Unit = { + try { + sessions.clear() + if (spark != null) { + spark.stop() + spark = null + } + } finally { + super.afterAll() + } + } + + private def nativeTest(name: String)(f: => Unit): Unit = test(name) { + assume(toolchain.isDefined, "A C++ compiler and the JNI headers are required.") + try f finally SparkSession.setActiveSession(spark) + } + + // Like Python data sources, native data sources are found in the active session. + private def newSession(): SparkSession = { + val session = spark.newSession() + SparkSession.setActiveSession(session) + sessions += session + session + } + + private def checkRows(df: => DataFrame, expected: Seq[Row]): Unit = { + QueryTest.checkAnswer(df, expected, checkToRDD = false) + } + + private def written(dir: File): Seq[String] = { + dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala).toSeq.sorted + } + + nativeTest("read and write with a package added as an artifact") { + val session = newSession() + session.addArtifact(createPackage(library, dataSources("")).getPath) + // 4 partitions, so that both executors load the library. + val df = session.read.format("native_range") + .option("end", "20").option("partitions", "4").load().select("id") + checkRows(df, (0L until 20L).map(Row(_))) + + val dir = Utils.createTempDir() + session.range(10).repartition(4).selectExpr("id", "concat('v', id) AS name") + .write.format("native_sink").option("path", dir.getPath).mode("append").save() + assert(written(dir) == (0 until 10).map(i => s"$i,v$i").sorted) + + val streamDir = Utils.createTempDir() + val query = session.readStream.format("native_counter") + .option("max_offset", "12").option("step", "5").load() + .writeStream.format("native_sink") + .option("path", streamDir.getPath) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try { + query.processAllAvailable() + } finally { + query.stop() + } + assert(written(streamDir) == (0 until 12).map(_.toString).sorted) + } + + nativeTest("executors load the distributed copy of a configured package") { + val session = newSession() + val pkg = createPackage(otherLibrary, dataSources("cluster_")) + session.conf.set(SQLConf.NATIVE_DATA_SOURCE_PATHS.key, pkg.getPath) + // Prepares the package on the driver as a query would: Spark reads its manifest, loads its + // library and copies it to the artifacts of the session under a name with its checksum. + val found = NativeDataSourceRegistry.lookup("cluster_native_range", session.sessionState.conf) + .map(_._2).get.asInstanceOf[NativeDataSourcePackage] + NativeLibraries.get(found) + val distributed = + NativeDataSourceRegistry.distribute(found).asInstanceOf[NativeDataSourcePackage] + assert(distributed.artifactUUID.contains(session.sessionUUID)) + assert(distributed.distributedFileName.contains(s"${pkg.getName}.${found.sha256}")) + + // Then the package at the path of the driver becomes unusable for the executors, while the + // driver keeps the package it read, since the file keeps its length and modification time. + // So the query only succeeds if the executors find the distributed copy, and check that it + // is the same package by its checksum. + val lastModified = pkg.lastModified() + Files.write(pkg.toPath, new Array[Byte](pkg.length().toInt)) + assert(pkg.setLastModified(lastModified)) + val df = session.read.format("cluster_native_range") + .option("end", "20").option("partitions", "4").load().select("id") + checkRows(df, (0L until 20L).map(Row(_))) + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 07db1bb8c47cf..694c4d777db68 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -177,13 +177,21 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { """{"type":"literal","dataType":"bigint","value":5}]}""")) assert(toJson(new And( predicate("IS_NULL", FieldReference(Seq("a", "b"))), - predicate("=", id, new GeneralScalarExpression("+", Array(id, LiteralValue(1, IntegerType)))) + predicate("IN", id, LiteralValue(1L, LongType), LiteralValue(2L, LongType)) )).contains( """{"type":"function","name":"AND","children":[{"type":"function","name":"IS_NULL",""" + - """"children":[{"type":"column","name":["a","b"]}]},{"type":"function","name":"=",""" + - """"children":[{"type":"column","name":["id"]},{"type":"function","name":"+",""" + + """"children":[{"type":"column","name":["a","b"]}]},{"type":"function","name":"IN",""" + """"children":[{"type":"column","name":["id"]},""" + - """{"type":"literal","dataType":"int","value":1}]}]}]}""")) + """{"type":"literal","dataType":"bigint","value":1},""" + + """{"type":"literal","dataType":"bigint","value":2}]}]}""")) + // Only the functions listed in NativeBridge.pushPredicates are passed, so not arithmetic. + NativePredicates.FUNCTION_NAMES.foreach { name => + assert(toJson(predicate(name, id, id)).nonEmpty, name) + } + val plusOne = new GeneralScalarExpression("+", Array(id, LiteralValue(1, IntegerType))) + assert(toJson(predicate("=", id, plusOne)).isEmpty) + assert(toJson(new And(predicate("IS_NULL", id), predicate("=", id, plusOne))).isEmpty) + assert(toJson(predicate("UPPER", FieldReference(Seq("s")))).isEmpty) def literal(value: Any, dataType: DataType): Option[String] = { toJson(predicate("=", id, LiteralValue(value, dataType))) @@ -360,6 +368,10 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } check("The column dec contains the Arrow type Decimal(10, 2, 256), which is not supported.", "extra" -> "decimal256") + // Arrow Java would also extend the elements of a list view with nulls, and Spark does not + // check their lengths. + check("The column lv contains the Arrow type ListView, which is not supported.", + "extra" -> "listview") } nativeTest("batch write") { @@ -819,18 +831,34 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { // An executor finds the library in its native library path, if it is not where the driver // found it. + val checksum = InstalledNativeLibrary(installed).sha256 + assert(checksum == NativeDataSourcePackage.sha256(installed)) val elsewhere = new File(Utils.createTempDir(), installed.getName).getPath - assert(NativeLibraries.get(InstalledNativeLibrary(elsewhere)) eq - NativeLibraries.get(InstalledNativeLibrary(installed.getPath))) + assert(NativeLibraries.get(InstalledNativeLibrary(elsewhere, checksum)) eq + NativeLibraries.get(InstalledNativeLibrary(installed))) + + // An executor only loads the same build as the driver, which serialized the state for its + // build: neither the library at the same path nor the one in its native library path is + // loaded if the driver found another build. + val otherBuild = NativeDataSourcePackage.sha256(installedNegatingLibrary) + Seq(installed.getPath, elsewhere).foreach { path => + checkError( + exception = intercept[SparkRuntimeException] { + NativeLibraries.get(InstalledNativeLibrary(path, otherBuild)) + }, + condition = "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND", + parameters = + Map("fileName" -> installed.getName, "checksum" -> otherBuild, "path" -> path)) + } } val missing = new File(Utils.createTempDir(), installed.getName).getPath checkError( exception = intercept[SparkRuntimeException] { - NativeLibraries.get(InstalledNativeLibrary(missing)) + NativeLibraries.get(InstalledNativeLibrary(missing, "0")) }, condition = "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND", - parameters = Map("fileName" -> installed.getName, "path" -> missing)) + parameters = Map("fileName" -> installed.getName, "checksum" -> "0", "path" -> missing)) checkError( exception = intercept[SparkClassNotFoundException] { newSession().read.format("installed_native_range").load() From b9ad01ac93beb9aaaed9c25a3ad888d7a0e475ba Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Tue, 13 Oct 2026 15:29:46 +0900 Subject: [PATCH 14/27] Call native data sources through the FFM API instead of JNI Replace the JNI interface, NativeBridge, with a plain C interface, spark_datasource.h, which Spark calls through the Foreign Function and Memory API. Native data sources then require Java 22 or later; with an earlier version, a query that uses one fails with NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION. Co-authored-by: Isaac --- .../resources/error/error-conditions.json | 6 + docs/sql-native-data-sources.md | 391 +++++------ .../sql/errors/QueryExecutionErrors.scala | 7 + .../CheckConnectJvmClientCompatibility.scala | 1 - .../NativeDataSourceConnectSuite.scala | 4 +- .../spark/sql/datasource/NativeBridge.java | 442 ------------ .../datasources/v2/ffi/NativeLibrary.java | 648 ++++++++++++++---- .../spark/sql/datasource/spark_datasource.h | 505 ++++++++++++++ .../datasources/v2/ffi/NativeDataSource.scala | 10 +- .../v2/ffi/NativeDataSourcePackage.scala | 28 +- .../v2/ffi/NativeDataSourceRegistry.scala | 5 +- .../v2/ffi/NativeDataSourceV2.scala | 2 +- .../datasources/v2/ffi/NativePredicates.scala | 4 +- .../test_native_datasource.cc | 409 ++++------- .../v2/ffi/NativeDataSourceClusterSuite.scala | 2 +- .../v2/ffi/NativeDataSourceSuite.scala | 35 +- .../v2/ffi/NativeDataSourceTestUtils.scala | 22 +- 17 files changed, 1444 insertions(+), 1077 deletions(-) delete mode 100644 sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java create mode 100644 sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index f08fcf3748be4..8ec5072e34e6b 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -6332,6 +6332,12 @@ ], "sqlState" : "42K02" }, + "NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION" : { + "message" : [ + "Native data sources require Java or later, but Spark runs on Java ." + ], + "sqlState" : "0A000" + }, "NEAREST_BY_JOIN" : { "message" : [ "Invalid nearest-by join." diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 8dcfa59694fd1..4276305a92a1b 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -25,10 +25,10 @@ license: | ## Overview A native data source is a data source implemented in native code, such as Rust or C++, without -any JVM code. It is a shared library that implements the -[`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) binary interface. -Spark finds the library automatically when it is installed like other native libraries, or when a -session adds it in a package, and loads it when a query uses the data source. Spark adapts it to +any JVM code. It is a shared library that exports the C functions declared in the header +[`spark_datasource.h`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h). Spark finds the library +automatically when it is installed like other native libraries, or when a session adds it in a +package, and loads it when a query uses the data source. Spark adapts it to [Data Source V2](sql-v2-data-sources.html), so it supports: - batch and micro-batch streaming reads, @@ -39,53 +39,64 @@ Data is exchanged only as columnar batches, through the [Arrow C data interface](https://arrow.apache.org/docs/format/CDataInterface.html), so Spark reads the data that a library returns without copying it. +Spark calls native data sources with the +[Foreign Function and Memory API](https://openjdk.org/jeps/454) of Java, which is final from Java +22, so they require Java 22 or later on the driver and the executors. With an earlier version of +Java, a query that uses a native data source fails with +`NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION`. + ## How It Works -A native data source library implements the `native` methods of `NativeBridge` as JNI functions, -named after the class, such as `Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource`. -Spark loads each library into its own copy of the class, so any number of native data sources can -be loaded side by side. +A native data source library exports the functions of `spark_datasource.h`, such as +`spark_ds_create`, as plain C functions: it needs no JVM headers or bindings. Spark looks the +functions up in each library separately, so any number of native data sources can be loaded side +by side, although they export the same functions. Spark calls the functions of the library as follows: - **Data sources**: for each schema inference, scan and write, Spark creates a data source on the - driver with `createDataSource`, which receives the options of the user, and closes it right - after. When the user does not specify a schema, Spark asks for it with `schema`. -- **Batch reads**: on the driver, Spark creates a reader with `createReader`, pushes operations down - with `pushPredicates`, `pushLimit` and `pruneColumns`, plans the partitions with `partitions`, - and serializes the state of the reader with `serializeReader`. Then it calls `read` on the - executors, in one task for each partition, and `read` exports an `ArrowArrayStream`. -- **Streaming reads**: on the driver, Spark creates a stream reader with `createStreamReader`. For - each micro-batch, it asks for the `latestOffset`, plans the partitions with `streamPartitions`, - and calls `read` on the executors. Offsets are JSON strings defined by the library. -- **Writes**: on the driver, Spark creates a writer with `createWriter` or `createStreamWriter`. On - the executors, each task creates a data writer with `createDataWriter`, writes the batches of its - data with `write`, and calls `commitDataWriter`, or `abortDataWriter` if the task or the commit - fails. Then the driver calls `commit` with the messages of the tasks, or `abort` if the write - failed. A micro-batch streaming query creates a new stream writer for each micro-batch, and Spark - closes the previous one when it creates the next one. Because Spark is not told when a streaming - query stops, the last stream writer of a query is closed once it is no longer referenced, or - when the query is restarted. + driver with `spark_ds_create`, which receives the options of the user, and closes it right after + with `spark_ds_close`. When the user does not specify a schema, Spark asks for it with + `spark_ds_schema`. +- **Batch reads**: on the driver, Spark creates a reader with `spark_ds_create_reader`, pushes + operations down with `spark_ds_push_predicates`, `spark_ds_push_limit` and + `spark_ds_prune_columns`, plans the partitions with `spark_ds_partitions`, and serializes the + state of the reader with `spark_ds_serialize_reader`. Then it calls `spark_ds_read` on the + executors, in one task for each partition, and `spark_ds_read` exports an `ArrowArrayStream`. +- **Streaming reads**: on the driver, Spark creates a stream reader with + `spark_ds_create_stream_reader`. For each micro-batch, it asks for the `spark_ds_latest_offset`, + plans the partitions with `spark_ds_stream_partitions`, and calls `spark_ds_read` on the + executors. Offsets are JSON strings defined by the library. +- **Writes**: on the driver, Spark creates a writer with `spark_ds_create_writer` or + `spark_ds_create_stream_writer`. On the executors, each task creates a data writer with + `spark_ds_create_data_writer`, writes the batches of its data with `spark_ds_write`, and calls + `spark_ds_commit_data_writer`, or `spark_ds_abort_data_writer` if the task or the commit fails. + Then the driver calls `spark_ds_commit` with the messages of the tasks, or `spark_ds_abort` if + the write failed. A micro-batch streaming query creates a new stream writer for each + micro-batch, and Spark closes the previous one when it creates the next one. Because Spark is not + told when a streaming query stops, the last stream writer of a query is closed once it is no + longer referenced, or when the query is restarted. The library defines the format of the partitions, of the state of its readers and writers, and of -the commit messages, which Spark sends to the executors and back as byte arrays. A library reports -an error by throwing a Java exception, for example with the JNI function `ThrowNew`, and Spark -reports it as a `NATIVE_DATA_SOURCE_ERROR`. +the commit messages, which Spark sends to the executors and back as byte arrays. The library +returns them, and the offsets, by passing them to a `SparkDsSink` callback, which copies them. A +function reports an error by returning a non-zero value, with the message returned by +`spark_ds_last_error`, and Spark reports it as a `NATIVE_DATA_SOURCE_ERROR`. A C++ exception or a +Rust panic must not cross the interface. A library only exports the functions of the operations it supports. If a function is missing, -Spark treats the operation as unsupported, except for the pushdown and `serialize` functions, which -are optional: +Spark treats the operation as unsupported, except for the optional functions: | Operation | Required functions | Optional functions | |-----------|--------------------|--------------------| -| All | `abiVersion`, `createDataSource`, `closeDataSource` | `schema` | -| Batch read | `createReader`, `partitions`, `read`, `closeReader` | `pushPredicates`, `pushLimit`, `pruneColumns`, `serializeReader` | -| Streaming read | `createStreamReader`, `initialOffset`, `latestOffset`, `streamPartitions`, `read`, `closeStreamReader` | `serializeStreamReader`, `commitOffset` | -| Batch write | `createWriter`, `commit`, `abort`, `closeWriter`, `createDataWriter`, `write`, `commitDataWriter`, `abortDataWriter` | `serializeWriter` | -| Streaming write | `createStreamWriter`, `commit`, `abort`, `closeWriter`, `createDataWriter`, `write`, `commitDataWriter`, `abortDataWriter` | `serializeWriter` | - -The [Javadoc of `NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) -specifies each function, including the ownership of the Arrow structs and the JSON format of the +| All | `spark_ds_abi_version`, `spark_ds_create`, `spark_ds_close` | `spark_ds_last_error`, `spark_ds_schema` | +| Batch read | `spark_ds_create_reader`, `spark_ds_partitions`, `spark_ds_read`, `spark_ds_close_reader` | `spark_ds_push_predicates`, `spark_ds_push_limit`, `spark_ds_prune_columns`, `spark_ds_serialize_reader` | +| Streaming read | `spark_ds_create_stream_reader`, `spark_ds_initial_offset`, `spark_ds_latest_offset`, `spark_ds_stream_partitions`, `spark_ds_read`, `spark_ds_close_stream_reader` | `spark_ds_serialize_stream_reader`, `spark_ds_commit_offset` | +| Batch write | `spark_ds_create_writer`, `spark_ds_commit`, `spark_ds_abort`, `spark_ds_close_writer`, `spark_ds_create_data_writer`, `spark_ds_write`, `spark_ds_commit_data_writer`, `spark_ds_abort_data_writer` | `spark_ds_serialize_writer` | +| Streaming write | `spark_ds_create_stream_writer`, `spark_ds_commit`, `spark_ds_abort`, `spark_ds_close_writer`, `spark_ds_create_data_writer`, `spark_ds_write`, `spark_ds_commit_data_writer`, `spark_ds_abort_data_writer` | `spark_ds_serialize_writer` | + +The comments of [`spark_datasource.h`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h) +specify each function, including the ownership of the Arrow structs and the JSON format of the pushed predicates. A predicate that a library reports as evaluated completely must follow the semantics of Spark SQL, for example for NaN values, which are equal to each other and greater than any other value. The arrays that a library exports must have an offset of 0: copy a sliced @@ -94,23 +105,23 @@ checks that each array has at least as many values as its parent refers to, but that the buffers of the arrays hold those values, because the C data interface does not carry the sizes of buffers: a library that exports a buffer that is too short makes Spark read past its end. -The `get_next` function of the stream that `read` exports may block, for example on network I/O, -but Spark cannot interrupt native code: when a task is killed, Spark stops reading only once -`get_next` returns. Use timeouts for blocking operations, or enable the task reaper with +The `get_next` function of the stream that `spark_ds_read` exports may block, for example on +network I/O, but Spark cannot interrupt native code: when a task is killed, Spark stops reading +only once `get_next` returns. Use timeouts for blocking operations, or enable the task reaper with `spark.task.reaper.enabled` and `spark.task.reaper.killTimeout`, which kills an executor whose killed task does not end. ### Compatibility -The JVM binds the functions of a library by their names only, so the names and the signatures of -the methods of `NativeBridge` are the interface. Its version, 1, is the `abiVersion` that a -library returns and that the manifest of a package declares. Within a version, the signature and -the semantics of an existing function never change, and new functions are only added as optional -functions, so a library keeps working with later releases of Spark. New values may appear where -the Javadoc says so: for example, later releases may push down predicates with other names, which -a library must not report as evaluated completely. An incompatible change comes with a new -version, with new names or a new class. Spark supports version 1 at least in all the releases of -Spark 4 from 4.4.0. +Spark binds the functions of a library by their names, so the names and the signatures of the +functions of `spark_datasource.h` are the interface. Its version, 1, is `SPARK_DS_ABI_VERSION`, +which `spark_ds_abi_version` returns and the manifest of a package declares. Within a version, the +signature and the semantics of an existing function never change, and new functions are only +added as optional functions, so a library keeps working with later releases of Spark. New values +may appear where the header says so: for example, later releases may push down predicates with +other names, which a library must not report as evaluated completely. An incompatible change comes +with a new version, with new names. Spark supports version 1 at least in all the releases of Spark +4 from 4.4.0. ### Limitations @@ -139,7 +150,7 @@ like other native libraries is found automatically, without any configuration. T named after the data source, in lower case, with the prefix `spark_datasource_`: the library of `rust_range` is `libspark_datasource_rust_range.so` on Linux, `libspark_datasource_rust_range.dylib` on macOS and `spark_datasource_rust_range.dll` on Windows. -Spark passes the name of the data source to `createDataSource`. A library that implements several +Spark passes the name of the data source to `spark_ds_create`. A library that implements several data sources is installed under each of their names, for example with symbolic links. Spark looks for the library in the native library path, which consists of, in order: @@ -190,9 +201,9 @@ is a zip file with the extension `.sparkpkg`, which contains a manifest named } ``` -- `abiVersion` is the version of `NativeBridge` that the library implements, currently 1. +- `abiVersion` is the version of `spark_datasource.h` that the library implements, currently 1. - `dataSources` lists the names of the data sources that the library implements. The names are - case-insensitive, and Spark passes the name to `createDataSource`. + case-insensitive, and Spark passes the name to `spark_ds_create`. - `libraries` maps each platform, `-` with an `os` of `linux`, `osx` or `windows` and an `arch` of `x86_64` or `aarch64`, to the path of the library in the package. @@ -218,11 +229,11 @@ installed library provides it. ## Example in Rust -This data source reads the numbers `[0, end)` as a column `id`, in two partitions. It uses the -[`jni`](https://crates.io/crates/jni) crate, and the `arrow-array` and `arrow-schema` crates of -[arrow-rs](https://github.com/apache/arrow-rs), which export schemas and streams with -`FFI_ArrowSchema` and `FFI_ArrowArrayStream`. `Cargo.toml`, where the name of the library is the -name that Spark finds when it is installed: +This data source reads the numbers `[0, end)` as a column `id`, in two partitions. It only depends +on the `arrow-array` and `arrow-schema` crates of [arrow-rs](https://github.com/apache/arrow-rs), +which export schemas and streams with `FFI_ArrowSchema` and `FFI_ArrowArrayStream`, and declares +the functions of `spark_datasource.h` that it implements as `extern "C"` functions. `Cargo.toml`, +where the name of the library is the name that Spark finds when it is installed: ```toml [package] @@ -237,7 +248,6 @@ crate-type = ["cdylib"] [dependencies] arrow-array = { version = "58", features = ["ffi"] } arrow-schema = { version = "58", features = ["ffi"] } -jni = "0.21" ``` `src/lib.rs`: @@ -247,6 +257,8 @@ jni = "0.21" //! A native data source for Apache Spark, written in Rust. It reads the numbers [0, end) as a //! column `id`, in two partitions. The option `end` defaults to 10. +use std::cell::RefCell; +use std::ffi::{c_char, c_int, c_void, CStr, CString}; use std::panic::{catch_unwind, AssertUnwindSafe}; use std::sync::Arc; @@ -254,12 +266,12 @@ use arrow_array::ffi_stream::FFI_ArrowArrayStream; use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator}; use arrow_schema::ffi::FFI_ArrowSchema; use arrow_schema::{DataType, Field, Schema, SchemaRef}; -use jni::objects::{JByteArray, JClass, JObject, JObjectArray, JString}; -use jni::sys::{jint, jlong, jobjectArray}; -use jni::JNIEnv; type Result = std::result::Result>; +/// `SparkDsSink`: receives a byte array returned to Spark, and copies it. +type Sink = extern "C" fn(sink: *mut c_void, data: *const u8, length: i64); + struct RangeSource { end: i64, } @@ -268,140 +280,118 @@ struct RangeReader { end: i64, } +thread_local! { + static LAST_ERROR: RefCell = RefCell::new(CString::default()); +} + fn schema() -> SchemaRef { Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])) } -/// Runs `f`, and reports its error or panic as a Java exception: neither may cross into the JVM. -fn call(env: &mut JNIEnv, default: T, f: impl FnOnce(&mut JNIEnv) -> Result) -> T { - let message = match catch_unwind(AssertUnwindSafe(|| f(env))) { - Ok(Ok(value)) => return value, +/// Runs `f`, and reports its error or panic to Spark: neither may cross into the JVM. +fn call(f: impl FnOnce() -> Result<()>) -> c_int { + let message = match catch_unwind(AssertUnwindSafe(f)) { + Ok(Ok(())) => return 0, Ok(Err(e)) => e.to_string(), Err(_) => "the data source panicked".to_string(), }; - if !env.exception_check().unwrap_or(true) { - let _ = env.throw_new("java/lang/RuntimeException", message); - } - default + LAST_ERROR.with(|error| *error.borrow_mut() = CString::new(message).unwrap_or_default()); + 1 } -fn strings(env: &mut JNIEnv, array: &JObjectArray) -> Result> { - let mut result = Vec::new(); - for i in 0..env.get_array_length(array)? { - let value = JString::from(env.get_object_array_element(array, i)?); - result.push(env.get_string(&value)?.into()); - } - Ok(result) +#[no_mangle] +pub extern "C" fn spark_ds_last_error() -> *const c_char { + LAST_ERROR.with(|error| error.borrow().as_ptr()) } #[no_mangle] -pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_abiVersion( - _env: JNIEnv, - _class: JClass, -) -> jint { +pub extern "C" fn spark_ds_abi_version() -> i32 { 1 } #[no_mangle] -pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource( - mut env: JNIEnv, - _class: JClass, - _name: JString, - keys: JObjectArray, - values: JObjectArray, -) -> jlong { - call(&mut env, 0, |env| { - let mut source = RangeSource { end: 10 }; - for (key, value) in strings(env, &keys)?.into_iter().zip(strings(env, &values)?) { - if key == "end" { - source.end = value.parse()?; +pub unsafe extern "C" fn spark_ds_create( + _name: *const c_char, + option_keys: *const *const c_char, + option_values: *const *const c_char, + num_options: i64, + source: *mut *mut c_void, +) -> c_int { + call(|| { + let mut range = RangeSource { end: 10 }; + for i in 0..num_options as usize { + if CStr::from_ptr(*option_keys.add(i)).to_str()? == "end" { + range.end = CStr::from_ptr(*option_values.add(i)).to_str()?.parse()?; } } - Ok(Box::into_raw(Box::new(source)) as jlong) + *source = Box::into_raw(Box::new(range)) as *mut c_void; + Ok(()) }) } #[no_mangle] -pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_schema( - mut env: JNIEnv, - _class: JClass, - _source: jlong, - schema_address: jlong, -) { - call(&mut env, (), |_| { - let schema = FFI_ArrowSchema::try_from(schema().as_ref())?; - unsafe { std::ptr::write(schema_address as *mut FFI_ArrowSchema, schema) }; +pub unsafe extern "C" fn spark_ds_schema(_source: *mut c_void, out: *mut FFI_ArrowSchema) -> c_int { + call(|| { + std::ptr::write(out, FFI_ArrowSchema::try_from(schema().as_ref())?); Ok(()) }) } #[no_mangle] -pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_createReader( - mut env: JNIEnv, - _class: JClass, - source: jlong, - _schema_address: jlong, -) -> jlong { - call(&mut env, 0, |_| { - let source = unsafe { &*(source as *const RangeSource) }; - Ok(Box::into_raw(Box::new(RangeReader { end: source.end })) as jlong) +pub unsafe extern "C" fn spark_ds_create_reader( + source: *mut c_void, + _schema: *mut FFI_ArrowSchema, + reader: *mut *mut c_void, +) -> c_int { + call(|| { + let end = (*(source as *const RangeSource)).end; + *reader = Box::into_raw(Box::new(RangeReader { end })) as *mut c_void; + Ok(()) }) } #[no_mangle] -pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_closeDataSource( - _env: JNIEnv, - _class: JClass, - source: jlong, -) { - drop(unsafe { Box::from_raw(source as *mut RangeSource) }); +pub unsafe extern "C" fn spark_ds_close(source: *mut c_void) { + drop(Box::from_raw(source as *mut RangeSource)); } #[no_mangle] -pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_partitions( - mut env: JNIEnv, - _class: JClass, - reader: jlong, -) -> jobjectArray { - call(&mut env, std::ptr::null_mut(), |env| { - let end = unsafe { &*(reader as *const RangeReader) }.end; - let bounds = [(0, end / 2), (end / 2, end)]; - let partitions = env.new_object_array(bounds.len() as i32, "[B", JObject::null())?; - for (i, (start, end)) in bounds.into_iter().enumerate() { +pub unsafe extern "C" fn spark_ds_partitions( + reader: *mut c_void, + add_partition: Sink, + sink: *mut c_void, +) -> c_int { + call(|| { + let end = (*(reader as *const RangeReader)).end; + for (start, end) in [(0, end / 2), (end / 2, end)] { let bytes = [start.to_le_bytes(), end.to_le_bytes()].concat(); - let partition = env.byte_array_from_slice(&bytes)?; - env.set_object_array_element(&partitions, i as i32, partition)?; + add_partition(sink, bytes.as_ptr(), bytes.len() as i64); } - Ok(partitions.into_raw()) + Ok(()) }) } #[no_mangle] -pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_closeReader( - _env: JNIEnv, - _class: JClass, - reader: jlong, -) { - drop(unsafe { Box::from_raw(reader as *mut RangeReader) }); +pub unsafe extern "C" fn spark_ds_close_reader(reader: *mut c_void) { + drop(Box::from_raw(reader as *mut RangeReader)); } #[no_mangle] -pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_read( - mut env: JNIEnv, - _class: JClass, - _reader_state: JByteArray, - partition: JByteArray, - stream_address: jlong, -) { - call(&mut env, (), |env| { - let bytes = env.convert_byte_array(&partition)?; +pub unsafe extern "C" fn spark_ds_read( + _reader_state: *const u8, + _reader_state_length: i64, + partition: *const u8, + partition_length: i64, + out: *mut FFI_ArrowArrayStream, +) -> c_int { + call(|| { + let bytes = std::slice::from_raw_parts(partition, partition_length as usize); let start = i64::from_le_bytes(bytes[0..8].try_into()?); let end = i64::from_le_bytes(bytes[8..16].try_into()?); let ids = Int64Array::from_iter_values(start..end); let batch = RecordBatch::try_new(schema(), vec![Arc::new(ids)])?; let batches = RecordBatchIterator::new(vec![Ok(batch)], schema()); - let stream = FFI_ArrowArrayStream::new(Box::new(batches)); - unsafe { std::ptr::write(stream_address as *mut FFI_ArrowArrayStream, stream) }; + std::ptr::write(out, FFI_ArrowArrayStream::new(Box::new(batches))); Ok(()) }) } @@ -439,26 +429,25 @@ EOF ## Example in C++ The same data source in C++, with [Arrow C++](https://arrow.apache.org/docs/cpp/), which exports -schemas and streams with `arrow::ExportSchema` and `arrow::ExportRecordBatchReader`: +schemas and streams with `arrow::ExportSchema` and `arrow::ExportRecordBatchReader`. It includes +`spark_datasource.h`, which is also in the `spark-sql` JAR, under +`org/apache/spark/sql/datasource/`: {% raw %} ```cpp // A native data source for Apache Spark, written in C++ with Arrow C++. It reads the numbers // [0, end) as a column `id`, in two partitions. The option `end` defaults to 10. -#include - #include +#include #include #include #include -#include #include #include -#define SPARK_JNI(ret, method) \ - extern "C" JNIEXPORT ret JNICALL Java_org_apache_spark_sql_datasource_NativeBridge_##method +#include "spark_datasource.h" namespace { @@ -470,18 +459,21 @@ struct RangeReader { int64_t end; }; +thread_local std::string last_error; + std::shared_ptr Schema() { return arrow::schema({arrow::field("id", arrow::int64(), /*nullable=*/false)}); } -// Runs `f`, and reports its error as a Java exception: no C++ exception may cross into the JVM. +// Runs `f`, and reports its error to Spark: no C++ exception may cross into the JVM. template -auto Call(JNIEnv* env, F&& f) -> decltype(f()) { +int Call(F&& f) { try { - return f(); + f(); + return 0; } catch (const std::exception& e) { - env->ThrowNew(env->FindClass("java/lang/RuntimeException"), e.what()); - if constexpr (!std::is_void_v) return {}; + last_error = e.what(); + return 1; } } @@ -489,81 +481,63 @@ void Check(const arrow::Status& status) { if (!status.ok()) throw std::runtime_error(status.ToString()); } -std::string ToString(JNIEnv* env, jobject value) { - const char* chars = env->GetStringUTFChars(static_cast(value), nullptr); - std::string result(chars); - env->ReleaseStringUTFChars(static_cast(value), chars); - return result; -} - } // namespace -SPARK_JNI(jint, abiVersion)(JNIEnv*, jclass) { return 1; } +extern "C" { + +const char* spark_ds_last_error(void) { return last_error.c_str(); } + +int32_t spark_ds_abi_version(void) { return SPARK_DS_ABI_VERSION; } -SPARK_JNI(jlong, createDataSource)(JNIEnv* env, jclass, jstring, jobjectArray keys, - jobjectArray values) { - return Call(env, [&] { - auto source = std::make_unique(); - for (jsize i = 0; i < env->GetArrayLength(keys); i++) { - if (ToString(env, env->GetObjectArrayElement(keys, i)) == "end") { - source->end = std::stoll(ToString(env, env->GetObjectArrayElement(values, i))); - } +int spark_ds_create(const char*, const char* const* option_keys, + const char* const* option_values, int64_t num_options, void** source) { + return Call([&] { + auto range = std::make_unique(); + for (int64_t i = 0; i < num_options; i++) { + if (std::string(option_keys[i]) == "end") range->end = std::stoll(option_values[i]); } - return reinterpret_cast(source.release()); + *source = range.release(); }); } -SPARK_JNI(void, schema)(JNIEnv* env, jclass, jlong, jlong schema_address) { - Call(env, [&] { - Check(arrow::ExportSchema(*Schema(), reinterpret_cast(schema_address))); - }); +int spark_ds_schema(void*, struct ArrowSchema* out) { + return Call([&] { Check(arrow::ExportSchema(*Schema(), out)); }); } -SPARK_JNI(jlong, createReader)(JNIEnv* env, jclass, jlong source, jlong) { - return Call(env, [&] { - auto* reader = new RangeReader{reinterpret_cast(source)->end}; - return reinterpret_cast(reader); - }); +int spark_ds_create_reader(void* source, struct ArrowSchema*, void** reader) { + return Call([&] { *reader = new RangeReader{static_cast(source)->end}; }); } -SPARK_JNI(void, closeDataSource)(JNIEnv*, jclass, jlong source) { - delete reinterpret_cast(source); -} +void spark_ds_close(void* source) { delete static_cast(source); } -SPARK_JNI(jobjectArray, partitions)(JNIEnv* env, jclass, jlong reader) { - return Call(env, [&] { - int64_t end = reinterpret_cast(reader)->end; +int spark_ds_partitions(void* reader, SparkDsSink add_partition, void* sink) { + return Call([&] { + int64_t end = static_cast(reader)->end; int64_t bounds[2][2] = {{0, end / 2}, {end / 2, end}}; - jobjectArray partitions = env->NewObjectArray(2, env->FindClass("[B"), nullptr); - for (jsize i = 0; i < 2; i++) { - jbyteArray partition = env->NewByteArray(sizeof(bounds[i])); - env->SetByteArrayRegion(partition, 0, sizeof(bounds[i]), - reinterpret_cast(bounds[i])); - env->SetObjectArrayElement(partitions, i, partition); + for (const auto& partition : bounds) { + add_partition(sink, reinterpret_cast(partition), sizeof(partition)); } - return partitions; }); } -SPARK_JNI(void, closeReader)(JNIEnv*, jclass, jlong reader) { - delete reinterpret_cast(reader); -} +void spark_ds_close_reader(void* reader) { delete static_cast(reader); } -SPARK_JNI(void, read)(JNIEnv* env, jclass, jbyteArray, jbyteArray partition, - jlong stream_address) { - Call(env, [&] { +int spark_ds_read(const uint8_t*, int64_t, const uint8_t* partition, int64_t, + struct ArrowArrayStream* out) { + return Call([&] { int64_t bounds[2]; - env->GetByteArrayRegion(partition, 0, sizeof(bounds), reinterpret_cast(bounds)); + std::memcpy(bounds, partition, sizeof(bounds)); arrow::Int64Builder ids; for (int64_t id = bounds[0]; id < bounds[1]; id++) Check(ids.Append(id)); std::shared_ptr array; Check(ids.Finish(&array)); auto batch = arrow::RecordBatch::Make(Schema(), array->length(), {array}); auto batches = arrow::RecordBatchReader::Make({batch}, Schema()).ValueOrDie(); - Check(arrow::ExportRecordBatchReader(batches, - reinterpret_cast(stream_address))); + Check(arrow::ExportRecordBatchReader(batches, out)); }); } + +} // extern "C" ``` {% endraw %} @@ -572,8 +546,7 @@ package it as above: ```bash c++ -std=c++20 -O2 -shared -fPIC \ - -I"$JAVA_HOME/include" -I"$JAVA_HOME/include/darwin" \ - -I"$ARROW_HOME/include" -L"$ARROW_HOME/lib" -larrow \ + -I/path/to/spark_datasource.h/dir -I"$ARROW_HOME/include" -L"$ARROW_HOME/lib" -larrow \ -o libspark_datasource_cpp_range.dylib cpp_range.cc cp libspark_datasource_cpp_range.dylib /usr/local/lib/ ``` diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala index 30dcebcd85474..5f940e5126473 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala @@ -846,6 +846,13 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE messageParameters = Map("fileName" -> fileName, "checksum" -> checksum, "path" -> path)) } + def nativeDataSourceUnsupportedJavaVersionError( + required: Int, current: Int): SparkUnsupportedOperationException = { + new SparkUnsupportedOperationException( + errorClass = "NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION", + messageParameters = Map("required" -> required.toString, "current" -> current.toString)) + } + def removedClassInSpark2Error(className: String, e: Throwable): SparkClassNotFoundException = { new SparkClassNotFoundException( errorClass = "_LEGACY_ERROR_TEMP_2052", diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala index 4081e73ccf757..2ce4cb0166a19 100644 --- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala +++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/CheckConnectJvmClientCompatibility.scala @@ -162,7 +162,6 @@ object CheckConnectJvmClientCompatibility { ProblemFilters.exclude[Problem]("org.apache.spark.sql.columnar.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.connector.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.classic.*"), - ProblemFilters.exclude[Problem]("org.apache.spark.sql.datasource.*"), // Native data sources ProblemFilters.exclude[Problem]("org.apache.spark.sql.execution.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.internal.*"), ProblemFilters.exclude[Problem]("org.apache.spark.sql.jdbc.*"), diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala index efef2b08240a7..f8aaa6b335fc7 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala @@ -29,7 +29,7 @@ import org.apache.spark.util.Utils /** * Tests native data sources with Spark Connect: a client adds packages to its session with * `spark.addArtifact`, and the server finds the data sources in the packages of the session or in - * the installed libraries. The tests are skipped when there is no C++ compiler or no JNI headers. + * the installed libraries. The tests are skipped when there is no C++ compiler, or before Java 22. */ class NativeDataSourceConnectSuite extends SparkConnectServerTest { @@ -38,7 +38,7 @@ class NativeDataSourceConnectSuite extends SparkConnectServerTest { compile("test_native_datasource_installed", "-DNAME_PREFIX=\"installed_\"") private def nativeTest(name: String)(f: => Unit): Unit = test(name) { - assume(toolchain.isDefined, "A C++ compiler and the JNI headers are required.") + assume(toolchain.isDefined, "A C++ compiler and Java 22 or later are required.") f } diff --git a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java deleted file mode 100644 index ad8b7fdb06e8f..0000000000000 --- a/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java +++ /dev/null @@ -1,442 +0,0 @@ -/* - * 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; - -import org.apache.spark.annotation.Evolving; - -/** - * The binary interface between Spark and data sources implemented in native code, such as Rust - * or C++. - *

    - * A native data source is a shared library that implements the {@code native} methods of this - * class as JNI functions, for example - * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each - * library into its own copy of this class, so any number of libraries that export the same - * functions can be loaded side by side. Spark adapts a native data source to Data Source V2, with - * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and - * limit pushdown. - *

    - * Spark finds the library of a data source automatically when it is installed like other native - * libraries and named after the data source: the library of {@code my_source} is - * {@code spark_datasource_my_source}, such as {@code libspark_datasource_my_source.so} in - * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of - * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux, and in the - * {@code lib} directory of each installation prefix in {@code PATH}. A library can also be - * distributed in a native data source package: a zip file with the extension {@code .sparkpkg} - * that contains a manifest named {@code spark-native-datasource.json} and the library built for - * one or more platforms, added to a session with {@code spark.addArtifact} or under the paths in - * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL - * guide for details and examples. - * - *

    Lifecycle

    - * Spark creates a data source with {@link #createDataSource} on the driver for each schema - * inference, scan and write, and closes it right after. - *
      - *
    • Batch reads. On the driver, Spark creates a reader with {@link #createReader}, - * pushes operations down by calling, each at most once and in this order, - * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions - * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, - * and closes it. Then it calls {@link #read} on the executors, in one task for each - * partition.
      • - *
    • Streaming reads. On the driver, Spark creates a stream reader with - * {@link #createStreamReader} and keeps it until the streaming query stops. For each - * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with - * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then - * calls {@link #read} on the executors. Offsets are JSON strings defined by the library. When - * a streaming query restarts, from its checkpoint, Spark creates a new stream reader and may - * call {@link #streamPartitions} again for the last range that it planned before, so the - * partitions of a range must read the same data each time.
      • - *
    • Writes. On the driver, once it computed the data to write, right before it starts - * the tasks of the write (or of a micro-batch), Spark creates a writer with {@link #createWriter} - * or {@link #createStreamWriter}, and serializes its state with {@link #serializeWriter}. If - * computing the data fails, Spark creates no writer. On the executors, each task creates a data - * writer with {@link #createDataWriter}, calls {@link #write} for each batch of at most - * {@code spark.sql.execution.arrow.maxRecordsPerBatch} rows, which also ends once it holds - * {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, and calls {@link #commitDataWriter}, - * or {@link #abortDataWriter} if the task or the commit fails. Then the driver calls - * {@link #commit} with the messages of the tasks if they all succeeded, or {@link #abort} - * otherwise, including when {@link #commit} fails, and closes the writer. Several attempts of a - * task may write the same partition, after a failure or with speculative execution, each with a - * data writer of its own. Spark lets the attempts of a partition commit one at a time, and passes - * to {@link #commit} the message of one committed attempt for each partition, so the write - * consists of the data of these attempts only. It aborts the data writers of the other attempts, - * except on an executor that is lost, where a data writer is neither committed nor aborted, and - * may even have committed before its message reached the driver. A streaming write commits or - * aborts each micro-batch (epoch). Spark may commit the same epoch more than once, for example - * when a streaming query restarts after a failure, so committing an epoch that is already - * committed must have no effect, for exactly-once output. A micro-batch streaming query creates a - * new stream writer for each micro-batch, and Spark closes the previous stream writer of the - * query when the query creates a new one. Spark is not told when a streaming query stops, so it - * closes the last stream writer of the query once the writer is no longer referenced, or when the - * query is restarted.
      • - *
    - * - *

    Conventions

    - *
      - *
    • Handles. The {@code create} methods return an opaque, non-zero handle, typically a - * pointer. Spark releases each handle exactly once with the matching {@code close} method, or - * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers, though not - * always right after the last use of a stream writer (see Writes above). A reader or - * writer handle must not depend on the data source handle it was created from, because Spark - * closes the data source handle right after creating it.
      • - *
    • State. Partitions, the state of readers and writers, and commit messages are byte - * arrays in a format defined by the library. Spark sends them to the executors and back.
      • - *
    • Errors. To report an error, throw a Java exception, for example with the JNI - * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and return any value. Spark - * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.
      • - *
    • Arrow data. Schemas and data are exchanged with the - * Arrow C data - * interface. A {@code long} argument named {@code ...Address} is the address of a C struct - * allocated by Spark. For an output struct, the library initializes it and Spark takes - * ownership. For an input struct, the library may move it to take ownership; otherwise, Spark - * releases it after the call returns, so the library must not use it afterwards. The arrays - * that the library exports must have an offset of 0, at every level: Arrow Java ignores the - * offsets of the arrays it imports, so Spark rejects an array with another offset, such as a - * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with - * fewer values than its parent refers to, such as a column shorter than its batch. The buffers - * of an array must hold all of its values, as the C data interface requires: it does not carry - * the sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that - * is too short.
      • - *
    • Schemas. Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: - * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, - * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is - * {@code timestamp[us]} without one. Data read from a library may also use - * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, and any names - * for the fields of the entries of a map, but not unsigned integers, {@code decimal256}, list - * views or dictionary encoding.
      • - *
    • Options. Option keys are passed in lower case, because Spark matches them - * case-insensitively.
      • - *
    • Optional methods. A library only needs to export the functions for the operations - * it supports. If a function is missing, Spark treats the operation as unsupported, except for - * the pushdown methods and the {@code serialize} methods, which then push nothing down and - * return no state.
      • - *
    • Threads. Spark calls the library from multiple threads, and calls on different - * handles may run concurrently. Calls on the same handle never run concurrently. A handle may - * be released from a thread other than the threads that used it, such as a JVM - * {@code Cleaner} thread when Spark releases a handle once it is no longer referenced, for - * example the last stream writer of a streaming query.
      • - *
    - * - *

    Compatibility

    - * The JVM binds the functions of a library by their names only, so the name of this class, the - * names of its methods and their signatures are the interface. Within an {@link #ABI_VERSION}: - *
      - *
    • the signature and the semantics of an existing method never change;
      • - *
    • new methods may be added, only as optional methods, so that a library that does not - * export them keeps working as before;
      • - *
    • new values may appear where this documentation says so, such as the names of pushed - * predicates.
      • - *
    - * An incompatible change uses a new {@link #ABI_VERSION}, with new names or a new class. A - * library reports the version it implements with {@link #abiVersion}, and Spark rejects a library - * of a version it does not support. Spark supports version 1 at least in all the releases of - * Spark 4 from 4.4.0. - * - * @since 4.4.0 - */ -@Evolving -public final class NativeBridge { - - /** - * The version of this interface. It changes when the interface changes incompatibly: see - * Compatibility above. - */ - public static final int ABI_VERSION = 1; - - private NativeBridge() {} - - /** - * Loads the library that implements the native methods of this copy of the class. Spark calls - * it on the copy of the class it creates for the library. - */ - static void load(String path) { - System.load(path); - } - - /** Returns the version of this interface that the library implements, {@link #ABI_VERSION}. */ - public static native int abiVersion(); - - // --------------------------------------------------------------------------------------------- - // Data source (driver) - // --------------------------------------------------------------------------------------------- - - /** - * Creates a data source with the options given by the user, such as - * {@code spark.read.option("key", "value")}. - * - * @param name the name of the data source, as listed in the package manifest. A library can - * implement several data sources. - * @param optionKeys the option keys, in lower case - * @param optionValues the option values, in the same order as the keys - * @return the data source handle - */ - public static native long createDataSource( - String name, String[] optionKeys, String[] optionValues); - - /** - * Exports the schema of the data. Spark calls it only when the user does not specify a schema. - * - * @param schemaAddress the address of the {@code ArrowSchema} to initialize with a struct type - * whose fields are the columns - */ - public static native void schema(long dataSource, long schemaAddress); - - /** - * Creates a reader for a batch scan. - * - * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the - * user specified, or the one returned by {@link #schema} - * @return the reader handle - */ - public static native long createReader(long dataSource, long schemaAddress); - - /** - * Creates a reader for a micro-batch streaming scan. - * - * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the - * user specified, or the one returned by {@link #schema} - * @return the stream reader handle - */ - public static native long createStreamReader(long dataSource, long schemaAddress); - - /** - * Creates a writer for a batch write. - * - * @param schemaAddress the address of the {@code ArrowSchema} of the data to write - * @param overwrite whether to replace the existing data (save mode "overwrite") instead of - * appending to it - * @return the writer handle - */ - public static native long createWriter(long dataSource, long schemaAddress, boolean overwrite); - - /** - * Creates a writer for a streaming write. - * - * @param schemaAddress the address of the {@code ArrowSchema} of the data to write - * @param overwrite whether to replace the existing data in each micro-batch (output mode - * "complete") instead of appending to it - * @return the writer handle - */ - public static native long createStreamWriter( - long dataSource, long schemaAddress, boolean overwrite); - - /** Releases a data source handle. */ - public static native void closeDataSource(long dataSource); - - // --------------------------------------------------------------------------------------------- - // Batch reader (driver) - // --------------------------------------------------------------------------------------------- - - /** - * Pushes predicates down to the reader. The predicates are combined with AND. - *

    - * Each predicate is a JSON expression tree. An expression is one of: - *

      - *
    • {@code {"type": "column", "name": ["a", "b"]}}: a column. A nested field has a name - * part for each level.
      • - *
    • {@code {"type": "literal", "dataType": "bigint", "value": 1}}: a literal. The data type - * is a Spark SQL type name such as {@code int}, {@code string} or {@code decimal(10,2)}. The - * value is {@code null}, a JSON boolean or number for boolean and numeric types, a string - * for {@code string}, a base64 string for {@code binary}, a decimal string for - * {@code decimal}, the days since the epoch for {@code date}, and the microseconds since the - * epoch for {@code timestamp} and {@code timestamp_ntz}.
      • - *
    • {@code {"type": "function", "name": ">", "children": [...]}}: a predicate, with the - * name it has in {@code org.apache.spark.sql.connector.expressions.filter.Predicate}. The - * names are {@code =}, {@code <>}, {@code <=>} (equal, with null equal to null), {@code <}, - * {@code <=}, {@code >} and {@code >=} with two children, {@code IN} with the value to look - * for and then the values to look in, {@code IS_NULL} and {@code IS_NOT_NULL} with one - * child, {@code STARTS_WITH}, {@code ENDS_WITH} and {@code CONTAINS} with the string and the - * string to look for, {@code AND} and {@code OR} with two predicates, and {@code NOT} with - * one. The children of the other functions are columns, literals or predicates.
      • - *
    - * Spark does not pass predicates that it cannot express this way, such as predicates that use - * arithmetic or other functions, nor predicates on columns that contain strings with a - * non-binary collation, such as {@code UTF8_LCASE}: strings are compared by their UTF-8 bytes. - * A later version of Spark may pass predicates with other names without changing - * {@link #ABI_VERSION}, so a reader must not report a predicate as evaluated completely if it - * contains a name that the reader does not know, or a data type that it does not support. - *

    - * A predicate that the reader evaluates completely must follow the semantics of Spark SQL. In - * particular, for floating-point columns, NaN is equal to NaN and greater than any other value, - * including positive infinity, and {@code -0.0} is equal to {@code 0.0}, unlike IEEE 754 - * comparisons. For example, {@code d > 1.0} is true for a NaN value of {@code d}. Comparisons - * with null are null, so a row only matches if the predicate is true. - * - * @return for each predicate, whether the reader evaluates it completely, so that Spark does - * not have to evaluate it again. The reader can still use the other predicates to skip - * data. - */ - public static native boolean[] pushPredicates(long reader, String[] predicates); - - /** - * Pushes a LIMIT down to the reader. - * - * @return whether the reader returns at most {@code limit} rows in total. Spark applies the - * limit again after reading either way. - */ - public static native boolean pushLimit(long reader, int limit); - - /** - * Prunes the columns to read. - * - * @param columnNames the names of the top-level columns that the query needs - * @return whether {@link #read} returns exactly these columns, in this order. If not, it - * returns all the columns of the schema to read. - */ - public static native boolean pruneColumns(long reader, String[] columnNames); - - /** - * Plans the partitions of the scan. Spark reads each partition in a separate task. - * - * @return the partitions, each serialized in a format defined by the library - */ - public static native byte[][] partitions(long reader); - - /** - * Serializes the state of the reader that {@link #read} needs on the executors, such as the - * pushed predicates. Spark calls it after {@link #partitions}. - */ - public static native byte[] serializeReader(long reader); - - /** Releases a batch reader handle. */ - public static native void closeReader(long reader); - - // --------------------------------------------------------------------------------------------- - // Stream reader (driver) - // --------------------------------------------------------------------------------------------- - - /** Returns the offset to start from when the streaming query has no checkpoint, as JSON. */ - public static native String initialOffset(long streamReader); - - /** Returns the most recent offset available, as JSON. */ - public static native String latestOffset(long streamReader); - - /** - * Plans the partitions that read the data after {@code start} up to and including {@code end}. - * - * @return the partitions, each serialized in a format defined by the library - */ - public static native byte[][] streamPartitions(long streamReader, String start, String end); - - /** - * Serializes the state of the stream reader that {@link #read} needs on the executors. Spark - * calls it after each call to {@link #streamPartitions}. - */ - public static native byte[] serializeStreamReader(long streamReader); - - /** - * Tells the library that Spark has processed all the data up to and including {@code end}, - * and will not ask for it again. - */ - public static native void commitOffset(long streamReader, String end); - - /** Releases a stream reader handle, when the streaming query stops. */ - public static native void closeStreamReader(long streamReader); - - // --------------------------------------------------------------------------------------------- - // Partition reader (executor) - // --------------------------------------------------------------------------------------------- - - /** - * Reads a partition of a batch or streaming scan. - * - * @param readerState the state returned by {@link #serializeReader} or - * {@link #serializeStreamReader}, or an empty array if there is none - * @param partition a partition returned by {@link #partitions} or {@link #streamPartitions} - * @param streamAddress the address of the {@code ArrowArrayStream} to initialize. Each array - * of the stream is a struct array whose fields are the columns to read: - * the columns of the schema to read, or the pruned columns. Spark calls - * the stream from the thread of the task that reads the partition. - * {@code get_next} may block, for example on network I/O, but Spark - * cannot interrupt native code: when the task is killed, Spark only - * stops reading once {@code get_next} returns, and then releases the - * stream. So a {@code get_next} that blocks forever keeps the task - * running, unless the task reaper ({@code spark.task.reaper.enabled} - * and {@code spark.task.reaper.killTimeout}) kills the executor. Use - * timeouts for blocking operations. - */ - public static native void read(byte[] readerState, byte[] partition, long streamAddress); - - // --------------------------------------------------------------------------------------------- - // Writer (driver) - // --------------------------------------------------------------------------------------------- - - /** - * Serializes the state of the writer that {@link #createDataWriter} needs on the executors. - * Spark calls it once, after creating the writer. - */ - public static native byte[] serializeWriter(long writer); - - /** - * Commits a write, when all the tasks succeeded. - * - * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write - * @param messages the messages returned by {@link #commitDataWriter} of each task - */ - public static native void commit(long writer, long epochId, byte[][] messages); - - /** - * Aborts a write, when it failed. - * - * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write - * @param messages the messages returned by {@link #commitDataWriter} of each task. An element - * is null if its task did not commit. - */ - public static native void abort(long writer, long epochId, byte[][] messages); - - /** Releases a writer handle. */ - public static native void closeWriter(long writer); - - // --------------------------------------------------------------------------------------------- - // Data writer (executor) - // --------------------------------------------------------------------------------------------- - - /** - * Creates the writer for one task. - * - * @param writerState the state returned by {@link #serializeWriter}, or an empty array if there - * is none - * @param partitionId the partition that the task writes - * @param taskId the ID of the task attempt - * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write - * @return the data writer handle - */ - public static native long createDataWriter( - byte[] writerState, int partitionId, long taskId, long epochId); - - /** - * Writes a batch of data. - * - * @param arrayAddress the address of the {@code ArrowArray}: a struct array whose fields are - * the columns - * @param schemaAddress the address of the {@code ArrowSchema} of the array - */ - public static native void write(long dataWriter, long arrayAddress, long schemaAddress); - - /** - * Commits the data written by the task and releases the data writer handle. If it fails, the - * handle is not released: Spark then calls {@link #abortDataWriter} with it, to abort the data - * written by the task and release the handle. - * - * @return the message to pass to {@link #commit} on the driver - */ - public static native byte[] commitDataWriter(long dataWriter); - - /** Aborts the data written by the task and releases the data writer handle. */ - public static native void abortDataWriter(long dataWriter); -} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java index 61145e186879e..d2337d8748e12 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -17,39 +17,62 @@ package org.apache.spark.sql.execution.datasources.v2.ffi; -import java.io.IOException; -import java.io.InputStream; import java.lang.invoke.MethodHandle; import java.lang.invoke.MethodHandles; import java.lang.invoke.MethodType; -import java.lang.reflect.InvocationTargetException; -import java.lang.reflect.Method; +import java.lang.reflect.Array; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; -import org.apache.spark.sql.datasource.NativeBridge; +import org.apache.spark.unsafe.Platform; /** - * A native data source library, loaded into its own copy of {@link NativeBridge}. + * A native data source library, which implements the C interface in + * {@code org/apache/spark/sql/datasource/spark_datasource.h}, called through the Foreign Function + * and Memory (FFM) API of Java 22 and later. *

    - * JNI links the native methods of a class to the libraries loaded by the class loader of that - * class, and a library belongs to the class loader of the class that loads it. Defining a copy of - * {@link NativeBridge} in a separate class loader for each library therefore lets all the - * libraries export the same JNI functions. Spark calls the methods of a copy through method - * handles, because the copy is a different class than the {@link NativeBridge} it is compiled - * against. Calling a method that the library does not export throws - * {@link UnsatisfiedLinkError}. + * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached + * through method handles, looked up once when the first library is loaded. Each function of a + * library is bound to a method handle that takes and returns primitive values only, with + * pointers as {@code long} addresses, so the calls themselves use no reflection. The memory that + * Spark passes to the library is allocated with {@link Platform}. + *

    + * Each library is looked up separately, so all the libraries can export the same functions. + * Calling a function that the library does not export throws {@link MissingFunctionException}. A + * function that fails throws a {@link RuntimeException} with the message of + * {@code spark_ds_last_error}. */ final class NativeLibrary { - private static final String BRIDGE_CLASS = NativeBridge.class.getName(); + /** The version of the interface that Spark supports, {@code SPARK_DS_ABI_VERSION}. */ + static final int ABI_VERSION = 1; + + /** The first version of Java with the final FFM API. */ + static final int MINIMUM_JAVA_VERSION = 22; + + /** Whether this JVM can load native data source libraries. */ + static boolean isSupported() { + return Runtime.version().feature() >= MINIMUM_JAVA_VERSION; + } + // The C types of the parameters and the results of the functions. + private enum C { VOID, BOOL, INT32, INT64, POINTER } + + private final String path; + private final MethodHandle lastError; private final MethodHandle abiVersion; - private final MethodHandle createDataSource; + private final MethodHandle create; private final MethodHandle schema; private final MethodHandle createReader; private final MethodHandle createStreamReader; private final MethodHandle createWriter; private final MethodHandle createStreamWriter; - private final MethodHandle closeDataSource; + private final MethodHandle close; private final MethodHandle pushPredicates; private final MethodHandle pushLimit; private final MethodHandle pruneColumns; @@ -72,225 +95,594 @@ final class NativeLibrary { private final MethodHandle commitDataWriter; private final MethodHandle abortDataWriter; - private NativeLibrary(Class bridge) throws ReflectiveOperationException { - Class jlong = long.class; - Class bytes = byte[].class; - Class strings = String[].class; - abiVersion = find(bridge, "abiVersion", int.class); - createDataSource = find(bridge, "createDataSource", jlong, String.class, strings, strings); - schema = find(bridge, "schema", void.class, jlong, jlong); - createReader = find(bridge, "createReader", jlong, jlong, jlong); - createStreamReader = find(bridge, "createStreamReader", jlong, jlong, jlong); - createWriter = find(bridge, "createWriter", jlong, jlong, jlong, boolean.class); - createStreamWriter = find(bridge, "createStreamWriter", jlong, jlong, jlong, boolean.class); - closeDataSource = find(bridge, "closeDataSource", void.class, jlong); - pushPredicates = find(bridge, "pushPredicates", boolean[].class, jlong, strings); - pushLimit = find(bridge, "pushLimit", boolean.class, jlong, int.class); - pruneColumns = find(bridge, "pruneColumns", boolean.class, jlong, strings); - partitions = find(bridge, "partitions", byte[][].class, jlong); - serializeReader = find(bridge, "serializeReader", bytes, jlong); - closeReader = find(bridge, "closeReader", void.class, jlong); - initialOffset = find(bridge, "initialOffset", String.class, jlong); - latestOffset = find(bridge, "latestOffset", String.class, jlong); + private NativeLibrary(String path, Object lookup) throws Throwable { + this.path = path; + C i32 = C.INT32; + C i64 = C.INT64; + C ptr = C.POINTER; + lastError = bind(lookup, "spark_ds_last_error", ptr); + abiVersion = bind(lookup, "spark_ds_abi_version", i32); + create = bind(lookup, "spark_ds_create", i32, ptr, ptr, ptr, i64, ptr); + schema = bind(lookup, "spark_ds_schema", i32, ptr, ptr); + createReader = bind(lookup, "spark_ds_create_reader", i32, ptr, ptr, ptr); + createStreamReader = bind(lookup, "spark_ds_create_stream_reader", i32, ptr, ptr, ptr); + createWriter = bind(lookup, "spark_ds_create_writer", i32, ptr, ptr, C.BOOL, ptr); + createStreamWriter = + bind(lookup, "spark_ds_create_stream_writer", i32, ptr, ptr, C.BOOL, ptr); + close = bind(lookup, "spark_ds_close", C.VOID, ptr); + pushPredicates = bind(lookup, "spark_ds_push_predicates", i32, ptr, ptr, i64, ptr); + pushLimit = bind(lookup, "spark_ds_push_limit", i32, ptr, i32, ptr); + pruneColumns = bind(lookup, "spark_ds_prune_columns", i32, ptr, ptr, i64, ptr); + partitions = bind(lookup, "spark_ds_partitions", i32, ptr, ptr, ptr); + serializeReader = bind(lookup, "spark_ds_serialize_reader", i32, ptr, ptr, ptr); + closeReader = bind(lookup, "spark_ds_close_reader", C.VOID, ptr); + initialOffset = bind(lookup, "spark_ds_initial_offset", i32, ptr, ptr, ptr); + latestOffset = bind(lookup, "spark_ds_latest_offset", i32, ptr, ptr, ptr); streamPartitions = - find(bridge, "streamPartitions", byte[][].class, jlong, String.class, String.class); - serializeStreamReader = find(bridge, "serializeStreamReader", bytes, jlong); - commitOffset = find(bridge, "commitOffset", void.class, jlong, String.class); - closeStreamReader = find(bridge, "closeStreamReader", void.class, jlong); - read = find(bridge, "read", void.class, bytes, bytes, jlong); - serializeWriter = find(bridge, "serializeWriter", bytes, jlong); - commit = find(bridge, "commit", void.class, jlong, jlong, byte[][].class); - abort = find(bridge, "abort", void.class, jlong, jlong, byte[][].class); - closeWriter = find(bridge, "closeWriter", void.class, jlong); - createDataWriter = find(bridge, "createDataWriter", jlong, bytes, int.class, jlong, jlong); - write = find(bridge, "write", void.class, jlong, jlong, jlong); - commitDataWriter = find(bridge, "commitDataWriter", bytes, jlong); - abortDataWriter = find(bridge, "abortDataWriter", void.class, jlong); - } - - private static MethodHandle find( - Class bridge, String name, Class returnType, Class... parameterTypes) - throws ReflectiveOperationException { - return MethodHandles.publicLookup() - .findStatic(bridge, name, MethodType.methodType(returnType, parameterTypes)); + bind(lookup, "spark_ds_stream_partitions", i32, ptr, ptr, ptr, ptr, ptr); + serializeStreamReader = + bind(lookup, "spark_ds_serialize_stream_reader", i32, ptr, ptr, ptr); + commitOffset = bind(lookup, "spark_ds_commit_offset", i32, ptr, ptr); + closeStreamReader = bind(lookup, "spark_ds_close_stream_reader", C.VOID, ptr); + read = bind(lookup, "spark_ds_read", i32, ptr, i64, ptr, i64, ptr); + serializeWriter = bind(lookup, "spark_ds_serialize_writer", i32, ptr, ptr, ptr); + commit = bind(lookup, "spark_ds_commit", i32, ptr, i64, ptr, ptr, i64); + abort = bind(lookup, "spark_ds_abort", i32, ptr, i64, ptr, ptr, i64); + closeWriter = bind(lookup, "spark_ds_close_writer", C.VOID, ptr); + createDataWriter = + bind(lookup, "spark_ds_create_data_writer", i32, ptr, i64, i32, i64, i64, ptr); + write = bind(lookup, "spark_ds_write", i32, ptr, ptr, ptr); + commitDataWriter = bind(lookup, "spark_ds_commit_data_writer", i32, ptr, ptr, ptr); + abortDataWriter = bind(lookup, "spark_ds_abort_data_writer", i32, ptr); } /** - * Loads the library at the given path into a new copy of {@link NativeBridge}. + * Loads the library at the given path. The library stays loaded until the JVM exits. + * + * @throws LoadException if the library cannot be loaded */ static NativeLibrary load(String path) throws Throwable { - Class bridge = new BridgeClassLoader(NativeBridge.class.getClassLoader()) - .loadClass(BRIDGE_CLASS); - // The copy itself has to call System.load, because a library belongs to the class loader of - // the class that loads it. NativeBridge.load is package-private, so that it is not part of - // the public interface, so it is called reflectively. Java always allows this for the classes - // on the class path: they are in unnamed modules, whose packages are all open. - Method load = bridge.getDeclaredMethod("load", String.class); - load.setAccessible(true); + if (!isSupported()) { + throw new UnsupportedOperationException( + "Native data sources require Java " + MINIMUM_JAVA_VERSION + " or later."); + } + Object lookup; try { - load.invoke(null, path); - } catch (InvocationTargetException e) { - throw e.getCause(); + lookup = Ffm.LIBRARY_LOOKUP.invoke(Path.of(path), Ffm.GLOBAL_ARENA); + } catch (IllegalArgumentException e) { + throw new LoadException(e.getMessage(), e); } - return new NativeLibrary(bridge); + return new NativeLibrary(path, lookup); } /** - * A class loader that defines its own copy of {@link NativeBridge}, and delegates all the other - * classes to its parent. + * Binds a function of the library to a method handle whose parameters and result have the + * Java types of the given C types, with {@code long} for pointers. Returns null if the library + * does not export the function. */ - private static final class BridgeClassLoader extends ClassLoader { - static { - ClassLoader.registerAsParallelCapable(); + private static MethodHandle bind(Object lookup, String name, C result, C... parameters) + throws Throwable { + Optional symbol = (Optional) Ffm.FIND.invoke(lookup, name); + if (symbol.isEmpty()) { + return null; } - - BridgeClassLoader(ClassLoader parent) { - super(parent); + Object layouts = Ffm.layouts(parameters); + Object descriptor = result == C.VOID + ? Ffm.FUNCTION_OF_VOID.invoke(layouts) + : Ffm.FUNCTION_OF.invoke(Ffm.layout(result), layouts); + MethodHandle handle = (MethodHandle) Ffm.DOWNCALL.invoke( + Ffm.LINKER, symbol.get(), descriptor, Ffm.NO_OPTIONS); + // Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses. + for (int i = 0; i < parameters.length; i++) { + if (parameters[i] == C.POINTER) { + handle = MethodHandles.filterArguments(handle, i, Ffm.OF_ADDRESS); + } } + if (result == C.POINTER) { + handle = MethodHandles.filterReturnValue(handle, Ffm.ADDRESS_OF); + } + return handle; + } + + /** The method handles of the FFM API, initialized when the first library is loaded. */ + private static final class Ffm { + static final Class LAYOUT_CLASS; + static final Object LINKER; + static final Object GLOBAL_ARENA; + static final Object NO_OPTIONS; + static final MethodHandle LIBRARY_LOOKUP; + static final MethodHandle FIND; + static final MethodHandle DOWNCALL; + static final MethodHandle FUNCTION_OF; + static final MethodHandle FUNCTION_OF_VOID; + // (long) -> MemorySegment and (MemorySegment) -> long. + static final MethodHandle OF_ADDRESS; + static final MethodHandle ADDRESS_OF; + static final Object BOOL; + static final Object INT32; + static final Object INT64; + static final Object POINTER; + // The address of the C function SparkDsSink, which calls NativeLibrary.sink. + static final long SINK; - @Override - protected Class loadClass(String name, boolean resolve) throws ClassNotFoundException { - if (!BRIDGE_CLASS.equals(name)) { - return super.loadClass(name, resolve); + static { + try { + // A full-privilege lookup, so that the restricted methods of the FFM API see this class + // as their caller, in the unnamed module that Spark enables native access for. + MethodHandles.Lookup lookup = MethodHandles.lookup(); + Class linker = Class.forName("java.lang.foreign.Linker"); + Class option = Class.forName("java.lang.foreign.Linker$Option"); + Class symbolLookup = Class.forName("java.lang.foreign.SymbolLookup"); + Class arena = Class.forName("java.lang.foreign.Arena"); + Class segment = Class.forName("java.lang.foreign.MemorySegment"); + Class layout = Class.forName("java.lang.foreign.MemoryLayout"); + Class layoutArray = Array.newInstance(layout, 0).getClass(); + Class valueLayout = Class.forName("java.lang.foreign.ValueLayout"); + Class descriptor = Class.forName("java.lang.foreign.FunctionDescriptor"); + Class options = Array.newInstance(option, 0).getClass(); + + LAYOUT_CLASS = layout; + LINKER = lookup.findStatic(linker, "nativeLinker", MethodType.methodType(linker)) + .invoke(); + GLOBAL_ARENA = lookup.findStatic(arena, "global", MethodType.methodType(arena)).invoke(); + NO_OPTIONS = Array.newInstance(option, 0); + LIBRARY_LOOKUP = lookup.findStatic(symbolLookup, "libraryLookup", + MethodType.methodType(symbolLookup, Path.class, arena)); + FIND = lookup.findVirtual(symbolLookup, "find", + MethodType.methodType(Optional.class, String.class)); + DOWNCALL = lookup.findVirtual(linker, "downcallHandle", + MethodType.methodType(MethodHandle.class, segment, descriptor, options)) + .asFixedArity(); + FUNCTION_OF = lookup.findStatic(descriptor, "of", + MethodType.methodType(descriptor, layout, layoutArray)).asFixedArity(); + FUNCTION_OF_VOID = lookup.findStatic(descriptor, "ofVoid", + MethodType.methodType(descriptor, layoutArray)).asFixedArity(); + OF_ADDRESS = lookup.findStatic(segment, "ofAddress", + MethodType.methodType(segment, long.class)); + ADDRESS_OF = lookup.findVirtual(segment, "address", MethodType.methodType(long.class)); + BOOL = valueLayout.getField("JAVA_BOOLEAN").get(null); + INT32 = valueLayout.getField("JAVA_INT").get(null); + INT64 = valueLayout.getField("JAVA_LONG").get(null); + POINTER = valueLayout.getField("ADDRESS").get(null); + + // The upcall stub of the sink, which lives as long as the JVM. + MethodHandle target = lookup.findStatic(NativeLibrary.class, "sink", + MethodType.methodType(void.class, long.class, long.class, long.class)); + target = MethodHandles.filterArguments(target, 0, ADDRESS_OF, ADDRESS_OF); + Object sinkDescriptor = + FUNCTION_OF_VOID.invoke(layouts(C.POINTER, C.POINTER, C.INT64)); + Object stub = lookup.findVirtual(linker, "upcallStub", + MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options)) + .asFixedArity() + .invoke(LINKER, target, sinkDescriptor, GLOBAL_ARENA, NO_OPTIONS); + SINK = (long) ADDRESS_OF.invoke(stub); + } catch (Throwable e) { + throw new IllegalStateException("Cannot initialize the FFM API.", e); } - synchronized (getClassLoadingLock(name)) { - Class loaded = findLoadedClass(name); - if (loaded == null) { - String resource = name.replace('.', '/') + ".class"; - try (InputStream in = getParent().getResourceAsStream(resource)) { - if (in == null) { - throw new ClassNotFoundException(name); - } - byte[] bytes = in.readAllBytes(); - loaded = defineClass(name, bytes, 0, bytes.length); - } catch (IOException e) { - throw new ClassNotFoundException(name, e); - } - } - if (resolve) { - resolveClass(loaded); - } - return loaded; + } + + /** Returns a MemoryLayout[] with the layouts of the given types. */ + static Object layouts(C... types) { + Object layouts = Array.newInstance(LAYOUT_CLASS, types.length); + for (int i = 0; i < types.length; i++) { + Array.set(layouts, i, layout(types[i])); } + return layouts; + } + + static Object layout(C type) { + return switch (type) { + case BOOL -> BOOL; + case INT32 -> INT32; + case INT64 -> INT64; + case POINTER -> POINTER; + case VOID -> throw new IllegalArgumentException("void has no layout"); + }; + } + } + + /** Thrown when a library cannot be loaded. */ + static final class LoadException extends RuntimeException { + LoadException(String message, Throwable cause) { + super(message, cause); + } + } + + /** Thrown when a function is called that the library does not export, with its name. */ + static final class MissingFunctionException extends RuntimeException { + MissingFunctionException(String name) { + super(name); + } + } + + // --------------------------------------------------------------------------------------------- + // Calls + // --------------------------------------------------------------------------------------------- + + private static MethodHandle require(MethodHandle function, String name) { + if (function == null) { + throw new MissingFunctionException(name); } + return function; + } + + /** Throws the error that the library reported if a function failed. */ + private void check(int status, String name) throws Throwable { + if (status == 0) { + return; + } + long message = lastError == null ? 0 : (long) lastError.invokeExact(); + throw new RuntimeException(message == 0 + ? name + " failed with the error code " + status + "." + : new String(readCString(message), StandardCharsets.UTF_8)); + } + + private static long readHandle(long out, String name) { + long handle = Platform.getLong(null, out); + if (handle == 0) { + throw new IllegalStateException(name + " returned a NULL handle."); + } + return handle; } int abiVersion() throws Throwable { - return (int) abiVersion.invokeExact(); + return (int) require(abiVersion, "spark_ds_abi_version").invokeExact(); } long createDataSource(String name, String[] optionKeys, String[] optionValues) throws Throwable { - return (long) createDataSource.invokeExact(name, optionKeys, optionValues); + MethodHandle f = require(create, "spark_ds_create"); + try (Memory memory = new Memory()) { + long out = memory.zeroed(8); + check((int) f.invokeExact(memory.string(name), memory.strings(optionKeys), + memory.strings(optionValues), (long) optionKeys.length, out), "spark_ds_create"); + return readHandle(out, "spark_ds_create"); + } } void schema(long dataSource, long schemaAddress) throws Throwable { - schema.invokeExact(dataSource, schemaAddress); + check((int) require(schema, "spark_ds_schema").invokeExact(dataSource, schemaAddress), + "spark_ds_schema"); } long createReader(long dataSource, long schemaAddress) throws Throwable { - return (long) createReader.invokeExact(dataSource, schemaAddress); + return createWithSchema(createReader, "spark_ds_create_reader", dataSource, schemaAddress); } long createStreamReader(long dataSource, long schemaAddress) throws Throwable { - return (long) createStreamReader.invokeExact(dataSource, schemaAddress); + return createWithSchema( + createStreamReader, "spark_ds_create_stream_reader", dataSource, schemaAddress); + } + + private long createWithSchema( + MethodHandle function, String name, long dataSource, long schemaAddress) throws Throwable { + MethodHandle f = require(function, name); + try (Memory memory = new Memory()) { + long out = memory.zeroed(8); + check((int) f.invokeExact(dataSource, schemaAddress, out), name); + return readHandle(out, name); + } } long createWriter(long dataSource, long schemaAddress, boolean overwrite) throws Throwable { - return (long) createWriter.invokeExact(dataSource, schemaAddress, overwrite); + return createAnyWriter( + createWriter, "spark_ds_create_writer", dataSource, schemaAddress, overwrite); } long createStreamWriter(long dataSource, long schemaAddress, boolean overwrite) throws Throwable { - return (long) createStreamWriter.invokeExact(dataSource, schemaAddress, overwrite); + return createAnyWriter( + createStreamWriter, "spark_ds_create_stream_writer", dataSource, schemaAddress, overwrite); + } + + private long createAnyWriter( + MethodHandle function, + String name, + long dataSource, + long schemaAddress, + boolean overwrite) throws Throwable { + MethodHandle f = require(function, name); + try (Memory memory = new Memory()) { + long out = memory.zeroed(8); + check((int) f.invokeExact(dataSource, schemaAddress, overwrite, out), name); + return readHandle(out, name); + } } void closeDataSource(long dataSource) throws Throwable { - closeDataSource.invokeExact(dataSource); + require(close, "spark_ds_close").invokeExact(dataSource); } boolean[] pushPredicates(long reader, String[] predicates) throws Throwable { - return (boolean[]) pushPredicates.invokeExact(reader, predicates); + MethodHandle f = require(pushPredicates, "spark_ds_push_predicates"); + try (Memory memory = new Memory()) { + long evaluated = memory.zeroed(Math.max(1, predicates.length)); + check((int) f.invokeExact(reader, memory.strings(predicates), (long) predicates.length, + evaluated), "spark_ds_push_predicates"); + boolean[] result = new boolean[predicates.length]; + for (int i = 0; i < result.length; i++) { + result[i] = Platform.getByte(null, evaluated + i) != 0; + } + return result; + } } boolean pushLimit(long reader, int limit) throws Throwable { - return (boolean) pushLimit.invokeExact(reader, limit); + MethodHandle f = require(pushLimit, "spark_ds_push_limit"); + try (Memory memory = new Memory()) { + long applied = memory.zeroed(1); + check((int) f.invokeExact(reader, limit, applied), "spark_ds_push_limit"); + return Platform.getByte(null, applied) != 0; + } } boolean pruneColumns(long reader, String[] columnNames) throws Throwable { - return (boolean) pruneColumns.invokeExact(reader, columnNames); + MethodHandle f = require(pruneColumns, "spark_ds_prune_columns"); + try (Memory memory = new Memory()) { + long pruned = memory.zeroed(1); + check((int) f.invokeExact(reader, memory.strings(columnNames), (long) columnNames.length, + pruned), "spark_ds_prune_columns"); + return Platform.getByte(null, pruned) != 0; + } } byte[][] partitions(long reader) throws Throwable { - return (byte[][]) partitions.invokeExact(reader); + MethodHandle f = require(partitions, "spark_ds_partitions"); + try (Sink sink = new Sink()) { + check((int) f.invokeExact(reader, Ffm.SINK, sink.id), "spark_ds_partitions"); + return sink.all(); + } } byte[] serializeReader(long reader) throws Throwable { - return (byte[]) serializeReader.invokeExact(reader); + return serialize(serializeReader, "spark_ds_serialize_reader", reader); } void closeReader(long reader) throws Throwable { - closeReader.invokeExact(reader); + require(closeReader, "spark_ds_close_reader").invokeExact(reader); } String initialOffset(long streamReader) throws Throwable { - return (String) initialOffset.invokeExact(streamReader); + return offset(initialOffset, "spark_ds_initial_offset", streamReader); } String latestOffset(long streamReader) throws Throwable { - return (String) latestOffset.invokeExact(streamReader); + return offset(latestOffset, "spark_ds_latest_offset", streamReader); + } + + private String offset(MethodHandle function, String name, long streamReader) + throws Throwable { + byte[] offset = serialize(function, name, streamReader); + if (offset == null) { + throw new IllegalStateException(name + " returned no offset."); + } + return new String(offset, StandardCharsets.UTF_8); } byte[][] streamPartitions(long streamReader, String start, String end) throws Throwable { - return (byte[][]) streamPartitions.invokeExact(streamReader, start, end); + MethodHandle f = require(streamPartitions, "spark_ds_stream_partitions"); + try (Memory memory = new Memory(); Sink sink = new Sink()) { + check((int) f.invokeExact(streamReader, memory.string(start), memory.string(end), + Ffm.SINK, sink.id), "spark_ds_stream_partitions"); + return sink.all(); + } } byte[] serializeStreamReader(long streamReader) throws Throwable { - return (byte[]) serializeStreamReader.invokeExact(streamReader); + return serialize(serializeStreamReader, "spark_ds_serialize_stream_reader", streamReader); } void commitOffset(long streamReader, String end) throws Throwable { - commitOffset.invokeExact(streamReader, end); + MethodHandle f = require(commitOffset, "spark_ds_commit_offset"); + try (Memory memory = new Memory()) { + check((int) f.invokeExact(streamReader, memory.string(end)), "spark_ds_commit_offset"); + } } void closeStreamReader(long streamReader) throws Throwable { - closeStreamReader.invokeExact(streamReader); + require(closeStreamReader, "spark_ds_close_stream_reader").invokeExact(streamReader); } void read(byte[] readerState, byte[] partition, long streamAddress) throws Throwable { - read.invokeExact(readerState, partition, streamAddress); + MethodHandle f = require(read, "spark_ds_read"); + try (Memory memory = new Memory()) { + check((int) f.invokeExact(memory.bytes(readerState), (long) readerState.length, + memory.bytes(partition), (long) partition.length, streamAddress), "spark_ds_read"); + } } byte[] serializeWriter(long writer) throws Throwable { - return (byte[]) serializeWriter.invokeExact(writer); + return serialize(serializeWriter, "spark_ds_serialize_writer", writer); } void commit(long writer, long epochId, byte[][] messages) throws Throwable { - commit.invokeExact(writer, epochId, messages); + finishWrite(commit, "spark_ds_commit", writer, epochId, messages); } void abort(long writer, long epochId, byte[][] messages) throws Throwable { - abort.invokeExact(writer, epochId, messages); + finishWrite(abort, "spark_ds_abort", writer, epochId, messages); + } + + private void finishWrite( + MethodHandle function, + String name, + long writer, + long epochId, + byte[][] messages) throws Throwable { + MethodHandle f = require(function, name); + try (Memory memory = new Memory()) { + long pointers = memory.zeroed(8L * Math.max(1, messages.length)); + long lengths = memory.zeroed(8L * Math.max(1, messages.length)); + for (int i = 0; i < messages.length; i++) { + if (messages[i] != null) { + Platform.putLong(null, pointers + 8L * i, memory.bytes(messages[i])); + Platform.putLong(null, lengths + 8L * i, messages[i].length); + } + } + check((int) f.invokeExact(writer, epochId, pointers, lengths, (long) messages.length), + name); + } } void closeWriter(long writer) throws Throwable { - closeWriter.invokeExact(writer); + require(closeWriter, "spark_ds_close_writer").invokeExact(writer); } long createDataWriter(byte[] writerState, int partitionId, long taskId, long epochId) throws Throwable { - return (long) createDataWriter.invokeExact(writerState, partitionId, taskId, epochId); + MethodHandle f = require(createDataWriter, "spark_ds_create_data_writer"); + try (Memory memory = new Memory()) { + long out = memory.zeroed(8); + check((int) f.invokeExact(memory.bytes(writerState), (long) writerState.length, + partitionId, taskId, epochId, out), "spark_ds_create_data_writer"); + return readHandle(out, "spark_ds_create_data_writer"); + } } void write(long dataWriter, long arrayAddress, long schemaAddress) throws Throwable { - write.invokeExact(dataWriter, arrayAddress, schemaAddress); + check((int) require(write, "spark_ds_write").invokeExact( + dataWriter, arrayAddress, schemaAddress), "spark_ds_write"); } byte[] commitDataWriter(long dataWriter) throws Throwable { - return (byte[]) commitDataWriter.invokeExact(dataWriter); + byte[] message = serialize(commitDataWriter, "spark_ds_commit_data_writer", dataWriter); + return message == null ? new byte[0] : message; } void abortDataWriter(long dataWriter) throws Throwable { - abortDataWriter.invokeExact(dataWriter); + check((int) require(abortDataWriter, "spark_ds_abort_data_writer").invokeExact(dataWriter), + "spark_ds_abort_data_writer"); + } + + /** + * Calls a function that returns a single byte array through a sink. Returns null if the + * function does not call the sink. + */ + private byte[] serialize(MethodHandle function, String name, long handle) throws Throwable { + MethodHandle f = require(function, name); + try (Sink sink = new Sink()) { + check((int) f.invokeExact(handle, Ffm.SINK, sink.id), name); + byte[][] values = sink.all(); + if (values.length > 1) { + throw new IllegalStateException( + name + " returned " + values.length + " values instead of one."); + } + return values.length == 0 ? null : values[0]; + } + } + + @Override + public String toString() { + return "NativeLibrary(" + path + ")"; + } + + // --------------------------------------------------------------------------------------------- + // Memory + // --------------------------------------------------------------------------------------------- + + /** Native memory for the arguments of a call, freed after the call. */ + private static final class Memory implements AutoCloseable { + private final List allocations = new ArrayList<>(); + + long zeroed(long size) { + long address = Platform.allocateMemory(size); + allocations.add(address); + Platform.setMemory(address, (byte) 0, size); + return address; + } + + /** Copies the bytes. Empty bytes are passed as a valid pointer too. */ + long bytes(byte[] bytes) { + long address = zeroed(Math.max(1, bytes.length)); + Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address, bytes.length); + return address; + } + + /** Copies the string as a NUL-terminated UTF-8 string. */ + long string(String value) { + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + long address = zeroed(bytes.length + 1L); + Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address, bytes.length); + return address; + } + + /** Copies the strings, and returns the address of an array of pointers to them. */ + long strings(String[] values) { + long array = zeroed(8L * Math.max(1, values.length)); + for (int i = 0; i < values.length; i++) { + Platform.putLong(null, array + 8L * i, string(values[i])); + } + return array; + } + + @Override + public void close() { + allocations.forEach(Platform::freeMemory); + } + } + + private static byte[] readCString(long address) { + int length = 0; + while (Platform.getByte(null, address + length) != 0) { + length++; + } + byte[] bytes = new byte[length]; + Platform.copyMemory(null, address, bytes, Platform.BYTE_ARRAY_OFFSET, length); + return bytes; + } + + // --------------------------------------------------------------------------------------------- + // Sink + // --------------------------------------------------------------------------------------------- + + private static final AtomicLong nextSinkId = new AtomicLong(1); + private static final ConcurrentHashMap sinks = new ConcurrentHashMap<>(); + + /** + * Collects the byte arrays that a library returns during a call, through the C function + * {@link Ffm#SINK}. The library passes {@link #id} back to identify it. + */ + private static final class Sink implements AutoCloseable { + final long id = nextSinkId.getAndIncrement(); + final List values = new ArrayList<>(); + Throwable error; + + Sink() { + sinks.put(id, this); + } + + byte[][] all() throws Throwable { + if (error != null) { + throw error; + } + return values.toArray(new byte[0][]); + } + + @Override + public void close() { + sinks.remove(id); + } + } + + /** + * Called by the library, through the upcall stub {@link Ffm#SINK}, with a byte array to copy. + * It must not throw: an exception thrown by an upcall crashes the JVM. + */ + private static void sink(long id, long data, long length) { + Sink sink = sinks.get(id); + if (sink == null || sink.error != null) { + return; + } + try { + if (length < 0 || length > Integer.MAX_VALUE - 8 || (data == 0 && length > 0)) { + throw new IllegalStateException("The library returned an invalid byte array, of " + + length + " bytes at the address " + data + "."); + } + byte[] bytes = new byte[(int) length]; + Platform.copyMemory(null, data, bytes, Platform.BYTE_ARRAY_OFFSET, length); + sink.values.add(bytes); + } catch (Throwable e) { + sink.error = e; + } } } diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h new file mode 100644 index 0000000000000..679a4d99e3349 --- /dev/null +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -0,0 +1,505 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * spark_datasource.h: the binary interface between Spark and data sources implemented in native + * code, such as Rust or C++. Since Spark 4.4.0. + * + * A native data source is a shared library that exports the C functions declared here. Spark + * calls them with the Foreign Function and Memory API of Java, so native data sources require + * Java 22 or later. Spark looks the functions up in each library separately, so any number of + * libraries that export the same functions can be loaded side by side. Spark adapts a native data + * source to Data Source V2, with batch and micro-batch streaming reads, batch and streaming + * writes, and predicate, column and limit pushdown. + * + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of my_source is + * spark_datasource_my_source, such as libspark_datasource_my_source.so in /usr/local/lib on Linux. + * Spark looks for it in the directories of java.library.path, which include LD_LIBRARY_PATH on + * Linux, and in the lib directory of each installation prefix in PATH. A library can also be + * distributed in a native data source package: a zip file with the extension .sparkpkg that + * contains a manifest named spark-native-datasource.json and the library built for one or more + * platforms, added to a session with spark.addArtifact or under the paths in + * spark.sql.dataSource.native.paths. See the "Native Data Sources" page of the Spark SQL guide for + * details and examples. + * + * Lifecycle + * --------- + * Spark creates a data source with spark_ds_create on the driver for each schema inference, scan + * and write, and closes it right after. + * + * - Batch reads. On the driver, Spark creates a reader with spark_ds_create_reader, pushes + * operations down by calling, each at most once and in this order, spark_ds_push_predicates, + * spark_ds_push_limit and spark_ds_prune_columns, plans the partitions with + * spark_ds_partitions, serializes the state of the reader with spark_ds_serialize_reader, and + * closes it. Then it calls spark_ds_read on the executors, in one task for each partition. + * - Streaming reads. On the driver, Spark creates a stream reader with + * spark_ds_create_stream_reader and keeps it until the streaming query stops. For each + * micro-batch, it asks for the spark_ds_latest_offset, plans the partitions with + * spark_ds_stream_partitions and serializes the state with spark_ds_serialize_stream_reader, + * then calls spark_ds_read on the executors. Offsets are JSON strings defined by the library. + * When a streaming query restarts, from its checkpoint, Spark creates a new stream reader and + * may call spark_ds_stream_partitions again for the last range that it planned before, so the + * partitions of a range must read the same data each time. + * - Writes. On the driver, once it computed the data to write, right before it starts the tasks + * of the write (or of a micro-batch), Spark creates a writer with spark_ds_create_writer or + * spark_ds_create_stream_writer, and serializes its state with spark_ds_serialize_writer. If + * computing the data fails, Spark creates no writer. On the executors, each task creates a data + * writer with spark_ds_create_data_writer, calls spark_ds_write for each batch of at most + * spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends once it holds + * spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls spark_ds_commit_data_writer, or + * spark_ds_abort_data_writer if the task or the commit fails. Then the driver calls + * spark_ds_commit with the messages of the tasks if they all succeeded, or spark_ds_abort + * otherwise, including when spark_ds_commit fails, and closes the writer. Several attempts of + * a task may write the same partition, after a failure or with speculative execution, each + * with a data writer of its own. Spark lets the attempts of a partition commit one at a time, + * and passes to spark_ds_commit the message of one committed attempt for each partition, so + * the write consists of the data of these attempts only. It aborts the data writers of the + * other attempts, except on an executor that is lost, where a data writer is neither committed + * nor aborted, and may even have committed before its message reached the driver. A streaming + * write commits or aborts each micro-batch (epoch). Spark may commit the same epoch more than + * once, for example when a streaming query restarts after a failure, so committing an epoch + * that is already committed must have no effect, for exactly-once output. A micro-batch + * streaming query creates a new stream writer for each micro-batch, and Spark closes the + * previous stream writer of the query when the query creates a new one. Spark is not told when + * a streaming query stops, so it closes the last stream writer of the query once the writer is + * no longer referenced, or when the query is restarted. + * + * Conventions + * ----------- + * - Errors. A function that can fail returns 0 on success, or any other value on failure. It then + * sets the message that spark_ds_last_error returns, and Spark reports it as a + * NATIVE_DATA_SOURCE_ERROR. A C++ exception or a Rust panic must not cross the interface: catch + * it and return an error instead. + * - Handles. The create functions return an opaque, non-NULL handle in their last argument, + * typically a pointer. Spark releases each handle exactly once with the matching close + * function, or with spark_ds_commit_data_writer or spark_ds_abort_data_writer for data writers, + * though not always right after the last use of a stream writer (see Writes above). A reader or + * writer handle must not depend on the data source handle it was created from, because Spark + * closes the data source handle right after creating it. + * - Memory. What Spark passes to a function, such as strings and byte arrays, is only valid + * during the call. Strings are UTF-8 and NUL-terminated. The library returns strings and byte + * arrays by passing them to a SparkDsSink, which copies them. + * - State. Partitions, the state of readers and writers, and commit messages are byte arrays in a + * format defined by the library. Spark sends them to the executors and back. + * - Arrow data. Schemas and data are exchanged with the Arrow C data interface + * (https://arrow.apache.org/docs/format/CDataInterface.html), as structs allocated by Spark. + * For an output struct, the library initializes it and Spark takes ownership. For an input + * struct, the library may move it to take ownership; otherwise, Spark releases it after the + * call returns, so the library must not use it afterwards. The arrays that the library exports + * must have an offset of 0, at every level: Arrow Java ignores the offsets of the arrays it + * imports, so Spark rejects an array with another offset, such as a zero-copy slice. Copy a + * sliced array before exporting it. Spark also rejects an array with fewer values than its + * parent refers to, such as a column shorter than its batch. The buffers of an array must hold + * all of its values, as the C data interface requires: it does not carry the sizes of the + * buffers, so Spark cannot check them, and reads past the end of a buffer that is too short. + * - Schemas. Spark types map to Arrow types as in spark.sql.execution.arrow: for example, bigint + * is int64, string is utf8, timestamp is timestamp[us] with a time zone and timestamp_ntz is + * timestamp[us] without one. Data read from a library may also use large_utf8, utf8_view and + * the other variants of the same type, and any names for the fields of the entries of a map, + * but not unsigned integers, decimal256, list views or dictionary encoding. + * - Options. Option keys are passed in lower case, because Spark matches them + * case-insensitively. + * - Optional functions. A library only needs to export the functions for the operations it + * supports. If a function is missing, Spark treats the operation as unsupported, except for the + * pushdown functions and the serialize functions, which then push nothing down and return no + * state, and spark_ds_commit_offset, which then does nothing. spark_ds_abi_version is required. + * - Threads. Spark calls the library from multiple threads, and calls on different handles may + * run concurrently. Calls on the same handle never run concurrently. A handle may be released + * from a thread other than the threads that used it, such as a JVM Cleaner thread when Spark + * releases a handle once it is no longer referenced, for example the last stream writer of a + * streaming query. A SparkDsSink must be called on the thread of the call that it was passed + * to, during that call. + * + * Compatibility + * ------------- + * Spark binds the functions of a library by their names, so the names of the functions and their + * signatures are the interface. Within a SPARK_DS_ABI_VERSION: + * - the signature and the semantics of an existing function never change; + * - new functions may be added, only as optional functions, so that a library that does not + * export them keeps working as before; + * - new values may appear where this documentation says so, such as the names of pushed + * predicates. + * An incompatible change uses a new SPARK_DS_ABI_VERSION, with new names. A library reports the + * version it implements with spark_ds_abi_version, and Spark rejects a library of a version it + * does not support. Spark supports version 1 at least in all the releases of Spark 4 from 4.4.0. + */ +#ifndef SPARK_DATASOURCE_H +#define SPARK_DATASOURCE_H + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* The version of this interface. */ +#define SPARK_DS_ABI_VERSION 1 + +/* The Arrow C data interface: https://arrow.apache.org/docs/format/CDataInterface.html */ +#ifndef ARROW_C_DATA_INTERFACE +#define ARROW_C_DATA_INTERFACE + +#define ARROW_FLAG_DICTIONARY_ORDERED 1 +#define ARROW_FLAG_NULLABLE 2 +#define ARROW_FLAG_MAP_KEYS_SORTED 4 + +struct ArrowSchema { + const char* format; + const char* name; + const char* metadata; + int64_t flags; + int64_t n_children; + struct ArrowSchema** children; + struct ArrowSchema* dictionary; + void (*release)(struct ArrowSchema*); + void* private_data; +}; + +struct ArrowArray { + int64_t length; + int64_t null_count; + int64_t offset; + int64_t n_buffers; + int64_t n_children; + const void** buffers; + struct ArrowArray** children; + struct ArrowArray* dictionary; + void (*release)(struct ArrowArray*); + void* private_data; +}; + +#endif /* ARROW_C_DATA_INTERFACE */ + +#ifndef ARROW_C_STREAM_INTERFACE +#define ARROW_C_STREAM_INTERFACE + +struct ArrowArrayStream { + int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out); + int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out); + const char* (*get_last_error)(struct ArrowArrayStream*); + void (*release)(struct ArrowArrayStream*); + void* private_data; +}; + +#endif /* ARROW_C_STREAM_INTERFACE */ + +/* + * Receives a string or a byte array that the library returns, and copies it. `sink` is the + * argument passed with the sink to the function that returns it. + */ +typedef void (*SparkDsSink)(void* sink, const uint8_t* data, int64_t length); + +/* + * Returns the message of the last error on the calling thread, as a NUL-terminated UTF-8 string + * that stays valid until the next call on that thread. Optional: without it, Spark only reports + * the value that the function returned. + */ +const char* spark_ds_last_error(void); + +/* Returns the version of this interface that the library implements, SPARK_DS_ABI_VERSION. */ +int32_t spark_ds_abi_version(void); + +/* --------------------------------------------------------------------------------------------- + * Data source (driver) + * --------------------------------------------------------------------------------------------- */ + +/* + * Creates a data source with the options given by the user, such as + * spark.read.option("key", "value"). + * + * name: the name of the data source, as listed in the package manifest. A library can implement + * several data sources. + * option_keys, option_values: the num_options options, with the keys in lower case + * source: set to the data source handle + */ +int spark_ds_create(const char* name, const char* const* option_keys, + const char* const* option_values, int64_t num_options, void** source); + +/* + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. + * + * out: initialized with a struct type whose fields are the columns + */ +int spark_ds_schema(void* source, struct ArrowSchema* out); + +/* + * Creates a reader for a batch scan. + * + * schema: the schema to read: the one the user specified, or the one returned by spark_ds_schema + * reader: set to the reader handle + */ +int spark_ds_create_reader(void* source, struct ArrowSchema* schema, void** reader); + +/* + * Creates a reader for a micro-batch streaming scan. + * + * schema: the schema to read: the one the user specified, or the one returned by spark_ds_schema + * reader: set to the stream reader handle + */ +int spark_ds_create_stream_reader(void* source, struct ArrowSchema* schema, void** reader); + +/* + * Creates a writer for a batch write. + * + * schema: the schema of the data to write + * overwrite: whether to replace the existing data (save mode "overwrite") instead of appending to + * it + * writer: set to the writer handle + */ +int spark_ds_create_writer(void* source, struct ArrowSchema* schema, bool overwrite, + void** writer); + +/* + * Creates a writer for a streaming write. + * + * schema: the schema of the data to write + * overwrite: whether to replace the existing data in each micro-batch (output mode "complete") + * instead of appending to it + * writer: set to the writer handle + */ +int spark_ds_create_stream_writer(void* source, struct ArrowSchema* schema, bool overwrite, + void** writer); + +/* Releases a data source handle. */ +void spark_ds_close(void* source); + +/* --------------------------------------------------------------------------------------------- + * Batch reader (driver) + * --------------------------------------------------------------------------------------------- */ + +/* + * Pushes predicates down to the reader. The predicates are combined with AND. + * + * Each predicate is a JSON expression tree. An expression is one of: + * - {"type": "column", "name": ["a", "b"]}: a column. A nested field has a name part for each + * level. + * - {"type": "literal", "dataType": "bigint", "value": 1}: a literal. The data type is a Spark SQL + * type name such as int, string or decimal(10,2). The value is null, a JSON boolean or number + * for boolean and numeric types, a string for string, a base64 string for binary, a decimal + * string for decimal, the days since the epoch for date, and the microseconds since the epoch + * for timestamp and timestamp_ntz. + * - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in + * org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=> (equal, + * with null equal to null), <, <=, > and >= with two children, IN with the value to look for + * and then the values to look in, IS_NULL and IS_NOT_NULL with one child, STARTS_WITH, + * ENDS_WITH and CONTAINS with the string and the string to look for, AND and OR with two + * predicates, and NOT with one. The children of the other functions are columns, literals or + * predicates. + * Spark does not pass predicates that it cannot express this way, such as predicates that use + * arithmetic or other functions, nor predicates on columns that contain strings with a + * non-binary collation, such as UTF8_LCASE: strings are compared by their UTF-8 bytes. A later + * version of Spark may pass predicates with other names without changing SPARK_DS_ABI_VERSION, + * so a reader must not report a predicate as evaluated completely if it contains a name that the + * reader does not know, or a data type that it does not support. + * + * A predicate that the reader evaluates completely must follow the semantics of Spark SQL. In + * particular, for floating-point columns, NaN is equal to NaN and greater than any other value, + * including positive infinity, and -0.0 is equal to 0.0, unlike IEEE 754 comparisons. For + * example, d > 1.0 is true for a NaN value of d. Comparisons with null are null, so a row only + * matches if the predicate is true. + * + * predicates: the count predicates, as JSON + * evaluated: an array of count values, all false, to set to whether the reader evaluates each + * predicate completely, so that Spark does not have to evaluate it again. The reader + * can still use the other predicates to skip data. + */ +int spark_ds_push_predicates(void* reader, const char* const* predicates, int64_t count, + bool* evaluated); + +/* + * Pushes a LIMIT down to the reader. + * + * applied: set to whether the reader returns at most `limit` rows in total. Spark applies the + * limit again after reading either way. + */ +int spark_ds_push_limit(void* reader, int32_t limit, bool* applied); + +/* + * Prunes the columns to read. + * + * columns: the names of the count top-level columns that the query needs + * pruned: set to whether spark_ds_read returns exactly these columns, in this order. If not, it + * returns all the columns of the schema to read. + */ +int spark_ds_prune_columns(void* reader, const char* const* columns, int64_t count, + bool* pruned); + +/* + * Plans the partitions of the scan. Spark reads each partition in a separate task. + * + * add_partition: called with each partition, serialized in a format defined by the library + */ +int spark_ds_partitions(void* reader, SparkDsSink add_partition, void* sink); + +/* + * Serializes the state of the reader that spark_ds_read needs on the executors, such as the + * pushed predicates. Spark calls it after spark_ds_partitions. + * + * set_state: called once with the state, or not at all if there is none + */ +int spark_ds_serialize_reader(void* reader, SparkDsSink set_state, void* sink); + +/* Releases a batch reader handle. */ +void spark_ds_close_reader(void* reader); + +/* --------------------------------------------------------------------------------------------- + * Stream reader (driver) + * --------------------------------------------------------------------------------------------- */ + +/* + * Returns the offset to start from when the streaming query has no checkpoint. + * + * set_offset: called once with the offset, as JSON + */ +int spark_ds_initial_offset(void* reader, SparkDsSink set_offset, void* sink); + +/* + * Returns the most recent offset available. + * + * set_offset: called once with the offset, as JSON + */ +int spark_ds_latest_offset(void* reader, SparkDsSink set_offset, void* sink); + +/* + * Plans the partitions that read the data after `start` up to and including `end`, two offsets + * as JSON. + * + * add_partition: called with each partition, serialized in a format defined by the library + */ +int spark_ds_stream_partitions(void* reader, const char* start, const char* end, + SparkDsSink add_partition, void* sink); + +/* + * Serializes the state of the stream reader that spark_ds_read needs on the executors. Spark + * calls it after each call to spark_ds_stream_partitions. + * + * set_state: called once with the state, or not at all if there is none + */ +int spark_ds_serialize_stream_reader(void* reader, SparkDsSink set_state, void* sink); + +/* + * Tells the library that Spark has processed all the data up to and including `end`, an offset + * as JSON, and will not ask for it again. + */ +int spark_ds_commit_offset(void* reader, const char* end); + +/* Releases a stream reader handle, when the streaming query stops. */ +void spark_ds_close_stream_reader(void* reader); + +/* --------------------------------------------------------------------------------------------- + * Partition reader (executor) + * --------------------------------------------------------------------------------------------- */ + +/* + * Reads a partition of a batch or streaming scan. + * + * reader_state: the state returned by spark_ds_serialize_reader or + * spark_ds_serialize_stream_reader, empty if there is none + * partition: a partition returned by spark_ds_partitions or spark_ds_stream_partitions + * out: the stream to initialize. Each array of the stream is a struct array whose fields are the + * columns to read: the columns of the schema to read, or the pruned columns. Spark calls the + * stream from the thread of the task that reads the partition. get_next may block, for + * example on network I/O, but Spark cannot interrupt native code: when the task is killed, + * Spark only stops reading once get_next returns, and then releases the stream. So a + * get_next that blocks forever keeps the task running, unless the task reaper + * (spark.task.reaper.enabled and spark.task.reaper.killTimeout) kills the executor. Use + * timeouts for blocking operations. + */ +int spark_ds_read(const uint8_t* reader_state, int64_t reader_state_length, + const uint8_t* partition, int64_t partition_length, + struct ArrowArrayStream* out); + +/* --------------------------------------------------------------------------------------------- + * Writer (driver) + * --------------------------------------------------------------------------------------------- */ + +/* + * Serializes the state of the writer that spark_ds_create_data_writer needs on the executors. + * Spark calls it once, after creating the writer. + * + * set_state: called once with the state, or not at all if there is none + */ +int spark_ds_serialize_writer(void* writer, SparkDsSink set_state, void* sink); + +/* + * Commits a write, when all the tasks succeeded. + * + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * messages, message_lengths: the num_messages messages returned by spark_ds_commit_data_writer + * of each task + */ +int spark_ds_commit(void* writer, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages); + +/* + * Aborts a write, when it failed. + * + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * messages, message_lengths: the num_messages messages returned by spark_ds_commit_data_writer + * of each task. A message is NULL if its task did not commit. + */ +int spark_ds_abort(void* writer, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages); + +/* Releases a writer handle. */ +void spark_ds_close_writer(void* writer); + +/* --------------------------------------------------------------------------------------------- + * Data writer (executor) + * --------------------------------------------------------------------------------------------- */ + +/* + * Creates the writer for one task. + * + * writer_state: the state returned by spark_ds_serialize_writer, empty if there is none + * partition_id: the partition that the task writes + * task_id: the ID of the task attempt + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * data_writer: set to the data writer handle + */ +int spark_ds_create_data_writer(const uint8_t* writer_state, int64_t writer_state_length, + int32_t partition_id, int64_t task_id, int64_t epoch_id, + void** data_writer); + +/* + * Writes a batch of data. + * + * array: a struct array whose fields are the columns + * schema: the schema of the array + */ +int spark_ds_write(void* data_writer, struct ArrowArray* array, struct ArrowSchema* schema); + +/* + * Commits the data written by the task and releases the data writer handle. If it fails, the + * handle is not released: Spark then calls spark_ds_abort_data_writer with it, to abort the data + * written by the task and release the handle. + * + * set_message: called once with the message to pass to spark_ds_commit on the driver, or not at + * all for an empty message + */ +int spark_ds_commit_data_writer(void* data_writer, SparkDsSink set_message, void* sink); + +/* Aborts the data written by the task and releases the data writer handle. */ +int spark_ds_abort_data_writer(void* data_writer); + +#ifdef __cplusplus +} +#endif + +#endif /* SPARK_DATASOURCE_H */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index e2307c70c3a42..943da94be272e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -43,8 +43,8 @@ import org.apache.spark.unsafe.Platform import org.apache.spark.util.Utils /** - * A [[ColumnarDataSource]] implemented by a native library, through - * `org.apache.spark.sql.datasource.NativeBridge`. + * A [[ColumnarDataSource]] implemented by a native library, through the C interface in + * `org/apache/spark/sql/datasource/spark_datasource.h`. * * @param location where the library is on the driver. The readers and the writers distribute it * to the executors of the active session when they are created. @@ -684,7 +684,7 @@ object NativeHandle extends Logging { try { release(value) } catch { - case e @ (NonFatal(_) | _: LinkageError) => + case NonFatal(e) => logWarning("Failed to release a native data source handle.", e) } } @@ -720,7 +720,7 @@ object NativeCalls { f } catch { case e: Throwable with SparkThrowable => throw e - case e: UnsatisfiedLinkError => + case e: NativeLibrary.MissingFunctionException => throw QueryExecutionErrors.nativeDataSourceError( action, name, s"The library does not implement ${e.getMessage}.", e) case NonFatal(e) => @@ -743,7 +743,7 @@ object NativeCalls { try { Some(f) } catch { - case _: UnsatisfiedLinkError => None + case _: NativeLibrary.MissingFunctionException => None } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala index 8db977cb1fc8e..de14ec5409867 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -33,7 +33,6 @@ import com.google.common.cache.{Cache, CacheBuilder} import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.PATH import org.apache.spark.sql.Artifact -import org.apache.spark.sql.datasource.NativeBridge import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.util.Utils @@ -50,7 +49,7 @@ import org.apache.spark.util.Utils * } * }}} * - * @param abiVersion the version of [[NativeBridge]] that the libraries implement + * @param abiVersion the version of the native data source interface that the libraries implement * @param dataSources the names of the data sources that the libraries implement * @param libraries the path of the library in the package for each platform, as returned by * [[NativePlatform.current]] @@ -292,15 +291,24 @@ object NativePlatform { */ object NativeLibraries extends Logging { // The libraries of the packages, by the checksum of the package, and the installed libraries, - // by their canonical path: the JVM loads a library file in a single class loader. A library - // that is loaded but rejected stays loaded, so its error is kept: it cannot be loaded again. + // by their canonical path, so that each library file is loaded once. A library that is loaded + // but rejected stays loaded, so its error is kept: it is not loaded again. private val packagedLibraries = new ConcurrentHashMap[String, Try[NativeLibrary]]() private val installedLibraries = new ConcurrentHashMap[String, Try[NativeLibrary]]() private lazy val extractionDir: File = Utils.createTempDir(namePrefix = "native-datasources") /** Returns the library at the given location, and loads it if needed. */ - private[ffi] def get(location: NativeLibraryLocation): NativeLibrary = location match { + private[ffi] def get(location: NativeLibraryLocation): NativeLibrary = { + // The libraries are called with the FFM API, which is final only from Java 22. + if (!NativeLibrary.isSupported) { + throw QueryExecutionErrors.nativeDataSourceUnsupportedJavaVersionError( + NativeLibrary.MINIMUM_JAVA_VERSION, Runtime.version().feature()) + } + getSupported(location) + } + + private def getSupported(location: NativeLibraryLocation): NativeLibrary = location match { case pkg: NativeDataSourcePackage => getOrLoad(packagedLibraries, pkg.sha256)(loadPackaged(pkg)) case installed: InstalledNativeLibrary => @@ -362,18 +370,18 @@ object NativeLibraries extends Logging { val library = try { NativeLibrary.load(file.getPath) } catch { - case e: UnsatisfiedLinkError => throw invalidLibrary(e.getMessage, e) + case e: NativeLibrary.LoadException => throw invalidLibrary(e.getMessage, e) } Try { val abiVersion = try { library.abiVersion() } catch { - case e: UnsatisfiedLinkError => - throw invalidLibrary("It does not implement NativeBridge.abiVersion.", e) + case e: NativeLibrary.MissingFunctionException => + throw invalidLibrary("It does not implement spark_ds_abi_version.", e) } - if (abiVersion != NativeBridge.ABI_VERSION) { + if (abiVersion != NativeLibrary.ABI_VERSION) { throw unsupportedAbiVersion( - Map("version" -> abiVersion.toString, "supported" -> NativeBridge.ABI_VERSION.toString)) + Map("version" -> abiVersion.toString, "supported" -> NativeLibrary.ABI_VERSION.toString)) } library } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala index d5b09ce578a5b..e350b2ba7211b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -25,7 +25,6 @@ import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.PATH import org.apache.spark.sql.Artifact import org.apache.spark.sql.classic.SparkSession -import org.apache.spark.sql.datasource.NativeBridge import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.util.Utils @@ -62,7 +61,7 @@ object NativeDataSourceRegistry extends Logging { .map(NativeDataSourcePackage.tryRead).partitionMap(identity) // A package of a version of the interface that Spark does not support is unusable too. val (unsupported, usable) = valid.partitionMap { pkg => - if (pkg.manifest.abiVersion == NativeBridge.ABI_VERSION) Right(pkg) else Left(pkg) + if (pkg.manifest.abiVersion == NativeLibrary.ABI_VERSION) Right(pkg) else Left(pkg) } lookupPackage(name, usable) .orElse(lookupInstalledLibrary(name)) @@ -76,7 +75,7 @@ object NativeDataSourceRegistry extends Logging { "UNSUPPORTED_ABI_VERSION", Map( "version" -> pkg.manifest.abiVersion.toString, - "supported" -> NativeBridge.ABI_VERSION.toString)) + "supported" -> NativeLibrary.ABI_VERSION.toString)) } invalid.find(_.dataSources.exists(_.equalsIgnoreCase(name))).foreach(p => throw p.error) None diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala index 93876788182f2..e9e6f3815594f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala @@ -29,7 +29,7 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap /** * The Data Source V2 provider of the native data sources: data sources implemented by a native - * library through `org.apache.spark.sql.datasource.NativeBridge`. + * library through the C interface in `org/apache/spark/sql/datasource/spark_datasource.h`. */ class NativeDataSourceV2 extends NamedTableProvider { private var shortName: String = _ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala index 5a1c5f5892a70..2ab8497a0476a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala @@ -27,14 +27,14 @@ import org.apache.spark.sql.types._ /** * Encodes Data Source V2 predicates as the JSON expression trees passed to - * `NativeBridge.pushPredicates`. + * `spark_ds_push_predicates`. */ object NativePredicates { private val mapper = new ObjectMapper() /** * The names of the predicates that are passed to the libraries, as listed in - * `NativeBridge.pushPredicates`. They are part of the interface, so a predicate with any other + * `spark_ds_push_predicates`. They are part of the interface, so a predicate with any other * name, such as an arithmetic function, is not passed: its semantics, for example on overflow, * would be part of the interface too, and the names of Data Source V2 grow with new releases. */ diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index 4aba24753fca4..53468b84a8173 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -15,9 +15,9 @@ * limitations under the License. */ -// A native data source library used by NativeDataSourceSuite. It implements the JNI functions of -// org.apache.spark.sql.datasource.NativeBridge without any dependency, building the Arrow C data -// interface structs by hand. It provides three data sources: +// A native data source library used by NativeDataSourceSuite. It implements the C interface in +// spark_datasource.h without any other dependency, building the Arrow C data interface structs by +// hand. It provides three data sources: // // native_range Reads rows [0, end) with columns of several types. Supports predicate, // column and limit pushdown. With the option "extra", the schema has one more @@ -37,9 +37,10 @@ // NAME_PREFIX A prefix for the names of the data sources, "" by default. // ID_SIGN Multiplies the ids produced by native_range, 1 by default. // READ_ONLY Leaves out the functions for writes, streaming reads and pushdown. -// ABI_VERSION The version returned by abiVersion(), 1 by default. +// ABI_VERSION The version returned by spark_ds_abi_version, 1 by default. +// NO_LAST_ERROR Leaves out spark_ds_last_error. -#include +#include "spark_datasource.h" #include #include @@ -70,58 +71,6 @@ #define ABI_VERSION 1 #endif -#define SPARK_JNI(ret, method) \ - extern "C" JNIEXPORT ret JNICALL Java_org_apache_spark_sql_datasource_NativeBridge_##method - -// --------------------------------------------------------------------------------------------- -// The Arrow C data interface: https://arrow.apache.org/docs/format/CDataInterface.html -// --------------------------------------------------------------------------------------------- - -#ifndef ARROW_C_DATA_INTERFACE -#define ARROW_C_DATA_INTERFACE - -#define ARROW_FLAG_NULLABLE 2 - -struct ArrowSchema { - const char* format; - const char* name; - const char* metadata; - int64_t flags; - int64_t n_children; - struct ArrowSchema** children; - struct ArrowSchema* dictionary; - void (*release)(struct ArrowSchema*); - void* private_data; -}; - -struct ArrowArray { - int64_t length; - int64_t null_count; - int64_t offset; - int64_t n_buffers; - int64_t n_children; - const void** buffers; - struct ArrowArray** children; - struct ArrowArray* dictionary; - void (*release)(struct ArrowArray*); - void* private_data; -}; - -#endif // ARROW_C_DATA_INTERFACE - -#ifndef ARROW_C_STREAM_INTERFACE -#define ARROW_C_STREAM_INTERFACE - -struct ArrowArrayStream { - int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out); - int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out); - const char* (*get_last_error)(struct ArrowArrayStream*); - void (*release)(struct ArrowArrayStream*); - void* private_data; -}; - -#endif // ARROW_C_STREAM_INTERFACE - namespace { // --------------------------------------------------------------------------------------------- @@ -601,103 +550,53 @@ std::vector split(const std::string& text) { } // --------------------------------------------------------------------------------------------- -// JNI helpers +// Helpers for the C interface // --------------------------------------------------------------------------------------------- -void throwJava(JNIEnv* env, const std::string& message) { - jclass cls = env->FindClass("java/lang/RuntimeException"); - if (cls != nullptr) env->ThrowNew(cls, message.c_str()); -} - -// Runs `f`, and turns a C++ exception into a Java exception: none may cross the boundary. -template -T guarded(JNIEnv* env, T fallback, F&& f) { - try { - return f(); - } catch (const std::exception& e) { - throwJava(env, e.what()); - return fallback; - } -} +thread_local std::string last_error; +// Runs `f`, and turns a C++ exception into an error: none may cross the boundary. template -void guardedVoid(JNIEnv* env, F&& f) { +int guarded(F&& f) { try { f(); + return 0; } catch (const std::exception& e) { - throwJava(env, e.what()); - } -} - -std::string toString(JNIEnv* env, jstring value) { - if (value == nullptr) return ""; - const char* chars = env->GetStringUTFChars(value, nullptr); - std::string result(chars); - env->ReleaseStringUTFChars(value, chars); - return result; -} - -std::vector toStrings(JNIEnv* env, jobjectArray values) { - std::vector result; - jsize n = env->GetArrayLength(values); - for (jsize i = 0; i < n; i++) { - auto value = static_cast(env->GetObjectArrayElement(values, i)); - result.push_back(toString(env, value)); - env->DeleteLocalRef(value); + last_error = e.what(); + return 1; } - return result; } -std::string toBytes(JNIEnv* env, jbyteArray bytes) { - if (bytes == nullptr) return ""; - jsize n = env->GetArrayLength(bytes); - std::string result(static_cast(n), '\0'); - if (n > 0) env->GetByteArrayRegion(bytes, 0, n, reinterpret_cast(&result[0])); - return result; +std::vector toStrings(const char* const* values, int64_t count) { + return std::vector(values, values + count); } -jbyteArray newBytes(JNIEnv* env, const std::string& bytes) { - jbyteArray result = env->NewByteArray(static_cast(bytes.size())); - env->SetByteArrayRegion(result, 0, static_cast(bytes.size()), - reinterpret_cast(bytes.data())); - return result; +std::string toBytes(const uint8_t* data, int64_t length) { + return std::string(reinterpret_cast(data), static_cast(length)); } -jobjectArray newByteArrays(JNIEnv* env, const std::vector& values) { - jobjectArray result = - env->NewObjectArray(static_cast(values.size()), env->FindClass("[B"), nullptr); - for (size_t i = 0; i < values.size(); i++) { - jbyteArray value = newBytes(env, values[i]); - env->SetObjectArrayElement(result, static_cast(i), value); - env->DeleteLocalRef(value); - } - return result; +void sendBytes(SparkDsSink sink_function, void* sink, const std::string& bytes) { + sink_function(sink, reinterpret_cast(bytes.data()), + static_cast(bytes.size())); } [[maybe_unused]] -std::vector toByteStrings(JNIEnv* env, jobjectArray values, bool* has_null) { +std::vector toByteStrings(const uint8_t* const* messages, const int64_t* lengths, + int64_t count, bool* has_null) { std::vector result; - jsize n = env->GetArrayLength(values); - for (jsize i = 0; i < n; i++) { - auto value = static_cast(env->GetObjectArrayElement(values, i)); - if (value == nullptr) { + for (int64_t i = 0; i < count; i++) { + if (messages[i] == nullptr) { *has_null = true; } else { - result.push_back(toBytes(env, value)); - env->DeleteLocalRef(value); + result.push_back(toBytes(messages[i], lengths[i])); } } return result; } template -jlong toHandle(T* object) { - return reinterpret_cast(object); -} - -template -T* fromHandle(jlong handle) { - return reinterpret_cast(handle); +T* fromHandle(void* handle) { + return static_cast(handle); } // --------------------------------------------------------------------------------------------- @@ -771,7 +670,7 @@ struct Writer { std::string closed_marker; }; -[[maybe_unused]] std::string closedMarker(jlong epoch_id) { +[[maybe_unused]] std::string closedMarker(int64_t epoch_id) { return epoch_id < 0 ? "/_CLOSED" : "/_CLOSED_" + std::to_string(epoch_id); } @@ -843,51 +742,54 @@ std::vector listFiles(const std::string& dir, const std::string& pr } // namespace + // --------------------------------------------------------------------------------------------- -// NativeBridge +// spark_datasource.h // --------------------------------------------------------------------------------------------- -SPARK_JNI(jint, abiVersion)(JNIEnv*, jclass) { return ABI_VERSION; } - -SPARK_JNI(jlong, createDataSource)(JNIEnv* env, jclass, jstring name_value, jobjectArray keys, - jobjectArray values) { - return guarded(env, 0, [&] { - auto* source = new DataSource{toString(env, name_value), {}}; - std::unique_ptr owner(source); - auto key_strings = toStrings(env, keys); - auto value_strings = toStrings(env, values); - for (size_t i = 0; i < key_strings.size(); i++) { - source->options[key_strings[i]] = value_strings[i]; - } +extern "C" { + +#ifndef NO_LAST_ERROR +const char* spark_ds_last_error(void) { return last_error.c_str(); } +#endif + +int32_t spark_ds_abi_version(void) { return ABI_VERSION; } + +int spark_ds_create(const char* name_value, const char* const* option_keys, + const char* const* option_values, int64_t num_options, void** handle) { + return guarded([&] { + auto source = std::make_unique(DataSource{name_value, {}}); + auto keys = toStrings(option_keys, num_options); + auto values = toStrings(option_values, num_options); + for (size_t i = 0; i < keys.size(); i++) source->options[keys[i]] = values[i]; maybeFail(source->options, "create"); if (source->name != name("native_range") && source->name != name("native_sink") && source->name != name("native_counter")) { throw std::runtime_error("unknown data source " + source->name); } - return toHandle(owner.release()); + *handle = source.release(); }); } -SPARK_JNI(void, schema)(JNIEnv* env, jclass, jlong handle, jlong schema_address) { - guardedVoid(env, [&] { +int spark_ds_schema(void* handle, struct ArrowSchema* out) { + return guarded([&] { auto* source = fromHandle(handle); maybeFail(source->options, "schema"); if (source->name == name("native_sink")) { throw std::runtime_error("native_sink has no schema; it can only be written"); } bool counter = source->name == name("native_counter"); - exportSchema(allColumns(counter, get(source->options, "extra", "")), - reinterpret_cast(schema_address)); + exportSchema(allColumns(counter, get(source->options, "extra", "")), out); }); } -SPARK_JNI(jlong, createReader)(JNIEnv* env, jclass, jlong handle, jlong schema_address) { - return guarded(env, 0, [&] { +int spark_ds_create_reader(void* handle, struct ArrowSchema* schema, void** out) { + return guarded([&] { auto* source = fromHandle(handle); if (source->name != name("native_range")) { throw std::runtime_error(source->name + " does not support batch reads"); } - auto columns = importColumns(reinterpret_cast(schema_address)); + auto columns = importColumns(schema); for (const auto& column : columns) findColumn(column); auto* reader = new Reader(); reader->state["columns"] = join(columns); @@ -898,48 +800,45 @@ SPARK_JNI(jlong, createReader)(JNIEnv* env, jclass, jlong handle, jlong schema_a reader->state["fail"] = get(source->options, "fail", ""); reader->state["slice"] = get(source->options, "slice", ""); reader->state["short"] = get(source->options, "short", ""); - return toHandle(reader); + *out = reader; }); } #ifndef READ_ONLY -SPARK_JNI(jlong, createStreamReader)(JNIEnv* env, jclass, jlong handle, jlong schema_address) { - return guarded(env, 0, [&] { +int spark_ds_create_stream_reader(void* handle, struct ArrowSchema* schema, void** out) { + return guarded([&] { auto* source = fromHandle(handle); if (source->name != name("native_counter")) { throw std::runtime_error(source->name + " does not support streaming reads"); } auto* reader = new Reader(); - reader->state["columns"] = join(importColumns(reinterpret_cast(schema_address))); + reader->state["columns"] = join(importColumns(schema)); reader->state["max_offset"] = get(source->options, "max_offset", "10"); reader->state["step"] = get(source->options, "step", reader->state["max_offset"]); reader->state["current"] = "0"; reader->state["batch_size"] = get(source->options, "batch_size", "1000"); reader->state["commit_path"] = get(source->options, "commit_path", ""); - return toHandle(reader); + *out = reader; }); } #endif // READ_ONLY -SPARK_JNI(void, closeDataSource)(JNIEnv*, jclass, jlong handle) { - delete fromHandle(handle); -} +void spark_ds_close(void* handle) { delete fromHandle(handle); } #ifndef READ_ONLY -SPARK_JNI(jbooleanArray, pushPredicates)(JNIEnv* env, jclass, jlong handle, - jobjectArray predicates) { - return guarded(env, nullptr, [&] { +int spark_ds_push_predicates(void* handle, const char* const* predicates, int64_t count, + bool* evaluated) { + return guarded([&] { auto* reader = fromHandle(handle); - auto json = toStrings(env, predicates); - std::vector accepted; + auto json = toStrings(predicates, count); int64_t lo = getLong(reader->state, "lo", 0); int64_t hi = getLong(reader->state, "hi", 0); - for (const auto& text : json) { + for (size_t i = 0; i < json.size(); i++) { // Accepts "id " for the comparison operators. - Json predicate = JsonParser(text).parse(); + Json predicate = JsonParser(json[i]).parse(); const Json& children = predicate["children"]; std::string op = predicate["name"].text; bool comparison = op == ">" || op == ">=" || op == "<" || op == "<=" || op == "="; @@ -960,40 +859,38 @@ SPARK_JNI(jbooleanArray, pushPredicates)(JNIEnv* env, jclass, jlong handle, hi = std::min(hi, value + 1); } } - accepted.push_back(ok ? JNI_TRUE : JNI_FALSE); + evaluated[i] = ok; } reader->state["lo"] = std::to_string(lo); reader->state["hi"] = std::to_string(std::max(lo, hi)); reader->state["predicates"] = std::to_string(json.size()); - jbooleanArray result = env->NewBooleanArray(static_cast(accepted.size())); - env->SetBooleanArrayRegion(result, 0, static_cast(accepted.size()), accepted.data()); - return result; }); } -SPARK_JNI(jboolean, pushLimit)(JNIEnv* env, jclass, jlong handle, jint limit) { - return guarded(env, JNI_FALSE, [&] { +int spark_ds_push_limit(void* handle, int32_t limit, bool* applied) { + return guarded([&] { auto* reader = fromHandle(handle); - // NativeBridge calls pushLimit at most once. + // Spark calls spark_ds_push_limit at most once. if (reader->state.count("limit") > 0) throw std::runtime_error("pushLimit was called twice"); reader->state["limit"] = std::to_string(limit); - return JNI_TRUE; + *applied = true; }); } -SPARK_JNI(jboolean, pruneColumns)(JNIEnv* env, jclass, jlong handle, jobjectArray columns) { - return guarded(env, JNI_FALSE, [&] { - auto names = toStrings(env, columns); +int spark_ds_prune_columns(void* handle, const char* const* columns, int64_t count, + bool* pruned) { + return guarded([&] { + auto names = toStrings(columns, count); for (const auto& column : names) findColumn(column); fromHandle(handle)->state["columns"] = join(names); - return JNI_TRUE; + *pruned = true; }); } #endif // READ_ONLY -SPARK_JNI(jobjectArray, partitions)(JNIEnv* env, jclass, jlong handle) { - return guarded(env, nullptr, [&] { +int spark_ds_partitions(void* handle, SparkDsSink add_partition, void* sink) { + return guarded([&] { auto* reader = fromHandle(handle); maybeFail(reader->state, "plan"); int64_t lo = getLong(reader->state, "lo", 0); @@ -1004,79 +901,69 @@ SPARK_JNI(jobjectArray, partitions)(JNIEnv* env, jclass, jlong handle) { hi = std::min(hi, lo + getLong(reader->state, "limit", 0)); count = 1; } - std::vector partitions; int64_t step = std::max(1, (hi - lo + count - 1) / count); for (int64_t start = lo; start < hi; start += step) { int64_t stop = std::min(hi, start + step); - partitions.push_back(std::to_string(start) + "," + std::to_string(stop)); + sendBytes(add_partition, sink, std::to_string(start) + "," + std::to_string(stop)); } - return newByteArrays(env, partitions); }); } -SPARK_JNI(jbyteArray, serializeReader)(JNIEnv* env, jclass, jlong handle) { - return guarded(env, nullptr, [&] { - return newBytes(env, serialize(fromHandle(handle)->state)); - }); +int spark_ds_serialize_reader(void* handle, SparkDsSink set_state, void* sink) { + return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); } -SPARK_JNI(void, closeReader)(JNIEnv*, jclass, jlong handle) { delete fromHandle(handle); } +void spark_ds_close_reader(void* handle) { delete fromHandle(handle); } #ifndef READ_ONLY -SPARK_JNI(jstring, initialOffset)(JNIEnv* env, jclass, jlong) { - return env->NewStringUTF("{\"offset\":0}"); +int spark_ds_initial_offset(void*, SparkDsSink set_offset, void* sink) { + return guarded([&] { sendBytes(set_offset, sink, "{\"offset\":0}"); }); } // Each call makes `step` more rows available, up to `max_offset`. -SPARK_JNI(jstring, latestOffset)(JNIEnv* env, jclass, jlong handle) { - return guarded(env, nullptr, [&] { +int spark_ds_latest_offset(void* handle, SparkDsSink set_offset, void* sink) { + return guarded([&] { Options& state = fromHandle(handle)->state; int64_t latest = std::min(getLong(state, "max_offset", 10), getLong(state, "current", 0) + getLong(state, "step", 10)); state["current"] = std::to_string(latest); - return env->NewStringUTF(("{\"offset\":" + std::to_string(latest) + "}").c_str()); + sendBytes(set_offset, sink, "{\"offset\":" + std::to_string(latest) + "}"); }); } -SPARK_JNI(jobjectArray, streamPartitions)(JNIEnv* env, jclass, jlong, jstring start, - jstring end) { - return guarded(env, nullptr, [&] { - auto lo = static_cast(JsonParser(toString(env, start)).parse()["offset"].number); - auto hi = static_cast(JsonParser(toString(env, end)).parse()["offset"].number); - std::vector partitions; - if (hi > lo) partitions.push_back(std::to_string(lo) + "," + std::to_string(hi)); - return newByteArrays(env, partitions); +int spark_ds_stream_partitions(void*, const char* start, const char* end, + SparkDsSink add_partition, void* sink) { + return guarded([&] { + auto lo = static_cast(JsonParser(start).parse()["offset"].number); + auto hi = static_cast(JsonParser(end).parse()["offset"].number); + if (hi > lo) sendBytes(add_partition, sink, std::to_string(lo) + "," + std::to_string(hi)); }); } -SPARK_JNI(jbyteArray, serializeStreamReader)(JNIEnv* env, jclass, jlong handle) { - return guarded(env, nullptr, [&] { - return newBytes(env, serialize(fromHandle(handle)->state)); - }); +int spark_ds_serialize_stream_reader(void* handle, SparkDsSink set_state, void* sink) { + return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); } -SPARK_JNI(void, commitOffset)(JNIEnv* env, jclass, jlong handle, jstring end) { - guardedVoid(env, [&] { +int spark_ds_commit_offset(void* handle, const char* end) { + return guarded([&] { std::string path = get(fromHandle(handle)->state, "commit_path", ""); - if (!path.empty()) writeFile(path, toString(env, end)); + if (!path.empty()) writeFile(path, end); }); } -SPARK_JNI(void, closeStreamReader)(JNIEnv*, jclass, jlong handle) { - delete fromHandle(handle); -} +void spark_ds_close_stream_reader(void* handle) { delete fromHandle(handle); } #endif // READ_ONLY -SPARK_JNI(void, read)(JNIEnv* env, jclass, jbyteArray reader_state, jbyteArray partition, - jlong stream_address) { - guardedVoid(env, [&] { - Options state = deserialize(toBytes(env, reader_state)); +int spark_ds_read(const uint8_t* reader_state, int64_t reader_state_length, + const uint8_t* partition, int64_t partition_length, + struct ArrowArrayStream* stream) { + return guarded([&] { + Options state = deserialize(toBytes(reader_state, reader_state_length)); maybeFail(state, "open"); - auto bounds = split(toBytes(env, partition)); + auto bounds = split(toBytes(partition, partition_length)); if (bounds.size() != 2) throw std::runtime_error("invalid partition"); - auto* stream = reinterpret_cast(stream_address); stream->get_schema = streamGetSchema; stream->get_next = streamGetNext; stream->get_last_error = streamGetLastError; @@ -1091,18 +978,18 @@ SPARK_JNI(void, read)(JNIEnv* env, jclass, jbyteArray reader_state, jbyteArray p #ifndef READ_ONLY +} // extern "C" + namespace { -jlong createAnyWriter(JNIEnv* env, jlong handle, jlong schema_address, jboolean overwrite) { - return guarded(env, 0, [&] { +int createAnyWriter(void* handle, struct ArrowSchema* schema, bool overwrite, void** out) { + return guarded([&] { auto* source = fromHandle(handle); if (source->name != name("native_sink")) { throw std::runtime_error(source->name + " does not support writes"); } // Spark releases the schema after the call, since it is not moved here. - auto* schema = reinterpret_cast(schema_address); if (schema->n_children == 0) throw std::runtime_error("no columns to write"); - auto* writer = new Writer(); // The location of a table in a catalog is passed as a URI. std::string path = get(source->options, "path", ""); if (path.rfind("file://", 0) == 0) { @@ -1110,46 +997,44 @@ jlong createAnyWriter(JNIEnv* env, jlong handle, jlong schema_address, jboolean } else if (path.rfind("file:", 0) == 0) { path = path.substr(5); } + if (path.empty()) throw std::runtime_error("the option 'path' is required"); + auto* writer = new Writer(); writer->state["path"] = path; writer->state["overwrite"] = overwrite ? "true" : "false"; writer->state["fail"] = get(source->options, "fail", ""); - if (writer->state["path"].empty()) { - delete writer; - throw std::runtime_error("the option 'path' is required"); - } // Records that a writer was created, which Spark then commits or aborts. makeDirs(path); writeFile(path + "/_CREATED", ""); - return toHandle(writer); + *out = writer; }); } } // namespace -SPARK_JNI(jlong, createWriter)(JNIEnv* env, jclass, jlong handle, jlong schema_address, - jboolean overwrite) { - return createAnyWriter(env, handle, schema_address, overwrite); +extern "C" { + +int spark_ds_create_writer(void* handle, struct ArrowSchema* schema, bool overwrite, + void** writer) { + return createAnyWriter(handle, schema, overwrite, writer); } -SPARK_JNI(jlong, createStreamWriter)(JNIEnv* env, jclass, jlong handle, jlong schema_address, - jboolean overwrite) { - return createAnyWriter(env, handle, schema_address, overwrite); +int spark_ds_create_stream_writer(void* handle, struct ArrowSchema* schema, bool overwrite, + void** writer) { + return createAnyWriter(handle, schema, overwrite, writer); } -SPARK_JNI(jbyteArray, serializeWriter)(JNIEnv* env, jclass, jlong handle) { - return guarded(env, nullptr, [&] { - return newBytes(env, serialize(fromHandle(handle)->state)); - }); +int spark_ds_serialize_writer(void* handle, SparkDsSink set_state, void* sink) { + return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); } -SPARK_JNI(void, commit)(JNIEnv* env, jclass, jlong handle, jlong epoch_id, - jobjectArray messages) { - guardedVoid(env, [&] { +int spark_ds_commit(void* handle, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages) { + return guarded([&] { auto* writer = fromHandle(handle); std::string path = writer->state["path"]; maybeFail(writer->state, "commit"); bool has_null = false; - auto temp_files = toByteStrings(env, messages, &has_null); + auto temp_files = toByteStrings(messages, message_lengths, num_messages, &has_null); if (has_null) throw std::runtime_error("commit got a null message"); if (writer->state["overwrite"] == "true") { for (const auto& file : listFiles(path, "part-")) unlink((path + "/" + file).c_str()); @@ -1164,48 +1049,47 @@ SPARK_JNI(void, commit)(JNIEnv* env, jclass, jlong handle, jlong epoch_id, }); } -SPARK_JNI(void, abort)(JNIEnv* env, jclass, jlong handle, jlong epoch_id, - jobjectArray messages) { - guardedVoid(env, [&] { +int spark_ds_abort(void* handle, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages) { + return guarded([&] { auto* writer = fromHandle(handle); writer->closed_marker = closedMarker(epoch_id); bool has_null = false; - for (const auto& temp_file : toByteStrings(env, messages, &has_null)) { + for (const auto& temp_file : + toByteStrings(messages, message_lengths, num_messages, &has_null)) { unlink(temp_file.c_str()); } writeFile(writer->state["path"] + "/_ABORTED", ""); }); } -SPARK_JNI(void, closeWriter)(JNIEnv* env, jclass, jlong handle) { +void spark_ds_close_writer(void* handle) { std::unique_ptr writer(fromHandle(handle)); - guardedVoid(env, [&] { + guarded([&] { if (!writer->closed_marker.empty()) { writeFile(writer->state["path"] + writer->closed_marker, ""); } }); } -SPARK_JNI(jlong, createDataWriter)(JNIEnv* env, jclass, jbyteArray writer_state, - jint partition_id, jlong task_id, jlong epoch_id) { - return guarded(env, 0, [&] { - Options state = deserialize(toBytes(env, writer_state)); +int spark_ds_create_data_writer(const uint8_t* writer_state, int64_t writer_state_length, + int32_t partition_id, int64_t task_id, int64_t epoch_id, + void** out) { + return guarded([&] { + Options state = deserialize(toBytes(writer_state, writer_state_length)); std::string dir = state["path"] + "/_temporary"; makeDirs(dir); std::string temp_path = dir + "/" + std::to_string(epoch_id) + "-" + std::to_string(partition_id) + "-" + std::to_string(task_id) + ".csv"; FILE* file = std::fopen(temp_path.c_str(), "w"); if (file == nullptr) throw std::runtime_error("cannot write " + temp_path); - return toHandle(new DataWriter{temp_path, file, state}); + *out = new DataWriter{temp_path, file, state}; }); } -SPARK_JNI(void, write)(JNIEnv* env, jclass, jlong handle, jlong array_address, - jlong schema_address) { - guardedVoid(env, [&] { +int spark_ds_write(void* handle, struct ArrowArray* array, struct ArrowSchema* schema) { + return guarded([&] { auto* writer = fromHandle(handle); - auto* array = reinterpret_cast(array_address); - auto* schema = reinterpret_cast(schema_address); // The batch is released here, so Spark does not release it after the call. std::unique_ptr array_owner( array, [](ArrowArray* a) { a->release(a); }); @@ -1224,22 +1108,25 @@ SPARK_JNI(void, write)(JNIEnv* env, jclass, jlong handle, jlong array_address, }); } -SPARK_JNI(jbyteArray, commitDataWriter)(JNIEnv* env, jclass, jlong handle) { - return guarded(env, nullptr, [&] { +int spark_ds_commit_data_writer(void* handle, SparkDsSink set_message, void* sink) { + return guarded([&] { auto* writer = fromHandle(handle); - // If the commit fails, the data writer stays valid: Spark aborts it with abortDataWriter. + // If the commit fails, the data writer stays valid: Spark aborts it with + // spark_ds_abort_data_writer. maybeFail(writer->state, "commitDataWriter"); - jbyteArray message = newBytes(env, writer->temp_path); + sendBytes(set_message, sink, writer->temp_path); std::fclose(writer->file); delete writer; - return message; }); } -SPARK_JNI(void, abortDataWriter)(JNIEnv*, jclass, jlong handle) { +int spark_ds_abort_data_writer(void* handle) { std::unique_ptr writer(fromHandle(handle)); std::fclose(writer->file); unlink(writer->temp_path.c_str()); + return 0; } #endif // READ_ONLY + +} // extern "C" diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala index d261654f7cede..fa594f2c20ff6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala @@ -70,7 +70,7 @@ class NativeDataSourceClusterSuite extends QueryTest { } private def nativeTest(name: String)(f: => Unit): Unit = test(name) { - assume(toolchain.isDefined, "A C++ compiler and the JNI headers are required.") + assume(toolchain.isDefined, "A C++ compiler and Java 22 or later are required.") try f finally SparkSession.setActiveSession(spark) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 694c4d777db68..70842f4add1f3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -48,7 +48,7 @@ import org.apache.spark.util.Utils /** * Tests native data sources with the library in `native-datasource/test_native_datasource.cc`, * compiled with the C++ compiler of the machine by [[NativeDataSourceTestUtils]]. The tests that - * load a library are skipped when there is no C++ compiler or no JNI headers. + * load a library are skipped when there is no C++ compiler, or before Java 22. */ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { import NativeDataSourceTestUtils._ @@ -65,6 +65,8 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { compile("test_native_datasource_b", "-DNAME_PREFIX=\"b_\"", "-DID_SIGN=-1") private lazy val abiV2Library = compile("test_native_datasource_v2", "-DABI_VERSION=2", "-DNAME_PREFIX=\"v2_\"") + private lazy val noLastErrorLibrary = + compile("test_native_datasource_ne", "-DNO_LAST_ERROR", "-DNAME_PREFIX=\"ne_\"") private lazy val installedLibrary = compile("test_native_datasource_installed", "-DNAME_PREFIX=\"installed_\"") private lazy val installedNegatingLibrary = compile( @@ -83,7 +85,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } private def nativeTest(name: String)(f: => Unit): Unit = test(name) { - assume(toolchain.isDefined, "A C++ compiler and the JNI headers are required.") + assume(toolchain.isDefined, "A C++ compiler and Java 22 or later are required.") try f finally SparkSession.setActiveSession(spark) } @@ -250,6 +252,23 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(released == Seq(1L, 3L, 2L)) } + test("native data sources require Java 22 or later") { + assume(!NativeLibrary.isSupported, "The Java version supports native data sources.") + val dir = Utils.createTempDir() + // Not loaded: Spark fails before it loads the library. + Files.writeString( + new File(dir, NativeDataSourceRegistry.libraryFileName("old_java")).toPath, "unused") + withLibraryPath(dir) { + checkError( + exception = intercept[SparkUnsupportedOperationException] { + spark.read.format("old_java").load() + }, + condition = "NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION", + parameters = Map( + "required" -> "22", "current" -> Runtime.version().feature().toString)) + } + } + nativeTest("disable native data sources") { val conf = new SQLConf conf.setConf(SQLConf.NATIVE_DATA_SOURCE_PATHS, Seq(defaultPackage.getPath)) @@ -569,6 +588,18 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { check("plan a scan of", "native_sink does not support batch reads", "native_sink") { session => session.read.schema("id BIGINT").format("native_sink").load().collect() } + // Without spark_ds_last_error, only the value that the function returned is reported. + val e = intercept[Exception] { + newSession(createPackage(noLastErrorLibrary, dataSources("ne_"))) + .read.format("ne_native_range").option("fail", "create").load() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "create", + "name" -> "ne_native_range", + "msg" -> "spark_ds_create failed with the error code 1.")) } nativeTest("operations that the library does not implement") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala index d638f782bc6cf..cc2ec905b06ad 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala @@ -29,7 +29,7 @@ import org.apache.spark.util.Utils * Builds the native data source library of the tests, * `native-datasource/test_native_datasource.cc`, with the C++ compiler of the machine, and * packages or installs it. The tests that load the library are skipped when [[toolchain]] is not - * defined. + * defined, which includes the versions of Java that cannot load native data sources. */ object NativeDataSourceTestUtils { @@ -41,17 +41,19 @@ object NativeDataSourceTestUtils { // each package has its own name. private val packageCount = new AtomicInteger() - /** The C++ compiler and the JNI include directories, if available. */ + /** + * The C++ compiler, if available on a version of Java that can load native data sources, and + * the include directory with `spark_datasource.h`. + */ lazy val toolchain: Option[(String, Seq[String])] = { - val include = new File(System.getProperty("java.home"), "include") - val platformInclude = Option(include.listFiles()).toSeq.flatten - .find(dir => new File(dir, "jni_md.h").isFile) val compiler = Seq("c++", "g++", "clang++").find(Utils.checkCommandAvailable) - for { - compiler <- compiler - platformInclude <- platformInclude - if new File(include, "jni.h").isFile - } yield (compiler, Seq(include.getPath, platformInclude.getPath)) + compiler.filter(_ => NativeLibrary.isSupported).map { compiler => + val header = new File(buildDir, "spark_datasource.h") + val in = getClass.getClassLoader + .getResourceAsStream("org/apache/spark/sql/datasource/spark_datasource.h") + Utils.tryWithResource(in)(Files.copy(_, header.toPath)) + (compiler, Seq(buildDir.getPath)) + } } /** From 56febbb50fff81000221982e2fc27e7c0e12d69a Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Tue, 13 Oct 2026 15:43:10 +0900 Subject: [PATCH 15/27] Report FFM loading failures clearly and clarify the C interface - Report a library that cannot be loaded because native access is denied, or because the FFM API cannot be initialized, as a load error with its root cause. - Report an empty spark_ds_last_error message with the error code. - Never throw from the sink upcall, and remove an unused field. - Document in spark_datasource.h that the close functions cannot fail, that NULL handles are errors, and the name passed for an installed library. - Show how to get spark_datasource.h from the spark-sql JAR, and catch all C++ exceptions in the C++ example. Co-authored-by: Isaac --- docs/sql-native-data-sources.md | 17 +++++-- .../datasources/v2/ffi/NativeLibrary.java | 51 ++++++++++++------- .../spark/sql/datasource/spark_datasource.h | 26 +++++----- 3 files changed, 58 insertions(+), 36 deletions(-) diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 4276305a92a1b..2ce06950bcb55 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -50,7 +50,7 @@ Java, a query that uses a native data source fails with A native data source library exports the functions of `spark_datasource.h`, such as `spark_ds_create`, as plain C functions: it needs no JVM headers or bindings. Spark looks the functions up in each library separately, so any number of native data sources can be loaded side -by side, although they export the same functions. +by side, even though they export the same functions. Spark calls the functions of the library as follows: @@ -430,8 +430,13 @@ EOF The same data source in C++, with [Arrow C++](https://arrow.apache.org/docs/cpp/), which exports schemas and streams with `arrow::ExportSchema` and `arrow::ExportRecordBatchReader`. It includes -`spark_datasource.h`, which is also in the `spark-sql` JAR, under -`org/apache/spark/sql/datasource/`: +`spark_datasource.h`, which is also in the `spark-sql` JAR of a Spark distribution: + +```bash +unzip -j "$SPARK_HOME"/jars/spark-sql_*.jar org/apache/spark/sql/datasource/spark_datasource.h +``` + +`cpp_range.cc`: {% raw %} ```cpp @@ -474,6 +479,9 @@ int Call(F&& f) { } catch (const std::exception& e) { last_error = e.what(); return 1; + } catch (...) { + last_error = "unknown error"; + return 1; } } @@ -545,8 +553,7 @@ Build the library, here on macOS with Arrow C++ 24, which requires C++20, and in package it as above: ```bash -c++ -std=c++20 -O2 -shared -fPIC \ - -I/path/to/spark_datasource.h/dir -I"$ARROW_HOME/include" -L"$ARROW_HOME/lib" -larrow \ +c++ -std=c++20 -O2 -shared -fPIC -I. -I"$ARROW_HOME/include" -L"$ARROW_HOME/lib" -larrow \ -o libspark_datasource_cpp_range.dylib cpp_range.cc cp libspark_datasource_cpp_range.dylib /usr/local/lib/ ``` diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java index d2337d8748e12..e7c9129459fa1 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -63,7 +63,6 @@ static boolean isSupported() { // The C types of the parameters and the results of the functions. private enum C { VOID, BOOL, INT32, INT64, POINTER } - private final String path; private final MethodHandle lastError; private final MethodHandle abiVersion; private final MethodHandle create; @@ -95,8 +94,7 @@ private enum C { VOID, BOOL, INT32, INT64, POINTER } private final MethodHandle commitDataWriter; private final MethodHandle abortDataWriter; - private NativeLibrary(String path, Object lookup) throws Throwable { - this.path = path; + private NativeLibrary(Object lookup) throws Throwable { C i32 = C.INT32; C i64 = C.INT64; C ptr = C.POINTER; @@ -149,10 +147,16 @@ static NativeLibrary load(String path) throws Throwable { Object lookup; try { lookup = Ffm.LIBRARY_LOOKUP.invoke(Path.of(path), Ffm.GLOBAL_ARENA); - } catch (IllegalArgumentException e) { + } catch (IllegalArgumentException | IllegalCallerException e) { + // The library is not a valid library, or native access is not enabled for Spark. throw new LoadException(e.getMessage(), e); + } catch (LinkageError e) { + // The FFM API could not be initialized, see Ffm: report why, such as an + // IllegalCallerException when native access is not enabled for Spark. + throw new LoadException( + "Cannot use the Foreign Function and Memory API: " + rootCause(e), e); } - return new NativeLibrary(path, lookup); + return new NativeLibrary(lookup); } /** @@ -257,7 +261,9 @@ private static final class Ffm { .invoke(LINKER, target, sinkDescriptor, GLOBAL_ARENA, NO_OPTIONS); SINK = (long) ADDRESS_OF.invoke(stub); } catch (Throwable e) { - throw new IllegalStateException("Cannot initialize the FFM API.", e); + // The message names the cause, because the JVM only keeps this message for the later + // uses of the class. + throw new IllegalStateException("Cannot initialize the FFM API: " + rootCause(e), e); } } @@ -281,6 +287,14 @@ static Object layout(C type) { } } + private static Throwable rootCause(Throwable e) { + Throwable cause = e; + while (cause.getCause() != null && cause.getCause() != cause) { + cause = cause.getCause(); + } + return cause; + } + /** Thrown when a library cannot be loaded. */ static final class LoadException extends RuntimeException { LoadException(String message, Throwable cause) { @@ -311,10 +325,11 @@ private void check(int status, String name) throws Throwable { if (status == 0) { return; } - long message = lastError == null ? 0 : (long) lastError.invokeExact(); - throw new RuntimeException(message == 0 + long address = lastError == null ? 0 : (long) lastError.invokeExact(); + String message = address == 0 ? "" : new String(readCString(address), StandardCharsets.UTF_8); + throw new RuntimeException(message.isEmpty() ? name + " failed with the error code " + status + "." - : new String(readCString(message), StandardCharsets.UTF_8)); + : message); } private static long readHandle(long out, String name) { @@ -571,11 +586,6 @@ private byte[] serialize(MethodHandle function, String name, long handle) throws } } - @Override - public String toString() { - return "NativeLibrary(" + path + ")"; - } - // --------------------------------------------------------------------------------------------- // Memory // --------------------------------------------------------------------------------------------- @@ -669,11 +679,12 @@ public void close() { * It must not throw: an exception thrown by an upcall crashes the JVM. */ private static void sink(long id, long data, long length) { - Sink sink = sinks.get(id); - if (sink == null || sink.error != null) { - return; - } + Sink sink = null; try { + sink = sinks.get(id); + if (sink == null || sink.error != null) { + return; + } if (length < 0 || length > Integer.MAX_VALUE - 8 || (data == 0 && length > 0)) { throw new IllegalStateException("The library returned an invalid byte array, of " + length + " bytes at the address " + data + "."); @@ -682,7 +693,9 @@ private static void sink(long id, long data, long length) { Platform.copyMemory(null, data, bytes, Platform.BYTE_ARRAY_OFFSET, length); sink.values.add(bytes); } catch (Throwable e) { - sink.error = e; + if (sink != null) { + sink.error = e; + } } } } diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h index 679a4d99e3349..16ced5dd72232 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -81,16 +81,17 @@ * * Conventions * ----------- - * - Errors. A function that can fail returns 0 on success, or any other value on failure. It then - * sets the message that spark_ds_last_error returns, and Spark reports it as a - * NATIVE_DATA_SOURCE_ERROR. A C++ exception or a Rust panic must not cross the interface: catch - * it and return an error instead. + * - Errors. The functions that return an int return 0 on success, or any other value on failure. + * They then set the message that spark_ds_last_error returns, and Spark reports it as a + * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail. A C++ exception or a Rust panic + * must not cross the interface: catch it and return an error instead. * - Handles. The create functions return an opaque, non-NULL handle in their last argument, - * typically a pointer. Spark releases each handle exactly once with the matching close - * function, or with spark_ds_commit_data_writer or spark_ds_abort_data_writer for data writers, - * though not always right after the last use of a stream writer (see Writes above). A reader or - * writer handle must not depend on the data source handle it was created from, because Spark - * closes the data source handle right after creating it. + * typically a pointer: Spark reports a NULL handle as an error. Spark releases each handle + * exactly once with the matching close function, or with spark_ds_commit_data_writer or + * spark_ds_abort_data_writer for data writers, though not always right after the last use of a + * stream writer (see Writes above). A reader or writer handle must not depend on the data + * source handle it was created from, because Spark closes the data source handle right after + * creating it. * - Memory. What Spark passes to a function, such as strings and byte arrays, is only valid * during the call. Strings are UTF-8 and NUL-terminated. The library returns strings and byte * arrays by passing them to a SparkDsSink, which copies them. @@ -117,7 +118,8 @@ * - Optional functions. A library only needs to export the functions for the operations it * supports. If a function is missing, Spark treats the operation as unsupported, except for the * pushdown functions and the serialize functions, which then push nothing down and return no - * state, and spark_ds_commit_offset, which then does nothing. spark_ds_abi_version is required. + * state, spark_ds_commit_offset, which then does nothing, and spark_ds_last_error. + * spark_ds_abi_version is required. * - Threads. Spark calls the library from multiple threads, and calls on different handles may * run concurrently. Calls on the same handle never run concurrently. A handle may be released * from a thread other than the threads that used it, such as a JVM Cleaner thread when Spark @@ -223,8 +225,8 @@ int32_t spark_ds_abi_version(void); * Creates a data source with the options given by the user, such as * spark.read.option("key", "value"). * - * name: the name of the data source, as listed in the package manifest. A library can implement - * several data sources. + * name: the name of the data source, as listed in the package manifest, or in lower case for an + * installed library. A library can implement several data sources. * option_keys, option_values: the num_options options, with the keys in lower case * source: set to the data source handle */ From f55105d4644b30d53177bf82bdcf4e3841a7987c Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Tue, 13 Oct 2026 16:20:55 +0900 Subject: [PATCH 16/27] Implement the Arrow C data interface with the FFM API, without arrow-c-data arrow-c-data calls the callbacks of the Arrow C structs through a JNI library that it bundles for a few platforms only. Spark now exports and imports the schemas and the arrays, and reads the streams of native data sources, with the FFM API itself, in ArrowCData: arrays are still imported without copying their buffers, through ForeignAllocation, and the release callbacks of the structs that Spark exports are FFM upcalls. The FFM plumbing moves to Ffm, shared with NativeLibrary. This removes the arrow-c-data dependency that this change added, and rejects a stream that the library did not initialize instead of calling its NULL callbacks. Co-authored-by: Isaac --- LICENSE-binary | 1 - dev/deps/spark-deps-hadoop-3-hive-2.3 | 1 - pom.xml | 19 - sql/core/pom.xml | 4 - .../datasources/v2/ffi/ArrowCData.java | 1062 +++++++++++++++++ .../sql/execution/datasources/v2/ffi/Ffm.java | 208 ++++ .../datasources/v2/ffi/NativeLibrary.java | 248 +--- .../spark/sql/datasource/spark_datasource.h | 6 +- .../datasources/v2/ffi/NativeDataSource.scala | 177 ++- .../test_native_datasource.cc | 2 + .../datasources/v2/ffi/ArrowCDataSuite.scala | 463 +++++++ .../v2/ffi/NativeDataSourceSuite.scala | 5 +- 12 files changed, 1880 insertions(+), 316 deletions(-) create mode 100644 sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCData.java create mode 100644 sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/Ffm.java create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCDataSuite.scala diff --git a/LICENSE-binary b/LICENSE-binary index 53dec42594efa..7051e09a6383a 100644 --- a/LICENSE-binary +++ b/LICENSE-binary @@ -301,7 +301,6 @@ joda-time:joda-time net.sf.opencsv:opencsv net.sf.supercsv:super-csv net.sf.jpam:jpam -org.apache.arrow:arrow-c-data org.apache.arrow:arrow-compression org.apache.arrow:arrow-format org.apache.arrow:arrow-memory-core diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index b088d623ea7cb..50f978819c3da 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -21,7 +21,6 @@ antlr4-runtime/4.13.1//antlr4-runtime-4.13.1.jar aopalliance-repackaged/3.0.6//aopalliance-repackaged-3.0.6.jar arpack/3.2.0//arpack-3.2.0.jar arpack_combined_all/0.1//arpack_combined_all-0.1.jar -arrow-c-data/19.0.0//arrow-c-data-19.0.0.jar arrow-compression/19.0.0//arrow-compression-19.0.0.jar arrow-format/19.0.0//arrow-format-19.0.0.jar arrow-memory-core/19.0.0//arrow-memory-core-19.0.0.jar diff --git a/pom.xml b/pom.xml index c2eb3de50b8cd..da05308883570 100644 --- a/pom.xml +++ b/pom.xml @@ -2570,25 +2570,6 @@ - - org.apache.arrow - arrow-c-data - ${arrow.version} - - - com.fasterxml.jackson.core - jackson-annotations - - - com.fasterxml.jackson.core - jackson-core - - - io.netty - netty-common - - - org.apache.arrow arrow-compression diff --git a/sql/core/pom.xml b/sql/core/pom.xml index 7f3db15f70fdf..281546b901a7a 100644 --- a/sql/core/pom.xml +++ b/sql/core/pom.xml @@ -298,10 +298,6 @@ org.apache.arrow arrow-compression - - org.apache.arrow - arrow-c-data - org.apache.spark spark-udf-worker-proto_${scala.binary.version} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCData.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCData.java new file mode 100644 index 0000000000000..803949d08fd1c --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCData.java @@ -0,0 +1,1062 @@ +/* + * 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.execution.datasources.v2.ffi; + +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.MethodType; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.ForeignAllocation; +import org.apache.arrow.vector.BitVectorHelper; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.message.ArrowFieldNode; +import org.apache.arrow.vector.types.DateUnit; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.IntervalUnit; +import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.types.UnionMode; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.DictionaryEncoding; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; + +import org.apache.spark.internal.SparkLogger; +import org.apache.spark.internal.SparkLoggerFactory; +import org.apache.spark.sql.execution.datasources.v2.ffi.Ffm.C; +import org.apache.spark.unsafe.Platform; + +/** + * The Arrow C data + * interface, implemented with the FFM API, see {@link Ffm}: Spark exports and imports the + * schemas and the arrays that it exchanges with native data source libraries, and reads the + * Arrow streams that they export. + *

    + * The C structs are in native memory, at {@code long} addresses, with the layout of 64-bit + * platforms. A struct is released once its {@code release} callback is NULL. Arrays are imported + * without copying their buffers: the imported vectors refer to the memory of the array, which is + * released once the vectors no longer use any of its buffers. Imported arrays must have an offset + * of 0, at every level. + */ +final class ArrowCData { + + private ArrowCData() {} + + private static final SparkLogger LOG = SparkLoggerFactory.getLogger(ArrowCData.class); + + static final long SCHEMA_SIZE = 72; + static final long ARRAY_SIZE = 80; + static final long STREAM_SIZE = 40; + + // The fields of struct ArrowSchema. + private static final long SCHEMA_FORMAT = 0; + private static final long SCHEMA_NAME = 8; + private static final long SCHEMA_METADATA = 16; + private static final long SCHEMA_FLAGS = 24; + private static final long SCHEMA_N_CHILDREN = 32; + private static final long SCHEMA_CHILDREN = 40; + private static final long SCHEMA_DICTIONARY = 48; + private static final long SCHEMA_RELEASE = 56; + private static final long SCHEMA_PRIVATE_DATA = 64; + + // The fields of struct ArrowArray. + private static final long ARRAY_LENGTH = 0; + private static final long ARRAY_NULL_COUNT = 8; + private static final long ARRAY_OFFSET = 16; + private static final long ARRAY_N_BUFFERS = 24; + private static final long ARRAY_N_CHILDREN = 32; + private static final long ARRAY_BUFFERS = 40; + private static final long ARRAY_CHILDREN = 48; + private static final long ARRAY_DICTIONARY = 56; + private static final long ARRAY_RELEASE = 64; + private static final long ARRAY_PRIVATE_DATA = 72; + + // The fields of struct ArrowArrayStream. + private static final long STREAM_GET_SCHEMA = 0; + private static final long STREAM_GET_NEXT = 8; + private static final long STREAM_GET_LAST_ERROR = 16; + private static final long STREAM_RELEASE = 24; + + private static final long FLAG_DICTIONARY_ORDERED = 1; + private static final long FLAG_NULLABLE = 2; + private static final long FLAG_MAP_KEYS_SORTED = 4; + + // Deeper schemas and arrays are rejected, rather than overflowing the stack. + private static final int MAX_DEPTH = 64; + + /** The callbacks of the C structs, created when they are first used. */ + private static final class Callbacks { + // int (*)(void*, void*), such as get_next. + static final MethodHandle INT_FUNCTION; + // const char* (*)(void*), such as get_last_error. + static final MethodHandle POINTER_FUNCTION; + // void (*)(void*), such as release. + static final MethodHandle VOID_FUNCTION; + // The release callbacks of the structs that Spark exports. + static final long RELEASE_SCHEMA; + static final long RELEASE_ARRAY; + + static { + try { + INT_FUNCTION = Ffm.bindPointer(C.INT32, C.POINTER, C.POINTER); + POINTER_FUNCTION = Ffm.bindPointer(C.POINTER, C.POINTER); + VOID_FUNCTION = Ffm.bindPointer(C.VOID, C.POINTER); + MethodHandles.Lookup lookup = MethodHandles.lookup(); + MethodType release = MethodType.methodType(void.class, long.class); + RELEASE_SCHEMA = Ffm.upcall( + lookup.findStatic(ArrowCData.class, "releaseExportedSchema", release), + C.VOID, C.POINTER); + RELEASE_ARRAY = Ffm.upcall( + lookup.findStatic(ArrowCData.class, "releaseExportedArray", release), + C.VOID, C.POINTER); + } catch (Throwable e) { + throw new IllegalStateException("Cannot create the Arrow C data callbacks.", e); + } + } + } + + // --------------------------------------------------------------------------------------------- + // Structs + // --------------------------------------------------------------------------------------------- + + /** Allocates a struct in native memory, initialized to zero, which is a released struct. */ + static long allocate(long size) { + long address = Platform.allocateMemory(size); + Platform.setMemory(address, (byte) 0, size); + return address; + } + + static void free(long address) { + Platform.freeMemory(address); + } + + static boolean isSchemaReleased(long schema) { + return Platform.getLong(null, schema + SCHEMA_RELEASE) == 0; + } + + static boolean isArrayReleased(long array) { + return Platform.getLong(null, array + ARRAY_RELEASE) == 0; + } + + static long length(long array) { + return Platform.getLong(null, array + ARRAY_LENGTH); + } + + static long offset(long array) { + return Platform.getLong(null, array + ARRAY_OFFSET); + } + + static long bufferCount(long array) { + return Platform.getLong(null, array + ARRAY_N_BUFFERS); + } + + /** The address of the buffer at the given index of an array, which may be 0 (NULL). */ + static long buffer(long array, int index) { + return Platform.getLong(null, Platform.getLong(null, array + ARRAY_BUFFERS) + 8L * index); + } + + static long childCount(long array) { + return Platform.getLong(null, array + ARRAY_N_CHILDREN); + } + + static long child(long array, int index) { + return Platform.getLong(null, Platform.getLong(null, array + ARRAY_CHILDREN) + 8L * index); + } + + /** Releases a schema with its release callback, unless it is released already. */ + static void releaseSchema(long schema) throws Throwable { + release(schema, SCHEMA_RELEASE); + } + + /** Releases an array with its release callback, unless it is released already. */ + static void releaseArray(long array) throws Throwable { + release(array, ARRAY_RELEASE); + } + + /** Releases a stream with its release callback, unless it is released already. */ + static void releaseStream(long stream) throws Throwable { + release(stream, STREAM_RELEASE); + } + + private static void release(long struct, long releaseField) throws Throwable { + long release = Platform.getLong(null, struct + releaseField); + if (release != 0) { + Callbacks.VOID_FUNCTION.invokeExact(release, struct); + // The callback marks the struct as released. This makes sure that it is not released again. + Platform.putLong(null, struct + releaseField, 0L); + } + } + + // --------------------------------------------------------------------------------------------- + // Streams + // --------------------------------------------------------------------------------------------- + + /** Calls get_schema of a stream, which initializes the schema at {@code out}. */ + static void streamSchema(long stream, long out) throws Throwable { + long getSchema = callback(stream, STREAM_GET_SCHEMA); + checkStream(stream, (int) Callbacks.INT_FUNCTION.invokeExact(getSchema, stream, out)); + } + + /** + * Calls get_next of a stream, which initializes the array at {@code out}, or marks it as + * released at the end of the stream. + */ + static void streamNext(long stream, long out) throws Throwable { + long getNext = callback(stream, STREAM_GET_NEXT); + checkStream(stream, (int) Callbacks.INT_FUNCTION.invokeExact(getNext, stream, out)); + } + + /** A callback of a stream, which must not be NULL: calling NULL would crash the JVM. */ + private static long callback(long stream, long field) { + if (Platform.getLong(null, stream + STREAM_RELEASE) == 0) { + throw new IllegalStateException("The stream is released, or was not initialized."); + } + long callback = Platform.getLong(null, stream + field); + if (callback == 0) { + throw new IllegalStateException("A callback of the stream is NULL."); + } + return callback; + } + + private static void checkStream(long stream, int errno) throws Throwable { + if (errno == 0) { + return; + } + long getLastError = Platform.getLong(null, stream + STREAM_GET_LAST_ERROR); + long message = getLastError == 0 + ? 0 : (long) Callbacks.POINTER_FUNCTION.invokeExact(getLastError, stream); + String text = message == 0 ? "" : readString(message); + throw new RuntimeException(text.isEmpty() ? "The stream failed with the error " + errno + "." + : text); + } + + // --------------------------------------------------------------------------------------------- + // Exporting schemas + // --------------------------------------------------------------------------------------------- + + /** + * The memory of an exported struct, freed by its release callback. It does not include the + * struct itself, which belongs to its parent or to whoever allocated it. + */ + private static final class Exported { + final List allocations = new ArrayList<>(); + final List buffers = new ArrayList<>(); + + long allocate(long size) { + long address = ArrowCData.allocate(size); + allocations.add(address); + return address; + } + + long string(String value) { + return bytes(value.getBytes(StandardCharsets.UTF_8), true); + } + + long bytes(byte[] bytes, boolean terminate) { + long address = allocate(bytes.length + (terminate ? 1 : 0)); + Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address, bytes.length); + return address; + } + + void free() { + for (ArrowBuf buffer : buffers) { + // Null for a buffer exported as NULL. + if (buffer != null) { + buffer.getReferenceManager().release(); + } + } + allocations.forEach(Platform::freeMemory); + } + } + + private static final AtomicLong nextExportedId = new AtomicLong(1); + private static final ConcurrentHashMap exported = new ConcurrentHashMap<>(); + + /** Exports a schema as a struct type whose fields are the fields of the schema. */ + static void exportSchema(Schema schema, long out) throws Throwable { + exportField( + new Field("", new FieldType(false, ArrowType.Struct.INSTANCE, null, + schema.getCustomMetadata()), schema.getFields()), + out); + } + + private static void exportField(Field field, long out) throws Throwable { + if (field.getDictionary() != null) { + throw new UnsupportedOperationException( + "Dictionary-encoded fields cannot be exported: " + field); + } + String format = format(field.getType()); + List children = field.getChildren(); + long flags = field.isNullable() ? FLAG_NULLABLE : 0; + if (field.getType() instanceof ArrowType.Map map && map.getKeysSorted()) { + flags |= FLAG_MAP_KEYS_SORTED; + } + // A released struct with its callback, so that a failure below releases what was exported. + Platform.setMemory(out, (byte) 0, SCHEMA_SIZE); + Exported data = register(out, SCHEMA_PRIVATE_DATA); + Platform.putLong(null, out + SCHEMA_FLAGS, flags); + Platform.putLong(null, out + SCHEMA_RELEASE, Callbacks.RELEASE_SCHEMA); + try { + Platform.putLong(null, out + SCHEMA_FORMAT, data.string(format)); + String name = field.getName(); + Platform.putLong(null, out + SCHEMA_NAME, name == null ? 0L : data.string(name)); + Platform.putLong(null, out + SCHEMA_METADATA, encodeMetadata(data, field.getMetadata())); + Platform.putLong( + null, out + SCHEMA_CHILDREN, childStructs(data, children.size(), SCHEMA_SIZE)); + for (int i = 0; i < children.size(); i++) { + exportField(children.get(i), schemaChild(out, i)); + // Counted once exported, so that a failure releases the exported children only. + Platform.putLong(null, out + SCHEMA_N_CHILDREN, i + 1L); + } + } catch (Throwable e) { + releaseExportedSchema(out); + throw e; + } + } + + private static long schemaChild(long schema, int index) { + return Platform.getLong(null, Platform.getLong(null, schema + SCHEMA_CHILDREN) + 8L * index); + } + + /** Registers the private data of a struct being exported. */ + private static Exported register(long struct, long privateDataField) { + long id = nextExportedId.getAndIncrement(); + Exported data = new Exported(); + exported.put(id, data); + Platform.putLong(null, struct + privateDataField, id); + return data; + } + + /** + * Allocates the children of an exported struct: an array of pointers to the children, followed + * by the children themselves, initialized to zero. Returns the address of the pointers. + */ + private static long childStructs(Exported data, int count, long structSize) { + if (count == 0) { + return 0; + } + long pointers = data.allocate(count * (8 + structSize)); + for (int i = 0; i < count; i++) { + Platform.putLong(null, pointers + 8L * i, pointers + 8L * count + structSize * i); + } + return pointers; + } + + /** The release callback of the schemas that Spark exports. */ + private static void releaseExportedSchema(long schema) { + try { + long children = Platform.getLong(null, schema + SCHEMA_CHILDREN); + long count = Platform.getLong(null, schema + SCHEMA_N_CHILDREN); + for (int i = 0; i < count; i++) { + // A child that was moved is marked as released here, and released where it was moved. + long child = Platform.getLong(null, children + 8L * i); + if (!isSchemaReleased(child)) { + releaseExportedSchema(child); + } + } + Exported data = exported.remove(Platform.getLong(null, schema + SCHEMA_PRIVATE_DATA)); + if (data != null) { + data.free(); + } + } catch (Throwable e) { + // An exception must not be thrown to the caller of the callback, which is native code. + LOG.warn("Failed to release an exported Arrow schema.", e); + } + Platform.putLong(null, schema + SCHEMA_RELEASE, 0L); + } + + // --------------------------------------------------------------------------------------------- + // Importing schemas + // --------------------------------------------------------------------------------------------- + + /** + * Imports a schema, which must be a struct type whose fields are the fields of the schema, and + * releases it. + */ + static Schema importSchema(long schema) throws Throwable { + try { + Field field = importField(schema, 0, new AtomicLong()); + if (!(field.getType() instanceof ArrowType.Struct)) { + throw new IllegalArgumentException( + "The schema is not a struct type, but " + field.getType() + "."); + } + return new Schema(field.getChildren(), field.getMetadata()); + } finally { + releaseSchema(schema); + } + } + + private static Field importField(long schema, int depth, AtomicLong dictionaryIds) { + if (depth > MAX_DEPTH) { + throw new IllegalArgumentException("The schema is nested more than " + MAX_DEPTH + + " levels deep."); + } + long formatAddress = Platform.getLong(null, schema + SCHEMA_FORMAT); + if (formatAddress == 0) { + throw new IllegalArgumentException("The schema has no format."); + } + String format = readString(formatAddress); + long nameAddress = Platform.getLong(null, schema + SCHEMA_NAME); + String name = nameAddress == 0 ? "" : readString(nameAddress); + Map metadata = + decodeMetadata(Platform.getLong(null, schema + SCHEMA_METADATA)); + long flags = Platform.getLong(null, schema + SCHEMA_FLAGS); + long count = Platform.getLong(null, schema + SCHEMA_N_CHILDREN); + List children = new ArrayList<>(); + for (int i = 0; i < count; i++) { + children.add(importField(schemaChild(schema, i), depth + 1, dictionaryIds)); + } + ArrowType type = parseFormat(format, flags); + DictionaryEncoding dictionary = null; + long dictionarySchema = Platform.getLong(null, schema + SCHEMA_DICTIONARY); + if (dictionarySchema != 0) { + if (!(type instanceof ArrowType.Int indexType)) { + throw new IllegalArgumentException( + "The indices of the dictionary of " + name + " are not integers: " + format + "."); + } + // Only the encoding is kept: Spark does not read dictionary-encoded data. + importField(dictionarySchema, depth + 1, dictionaryIds); + dictionary = new DictionaryEncoding(dictionaryIds.getAndIncrement(), + (flags & FLAG_DICTIONARY_ORDERED) != 0, indexType); + } + boolean nullable = (flags & FLAG_NULLABLE) != 0; + return new Field(name, new FieldType(nullable, type, dictionary, metadata), children); + } + + // --------------------------------------------------------------------------------------------- + // Formats + // --------------------------------------------------------------------------------------------- + + /** The format string of an Arrow type. */ + static String format(ArrowType type) { + switch (type.getTypeID()) { + case Null: return "n"; + case Bool: return "b"; + case Int: { + ArrowType.Int t = (ArrowType.Int) type; + String formats = t.getIsSigned() ? "csil" : "CSIL"; + switch (t.getBitWidth()) { + case 8: return formats.substring(0, 1); + case 16: return formats.substring(1, 2); + case 32: return formats.substring(2, 3); + case 64: return formats.substring(3, 4); + default: break; + } + break; + } + case FloatingPoint: { + switch (((ArrowType.FloatingPoint) type).getPrecision()) { + case HALF: return "e"; + case SINGLE: return "f"; + case DOUBLE: return "g"; + default: break; + } + break; + } + case Binary: return "z"; + case LargeBinary: return "Z"; + case BinaryView: return "vz"; + case Utf8: return "u"; + case LargeUtf8: return "U"; + case Utf8View: return "vu"; + case FixedSizeBinary: + return "w:" + ((ArrowType.FixedSizeBinary) type).getByteWidth(); + case Decimal: { + ArrowType.Decimal t = (ArrowType.Decimal) type; + return "d:" + t.getPrecision() + "," + t.getScale() + + (t.getBitWidth() == 128 ? "" : "," + t.getBitWidth()); + } + case Date: + return ((ArrowType.Date) type).getUnit() == DateUnit.DAY ? "tdD" : "tdm"; + case Time: { + ArrowType.Time t = (ArrowType.Time) type; + return "tt" + unit(t.getUnit()); + } + case Timestamp: { + ArrowType.Timestamp t = (ArrowType.Timestamp) type; + String timezone = t.getTimezone(); + return "ts" + unit(t.getUnit()) + ":" + (timezone == null ? "" : timezone); + } + case Duration: + return "tD" + unit(((ArrowType.Duration) type).getUnit()); + case Interval: { + switch (((ArrowType.Interval) type).getUnit()) { + case YEAR_MONTH: return "tiM"; + case DAY_TIME: return "tiD"; + case MONTH_DAY_NANO: return "tin"; + default: break; + } + break; + } + case List: return "+l"; + case LargeList: return "+L"; + case ListView: return "+vl"; + case LargeListView: return "+vL"; + case FixedSizeList: + return "+w:" + ((ArrowType.FixedSizeList) type).getListSize(); + case Struct: return "+s"; + case Map: return "+m"; + default: break; + } + throw new UnsupportedOperationException("The Arrow type " + type + " cannot be exported."); + } + + private static char unit(TimeUnit unit) { + return switch (unit) { + case SECOND -> 's'; + case MILLISECOND -> 'm'; + case MICROSECOND -> 'u'; + case NANOSECOND -> 'n'; + }; + } + + /** The Arrow type of a format string. */ + static ArrowType parseFormat(String format, long flags) { + switch (format) { + case "n": return ArrowType.Null.INSTANCE; + case "b": return ArrowType.Bool.INSTANCE; + case "c": return new ArrowType.Int(8, true); + case "C": return new ArrowType.Int(8, false); + case "s": return new ArrowType.Int(16, true); + case "S": return new ArrowType.Int(16, false); + case "i": return new ArrowType.Int(32, true); + case "I": return new ArrowType.Int(32, false); + case "l": return new ArrowType.Int(64, true); + case "L": return new ArrowType.Int(64, false); + case "e": return new ArrowType.FloatingPoint(FloatingPointPrecision.HALF); + case "f": return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); + case "g": return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); + case "z": return ArrowType.Binary.INSTANCE; + case "Z": return ArrowType.LargeBinary.INSTANCE; + case "vz": return ArrowType.BinaryView.INSTANCE; + case "u": return ArrowType.Utf8.INSTANCE; + case "U": return ArrowType.LargeUtf8.INSTANCE; + case "vu": return ArrowType.Utf8View.INSTANCE; + case "tdD": return new ArrowType.Date(DateUnit.DAY); + case "tdm": return new ArrowType.Date(DateUnit.MILLISECOND); + case "tts": return new ArrowType.Time(TimeUnit.SECOND, 32); + case "ttm": return new ArrowType.Time(TimeUnit.MILLISECOND, 32); + case "ttu": return new ArrowType.Time(TimeUnit.MICROSECOND, 64); + case "ttn": return new ArrowType.Time(TimeUnit.NANOSECOND, 64); + case "tDs": return new ArrowType.Duration(TimeUnit.SECOND); + case "tDm": return new ArrowType.Duration(TimeUnit.MILLISECOND); + case "tDu": return new ArrowType.Duration(TimeUnit.MICROSECOND); + case "tDn": return new ArrowType.Duration(TimeUnit.NANOSECOND); + case "tiM": return new ArrowType.Interval(IntervalUnit.YEAR_MONTH); + case "tiD": return new ArrowType.Interval(IntervalUnit.DAY_TIME); + case "tin": return new ArrowType.Interval(IntervalUnit.MONTH_DAY_NANO); + case "+l": return ArrowType.List.INSTANCE; + case "+L": return ArrowType.LargeList.INSTANCE; + case "+vl": return ArrowType.ListView.INSTANCE; + case "+vL": return ArrowType.LargeListView.INSTANCE; + case "+s": return ArrowType.Struct.INSTANCE; + case "+m": return new ArrowType.Map((flags & FLAG_MAP_KEYS_SORTED) != 0); + case "+r": return ArrowType.RunEndEncoded.INSTANCE; + default: break; + } + try { + if (format.startsWith("ts") && format.length() >= 4 && format.charAt(3) == ':') { + String timezone = format.substring(4); + return new ArrowType.Timestamp( + timeUnit(format.charAt(2)), timezone.isEmpty() ? null : timezone); + } else if (format.startsWith("d:")) { + String[] parts = format.substring(2).split(",", -1); + if (parts.length == 2 || parts.length == 3) { + return new ArrowType.Decimal(Integer.parseInt(parts[0]), Integer.parseInt(parts[1]), + parts.length == 3 ? Integer.parseInt(parts[2]) : 128); + } + } else if (format.startsWith("w:")) { + return new ArrowType.FixedSizeBinary(Integer.parseInt(format.substring(2))); + } else if (format.startsWith("+w:")) { + return new ArrowType.FixedSizeList(Integer.parseInt(format.substring(3))); + } else if (format.startsWith("+ud:") || format.startsWith("+us:")) { + String ids = format.substring(4); + int[] typeIds = ids.isEmpty() ? new int[0] + : Arrays.stream(ids.split(",", -1)).mapToInt(Integer::parseInt).toArray(); + return new ArrowType.Union( + format.charAt(2) == 'd' ? UnionMode.Dense : UnionMode.Sparse, typeIds); + } + } catch (IllegalArgumentException e) { + // Not a valid format, reported below. NumberFormatException is an IllegalArgumentException. + } + throw new IllegalArgumentException("Unsupported Arrow format: " + format); + } + + private static TimeUnit timeUnit(char unit) { + return switch (unit) { + case 's' -> TimeUnit.SECOND; + case 'm' -> TimeUnit.MILLISECOND; + case 'u' -> TimeUnit.MICROSECOND; + case 'n' -> TimeUnit.NANOSECOND; + default -> throw new IllegalArgumentException("Unsupported time unit: " + unit); + }; + } + + // --------------------------------------------------------------------------------------------- + // Metadata + // --------------------------------------------------------------------------------------------- + + // The metadata are an int32 number of entries, followed by the length and the bytes of the key + // and of the value of each entry, with lengths as int32s in the native byte order. + private static long encodeMetadata(Exported data, Map metadata) { + if (metadata == null || metadata.isEmpty()) { + return 0; + } + List parts = new ArrayList<>(); + long size = 4; + for (Map.Entry entry : metadata.entrySet()) { + byte[] key = entry.getKey().getBytes(StandardCharsets.UTF_8); + byte[] value = entry.getValue().getBytes(StandardCharsets.UTF_8); + parts.add(key); + parts.add(value); + size += 8L + key.length + value.length; + } + long address = data.allocate(size); + Platform.putInt(null, address, metadata.size()); + long position = address + 4; + for (byte[] part : parts) { + Platform.putInt(null, position, part.length); + Platform.copyMemory(part, Platform.BYTE_ARRAY_OFFSET, null, position + 4, part.length); + position += 4L + part.length; + } + return address; + } + + private static Map decodeMetadata(long address) { + if (address == 0) { + return null; + } + int count = Platform.getInt(null, address); + if (count < 0) { + throw new IllegalArgumentException("The metadata of the schema are invalid."); + } + Map metadata = new LinkedHashMap<>(); + long position = address + 4; + for (int i = 0; i < count; i++) { + String[] entry = new String[2]; + for (int j = 0; j < 2; j++) { + int length = Platform.getInt(null, position); + if (length < 0) { + throw new IllegalArgumentException("The metadata of the schema are invalid."); + } + byte[] bytes = new byte[length]; + Platform.copyMemory(null, position + 4, bytes, Platform.BYTE_ARRAY_OFFSET, length); + entry[j] = new String(bytes, StandardCharsets.UTF_8); + position += 4L + length; + } + metadata.put(entry[0], entry[1]); + } + return metadata; + } + + // --------------------------------------------------------------------------------------------- + // Exporting arrays + // --------------------------------------------------------------------------------------------- + + /** + * Exports the vectors of a root as a struct array whose children are the vectors. The array + * refers to the buffers of the vectors, which it retains until it is released. + */ + static void exportRoot(BufferAllocator allocator, VectorSchemaRoot root, long out) + throws Throwable { + Exported data = startArray(out, root.getRowCount(), 0); + try { + // A struct array has a single buffer, the validity buffer, NULL as there are no nulls. + Platform.putLong(null, out + ARRAY_BUFFERS, data.allocate(8)); + Platform.putLong(null, out + ARRAY_N_BUFFERS, 1L); + exportChildren(allocator, data, root.getFieldVectors(), out); + } catch (Throwable e) { + releaseExportedArray(out); + throw e; + } + } + + private static void exportVector(BufferAllocator allocator, FieldVector vector, long out) + throws Throwable { + if (vector.getField().getDictionary() != null) { + throw new UnsupportedOperationException( + "Dictionary-encoded vectors cannot be exported: " + vector.getField()); + } + Exported data = startArray(out, vector.getValueCount(), vector.getNullCount()); + try { + int count = vector.getExportedCDataBufferCount(); + ArrowBuf pointers = allocator.buffer(8L * Math.max(1, count)); + data.buffers.add(pointers); + // Writes the addresses of the buffers, or NULL, and adds them to the buffers to release. + vector.exportCDataBuffers(data.buffers, pointers, 0L); + Platform.putLong(null, out + ARRAY_BUFFERS, pointers.memoryAddress()); + Platform.putLong(null, out + ARRAY_N_BUFFERS, count); + exportChildren(allocator, data, vector.getChildrenFromFields(), out); + } catch (Throwable e) { + releaseExportedArray(out); + throw e; + } + } + + /** + * Initializes an exported array without buffers or children, with its callback, so that a + * failure while exporting them releases what was exported. + */ + private static Exported startArray(long out, long length, long nullCount) { + Platform.setMemory(out, (byte) 0, ARRAY_SIZE); + Exported data = register(out, ARRAY_PRIVATE_DATA); + Platform.putLong(null, out + ARRAY_LENGTH, length); + Platform.putLong(null, out + ARRAY_NULL_COUNT, nullCount); + Platform.putLong(null, out + ARRAY_RELEASE, Callbacks.RELEASE_ARRAY); + return data; + } + + private static void exportChildren( + BufferAllocator allocator, Exported data, List children, long out) + throws Throwable { + Platform.putLong(null, out + ARRAY_CHILDREN, childStructs(data, children.size(), ARRAY_SIZE)); + for (int i = 0; i < children.size(); i++) { + exportVector(allocator, children.get(i), child(out, i)); + // Counted once exported, so that a failure releases the exported children only. + Platform.putLong(null, out + ARRAY_N_CHILDREN, i + 1L); + } + } + + /** The release callback of the arrays that Spark exports. */ + private static void releaseExportedArray(long array) { + try { + long count = childCount(array); + for (int i = 0; i < count; i++) { + // A child that was moved is marked as released here, and released where it was moved. + long child = child(array, i); + if (!isArrayReleased(child)) { + releaseExportedArray(child); + } + } + Exported data = exported.remove(Platform.getLong(null, array + ARRAY_PRIVATE_DATA)); + if (data != null) { + data.free(); + } + } catch (Throwable e) { + // An exception must not be thrown to the caller of the callback, which is native code. + LOG.warn("Failed to release an exported Arrow array.", e); + } + Platform.putLong(null, array + ARRAY_RELEASE, 0L); + } + + // --------------------------------------------------------------------------------------------- + // Importing arrays + // --------------------------------------------------------------------------------------------- + + /** + * An imported array, moved to memory that Spark owns. It is released once the imported buffers + * that refer to its memory are all released, and the importer released its own reference. + */ + private static final class Imported { + private final long array; + private final AtomicInteger references = new AtomicInteger(1); + + Imported(long array) { + this.array = array; + } + + void retain() { + references.incrementAndGet(); + } + + void release() { + if (references.decrementAndGet() == 0) { + try { + releaseArray(array); + } catch (Throwable e) { + LOG.warn("Failed to release an imported Arrow array.", e); + } finally { + free(array); + } + } + } + } + + /** + * Imports a struct array into the vectors of a root, one child for each vector, without copying + * the buffers. The array is moved, and marked as released: the root releases it once its vectors + * no longer use its buffers, such as when the next array is imported or the root is closed. + */ + static void importIntoRoot(BufferAllocator allocator, long array, VectorSchemaRoot root) + throws Throwable { + long moved = allocate(ARRAY_SIZE); + Platform.copyMemory(null, array, null, moved, ARRAY_SIZE); + Platform.putLong(null, array + ARRAY_RELEASE, 0L); + Imported imported = new Imported(moved); + try { + List vectors = root.getFieldVectors(); + long length = length(moved); + if (childCount(moved) != vectors.size()) { + throw new IllegalArgumentException("The array has " + childCount(moved) + + " columns, but the schema has " + vectors.size() + "."); + } + checkLength(length); + for (int i = 0; i < vectors.size(); i++) { + importVector(allocator, imported, child(moved, i), vectors.get(i), 1); + } + root.setRowCount((int) length); + } finally { + imported.release(); + } + } + + private static void importVector( + BufferAllocator allocator, + Imported imported, + long array, + FieldVector vector, + int depth) throws Throwable { + if (depth > MAX_DEPTH) { + throw new IllegalArgumentException("The array is nested more than " + MAX_DEPTH + + " levels deep."); + } + if (offset(array) != 0) { + throw new IllegalArgumentException("The array has an offset of " + offset(array) + + ", but only arrays with an offset of 0 are supported."); + } + ArrowType type = vector.getField().getType(); + long length = length(array); + checkLength(length); + List buffers = new ArrayList<>(); + try { + importBuffers(allocator, imported, array, type, length, buffers); + long nullCount = Platform.getLong(null, array + ARRAY_NULL_COUNT); + if (nullCount < 0) { + // Unknown: computed from the validity buffer, which is NULL when there are no nulls. + ArrowBuf validity = buffers.isEmpty() ? null : buffers.get(0); + nullCount = type instanceof ArrowType.Null ? length + : validity == null || validity.capacity() == 0 ? 0 + : BitVectorHelper.getNullCount(validity, (int) length); + } else if (nullCount > 0 && !buffers.isEmpty() && buffers.get(0) == null && + !(type instanceof ArrowType.Null)) { + throw new IllegalArgumentException("The array has " + nullCount + " nulls, but no " + + "validity buffer."); + } + vector.loadFieldBuffers(new ArrowFieldNode(length, nullCount), buffers); + } finally { + // The vector retains the buffers that it uses. + for (ArrowBuf buffer : buffers) { + if (buffer != null) { + buffer.getReferenceManager().release(); + } + } + } + List children = vector.getChildrenFromFields(); + if (childCount(array) != children.size()) { + throw new IllegalArgumentException("The array of type " + type + " has " + + childCount(array) + " children, but " + children.size() + " are expected."); + } + for (int i = 0; i < children.size(); i++) { + importVector(allocator, imported, child(array, i), children.get(i), depth + 1); + } + } + + private static void checkLength(long length) { + if (length < 0 || length > Integer.MAX_VALUE) { + throw new IllegalArgumentException("The array has an unsupported length of " + length + + "."); + } + } + + /** Wraps the buffers of an array for the vector of the given type, as Arrow Java loads them. */ + private static void importBuffers( + BufferAllocator allocator, + Imported imported, + long array, + ArrowType type, + long length, + List buffers) { + long count = bufferCount(array); + switch (type.getTypeID()) { + case Null: + expectBuffers(type, count, 0); + return; + case Struct: + case FixedSizeList: + expectBuffers(type, count, 1); + buffers.add(validity(allocator, imported, array, length)); + return; + case List: + case Map: + case LargeList: { + expectBuffers(type, count, 2); + int width = type instanceof ArrowType.LargeList ? 8 : 4; + buffers.add(validity(allocator, imported, array, length)); + buffers.add(offsets(allocator, imported, array, length, width)); + return; + } + case Utf8: + case Binary: + case LargeUtf8: + case LargeBinary: { + expectBuffers(type, count, 3); + int width = type instanceof ArrowType.LargeUtf8 || type instanceof ArrowType.LargeBinary + ? 8 : 4; + buffers.add(validity(allocator, imported, array, length)); + ArrowBuf offsets = offsets(allocator, imported, array, length, width); + buffers.add(offsets); + long size = width == 8 ? offsets.getLong(8L * length) : offsets.getInt(4L * length); + buffers.add(wrap(allocator, imported, buffer(array, 2), size, "data")); + return; + } + case Utf8View: + case BinaryView: { + // The validity, the views, the data buffers, and the sizes of the data buffers. + if (count < 3) { + expectBuffers(type, count, 3); + } + buffers.add(validity(allocator, imported, array, length)); + buffers.add(wrap(allocator, imported, buffer(array, 1), 16L * length, "views")); + long sizes = buffer(array, (int) count - 1); + if (sizes == 0 && count > 3) { + throw new IllegalArgumentException("The array of type " + type + " has no sizes of " + + "its data buffers."); + } + for (int i = 2; i < count - 1; i++) { + long size = Platform.getLong(null, sizes + 8L * (i - 2)); + buffers.add(wrap(allocator, imported, buffer(array, i), size, "data")); + } + return; + } + case Bool: + expectBuffers(type, count, 2); + buffers.add(validity(allocator, imported, array, length)); + buffers.add(wrap(allocator, imported, buffer(array, 1), (length + 7) / 8, "values")); + return; + default: { + int width = fixedWidth(type); + expectBuffers(type, count, 2); + buffers.add(validity(allocator, imported, array, length)); + buffers.add(wrap(allocator, imported, buffer(array, 1), width * length, "values")); + } + } + } + + /** The width in bytes of the values of a fixed-width type. */ + private static int fixedWidth(ArrowType type) { + switch (type.getTypeID()) { + case Int: return ((ArrowType.Int) type).getBitWidth() / 8; + case FloatingPoint: + switch (((ArrowType.FloatingPoint) type).getPrecision()) { + case HALF: return 2; + case SINGLE: return 4; + case DOUBLE: return 8; + default: break; + } + break; + case Decimal: return ((ArrowType.Decimal) type).getBitWidth() / 8; + case Date: return ((ArrowType.Date) type).getUnit() == DateUnit.DAY ? 4 : 8; + case Time: return ((ArrowType.Time) type).getBitWidth() / 8; + case Timestamp: + case Duration: + return 8; + case Interval: + switch (((ArrowType.Interval) type).getUnit()) { + case YEAR_MONTH: return 4; + case DAY_TIME: return 8; + case MONTH_DAY_NANO: return 16; + default: break; + } + break; + case FixedSizeBinary: return ((ArrowType.FixedSizeBinary) type).getByteWidth(); + default: break; + } + throw new UnsupportedOperationException("Arrays of the Arrow type " + type + " cannot be " + + "imported."); + } + + private static void expectBuffers(ArrowType type, long count, int expected) { + if (count != expected) { + throw new IllegalArgumentException("The array of type " + type + " has " + count + + " buffers, but " + expected + " are expected."); + } + } + + /** The validity buffer, or null if it is NULL, which Arrow Java loads as all valid or null. */ + private static ArrowBuf validity( + BufferAllocator allocator, Imported imported, long array, long length) { + long address = buffer(array, 0); + return address == 0 ? null : wrap(allocator, imported, address, (length + 7) / 8, "validity"); + } + + /** The offsets buffer, which may be NULL in an empty array. */ + private static ArrowBuf offsets( + BufferAllocator allocator, Imported imported, long array, long length, int width) { + long address = buffer(array, 1); + long size = width * (length + 1); + if (address == 0 && length == 0) { + ArrowBuf offsets = allocator.buffer(size); + offsets.setZero(0, size); + return offsets; + } + return wrap(allocator, imported, address, size, "offsets"); + } + + /** Wraps a buffer of an imported array, which keeps the array alive until it is released. */ + private static ArrowBuf wrap( + BufferAllocator allocator, Imported imported, long address, long size, String buffer) { + if (size < 0) { + throw new IllegalArgumentException("The " + buffer + " buffer has a negative size."); + } + if (size == 0) { + return allocator.getEmpty(); + } + if (address == 0) { + throw new IllegalArgumentException("The " + buffer + " buffer is NULL."); + } + imported.retain(); + try { + return allocator.wrapForeignAllocation(new ForeignAllocation(size, address) { + @Override + protected void release0() { + imported.release(); + } + }); + } catch (Throwable e) { + imported.release(); + throw e; + } + } + + // --------------------------------------------------------------------------------------------- + // Strings + // --------------------------------------------------------------------------------------------- + + /** Reads a NUL-terminated UTF-8 string. */ + static String readString(long address) { + int length = 0; + while (Platform.getByte(null, address + length) != 0) { + length++; + } + byte[] bytes = new byte[length]; + Platform.copyMemory(null, address, bytes, Platform.BYTE_ARRAY_OFFSET, length); + return new String(bytes, StandardCharsets.UTF_8); + } +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/Ffm.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/Ffm.java new file mode 100644 index 0000000000000..75ce65b06d119 --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/Ffm.java @@ -0,0 +1,208 @@ +/* + * 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.execution.datasources.v2.ffi; + +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.MethodType; +import java.lang.reflect.Array; +import java.nio.file.Path; +import java.util.Optional; + +/** + * The Foreign Function and Memory (FFM) API of Java 22 and later, for native data sources. + *

    + * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached + * through method handles, looked up when this class is initialized, the first time that a native + * data source library is loaded. The C functions are bound to method handles that take and return + * primitive values only, with pointers as {@code long} addresses, so calling them uses no + * reflection. + */ +final class Ffm { + + private Ffm() {} + + /** The C types of the parameters and the results of functions. */ + enum C { VOID, BOOL, INT32, INT64, POINTER } + + private static final Class LAYOUT_CLASS; + private static final Object LINKER; + private static final Object GLOBAL_ARENA; + private static final Object NO_OPTIONS; + private static final MethodHandle LIBRARY_LOOKUP; + private static final MethodHandle FIND; + private static final MethodHandle DOWNCALL; + private static final MethodHandle POINTER_DOWNCALL; + private static final MethodHandle UPCALL; + private static final MethodHandle FUNCTION_OF; + private static final MethodHandle FUNCTION_OF_VOID; + // (long) -> MemorySegment and (MemorySegment) -> long. + private static final MethodHandle OF_ADDRESS; + private static final MethodHandle ADDRESS_OF; + private static final Object BOOL; + private static final Object INT32; + private static final Object INT64; + private static final Object POINTER; + + static { + try { + // A full-privilege lookup, so that the restricted methods of the FFM API see this class as + // their caller, in the unnamed module that Spark enables native access for. + MethodHandles.Lookup lookup = MethodHandles.lookup(); + Class linker = Class.forName("java.lang.foreign.Linker"); + Class option = Class.forName("java.lang.foreign.Linker$Option"); + Class symbolLookup = Class.forName("java.lang.foreign.SymbolLookup"); + Class arena = Class.forName("java.lang.foreign.Arena"); + Class segment = Class.forName("java.lang.foreign.MemorySegment"); + Class layout = Class.forName("java.lang.foreign.MemoryLayout"); + Class layoutArray = Array.newInstance(layout, 0).getClass(); + Class valueLayout = Class.forName("java.lang.foreign.ValueLayout"); + Class descriptor = Class.forName("java.lang.foreign.FunctionDescriptor"); + Class options = Array.newInstance(option, 0).getClass(); + + LAYOUT_CLASS = layout; + LINKER = + lookup.findStatic(linker, "nativeLinker", MethodType.methodType(linker)).invoke(); + GLOBAL_ARENA = lookup.findStatic(arena, "global", MethodType.methodType(arena)).invoke(); + NO_OPTIONS = Array.newInstance(option, 0); + LIBRARY_LOOKUP = lookup.findStatic(symbolLookup, "libraryLookup", + MethodType.methodType(symbolLookup, Path.class, arena)); + FIND = lookup.findVirtual(symbolLookup, "find", + MethodType.methodType(Optional.class, String.class)); + DOWNCALL = lookup.findVirtual(linker, "downcallHandle", + MethodType.methodType(MethodHandle.class, segment, descriptor, options)).asFixedArity(); + POINTER_DOWNCALL = lookup.findVirtual(linker, "downcallHandle", + MethodType.methodType(MethodHandle.class, descriptor, options)).asFixedArity(); + UPCALL = lookup.findVirtual(linker, "upcallStub", + MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options)) + .asFixedArity(); + FUNCTION_OF = lookup.findStatic(descriptor, "of", + MethodType.methodType(descriptor, layout, layoutArray)).asFixedArity(); + FUNCTION_OF_VOID = lookup.findStatic(descriptor, "ofVoid", + MethodType.methodType(descriptor, layoutArray)).asFixedArity(); + OF_ADDRESS = + lookup.findStatic(segment, "ofAddress", MethodType.methodType(segment, long.class)); + ADDRESS_OF = lookup.findVirtual(segment, "address", MethodType.methodType(long.class)); + BOOL = valueLayout.getField("JAVA_BOOLEAN").get(null); + INT32 = valueLayout.getField("JAVA_INT").get(null); + INT64 = valueLayout.getField("JAVA_LONG").get(null); + POINTER = valueLayout.getField("ADDRESS").get(null); + } catch (Throwable e) { + // The message names the cause, because the JVM only keeps this message for the later uses + // of the class. + throw new IllegalStateException("Cannot initialize the FFM API: " + rootCause(e), e); + } + } + + static Throwable rootCause(Throwable e) { + Throwable cause = e; + while (cause.getCause() != null && cause.getCause() != cause) { + cause = cause.getCause(); + } + return cause; + } + + /** + * Returns the symbols of the library at the given path, which stays loaded until the JVM exits. + * + * @throws IllegalArgumentException if the file is not a library + * @throws IllegalCallerException if native access is not enabled for Spark + */ + static Object libraryLookup(Path path) throws Throwable { + return LIBRARY_LOOKUP.invoke(path, GLOBAL_ARENA); + } + + /** + * Binds a function of a library, looked up with {@link #libraryLookup}, to a method handle with + * the Java types of the given C types. Returns null if the library does not export the function. + */ + static MethodHandle bind(Object library, String name, C result, C... parameters) + throws Throwable { + Optional symbol = (Optional) FIND.invoke(library, name); + if (symbol.isEmpty()) { + return null; + } + MethodHandle handle = (MethodHandle) + DOWNCALL.invoke(LINKER, symbol.get(), descriptor(result, parameters), NO_OPTIONS); + return withAddresses(handle, 0, result, parameters); + } + + /** + * Returns a method handle that calls the C function whose address is its first argument, with + * the Java types of the given C types, such as a callback of an Arrow C struct. + */ + static MethodHandle bindPointer(C result, C... parameters) throws Throwable { + MethodHandle handle = (MethodHandle) + POINTER_DOWNCALL.invoke(LINKER, descriptor(result, parameters), NO_OPTIONS); + handle = MethodHandles.filterArguments(handle, 0, OF_ADDRESS); + return withAddresses(handle, 1, result, parameters); + } + + /** + * Returns the address of a C function with the given C types that calls the given method handle, + * whose parameters and result have the Java types of the C types. The function lives as long as + * the JVM. The method handle must not throw: an exception thrown by an upcall crashes the JVM. + */ + static long upcall(MethodHandle target, C result, C... parameters) throws Throwable { + for (int i = 0; i < parameters.length; i++) { + if (parameters[i] == C.POINTER) { + target = MethodHandles.filterArguments(target, i, ADDRESS_OF); + } + } + if (result == C.POINTER) { + target = MethodHandles.filterReturnValue(target, OF_ADDRESS); + } + Object stub = + UPCALL.invoke(LINKER, target, descriptor(result, parameters), GLOBAL_ARENA, NO_OPTIONS); + return (long) ADDRESS_OF.invoke(stub); + } + + /** Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses. */ + private static MethodHandle withAddresses( + MethodHandle handle, int firstParameter, C result, C... parameters) { + for (int i = 0; i < parameters.length; i++) { + if (parameters[i] == C.POINTER) { + handle = MethodHandles.filterArguments(handle, firstParameter + i, OF_ADDRESS); + } + } + if (result == C.POINTER) { + handle = MethodHandles.filterReturnValue(handle, ADDRESS_OF); + } + return handle; + } + + private static Object descriptor(C result, C... parameters) throws Throwable { + Object layouts = Array.newInstance(LAYOUT_CLASS, parameters.length); + for (int i = 0; i < parameters.length; i++) { + Array.set(layouts, i, layout(parameters[i])); + } + return result == C.VOID + ? FUNCTION_OF_VOID.invoke(layouts) + : FUNCTION_OF.invoke(layout(result), layouts); + } + + private static Object layout(C type) { + return switch (type) { + case BOOL -> BOOL; + case INT32 -> INT32; + case INT64 -> INT64; + case POINTER -> POINTER; + case VOID -> throw new IllegalArgumentException("void has no layout"); + }; + } +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java index e7c9129459fa1..33492ff80ed34 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -20,27 +20,22 @@ import java.lang.invoke.MethodHandle; import java.lang.invoke.MethodHandles; import java.lang.invoke.MethodType; -import java.lang.reflect.Array; import java.nio.charset.StandardCharsets; import java.nio.file.Path; import java.util.ArrayList; import java.util.List; -import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicLong; +import org.apache.spark.sql.execution.datasources.v2.ffi.Ffm.C; import org.apache.spark.unsafe.Platform; /** * A native data source library, which implements the C interface in * {@code org/apache/spark/sql/datasource/spark_datasource.h}, called through the Foreign Function - * and Memory (FFM) API of Java 22 and later. - *

    - * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached - * through method handles, looked up once when the first library is loaded. Each function of a - * library is bound to a method handle that takes and returns primitive values only, with - * pointers as {@code long} addresses, so the calls themselves use no reflection. The memory that - * Spark passes to the library is allocated with {@link Platform}. + * and Memory (FFM) API of Java 22 and later, see {@link Ffm}. Each function of the library is + * bound once, when the library is loaded. The memory that Spark passes to the library is allocated + * with {@link Platform}. *

    * Each library is looked up separately, so all the libraries can export the same functions. * Calling a function that the library does not export throws {@link MissingFunctionException}. A @@ -60,9 +55,6 @@ static boolean isSupported() { return Runtime.version().feature() >= MINIMUM_JAVA_VERSION; } - // The C types of the parameters and the results of the functions. - private enum C { VOID, BOOL, INT32, INT64, POINTER } - private final MethodHandle lastError; private final MethodHandle abiVersion; private final MethodHandle create; @@ -94,44 +86,44 @@ private enum C { VOID, BOOL, INT32, INT64, POINTER } private final MethodHandle commitDataWriter; private final MethodHandle abortDataWriter; - private NativeLibrary(Object lookup) throws Throwable { + private NativeLibrary(Object library) throws Throwable { C i32 = C.INT32; C i64 = C.INT64; C ptr = C.POINTER; - lastError = bind(lookup, "spark_ds_last_error", ptr); - abiVersion = bind(lookup, "spark_ds_abi_version", i32); - create = bind(lookup, "spark_ds_create", i32, ptr, ptr, ptr, i64, ptr); - schema = bind(lookup, "spark_ds_schema", i32, ptr, ptr); - createReader = bind(lookup, "spark_ds_create_reader", i32, ptr, ptr, ptr); - createStreamReader = bind(lookup, "spark_ds_create_stream_reader", i32, ptr, ptr, ptr); - createWriter = bind(lookup, "spark_ds_create_writer", i32, ptr, ptr, C.BOOL, ptr); + lastError = Ffm.bind(library, "spark_ds_last_error", ptr); + abiVersion = Ffm.bind(library, "spark_ds_abi_version", i32); + create = Ffm.bind(library, "spark_ds_create", i32, ptr, ptr, ptr, i64, ptr); + schema = Ffm.bind(library, "spark_ds_schema", i32, ptr, ptr); + createReader = Ffm.bind(library, "spark_ds_create_reader", i32, ptr, ptr, ptr); + createStreamReader = Ffm.bind(library, "spark_ds_create_stream_reader", i32, ptr, ptr, ptr); + createWriter = Ffm.bind(library, "spark_ds_create_writer", i32, ptr, ptr, C.BOOL, ptr); createStreamWriter = - bind(lookup, "spark_ds_create_stream_writer", i32, ptr, ptr, C.BOOL, ptr); - close = bind(lookup, "spark_ds_close", C.VOID, ptr); - pushPredicates = bind(lookup, "spark_ds_push_predicates", i32, ptr, ptr, i64, ptr); - pushLimit = bind(lookup, "spark_ds_push_limit", i32, ptr, i32, ptr); - pruneColumns = bind(lookup, "spark_ds_prune_columns", i32, ptr, ptr, i64, ptr); - partitions = bind(lookup, "spark_ds_partitions", i32, ptr, ptr, ptr); - serializeReader = bind(lookup, "spark_ds_serialize_reader", i32, ptr, ptr, ptr); - closeReader = bind(lookup, "spark_ds_close_reader", C.VOID, ptr); - initialOffset = bind(lookup, "spark_ds_initial_offset", i32, ptr, ptr, ptr); - latestOffset = bind(lookup, "spark_ds_latest_offset", i32, ptr, ptr, ptr); + Ffm.bind(library, "spark_ds_create_stream_writer", i32, ptr, ptr, C.BOOL, ptr); + close = Ffm.bind(library, "spark_ds_close", C.VOID, ptr); + pushPredicates = Ffm.bind(library, "spark_ds_push_predicates", i32, ptr, ptr, i64, ptr); + pushLimit = Ffm.bind(library, "spark_ds_push_limit", i32, ptr, i32, ptr); + pruneColumns = Ffm.bind(library, "spark_ds_prune_columns", i32, ptr, ptr, i64, ptr); + partitions = Ffm.bind(library, "spark_ds_partitions", i32, ptr, ptr, ptr); + serializeReader = Ffm.bind(library, "spark_ds_serialize_reader", i32, ptr, ptr, ptr); + closeReader = Ffm.bind(library, "spark_ds_close_reader", C.VOID, ptr); + initialOffset = Ffm.bind(library, "spark_ds_initial_offset", i32, ptr, ptr, ptr); + latestOffset = Ffm.bind(library, "spark_ds_latest_offset", i32, ptr, ptr, ptr); streamPartitions = - bind(lookup, "spark_ds_stream_partitions", i32, ptr, ptr, ptr, ptr, ptr); + Ffm.bind(library, "spark_ds_stream_partitions", i32, ptr, ptr, ptr, ptr, ptr); serializeStreamReader = - bind(lookup, "spark_ds_serialize_stream_reader", i32, ptr, ptr, ptr); - commitOffset = bind(lookup, "spark_ds_commit_offset", i32, ptr, ptr); - closeStreamReader = bind(lookup, "spark_ds_close_stream_reader", C.VOID, ptr); - read = bind(lookup, "spark_ds_read", i32, ptr, i64, ptr, i64, ptr); - serializeWriter = bind(lookup, "spark_ds_serialize_writer", i32, ptr, ptr, ptr); - commit = bind(lookup, "spark_ds_commit", i32, ptr, i64, ptr, ptr, i64); - abort = bind(lookup, "spark_ds_abort", i32, ptr, i64, ptr, ptr, i64); - closeWriter = bind(lookup, "spark_ds_close_writer", C.VOID, ptr); + Ffm.bind(library, "spark_ds_serialize_stream_reader", i32, ptr, ptr, ptr); + commitOffset = Ffm.bind(library, "spark_ds_commit_offset", i32, ptr, ptr); + closeStreamReader = Ffm.bind(library, "spark_ds_close_stream_reader", C.VOID, ptr); + read = Ffm.bind(library, "spark_ds_read", i32, ptr, i64, ptr, i64, ptr); + serializeWriter = Ffm.bind(library, "spark_ds_serialize_writer", i32, ptr, ptr, ptr); + commit = Ffm.bind(library, "spark_ds_commit", i32, ptr, i64, ptr, ptr, i64); + abort = Ffm.bind(library, "spark_ds_abort", i32, ptr, i64, ptr, ptr, i64); + closeWriter = Ffm.bind(library, "spark_ds_close_writer", C.VOID, ptr); createDataWriter = - bind(lookup, "spark_ds_create_data_writer", i32, ptr, i64, i32, i64, i64, ptr); - write = bind(lookup, "spark_ds_write", i32, ptr, ptr, ptr); - commitDataWriter = bind(lookup, "spark_ds_commit_data_writer", i32, ptr, ptr, ptr); - abortDataWriter = bind(lookup, "spark_ds_abort_data_writer", i32, ptr); + Ffm.bind(library, "spark_ds_create_data_writer", i32, ptr, i64, i32, i64, i64, ptr); + write = Ffm.bind(library, "spark_ds_write", i32, ptr, ptr, ptr); + commitDataWriter = Ffm.bind(library, "spark_ds_commit_data_writer", i32, ptr, ptr, ptr); + abortDataWriter = Ffm.bind(library, "spark_ds_abort_data_writer", i32, ptr); } /** @@ -144,9 +136,9 @@ static NativeLibrary load(String path) throws Throwable { throw new UnsupportedOperationException( "Native data sources require Java " + MINIMUM_JAVA_VERSION + " or later."); } - Object lookup; + Object library; try { - lookup = Ffm.LIBRARY_LOOKUP.invoke(Path.of(path), Ffm.GLOBAL_ARENA); + library = Ffm.libraryLookup(Path.of(path)); } catch (IllegalArgumentException | IllegalCallerException e) { // The library is not a valid library, or native access is not enabled for Spark. throw new LoadException(e.getMessage(), e); @@ -154,145 +146,9 @@ static NativeLibrary load(String path) throws Throwable { // The FFM API could not be initialized, see Ffm: report why, such as an // IllegalCallerException when native access is not enabled for Spark. throw new LoadException( - "Cannot use the Foreign Function and Memory API: " + rootCause(e), e); + "Cannot use the Foreign Function and Memory API: " + Ffm.rootCause(e), e); } - return new NativeLibrary(lookup); - } - - /** - * Binds a function of the library to a method handle whose parameters and result have the - * Java types of the given C types, with {@code long} for pointers. Returns null if the library - * does not export the function. - */ - private static MethodHandle bind(Object lookup, String name, C result, C... parameters) - throws Throwable { - Optional symbol = (Optional) Ffm.FIND.invoke(lookup, name); - if (symbol.isEmpty()) { - return null; - } - Object layouts = Ffm.layouts(parameters); - Object descriptor = result == C.VOID - ? Ffm.FUNCTION_OF_VOID.invoke(layouts) - : Ffm.FUNCTION_OF.invoke(Ffm.layout(result), layouts); - MethodHandle handle = (MethodHandle) Ffm.DOWNCALL.invoke( - Ffm.LINKER, symbol.get(), descriptor, Ffm.NO_OPTIONS); - // Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses. - for (int i = 0; i < parameters.length; i++) { - if (parameters[i] == C.POINTER) { - handle = MethodHandles.filterArguments(handle, i, Ffm.OF_ADDRESS); - } - } - if (result == C.POINTER) { - handle = MethodHandles.filterReturnValue(handle, Ffm.ADDRESS_OF); - } - return handle; - } - - /** The method handles of the FFM API, initialized when the first library is loaded. */ - private static final class Ffm { - static final Class LAYOUT_CLASS; - static final Object LINKER; - static final Object GLOBAL_ARENA; - static final Object NO_OPTIONS; - static final MethodHandle LIBRARY_LOOKUP; - static final MethodHandle FIND; - static final MethodHandle DOWNCALL; - static final MethodHandle FUNCTION_OF; - static final MethodHandle FUNCTION_OF_VOID; - // (long) -> MemorySegment and (MemorySegment) -> long. - static final MethodHandle OF_ADDRESS; - static final MethodHandle ADDRESS_OF; - static final Object BOOL; - static final Object INT32; - static final Object INT64; - static final Object POINTER; - // The address of the C function SparkDsSink, which calls NativeLibrary.sink. - static final long SINK; - - static { - try { - // A full-privilege lookup, so that the restricted methods of the FFM API see this class - // as their caller, in the unnamed module that Spark enables native access for. - MethodHandles.Lookup lookup = MethodHandles.lookup(); - Class linker = Class.forName("java.lang.foreign.Linker"); - Class option = Class.forName("java.lang.foreign.Linker$Option"); - Class symbolLookup = Class.forName("java.lang.foreign.SymbolLookup"); - Class arena = Class.forName("java.lang.foreign.Arena"); - Class segment = Class.forName("java.lang.foreign.MemorySegment"); - Class layout = Class.forName("java.lang.foreign.MemoryLayout"); - Class layoutArray = Array.newInstance(layout, 0).getClass(); - Class valueLayout = Class.forName("java.lang.foreign.ValueLayout"); - Class descriptor = Class.forName("java.lang.foreign.FunctionDescriptor"); - Class options = Array.newInstance(option, 0).getClass(); - - LAYOUT_CLASS = layout; - LINKER = lookup.findStatic(linker, "nativeLinker", MethodType.methodType(linker)) - .invoke(); - GLOBAL_ARENA = lookup.findStatic(arena, "global", MethodType.methodType(arena)).invoke(); - NO_OPTIONS = Array.newInstance(option, 0); - LIBRARY_LOOKUP = lookup.findStatic(symbolLookup, "libraryLookup", - MethodType.methodType(symbolLookup, Path.class, arena)); - FIND = lookup.findVirtual(symbolLookup, "find", - MethodType.methodType(Optional.class, String.class)); - DOWNCALL = lookup.findVirtual(linker, "downcallHandle", - MethodType.methodType(MethodHandle.class, segment, descriptor, options)) - .asFixedArity(); - FUNCTION_OF = lookup.findStatic(descriptor, "of", - MethodType.methodType(descriptor, layout, layoutArray)).asFixedArity(); - FUNCTION_OF_VOID = lookup.findStatic(descriptor, "ofVoid", - MethodType.methodType(descriptor, layoutArray)).asFixedArity(); - OF_ADDRESS = lookup.findStatic(segment, "ofAddress", - MethodType.methodType(segment, long.class)); - ADDRESS_OF = lookup.findVirtual(segment, "address", MethodType.methodType(long.class)); - BOOL = valueLayout.getField("JAVA_BOOLEAN").get(null); - INT32 = valueLayout.getField("JAVA_INT").get(null); - INT64 = valueLayout.getField("JAVA_LONG").get(null); - POINTER = valueLayout.getField("ADDRESS").get(null); - - // The upcall stub of the sink, which lives as long as the JVM. - MethodHandle target = lookup.findStatic(NativeLibrary.class, "sink", - MethodType.methodType(void.class, long.class, long.class, long.class)); - target = MethodHandles.filterArguments(target, 0, ADDRESS_OF, ADDRESS_OF); - Object sinkDescriptor = - FUNCTION_OF_VOID.invoke(layouts(C.POINTER, C.POINTER, C.INT64)); - Object stub = lookup.findVirtual(linker, "upcallStub", - MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options)) - .asFixedArity() - .invoke(LINKER, target, sinkDescriptor, GLOBAL_ARENA, NO_OPTIONS); - SINK = (long) ADDRESS_OF.invoke(stub); - } catch (Throwable e) { - // The message names the cause, because the JVM only keeps this message for the later - // uses of the class. - throw new IllegalStateException("Cannot initialize the FFM API: " + rootCause(e), e); - } - } - - /** Returns a MemoryLayout[] with the layouts of the given types. */ - static Object layouts(C... types) { - Object layouts = Array.newInstance(LAYOUT_CLASS, types.length); - for (int i = 0; i < types.length; i++) { - Array.set(layouts, i, layout(types[i])); - } - return layouts; - } - - static Object layout(C type) { - return switch (type) { - case BOOL -> BOOL; - case INT32 -> INT32; - case INT64 -> INT64; - case POINTER -> POINTER; - case VOID -> throw new IllegalArgumentException("void has no layout"); - }; - } - } - - private static Throwable rootCause(Throwable e) { - Throwable cause = e; - while (cause.getCause() != null && cause.getCause() != cause) { - cause = cause.getCause(); - } - return cause; + return new NativeLibrary(library); } /** Thrown when a library cannot be loaded. */ @@ -444,7 +300,7 @@ boolean pruneColumns(long reader, String[] columnNames) throws Throwable { byte[][] partitions(long reader) throws Throwable { MethodHandle f = require(partitions, "spark_ds_partitions"); try (Sink sink = new Sink()) { - check((int) f.invokeExact(reader, Ffm.SINK, sink.id), "spark_ds_partitions"); + check((int) f.invokeExact(reader, Sinks.SINK, sink.id), "spark_ds_partitions"); return sink.all(); } } @@ -478,7 +334,7 @@ byte[][] streamPartitions(long streamReader, String start, String end) throws Th MethodHandle f = require(streamPartitions, "spark_ds_stream_partitions"); try (Memory memory = new Memory(); Sink sink = new Sink()) { check((int) f.invokeExact(streamReader, memory.string(start), memory.string(end), - Ffm.SINK, sink.id), "spark_ds_stream_partitions"); + Sinks.SINK, sink.id), "spark_ds_stream_partitions"); return sink.all(); } } @@ -576,7 +432,7 @@ void abortDataWriter(long dataWriter) throws Throwable { private byte[] serialize(MethodHandle function, String name, long handle) throws Throwable { MethodHandle f = require(function, name); try (Sink sink = new Sink()) { - check((int) f.invokeExact(handle, Ffm.SINK, sink.id), name); + check((int) f.invokeExact(handle, Sinks.SINK, sink.id), name); byte[][] values = sink.all(); if (values.length > 1) { throw new IllegalStateException( @@ -645,12 +501,28 @@ private static byte[] readCString(long address) { // Sink // --------------------------------------------------------------------------------------------- + /** The C function SparkDsSink, which calls {@link #sink}. */ + private static final class Sinks { + static final long SINK; + + static { + try { + SINK = Ffm.upcall( + MethodHandles.lookup().findStatic(NativeLibrary.class, "sink", + MethodType.methodType(void.class, long.class, long.class, long.class)), + C.VOID, C.POINTER, C.POINTER, C.INT64); + } catch (Throwable e) { + throw new IllegalStateException("Cannot create the sink of native data sources.", e); + } + } + } + private static final AtomicLong nextSinkId = new AtomicLong(1); private static final ConcurrentHashMap sinks = new ConcurrentHashMap<>(); /** * Collects the byte arrays that a library returns during a call, through the C function - * {@link Ffm#SINK}. The library passes {@link #id} back to identify it. + * {@link Sinks#SINK}. The library passes {@link #id} back to identify it. */ private static final class Sink implements AutoCloseable { final long id = nextSinkId.getAndIncrement(); @@ -675,7 +547,7 @@ public void close() { } /** - * Called by the library, through the upcall stub {@link Ffm#SINK}, with a byte array to copy. + * Called by the library, through the C function {@link Sinks#SINK}, with a byte array to copy. * It must not throw: an exception thrown by an upcall crashes the JVM. */ private static void sink(long id, long data, long length) { diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h index 16ced5dd72232..4cf8b464bf4fd 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -102,9 +102,9 @@ * For an output struct, the library initializes it and Spark takes ownership. For an input * struct, the library may move it to take ownership; otherwise, Spark releases it after the * call returns, so the library must not use it afterwards. The arrays that the library exports - * must have an offset of 0, at every level: Arrow Java ignores the offsets of the arrays it - * imports, so Spark rejects an array with another offset, such as a zero-copy slice. Copy a - * sliced array before exporting it. Spark also rejects an array with fewer values than its + * must have an offset of 0, at every level: Spark reads them with Arrow Java vectors, which + * have no offsets, so it rejects an array with another offset, such as a zero-copy slice. Copy + * a sliced array before exporting it. Spark also rejects an array with fewer values than its * parent refers to, such as a column shorter than its batch. The buffers of an array must hold * all of its values, as the C data interface requires: it does not carry the sizes of the * buffers, so Spark cannot check them, and reads past the end of a buffer that is too short. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index 943da94be272e..7d54926216439 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -23,7 +23,6 @@ import java.util.concurrent.atomic.AtomicLong import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal -import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, CDataDictionaryProvider, Data} import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} import org.apache.arrow.vector.complex.MapVector import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} @@ -376,16 +375,16 @@ class NativePartitionReader( private val allocator = ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) - // Read here rather than with `Data.importArrayStream`, to check each array before importing it. - private val stream = ArrowArrayStream.allocateNew(allocator) - private val dictionaries = new CDataDictionaryProvider() + // The ArrowArrayStream that the library exports. + private val stream = ArrowCData.allocate(ArrowCData.STREAM_SIZE) private var root: VectorSchemaRoot = _ private var batch: ColumnarBatch = _ try { openStream() - root = VectorSchemaRoot.create(importSchema(), allocator) - checkSchema() + val arrowSchema = importSchema() + checkSchema(arrowSchema) + root = VectorSchemaRoot.create(arrowSchema, allocator) batch = new ColumnarBatch(root.getFieldVectors.asScala.map(columnVector).toArray) } catch { case e: Throwable => @@ -395,22 +394,21 @@ class NativePartitionReader( private def openStream(): Unit = { val library = NativeLibraries.get(location) - call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) - .getOrElse { - throw new SparkUnsupportedOperationException( - "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) - } + call(name, "read")(optional(library.read(readerState, partition, stream))).getOrElse { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) + } } private def importSchema(): Schema = { - val arrowSchema = ArrowSchema.allocateNew(allocator) + val arrowSchema = ArrowCData.allocate(ArrowCData.SCHEMA_SIZE) try { - call(name, "read")(stream.getSchema(arrowSchema)) - val imported = call(name, "read")(Data.importSchema(allocator, arrowSchema, dictionaries)) + call(name, "read")(ArrowCData.streamSchema(stream, arrowSchema)) + val imported = call(name, "read")(ArrowCData.importSchema(arrowSchema)) new Schema( imported.getFields.asScala.map(withMapEntryNames).asJava, imported.getCustomMetadata) } finally { - arrowSchema.close() + ArrowCData.free(arrowSchema) } } @@ -432,11 +430,12 @@ class NativePartitionReader( new Field(field.getName, field.getFieldType, renamed.asJava) } - private def checkSchema(): Unit = { + /** Checks the schema of the data, before Spark creates the vectors to import it into. */ + private def checkSchema(arrowSchema: Schema): Unit = { def isDictionaryEncoded(field: Field): Boolean = { field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) } - root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + arrowSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + "is not supported.") } @@ -448,13 +447,13 @@ class NativePartitionReader( case listView: ArrowType.ListView => Some(listView) case _ => field.getChildren.asScala.view.flatMap(unsupportedType).headOption } - root.getSchema.getFields.asScala.foreach { field => + arrowSchema.getFields.asScala.foreach { field => unsupportedType(field).foreach { arrowType => throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + s"$arrowType, which is not supported.") } } - val actual = ArrowUtils.fromArrowSchema(root.getSchema) + val actual = ArrowUtils.fromArrowSchema(arrowSchema) // The schema to read as the library got it: Arrow does not carry everything, such as the // collations of strings, the fields of intervals and user-defined types, which still apply to // the data that Spark reads. @@ -483,88 +482,80 @@ class NativePartitionReader( } override def next(): Boolean = { - val array = ArrowArray.allocateNew(allocator) + val array = ArrowCData.allocate(ArrowCData.ARRAY_SIZE) try { - call(name, "read")(stream.getNext(array)) - val snapshot = array.snapshot() - if (snapshot.release == 0) { + call(name, "read")(ArrowCData.streamNext(stream, array)) + if (ArrowCData.isArrayReleased(array)) { // The end of the stream. false } else { - try { - checkArrays(snapshot) - } catch { - case e: Throwable => - array.release() - throw e - } - call(name, "read")(Data.importIntoVectorSchemaRoot(allocator, array, root, dictionaries)) + checkArrays(array) + // Moves the array, which the root then releases. + call(name, "read")(ArrowCData.importIntoRoot(allocator, array, root)) batch.setNumRows(root.getRowCount) true } } finally { - array.close() + // Releases the array unless it was imported. + Utils.tryWithSafeFinally(ArrowCData.releaseArray(array))(ArrowCData.free(array)) } } /** * Checks the array of a batch, and the arrays of its columns at every level, before importing - * them. Arrow Java trusts the arrays, and would read them wrongly otherwise: - * - An array must have an offset of 0. Arrow Java ignores the offsets of the arrays it imports, - * so it would read a sliced array from the start of its buffers. + * them into Arrow Java vectors, which would read them wrongly otherwise: + * - An array must have an offset of 0, since the vectors have no offsets. The import rejects + * other offsets too, but this check names the column. * - An array must have at least as many values as its parent refers to: a column as many as * the batch has rows, a field as many as its struct has values, and the elements of a list or - * a map as many as its last offset says. Arrow Java would extend a shorter array with nulls. + * a map as many as its last offset says. The vectors would extend a shorter array with nulls. * Only the lengths that the library exported show this: the import sets the length of each - * array to what its parent refers to. - * The buffers are not checked, since the C data interface does not carry their sizes: Arrow - * Java takes them from the lengths, so the library must export buffers that hold every value. + * vector to what its parent refers to. + * The buffers are not checked, since the C data interface does not carry their sizes: the + * import takes them from the lengths, so the library must export buffers that hold every value. */ - private def checkArrays(batch: ArrowArray.Snapshot): Unit = { - def children(array: ArrowArray.Snapshot): Seq[ArrowArray.Snapshot] = { - (0 until array.n_children.toInt).map { i => - ArrowArray.wrap(Platform.getLong(null, array.children + 8L * i)).snapshot() - } - } - // The last offset of a list or a map, in its second buffer. Arrow Java rejects the array if + private def checkArrays(batch: Long): Unit = { + import ArrowCData.{child, childCount, length, offset} + def children(array: Long): Seq[Long] = (0 until childCount(array).toInt).map(child(array, _)) + // The last offset of a list or a map, in its second buffer. The import rejects the array if // the buffer is missing. - def lastOffset(array: ArrowArray.Snapshot): Long = { - val offsets = if (array.length > 0 && array.n_buffers >= 2) { - Platform.getLong(null, array.buffers + 8L) + def lastOffset(array: Long): Long = { + val offsets = if (length(array) > 0 && ArrowCData.bufferCount(array) >= 2) { + ArrowCData.buffer(array, 1) } else { 0L } - if (offsets == 0) 0 else Platform.getInt(null, offsets + 4L * array.length) + if (offsets == 0) 0 else Platform.getInt(null, offsets + 4L * length(array)) } // The number of values that an array refers to in each of its children. - def referenced(array: ArrowArray.Snapshot, arrowType: ArrowType): Long = arrowType match { - case _: ArrowType.Struct => array.length + def referenced(array: Long, arrowType: ArrowType): Long = arrowType match { + case _: ArrowType.Struct => length(array) case _: ArrowType.List | _: ArrowType.Map => lastOffset(array) // The other types that Spark reads have no children. List views are rejected before. case _ => 0 } - def checkColumn(array: ArrowArray.Snapshot, field: Field, column: String): Unit = { - if (array.offset != 0) { - throw error(name, "read", s"The column $column has an offset of ${array.offset}, but " + + def checkColumn(array: Long, field: Field, column: String): Unit = { + if (offset(array) != 0) { + throw error(name, "read", s"The column $column has an offset of ${offset(array)}, but " + "Spark only reads arrays with an offset of 0. Copy sliced arrays before exporting them.") } - val length = referenced(array, field.getType) + val referencedLength = referenced(array, field.getType) children(array).zip(field.getChildren.asScala).foreach { case (child, childField) => - if (child.length < length) { + if (length(child) < referencedLength) { throw error(name, "read", s"The column $column contains an array with a length of " + - s"${child.length}, but its parent refers to $length values.") + s"${length(child)}, but its parent refers to $referencedLength values.") } checkColumn(child, childField, column) } } - if (batch.offset != 0) { - throw error(name, "read", s"The batch has an offset of ${batch.offset}, but Spark only " + + if (offset(batch) != 0) { + throw error(name, "read", s"The batch has an offset of ${offset(batch)}, but Spark only " + "reads arrays with an offset of 0. Copy sliced arrays before exporting them.") } children(batch).zip(root.getSchema.getFields.asScala).foreach { case (column, field) => - if (column.length < batch.length) { + if (length(column) < length(batch)) { throw error(name, "read", s"The column ${field.getName} has a length of " + - s"${column.length}, but the batch has a length of ${batch.length}.") + s"${length(column)}, but the batch has a length of ${length(batch)}.") } checkColumn(column, field, field.getName) } @@ -575,12 +566,11 @@ class NativePartitionReader( override def close(): Unit = { Utils.tryWithSafeFinally { if (root != null) { + // Releases the last array imported. root.close() } - dictionaries.close() // Releases the stream of the library, unless it was not exported. - stream.release() - stream.close() + Utils.tryWithSafeFinally(ArrowCData.releaseStream(stream))(ArrowCData.free(stream)) }(Utils.tryLogNonFatalError(allocator.close())) } } @@ -616,19 +606,21 @@ class NativeDataWriter( } val root = new VectorSchemaRoot( vectors.map(_.getField).asJava, vectors.asJava, batch.numRows()) - val array = ArrowArray.allocateNew(ArrowUtils.rootAllocator) - val arrowSchema = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) + val array = ArrowCData.allocate(ArrowCData.ARRAY_SIZE) + val arrowSchema = ArrowCData.allocate(ArrowCData.SCHEMA_SIZE) try { - Data.exportVectorSchemaRoot(ArrowUtils.rootAllocator, root, null, array, arrowSchema) - call(name, "write to") { - library.write(handle, array.memoryAddress(), arrowSchema.memoryAddress()) - } + ArrowCData.exportSchema(root.getSchema, arrowSchema) + ArrowCData.exportRoot(ArrowUtils.rootAllocator, root, array) + call(name, "write to")(library.write(handle, array, arrowSchema)) } finally { // Releases the batch unless the library moved it to take ownership. - array.release() - array.close() - arrowSchema.release() - arrowSchema.close() + Utils.tryWithSafeFinally { + ArrowCData.releaseArray(array) + ArrowCData.releaseSchema(arrowSchema) + } { + ArrowCData.free(array) + ArrowCData.free(arrowSchema) + } } } @@ -728,15 +720,7 @@ object NativeCalls { } } - // How Arrow reports the error message of an Arrow C stream. - private val StreamError = """(?s)\[errno \d+\] CDataJniException\{errno=\d+, message=(.*)\}""".r - - private[ffi] def message(e: Throwable): String = { - Option(e.getMessage).getOrElse(e.toString) match { - case StreamError(message) => message - case message => message - } - } + private[ffi] def message(e: Throwable): String = Option(e.getMessage).getOrElse(e.toString) /** Calls a function that the library may not implement. Returns None if it does not. */ def optional[T](f: => T): Option[T] = { @@ -764,31 +748,26 @@ object NativeArrow { val arrowSchema = ArrowUtils.toArrowSchema( schema, SQLConf.get.sessionLocalTimeZone, errorOnDuplicatedFieldNames = true, largeVarTypes = false) - val struct = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) + val struct = ArrowCData.allocate(ArrowCData.SCHEMA_SIZE) try { - Data.exportSchema(ArrowUtils.rootAllocator, arrowSchema, null, struct) - f(struct.memoryAddress()) + ArrowCData.exportSchema(arrowSchema, struct) + f(struct) } finally { // Releases the schema unless the library moved it to take ownership. - struct.release() - struct.close() + Utils.tryWithSafeFinally(ArrowCData.releaseSchema(struct))(ArrowCData.free(struct)) } } /** Calls `f` with the address of an `ArrowSchema` to initialize, and imports it. */ def importSchema(f: Long => Unit): StructType = { - val struct = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) - var imported = false + val struct = ArrowCData.allocate(ArrowCData.SCHEMA_SIZE) try { - f(struct.memoryAddress()) - imported = true - // Releases and closes the struct. - ArrowUtils.fromArrowSchema(Data.importSchema(ArrowUtils.rootAllocator, struct, null)) + f(struct) + // Releases the schema. + ArrowUtils.fromArrowSchema(ArrowCData.importSchema(struct)) } finally { - if (!imported) { - struct.release() - struct.close() - } + // Releases the schema if the library initialized it but it was not imported. + Utils.tryWithSafeFinally(ArrowCData.releaseSchema(struct))(ArrowCData.free(struct)) } } } diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index 53468b84a8173..782d9af407850 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -962,6 +962,8 @@ int spark_ds_read(const uint8_t* reader_state, int64_t reader_state_length, return guarded([&] { Options state = deserialize(toBytes(reader_state, reader_state_length)); maybeFail(state, "open"); + // Returns success without initializing the stream, which Spark must reject. + if (get(state, "fail", "") == "uninitialized") return; auto bounds = split(toBytes(partition, partition_length)); if (bounds.size() != 2) throw std::runtime_error("invalid partition"); stream->get_schema = streamGetSchema; diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCDataSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCDataSuite.scala new file mode 100644 index 0000000000000..bf04116a6ad1a --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCDataSuite.scala @@ -0,0 +1,463 @@ +/* + * 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.execution.datasources.v2.ffi + +import java.nio.charset.StandardCharsets.UTF_8 + +import scala.jdk.CollectionConverters._ +import scala.util.Random + +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.{FieldVector, LargeVarBinaryVector, LargeVarCharVector, VectorSchemaRoot, ViewVarBinaryVector, ViewVarCharVector} +import org.apache.arrow.vector.compare.VectorEqualsVisitor +import org.apache.arrow.vector.types.{DateUnit, FloatingPointPrecision, IntervalUnit, TimeUnit, UnionMode} +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.{RandomDataGenerator, Row} +import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow} +import org.apache.spark.sql.execution.arrow.ArrowWriter +import org.apache.spark.sql.types._ +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.unsafe.Platform + +/** + * Tests the Arrow C data interface of [[ArrowCData]], by exporting and importing in the JVM. The + * tests are skipped before Java 22, where the FFM API is not available. + */ +class ArrowCDataSuite extends SparkFunSuite { + + private def ffmTest(name: String)(f: => Unit): Unit = test(name) { + assume(NativeLibrary.isSupported, "Java 22 or later is required.") + f + } + + /** Runs `f` with an allocator, and checks that all its memory is released afterwards. */ + private def withAllocator[T](f: BufferAllocator => T): T = { + val allocator = ArrowUtils.rootAllocator.newChildAllocator("test", 0, Long.MaxValue) + try { + val result = f(allocator) + assert(allocator.getAllocatedMemory == 0, "Memory was not released.") + result + } finally { + allocator.close() + } + } + + private def withStruct[T](size: Long)(f: Long => T): T = { + val address = ArrowCData.allocate(size) + try f(address) finally ArrowCData.free(address) + } + + /** Exports and imports a schema. */ + private def roundTrip(schema: Schema): Schema = withStruct(ArrowCData.SCHEMA_SIZE) { struct => + ArrowCData.exportSchema(schema, struct) + ArrowCData.importSchema(struct) + } + + /** + * Exports the root, imports it into a new root with the imported schema, and passes the new + * root to `f`. The exported root is closed first, so the imported buffers must stay valid. + */ + private def roundTrip[T](allocator: BufferAllocator, root: VectorSchemaRoot)( + f: VectorSchemaRoot => T): T = { + val imported = VectorSchemaRoot.create(roundTrip(root.getSchema), allocator) + try { + withStruct(ArrowCData.ARRAY_SIZE) { array => + ArrowCData.exportRoot(allocator, root, array) + ArrowCData.importIntoRoot(allocator, array, imported) + assert(ArrowCData.isArrayReleased(array), "The array was not moved.") + } + f(imported) + } finally { + imported.close() + } + } + + private def assertSameVectors(expected: VectorSchemaRoot, actual: VectorSchemaRoot): Unit = { + assert(actual.getRowCount == expected.getRowCount) + expected.getFieldVectors.asScala.zip(actual.getFieldVectors.asScala).foreach { + case (e, a) => assert(VectorEqualsVisitor.vectorEquals(e, a), s"Different ${e.getName}.") + } + } + + private def rows(root: VectorSchemaRoot, schema: StructType): Seq[Row] = { + val vectors = root.getFieldVectors.asScala.map(v => new ArrowColumnVector(v): ColumnVector) + val batch = new ColumnarBatch(vectors.toArray, root.getRowCount) + val toScala = CatalystTypeConverters.createToScalaConverter(schema) + // Converted without copying the rows: ColumnarArray.copy fails on primitive arrays with nulls. + batch.rowIterator().asScala.map(row => toScala(row).asInstanceOf[Row]).toList + } + + private val schema = StructType(Seq( + StructField("boolean", BooleanType), + StructField("byte", ByteType), + StructField("short", ShortType), + StructField("int", IntegerType, nullable = false), + StructField("long", LongType), + StructField("float", FloatType), + StructField("double", DoubleType), + StructField("decimal", DecimalType(10, 2)), + StructField("wide_decimal", DecimalType(38, 18)), + StructField("string", StringType), + StructField("binary", BinaryType), + StructField("date", DateType), + StructField("timestamp", TimestampType), + StructField("timestamp_ntz", TimestampNTZType), + StructField("year_month", YearMonthIntervalType()), + StructField("day_time", DayTimeIntervalType()), + StructField("null", NullType), + StructField("ints", ArrayType(IntegerType)), + StructField("strings", ArrayType(StringType, containsNull = false)), + StructField("map", MapType(StringType, IntegerType)), + StructField("nested_map", MapType(IntegerType, ArrayType(LongType))), + StructField("struct", StructType(Seq( + StructField("a", IntegerType), StructField("b", StringType)))), + StructField("structs", ArrayType(StructType(Seq( + StructField("x", DoubleType), StructField("y", ArrayType(StringType)))))))) + + private def write(allocator: BufferAllocator, schema: StructType, rows: Seq[Row]) + : VectorSchemaRoot = { + val root = VectorSchemaRoot.create( + ArrowUtils.toArrowSchema(schema, "UTC", errorOnDuplicatedFieldNames = true, + largeVarTypes = false), + allocator) + val writer = ArrowWriter.create(root) + val toCatalyst = CatalystTypeConverters.createToCatalystConverter(schema) + rows.foreach(row => writer.write(toCatalyst(row).asInstanceOf[InternalRow])) + writer.finish() + root + } + + private def randomRows(schema: StructType, count: Int): Seq[Row] = { + val random = new Random(42) + Seq.fill(count)(RandomDataGenerator.randomRow(random, schema)) + } + + ffmTest("round trip of schemas") { + // scalastyle:off nonascii + val metadata = Map("key" -> "value", "unicode \u00e9" -> "\u4e2d").asJava + // scalastyle:on nonascii + val arrowSchema = new Schema(Seq( + new Field("a", new FieldType(true, new ArrowType.Int(32, true), null, metadata), null), + new Field("m", FieldType.notNullable(new ArrowType.Map(true)), Seq( + new Field("entries", FieldType.notNullable(ArrowType.Struct.INSTANCE), Seq( + new Field("key", FieldType.notNullable(ArrowType.Utf8.INSTANCE), null), + new Field("value", FieldType.nullable(ArrowType.Utf8.INSTANCE), null)).asJava)).asJava), + new Field("ts", FieldType.nullable( + new ArrowType.Timestamp(TimeUnit.MICROSECOND, "America/Los_Angeles")), null)).asJava, + Map("schema" -> "metadata").asJava) + assert(roundTrip(arrowSchema) == arrowSchema) + val sparkSchema = ArrowUtils.toArrowSchema( + schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = false) + assert(roundTrip(sparkSchema) == sparkSchema) + } + + ffmTest("round trip of arrays of the Spark types") { + withAllocator { allocator => + val expected = randomRows(schema, 100) + val root = write(allocator, schema, expected) + try { + roundTrip(allocator, root) { imported => + assertSameVectors(root, imported) + assert(rows(imported, schema) == expected) + } + } finally { + root.close() + } + } + } + + ffmTest("the imported buffers outlive the exported vectors") { + withAllocator { allocator => + val schema = StructType(Seq(StructField("s", StringType), StructField("l", LongType))) + val expected = randomRows(schema, 50) + val root = write(allocator, schema, expected) + val imported = VectorSchemaRoot.create(roundTrip(root.getSchema), allocator) + try { + withStruct(ArrowCData.ARRAY_SIZE) { array => + ArrowCData.exportRoot(allocator, root, array) + // The array retains the buffers of the vectors. + root.close() + ArrowCData.importIntoRoot(allocator, array, imported) + } + assert(rows(imported, schema) == expected) + } finally { + imported.close() + } + } + } + + ffmTest("round trip of empty arrays") { + withAllocator { allocator => + val root = write(allocator, schema, Nil) + try { + roundTrip(allocator, root) { imported => + assert(imported.getRowCount == 0) + assertSameVectors(root, imported) + } + } finally { + root.close() + } + } + } + + ffmTest("round trip of views and large variable-width types") { + withAllocator { allocator => + // scalastyle:off nonascii + val values = Seq("short", null, "a string longer than twelve bytes", "", "\u00e9\u4e2d") + // scalastyle:on nonascii + val utf8View = new ViewVarCharVector("utf8_view", allocator) + val binaryView = new ViewVarBinaryVector("binary_view", allocator) + val largeUtf8 = new LargeVarCharVector("large_utf8", allocator) + val largeBinary = new LargeVarBinaryVector("large_binary", allocator) + val vectors = Seq[FieldVector](utf8View, binaryView, largeUtf8, largeBinary) + vectors.foreach(_.allocateNew()) + values.zipWithIndex.foreach { + case (null, i) => + utf8View.setNull(i) + binaryView.setNull(i) + largeUtf8.setNull(i) + largeBinary.setNull(i) + case (value, i) => + val bytes = value.getBytes(UTF_8) + utf8View.setSafe(i, bytes) + binaryView.setSafe(i, bytes) + largeUtf8.setSafe(i, bytes) + largeBinary.setSafe(i, bytes) + } + vectors.foreach(_.setValueCount(values.length)) + val root = new VectorSchemaRoot(vectors.asJava) + try { + roundTrip(allocator, root) { imported => + assertSameVectors(root, imported) + val read = rows(imported, StructType(Seq( + StructField("utf8_view", StringType), StructField("binary_view", BinaryType), + StructField("large_utf8", StringType), StructField("large_binary", BinaryType)))) + assert(read.map(_.getString(0)) == values) + assert(read.map(_.getString(2)) == values) + } + } finally { + root.close() + } + } + } + + ffmTest("release a child that the consumer moved") { + withAllocator { allocator => + val schema = StructType(Seq(StructField("a", LongType), StructField("b", StringType))) + val root = write(allocator, schema, randomRows(schema, 10)) + try { + withStruct(ArrowCData.ARRAY_SIZE) { array => + ArrowCData.exportRoot(allocator, root, array) + withStruct(ArrowCData.ARRAY_SIZE) { moved => + // Moves the first child, as the C data interface allows: copies it, and marks the + // original as released, by setting its release callback, at offset 64, to NULL. + val child = ArrowCData.child(array, 0) + Platform.copyMemory(null, child, null, moved, ArrowCData.ARRAY_SIZE) + Platform.putLong(null, child + 64, 0L) + ArrowCData.releaseArray(array) + assert(ArrowCData.isArrayReleased(array)) + assert(!ArrowCData.isArrayReleased(moved)) + root.close() + // The moved child still retains its buffers. + assert(allocator.getAllocatedMemory > 0) + ArrowCData.releaseArray(moved) + assert(ArrowCData.isArrayReleased(moved)) + } + } + } finally { + root.close() + } + } + } + + ffmTest("formats") { + val types = Seq( + ArrowType.Null.INSTANCE, + ArrowType.Bool.INSTANCE, + new ArrowType.Int(8, true), new ArrowType.Int(8, false), + new ArrowType.Int(16, true), new ArrowType.Int(16, false), + new ArrowType.Int(32, true), new ArrowType.Int(32, false), + new ArrowType.Int(64, true), new ArrowType.Int(64, false), + new ArrowType.FloatingPoint(FloatingPointPrecision.HALF), + new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), + ArrowType.Binary.INSTANCE, ArrowType.LargeBinary.INSTANCE, ArrowType.BinaryView.INSTANCE, + ArrowType.Utf8.INSTANCE, ArrowType.LargeUtf8.INSTANCE, ArrowType.Utf8View.INSTANCE, + new ArrowType.FixedSizeBinary(5), + new ArrowType.Decimal(10, 2, 128), new ArrowType.Decimal(40, 5, 256), + new ArrowType.Date(DateUnit.DAY), new ArrowType.Date(DateUnit.MILLISECOND), + new ArrowType.Time(TimeUnit.SECOND, 32), new ArrowType.Time(TimeUnit.MILLISECOND, 32), + new ArrowType.Time(TimeUnit.MICROSECOND, 64), new ArrowType.Time(TimeUnit.NANOSECOND, 64), + new ArrowType.Interval(IntervalUnit.YEAR_MONTH), + new ArrowType.Interval(IntervalUnit.DAY_TIME), + new ArrowType.Interval(IntervalUnit.MONTH_DAY_NANO), + ArrowType.List.INSTANCE, ArrowType.LargeList.INSTANCE, + ArrowType.ListView.INSTANCE, ArrowType.LargeListView.INSTANCE, + new ArrowType.FixedSizeList(3), + ArrowType.Struct.INSTANCE, + new ArrowType.Map(false)) ++ + TimeUnit.values().toSeq.flatMap { unit => + Seq(new ArrowType.Timestamp(unit, null), new ArrowType.Timestamp(unit, "UTC"), + new ArrowType.Duration(unit)) + } + types.foreach { arrowType => + assert(ArrowCData.parseFormat(ArrowCData.format(arrowType), 0) == arrowType) + } + // The flags of a map say whether its keys are sorted. + assert(ArrowCData.parseFormat("+m", 4) == new ArrowType.Map(true)) + assert(ArrowCData.parseFormat("+ud:1,2", 0) == + new ArrowType.Union(UnionMode.Dense, Array(1, 2))) + assert(ArrowCData.parseFormat("tsu:+01:00", 0) == + new ArrowType.Timestamp(TimeUnit.MICROSECOND, "+01:00")) + Seq("", "x", "d:1", "d:a,b", "w:x", "+w:", "tsx:", "tsu").foreach { format => + val e = intercept[IllegalArgumentException](ArrowCData.parseFormat(format, 0)) + assert(e.getMessage == s"Unsupported Arrow format: $format") + } + } + + /** + * Builds ArrowSchema structs by hand, without release callbacks, so that importing them does not + * release them: the memory is freed afterwards with `free`. The fields are at the offsets of the + * struct: format 0, name 8, metadata 16, flags 24, n_children 32, children 40, dictionary 48. + */ + private class RawSchemas { + private val allocations = scala.collection.mutable.ArrayBuffer.empty[Long] + + private def allocate(size: Long): Long = { + val address = ArrowCData.allocate(size) + allocations += address + address + } + + private def string(value: String): Long = { + val bytes = value.getBytes(UTF_8) + val address = allocate(bytes.length + 1) + Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address, bytes.length) + address + } + + def schema( + format: String, + name: String, + children: Seq[Long] = Nil, + dictionary: Long = 0, + flags: Long = 2): Long = { + val address = allocate(ArrowCData.SCHEMA_SIZE) + Platform.putLong(null, address, string(format)) + Platform.putLong(null, address + 8, string(name)) + Platform.putLong(null, address + 24, flags) + Platform.putLong(null, address + 32, children.length) + if (children.nonEmpty) { + val pointers = allocate(8L * children.length) + children.zipWithIndex.foreach { case (child, i) => + Platform.putLong(null, pointers + 8L * i, child) + } + Platform.putLong(null, address + 40, pointers) + } + Platform.putLong(null, address + 48, dictionary) + address + } + + def free(): Unit = allocations.foreach(ArrowCData.free) + } + + ffmTest("import dictionary-encoded fields") { + val raw = new RawSchemas + try { + val dictionary = raw.schema("u", "") + val struct = raw.schema("+s", "", Seq(raw.schema("i", "x", dictionary = dictionary, + flags = 3)), flags = 0) + val field = ArrowCData.importSchema(struct).getFields.get(0) + assert(field.getName == "x") + assert(field.getDictionary != null) + assert(field.getDictionary.isOrdered) + assert(field.getDictionary.getIndexType == new ArrowType.Int(32, true)) + } finally { + raw.free() + } + } + + ffmTest("reject invalid schemas") { + val raw = new RawSchemas + try { + def checkInvalid(schema: Long, message: String): Unit = { + val e = intercept[IllegalArgumentException](ArrowCData.importSchema(schema)) + assert(e.getMessage == message) + } + checkInvalid(raw.schema("l", "x"), + "The schema is not a struct type, but Int(64, true).") + checkInvalid(raw.schema("+s", "", Seq(raw.schema("?", "x"))), + "Unsupported Arrow format: ?") + val dictionary = raw.schema("u", "") + checkInvalid(raw.schema("+s", "", Seq(raw.schema("u", "x", dictionary = dictionary))), + "The indices of the dictionary of x are not integers: u.") + var nested = raw.schema("i", "x") + (0 to 64).foreach(_ => nested = raw.schema("+s", "s", Seq(nested))) + checkInvalid(nested, "The schema is nested more than 64 levels deep.") + } finally { + raw.free() + } + } + + ffmTest("reject invalid arrays") { + withAllocator { allocator => + val schema = StructType(Seq(StructField("a", IntegerType), StructField("s", StringType))) + // The mutations keep what the release callback of the array frees, so nothing leaks. + def checkInvalid(message: String, importInto: Option[Schema] = None)( + mutate: Long => Unit): Unit = { + val root = write(allocator, schema, Seq(Row(1, "a"), Row(2, "b"))) + val imported = VectorSchemaRoot.create(importInto.getOrElse(root.getSchema), allocator) + try { + withStruct(ArrowCData.ARRAY_SIZE) { array => + ArrowCData.exportRoot(allocator, root, array) + mutate(array) + val e = intercept[IllegalArgumentException] { + ArrowCData.importIntoRoot(allocator, array, imported) + } + assert(e.getMessage == message) + assert(ArrowCData.isArrayReleased(array), "The array was not moved.") + } + } finally { + imported.close() + root.close() + } + } + val oneColumn = new Schema(Seq(new Field("a", FieldType.nullable( + new ArrowType.Int(32, true)), null)).asJava) + checkInvalid("The array has 2 columns, but the schema has 1.", Some(oneColumn))(_ => ()) + checkInvalid("The array has an offset of 1, but only arrays with an offset of 0 are " + + "supported.") { array => + Platform.putLong(null, ArrowCData.child(array, 0) + 16, 1L) + } + checkInvalid("The array has 1 nulls, but no validity buffer.") { array => + val column = ArrowCData.child(array, 0) + Platform.putLong(null, column + 8, 1L) + Platform.putLong(null, Platform.getLong(null, column + 40), 0L) + } + checkInvalid("The array of type Utf8 has 2 buffers, but 3 are expected.") { array => + Platform.putLong(null, ArrowCData.child(array, 1) + 24, 2L) + } + checkInvalid("The data buffer is NULL.") { array => + val column = ArrowCData.child(array, 1) + Platform.putLong(null, Platform.getLong(null, column + 40) + 16, 0L) + } + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 70842f4add1f3..0cf4c734aa85f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -186,7 +186,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { """"children":[{"type":"column","name":["id"]},""" + """{"type":"literal","dataType":"bigint","value":1},""" + """{"type":"literal","dataType":"bigint","value":2}]}]}""")) - // Only the functions listed in NativeBridge.pushPredicates are passed, so not arithmetic. + // Only the functions listed in spark_ds_push_predicates are passed, so not arithmetic. NativePredicates.FUNCTION_NAMES.foreach { name => assert(toJson(predicate(name, id, id)).nonEmpty, name) } @@ -584,6 +584,9 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { check("read", "injected failure in get_next") { session => range(session, "fail" -> "read").collect() } + check("read", "The stream is released, or was not initialized.") { session => + range(session, "fail" -> "uninitialized").collect() + } // An error of the library is not an unsupported operation. check("plan a scan of", "native_sink does not support batch reads", "native_sink") { session => session.read.schema("id BIGINT").format("native_sink").load().collect() From 45f428dd27895b95ec5d57f0b3e5661333b5e5cf Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Tue, 13 Oct 2026 16:27:20 +0900 Subject: [PATCH 17/27] Revert "Implement the Arrow C data interface with the FFM API, without arrow-c-data" This reverts commit f55105d4644b30d53177bf82bdcf4e3841a7987c, to keep using the implementation of the Arrow C data interface that Arrow maintains, arrow-c-data, rather than one of Spark's own. Co-authored-by: Isaac --- LICENSE-binary | 1 + dev/deps/spark-deps-hadoop-3-hive-2.3 | 1 + pom.xml | 19 + sql/core/pom.xml | 4 + .../datasources/v2/ffi/ArrowCData.java | 1062 ----------------- .../sql/execution/datasources/v2/ffi/Ffm.java | 208 ---- .../datasources/v2/ffi/NativeLibrary.java | 248 +++- .../spark/sql/datasource/spark_datasource.h | 6 +- .../datasources/v2/ffi/NativeDataSource.scala | 177 +-- .../test_native_datasource.cc | 2 - .../datasources/v2/ffi/ArrowCDataSuite.scala | 463 ------- .../v2/ffi/NativeDataSourceSuite.scala | 5 +- 12 files changed, 316 insertions(+), 1880 deletions(-) delete mode 100644 sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCData.java delete mode 100644 sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/Ffm.java delete mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCDataSuite.scala diff --git a/LICENSE-binary b/LICENSE-binary index 7051e09a6383a..53dec42594efa 100644 --- a/LICENSE-binary +++ b/LICENSE-binary @@ -301,6 +301,7 @@ joda-time:joda-time net.sf.opencsv:opencsv net.sf.supercsv:super-csv net.sf.jpam:jpam +org.apache.arrow:arrow-c-data org.apache.arrow:arrow-compression org.apache.arrow:arrow-format org.apache.arrow:arrow-memory-core diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index 50f978819c3da..b088d623ea7cb 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -21,6 +21,7 @@ antlr4-runtime/4.13.1//antlr4-runtime-4.13.1.jar aopalliance-repackaged/3.0.6//aopalliance-repackaged-3.0.6.jar arpack/3.2.0//arpack-3.2.0.jar arpack_combined_all/0.1//arpack_combined_all-0.1.jar +arrow-c-data/19.0.0//arrow-c-data-19.0.0.jar arrow-compression/19.0.0//arrow-compression-19.0.0.jar arrow-format/19.0.0//arrow-format-19.0.0.jar arrow-memory-core/19.0.0//arrow-memory-core-19.0.0.jar diff --git a/pom.xml b/pom.xml index da05308883570..c2eb3de50b8cd 100644 --- a/pom.xml +++ b/pom.xml @@ -2570,6 +2570,25 @@ + + org.apache.arrow + arrow-c-data + ${arrow.version} + + + com.fasterxml.jackson.core + jackson-annotations + + + com.fasterxml.jackson.core + jackson-core + + + io.netty + netty-common + + + org.apache.arrow arrow-compression diff --git a/sql/core/pom.xml b/sql/core/pom.xml index 281546b901a7a..7f3db15f70fdf 100644 --- a/sql/core/pom.xml +++ b/sql/core/pom.xml @@ -298,6 +298,10 @@ org.apache.arrow arrow-compression + + org.apache.arrow + arrow-c-data + org.apache.spark spark-udf-worker-proto_${scala.binary.version} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCData.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCData.java deleted file mode 100644 index 803949d08fd1c..0000000000000 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCData.java +++ /dev/null @@ -1,1062 +0,0 @@ -/* - * 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.execution.datasources.v2.ffi; - -import java.lang.invoke.MethodHandle; -import java.lang.invoke.MethodHandles; -import java.lang.invoke.MethodType; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; - -import org.apache.arrow.memory.ArrowBuf; -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.ForeignAllocation; -import org.apache.arrow.vector.BitVectorHelper; -import org.apache.arrow.vector.FieldVector; -import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.ipc.message.ArrowFieldNode; -import org.apache.arrow.vector.types.DateUnit; -import org.apache.arrow.vector.types.FloatingPointPrecision; -import org.apache.arrow.vector.types.IntervalUnit; -import org.apache.arrow.vector.types.TimeUnit; -import org.apache.arrow.vector.types.UnionMode; -import org.apache.arrow.vector.types.pojo.ArrowType; -import org.apache.arrow.vector.types.pojo.DictionaryEncoding; -import org.apache.arrow.vector.types.pojo.Field; -import org.apache.arrow.vector.types.pojo.FieldType; -import org.apache.arrow.vector.types.pojo.Schema; - -import org.apache.spark.internal.SparkLogger; -import org.apache.spark.internal.SparkLoggerFactory; -import org.apache.spark.sql.execution.datasources.v2.ffi.Ffm.C; -import org.apache.spark.unsafe.Platform; - -/** - * The Arrow C data - * interface, implemented with the FFM API, see {@link Ffm}: Spark exports and imports the - * schemas and the arrays that it exchanges with native data source libraries, and reads the - * Arrow streams that they export. - *

    - * The C structs are in native memory, at {@code long} addresses, with the layout of 64-bit - * platforms. A struct is released once its {@code release} callback is NULL. Arrays are imported - * without copying their buffers: the imported vectors refer to the memory of the array, which is - * released once the vectors no longer use any of its buffers. Imported arrays must have an offset - * of 0, at every level. - */ -final class ArrowCData { - - private ArrowCData() {} - - private static final SparkLogger LOG = SparkLoggerFactory.getLogger(ArrowCData.class); - - static final long SCHEMA_SIZE = 72; - static final long ARRAY_SIZE = 80; - static final long STREAM_SIZE = 40; - - // The fields of struct ArrowSchema. - private static final long SCHEMA_FORMAT = 0; - private static final long SCHEMA_NAME = 8; - private static final long SCHEMA_METADATA = 16; - private static final long SCHEMA_FLAGS = 24; - private static final long SCHEMA_N_CHILDREN = 32; - private static final long SCHEMA_CHILDREN = 40; - private static final long SCHEMA_DICTIONARY = 48; - private static final long SCHEMA_RELEASE = 56; - private static final long SCHEMA_PRIVATE_DATA = 64; - - // The fields of struct ArrowArray. - private static final long ARRAY_LENGTH = 0; - private static final long ARRAY_NULL_COUNT = 8; - private static final long ARRAY_OFFSET = 16; - private static final long ARRAY_N_BUFFERS = 24; - private static final long ARRAY_N_CHILDREN = 32; - private static final long ARRAY_BUFFERS = 40; - private static final long ARRAY_CHILDREN = 48; - private static final long ARRAY_DICTIONARY = 56; - private static final long ARRAY_RELEASE = 64; - private static final long ARRAY_PRIVATE_DATA = 72; - - // The fields of struct ArrowArrayStream. - private static final long STREAM_GET_SCHEMA = 0; - private static final long STREAM_GET_NEXT = 8; - private static final long STREAM_GET_LAST_ERROR = 16; - private static final long STREAM_RELEASE = 24; - - private static final long FLAG_DICTIONARY_ORDERED = 1; - private static final long FLAG_NULLABLE = 2; - private static final long FLAG_MAP_KEYS_SORTED = 4; - - // Deeper schemas and arrays are rejected, rather than overflowing the stack. - private static final int MAX_DEPTH = 64; - - /** The callbacks of the C structs, created when they are first used. */ - private static final class Callbacks { - // int (*)(void*, void*), such as get_next. - static final MethodHandle INT_FUNCTION; - // const char* (*)(void*), such as get_last_error. - static final MethodHandle POINTER_FUNCTION; - // void (*)(void*), such as release. - static final MethodHandle VOID_FUNCTION; - // The release callbacks of the structs that Spark exports. - static final long RELEASE_SCHEMA; - static final long RELEASE_ARRAY; - - static { - try { - INT_FUNCTION = Ffm.bindPointer(C.INT32, C.POINTER, C.POINTER); - POINTER_FUNCTION = Ffm.bindPointer(C.POINTER, C.POINTER); - VOID_FUNCTION = Ffm.bindPointer(C.VOID, C.POINTER); - MethodHandles.Lookup lookup = MethodHandles.lookup(); - MethodType release = MethodType.methodType(void.class, long.class); - RELEASE_SCHEMA = Ffm.upcall( - lookup.findStatic(ArrowCData.class, "releaseExportedSchema", release), - C.VOID, C.POINTER); - RELEASE_ARRAY = Ffm.upcall( - lookup.findStatic(ArrowCData.class, "releaseExportedArray", release), - C.VOID, C.POINTER); - } catch (Throwable e) { - throw new IllegalStateException("Cannot create the Arrow C data callbacks.", e); - } - } - } - - // --------------------------------------------------------------------------------------------- - // Structs - // --------------------------------------------------------------------------------------------- - - /** Allocates a struct in native memory, initialized to zero, which is a released struct. */ - static long allocate(long size) { - long address = Platform.allocateMemory(size); - Platform.setMemory(address, (byte) 0, size); - return address; - } - - static void free(long address) { - Platform.freeMemory(address); - } - - static boolean isSchemaReleased(long schema) { - return Platform.getLong(null, schema + SCHEMA_RELEASE) == 0; - } - - static boolean isArrayReleased(long array) { - return Platform.getLong(null, array + ARRAY_RELEASE) == 0; - } - - static long length(long array) { - return Platform.getLong(null, array + ARRAY_LENGTH); - } - - static long offset(long array) { - return Platform.getLong(null, array + ARRAY_OFFSET); - } - - static long bufferCount(long array) { - return Platform.getLong(null, array + ARRAY_N_BUFFERS); - } - - /** The address of the buffer at the given index of an array, which may be 0 (NULL). */ - static long buffer(long array, int index) { - return Platform.getLong(null, Platform.getLong(null, array + ARRAY_BUFFERS) + 8L * index); - } - - static long childCount(long array) { - return Platform.getLong(null, array + ARRAY_N_CHILDREN); - } - - static long child(long array, int index) { - return Platform.getLong(null, Platform.getLong(null, array + ARRAY_CHILDREN) + 8L * index); - } - - /** Releases a schema with its release callback, unless it is released already. */ - static void releaseSchema(long schema) throws Throwable { - release(schema, SCHEMA_RELEASE); - } - - /** Releases an array with its release callback, unless it is released already. */ - static void releaseArray(long array) throws Throwable { - release(array, ARRAY_RELEASE); - } - - /** Releases a stream with its release callback, unless it is released already. */ - static void releaseStream(long stream) throws Throwable { - release(stream, STREAM_RELEASE); - } - - private static void release(long struct, long releaseField) throws Throwable { - long release = Platform.getLong(null, struct + releaseField); - if (release != 0) { - Callbacks.VOID_FUNCTION.invokeExact(release, struct); - // The callback marks the struct as released. This makes sure that it is not released again. - Platform.putLong(null, struct + releaseField, 0L); - } - } - - // --------------------------------------------------------------------------------------------- - // Streams - // --------------------------------------------------------------------------------------------- - - /** Calls get_schema of a stream, which initializes the schema at {@code out}. */ - static void streamSchema(long stream, long out) throws Throwable { - long getSchema = callback(stream, STREAM_GET_SCHEMA); - checkStream(stream, (int) Callbacks.INT_FUNCTION.invokeExact(getSchema, stream, out)); - } - - /** - * Calls get_next of a stream, which initializes the array at {@code out}, or marks it as - * released at the end of the stream. - */ - static void streamNext(long stream, long out) throws Throwable { - long getNext = callback(stream, STREAM_GET_NEXT); - checkStream(stream, (int) Callbacks.INT_FUNCTION.invokeExact(getNext, stream, out)); - } - - /** A callback of a stream, which must not be NULL: calling NULL would crash the JVM. */ - private static long callback(long stream, long field) { - if (Platform.getLong(null, stream + STREAM_RELEASE) == 0) { - throw new IllegalStateException("The stream is released, or was not initialized."); - } - long callback = Platform.getLong(null, stream + field); - if (callback == 0) { - throw new IllegalStateException("A callback of the stream is NULL."); - } - return callback; - } - - private static void checkStream(long stream, int errno) throws Throwable { - if (errno == 0) { - return; - } - long getLastError = Platform.getLong(null, stream + STREAM_GET_LAST_ERROR); - long message = getLastError == 0 - ? 0 : (long) Callbacks.POINTER_FUNCTION.invokeExact(getLastError, stream); - String text = message == 0 ? "" : readString(message); - throw new RuntimeException(text.isEmpty() ? "The stream failed with the error " + errno + "." - : text); - } - - // --------------------------------------------------------------------------------------------- - // Exporting schemas - // --------------------------------------------------------------------------------------------- - - /** - * The memory of an exported struct, freed by its release callback. It does not include the - * struct itself, which belongs to its parent or to whoever allocated it. - */ - private static final class Exported { - final List allocations = new ArrayList<>(); - final List buffers = new ArrayList<>(); - - long allocate(long size) { - long address = ArrowCData.allocate(size); - allocations.add(address); - return address; - } - - long string(String value) { - return bytes(value.getBytes(StandardCharsets.UTF_8), true); - } - - long bytes(byte[] bytes, boolean terminate) { - long address = allocate(bytes.length + (terminate ? 1 : 0)); - Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address, bytes.length); - return address; - } - - void free() { - for (ArrowBuf buffer : buffers) { - // Null for a buffer exported as NULL. - if (buffer != null) { - buffer.getReferenceManager().release(); - } - } - allocations.forEach(Platform::freeMemory); - } - } - - private static final AtomicLong nextExportedId = new AtomicLong(1); - private static final ConcurrentHashMap exported = new ConcurrentHashMap<>(); - - /** Exports a schema as a struct type whose fields are the fields of the schema. */ - static void exportSchema(Schema schema, long out) throws Throwable { - exportField( - new Field("", new FieldType(false, ArrowType.Struct.INSTANCE, null, - schema.getCustomMetadata()), schema.getFields()), - out); - } - - private static void exportField(Field field, long out) throws Throwable { - if (field.getDictionary() != null) { - throw new UnsupportedOperationException( - "Dictionary-encoded fields cannot be exported: " + field); - } - String format = format(field.getType()); - List children = field.getChildren(); - long flags = field.isNullable() ? FLAG_NULLABLE : 0; - if (field.getType() instanceof ArrowType.Map map && map.getKeysSorted()) { - flags |= FLAG_MAP_KEYS_SORTED; - } - // A released struct with its callback, so that a failure below releases what was exported. - Platform.setMemory(out, (byte) 0, SCHEMA_SIZE); - Exported data = register(out, SCHEMA_PRIVATE_DATA); - Platform.putLong(null, out + SCHEMA_FLAGS, flags); - Platform.putLong(null, out + SCHEMA_RELEASE, Callbacks.RELEASE_SCHEMA); - try { - Platform.putLong(null, out + SCHEMA_FORMAT, data.string(format)); - String name = field.getName(); - Platform.putLong(null, out + SCHEMA_NAME, name == null ? 0L : data.string(name)); - Platform.putLong(null, out + SCHEMA_METADATA, encodeMetadata(data, field.getMetadata())); - Platform.putLong( - null, out + SCHEMA_CHILDREN, childStructs(data, children.size(), SCHEMA_SIZE)); - for (int i = 0; i < children.size(); i++) { - exportField(children.get(i), schemaChild(out, i)); - // Counted once exported, so that a failure releases the exported children only. - Platform.putLong(null, out + SCHEMA_N_CHILDREN, i + 1L); - } - } catch (Throwable e) { - releaseExportedSchema(out); - throw e; - } - } - - private static long schemaChild(long schema, int index) { - return Platform.getLong(null, Platform.getLong(null, schema + SCHEMA_CHILDREN) + 8L * index); - } - - /** Registers the private data of a struct being exported. */ - private static Exported register(long struct, long privateDataField) { - long id = nextExportedId.getAndIncrement(); - Exported data = new Exported(); - exported.put(id, data); - Platform.putLong(null, struct + privateDataField, id); - return data; - } - - /** - * Allocates the children of an exported struct: an array of pointers to the children, followed - * by the children themselves, initialized to zero. Returns the address of the pointers. - */ - private static long childStructs(Exported data, int count, long structSize) { - if (count == 0) { - return 0; - } - long pointers = data.allocate(count * (8 + structSize)); - for (int i = 0; i < count; i++) { - Platform.putLong(null, pointers + 8L * i, pointers + 8L * count + structSize * i); - } - return pointers; - } - - /** The release callback of the schemas that Spark exports. */ - private static void releaseExportedSchema(long schema) { - try { - long children = Platform.getLong(null, schema + SCHEMA_CHILDREN); - long count = Platform.getLong(null, schema + SCHEMA_N_CHILDREN); - for (int i = 0; i < count; i++) { - // A child that was moved is marked as released here, and released where it was moved. - long child = Platform.getLong(null, children + 8L * i); - if (!isSchemaReleased(child)) { - releaseExportedSchema(child); - } - } - Exported data = exported.remove(Platform.getLong(null, schema + SCHEMA_PRIVATE_DATA)); - if (data != null) { - data.free(); - } - } catch (Throwable e) { - // An exception must not be thrown to the caller of the callback, which is native code. - LOG.warn("Failed to release an exported Arrow schema.", e); - } - Platform.putLong(null, schema + SCHEMA_RELEASE, 0L); - } - - // --------------------------------------------------------------------------------------------- - // Importing schemas - // --------------------------------------------------------------------------------------------- - - /** - * Imports a schema, which must be a struct type whose fields are the fields of the schema, and - * releases it. - */ - static Schema importSchema(long schema) throws Throwable { - try { - Field field = importField(schema, 0, new AtomicLong()); - if (!(field.getType() instanceof ArrowType.Struct)) { - throw new IllegalArgumentException( - "The schema is not a struct type, but " + field.getType() + "."); - } - return new Schema(field.getChildren(), field.getMetadata()); - } finally { - releaseSchema(schema); - } - } - - private static Field importField(long schema, int depth, AtomicLong dictionaryIds) { - if (depth > MAX_DEPTH) { - throw new IllegalArgumentException("The schema is nested more than " + MAX_DEPTH + - " levels deep."); - } - long formatAddress = Platform.getLong(null, schema + SCHEMA_FORMAT); - if (formatAddress == 0) { - throw new IllegalArgumentException("The schema has no format."); - } - String format = readString(formatAddress); - long nameAddress = Platform.getLong(null, schema + SCHEMA_NAME); - String name = nameAddress == 0 ? "" : readString(nameAddress); - Map metadata = - decodeMetadata(Platform.getLong(null, schema + SCHEMA_METADATA)); - long flags = Platform.getLong(null, schema + SCHEMA_FLAGS); - long count = Platform.getLong(null, schema + SCHEMA_N_CHILDREN); - List children = new ArrayList<>(); - for (int i = 0; i < count; i++) { - children.add(importField(schemaChild(schema, i), depth + 1, dictionaryIds)); - } - ArrowType type = parseFormat(format, flags); - DictionaryEncoding dictionary = null; - long dictionarySchema = Platform.getLong(null, schema + SCHEMA_DICTIONARY); - if (dictionarySchema != 0) { - if (!(type instanceof ArrowType.Int indexType)) { - throw new IllegalArgumentException( - "The indices of the dictionary of " + name + " are not integers: " + format + "."); - } - // Only the encoding is kept: Spark does not read dictionary-encoded data. - importField(dictionarySchema, depth + 1, dictionaryIds); - dictionary = new DictionaryEncoding(dictionaryIds.getAndIncrement(), - (flags & FLAG_DICTIONARY_ORDERED) != 0, indexType); - } - boolean nullable = (flags & FLAG_NULLABLE) != 0; - return new Field(name, new FieldType(nullable, type, dictionary, metadata), children); - } - - // --------------------------------------------------------------------------------------------- - // Formats - // --------------------------------------------------------------------------------------------- - - /** The format string of an Arrow type. */ - static String format(ArrowType type) { - switch (type.getTypeID()) { - case Null: return "n"; - case Bool: return "b"; - case Int: { - ArrowType.Int t = (ArrowType.Int) type; - String formats = t.getIsSigned() ? "csil" : "CSIL"; - switch (t.getBitWidth()) { - case 8: return formats.substring(0, 1); - case 16: return formats.substring(1, 2); - case 32: return formats.substring(2, 3); - case 64: return formats.substring(3, 4); - default: break; - } - break; - } - case FloatingPoint: { - switch (((ArrowType.FloatingPoint) type).getPrecision()) { - case HALF: return "e"; - case SINGLE: return "f"; - case DOUBLE: return "g"; - default: break; - } - break; - } - case Binary: return "z"; - case LargeBinary: return "Z"; - case BinaryView: return "vz"; - case Utf8: return "u"; - case LargeUtf8: return "U"; - case Utf8View: return "vu"; - case FixedSizeBinary: - return "w:" + ((ArrowType.FixedSizeBinary) type).getByteWidth(); - case Decimal: { - ArrowType.Decimal t = (ArrowType.Decimal) type; - return "d:" + t.getPrecision() + "," + t.getScale() + - (t.getBitWidth() == 128 ? "" : "," + t.getBitWidth()); - } - case Date: - return ((ArrowType.Date) type).getUnit() == DateUnit.DAY ? "tdD" : "tdm"; - case Time: { - ArrowType.Time t = (ArrowType.Time) type; - return "tt" + unit(t.getUnit()); - } - case Timestamp: { - ArrowType.Timestamp t = (ArrowType.Timestamp) type; - String timezone = t.getTimezone(); - return "ts" + unit(t.getUnit()) + ":" + (timezone == null ? "" : timezone); - } - case Duration: - return "tD" + unit(((ArrowType.Duration) type).getUnit()); - case Interval: { - switch (((ArrowType.Interval) type).getUnit()) { - case YEAR_MONTH: return "tiM"; - case DAY_TIME: return "tiD"; - case MONTH_DAY_NANO: return "tin"; - default: break; - } - break; - } - case List: return "+l"; - case LargeList: return "+L"; - case ListView: return "+vl"; - case LargeListView: return "+vL"; - case FixedSizeList: - return "+w:" + ((ArrowType.FixedSizeList) type).getListSize(); - case Struct: return "+s"; - case Map: return "+m"; - default: break; - } - throw new UnsupportedOperationException("The Arrow type " + type + " cannot be exported."); - } - - private static char unit(TimeUnit unit) { - return switch (unit) { - case SECOND -> 's'; - case MILLISECOND -> 'm'; - case MICROSECOND -> 'u'; - case NANOSECOND -> 'n'; - }; - } - - /** The Arrow type of a format string. */ - static ArrowType parseFormat(String format, long flags) { - switch (format) { - case "n": return ArrowType.Null.INSTANCE; - case "b": return ArrowType.Bool.INSTANCE; - case "c": return new ArrowType.Int(8, true); - case "C": return new ArrowType.Int(8, false); - case "s": return new ArrowType.Int(16, true); - case "S": return new ArrowType.Int(16, false); - case "i": return new ArrowType.Int(32, true); - case "I": return new ArrowType.Int(32, false); - case "l": return new ArrowType.Int(64, true); - case "L": return new ArrowType.Int(64, false); - case "e": return new ArrowType.FloatingPoint(FloatingPointPrecision.HALF); - case "f": return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); - case "g": return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); - case "z": return ArrowType.Binary.INSTANCE; - case "Z": return ArrowType.LargeBinary.INSTANCE; - case "vz": return ArrowType.BinaryView.INSTANCE; - case "u": return ArrowType.Utf8.INSTANCE; - case "U": return ArrowType.LargeUtf8.INSTANCE; - case "vu": return ArrowType.Utf8View.INSTANCE; - case "tdD": return new ArrowType.Date(DateUnit.DAY); - case "tdm": return new ArrowType.Date(DateUnit.MILLISECOND); - case "tts": return new ArrowType.Time(TimeUnit.SECOND, 32); - case "ttm": return new ArrowType.Time(TimeUnit.MILLISECOND, 32); - case "ttu": return new ArrowType.Time(TimeUnit.MICROSECOND, 64); - case "ttn": return new ArrowType.Time(TimeUnit.NANOSECOND, 64); - case "tDs": return new ArrowType.Duration(TimeUnit.SECOND); - case "tDm": return new ArrowType.Duration(TimeUnit.MILLISECOND); - case "tDu": return new ArrowType.Duration(TimeUnit.MICROSECOND); - case "tDn": return new ArrowType.Duration(TimeUnit.NANOSECOND); - case "tiM": return new ArrowType.Interval(IntervalUnit.YEAR_MONTH); - case "tiD": return new ArrowType.Interval(IntervalUnit.DAY_TIME); - case "tin": return new ArrowType.Interval(IntervalUnit.MONTH_DAY_NANO); - case "+l": return ArrowType.List.INSTANCE; - case "+L": return ArrowType.LargeList.INSTANCE; - case "+vl": return ArrowType.ListView.INSTANCE; - case "+vL": return ArrowType.LargeListView.INSTANCE; - case "+s": return ArrowType.Struct.INSTANCE; - case "+m": return new ArrowType.Map((flags & FLAG_MAP_KEYS_SORTED) != 0); - case "+r": return ArrowType.RunEndEncoded.INSTANCE; - default: break; - } - try { - if (format.startsWith("ts") && format.length() >= 4 && format.charAt(3) == ':') { - String timezone = format.substring(4); - return new ArrowType.Timestamp( - timeUnit(format.charAt(2)), timezone.isEmpty() ? null : timezone); - } else if (format.startsWith("d:")) { - String[] parts = format.substring(2).split(",", -1); - if (parts.length == 2 || parts.length == 3) { - return new ArrowType.Decimal(Integer.parseInt(parts[0]), Integer.parseInt(parts[1]), - parts.length == 3 ? Integer.parseInt(parts[2]) : 128); - } - } else if (format.startsWith("w:")) { - return new ArrowType.FixedSizeBinary(Integer.parseInt(format.substring(2))); - } else if (format.startsWith("+w:")) { - return new ArrowType.FixedSizeList(Integer.parseInt(format.substring(3))); - } else if (format.startsWith("+ud:") || format.startsWith("+us:")) { - String ids = format.substring(4); - int[] typeIds = ids.isEmpty() ? new int[0] - : Arrays.stream(ids.split(",", -1)).mapToInt(Integer::parseInt).toArray(); - return new ArrowType.Union( - format.charAt(2) == 'd' ? UnionMode.Dense : UnionMode.Sparse, typeIds); - } - } catch (IllegalArgumentException e) { - // Not a valid format, reported below. NumberFormatException is an IllegalArgumentException. - } - throw new IllegalArgumentException("Unsupported Arrow format: " + format); - } - - private static TimeUnit timeUnit(char unit) { - return switch (unit) { - case 's' -> TimeUnit.SECOND; - case 'm' -> TimeUnit.MILLISECOND; - case 'u' -> TimeUnit.MICROSECOND; - case 'n' -> TimeUnit.NANOSECOND; - default -> throw new IllegalArgumentException("Unsupported time unit: " + unit); - }; - } - - // --------------------------------------------------------------------------------------------- - // Metadata - // --------------------------------------------------------------------------------------------- - - // The metadata are an int32 number of entries, followed by the length and the bytes of the key - // and of the value of each entry, with lengths as int32s in the native byte order. - private static long encodeMetadata(Exported data, Map metadata) { - if (metadata == null || metadata.isEmpty()) { - return 0; - } - List parts = new ArrayList<>(); - long size = 4; - for (Map.Entry entry : metadata.entrySet()) { - byte[] key = entry.getKey().getBytes(StandardCharsets.UTF_8); - byte[] value = entry.getValue().getBytes(StandardCharsets.UTF_8); - parts.add(key); - parts.add(value); - size += 8L + key.length + value.length; - } - long address = data.allocate(size); - Platform.putInt(null, address, metadata.size()); - long position = address + 4; - for (byte[] part : parts) { - Platform.putInt(null, position, part.length); - Platform.copyMemory(part, Platform.BYTE_ARRAY_OFFSET, null, position + 4, part.length); - position += 4L + part.length; - } - return address; - } - - private static Map decodeMetadata(long address) { - if (address == 0) { - return null; - } - int count = Platform.getInt(null, address); - if (count < 0) { - throw new IllegalArgumentException("The metadata of the schema are invalid."); - } - Map metadata = new LinkedHashMap<>(); - long position = address + 4; - for (int i = 0; i < count; i++) { - String[] entry = new String[2]; - for (int j = 0; j < 2; j++) { - int length = Platform.getInt(null, position); - if (length < 0) { - throw new IllegalArgumentException("The metadata of the schema are invalid."); - } - byte[] bytes = new byte[length]; - Platform.copyMemory(null, position + 4, bytes, Platform.BYTE_ARRAY_OFFSET, length); - entry[j] = new String(bytes, StandardCharsets.UTF_8); - position += 4L + length; - } - metadata.put(entry[0], entry[1]); - } - return metadata; - } - - // --------------------------------------------------------------------------------------------- - // Exporting arrays - // --------------------------------------------------------------------------------------------- - - /** - * Exports the vectors of a root as a struct array whose children are the vectors. The array - * refers to the buffers of the vectors, which it retains until it is released. - */ - static void exportRoot(BufferAllocator allocator, VectorSchemaRoot root, long out) - throws Throwable { - Exported data = startArray(out, root.getRowCount(), 0); - try { - // A struct array has a single buffer, the validity buffer, NULL as there are no nulls. - Platform.putLong(null, out + ARRAY_BUFFERS, data.allocate(8)); - Platform.putLong(null, out + ARRAY_N_BUFFERS, 1L); - exportChildren(allocator, data, root.getFieldVectors(), out); - } catch (Throwable e) { - releaseExportedArray(out); - throw e; - } - } - - private static void exportVector(BufferAllocator allocator, FieldVector vector, long out) - throws Throwable { - if (vector.getField().getDictionary() != null) { - throw new UnsupportedOperationException( - "Dictionary-encoded vectors cannot be exported: " + vector.getField()); - } - Exported data = startArray(out, vector.getValueCount(), vector.getNullCount()); - try { - int count = vector.getExportedCDataBufferCount(); - ArrowBuf pointers = allocator.buffer(8L * Math.max(1, count)); - data.buffers.add(pointers); - // Writes the addresses of the buffers, or NULL, and adds them to the buffers to release. - vector.exportCDataBuffers(data.buffers, pointers, 0L); - Platform.putLong(null, out + ARRAY_BUFFERS, pointers.memoryAddress()); - Platform.putLong(null, out + ARRAY_N_BUFFERS, count); - exportChildren(allocator, data, vector.getChildrenFromFields(), out); - } catch (Throwable e) { - releaseExportedArray(out); - throw e; - } - } - - /** - * Initializes an exported array without buffers or children, with its callback, so that a - * failure while exporting them releases what was exported. - */ - private static Exported startArray(long out, long length, long nullCount) { - Platform.setMemory(out, (byte) 0, ARRAY_SIZE); - Exported data = register(out, ARRAY_PRIVATE_DATA); - Platform.putLong(null, out + ARRAY_LENGTH, length); - Platform.putLong(null, out + ARRAY_NULL_COUNT, nullCount); - Platform.putLong(null, out + ARRAY_RELEASE, Callbacks.RELEASE_ARRAY); - return data; - } - - private static void exportChildren( - BufferAllocator allocator, Exported data, List children, long out) - throws Throwable { - Platform.putLong(null, out + ARRAY_CHILDREN, childStructs(data, children.size(), ARRAY_SIZE)); - for (int i = 0; i < children.size(); i++) { - exportVector(allocator, children.get(i), child(out, i)); - // Counted once exported, so that a failure releases the exported children only. - Platform.putLong(null, out + ARRAY_N_CHILDREN, i + 1L); - } - } - - /** The release callback of the arrays that Spark exports. */ - private static void releaseExportedArray(long array) { - try { - long count = childCount(array); - for (int i = 0; i < count; i++) { - // A child that was moved is marked as released here, and released where it was moved. - long child = child(array, i); - if (!isArrayReleased(child)) { - releaseExportedArray(child); - } - } - Exported data = exported.remove(Platform.getLong(null, array + ARRAY_PRIVATE_DATA)); - if (data != null) { - data.free(); - } - } catch (Throwable e) { - // An exception must not be thrown to the caller of the callback, which is native code. - LOG.warn("Failed to release an exported Arrow array.", e); - } - Platform.putLong(null, array + ARRAY_RELEASE, 0L); - } - - // --------------------------------------------------------------------------------------------- - // Importing arrays - // --------------------------------------------------------------------------------------------- - - /** - * An imported array, moved to memory that Spark owns. It is released once the imported buffers - * that refer to its memory are all released, and the importer released its own reference. - */ - private static final class Imported { - private final long array; - private final AtomicInteger references = new AtomicInteger(1); - - Imported(long array) { - this.array = array; - } - - void retain() { - references.incrementAndGet(); - } - - void release() { - if (references.decrementAndGet() == 0) { - try { - releaseArray(array); - } catch (Throwable e) { - LOG.warn("Failed to release an imported Arrow array.", e); - } finally { - free(array); - } - } - } - } - - /** - * Imports a struct array into the vectors of a root, one child for each vector, without copying - * the buffers. The array is moved, and marked as released: the root releases it once its vectors - * no longer use its buffers, such as when the next array is imported or the root is closed. - */ - static void importIntoRoot(BufferAllocator allocator, long array, VectorSchemaRoot root) - throws Throwable { - long moved = allocate(ARRAY_SIZE); - Platform.copyMemory(null, array, null, moved, ARRAY_SIZE); - Platform.putLong(null, array + ARRAY_RELEASE, 0L); - Imported imported = new Imported(moved); - try { - List vectors = root.getFieldVectors(); - long length = length(moved); - if (childCount(moved) != vectors.size()) { - throw new IllegalArgumentException("The array has " + childCount(moved) + - " columns, but the schema has " + vectors.size() + "."); - } - checkLength(length); - for (int i = 0; i < vectors.size(); i++) { - importVector(allocator, imported, child(moved, i), vectors.get(i), 1); - } - root.setRowCount((int) length); - } finally { - imported.release(); - } - } - - private static void importVector( - BufferAllocator allocator, - Imported imported, - long array, - FieldVector vector, - int depth) throws Throwable { - if (depth > MAX_DEPTH) { - throw new IllegalArgumentException("The array is nested more than " + MAX_DEPTH + - " levels deep."); - } - if (offset(array) != 0) { - throw new IllegalArgumentException("The array has an offset of " + offset(array) + - ", but only arrays with an offset of 0 are supported."); - } - ArrowType type = vector.getField().getType(); - long length = length(array); - checkLength(length); - List buffers = new ArrayList<>(); - try { - importBuffers(allocator, imported, array, type, length, buffers); - long nullCount = Platform.getLong(null, array + ARRAY_NULL_COUNT); - if (nullCount < 0) { - // Unknown: computed from the validity buffer, which is NULL when there are no nulls. - ArrowBuf validity = buffers.isEmpty() ? null : buffers.get(0); - nullCount = type instanceof ArrowType.Null ? length - : validity == null || validity.capacity() == 0 ? 0 - : BitVectorHelper.getNullCount(validity, (int) length); - } else if (nullCount > 0 && !buffers.isEmpty() && buffers.get(0) == null && - !(type instanceof ArrowType.Null)) { - throw new IllegalArgumentException("The array has " + nullCount + " nulls, but no " + - "validity buffer."); - } - vector.loadFieldBuffers(new ArrowFieldNode(length, nullCount), buffers); - } finally { - // The vector retains the buffers that it uses. - for (ArrowBuf buffer : buffers) { - if (buffer != null) { - buffer.getReferenceManager().release(); - } - } - } - List children = vector.getChildrenFromFields(); - if (childCount(array) != children.size()) { - throw new IllegalArgumentException("The array of type " + type + " has " + - childCount(array) + " children, but " + children.size() + " are expected."); - } - for (int i = 0; i < children.size(); i++) { - importVector(allocator, imported, child(array, i), children.get(i), depth + 1); - } - } - - private static void checkLength(long length) { - if (length < 0 || length > Integer.MAX_VALUE) { - throw new IllegalArgumentException("The array has an unsupported length of " + length + - "."); - } - } - - /** Wraps the buffers of an array for the vector of the given type, as Arrow Java loads them. */ - private static void importBuffers( - BufferAllocator allocator, - Imported imported, - long array, - ArrowType type, - long length, - List buffers) { - long count = bufferCount(array); - switch (type.getTypeID()) { - case Null: - expectBuffers(type, count, 0); - return; - case Struct: - case FixedSizeList: - expectBuffers(type, count, 1); - buffers.add(validity(allocator, imported, array, length)); - return; - case List: - case Map: - case LargeList: { - expectBuffers(type, count, 2); - int width = type instanceof ArrowType.LargeList ? 8 : 4; - buffers.add(validity(allocator, imported, array, length)); - buffers.add(offsets(allocator, imported, array, length, width)); - return; - } - case Utf8: - case Binary: - case LargeUtf8: - case LargeBinary: { - expectBuffers(type, count, 3); - int width = type instanceof ArrowType.LargeUtf8 || type instanceof ArrowType.LargeBinary - ? 8 : 4; - buffers.add(validity(allocator, imported, array, length)); - ArrowBuf offsets = offsets(allocator, imported, array, length, width); - buffers.add(offsets); - long size = width == 8 ? offsets.getLong(8L * length) : offsets.getInt(4L * length); - buffers.add(wrap(allocator, imported, buffer(array, 2), size, "data")); - return; - } - case Utf8View: - case BinaryView: { - // The validity, the views, the data buffers, and the sizes of the data buffers. - if (count < 3) { - expectBuffers(type, count, 3); - } - buffers.add(validity(allocator, imported, array, length)); - buffers.add(wrap(allocator, imported, buffer(array, 1), 16L * length, "views")); - long sizes = buffer(array, (int) count - 1); - if (sizes == 0 && count > 3) { - throw new IllegalArgumentException("The array of type " + type + " has no sizes of " + - "its data buffers."); - } - for (int i = 2; i < count - 1; i++) { - long size = Platform.getLong(null, sizes + 8L * (i - 2)); - buffers.add(wrap(allocator, imported, buffer(array, i), size, "data")); - } - return; - } - case Bool: - expectBuffers(type, count, 2); - buffers.add(validity(allocator, imported, array, length)); - buffers.add(wrap(allocator, imported, buffer(array, 1), (length + 7) / 8, "values")); - return; - default: { - int width = fixedWidth(type); - expectBuffers(type, count, 2); - buffers.add(validity(allocator, imported, array, length)); - buffers.add(wrap(allocator, imported, buffer(array, 1), width * length, "values")); - } - } - } - - /** The width in bytes of the values of a fixed-width type. */ - private static int fixedWidth(ArrowType type) { - switch (type.getTypeID()) { - case Int: return ((ArrowType.Int) type).getBitWidth() / 8; - case FloatingPoint: - switch (((ArrowType.FloatingPoint) type).getPrecision()) { - case HALF: return 2; - case SINGLE: return 4; - case DOUBLE: return 8; - default: break; - } - break; - case Decimal: return ((ArrowType.Decimal) type).getBitWidth() / 8; - case Date: return ((ArrowType.Date) type).getUnit() == DateUnit.DAY ? 4 : 8; - case Time: return ((ArrowType.Time) type).getBitWidth() / 8; - case Timestamp: - case Duration: - return 8; - case Interval: - switch (((ArrowType.Interval) type).getUnit()) { - case YEAR_MONTH: return 4; - case DAY_TIME: return 8; - case MONTH_DAY_NANO: return 16; - default: break; - } - break; - case FixedSizeBinary: return ((ArrowType.FixedSizeBinary) type).getByteWidth(); - default: break; - } - throw new UnsupportedOperationException("Arrays of the Arrow type " + type + " cannot be " + - "imported."); - } - - private static void expectBuffers(ArrowType type, long count, int expected) { - if (count != expected) { - throw new IllegalArgumentException("The array of type " + type + " has " + count + - " buffers, but " + expected + " are expected."); - } - } - - /** The validity buffer, or null if it is NULL, which Arrow Java loads as all valid or null. */ - private static ArrowBuf validity( - BufferAllocator allocator, Imported imported, long array, long length) { - long address = buffer(array, 0); - return address == 0 ? null : wrap(allocator, imported, address, (length + 7) / 8, "validity"); - } - - /** The offsets buffer, which may be NULL in an empty array. */ - private static ArrowBuf offsets( - BufferAllocator allocator, Imported imported, long array, long length, int width) { - long address = buffer(array, 1); - long size = width * (length + 1); - if (address == 0 && length == 0) { - ArrowBuf offsets = allocator.buffer(size); - offsets.setZero(0, size); - return offsets; - } - return wrap(allocator, imported, address, size, "offsets"); - } - - /** Wraps a buffer of an imported array, which keeps the array alive until it is released. */ - private static ArrowBuf wrap( - BufferAllocator allocator, Imported imported, long address, long size, String buffer) { - if (size < 0) { - throw new IllegalArgumentException("The " + buffer + " buffer has a negative size."); - } - if (size == 0) { - return allocator.getEmpty(); - } - if (address == 0) { - throw new IllegalArgumentException("The " + buffer + " buffer is NULL."); - } - imported.retain(); - try { - return allocator.wrapForeignAllocation(new ForeignAllocation(size, address) { - @Override - protected void release0() { - imported.release(); - } - }); - } catch (Throwable e) { - imported.release(); - throw e; - } - } - - // --------------------------------------------------------------------------------------------- - // Strings - // --------------------------------------------------------------------------------------------- - - /** Reads a NUL-terminated UTF-8 string. */ - static String readString(long address) { - int length = 0; - while (Platform.getByte(null, address + length) != 0) { - length++; - } - byte[] bytes = new byte[length]; - Platform.copyMemory(null, address, bytes, Platform.BYTE_ARRAY_OFFSET, length); - return new String(bytes, StandardCharsets.UTF_8); - } -} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/Ffm.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/Ffm.java deleted file mode 100644 index 75ce65b06d119..0000000000000 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/Ffm.java +++ /dev/null @@ -1,208 +0,0 @@ -/* - * 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.execution.datasources.v2.ffi; - -import java.lang.invoke.MethodHandle; -import java.lang.invoke.MethodHandles; -import java.lang.invoke.MethodType; -import java.lang.reflect.Array; -import java.nio.file.Path; -import java.util.Optional; - -/** - * The Foreign Function and Memory (FFM) API of Java 22 and later, for native data sources. - *

    - * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached - * through method handles, looked up when this class is initialized, the first time that a native - * data source library is loaded. The C functions are bound to method handles that take and return - * primitive values only, with pointers as {@code long} addresses, so calling them uses no - * reflection. - */ -final class Ffm { - - private Ffm() {} - - /** The C types of the parameters and the results of functions. */ - enum C { VOID, BOOL, INT32, INT64, POINTER } - - private static final Class LAYOUT_CLASS; - private static final Object LINKER; - private static final Object GLOBAL_ARENA; - private static final Object NO_OPTIONS; - private static final MethodHandle LIBRARY_LOOKUP; - private static final MethodHandle FIND; - private static final MethodHandle DOWNCALL; - private static final MethodHandle POINTER_DOWNCALL; - private static final MethodHandle UPCALL; - private static final MethodHandle FUNCTION_OF; - private static final MethodHandle FUNCTION_OF_VOID; - // (long) -> MemorySegment and (MemorySegment) -> long. - private static final MethodHandle OF_ADDRESS; - private static final MethodHandle ADDRESS_OF; - private static final Object BOOL; - private static final Object INT32; - private static final Object INT64; - private static final Object POINTER; - - static { - try { - // A full-privilege lookup, so that the restricted methods of the FFM API see this class as - // their caller, in the unnamed module that Spark enables native access for. - MethodHandles.Lookup lookup = MethodHandles.lookup(); - Class linker = Class.forName("java.lang.foreign.Linker"); - Class option = Class.forName("java.lang.foreign.Linker$Option"); - Class symbolLookup = Class.forName("java.lang.foreign.SymbolLookup"); - Class arena = Class.forName("java.lang.foreign.Arena"); - Class segment = Class.forName("java.lang.foreign.MemorySegment"); - Class layout = Class.forName("java.lang.foreign.MemoryLayout"); - Class layoutArray = Array.newInstance(layout, 0).getClass(); - Class valueLayout = Class.forName("java.lang.foreign.ValueLayout"); - Class descriptor = Class.forName("java.lang.foreign.FunctionDescriptor"); - Class options = Array.newInstance(option, 0).getClass(); - - LAYOUT_CLASS = layout; - LINKER = - lookup.findStatic(linker, "nativeLinker", MethodType.methodType(linker)).invoke(); - GLOBAL_ARENA = lookup.findStatic(arena, "global", MethodType.methodType(arena)).invoke(); - NO_OPTIONS = Array.newInstance(option, 0); - LIBRARY_LOOKUP = lookup.findStatic(symbolLookup, "libraryLookup", - MethodType.methodType(symbolLookup, Path.class, arena)); - FIND = lookup.findVirtual(symbolLookup, "find", - MethodType.methodType(Optional.class, String.class)); - DOWNCALL = lookup.findVirtual(linker, "downcallHandle", - MethodType.methodType(MethodHandle.class, segment, descriptor, options)).asFixedArity(); - POINTER_DOWNCALL = lookup.findVirtual(linker, "downcallHandle", - MethodType.methodType(MethodHandle.class, descriptor, options)).asFixedArity(); - UPCALL = lookup.findVirtual(linker, "upcallStub", - MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options)) - .asFixedArity(); - FUNCTION_OF = lookup.findStatic(descriptor, "of", - MethodType.methodType(descriptor, layout, layoutArray)).asFixedArity(); - FUNCTION_OF_VOID = lookup.findStatic(descriptor, "ofVoid", - MethodType.methodType(descriptor, layoutArray)).asFixedArity(); - OF_ADDRESS = - lookup.findStatic(segment, "ofAddress", MethodType.methodType(segment, long.class)); - ADDRESS_OF = lookup.findVirtual(segment, "address", MethodType.methodType(long.class)); - BOOL = valueLayout.getField("JAVA_BOOLEAN").get(null); - INT32 = valueLayout.getField("JAVA_INT").get(null); - INT64 = valueLayout.getField("JAVA_LONG").get(null); - POINTER = valueLayout.getField("ADDRESS").get(null); - } catch (Throwable e) { - // The message names the cause, because the JVM only keeps this message for the later uses - // of the class. - throw new IllegalStateException("Cannot initialize the FFM API: " + rootCause(e), e); - } - } - - static Throwable rootCause(Throwable e) { - Throwable cause = e; - while (cause.getCause() != null && cause.getCause() != cause) { - cause = cause.getCause(); - } - return cause; - } - - /** - * Returns the symbols of the library at the given path, which stays loaded until the JVM exits. - * - * @throws IllegalArgumentException if the file is not a library - * @throws IllegalCallerException if native access is not enabled for Spark - */ - static Object libraryLookup(Path path) throws Throwable { - return LIBRARY_LOOKUP.invoke(path, GLOBAL_ARENA); - } - - /** - * Binds a function of a library, looked up with {@link #libraryLookup}, to a method handle with - * the Java types of the given C types. Returns null if the library does not export the function. - */ - static MethodHandle bind(Object library, String name, C result, C... parameters) - throws Throwable { - Optional symbol = (Optional) FIND.invoke(library, name); - if (symbol.isEmpty()) { - return null; - } - MethodHandle handle = (MethodHandle) - DOWNCALL.invoke(LINKER, symbol.get(), descriptor(result, parameters), NO_OPTIONS); - return withAddresses(handle, 0, result, parameters); - } - - /** - * Returns a method handle that calls the C function whose address is its first argument, with - * the Java types of the given C types, such as a callback of an Arrow C struct. - */ - static MethodHandle bindPointer(C result, C... parameters) throws Throwable { - MethodHandle handle = (MethodHandle) - POINTER_DOWNCALL.invoke(LINKER, descriptor(result, parameters), NO_OPTIONS); - handle = MethodHandles.filterArguments(handle, 0, OF_ADDRESS); - return withAddresses(handle, 1, result, parameters); - } - - /** - * Returns the address of a C function with the given C types that calls the given method handle, - * whose parameters and result have the Java types of the C types. The function lives as long as - * the JVM. The method handle must not throw: an exception thrown by an upcall crashes the JVM. - */ - static long upcall(MethodHandle target, C result, C... parameters) throws Throwable { - for (int i = 0; i < parameters.length; i++) { - if (parameters[i] == C.POINTER) { - target = MethodHandles.filterArguments(target, i, ADDRESS_OF); - } - } - if (result == C.POINTER) { - target = MethodHandles.filterReturnValue(target, OF_ADDRESS); - } - Object stub = - UPCALL.invoke(LINKER, target, descriptor(result, parameters), GLOBAL_ARENA, NO_OPTIONS); - return (long) ADDRESS_OF.invoke(stub); - } - - /** Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses. */ - private static MethodHandle withAddresses( - MethodHandle handle, int firstParameter, C result, C... parameters) { - for (int i = 0; i < parameters.length; i++) { - if (parameters[i] == C.POINTER) { - handle = MethodHandles.filterArguments(handle, firstParameter + i, OF_ADDRESS); - } - } - if (result == C.POINTER) { - handle = MethodHandles.filterReturnValue(handle, ADDRESS_OF); - } - return handle; - } - - private static Object descriptor(C result, C... parameters) throws Throwable { - Object layouts = Array.newInstance(LAYOUT_CLASS, parameters.length); - for (int i = 0; i < parameters.length; i++) { - Array.set(layouts, i, layout(parameters[i])); - } - return result == C.VOID - ? FUNCTION_OF_VOID.invoke(layouts) - : FUNCTION_OF.invoke(layout(result), layouts); - } - - private static Object layout(C type) { - return switch (type) { - case BOOL -> BOOL; - case INT32 -> INT32; - case INT64 -> INT64; - case POINTER -> POINTER; - case VOID -> throw new IllegalArgumentException("void has no layout"); - }; - } -} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java index 33492ff80ed34..e7c9129459fa1 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -20,22 +20,27 @@ import java.lang.invoke.MethodHandle; import java.lang.invoke.MethodHandles; import java.lang.invoke.MethodType; +import java.lang.reflect.Array; import java.nio.charset.StandardCharsets; import java.nio.file.Path; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicLong; -import org.apache.spark.sql.execution.datasources.v2.ffi.Ffm.C; import org.apache.spark.unsafe.Platform; /** * A native data source library, which implements the C interface in * {@code org/apache/spark/sql/datasource/spark_datasource.h}, called through the Foreign Function - * and Memory (FFM) API of Java 22 and later, see {@link Ffm}. Each function of the library is - * bound once, when the library is loaded. The memory that Spark passes to the library is allocated - * with {@link Platform}. + * and Memory (FFM) API of Java 22 and later. + *

    + * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached + * through method handles, looked up once when the first library is loaded. Each function of a + * library is bound to a method handle that takes and returns primitive values only, with + * pointers as {@code long} addresses, so the calls themselves use no reflection. The memory that + * Spark passes to the library is allocated with {@link Platform}. *

    * Each library is looked up separately, so all the libraries can export the same functions. * Calling a function that the library does not export throws {@link MissingFunctionException}. A @@ -55,6 +60,9 @@ static boolean isSupported() { return Runtime.version().feature() >= MINIMUM_JAVA_VERSION; } + // The C types of the parameters and the results of the functions. + private enum C { VOID, BOOL, INT32, INT64, POINTER } + private final MethodHandle lastError; private final MethodHandle abiVersion; private final MethodHandle create; @@ -86,44 +94,44 @@ static boolean isSupported() { private final MethodHandle commitDataWriter; private final MethodHandle abortDataWriter; - private NativeLibrary(Object library) throws Throwable { + private NativeLibrary(Object lookup) throws Throwable { C i32 = C.INT32; C i64 = C.INT64; C ptr = C.POINTER; - lastError = Ffm.bind(library, "spark_ds_last_error", ptr); - abiVersion = Ffm.bind(library, "spark_ds_abi_version", i32); - create = Ffm.bind(library, "spark_ds_create", i32, ptr, ptr, ptr, i64, ptr); - schema = Ffm.bind(library, "spark_ds_schema", i32, ptr, ptr); - createReader = Ffm.bind(library, "spark_ds_create_reader", i32, ptr, ptr, ptr); - createStreamReader = Ffm.bind(library, "spark_ds_create_stream_reader", i32, ptr, ptr, ptr); - createWriter = Ffm.bind(library, "spark_ds_create_writer", i32, ptr, ptr, C.BOOL, ptr); + lastError = bind(lookup, "spark_ds_last_error", ptr); + abiVersion = bind(lookup, "spark_ds_abi_version", i32); + create = bind(lookup, "spark_ds_create", i32, ptr, ptr, ptr, i64, ptr); + schema = bind(lookup, "spark_ds_schema", i32, ptr, ptr); + createReader = bind(lookup, "spark_ds_create_reader", i32, ptr, ptr, ptr); + createStreamReader = bind(lookup, "spark_ds_create_stream_reader", i32, ptr, ptr, ptr); + createWriter = bind(lookup, "spark_ds_create_writer", i32, ptr, ptr, C.BOOL, ptr); createStreamWriter = - Ffm.bind(library, "spark_ds_create_stream_writer", i32, ptr, ptr, C.BOOL, ptr); - close = Ffm.bind(library, "spark_ds_close", C.VOID, ptr); - pushPredicates = Ffm.bind(library, "spark_ds_push_predicates", i32, ptr, ptr, i64, ptr); - pushLimit = Ffm.bind(library, "spark_ds_push_limit", i32, ptr, i32, ptr); - pruneColumns = Ffm.bind(library, "spark_ds_prune_columns", i32, ptr, ptr, i64, ptr); - partitions = Ffm.bind(library, "spark_ds_partitions", i32, ptr, ptr, ptr); - serializeReader = Ffm.bind(library, "spark_ds_serialize_reader", i32, ptr, ptr, ptr); - closeReader = Ffm.bind(library, "spark_ds_close_reader", C.VOID, ptr); - initialOffset = Ffm.bind(library, "spark_ds_initial_offset", i32, ptr, ptr, ptr); - latestOffset = Ffm.bind(library, "spark_ds_latest_offset", i32, ptr, ptr, ptr); + bind(lookup, "spark_ds_create_stream_writer", i32, ptr, ptr, C.BOOL, ptr); + close = bind(lookup, "spark_ds_close", C.VOID, ptr); + pushPredicates = bind(lookup, "spark_ds_push_predicates", i32, ptr, ptr, i64, ptr); + pushLimit = bind(lookup, "spark_ds_push_limit", i32, ptr, i32, ptr); + pruneColumns = bind(lookup, "spark_ds_prune_columns", i32, ptr, ptr, i64, ptr); + partitions = bind(lookup, "spark_ds_partitions", i32, ptr, ptr, ptr); + serializeReader = bind(lookup, "spark_ds_serialize_reader", i32, ptr, ptr, ptr); + closeReader = bind(lookup, "spark_ds_close_reader", C.VOID, ptr); + initialOffset = bind(lookup, "spark_ds_initial_offset", i32, ptr, ptr, ptr); + latestOffset = bind(lookup, "spark_ds_latest_offset", i32, ptr, ptr, ptr); streamPartitions = - Ffm.bind(library, "spark_ds_stream_partitions", i32, ptr, ptr, ptr, ptr, ptr); + bind(lookup, "spark_ds_stream_partitions", i32, ptr, ptr, ptr, ptr, ptr); serializeStreamReader = - Ffm.bind(library, "spark_ds_serialize_stream_reader", i32, ptr, ptr, ptr); - commitOffset = Ffm.bind(library, "spark_ds_commit_offset", i32, ptr, ptr); - closeStreamReader = Ffm.bind(library, "spark_ds_close_stream_reader", C.VOID, ptr); - read = Ffm.bind(library, "spark_ds_read", i32, ptr, i64, ptr, i64, ptr); - serializeWriter = Ffm.bind(library, "spark_ds_serialize_writer", i32, ptr, ptr, ptr); - commit = Ffm.bind(library, "spark_ds_commit", i32, ptr, i64, ptr, ptr, i64); - abort = Ffm.bind(library, "spark_ds_abort", i32, ptr, i64, ptr, ptr, i64); - closeWriter = Ffm.bind(library, "spark_ds_close_writer", C.VOID, ptr); + bind(lookup, "spark_ds_serialize_stream_reader", i32, ptr, ptr, ptr); + commitOffset = bind(lookup, "spark_ds_commit_offset", i32, ptr, ptr); + closeStreamReader = bind(lookup, "spark_ds_close_stream_reader", C.VOID, ptr); + read = bind(lookup, "spark_ds_read", i32, ptr, i64, ptr, i64, ptr); + serializeWriter = bind(lookup, "spark_ds_serialize_writer", i32, ptr, ptr, ptr); + commit = bind(lookup, "spark_ds_commit", i32, ptr, i64, ptr, ptr, i64); + abort = bind(lookup, "spark_ds_abort", i32, ptr, i64, ptr, ptr, i64); + closeWriter = bind(lookup, "spark_ds_close_writer", C.VOID, ptr); createDataWriter = - Ffm.bind(library, "spark_ds_create_data_writer", i32, ptr, i64, i32, i64, i64, ptr); - write = Ffm.bind(library, "spark_ds_write", i32, ptr, ptr, ptr); - commitDataWriter = Ffm.bind(library, "spark_ds_commit_data_writer", i32, ptr, ptr, ptr); - abortDataWriter = Ffm.bind(library, "spark_ds_abort_data_writer", i32, ptr); + bind(lookup, "spark_ds_create_data_writer", i32, ptr, i64, i32, i64, i64, ptr); + write = bind(lookup, "spark_ds_write", i32, ptr, ptr, ptr); + commitDataWriter = bind(lookup, "spark_ds_commit_data_writer", i32, ptr, ptr, ptr); + abortDataWriter = bind(lookup, "spark_ds_abort_data_writer", i32, ptr); } /** @@ -136,9 +144,9 @@ static NativeLibrary load(String path) throws Throwable { throw new UnsupportedOperationException( "Native data sources require Java " + MINIMUM_JAVA_VERSION + " or later."); } - Object library; + Object lookup; try { - library = Ffm.libraryLookup(Path.of(path)); + lookup = Ffm.LIBRARY_LOOKUP.invoke(Path.of(path), Ffm.GLOBAL_ARENA); } catch (IllegalArgumentException | IllegalCallerException e) { // The library is not a valid library, or native access is not enabled for Spark. throw new LoadException(e.getMessage(), e); @@ -146,9 +154,145 @@ static NativeLibrary load(String path) throws Throwable { // The FFM API could not be initialized, see Ffm: report why, such as an // IllegalCallerException when native access is not enabled for Spark. throw new LoadException( - "Cannot use the Foreign Function and Memory API: " + Ffm.rootCause(e), e); + "Cannot use the Foreign Function and Memory API: " + rootCause(e), e); } - return new NativeLibrary(library); + return new NativeLibrary(lookup); + } + + /** + * Binds a function of the library to a method handle whose parameters and result have the + * Java types of the given C types, with {@code long} for pointers. Returns null if the library + * does not export the function. + */ + private static MethodHandle bind(Object lookup, String name, C result, C... parameters) + throws Throwable { + Optional symbol = (Optional) Ffm.FIND.invoke(lookup, name); + if (symbol.isEmpty()) { + return null; + } + Object layouts = Ffm.layouts(parameters); + Object descriptor = result == C.VOID + ? Ffm.FUNCTION_OF_VOID.invoke(layouts) + : Ffm.FUNCTION_OF.invoke(Ffm.layout(result), layouts); + MethodHandle handle = (MethodHandle) Ffm.DOWNCALL.invoke( + Ffm.LINKER, symbol.get(), descriptor, Ffm.NO_OPTIONS); + // Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses. + for (int i = 0; i < parameters.length; i++) { + if (parameters[i] == C.POINTER) { + handle = MethodHandles.filterArguments(handle, i, Ffm.OF_ADDRESS); + } + } + if (result == C.POINTER) { + handle = MethodHandles.filterReturnValue(handle, Ffm.ADDRESS_OF); + } + return handle; + } + + /** The method handles of the FFM API, initialized when the first library is loaded. */ + private static final class Ffm { + static final Class LAYOUT_CLASS; + static final Object LINKER; + static final Object GLOBAL_ARENA; + static final Object NO_OPTIONS; + static final MethodHandle LIBRARY_LOOKUP; + static final MethodHandle FIND; + static final MethodHandle DOWNCALL; + static final MethodHandle FUNCTION_OF; + static final MethodHandle FUNCTION_OF_VOID; + // (long) -> MemorySegment and (MemorySegment) -> long. + static final MethodHandle OF_ADDRESS; + static final MethodHandle ADDRESS_OF; + static final Object BOOL; + static final Object INT32; + static final Object INT64; + static final Object POINTER; + // The address of the C function SparkDsSink, which calls NativeLibrary.sink. + static final long SINK; + + static { + try { + // A full-privilege lookup, so that the restricted methods of the FFM API see this class + // as their caller, in the unnamed module that Spark enables native access for. + MethodHandles.Lookup lookup = MethodHandles.lookup(); + Class linker = Class.forName("java.lang.foreign.Linker"); + Class option = Class.forName("java.lang.foreign.Linker$Option"); + Class symbolLookup = Class.forName("java.lang.foreign.SymbolLookup"); + Class arena = Class.forName("java.lang.foreign.Arena"); + Class segment = Class.forName("java.lang.foreign.MemorySegment"); + Class layout = Class.forName("java.lang.foreign.MemoryLayout"); + Class layoutArray = Array.newInstance(layout, 0).getClass(); + Class valueLayout = Class.forName("java.lang.foreign.ValueLayout"); + Class descriptor = Class.forName("java.lang.foreign.FunctionDescriptor"); + Class options = Array.newInstance(option, 0).getClass(); + + LAYOUT_CLASS = layout; + LINKER = lookup.findStatic(linker, "nativeLinker", MethodType.methodType(linker)) + .invoke(); + GLOBAL_ARENA = lookup.findStatic(arena, "global", MethodType.methodType(arena)).invoke(); + NO_OPTIONS = Array.newInstance(option, 0); + LIBRARY_LOOKUP = lookup.findStatic(symbolLookup, "libraryLookup", + MethodType.methodType(symbolLookup, Path.class, arena)); + FIND = lookup.findVirtual(symbolLookup, "find", + MethodType.methodType(Optional.class, String.class)); + DOWNCALL = lookup.findVirtual(linker, "downcallHandle", + MethodType.methodType(MethodHandle.class, segment, descriptor, options)) + .asFixedArity(); + FUNCTION_OF = lookup.findStatic(descriptor, "of", + MethodType.methodType(descriptor, layout, layoutArray)).asFixedArity(); + FUNCTION_OF_VOID = lookup.findStatic(descriptor, "ofVoid", + MethodType.methodType(descriptor, layoutArray)).asFixedArity(); + OF_ADDRESS = lookup.findStatic(segment, "ofAddress", + MethodType.methodType(segment, long.class)); + ADDRESS_OF = lookup.findVirtual(segment, "address", MethodType.methodType(long.class)); + BOOL = valueLayout.getField("JAVA_BOOLEAN").get(null); + INT32 = valueLayout.getField("JAVA_INT").get(null); + INT64 = valueLayout.getField("JAVA_LONG").get(null); + POINTER = valueLayout.getField("ADDRESS").get(null); + + // The upcall stub of the sink, which lives as long as the JVM. + MethodHandle target = lookup.findStatic(NativeLibrary.class, "sink", + MethodType.methodType(void.class, long.class, long.class, long.class)); + target = MethodHandles.filterArguments(target, 0, ADDRESS_OF, ADDRESS_OF); + Object sinkDescriptor = + FUNCTION_OF_VOID.invoke(layouts(C.POINTER, C.POINTER, C.INT64)); + Object stub = lookup.findVirtual(linker, "upcallStub", + MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options)) + .asFixedArity() + .invoke(LINKER, target, sinkDescriptor, GLOBAL_ARENA, NO_OPTIONS); + SINK = (long) ADDRESS_OF.invoke(stub); + } catch (Throwable e) { + // The message names the cause, because the JVM only keeps this message for the later + // uses of the class. + throw new IllegalStateException("Cannot initialize the FFM API: " + rootCause(e), e); + } + } + + /** Returns a MemoryLayout[] with the layouts of the given types. */ + static Object layouts(C... types) { + Object layouts = Array.newInstance(LAYOUT_CLASS, types.length); + for (int i = 0; i < types.length; i++) { + Array.set(layouts, i, layout(types[i])); + } + return layouts; + } + + static Object layout(C type) { + return switch (type) { + case BOOL -> BOOL; + case INT32 -> INT32; + case INT64 -> INT64; + case POINTER -> POINTER; + case VOID -> throw new IllegalArgumentException("void has no layout"); + }; + } + } + + private static Throwable rootCause(Throwable e) { + Throwable cause = e; + while (cause.getCause() != null && cause.getCause() != cause) { + cause = cause.getCause(); + } + return cause; } /** Thrown when a library cannot be loaded. */ @@ -300,7 +444,7 @@ boolean pruneColumns(long reader, String[] columnNames) throws Throwable { byte[][] partitions(long reader) throws Throwable { MethodHandle f = require(partitions, "spark_ds_partitions"); try (Sink sink = new Sink()) { - check((int) f.invokeExact(reader, Sinks.SINK, sink.id), "spark_ds_partitions"); + check((int) f.invokeExact(reader, Ffm.SINK, sink.id), "spark_ds_partitions"); return sink.all(); } } @@ -334,7 +478,7 @@ byte[][] streamPartitions(long streamReader, String start, String end) throws Th MethodHandle f = require(streamPartitions, "spark_ds_stream_partitions"); try (Memory memory = new Memory(); Sink sink = new Sink()) { check((int) f.invokeExact(streamReader, memory.string(start), memory.string(end), - Sinks.SINK, sink.id), "spark_ds_stream_partitions"); + Ffm.SINK, sink.id), "spark_ds_stream_partitions"); return sink.all(); } } @@ -432,7 +576,7 @@ void abortDataWriter(long dataWriter) throws Throwable { private byte[] serialize(MethodHandle function, String name, long handle) throws Throwable { MethodHandle f = require(function, name); try (Sink sink = new Sink()) { - check((int) f.invokeExact(handle, Sinks.SINK, sink.id), name); + check((int) f.invokeExact(handle, Ffm.SINK, sink.id), name); byte[][] values = sink.all(); if (values.length > 1) { throw new IllegalStateException( @@ -501,28 +645,12 @@ private static byte[] readCString(long address) { // Sink // --------------------------------------------------------------------------------------------- - /** The C function SparkDsSink, which calls {@link #sink}. */ - private static final class Sinks { - static final long SINK; - - static { - try { - SINK = Ffm.upcall( - MethodHandles.lookup().findStatic(NativeLibrary.class, "sink", - MethodType.methodType(void.class, long.class, long.class, long.class)), - C.VOID, C.POINTER, C.POINTER, C.INT64); - } catch (Throwable e) { - throw new IllegalStateException("Cannot create the sink of native data sources.", e); - } - } - } - private static final AtomicLong nextSinkId = new AtomicLong(1); private static final ConcurrentHashMap sinks = new ConcurrentHashMap<>(); /** * Collects the byte arrays that a library returns during a call, through the C function - * {@link Sinks#SINK}. The library passes {@link #id} back to identify it. + * {@link Ffm#SINK}. The library passes {@link #id} back to identify it. */ private static final class Sink implements AutoCloseable { final long id = nextSinkId.getAndIncrement(); @@ -547,7 +675,7 @@ public void close() { } /** - * Called by the library, through the C function {@link Sinks#SINK}, with a byte array to copy. + * Called by the library, through the upcall stub {@link Ffm#SINK}, with a byte array to copy. * It must not throw: an exception thrown by an upcall crashes the JVM. */ private static void sink(long id, long data, long length) { diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h index 4cf8b464bf4fd..16ced5dd72232 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -102,9 +102,9 @@ * For an output struct, the library initializes it and Spark takes ownership. For an input * struct, the library may move it to take ownership; otherwise, Spark releases it after the * call returns, so the library must not use it afterwards. The arrays that the library exports - * must have an offset of 0, at every level: Spark reads them with Arrow Java vectors, which - * have no offsets, so it rejects an array with another offset, such as a zero-copy slice. Copy - * a sliced array before exporting it. Spark also rejects an array with fewer values than its + * must have an offset of 0, at every level: Arrow Java ignores the offsets of the arrays it + * imports, so Spark rejects an array with another offset, such as a zero-copy slice. Copy a + * sliced array before exporting it. Spark also rejects an array with fewer values than its * parent refers to, such as a column shorter than its batch. The buffers of an array must hold * all of its values, as the C data interface requires: it does not carry the sizes of the * buffers, so Spark cannot check them, and reads past the end of a buffer that is too short. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index 7d54926216439..943da94be272e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -23,6 +23,7 @@ import java.util.concurrent.atomic.AtomicLong import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, CDataDictionaryProvider, Data} import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} import org.apache.arrow.vector.complex.MapVector import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} @@ -375,16 +376,16 @@ class NativePartitionReader( private val allocator = ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) - // The ArrowArrayStream that the library exports. - private val stream = ArrowCData.allocate(ArrowCData.STREAM_SIZE) + // Read here rather than with `Data.importArrayStream`, to check each array before importing it. + private val stream = ArrowArrayStream.allocateNew(allocator) + private val dictionaries = new CDataDictionaryProvider() private var root: VectorSchemaRoot = _ private var batch: ColumnarBatch = _ try { openStream() - val arrowSchema = importSchema() - checkSchema(arrowSchema) - root = VectorSchemaRoot.create(arrowSchema, allocator) + root = VectorSchemaRoot.create(importSchema(), allocator) + checkSchema() batch = new ColumnarBatch(root.getFieldVectors.asScala.map(columnVector).toArray) } catch { case e: Throwable => @@ -394,21 +395,22 @@ class NativePartitionReader( private def openStream(): Unit = { val library = NativeLibraries.get(location) - call(name, "read")(optional(library.read(readerState, partition, stream))).getOrElse { - throw new SparkUnsupportedOperationException( - "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) - } + call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) + .getOrElse { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) + } } private def importSchema(): Schema = { - val arrowSchema = ArrowCData.allocate(ArrowCData.SCHEMA_SIZE) + val arrowSchema = ArrowSchema.allocateNew(allocator) try { - call(name, "read")(ArrowCData.streamSchema(stream, arrowSchema)) - val imported = call(name, "read")(ArrowCData.importSchema(arrowSchema)) + call(name, "read")(stream.getSchema(arrowSchema)) + val imported = call(name, "read")(Data.importSchema(allocator, arrowSchema, dictionaries)) new Schema( imported.getFields.asScala.map(withMapEntryNames).asJava, imported.getCustomMetadata) } finally { - ArrowCData.free(arrowSchema) + arrowSchema.close() } } @@ -430,12 +432,11 @@ class NativePartitionReader( new Field(field.getName, field.getFieldType, renamed.asJava) } - /** Checks the schema of the data, before Spark creates the vectors to import it into. */ - private def checkSchema(arrowSchema: Schema): Unit = { + private def checkSchema(): Unit = { def isDictionaryEncoded(field: Field): Boolean = { field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) } - arrowSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + "is not supported.") } @@ -447,13 +448,13 @@ class NativePartitionReader( case listView: ArrowType.ListView => Some(listView) case _ => field.getChildren.asScala.view.flatMap(unsupportedType).headOption } - arrowSchema.getFields.asScala.foreach { field => + root.getSchema.getFields.asScala.foreach { field => unsupportedType(field).foreach { arrowType => throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + s"$arrowType, which is not supported.") } } - val actual = ArrowUtils.fromArrowSchema(arrowSchema) + val actual = ArrowUtils.fromArrowSchema(root.getSchema) // The schema to read as the library got it: Arrow does not carry everything, such as the // collations of strings, the fields of intervals and user-defined types, which still apply to // the data that Spark reads. @@ -482,80 +483,88 @@ class NativePartitionReader( } override def next(): Boolean = { - val array = ArrowCData.allocate(ArrowCData.ARRAY_SIZE) + val array = ArrowArray.allocateNew(allocator) try { - call(name, "read")(ArrowCData.streamNext(stream, array)) - if (ArrowCData.isArrayReleased(array)) { + call(name, "read")(stream.getNext(array)) + val snapshot = array.snapshot() + if (snapshot.release == 0) { // The end of the stream. false } else { - checkArrays(array) - // Moves the array, which the root then releases. - call(name, "read")(ArrowCData.importIntoRoot(allocator, array, root)) + try { + checkArrays(snapshot) + } catch { + case e: Throwable => + array.release() + throw e + } + call(name, "read")(Data.importIntoVectorSchemaRoot(allocator, array, root, dictionaries)) batch.setNumRows(root.getRowCount) true } } finally { - // Releases the array unless it was imported. - Utils.tryWithSafeFinally(ArrowCData.releaseArray(array))(ArrowCData.free(array)) + array.close() } } /** * Checks the array of a batch, and the arrays of its columns at every level, before importing - * them into Arrow Java vectors, which would read them wrongly otherwise: - * - An array must have an offset of 0, since the vectors have no offsets. The import rejects - * other offsets too, but this check names the column. + * them. Arrow Java trusts the arrays, and would read them wrongly otherwise: + * - An array must have an offset of 0. Arrow Java ignores the offsets of the arrays it imports, + * so it would read a sliced array from the start of its buffers. * - An array must have at least as many values as its parent refers to: a column as many as * the batch has rows, a field as many as its struct has values, and the elements of a list or - * a map as many as its last offset says. The vectors would extend a shorter array with nulls. + * a map as many as its last offset says. Arrow Java would extend a shorter array with nulls. * Only the lengths that the library exported show this: the import sets the length of each - * vector to what its parent refers to. - * The buffers are not checked, since the C data interface does not carry their sizes: the - * import takes them from the lengths, so the library must export buffers that hold every value. + * array to what its parent refers to. + * The buffers are not checked, since the C data interface does not carry their sizes: Arrow + * Java takes them from the lengths, so the library must export buffers that hold every value. */ - private def checkArrays(batch: Long): Unit = { - import ArrowCData.{child, childCount, length, offset} - def children(array: Long): Seq[Long] = (0 until childCount(array).toInt).map(child(array, _)) - // The last offset of a list or a map, in its second buffer. The import rejects the array if + private def checkArrays(batch: ArrowArray.Snapshot): Unit = { + def children(array: ArrowArray.Snapshot): Seq[ArrowArray.Snapshot] = { + (0 until array.n_children.toInt).map { i => + ArrowArray.wrap(Platform.getLong(null, array.children + 8L * i)).snapshot() + } + } + // The last offset of a list or a map, in its second buffer. Arrow Java rejects the array if // the buffer is missing. - def lastOffset(array: Long): Long = { - val offsets = if (length(array) > 0 && ArrowCData.bufferCount(array) >= 2) { - ArrowCData.buffer(array, 1) + def lastOffset(array: ArrowArray.Snapshot): Long = { + val offsets = if (array.length > 0 && array.n_buffers >= 2) { + Platform.getLong(null, array.buffers + 8L) } else { 0L } - if (offsets == 0) 0 else Platform.getInt(null, offsets + 4L * length(array)) + if (offsets == 0) 0 else Platform.getInt(null, offsets + 4L * array.length) } // The number of values that an array refers to in each of its children. - def referenced(array: Long, arrowType: ArrowType): Long = arrowType match { - case _: ArrowType.Struct => length(array) + def referenced(array: ArrowArray.Snapshot, arrowType: ArrowType): Long = arrowType match { + case _: ArrowType.Struct => array.length case _: ArrowType.List | _: ArrowType.Map => lastOffset(array) // The other types that Spark reads have no children. List views are rejected before. case _ => 0 } - def checkColumn(array: Long, field: Field, column: String): Unit = { - if (offset(array) != 0) { - throw error(name, "read", s"The column $column has an offset of ${offset(array)}, but " + + def checkColumn(array: ArrowArray.Snapshot, field: Field, column: String): Unit = { + if (array.offset != 0) { + throw error(name, "read", s"The column $column has an offset of ${array.offset}, but " + "Spark only reads arrays with an offset of 0. Copy sliced arrays before exporting them.") } - val referencedLength = referenced(array, field.getType) + val length = referenced(array, field.getType) children(array).zip(field.getChildren.asScala).foreach { case (child, childField) => - if (length(child) < referencedLength) { + if (child.length < length) { throw error(name, "read", s"The column $column contains an array with a length of " + - s"${length(child)}, but its parent refers to $referencedLength values.") + s"${child.length}, but its parent refers to $length values.") } checkColumn(child, childField, column) } } - if (offset(batch) != 0) { - throw error(name, "read", s"The batch has an offset of ${offset(batch)}, but Spark only " + + if (batch.offset != 0) { + throw error(name, "read", s"The batch has an offset of ${batch.offset}, but Spark only " + "reads arrays with an offset of 0. Copy sliced arrays before exporting them.") } children(batch).zip(root.getSchema.getFields.asScala).foreach { case (column, field) => - if (length(column) < length(batch)) { + if (column.length < batch.length) { throw error(name, "read", s"The column ${field.getName} has a length of " + - s"${length(column)}, but the batch has a length of ${length(batch)}.") + s"${column.length}, but the batch has a length of ${batch.length}.") } checkColumn(column, field, field.getName) } @@ -566,11 +575,12 @@ class NativePartitionReader( override def close(): Unit = { Utils.tryWithSafeFinally { if (root != null) { - // Releases the last array imported. root.close() } + dictionaries.close() // Releases the stream of the library, unless it was not exported. - Utils.tryWithSafeFinally(ArrowCData.releaseStream(stream))(ArrowCData.free(stream)) + stream.release() + stream.close() }(Utils.tryLogNonFatalError(allocator.close())) } } @@ -606,21 +616,19 @@ class NativeDataWriter( } val root = new VectorSchemaRoot( vectors.map(_.getField).asJava, vectors.asJava, batch.numRows()) - val array = ArrowCData.allocate(ArrowCData.ARRAY_SIZE) - val arrowSchema = ArrowCData.allocate(ArrowCData.SCHEMA_SIZE) + val array = ArrowArray.allocateNew(ArrowUtils.rootAllocator) + val arrowSchema = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) try { - ArrowCData.exportSchema(root.getSchema, arrowSchema) - ArrowCData.exportRoot(ArrowUtils.rootAllocator, root, array) - call(name, "write to")(library.write(handle, array, arrowSchema)) + Data.exportVectorSchemaRoot(ArrowUtils.rootAllocator, root, null, array, arrowSchema) + call(name, "write to") { + library.write(handle, array.memoryAddress(), arrowSchema.memoryAddress()) + } } finally { // Releases the batch unless the library moved it to take ownership. - Utils.tryWithSafeFinally { - ArrowCData.releaseArray(array) - ArrowCData.releaseSchema(arrowSchema) - } { - ArrowCData.free(array) - ArrowCData.free(arrowSchema) - } + array.release() + array.close() + arrowSchema.release() + arrowSchema.close() } } @@ -720,7 +728,15 @@ object NativeCalls { } } - private[ffi] def message(e: Throwable): String = Option(e.getMessage).getOrElse(e.toString) + // How Arrow reports the error message of an Arrow C stream. + private val StreamError = """(?s)\[errno \d+\] CDataJniException\{errno=\d+, message=(.*)\}""".r + + private[ffi] def message(e: Throwable): String = { + Option(e.getMessage).getOrElse(e.toString) match { + case StreamError(message) => message + case message => message + } + } /** Calls a function that the library may not implement. Returns None if it does not. */ def optional[T](f: => T): Option[T] = { @@ -748,26 +764,31 @@ object NativeArrow { val arrowSchema = ArrowUtils.toArrowSchema( schema, SQLConf.get.sessionLocalTimeZone, errorOnDuplicatedFieldNames = true, largeVarTypes = false) - val struct = ArrowCData.allocate(ArrowCData.SCHEMA_SIZE) + val struct = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) try { - ArrowCData.exportSchema(arrowSchema, struct) - f(struct) + Data.exportSchema(ArrowUtils.rootAllocator, arrowSchema, null, struct) + f(struct.memoryAddress()) } finally { // Releases the schema unless the library moved it to take ownership. - Utils.tryWithSafeFinally(ArrowCData.releaseSchema(struct))(ArrowCData.free(struct)) + struct.release() + struct.close() } } /** Calls `f` with the address of an `ArrowSchema` to initialize, and imports it. */ def importSchema(f: Long => Unit): StructType = { - val struct = ArrowCData.allocate(ArrowCData.SCHEMA_SIZE) + val struct = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) + var imported = false try { - f(struct) - // Releases the schema. - ArrowUtils.fromArrowSchema(ArrowCData.importSchema(struct)) + f(struct.memoryAddress()) + imported = true + // Releases and closes the struct. + ArrowUtils.fromArrowSchema(Data.importSchema(ArrowUtils.rootAllocator, struct, null)) } finally { - // Releases the schema if the library initialized it but it was not imported. - Utils.tryWithSafeFinally(ArrowCData.releaseSchema(struct))(ArrowCData.free(struct)) + if (!imported) { + struct.release() + struct.close() + } } } } diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index 782d9af407850..53468b84a8173 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -962,8 +962,6 @@ int spark_ds_read(const uint8_t* reader_state, int64_t reader_state_length, return guarded([&] { Options state = deserialize(toBytes(reader_state, reader_state_length)); maybeFail(state, "open"); - // Returns success without initializing the stream, which Spark must reject. - if (get(state, "fail", "") == "uninitialized") return; auto bounds = split(toBytes(partition, partition_length)); if (bounds.size() != 2) throw std::runtime_error("invalid partition"); stream->get_schema = streamGetSchema; diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCDataSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCDataSuite.scala deleted file mode 100644 index bf04116a6ad1a..0000000000000 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/ArrowCDataSuite.scala +++ /dev/null @@ -1,463 +0,0 @@ -/* - * 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.execution.datasources.v2.ffi - -import java.nio.charset.StandardCharsets.UTF_8 - -import scala.jdk.CollectionConverters._ -import scala.util.Random - -import org.apache.arrow.memory.BufferAllocator -import org.apache.arrow.vector.{FieldVector, LargeVarBinaryVector, LargeVarCharVector, VectorSchemaRoot, ViewVarBinaryVector, ViewVarCharVector} -import org.apache.arrow.vector.compare.VectorEqualsVisitor -import org.apache.arrow.vector.types.{DateUnit, FloatingPointPrecision, IntervalUnit, TimeUnit, UnionMode} -import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} - -import org.apache.spark.SparkFunSuite -import org.apache.spark.sql.{RandomDataGenerator, Row} -import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow} -import org.apache.spark.sql.execution.arrow.ArrowWriter -import org.apache.spark.sql.types._ -import org.apache.spark.sql.util.ArrowUtils -import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} -import org.apache.spark.unsafe.Platform - -/** - * Tests the Arrow C data interface of [[ArrowCData]], by exporting and importing in the JVM. The - * tests are skipped before Java 22, where the FFM API is not available. - */ -class ArrowCDataSuite extends SparkFunSuite { - - private def ffmTest(name: String)(f: => Unit): Unit = test(name) { - assume(NativeLibrary.isSupported, "Java 22 or later is required.") - f - } - - /** Runs `f` with an allocator, and checks that all its memory is released afterwards. */ - private def withAllocator[T](f: BufferAllocator => T): T = { - val allocator = ArrowUtils.rootAllocator.newChildAllocator("test", 0, Long.MaxValue) - try { - val result = f(allocator) - assert(allocator.getAllocatedMemory == 0, "Memory was not released.") - result - } finally { - allocator.close() - } - } - - private def withStruct[T](size: Long)(f: Long => T): T = { - val address = ArrowCData.allocate(size) - try f(address) finally ArrowCData.free(address) - } - - /** Exports and imports a schema. */ - private def roundTrip(schema: Schema): Schema = withStruct(ArrowCData.SCHEMA_SIZE) { struct => - ArrowCData.exportSchema(schema, struct) - ArrowCData.importSchema(struct) - } - - /** - * Exports the root, imports it into a new root with the imported schema, and passes the new - * root to `f`. The exported root is closed first, so the imported buffers must stay valid. - */ - private def roundTrip[T](allocator: BufferAllocator, root: VectorSchemaRoot)( - f: VectorSchemaRoot => T): T = { - val imported = VectorSchemaRoot.create(roundTrip(root.getSchema), allocator) - try { - withStruct(ArrowCData.ARRAY_SIZE) { array => - ArrowCData.exportRoot(allocator, root, array) - ArrowCData.importIntoRoot(allocator, array, imported) - assert(ArrowCData.isArrayReleased(array), "The array was not moved.") - } - f(imported) - } finally { - imported.close() - } - } - - private def assertSameVectors(expected: VectorSchemaRoot, actual: VectorSchemaRoot): Unit = { - assert(actual.getRowCount == expected.getRowCount) - expected.getFieldVectors.asScala.zip(actual.getFieldVectors.asScala).foreach { - case (e, a) => assert(VectorEqualsVisitor.vectorEquals(e, a), s"Different ${e.getName}.") - } - } - - private def rows(root: VectorSchemaRoot, schema: StructType): Seq[Row] = { - val vectors = root.getFieldVectors.asScala.map(v => new ArrowColumnVector(v): ColumnVector) - val batch = new ColumnarBatch(vectors.toArray, root.getRowCount) - val toScala = CatalystTypeConverters.createToScalaConverter(schema) - // Converted without copying the rows: ColumnarArray.copy fails on primitive arrays with nulls. - batch.rowIterator().asScala.map(row => toScala(row).asInstanceOf[Row]).toList - } - - private val schema = StructType(Seq( - StructField("boolean", BooleanType), - StructField("byte", ByteType), - StructField("short", ShortType), - StructField("int", IntegerType, nullable = false), - StructField("long", LongType), - StructField("float", FloatType), - StructField("double", DoubleType), - StructField("decimal", DecimalType(10, 2)), - StructField("wide_decimal", DecimalType(38, 18)), - StructField("string", StringType), - StructField("binary", BinaryType), - StructField("date", DateType), - StructField("timestamp", TimestampType), - StructField("timestamp_ntz", TimestampNTZType), - StructField("year_month", YearMonthIntervalType()), - StructField("day_time", DayTimeIntervalType()), - StructField("null", NullType), - StructField("ints", ArrayType(IntegerType)), - StructField("strings", ArrayType(StringType, containsNull = false)), - StructField("map", MapType(StringType, IntegerType)), - StructField("nested_map", MapType(IntegerType, ArrayType(LongType))), - StructField("struct", StructType(Seq( - StructField("a", IntegerType), StructField("b", StringType)))), - StructField("structs", ArrayType(StructType(Seq( - StructField("x", DoubleType), StructField("y", ArrayType(StringType)))))))) - - private def write(allocator: BufferAllocator, schema: StructType, rows: Seq[Row]) - : VectorSchemaRoot = { - val root = VectorSchemaRoot.create( - ArrowUtils.toArrowSchema(schema, "UTC", errorOnDuplicatedFieldNames = true, - largeVarTypes = false), - allocator) - val writer = ArrowWriter.create(root) - val toCatalyst = CatalystTypeConverters.createToCatalystConverter(schema) - rows.foreach(row => writer.write(toCatalyst(row).asInstanceOf[InternalRow])) - writer.finish() - root - } - - private def randomRows(schema: StructType, count: Int): Seq[Row] = { - val random = new Random(42) - Seq.fill(count)(RandomDataGenerator.randomRow(random, schema)) - } - - ffmTest("round trip of schemas") { - // scalastyle:off nonascii - val metadata = Map("key" -> "value", "unicode \u00e9" -> "\u4e2d").asJava - // scalastyle:on nonascii - val arrowSchema = new Schema(Seq( - new Field("a", new FieldType(true, new ArrowType.Int(32, true), null, metadata), null), - new Field("m", FieldType.notNullable(new ArrowType.Map(true)), Seq( - new Field("entries", FieldType.notNullable(ArrowType.Struct.INSTANCE), Seq( - new Field("key", FieldType.notNullable(ArrowType.Utf8.INSTANCE), null), - new Field("value", FieldType.nullable(ArrowType.Utf8.INSTANCE), null)).asJava)).asJava), - new Field("ts", FieldType.nullable( - new ArrowType.Timestamp(TimeUnit.MICROSECOND, "America/Los_Angeles")), null)).asJava, - Map("schema" -> "metadata").asJava) - assert(roundTrip(arrowSchema) == arrowSchema) - val sparkSchema = ArrowUtils.toArrowSchema( - schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = false) - assert(roundTrip(sparkSchema) == sparkSchema) - } - - ffmTest("round trip of arrays of the Spark types") { - withAllocator { allocator => - val expected = randomRows(schema, 100) - val root = write(allocator, schema, expected) - try { - roundTrip(allocator, root) { imported => - assertSameVectors(root, imported) - assert(rows(imported, schema) == expected) - } - } finally { - root.close() - } - } - } - - ffmTest("the imported buffers outlive the exported vectors") { - withAllocator { allocator => - val schema = StructType(Seq(StructField("s", StringType), StructField("l", LongType))) - val expected = randomRows(schema, 50) - val root = write(allocator, schema, expected) - val imported = VectorSchemaRoot.create(roundTrip(root.getSchema), allocator) - try { - withStruct(ArrowCData.ARRAY_SIZE) { array => - ArrowCData.exportRoot(allocator, root, array) - // The array retains the buffers of the vectors. - root.close() - ArrowCData.importIntoRoot(allocator, array, imported) - } - assert(rows(imported, schema) == expected) - } finally { - imported.close() - } - } - } - - ffmTest("round trip of empty arrays") { - withAllocator { allocator => - val root = write(allocator, schema, Nil) - try { - roundTrip(allocator, root) { imported => - assert(imported.getRowCount == 0) - assertSameVectors(root, imported) - } - } finally { - root.close() - } - } - } - - ffmTest("round trip of views and large variable-width types") { - withAllocator { allocator => - // scalastyle:off nonascii - val values = Seq("short", null, "a string longer than twelve bytes", "", "\u00e9\u4e2d") - // scalastyle:on nonascii - val utf8View = new ViewVarCharVector("utf8_view", allocator) - val binaryView = new ViewVarBinaryVector("binary_view", allocator) - val largeUtf8 = new LargeVarCharVector("large_utf8", allocator) - val largeBinary = new LargeVarBinaryVector("large_binary", allocator) - val vectors = Seq[FieldVector](utf8View, binaryView, largeUtf8, largeBinary) - vectors.foreach(_.allocateNew()) - values.zipWithIndex.foreach { - case (null, i) => - utf8View.setNull(i) - binaryView.setNull(i) - largeUtf8.setNull(i) - largeBinary.setNull(i) - case (value, i) => - val bytes = value.getBytes(UTF_8) - utf8View.setSafe(i, bytes) - binaryView.setSafe(i, bytes) - largeUtf8.setSafe(i, bytes) - largeBinary.setSafe(i, bytes) - } - vectors.foreach(_.setValueCount(values.length)) - val root = new VectorSchemaRoot(vectors.asJava) - try { - roundTrip(allocator, root) { imported => - assertSameVectors(root, imported) - val read = rows(imported, StructType(Seq( - StructField("utf8_view", StringType), StructField("binary_view", BinaryType), - StructField("large_utf8", StringType), StructField("large_binary", BinaryType)))) - assert(read.map(_.getString(0)) == values) - assert(read.map(_.getString(2)) == values) - } - } finally { - root.close() - } - } - } - - ffmTest("release a child that the consumer moved") { - withAllocator { allocator => - val schema = StructType(Seq(StructField("a", LongType), StructField("b", StringType))) - val root = write(allocator, schema, randomRows(schema, 10)) - try { - withStruct(ArrowCData.ARRAY_SIZE) { array => - ArrowCData.exportRoot(allocator, root, array) - withStruct(ArrowCData.ARRAY_SIZE) { moved => - // Moves the first child, as the C data interface allows: copies it, and marks the - // original as released, by setting its release callback, at offset 64, to NULL. - val child = ArrowCData.child(array, 0) - Platform.copyMemory(null, child, null, moved, ArrowCData.ARRAY_SIZE) - Platform.putLong(null, child + 64, 0L) - ArrowCData.releaseArray(array) - assert(ArrowCData.isArrayReleased(array)) - assert(!ArrowCData.isArrayReleased(moved)) - root.close() - // The moved child still retains its buffers. - assert(allocator.getAllocatedMemory > 0) - ArrowCData.releaseArray(moved) - assert(ArrowCData.isArrayReleased(moved)) - } - } - } finally { - root.close() - } - } - } - - ffmTest("formats") { - val types = Seq( - ArrowType.Null.INSTANCE, - ArrowType.Bool.INSTANCE, - new ArrowType.Int(8, true), new ArrowType.Int(8, false), - new ArrowType.Int(16, true), new ArrowType.Int(16, false), - new ArrowType.Int(32, true), new ArrowType.Int(32, false), - new ArrowType.Int(64, true), new ArrowType.Int(64, false), - new ArrowType.FloatingPoint(FloatingPointPrecision.HALF), - new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), - new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), - ArrowType.Binary.INSTANCE, ArrowType.LargeBinary.INSTANCE, ArrowType.BinaryView.INSTANCE, - ArrowType.Utf8.INSTANCE, ArrowType.LargeUtf8.INSTANCE, ArrowType.Utf8View.INSTANCE, - new ArrowType.FixedSizeBinary(5), - new ArrowType.Decimal(10, 2, 128), new ArrowType.Decimal(40, 5, 256), - new ArrowType.Date(DateUnit.DAY), new ArrowType.Date(DateUnit.MILLISECOND), - new ArrowType.Time(TimeUnit.SECOND, 32), new ArrowType.Time(TimeUnit.MILLISECOND, 32), - new ArrowType.Time(TimeUnit.MICROSECOND, 64), new ArrowType.Time(TimeUnit.NANOSECOND, 64), - new ArrowType.Interval(IntervalUnit.YEAR_MONTH), - new ArrowType.Interval(IntervalUnit.DAY_TIME), - new ArrowType.Interval(IntervalUnit.MONTH_DAY_NANO), - ArrowType.List.INSTANCE, ArrowType.LargeList.INSTANCE, - ArrowType.ListView.INSTANCE, ArrowType.LargeListView.INSTANCE, - new ArrowType.FixedSizeList(3), - ArrowType.Struct.INSTANCE, - new ArrowType.Map(false)) ++ - TimeUnit.values().toSeq.flatMap { unit => - Seq(new ArrowType.Timestamp(unit, null), new ArrowType.Timestamp(unit, "UTC"), - new ArrowType.Duration(unit)) - } - types.foreach { arrowType => - assert(ArrowCData.parseFormat(ArrowCData.format(arrowType), 0) == arrowType) - } - // The flags of a map say whether its keys are sorted. - assert(ArrowCData.parseFormat("+m", 4) == new ArrowType.Map(true)) - assert(ArrowCData.parseFormat("+ud:1,2", 0) == - new ArrowType.Union(UnionMode.Dense, Array(1, 2))) - assert(ArrowCData.parseFormat("tsu:+01:00", 0) == - new ArrowType.Timestamp(TimeUnit.MICROSECOND, "+01:00")) - Seq("", "x", "d:1", "d:a,b", "w:x", "+w:", "tsx:", "tsu").foreach { format => - val e = intercept[IllegalArgumentException](ArrowCData.parseFormat(format, 0)) - assert(e.getMessage == s"Unsupported Arrow format: $format") - } - } - - /** - * Builds ArrowSchema structs by hand, without release callbacks, so that importing them does not - * release them: the memory is freed afterwards with `free`. The fields are at the offsets of the - * struct: format 0, name 8, metadata 16, flags 24, n_children 32, children 40, dictionary 48. - */ - private class RawSchemas { - private val allocations = scala.collection.mutable.ArrayBuffer.empty[Long] - - private def allocate(size: Long): Long = { - val address = ArrowCData.allocate(size) - allocations += address - address - } - - private def string(value: String): Long = { - val bytes = value.getBytes(UTF_8) - val address = allocate(bytes.length + 1) - Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address, bytes.length) - address - } - - def schema( - format: String, - name: String, - children: Seq[Long] = Nil, - dictionary: Long = 0, - flags: Long = 2): Long = { - val address = allocate(ArrowCData.SCHEMA_SIZE) - Platform.putLong(null, address, string(format)) - Platform.putLong(null, address + 8, string(name)) - Platform.putLong(null, address + 24, flags) - Platform.putLong(null, address + 32, children.length) - if (children.nonEmpty) { - val pointers = allocate(8L * children.length) - children.zipWithIndex.foreach { case (child, i) => - Platform.putLong(null, pointers + 8L * i, child) - } - Platform.putLong(null, address + 40, pointers) - } - Platform.putLong(null, address + 48, dictionary) - address - } - - def free(): Unit = allocations.foreach(ArrowCData.free) - } - - ffmTest("import dictionary-encoded fields") { - val raw = new RawSchemas - try { - val dictionary = raw.schema("u", "") - val struct = raw.schema("+s", "", Seq(raw.schema("i", "x", dictionary = dictionary, - flags = 3)), flags = 0) - val field = ArrowCData.importSchema(struct).getFields.get(0) - assert(field.getName == "x") - assert(field.getDictionary != null) - assert(field.getDictionary.isOrdered) - assert(field.getDictionary.getIndexType == new ArrowType.Int(32, true)) - } finally { - raw.free() - } - } - - ffmTest("reject invalid schemas") { - val raw = new RawSchemas - try { - def checkInvalid(schema: Long, message: String): Unit = { - val e = intercept[IllegalArgumentException](ArrowCData.importSchema(schema)) - assert(e.getMessage == message) - } - checkInvalid(raw.schema("l", "x"), - "The schema is not a struct type, but Int(64, true).") - checkInvalid(raw.schema("+s", "", Seq(raw.schema("?", "x"))), - "Unsupported Arrow format: ?") - val dictionary = raw.schema("u", "") - checkInvalid(raw.schema("+s", "", Seq(raw.schema("u", "x", dictionary = dictionary))), - "The indices of the dictionary of x are not integers: u.") - var nested = raw.schema("i", "x") - (0 to 64).foreach(_ => nested = raw.schema("+s", "s", Seq(nested))) - checkInvalid(nested, "The schema is nested more than 64 levels deep.") - } finally { - raw.free() - } - } - - ffmTest("reject invalid arrays") { - withAllocator { allocator => - val schema = StructType(Seq(StructField("a", IntegerType), StructField("s", StringType))) - // The mutations keep what the release callback of the array frees, so nothing leaks. - def checkInvalid(message: String, importInto: Option[Schema] = None)( - mutate: Long => Unit): Unit = { - val root = write(allocator, schema, Seq(Row(1, "a"), Row(2, "b"))) - val imported = VectorSchemaRoot.create(importInto.getOrElse(root.getSchema), allocator) - try { - withStruct(ArrowCData.ARRAY_SIZE) { array => - ArrowCData.exportRoot(allocator, root, array) - mutate(array) - val e = intercept[IllegalArgumentException] { - ArrowCData.importIntoRoot(allocator, array, imported) - } - assert(e.getMessage == message) - assert(ArrowCData.isArrayReleased(array), "The array was not moved.") - } - } finally { - imported.close() - root.close() - } - } - val oneColumn = new Schema(Seq(new Field("a", FieldType.nullable( - new ArrowType.Int(32, true)), null)).asJava) - checkInvalid("The array has 2 columns, but the schema has 1.", Some(oneColumn))(_ => ()) - checkInvalid("The array has an offset of 1, but only arrays with an offset of 0 are " + - "supported.") { array => - Platform.putLong(null, ArrowCData.child(array, 0) + 16, 1L) - } - checkInvalid("The array has 1 nulls, but no validity buffer.") { array => - val column = ArrowCData.child(array, 0) - Platform.putLong(null, column + 8, 1L) - Platform.putLong(null, Platform.getLong(null, column + 40), 0L) - } - checkInvalid("The array of type Utf8 has 2 buffers, but 3 are expected.") { array => - Platform.putLong(null, ArrowCData.child(array, 1) + 24, 2L) - } - checkInvalid("The data buffer is NULL.") { array => - val column = ArrowCData.child(array, 1) - Platform.putLong(null, Platform.getLong(null, column + 40) + 16, 0L) - } - } - } -} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 0cf4c734aa85f..70842f4add1f3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -186,7 +186,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { """"children":[{"type":"column","name":["id"]},""" + """{"type":"literal","dataType":"bigint","value":1},""" + """{"type":"literal","dataType":"bigint","value":2}]}]}""")) - // Only the functions listed in spark_ds_push_predicates are passed, so not arithmetic. + // Only the functions listed in NativeBridge.pushPredicates are passed, so not arithmetic. NativePredicates.FUNCTION_NAMES.foreach { name => assert(toJson(predicate(name, id, id)).nonEmpty, name) } @@ -584,9 +584,6 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { check("read", "injected failure in get_next") { session => range(session, "fail" -> "read").collect() } - check("read", "The stream is released, or was not initialized.") { session => - range(session, "fail" -> "uninitialized").collect() - } // An error of the library is not an unsupported operation. check("plan a scan of", "native_sink does not support batch reads", "native_sink") { session => session.read.schema("id BIGINT").format("native_sink").load().collect() From dd3232fc48df4ba4ecec1549f33cdc6b944c5ac6 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Tue, 13 Oct 2026 16:31:29 +0900 Subject: [PATCH 18/27] Reject Arrow streams that the library did not initialize Arrow calls the callbacks of an Arrow C stream without checking them, so a library whose spark_ds_read returns success without initializing the stream would crash the JVM. Spark now reports it as an error instead. Co-authored-by: Isaac --- .../execution/datasources/v2/ffi/NativeDataSource.scala | 8 ++++++++ .../resources/native-datasource/test_native_datasource.cc | 2 ++ .../datasources/v2/ffi/NativeDataSourceSuite.scala | 6 +++++- 3 files changed, 15 insertions(+), 1 deletion(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index 943da94be272e..a8420f2817a89 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -403,6 +403,14 @@ class NativePartitionReader( } private def importSchema(): Schema = { + // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash + // the JVM. + val snapshot = stream.snapshot() + if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 || + snapshot.get_last_error == 0) { + throw error(name, "read", "The library did not initialize the stream, or a callback of " + + "the stream is NULL.") + } val arrowSchema = ArrowSchema.allocateNew(allocator) try { call(name, "read")(stream.getSchema(arrowSchema)) diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index 53468b84a8173..782d9af407850 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -962,6 +962,8 @@ int spark_ds_read(const uint8_t* reader_state, int64_t reader_state_length, return guarded([&] { Options state = deserialize(toBytes(reader_state, reader_state_length)); maybeFail(state, "open"); + // Returns success without initializing the stream, which Spark must reject. + if (get(state, "fail", "") == "uninitialized") return; auto bounds = split(toBytes(partition, partition_length)); if (bounds.size() != 2) throw std::runtime_error("invalid partition"); stream->get_schema = streamGetSchema; diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 70842f4add1f3..f922793d7c780 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -186,7 +186,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { """"children":[{"type":"column","name":["id"]},""" + """{"type":"literal","dataType":"bigint","value":1},""" + """{"type":"literal","dataType":"bigint","value":2}]}]}""")) - // Only the functions listed in NativeBridge.pushPredicates are passed, so not arithmetic. + // Only the functions listed in spark_ds_push_predicates are passed, so not arithmetic. NativePredicates.FUNCTION_NAMES.foreach { name => assert(toJson(predicate(name, id, id)).nonEmpty, name) } @@ -584,6 +584,10 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { check("read", "injected failure in get_next") { session => range(session, "fail" -> "read").collect() } + check("read", "The library did not initialize the stream, or a callback of the stream is " + + "NULL.") { session => + range(session, "fail" -> "uninitialized").collect() + } // An error of the library is not an unsupported operation. check("plan a scan of", "native_sink does not support batch reads", "native_sink") { session => session.read.schema("id BIGINT").format("native_sink").load().collect() From 6918d74f8774d235a400797b3a4c057c8de7f7a7 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Tue, 13 Oct 2026 19:32:36 +0900 Subject: [PATCH 19/27] Keep a doc comment of NativeDataSourceConnectSuite within the scalafmt line limit The Scala linter failed because a line of the doc comment had 99 characters, more than the 98 columns that scalafmt allows in sql/connect. Co-authored-by: Isaac --- .../apache/spark/sql/connect/NativeDataSourceConnectSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala index f8aaa6b335fc7..97c4b3ae61046 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala @@ -29,7 +29,7 @@ import org.apache.spark.util.Utils /** * Tests native data sources with Spark Connect: a client adds packages to its session with * `spark.addArtifact`, and the server finds the data sources in the packages of the session or in - * the installed libraries. The tests are skipped when there is no C++ compiler, or before Java 22. + * the installed libraries. The tests are skipped without a C++ compiler, or before Java 22. */ class NativeDataSourceConnectSuite extends SparkConnectServerTest { From a28dabe1c26cb8ac2f42a787ff934aee731cc6b7 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Wed, 14 Oct 2026 08:23:50 +0900 Subject: [PATCH 20/27] Apply the latest review round of #59313 to the C interface Mirrors aa0f3ba and e3ff80c of #59313, in the terms of spark_datasource.h: - Spark evaluates all the pushed predicates again after reading. The header states the rule for skipping data, and the jar contains predicate-cases.json, generated by a test from Spark's own evaluation. The JSON of the predicates is ASCII. - spark_ds_create_stream_writer gets the ID of the streaming query. - Offsets must be compact JSON on a single line. - An installed library replaced in place fails with INVALID_NATIVE_DATA_SOURCE_LIBRARY.REPLACED, and the caches identify files with their file key too. - Concurrent distributions of a configured package are serialized, a scan planned again after a failure fails with the same error, unsupported Arrow types and the link errors of Arrow are reported as NATIVE_DATA_SOURCE_ERROR, the column vectors of a read are not freeable between batches, packages of an unsupported version are logged, and PySpark addArtifact only passes file: URIs as they are. - Docs: the writer lifecycle, close functions, Arrow variants, a Memory section, broadcasts with AQE, replaced libraries, and the addArtifact Scaladoc. NativeLibrary already rejected NULL handles and missing offsets, so their tests use its errors. Co-authored-by: Isaac --- .../resources/error/error-conditions.json | 5 + docs/sql-native-data-sources.md | 82 +++- python/pyspark/sql/session.py | 10 +- python/pyspark/sql/tests/test_artifact.py | 7 +- .../org/apache/spark/sql/SparkSession.scala | 17 +- .../datasources/v2/ffi/NativeLibrary.java | 30 +- .../spark/sql/datasource/predicate-cases.json | 44 ++ .../spark/sql/datasource/spark_datasource.h | 114 +++-- .../v2/columnar/ColumnarDataSource.scala | 8 +- .../v2/columnar/ColumnarTable.scala | 5 +- .../v2/columnar/ColumnarWrite.scala | 11 +- .../datasources/v2/ffi/NativeDataSource.scala | 84 +++- .../v2/ffi/NativeDataSourcePackage.scala | 91 ++-- .../v2/ffi/NativeDataSourceRegistry.scala | 10 +- .../datasources/v2/ffi/NativePredicates.scala | 6 +- .../test_native_datasource.cc | 136 +++++- .../v2/ffi/NativeDataSourceClusterSuite.scala | 9 +- .../v2/ffi/NativeDataSourceSuite.scala | 417 +++++++++++++++++- 18 files changed, 921 insertions(+), 165 deletions(-) create mode 100644 sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 8ec5072e34e6b..fbc26caba7ee7 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -4802,6 +4802,11 @@ "" ] }, + "REPLACED" : { + "message" : [ + "The library was replaced after Spark loaded it: Spark loaded the build with the SHA-256 checksum , but the file now has the checksum . A JVM cannot load another build of a library from the same path, so restart the Spark application to use the new build." + ] + }, "UNSUPPORTED_ABI_VERSION" : { "message" : [ "It requires version of the native data source interface, but Spark supports version ." diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 2ce06950bcb55..8554b6cb6be22 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -66,23 +66,34 @@ Spark calls the functions of the library as follows: - **Streaming reads**: on the driver, Spark creates a stream reader with `spark_ds_create_stream_reader`. For each micro-batch, it asks for the `spark_ds_latest_offset`, plans the partitions with `spark_ds_stream_partitions`, and calls `spark_ds_read` on the - executors. Offsets are JSON strings defined by the library. -- **Writes**: on the driver, Spark creates a writer with `spark_ds_create_writer` or - `spark_ds_create_stream_writer`. On the executors, each task creates a data writer with + executors. Offsets are defined by the library, as compact JSON on a single line: Spark stores + each offset as a line of the offset log of the checkpoint, and rejects an offset that is not + valid JSON or contains a line break, such as pretty-printed JSON. +- **Writes**: on the driver, right before the tasks of the write (or of a micro-batch) start, Spark + creates a writer with `spark_ds_create_writer` or `spark_ds_create_stream_writer`. With adaptive + query execution, it first runs the shuffle stages of the data to write, and creates no writer if + they fail, but the last stage runs in the tasks of the write, so a writer can still be aborted + because computing the data failed: do irreversible work, such as removing the existing data for + an overwrite, in `spark_ds_commit`. On the executors, each task creates a data writer with `spark_ds_create_data_writer`, writes the batches of its data with `spark_ds_write`, and calls `spark_ds_commit_data_writer`, or `spark_ds_abort_data_writer` if the task or the commit fails. Then the driver calls `spark_ds_commit` with the messages of the tasks, or `spark_ds_abort` if - the write failed. A micro-batch streaming query creates a new stream writer for each - micro-batch, and Spark closes the previous one when it creates the next one. Because Spark is not - told when a streaming query stops, the last stream writer of a query is closed once it is no - longer referenced, or when the query is restarted. + the write failed. A streaming write commits each micro-batch with its epoch ID, and + `spark_ds_create_stream_writer` gets the ID of the query, which stays the same when the query + restarts from its checkpoint: together, they identify the output of a micro-batch, so that a + library can skip an epoch that it already committed. A micro-batch streaming query creates a new + stream writer for each micro-batch, and Spark closes the previous one when it creates the next + one. Because Spark is not told when a streaming query stops, the last stream writer of a query is + closed once it is no longer referenced, or when the query is restarted. The library defines the format of the partitions, of the state of its readers and writers, and of the commit messages, which Spark sends to the executors and back as byte arrays. The library returns them, and the offsets, by passing them to a `SparkDsSink` callback, which copies them. A function reports an error by returning a non-zero value, with the message returned by `spark_ds_last_error`, and Spark reports it as a `NATIVE_DATA_SOURCE_ERROR`. A C++ exception or a -Rust panic must not cross the interface. +Rust panic must not cross the interface. The close functions cannot fail, and they can run after +the query, so do the work that can fail, such as flushing data, in `spark_ds_commit`. Strings are +UTF-8, and the JSON of the predicates is ASCII, with the other characters escaped. A library only exports the functions of the operations it supports. If a function is missing, Spark treats the operation as unsupported, except for the optional functions: @@ -97,13 +108,25 @@ Spark treats the operation as unsupported, except for the optional functions: The comments of [`spark_datasource.h`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h) specify each function, including the ownership of the Arrow structs and the JSON format of the -pushed predicates. A predicate that a library reports as evaluated completely must follow the -semantics of Spark SQL, for example for NaN values, which are equal to each other and greater -than any other value. The arrays that a library exports must have an offset of 0: copy a sliced -array, such as the result of `RecordBatch::slice` in arrow-rs, before exporting it. Spark also -checks that each array has at least as many values as its parent refers to, but it cannot check -that the buffers of the arrays hold those values, because the C data interface does not carry the -sizes of buffers: a library that exports a buffer that is too short makes Spark read past its end. +pushed predicates. Spark evaluates all the pushed predicates again after reading, so a library can +use them to skip data, for example files or row groups whose statistics show that no row matches, +but it must never skip a row for which the predicates are true under the semantics of Spark SQL: +NaN values are equal to each other and greater than any other value, `-0.0` is equal to `0.0`, +comparisons with null are null, and strings compare by their UTF-8 bytes. The comparisons of IEEE +754 and the total orders of Arrow libraries, such as the comparison kernels of arrow-rs, differ for +some of these values. The jar of Spark SQL contains test cases for a library, +`org/apache/spark/sql/datasource/predicate-cases.json`: predicates as Spark pushes them, with the +rows of a table for which they are true: + +```bash +unzip -p "$SPARK_HOME"/jars/spark-sql_*.jar org/apache/spark/sql/datasource/predicate-cases.json +``` + +The arrays that a library exports must have an offset of 0: copy a sliced array, such as the +result of `RecordBatch::slice` in arrow-rs, before exporting it. Spark also checks that each array +has at least as many values as its parent refers to, but it cannot check that the buffers of the +arrays hold those values, because the C data interface does not carry the sizes of buffers: a +library that exports a buffer that is too short makes Spark read past its end. The `get_next` function of the stream that `spark_ds_read` exports may block, for example on network I/O, but Spark cannot interrupt native code: when a task is killed, Spark stops reading @@ -119,7 +142,7 @@ which `spark_ds_abi_version` returns and the manifest of a package declares. Wit signature and the semantics of an existing function never change, and new functions are only added as optional functions, so a library keeps working with later releases of Spark. New values may appear where the header says so: for example, later releases may push down predicates with -other names, which a library must not report as evaluated completely. An incompatible change comes +other names, which a library must ignore if it does not know them. An incompatible change comes with a new version, with new names. Spark supports version 1 at least in all the releases of Spark 4 from 4.4.0. @@ -128,13 +151,29 @@ with a new version, with new names. Spark supports version 1 at least in all the Version 1 of the interface does not cover these features of Data Source V2 yet. They can be added later as optional functions: -- statistics: Spark estimates the size of a native scan with `spark.sql.defaultSizeInBytes`, so a - native scan is never the broadcast side of a join, unless the query has a broadcast hint; +- statistics: Spark estimates the size of a native scan with `spark.sql.defaultSizeInBytes`, so it + does not plan a native scan as the broadcast side of a join, unless the query has a broadcast + hint. Adaptive query execution can still broadcast it at run time, from the size of a shuffle of + the scan; - custom metrics, shown in the Spark UI; - admission control for streaming reads, such as a maximum number of rows for each micro-batch, and continuous streaming; - the pushdown of aggregates, of the sort order and of the partitioning of the data; -- a function to cancel a read in progress. +- a function to cancel a read in progress; +- memory accounting: see below. + +### Memory + +Spark does not manage the memory of native data sources. The memory that a library allocates, and +the Arrow buffers that it exports, which Spark reads without copying them, are not accounted by the +memory manager of Spark, nor limited by `spark.memory.offHeap.size`. Neither are the Arrow batches +that Spark writes, which end at `spark.sql.execution.arrow.maxRecordsPerBatch` rows or once they +hold `spark.sql.execution.arrow.maxBytesPerBatch` bytes, while the size of the batches that a +library reads is up to the library. This memory comes out of the memory overhead of the driver and +the executors, `spark.driver.memoryOverhead` and `spark.executor.memoryOverhead`, which are 10% of +their memory, and at least 384 MiB, by default. Increase it for libraries that use much memory, and +bound the memory that a library uses, for example with the size of its batches: on YARN or +Kubernetes, an executor that uses more memory than it requested is killed. ## Finding Native Data Sources @@ -183,6 +222,11 @@ with the SHA-256 checksum of the file, because the partitions and the state that serializes are only meant for that build: a query fails with `NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND` on a node where a different build is installed. +A JVM loads a library from a path only once. So after a new build is installed at the same path, +the queries of an application that loaded the previous build fail with +`INVALID_NATIVE_DATA_SOURCE_LIBRARY.REPLACED`, on the driver or on the executors that loaded it, +rather than running mixed builds: restart the application to use the new build. + ### Packages A native data source package distributes a library with a session, for one or more platforms. It diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index ed1133bdcb04a..916d4a1640ea6 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -20,7 +20,6 @@ import warnings from collections.abc import Sized from functools import cached_property, reduce -from pathlib import PureWindowsPath from threading import RLock from types import TracebackType from typing import ( @@ -2299,12 +2298,13 @@ def addArtifacts( }, ) # Native data source packages are artifacts of the session in the JVM, whatever the - # flags, as with Spark Connect. The JVM resolves URIs, such as file:///tmp/a.sparkpkg, - # but resolves a relative path against its own working directory, so it gets an absolute - # path instead. A Windows path, such as C:/a.sparkpkg, is a path too. + # flags, as with Spark Connect. The JVM resolves file URIs, such as + # file:///tmp/a.sparkpkg, but resolves a relative path against its own working directory, + # so it gets an absolute path instead. Anything else is a path, such as C:/a.sparkpkg on + # Windows, or a file name with a colon, such as a:1.sparkpkg. for p in path: if p.endswith(".sparkpkg"): - if urlparse(p).scheme and not PureWindowsPath(p).drive: + if urlparse(p).scheme == "file": self._jsparkSession.addArtifact(p) else: self._jsparkSession.addArtifact(os.path.abspath(p)) diff --git a/python/pyspark/sql/tests/test_artifact.py b/python/pyspark/sql/tests/test_artifact.py index 5e65cf7f9866f..73f18abc9676e 100644 --- a/python/pyspark/sql/tests/test_artifact.py +++ b/python/pyspark/sql/tests/test_artifact.py @@ -122,14 +122,16 @@ def new_file(name): # together with other files. spark.addArtifact(new_file("my_file_source.sparkpkg"), file=True) spark.addArtifacts(new_file("my_other_source.sparkpkg"), new_file("other.txt")) - # A URI, and a path relative to the working directory of Python, which is not the - # one of the JVM. + # A URI, and paths relative to the working directory of Python, which is not the + # one of the JVM, also with a colon, which is not the scheme of a URI. spark.addArtifact(Path(new_file("my_uri_source.sparkpkg")).as_uri()) new_file("my_relative_source.sparkpkg") + new_file("my_source:1.sparkpkg") cwd = os.getcwd() os.chdir(d) try: spark.addArtifact("my_relative_source.sparkpkg") + spark.addArtifact("my_source:1.sparkpkg") finally: os.chdir(cwd) packages = spark._jsparkSession.artifactManager().getNativeDataSourcePackages()._1() @@ -140,6 +142,7 @@ def new_file(name): "my_other_source.sparkpkg", "my_relative_source.sparkpkg", "my_source.sparkpkg", + "my_source:1.sparkpkg", "my_uri_source.sparkpkg", ], ) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala index 8aa554a879a15..cce193fe50db8 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -568,6 +568,9 @@ abstract class SparkSession extends Serializable with Closeable { * Ivy settings. Otherwise, Ivy URIs are resolved by the server, which does not search Maven or * Ivy repositories local to the client. * + * Since 4.4.0, it also supports native data source packages: local files with the extension + * .sparkpkg, which make the native data sources that they contain available to the session. + * * @since 4.0.0 */ @Experimental @@ -583,6 +586,9 @@ abstract class SparkSession extends Serializable with Closeable { * Ivy settings. Otherwise, Ivy URIs are resolved by the server, which does not search Maven or * Ivy repositories local to the client. * + * Since 4.4.0, it also supports native data source packages: local files with the extension + * .sparkpkg, which make the native data sources that they contain available to the session. + * * @since 4.0.0 */ @Experimental @@ -593,7 +599,9 @@ abstract class SparkSession extends Serializable with Closeable { * specified by `target` under the session's working directory of that particular file * extension. * - * Supported target file extensions are .jar and .class. + * Supported target file extensions are .jar and .class, and since 4.4.0, .sparkpkg for native + * data source packages, which make the native data sources that they contain available to the + * session. * * ==Example== * {{{ @@ -613,7 +621,9 @@ abstract class SparkSession extends Serializable with Closeable { * Add a single artifact to the session while preserving the directory structure specified by * `target` under the session's working directory of that particular file extension. * - * Supported target file extensions are .jar and .class. + * Supported target file extensions are .jar and .class, and since 4.4.0, .sparkpkg for native + * data source packages, which make the native data sources that they contain available to the + * session. * * ==Example== * {{{ @@ -639,6 +649,9 @@ abstract class SparkSession extends Serializable with Closeable { * Ivy settings. Otherwise, Ivy URIs are resolved by the server, which does not search Maven or * Ivy repositories local to the client. * + * Since 4.4.0, it also supports native data source packages: local files with the extension + * .sparkpkg, which make the native data sources that they contain available to the session. + * * @since 4.0.0 */ @Experimental diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java index e7c9129459fa1..8e0f8754bac4d 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -106,7 +106,7 @@ private NativeLibrary(Object lookup) throws Throwable { createStreamReader = bind(lookup, "spark_ds_create_stream_reader", i32, ptr, ptr, ptr); createWriter = bind(lookup, "spark_ds_create_writer", i32, ptr, ptr, C.BOOL, ptr); createStreamWriter = - bind(lookup, "spark_ds_create_stream_writer", i32, ptr, ptr, C.BOOL, ptr); + bind(lookup, "spark_ds_create_stream_writer", i32, ptr, ptr, C.BOOL, ptr, ptr); close = bind(lookup, "spark_ds_close", C.VOID, ptr); pushPredicates = bind(lookup, "spark_ds_push_predicates", i32, ptr, ptr, i64, ptr); pushLimit = bind(lookup, "spark_ds_push_limit", i32, ptr, i32, ptr); @@ -380,27 +380,23 @@ private long createWithSchema( } long createWriter(long dataSource, long schemaAddress, boolean overwrite) throws Throwable { - return createAnyWriter( - createWriter, "spark_ds_create_writer", dataSource, schemaAddress, overwrite); + MethodHandle f = require(createWriter, "spark_ds_create_writer"); + try (Memory memory = new Memory()) { + long out = memory.zeroed(8); + check((int) f.invokeExact(dataSource, schemaAddress, overwrite, out), + "spark_ds_create_writer"); + return readHandle(out, "spark_ds_create_writer"); + } } - long createStreamWriter(long dataSource, long schemaAddress, boolean overwrite) + long createStreamWriter(long dataSource, long schemaAddress, boolean overwrite, String queryId) throws Throwable { - return createAnyWriter( - createStreamWriter, "spark_ds_create_stream_writer", dataSource, schemaAddress, overwrite); - } - - private long createAnyWriter( - MethodHandle function, - String name, - long dataSource, - long schemaAddress, - boolean overwrite) throws Throwable { - MethodHandle f = require(function, name); + MethodHandle f = require(createStreamWriter, "spark_ds_create_stream_writer"); try (Memory memory = new Memory()) { long out = memory.zeroed(8); - check((int) f.invokeExact(dataSource, schemaAddress, overwrite, out), name); - return readHandle(out, name); + check((int) f.invokeExact(dataSource, schemaAddress, overwrite, memory.string(queryId), out), + "spark_ds_create_stream_writer"); + return readHandle(out, "spark_ds_create_stream_writer"); } } diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json b/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json new file mode 100644 index 0000000000000..b6d74c3b6303a --- /dev/null +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json @@ -0,0 +1,44 @@ +{ + "description": "Predicates as Spark passes them to spark_ds_push_predicates, with the ids of the rows of the table for which they are true. A case lists the predicates that Spark pushes for the filter in its SQL, which are combined with AND. A row has a value for each column of the schema; NaN and infinite doubles are strings.", + "schema": "id INT,i INT,d DOUBLE,s STRING", + "rows": [ + [0, 1, 0.0, "a"], + [1, 2, -0.0, ""], + [2, null, "NaN", "\u00E9"], + [3, 3, "Infinity", "\uD83D\uDE00"], + [4, 1, "-Infinity", null], + [5, null, 1.5, "ab"], + [6, 2, null, "b\u00E9"], + [7, 4, 2.5, "\uFFFD"] + ], + "cases": [ + {"sql": "d = 0.0D", "predicates": [{"type":"function","name":"=","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":0.0}]}], "rows": [0, 1]}, + {"sql": "d = -0.0D", "predicates": [{"type":"function","name":"=","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":-0.0}]}], "rows": [0, 1]}, + {"sql": "d < 0.0D", "predicates": [{"type":"function","name":"<","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":0.0}]}], "rows": [4]}, + {"sql": "d > 1.0D", "predicates": [{"type":"function","name":">","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":1.0}]}], "rows": [2, 3, 5, 7]}, + {"sql": "d >= 2.5D", "predicates": [{"type":"function","name":">=","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":2.5}]}], "rows": [2, 3, 7]}, + {"sql": "d <> 0.0D", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":"=","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":0.0}]}]}], "rows": [2, 3, 4, 5, 7]}, + {"sql": "d <=> 0.0D", "predicates": [{"type":"function","name":"<=>","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":0.0}]}], "rows": [0, 1]}, + {"sql": "d IN (0.0D, 1.5D)", "predicates": [{"type":"function","name":"IN","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":0.0},{"type":"literal","dataType":"double","value":1.5}]}], "rows": [0, 1, 5]}, + {"sql": "NOT (d > 1.0D)", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":">","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":1.0}]}]}], "rows": [0, 1, 4]}, + {"sql": "NOT (d < 1.0D)", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":"<","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":1.0}]}]}], "rows": [2, 3, 5, 7]}, + {"sql": "d IS NULL", "predicates": [{"type":"function","name":"IS_NULL","children":[{"type":"column","name":["d"]}]}], "rows": [6]}, + {"sql": "d IS NOT NULL", "predicates": [{"type":"function","name":"IS_NOT_NULL","children":[{"type":"column","name":["d"]}]}], "rows": [0, 1, 2, 3, 4, 5, 7]}, + {"sql": "i = 1", "predicates": [{"type":"function","name":"=","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":1}]}], "rows": [0, 4]}, + {"sql": "i = NULL", "predicates": [{"type":"function","name":"=","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":null}]}], "rows": []}, + {"sql": "i <=> NULL", "predicates": [{"type":"function","name":"<=>","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":null}]}], "rows": [2, 5]}, + {"sql": "NOT (i <=> NULL)", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":"<=>","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":null}]}]}], "rows": [0, 1, 3, 4, 6, 7]}, + {"sql": "i IN (1, NULL)", "predicates": [{"type":"function","name":"IN","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":1},{"type":"literal","dataType":"int","value":null}]}], "rows": [0, 4]}, + {"sql": "NOT (i IN (1, NULL))", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":"IN","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":1},{"type":"literal","dataType":"int","value":null}]}]}], "rows": []}, + {"sql": "NOT (i IN (1, 3))", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":"IN","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":1},{"type":"literal","dataType":"int","value":3}]}]}], "rows": [1, 6, 7]}, + {"sql": "i IS NULL OR d > 1.0D", "predicates": [{"type":"function","name":"OR","children":[{"type":"function","name":"IS_NULL","children":[{"type":"column","name":["i"]}]},{"type":"function","name":">","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":1.0}]}]}], "rows": [2, 3, 5, 7]}, + {"sql": "i > 1 AND d IS NOT NULL", "predicates": [{"type":"function","name":">","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":1}]}, {"type":"function","name":"IS_NOT_NULL","children":[{"type":"column","name":["d"]}]}], "rows": [1, 3, 7]}, + {"sql": "s = '\uD83D\uDE00'", "predicates": [{"type":"function","name":"=","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"\uD83D\uDE00"}]}], "rows": [3]}, + {"sql": "s > '\uFFFD'", "predicates": [{"type":"function","name":">","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"\uFFFD"}]}], "rows": [3]}, + {"sql": "s < 'b'", "predicates": [{"type":"function","name":"<","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"b"}]}], "rows": [0, 1, 5]}, + {"sql": "startswith(s, 'a')", "predicates": [{"type":"function","name":"STARTS_WITH","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"a"}]}], "rows": [0, 5]}, + {"sql": "endswith(s, '\u00E9')", "predicates": [{"type":"function","name":"ENDS_WITH","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"\u00E9"}]}], "rows": [2, 6]}, + {"sql": "contains(s, '\uD83D\uDE00')", "predicates": [{"type":"function","name":"CONTAINS","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"\uD83D\uDE00"}]}], "rows": [3]}, + {"sql": "s IN ('', '\u00E9')", "predicates": [{"type":"function","name":"IN","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":""},{"type":"literal","dataType":"string","value":"\u00E9"}]}], "rows": [1, 2]} + ] +} diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h index 16ced5dd72232..db08aa28e851e 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -46,20 +46,28 @@ * operations down by calling, each at most once and in this order, spark_ds_push_predicates, * spark_ds_push_limit and spark_ds_prune_columns, plans the partitions with * spark_ds_partitions, serializes the state of the reader with spark_ds_serialize_reader, and - * closes it. Then it calls spark_ds_read on the executors, in one task for each partition. + * closes it. Then it calls spark_ds_read on the executors, in one task for each partition. If + * Spark never plans the partitions, for example to only explain the plan of the query, it + * closes the reader once the reader is no longer referenced. * - Streaming reads. On the driver, Spark creates a stream reader with * spark_ds_create_stream_reader and keeps it until the streaming query stops. For each * micro-batch, it asks for the spark_ds_latest_offset, plans the partitions with * spark_ds_stream_partitions and serializes the state with spark_ds_serialize_stream_reader, - * then calls spark_ds_read on the executors. Offsets are JSON strings defined by the library. - * When a streaming query restarts, from its checkpoint, Spark creates a new stream reader and - * may call spark_ds_stream_partitions again for the last range that it planned before, so the - * partitions of a range must read the same data each time. - * - Writes. On the driver, once it computed the data to write, right before it starts the tasks - * of the write (or of a micro-batch), Spark creates a writer with spark_ds_create_writer or - * spark_ds_create_stream_writer, and serializes its state with spark_ds_serialize_writer. If - * computing the data fails, Spark creates no writer. On the executors, each task creates a data - * writer with spark_ds_create_data_writer, calls spark_ds_write for each batch of at most + * then calls spark_ds_read on the executors. Offsets are defined by the library, as compact + * JSON on a single line: Spark stores each offset as a line of the offset log of the + * checkpoint, so it rejects an offset that is not valid JSON or contains a line break. When a + * streaming query restarts, from its checkpoint, Spark creates a new stream reader and may call + * spark_ds_stream_partitions again for the last range that it planned before, so the partitions + * of a range must read the same data each time. + * - Writes. On the driver, right before it starts the tasks of the write (or of a micro-batch), + * Spark creates a writer with spark_ds_create_writer or spark_ds_create_stream_writer, and + * serializes its state with spark_ds_serialize_writer. With adaptive query execution, Spark + * first runs the stages of the data to write that end with a shuffle, and creates no writer if + * they fail. The rest of the data, such as its last stage, is computed in the tasks of the + * write, so a writer may still be aborted because computing the data failed: do irreversible + * work, such as removing the existing data for an overwrite, in spark_ds_commit. On the + * executors, each task creates a data writer with spark_ds_create_data_writer, calls + * spark_ds_write for each batch of at most * spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends once it holds * spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls spark_ds_commit_data_writer, or * spark_ds_abort_data_writer if the task or the commit fails. Then the driver calls @@ -73,28 +81,36 @@ * nor aborted, and may even have committed before its message reached the driver. A streaming * write commits or aborts each micro-batch (epoch). Spark may commit the same epoch more than * once, for example when a streaming query restarts after a failure, so committing an epoch - * that is already committed must have no effect, for exactly-once output. A micro-batch - * streaming query creates a new stream writer for each micro-batch, and Spark closes the - * previous stream writer of the query when the query creates a new one. Spark is not told when - * a streaming query stops, so it closes the last stream writer of the query once the writer is - * no longer referenced, or when the query is restarted. + * that is already committed must have no effect, for exactly-once output. The ID of the query, + * which spark_ds_create_stream_writer gets, stays the same when the query restarts from its + * checkpoint, so the ID of the query and the epoch ID identify the output of a micro-batch, + * also when other queries write to the same destination. A micro-batch streaming query creates + * a new stream writer for each micro-batch, and Spark closes the previous stream writer of the + * query when the query creates a new one. Spark is not told when a streaming query stops, so it + * closes the last stream writer of the query once the writer is no longer referenced, or when + * the query is restarted. * * Conventions * ----------- * - Errors. The functions that return an int return 0 on success, or any other value on failure. * They then set the message that spark_ds_last_error returns, and Spark reports it as a - * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail. A C++ exception or a Rust panic - * must not cross the interface: catch it and return an error instead. + * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail, and they may run after the query, + * from another thread: do the work that can fail a query, such as flushing data, in + * spark_ds_commit. A C++ exception or a Rust panic must not cross the interface: catch it and + * return an error instead. * - Handles. The create functions return an opaque, non-NULL handle in their last argument, * typically a pointer: Spark reports a NULL handle as an error. Spark releases each handle * exactly once with the matching close function, or with spark_ds_commit_data_writer or - * spark_ds_abort_data_writer for data writers, though not always right after the last use of a - * stream writer (see Writes above). A reader or writer handle must not depend on the data + * spark_ds_abort_data_writer for data writers, though not always right after its last use: a + * batch reader whose partitions are never planned and the last stream writer of a query are + * released once they are no longer referenced (see Lifecycle above). A reader or writer handle + * must not depend on the data * source handle it was created from, because Spark closes the data source handle right after * creating it. * - Memory. What Spark passes to a function, such as strings and byte arrays, is only valid - * during the call. Strings are UTF-8 and NUL-terminated. The library returns strings and byte - * arrays by passing them to a SparkDsSink, which copies them. + * during the call. Strings are UTF-8 and NUL-terminated. The JSON of the predicates is ASCII: + * it escapes the other characters, which any JSON parser decodes. The library returns strings + * and byte arrays by passing them to a SparkDsSink, which copies them. * - State. Partitions, the state of readers and writers, and commit messages are byte arrays in a * format defined by the library. Spark sends them to the executors and back. * - Arrow data. Schemas and data are exchanged with the Arrow C data interface @@ -110,9 +126,10 @@ * buffers, so Spark cannot check them, and reads past the end of a buffer that is too short. * - Schemas. Spark types map to Arrow types as in spark.sql.execution.arrow: for example, bigint * is int64, string is utf8, timestamp is timestamp[us] with a time zone and timestamp_ntz is - * timestamp[us] without one. Data read from a library may also use large_utf8, utf8_view and - * the other variants of the same type, and any names for the fields of the entries of a map, - * but not unsigned integers, decimal256, list views or dictionary encoding. + * timestamp[us] without one. Data read from a library may also use large_utf8 and utf8_view for + * strings, large_binary and binary_view for binary data, and any names for the fields of the + * entries of a map. Spark rejects the other Arrow types, such as unsigned integers, decimal256, + * large lists, fixed-size lists and list views, and dictionary encoding. * - Options. Option keys are passed in lower case, because Spark matches them * case-insensitively. * - Optional functions. A library only needs to export the functions for the operations it @@ -273,10 +290,13 @@ int spark_ds_create_writer(void* source, struct ArrowSchema* schema, bool overwr * schema: the schema of the data to write * overwrite: whether to replace the existing data in each micro-batch (output mode "complete") * instead of appending to it + * query_id: the ID of the streaming query. It stays the same when the query restarts from its + * checkpoint, so with the epoch IDs, it identifies the commits of the query: see Writes + * above. * writer: set to the writer handle */ int spark_ds_create_stream_writer(void* source, struct ArrowSchema* schema, bool overwrite, - void** writer); + const char* query_id, void** writer); /* Releases a data source handle. */ void spark_ds_close(void* source); @@ -307,25 +327,39 @@ void spark_ds_close(void* source); * arithmetic or other functions, nor predicates on columns that contain strings with a * non-binary collation, such as UTF8_LCASE: strings are compared by their UTF-8 bytes. A later * version of Spark may pass predicates with other names without changing SPARK_DS_ABI_VERSION, - * so a reader must not report a predicate as evaluated completely if it contains a name that the - * reader does not know, or a data type that it does not support. - * - * A predicate that the reader evaluates completely must follow the semantics of Spark SQL. In - * particular, for floating-point columns, NaN is equal to NaN and greater than any other value, - * including positive infinity, and -0.0 is equal to 0.0, unlike IEEE 754 comparisons. For - * example, d > 1.0 is true for a NaN value of d. Comparisons with null are null, so a row only - * matches if the predicate is true. + * so a reader must ignore a predicate that contains a name that it does not know, or a data type + * that it does not support. + * + * Spark evaluates all the predicates again after reading, so a reader may use them only to skip + * data, for example the files whose statistics show that no row matches. But it must never drop + * or skip a row for which the predicates are true under the semantics of Spark SQL, which Spark + * cannot detect. In particular: + * - For floating-point columns, NaN is equal to NaN and greater than any other value, including + * positive infinity, and -0.0 is equal to 0.0. IEEE 754 comparisons treat NaN differently, and + * total orders, such as f64::total_cmp in Rust or the comparison kernels of arrow-rs, also + * order -0.0 before 0.0. For example, d > 1.0 is true for a NaN value of d, and d = 0.0 for + * -0.0. + * - Comparisons with null are null, and a row only matches if the predicates are true. For + * example, NOT (i IN (1, NULL)) matches no row. + * - Strings are compared by their UTF-8 bytes, which order the characters outside the Basic + * Multilingual Plane after all the others, unlike comparisons of UTF-16 strings. + * - To use only a part of a predicate, a reader may take the parts that it does not support as + * true, but not under a NOT, where that would drop rows. + * The file org/apache/spark/sql/datasource/predicate-cases.json in the jar of Spark SQL lists + * predicates as Spark passes them to this function, with the rows of a table for which they are + * true, to test a reader with. * * predicates: the count predicates, as JSON * evaluated: an array of count values, all false, to set to whether the reader evaluates each - * predicate completely, so that Spark does not have to evaluate it again. The reader - * can still use the other predicates to skip data. + * predicate completely. Spark shows these predicates as pushed in the plan of the + * query, and evaluates all the predicates again after reading anyway. */ int spark_ds_push_predicates(void* reader, const char* const* predicates, int64_t count, bool* evaluated); /* - * Pushes a LIMIT down to the reader. + * Pushes a LIMIT down to the reader. Spark only pushes a limit when the query does not filter the + * rows, since it evaluates the predicates after reading. * * applied: set to whether the reader returns at most `limit` rows in total. Spark applies the * limit again after reading either way. @@ -335,7 +369,9 @@ int spark_ds_push_limit(void* reader, int32_t limit, bool* applied); /* * Prunes the columns to read. * - * columns: the names of the count top-level columns that the query needs + * columns: the names of the count top-level columns that the query needs, which include the + * columns of the pushed predicates. There may be none, for example for + * SELECT count(*): then the batches only have a number of rows. * pruned: set to whether spark_ds_read returns exactly these columns, in this order. If not, it * returns all the columns of the schema to read. */ @@ -367,14 +403,14 @@ void spark_ds_close_reader(void* reader); /* * Returns the offset to start from when the streaming query has no checkpoint. * - * set_offset: called once with the offset, as JSON + * set_offset: called once with the offset, as compact JSON on a single line */ int spark_ds_initial_offset(void* reader, SparkDsSink set_offset, void* sink); /* * Returns the most recent offset available. * - * set_offset: called once with the offset, as JSON + * set_offset: called once with the offset, as compact JSON on a single line */ int spark_ds_latest_offset(void* reader, SparkDsSink set_offset, void* sink); diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala index 1a8ac1a968e87..d1bb7705c33f2 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala @@ -67,7 +67,10 @@ trait ColumnarDataSource { */ trait ColumnarReader extends Serializable { - /** Returns the predicates that Spark still has to evaluate after reading. */ + /** + * Returns the predicates that the reader does not evaluate completely. Spark evaluates all the + * predicates after reading anyway, and shows the others as pushed in the plan. + */ def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = predicates /** Returns whether the reader returns at most `limit` rows. Spark applies it again anyway. */ @@ -75,7 +78,8 @@ trait ColumnarReader extends Serializable { /** * Returns whether `read` returns exactly the columns of `requiredSchema`, the top-level columns - * that the query needs, instead of all the columns of the read schema. + * that the query needs, including the columns of the pushed predicates, instead of all the + * columns of the read schema. */ def pruneColumns(requiredSchema: StructType): Boolean = false diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala index eb973fe12a8f6..808025d01b7fd 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala @@ -105,7 +105,10 @@ class ColumnarScanBuilder( override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { val remaining = batchReader.pushPredicates(predicates) pushed = predicates.filterNot(remaining.contains) - remaining + // Spark evaluates all the predicates again after reading, even the ones that the reader + // evaluates completely: Spark SQL compares values in ways that a data source can easily get + // wrong, such as NaN with other values. The pushed predicates are shown in the plan. + predicates } override def pushedPredicates(): Array[Predicate] = pushed diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala index 4c9b814cbadc5..8759c2894ead9 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala @@ -69,8 +69,11 @@ class ColumnarWrite( } /** - * A batch write. The writer of the data source is created with the writer factory, once Spark - * computed the stages that the write depends on: Spark only aborts the write if it fails later. + * A batch write. The writer of the data source is created with the writer factory, right before + * the tasks of the write start: after the shuffle stages that adaptive query execution runs + * before the write, so no writer is created if they fail. The rest of the data, such as its last + * stage, is computed in the tasks of the write, and a failure there, as any later failure, aborts + * the write. So irreversible work belongs in the commit of the writer. */ class ColumnarBatchWrite(createWriter: () => ColumnarWriter, schema: StructType) extends BatchWrite { @@ -95,7 +98,9 @@ class ColumnarBatchWrite(createWriter: () => ColumnarWriter, schema: StructType) /** * A streaming write. As for a [[ColumnarBatchWrite]], the writer of the data source is created - * with the writer factory, once Spark computed the stages of the micro-batch. + * with the writer factory, right before the tasks of the micro-batch start. A stateful + * micro-batch runs without adaptive query execution, so all its stages, as those of any write + * without it, run after the writer is created. */ class ColumnarStreamingWrite(createWriter: () => ColumnarStreamWriter, schema: StructType) extends StreamingWrite { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index a8420f2817a89..797ca86bd8789 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -23,8 +23,11 @@ import java.util.concurrent.atomic.AtomicLong import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal +import com.fasterxml.jackson.core.JsonProcessingException +import com.fasterxml.jackson.databind.DeserializationFeature +import com.fasterxml.jackson.databind.json.JsonMapper import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, CDataDictionaryProvider, Data} -import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} +import org.apache.arrow.vector.{FieldVector, ValueVector, VectorSchemaRoot} import org.apache.arrow.vector.complex.MapVector import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} @@ -36,7 +39,7 @@ import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarDataSource, ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.types.{DataType, StructField, StructType} import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} import org.apache.spark.unsafe.Platform @@ -112,7 +115,7 @@ class NativeDataSource( create(library, "plan a write to", "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { dataSource => NativeArrow.exportSchema(schema) { address => - library.createStreamWriter(dataSource, address, overwrite) + library.createStreamWriter(dataSource, address, overwrite, queryId) } }, library.closeWriter) @@ -168,6 +171,9 @@ class NativeDataSourceReader( private var readerState: Array[Byte] = Array.emptyByteArray private var readSchema: StructType = schema @transient private var plannedPartitions: Array[InputPartition] = _ + // Why planning the partitions failed. The handle is released after the first try, but the + // scan can be planned again, for example by the lazy values of the physical plan that call it. + @transient private var planningError: Throwable = _ @transient private lazy val library = NativeLibraries.get(location) @@ -209,11 +215,18 @@ class NativeDataSourceReader( override def partitions(): Array[InputPartition] = synchronized { if (plannedPartitions == null) { + if (planningError != null) { + throw planningError + } try { val partitions = call(name, "plan a scan of")(library.partitions(handle.get)) readerState = bytesOrEmpty( call(name, "plan a scan of")(optional(library.serializeReader(handle.get)))) plannedPartitions = NativeInputPartition.from(partitions) + } catch { + case e: Throwable => + planningError = e + throw e } finally { // The executors only need the serialized state. handle.close() @@ -244,11 +257,35 @@ class NativeDataSourceStreamReader( @transient private lazy val library = NativeLibraries.get(location) override def initialOffset(): String = { - call(name, "plan a scan of")(library.initialOffset(handle.get)) + checkOffset("spark_ds_initial_offset", + call(name, "plan a scan of")(library.initialOffset(handle.get))) } override def latestOffset(): String = { - call(name, "plan a scan of")(library.latestOffset(handle.get)) + checkOffset("spark_ds_latest_offset", + call(name, "plan a scan of")(library.latestOffset(handle.get))) + } + + /** + * Checks an offset returned by the given function of the library: it must be JSON on a single + * line, because the offset log of the checkpoint stores each offset as a line. + */ + private def checkOffset(function: String, offset: String): String = { + if (offset.exists(c => c == '\n' || c == '\r')) { + throw error(name, "plan a scan of", s"$function returned an offset with a line break, " + + "but an offset must be JSON on a single line.") + } + val valid = try { + val json = NativeDataSourceStreamReader.mapper.readTree(offset) + json != null && !json.isMissingNode + } catch { + case _: JsonProcessingException => false + } + if (!valid) { + throw error(name, "plan a scan of", s"$function returned an offset that is not valid " + + s"JSON: $offset") + } + offset } override def partitions(start: String, end: String): Array[InputPartition] = { @@ -272,6 +309,12 @@ class NativeDataSourceStreamReader( override def stop(): Unit = handle.close() } +object NativeDataSourceStreamReader { + // Parses the offsets of the libraries to check them, and rejects anything after the JSON value. + private val mapper = + JsonMapper.builder().enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS).build() +} + /** Writes to a native data source for a batch write. */ class NativeDataSourceWriter( location: NativeLibraryLocation, @@ -462,7 +505,16 @@ class NativePartitionReader( s"$arrowType, which is not supported.") } } - val actual = ArrowUtils.fromArrowSchema(root.getSchema) + val actual = StructType(root.getSchema.getFields.asScala.map { field => + val dataType = try { + ArrowUtils.fromArrowField(field) + } catch { + case e: SparkUnsupportedOperationException if e.getCondition == "UNSUPPORTED_ARROWTYPE" => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"${e.getMessageParameters.get("typeName")}, which is not supported.") + } + StructField(field.getName, dataType, field.isNullable) + }.toSeq) // The schema to read as the library got it: Arrow does not carry everything, such as the // collations of strings, the fields of intervals and user-defined types, which still apply to // the data that Spark reads. @@ -481,7 +533,7 @@ class NativePartitionReader( private def columnVector(vector: FieldVector): ColumnVector = { try { - new ArrowColumnVector(vector) + new ReusedArrowColumnVector(vector) } catch { case NonFatal(e) => throw QueryExecutionErrors.nativeDataSourceError("read", name, @@ -593,6 +645,17 @@ class NativePartitionReader( } } +/** + * The Arrow column vector of a column of a [[NativePartitionReader]], which imports all its batches + * into the same Arrow vectors. It is not freeable between batches: Spark must not close it when it + * moves to the next batch, as `ColumnarToRowExec` does with the column vectors that are freeable, + * because closing a struct column also drops its child column vectors. The reader closes the + * Arrow vectors at the end. + */ +private class ReusedArrowColumnVector(vector: ValueVector) extends ArrowColumnVector(vector) { + override def closeIfFreeable(): Unit = {} +} + /** Writes the batches of a task to the data writer of a native library. */ class NativeDataWriter( location: NativeLibraryLocation, @@ -722,7 +785,10 @@ object NativeStreamWriters { /** Helpers to call native data source libraries. */ object NativeCalls { - /** Calls a native library, and rethrows the errors it reports as NATIVE_DATA_SOURCE_ERROR. */ + /** + * Calls a native library, and rethrows the errors it reports as NATIVE_DATA_SOURCE_ERROR, as + * well as the link errors, such as when Arrow cannot load its own JNI library. + */ def call[T](name: String, action: String)(f: => T): T = { try { f @@ -731,7 +797,7 @@ object NativeCalls { case e: NativeLibrary.MissingFunctionException => throw QueryExecutionErrors.nativeDataSourceError( action, name, s"The library does not implement ${e.getMessage}.", e) - case NonFatal(e) => + case e @ (NonFatal(_) | _: LinkageError) => throw QueryExecutionErrors.nativeDataSourceError(action, name, message(e), e) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala index de14ec5409867..938369ad6dd89 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.execution.datasources.v2.ffi import java.io.{File, FileInputStream, IOException} import java.nio.charset.StandardCharsets import java.nio.file.{Files, StandardCopyOption} +import java.nio.file.attribute.BasicFileAttributes import java.security.MessageDigest import java.util.{HexFormat, Locale} import java.util.concurrent.ConcurrentHashMap @@ -31,7 +32,7 @@ import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} import com.google.common.cache.{Cache, CacheBuilder} import org.apache.spark.internal.Logging -import org.apache.spark.internal.LogKeys.PATH +import org.apache.spark.internal.LogKeys.{PATH, VERSION_NUM} import org.apache.spark.sql.Artifact import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.util.Utils @@ -94,20 +95,33 @@ case class NativeDataSourcePackage( */ case class InvalidNativeDataSourcePackage(error: Throwable, dataSources: Seq[String]) +/** + * A version of a file, to cache what is read from it until it changes: the file is identified by + * its canonical path, and its version by its length, its modification time and its file key, such + * as the inode on Linux, which changes when another file replaces it, as `make install` does. + */ +private[ffi] case class FileVersion(path: String, length: Long, lastModified: Long, key: AnyRef) + +private[ffi] object FileVersion { + def apply(canonical: File): FileVersion = { + val attributes = Files.readAttributes(canonical.toPath, classOf[BasicFileAttributes]) + FileVersion(canonical.getPath, attributes.size(), attributes.lastModifiedTime().toMillis, + attributes.fileKey()) + } +} + object NativeDataSourcePackage extends Logging { val FILE_EXTENSION: String = Artifact.NATIVE_DATA_SOURCE_PACKAGE_EXTENSION val MANIFEST_NAME = "spark-native-datasource.json" - private case class CacheKey(path: String, length: Long, lastModified: Long) - private[ffi] type ReadResult = Either[InvalidNativeDataSourcePackage, NativeDataSourcePackage] // Bounded, because each session stores its artifacts in a directory of its own, so a // long-running server where many sessions add packages reads many different files. private val MAX_CACHED_PACKAGES = 1000 - private val packages: Cache[CacheKey, ReadResult] = - CacheBuilder.newBuilder().maximumSize(MAX_CACHED_PACKAGES).build[CacheKey, ReadResult]() + private val packages: Cache[FileVersion, ReadResult] = + CacheBuilder.newBuilder().maximumSize(MAX_CACHED_PACKAGES).build[FileVersion, ReadResult]() private val mapper = new ObjectMapper() /** Reads the package in the given file. The result is cached until the file changes. */ @@ -115,24 +129,31 @@ object NativeDataSourcePackage extends Logging { /** * Reads the package in the given file, or returns why it is invalid. The result is cached until - * the file changes, and an invalid package is logged when it is read. + * the file changes, and an invalid package, or a package of a version of the interface that + * Spark does not support, is logged when it is read. */ private[ffi] def tryRead(file: File): ReadResult = { - val canonical = try { - file.getCanonicalFile + val version = try { + FileVersion(file.getCanonicalFile) } catch { case e: IOException => return Left(InvalidNativeDataSourcePackage( invalidManifestError(file, s"The package cannot be read: $e", e), Nil)) } - val key = CacheKey(canonical.getPath, canonical.length(), canonical.lastModified()) - Option(packages.getIfPresent(key)).getOrElse { + Option(packages.getIfPresent(version)).getOrElse { + val canonical = new File(version.path) val result = parse(canonical) - result.swap.foreach { invalid => - logWarning(log"Invalid native data source package ${MDC(PATH, canonical)}.", - invalid.error) + result match { + case Left(invalid) => + logWarning(log"Invalid native data source package ${MDC(PATH, canonical)}.", + invalid.error) + case Right(pkg) if pkg.manifest.abiVersion != NativeLibrary.ABI_VERSION => + logWarning(log"The native data source package ${MDC(PATH, canonical)} implements " + + log"version ${MDC(VERSION_NUM, pkg.manifest.abiVersion)} of the native data source " + + log"interface, which this version of Spark does not support, so it is not used.") + case _ => } - packages.put(key, result) + packages.put(version, result) result } } @@ -243,22 +264,19 @@ object NativeDataSourcePackage extends Logging { case class InstalledNativeLibrary(path: String, sha256: String) extends NativeLibraryLocation object InstalledNativeLibrary { - private case class CacheKey(path: String, length: Long, lastModified: Long) - // Bounded, like the cache of the packages, although few libraries are installed. - private val checksums: Cache[CacheKey, String] = - CacheBuilder.newBuilder().maximumSize(1000).build[CacheKey, String]() + private val checksums: Cache[FileVersion, String] = + CacheBuilder.newBuilder().maximumSize(1000).build[FileVersion, String]() /** The installed library in the given file. */ def apply(file: File): InstalledNativeLibrary = InstalledNativeLibrary(file.getPath, sha256(file)) /** The SHA-256 checksum of the given file. The result is cached until the file changes. */ private[ffi] def sha256(file: File): String = { - val canonical = file.getCanonicalFile - val key = CacheKey(canonical.getPath, canonical.length(), canonical.lastModified()) - Option(checksums.getIfPresent(key)).getOrElse { - val checksum = NativeDataSourcePackage.sha256(canonical) - checksums.put(key, checksum) + val version = FileVersion(file.getCanonicalFile) + Option(checksums.getIfPresent(version)).getOrElse { + val checksum = NativeDataSourcePackage.sha256(new File(version.path)) + checksums.put(version, checksum) checksum } } @@ -291,10 +309,11 @@ object NativePlatform { */ object NativeLibraries extends Logging { // The libraries of the packages, by the checksum of the package, and the installed libraries, - // by their canonical path, so that each library file is loaded once. A library that is loaded - // but rejected stays loaded, so its error is kept: it is not loaded again. + // by their canonical path, with the checksum of the build that was loaded, so that each library + // file is loaded once. A library that is loaded but rejected stays loaded, so its error is kept: + // it is not loaded again. private val packagedLibraries = new ConcurrentHashMap[String, Try[NativeLibrary]]() - private val installedLibraries = new ConcurrentHashMap[String, Try[NativeLibrary]]() + private val installedLibraries = new ConcurrentHashMap[String, (String, Try[NativeLibrary])]() private lazy val extractionDir: File = Utils.createTempDir(namePrefix = "native-datasources") @@ -310,16 +329,26 @@ object NativeLibraries extends Logging { private def getSupported(location: NativeLibraryLocation): NativeLibrary = location match { case pkg: NativeDataSourcePackage => - getOrLoad(packagedLibraries, pkg.sha256)(loadPackaged(pkg)) + getOrLoad(packagedLibraries, pkg.sha256)(loadPackaged(pkg)).get case installed: InstalledNativeLibrary => + // The file has the checksum of the build of the driver. val file = locate(installed).getCanonicalFile - getOrLoad(installedLibraries, file.getPath)(loadInstalled(file)) + val (loadedChecksum, library) = + getOrLoad(installedLibraries, file.getPath)(installed.sha256 -> loadInstalled(file)) + // The JVM cannot load another build from the same path: the library was replaced since. + if (loadedChecksum != installed.sha256) { + throw QueryExecutionErrors.invalidNativeDataSourceLibraryError( + file.getPath, + "REPLACED", + Map("loadedChecksum" -> loadedChecksum, "checksum" -> installed.sha256)) + } + library.get } - private def getOrLoad(libraries: ConcurrentHashMap[String, Try[NativeLibrary]], key: String)( - load: => Try[NativeLibrary]): NativeLibrary = { + private def getOrLoad[T](libraries: ConcurrentHashMap[String, T], key: String)( + load: => T): T = { val library = libraries.get(key) - (if (library != null) library else libraries.computeIfAbsent(key, _ => load)).get + if (library != null) library else libraries.computeIfAbsent(key, _ => load) } private def loadPackaged(pkg: NativeDataSourcePackage): Try[NativeLibrary] = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala index e350b2ba7211b..71928a6c07647 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -207,9 +207,13 @@ object NativeDataSourceRegistry extends Logging { pkg.copy(artifactUUID = artifactUUID) } else { val copyName = s"${pkg.fileName}.${pkg.sha256}" - if (!artifactManager.hasFileArtifact(copyName)) { - artifactManager.addLocalArtifacts(Artifact.newFileArtifact( - Paths.get(copyName), new Artifact.LocalFile(Paths.get(pkg.path))) :: Nil) + // Queries of the session can use the package at the same time, and the artifact manager + // fails to add a file that another thread is adding. + artifactManager.synchronized { + if (!artifactManager.hasFileArtifact(copyName)) { + artifactManager.addLocalArtifacts(Artifact.newFileArtifact( + Paths.get(copyName), new Artifact.LocalFile(Paths.get(pkg.path))) :: Nil) + } } pkg.copy(artifactUUID = artifactUUID, distributedFileName = Some(copyName)) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala index 2ab8497a0476a..18167548c03e5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala @@ -18,7 +18,8 @@ package org.apache.spark.sql.execution.datasources.v2.ffi import java.util.Base64 -import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.core.json.JsonWriteFeature +import com.fasterxml.jackson.databind.json.JsonMapper import com.fasterxml.jackson.databind.node.ObjectNode import org.apache.spark.sql.connector.expressions.{Expression, GeneralScalarExpression, Literal, NamedReference} @@ -30,7 +31,8 @@ import org.apache.spark.sql.types._ * `spark_ds_push_predicates`. */ object NativePredicates { - private val mapper = new ObjectMapper() + // Writes ASCII only, escaping the other characters, which any JSON parser decodes. + private val mapper = JsonMapper.builder().enable(JsonWriteFeature.ESCAPE_NON_ASCII).build() /** * The names of the predicates that are passed to the libraries, as listed in diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index 782d9af407850..d1f6ae82bd187 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -20,17 +20,29 @@ // hand. It provides three data sources: // // native_range Reads rows [0, end) with columns of several types. Supports predicate, -// column and limit pushdown. With the option "extra", the schema has one more -// column: "map" adds a map whose entry fields are named "keys" and "values", -// as arrow-rs names them, "decimal256" adds a decimal256, and "listview" adds -// a list view. With the option "slice", the columns of the batches have an -// offset of 1. With the option "short", an array of the batches has one value -// fewer than its parent refers to, although its buffers hold them all: "id" -// shortens the column id, "entries" the entries of the map and their fields, -// and "values" the values of the entries. +// column and limit pushdown. With the option "claim_predicates", it reports +// all the predicates as evaluated completely, but ignores them. With the +// option "extra", the schema has one more column: "map" adds a map whose +// entry fields are named "keys" and "values", as arrow-rs names them, "struct" +// adds a struct, "decimal256" adds a decimal256, "listview" adds a list view, +// "largelist" adds a large list, and "dictionary" adds a dictionary-encoded +// string column, whose data Spark rejects before reading it. With the option +// "slice", the columns of the batches have an offset of 1. With the option +// "short", an array of the batches has one value fewer than its parent refers +// to, although its buffers hold them all: "id" shortens the column id, +// "entries" the entries of the map and their fields, and "values" the values +// of the entries. // native_sink Writes the rows as CSV files under the directory in the option "path". // native_counter Streams the rows [0, max_offset) with a single column "id", making "step" -// more rows available each time Spark asks for the latest offset. +// more rows available each time Spark asks for the latest offset. With the +// option "bad_offset", the latest offset is "none", "multiline" JSON or +// "invalid" JSON. With the option "closed_path", closing the stream reader +// creates that file. +// +// With the option "fail", a step fails with an error. With the option "zero", a function that +// creates a handle succeeds without setting it: "spark_ds_create", "spark_ds_create_reader", +// "spark_ds_create_stream_reader", "spark_ds_create_writer", "spark_ds_create_stream_writer" or +// "spark_ds_create_data_writer". // // The suite compiles it with these macros to get several distinct libraries: // @@ -95,9 +107,12 @@ const ColumnType kColumns[] = { // The columns that native_range adds with the option "extra", by its value. const std::map kExtraColumns = { - {"map", {"m", "+m"}}, // map: {"k": id} + {"map", {"m", "+m"}}, // map: {"k": id} + {"struct", {"st", "+s"}}, // struct: {id, "s"} {"decimal256", {"dec", "d:10,2,256"}}, // decimal256(10, 2): id - {"listview", {"lv", "+vl"}} // list_view: [id] + {"listview", {"lv", "+vl"}}, // list_view: [id] + {"largelist", {"ll", "+L"}}, // large_list, without data + {"dictionary", {"dict", "i"}} // dictionary, without data }; const ColumnType* findColumn(const std::string& name) { @@ -131,6 +146,7 @@ struct SchemaData { std::string name; std::vector> children; std::vector child_pointers; + std::unique_ptr dictionary; }; void releaseSchema(ArrowSchema* schema) { @@ -138,13 +154,16 @@ void releaseSchema(ArrowSchema* schema) { for (auto& child : data->children) { if (child->release != nullptr) child->release(child.get()); } + if (data->dictionary != nullptr && data->dictionary->release != nullptr) { + data->dictionary->release(data->dictionary.get()); + } delete data; schema->release = nullptr; } void initSchema(ArrowSchema* out, const std::string& format, const std::string& name, int64_t flags) { - auto* data = new SchemaData{format, name, {}, {}}; + auto* data = new SchemaData{format, name, {}, {}, nullptr}; out->format = data->format.c_str(); out->name = data->name.c_str(); out->metadata = nullptr; @@ -178,8 +197,16 @@ void exportSchema(const std::vector& columns, ArrowSchema* out) { ArrowSchema* entries = addChild(column, "+s", "entries", 0); addChild(entries, "u", "keys", 0); addChild(entries, "l", "values", ARROW_FLAG_NULLABLE); - } else if (name == "lv") { + } else if (name == "st") { + addChild(column, "l", "a", ARROW_FLAG_NULLABLE); + addChild(column, "u", "b", ARROW_FLAG_NULLABLE); + } else if (name == "lv" || name == "ll") { addChild(column, "l", "item", ARROW_FLAG_NULLABLE); + } else if (name == "dict") { + auto* data = static_cast(column->private_data); + data->dictionary = std::make_unique(); + initSchema(data->dictionary.get(), "u", "", ARROW_FLAG_NULLABLE); + column->dictionary = data->dictionary.get(); } } } @@ -318,6 +345,19 @@ void exportColumn(const std::string& name, int64_t lo, int64_t hi, ArrowArray* o initArray(addChild(entries), n, 0, {no_validity, toBytes(key_offsets), std::vector(keys.begin(), keys.end())}); initArray(addChild(entries), n, 0, {no_validity, toBytes(values)}); + } else if (name == "st") { + initArray(out, n, 0, {no_validity}); + std::vector ids; + std::vector offsets{0}; + std::string chars; + for (int64_t id = lo; id < hi; id++) { + ids.push_back(id); + chars += "s" + std::to_string(id); + offsets.push_back(static_cast(chars.size())); + } + initArray(addChild(out), n, 0, {no_validity, toBytes(ids)}); + initArray(addChild(out), n, 0, + {no_validity, toBytes(offsets), std::vector(chars.begin(), chars.end())}); } else if (name == "dec") { // 256-bit little-endian unscaled values of the scale 2. std::vector words; @@ -340,7 +380,8 @@ void exportColumn(const std::string& name, int64_t lo, int64_t hi, ArrowArray* o initArray(out, n, 0, {no_validity, toBytes(offsets), toBytes(sizes)}); initArray(addChild(out), n, 0, {no_validity, toBytes(values)}); } else { - throw std::runtime_error("unknown column " + name); + // Including "ll" and "dict", which Spark rejects from the schema, before reading data. + throw std::runtime_error("no data for the column " + name); } } @@ -614,6 +655,11 @@ void maybeFail(const Options& options, const std::string& step) { if (get(options, "fail", "") == step) throw std::runtime_error("injected failure in " + step); } +// Whether the given function succeeds without setting its handle, with the option "zero". +bool returnsZero(const Options& options, const std::string& function) { + return get(options, "zero", "") == function; +} + // The driver-side reader of native_range and native_counter. struct Reader { Options state; // Serialized to the executors. @@ -767,6 +813,7 @@ int spark_ds_create(const char* name_value, const char* const* option_keys, source->name != name("native_counter")) { throw std::runtime_error("unknown data source " + source->name); } + if (returnsZero(source->options, "spark_ds_create")) return; *handle = source.release(); }); } @@ -791,6 +838,7 @@ int spark_ds_create_reader(void* handle, struct ArrowSchema* schema, void** out) } auto columns = importColumns(schema); for (const auto& column : columns) findColumn(column); + if (returnsZero(source->options, "spark_ds_create_reader")) return; auto* reader = new Reader(); reader->state["columns"] = join(columns); reader->state["lo"] = "0"; @@ -800,6 +848,7 @@ int spark_ds_create_reader(void* handle, struct ArrowSchema* schema, void** out) reader->state["fail"] = get(source->options, "fail", ""); reader->state["slice"] = get(source->options, "slice", ""); reader->state["short"] = get(source->options, "short", ""); + reader->state["claim_predicates"] = get(source->options, "claim_predicates", ""); *out = reader; }); } @@ -812,13 +861,17 @@ int spark_ds_create_stream_reader(void* handle, struct ArrowSchema* schema, void if (source->name != name("native_counter")) { throw std::runtime_error(source->name + " does not support streaming reads"); } + auto columns = join(importColumns(schema)); + if (returnsZero(source->options, "spark_ds_create_stream_reader")) return; auto* reader = new Reader(); - reader->state["columns"] = join(importColumns(schema)); + reader->state["columns"] = columns; reader->state["max_offset"] = get(source->options, "max_offset", "10"); reader->state["step"] = get(source->options, "step", reader->state["max_offset"]); reader->state["current"] = "0"; reader->state["batch_size"] = get(source->options, "batch_size", "1000"); reader->state["commit_path"] = get(source->options, "commit_path", ""); + reader->state["bad_offset"] = get(source->options, "bad_offset", ""); + reader->state["closed_path"] = get(source->options, "closed_path", ""); *out = reader; }); } @@ -836,6 +889,12 @@ int spark_ds_push_predicates(void* handle, const char* const* predicates, int64_ auto json = toStrings(predicates, count); int64_t lo = getLong(reader->state, "lo", 0); int64_t hi = getLong(reader->state, "hi", 0); + // With the option "claim_predicates", reports every predicate as evaluated completely, but + // evaluates none of them. + if (reader->state["claim_predicates"] == "true") { + for (size_t i = 0; i < json.size(); i++) evaluated[i] = true; + return; + } for (size_t i = 0; i < json.size(); i++) { // Accepts "id " for the comparison operators. Json predicate = JsonParser(json[i]).parse(); @@ -928,7 +987,16 @@ int spark_ds_latest_offset(void* handle, SparkDsSink set_offset, void* sink) { int64_t latest = std::min(getLong(state, "max_offset", 10), getLong(state, "current", 0) + getLong(state, "step", 10)); state["current"] = std::to_string(latest); - sendBytes(set_offset, sink, "{\"offset\":" + std::to_string(latest) + "}"); + std::string bad_offset = get(state, "bad_offset", ""); + if (bad_offset == "none") return; + if (bad_offset == "multiline") { + // Pretty-printed JSON, as serde_json::to_string_pretty writes it. + sendBytes(set_offset, sink, "{\n \"offset\": " + std::to_string(latest) + "\n}"); + } else if (bad_offset == "invalid") { + sendBytes(set_offset, sink, "offset " + std::to_string(latest)); + } else { + sendBytes(set_offset, sink, "{\"offset\":" + std::to_string(latest) + "}"); + } }); } @@ -952,7 +1020,13 @@ int spark_ds_commit_offset(void* handle, const char* end) { }); } -void spark_ds_close_stream_reader(void* handle) { delete fromHandle(handle); } +void spark_ds_close_stream_reader(void* handle) { + std::unique_ptr reader(fromHandle(handle)); + guarded([&] { + std::string path = get(reader->state, "closed_path", ""); + if (!path.empty()) writeFile(path, ""); + }); +} #endif // READ_ONLY @@ -984,7 +1058,9 @@ int spark_ds_read(const uint8_t* reader_state, int64_t reader_state_length, namespace { -int createAnyWriter(void* handle, struct ArrowSchema* schema, bool overwrite, void** out) { +// Creates the writer of a batch write, or of a streaming write with the ID of its query. +int createAnyWriter(const char* function, void* handle, struct ArrowSchema* schema, + bool overwrite, const char* query_id, void** out) { return guarded([&] { auto* source = fromHandle(handle); if (source->name != name("native_sink")) { @@ -1000,13 +1076,17 @@ int createAnyWriter(void* handle, struct ArrowSchema* schema, bool overwrite, vo path = path.substr(5); } if (path.empty()) throw std::runtime_error("the option 'path' is required"); + if (returnsZero(source->options, function)) return; auto* writer = new Writer(); writer->state["path"] = path; writer->state["overwrite"] = overwrite ? "true" : "false"; writer->state["fail"] = get(source->options, "fail", ""); - // Records that a writer was created, which Spark then commits or aborts. + writer->state["zero"] = get(source->options, "zero", ""); + // Records that a writer was created, which Spark then commits or aborts, and the ID of the + // streaming query. makeDirs(path); writeFile(path + "/_CREATED", ""); + if (query_id != nullptr) writeFile(path + "/_QUERY_ID", query_id); *out = writer; }); } @@ -1017,12 +1097,13 @@ extern "C" { int spark_ds_create_writer(void* handle, struct ArrowSchema* schema, bool overwrite, void** writer) { - return createAnyWriter(handle, schema, overwrite, writer); + return createAnyWriter("spark_ds_create_writer", handle, schema, overwrite, nullptr, writer); } int spark_ds_create_stream_writer(void* handle, struct ArrowSchema* schema, bool overwrite, - void** writer) { - return createAnyWriter(handle, schema, overwrite, writer); + const char* query_id, void** writer) { + return createAnyWriter("spark_ds_create_stream_writer", handle, schema, overwrite, query_id, + writer); } int spark_ds_serialize_writer(void* handle, SparkDsSink set_state, void* sink) { @@ -1051,6 +1132,8 @@ int spark_ds_commit(void* handle, int64_t epoch_id, const uint8_t* const* messag }); } +// Records the aborted write in _ABORTED for a batch write, or _ABORTED_ for a streaming +// one, with the number of tasks that did not commit, whose messages are NULL. int spark_ds_abort(void* handle, int64_t epoch_id, const uint8_t* const* messages, const int64_t* message_lengths, int64_t num_messages) { return guarded([&] { @@ -1061,7 +1144,13 @@ int spark_ds_abort(void* handle, int64_t epoch_id, const uint8_t* const* message toByteStrings(messages, message_lengths, num_messages, &has_null)) { unlink(temp_file.c_str()); } - writeFile(writer->state["path"] + "/_ABORTED", ""); + int64_t null_messages = 0; + for (int64_t i = 0; i < num_messages; i++) { + if (messages[i] == nullptr) null_messages++; + } + writeFile(writer->state["path"] + + (epoch_id < 0 ? "/_ABORTED" : "/_ABORTED_" + std::to_string(epoch_id)), + std::to_string(null_messages)); }); } @@ -1079,6 +1168,7 @@ int spark_ds_create_data_writer(const uint8_t* writer_state, int64_t writer_stat void** out) { return guarded([&] { Options state = deserialize(toBytes(writer_state, writer_state_length)); + if (returnsZero(state, "spark_ds_create_data_writer")) return; std::string dir = state["path"] + "/_temporary"; makeDirs(dir); std::string temp_path = dir + "/" + std::to_string(epoch_id) + "-" + diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala index fa594f2c20ff6..bc7f369ba930f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala @@ -133,10 +133,11 @@ class NativeDataSourceClusterSuite extends QueryTest { assert(distributed.artifactUUID.contains(session.sessionUUID)) assert(distributed.distributedFileName.contains(s"${pkg.getName}.${found.sha256}")) - // Then the package at the path of the driver becomes unusable for the executors, while the - // driver keeps the package it read, since the file keeps its length and modification time. - // So the query only succeeds if the executors find the distributed copy, and check that it - // is the same package by its checksum. + // Then the file at the path of the driver is no longer a valid package, so the executors + // cannot use it, while the driver keeps the package it read, since the file is written in + // place and keeps its length and modification time. So the query only succeeds if the + // executors find the distributed copy. NativeDataSourceSuite checks that the executors only + // load a valid package with the checksum of the driver. val lastModified = pkg.lastModified() Files.write(pkg.toPath, new Array[Byte](pkg.length().toInt)) assert(pkg.setLastModified(lastModified)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index f922793d7c780..0797fe1212254 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -19,22 +19,31 @@ package org.apache.spark.sql.execution.datasources.v2.ffi import java.io.File import java.nio.charset.StandardCharsets.UTF_8 import java.nio.file.Files +import java.nio.file.StandardCopyOption.REPLACE_EXISTING import java.sql.{Date, Timestamp} import java.time.{Instant, LocalDate} import java.util.UUID +import java.util.concurrent.{Callable, CountDownLatch, TimeUnit} import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ +import com.fasterxml.jackson.core.json.JsonWriteFeature +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.json.JsonMapper + import org.apache.spark.{SparkClassNotFoundException, SparkFiles, SparkRuntimeException, SparkThrowable, SparkUnsupportedOperationException} import org.apache.spark.api.python.{PythonBroadcast, SimplePythonFunction} import org.apache.spark.broadcast.Broadcast import org.apache.spark.sql.{AnalysisException, QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.{Literal => CatalystLiteral} +import org.apache.spark.sql.catalyst.plans.logical.Filter import org.apache.spark.sql.classic.{DataFrame, SparkSession} import org.apache.spark.sql.connector.expressions.{Expression, FieldReference, GeneralScalarExpression, LiteralValue} import org.apache.spark.sql.connector.expressions.filter.{And, Predicate} +import org.apache.spark.sql.execution.WholeStageCodegenExec import org.apache.spark.sql.execution.datasources.DataSource -import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2ScanRelation, DataSourceV2Strategy} import org.apache.spark.sql.execution.datasources.v2.columnar.ColumnarScan import org.apache.spark.sql.execution.datasources.v2.python.{PythonDataSourceV2, UserDefinedPythonDataSource} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} @@ -43,7 +52,7 @@ import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.unsafe.types.UTF8String -import org.apache.spark.util.Utils +import org.apache.spark.util.{ThreadUtils, Utils} /** * Tests native data sources with the library in `native-datasource/test_native_datasource.cc`, @@ -207,6 +216,12 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(literal(2.5, DoubleType).contains(""""dataType":"double","value":2.5""")) assert(literal(UTF8String.fromString("a\"b"), StringType) .contains(""""dataType":"string","value":"a\"b"""")) + // The JSON is ASCII, with the other characters escaped, such as the characters outside the + // Basic Multilingual Plane. + val text = "x" + new String(Character.toChars(0x1F600)) + new String(Character.toChars(0xE9)) + val encoded = literal(UTF8String.fromString(text), StringType).get + assert(encoded.forall(_ < 128), encoded) + assert(new ObjectMapper().readTree(s"{$encoded}").get("value").asText() == text) assert(literal(Array[Byte](1, 2), BinaryType) .contains(""""dataType":"binary","value":"AQI="""")) assert(literal(Decimal("12.30"), DecimalType(10, 2)) @@ -238,6 +253,113 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(compare("id.b", "id").isEmpty) } + /** + * Generates `org/apache/spark/sql/datasource/predicate-cases.json`, which lists predicates as + * Spark passes them to `spark_ds_push_predicates`, with the rows of a table for which Spark + * finds them true, for libraries to test their pushdown with. Regenerate it with + * `SPARK_GENERATE_GOLDEN_FILES=1`. + */ + test("predicate cases") { + val schema = StructType.fromDDL("id INT, i INT, d DOUBLE, s STRING") + def char(codePoint: Int): String = new String(Character.toChars(codePoint)) + val (emoji, eAcute, replacement) = (char(0x1F600), char(0xE9), char(0xFFFD)) + val rows = Seq( + Row(0, 1, 0.0, "a"), + Row(1, 2, -0.0, ""), + Row(2, null, Double.NaN, eAcute), + Row(3, 3, Double.PositiveInfinity, emoji), + Row(4, 1, Double.NegativeInfinity, null), + Row(5, null, 1.5, "ab"), + Row(6, 2, null, "b" + eAcute), + Row(7, 4, 2.5, replacement)) + val filters = Seq( + // NaN is equal to NaN and greater than any other value, and -0.0 is equal to 0.0. + "d = 0.0D", "d = -0.0D", "d < 0.0D", "d > 1.0D", "d >= 2.5D", "d <> 0.0D", "d <=> 0.0D", + "d IN (0.0D, 1.5D)", "NOT (d > 1.0D)", "NOT (d < 1.0D)", "d IS NULL", "d IS NOT NULL", + // Comparisons with null are null, also under NOT. + "i = 1", "i = NULL", "i <=> NULL", "NOT (i <=> NULL)", "i IN (1, NULL)", + "NOT (i IN (1, NULL))", "NOT (i IN (1, 3))", "i IS NULL OR d > 1.0D", + "i > 1 AND d IS NOT NULL", + // Strings are compared by their UTF-8 bytes. + s"s = '$emoji'", s"s > '$replacement'", "s < 'b'", "startswith(s, 'a')", + s"endswith(s, '$eAcute')", s"contains(s, '$emoji')", s"s IN ('', '$eAcute')") + + val mapper = JsonMapper.builder().enable(JsonWriteFeature.ESCAPE_NON_ASCII).build() + def json(value: Any): String = mapper.writeValueAsString(value) + def jsonValue(value: Any): String = value match { + case d: Double if d.isNaN || d.isInfinite => json(d.toString) + case other => json(other) + } + val df = spark.createDataFrame(rows.asJava, schema) + val cases = filters.map { filter => + val filtered = df.where(filter) + // The predicates that Spark pushes for the filter: its conjuncts, with constants folded. + val condition = filtered.queryExecution.analyzed.collectFirst { + case Filter(condition, _) => condition + }.get + val predicates = splitConjunctivePredicates(condition.transformUp { + case e if e.foldable && !e.isInstanceOf[CatalystLiteral] => + CatalystLiteral.create(e.eval(), e.dataType) + }).map { predicate => + DataSourceV2Strategy.translateFilterV2(predicate) + .flatMap(NativePredicates.toJson(_, schema)) + .getOrElse(fail(s"$filter is not pushed down")) + } + val matches = filtered.select("id").collect().map(_.getInt(0)).sorted + s""" {"sql": ${json(filter)}, "predicates": [${predicates.mkString(", ")}], """ + + s""""rows": [${matches.mkString(", ")}]}""" + } + val generated = + s"""{ + | "description": ${json("Predicates as Spark passes them to spark_ds_push_" + + "predicates, with the ids of the rows of the table for which they are true. A " + + "case lists the predicates that Spark pushes for the filter in its SQL, which are " + + "combined with AND. A row has a value for each column of the schema; NaN and infinite " + + "doubles are strings.")}, + | "schema": ${json(schema.toDDL)}, + | "rows": [ + |${rows.map(row => row.toSeq.map(jsonValue).mkString(" [", ", ", "]")).mkString(",\n")} + | ], + | "cases": [ + |${cases.mkString(",\n")} + | ] + |} + |""".stripMargin + assert(generated.forall(_ < 128)) + val file = getWorkspaceFilePath("sql", "core", "src", "main", "resources", "org", "apache", + "spark", "sql", "datasource", "predicate-cases.json").toFile + if (regenerateGoldenFiles) { + Files.createDirectories(file.toPath.getParent) + Files.writeString(file.toPath, generated) + } else { + assert(Files.readString(file.toPath) == generated, + s"$file is out of date: regenerate it with SPARK_GENERATE_GOLDEN_FILES=1.") + } + // The file is in the jar of Spark SQL. + assert(getClass.getClassLoader.getResource( + "org/apache/spark/sql/datasource/predicate-cases.json") != null) + } + + test("tell missing functions from link errors") { + // A function that the library does not export is an unsupported operation. + val missing = new NativeLibrary.MissingFunctionException("spark_ds_f") + assert(NativeCalls.optional(throw missing).isEmpty) + checkError( + exception = intercept[SparkRuntimeException](NativeCalls.call("x", "read")(throw missing)), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "read", "name" -> "x", "msg" -> "The library does not implement spark_ds_f.")) + // Link errors are reported as they are, such as when Arrow, which still calls the Arrow C + // structs through JNI, cannot load its JNI library, which it extracts to java.io.tmpdir, for + // example if the directory does not allow it. + val linkError = new UnsatisfiedLinkError("/tmp/libarrow_cdata_jni.so: failed to map segment") + assert(intercept[UnsatisfiedLinkError](NativeCalls.optional(throw linkError)) eq linkError) + checkError( + exception = intercept[SparkRuntimeException](NativeCalls.call("x", "read")(throw linkError)), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map("action" -> "read", "name" -> "x", "msg" -> linkError.getMessage)) + } + test("close the previous stream writer of a streaming query") { val released = ArrayBuffer.empty[Long] def handle(value: Long): NativeHandle = new NativeHandle(value, released += _) @@ -334,6 +456,20 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { .filter("id < 10 AND name LIKE 'n1%'") checkRows(partly, Seq(1L).map(expectedRow)) assert(scan(partly).getMetaData()("PushedPredicates").contains("id < 10")) + + // Spark evaluates all the predicates again after reading, even the ones that the library + // reports as evaluated completely: here, a library that ignores them. + val claimed = range(newSession(defaultPackage), "end" -> "20", "claim_predicates" -> "true") + .filter("id > 15") + checkRows(claimed, (16L until 20L).map(expectedRow)) + assert(scan(claimed).getMetaData()("PushedPredicates").contains("id > 15")) + + // The answers of the library are for the predicates that Spark passes to it: here, not the + // first one, which uses another function. + val mapped = range(newSession(defaultPackage), "end" -> "20") + .filter("upper(name) = 'N1' AND id < 5") + checkRows(mapped, Seq(expectedRow(1L))) + assert(scan(mapped).getMetaData()("PushedPredicates") == "[id < 5]") } nativeTest("prune columns") { @@ -341,6 +477,13 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { checkRows(df, (0L until 10L).map(id => Row(if (id % 5 == 0) null else s"n$id", id))) assert(scan(df).readSchema() == StructType.fromDDL("id BIGINT, name STRING")) + // The columns of the pushed predicates are read too, since Spark evaluates the predicates + // again after reading, even the ones that the library evaluates completely. + val filtered = range(newSession(defaultPackage), "end" -> "10").filter("id > 7").select("name") + checkRows(filtered, Seq(Row("n8"), Row("n9"))) + assert(scan(filtered).getMetaData()("PushedPredicates").contains("id > 7")) + assert(scan(filtered).readSchema() == StructType.fromDDL("id BIGINT, name STRING")) + // A count reads no column at all. assert(range(newSession(defaultPackage), "end" -> "37", "batch_size" -> "10").count() == 37) } @@ -356,6 +499,11 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { .limit(10).offset(5) checkRows(withOffset, (5L until 10L).map(expectedRow)) assert(scan(withOffset).getMetaData()("PushedLimit") == "LIMIT 10") + + // Spark evaluates the predicates after reading, so it does not push a limit with them. + val filtered = range(newSession(defaultPackage), "end" -> "100").filter("id > 5").limit(3) + assert(filtered.collect().length == 3) + assert(!scan(filtered).getMetaData().contains("PushedLimit")) } nativeTest("read maps whose entry fields have other names") { @@ -365,10 +513,32 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { checkRows(df, (0L until 3L).map(id => Row(id, Map(s"k$id" -> id)))) } + nativeTest("read structs in several batches") { + // The reader imports the batches of a partition into the same vectors, which Spark must not + // close when it moves to the next batch: closing a struct drops its child vectors. + val df = range(newSession(defaultPackage), "end" -> "10", "partitions" -> "1", + "batch_size" -> "3", "extra" -> "struct").select("id", "st") + assert(df.queryExecution.executedPlan.exists(_.isInstanceOf[WholeStageCodegenExec])) + checkRows(df, (0L until 10L).map(id => Row(id, Row(id, s"s$id")))) + + // A nested field is read with its top-level column: the library only prunes the top-level + // columns, and Spark the nested fields. + val nested = range(newSession(defaultPackage), "end" -> "4", "extra" -> "struct") + .select("st.b") + checkRows(nested, (0L until 4L).map(id => Row(s"s$id"))) + assert(scan(nested).readSchema() == StructType.fromDDL("st STRUCT")) + } + nativeTest("reject the Arrow data that Spark cannot read") { def check(msg: String, options: (String, String)*): Unit = { + checkWithSchema(None, msg, options: _*) + } + // The schema to read is given for the types from which Spark cannot infer it either. + def checkWithSchema(schema: Option[String], msg: String, options: (String, String)*): Unit = { val e = intercept[Exception] { - range(newSession(defaultPackage), options :+ ("end" -> "4"): _*).collect() + val reader = newSession(defaultPackage).read.format("native_range") + schema.fold(reader)(s => reader.schema(s)).options(options.toMap + ("end" -> "4")).load() + .collect() } checkError( exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), @@ -391,6 +561,13 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { // check their lengths. check("The column lv contains the Arrow type ListView, which is not supported.", "extra" -> "listview") + // Types that Spark does not convert, and dictionary encoding. + checkWithSchema(Some("id BIGINT, ll ARRAY"), + "The column ll contains the Arrow type LargeList, which is not supported.", + "extra" -> "largelist") + checkWithSchema(Some("id BIGINT, dict STRING"), + "The column dict is dictionary-encoded, which is not supported.", + "extra" -> "dictionary") } nativeTest("batch write") { @@ -520,13 +697,40 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) } + nativeTest("abort a micro-batch whose write failed") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val e = intercept[Exception] { + val query = session.readStream.format("native_counter").option("max_offset", "4").load() + .writeStream + .format("native_sink") + .option("path", dir.getPath) + .option("fail", "write") + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try query.processAllAvailable() finally query.stop() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "write to", + "name" -> "native_sink", + "msg" -> "injected failure in write")) + // Spark aborted the first epoch, with a null message for its task, which did not commit. + assert(Files.readString(new File(dir, "_ABORTED_0").toPath) == "1") + assert(!new File(dir, "_EPOCH_0").exists()) + } + nativeTest("streaming read and write") { val session = newSession(defaultPackage) val dir = Utils.createTempDir() val commitPath = new File(Utils.createTempDir(), "committed") + val closedPath = new File(Utils.createTempDir(), "closed") val query = session.readStream.format("native_counter") .option("max_offset", "12").option("step", "5").option("batch_size", "2") .option("commit_path", commitPath.getPath) + .option("closed_path", closedPath.getPath) .load() .writeStream .format("native_sink") @@ -548,6 +752,8 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(written.sorted.toSeq == (0 until 12).map(_.toString).sorted) // Spark commits the offsets of a micro-batch when it runs the next one. assert(Files.readString(commitPath.toPath) == "{\"offset\":10}") + // The stream reader is closed when the query stops. + assert(closedPath.isFile) // The streaming source can also be read in micro-batches into a memory sink. val memoryQuery = session.readStream.format("native_counter").option("max_offset", "4") @@ -560,6 +766,56 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } } + nativeTest("reject offsets that are not JSON on a single line") { + // Spark writes each offset as a line of the offset log of the checkpoint. + def check(badOffset: String, msg: String): Unit = { + val e = intercept[Exception] { + val query = newSession(defaultPackage).readStream.format("native_counter") + .option("bad_offset", badOffset).load() + .writeStream + .format("noop") + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try query.processAllAvailable() finally query.stop() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map("action" -> "plan a scan of", "name" -> "native_counter", "msg" -> msg)) + } + check("none", "spark_ds_latest_offset returned no offset.") + check("multiline", "spark_ds_latest_offset returned an offset with a line break, but an " + + "offset must be JSON on a single line.") + check("invalid", "spark_ds_latest_offset returned an offset that is not valid JSON: offset 10") + } + + nativeTest("pass the ID of a streaming query to its stream writers") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val checkpoint = Utils.createTempDir().getPath + // Runs the query until it processed the rows [0, maxOffset). + def run(maxOffset: Int): StreamingQuery = { + val query = session.readStream.format("native_counter") + .option("max_offset", maxOffset.toString).load() + .writeStream + .format("native_sink") + .option("path", dir.getPath) + .option("checkpointLocation", checkpoint) + .start() + try query.processAllAvailable() finally query.stop() + query + } + val queryId = new File(dir, "_QUERY_ID") + val first = run(3) + assert(Files.readString(queryId.toPath) == first.id.toString) + // The ID stays the same when the query restarts from its checkpoint, unlike the ID of the + // run, so that the library can tell the commits of the query from those of other queries. + assert(queryId.delete()) + val restarted = run(6) + assert(restarted.id == first.id && restarted.runId != first.runId) + assert(Files.readString(queryId.toPath) == first.id.toString) + } + nativeTest("report the errors of native code") { def check(action: String, msg: String, name: String = "native_range")( f: SparkSession => Unit): Unit = { @@ -606,6 +862,78 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { "msg" -> "spark_ds_create failed with the error code 1.")) } + nativeTest("a scan whose planning failed fails again with the same error") { + // The reader is released after planning, even if it failed. The query execution of a + // DataFrame keeps the failure of its plan, but a scan can be planned again, for example by + // the lazy values of the physical plan that call it. + val session = newSession(defaultPackage) + val (name, location) = + NativeDataSourceRegistry.lookup("native_range", session.sessionState.conf).get + val options = new CaseInsensitiveStringMap(Map("fail" -> "plan").asJava) + val reader = + new NativeDataSource(location, name, options).reader(StructType.fromDDL("id BIGINT")) + (1 to 2).foreach { _ => + checkError( + exception = intercept[SparkRuntimeException](reader.partitions()), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "plan a scan of", + "name" -> "native_range", + "msg" -> "injected failure in plan")) + } + } + + nativeTest("reject NULL handles") { + def check(function: String, action: String, name: String)(f: SparkSession => Unit): Unit = { + val e = intercept[Exception](f(newSession(defaultPackage))) + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> action, + "name" -> name, + "msg" -> s"$function returned a NULL handle.")) + } + def runStream(start: => StreamingQuery): Unit = { + val query = start + try query.processAllAvailable() finally query.stop() + } + check("spark_ds_create", "create", "native_range") { session => + range(session, "zero" -> "spark_ds_create") + } + check("spark_ds_create_reader", "plan a scan of", "native_range") { session => + range(session, "zero" -> "spark_ds_create_reader").collect() + } + check("spark_ds_create_stream_reader", "plan a scan of", "native_counter") { session => + runStream { + session.readStream.format("native_counter") + .option("zero", "spark_ds_create_stream_reader").load() + .writeStream.format("noop").start() + } + } + Seq( + "spark_ds_create_writer" -> "plan a write to", + "spark_ds_create_data_writer" -> "write to" + ).foreach { case (function, action) => + check(function, action, "native_sink") { session => + session.range(2).write.format("native_sink") + .option("path", Utils.createTempDir().getPath).option("zero", function) + .mode("append").save() + } + } + check("spark_ds_create_stream_writer", "plan a write to", "native_sink") { session => + runStream { + session.readStream.format("native_counter").load() + .writeStream + .format("native_sink") + .option("path", Utils.createTempDir().getPath) + .option("zero", "spark_ds_create_stream_writer") + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + } + } + } + nativeTest("operations that the library does not implement") { val session = newSession(createPackage(readOnlyLibrary, dataSources("ro_"))) // The library does not push anything down, so Spark evaluates the filter. @@ -707,6 +1035,42 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } } + nativeTest("distribute a configured package for concurrent queries") { + // Queries of a session can use a configured package for the first time at the same time. + val configured = createPackage(library, dataSources(""), dir = Utils.createTempDir()) + val session = newSession(configured) + val pkg = NativeDataSourcePackage.read(configured) + val threads = 8 + val pool = ThreadUtils.newDaemonFixedThreadPool(threads, "distribute-package") + try { + val start = new CountDownLatch(1) + val results = (1 to threads).map { _ => + pool.submit(new Callable[NativeDataSourcePackage] { + override def call(): NativeDataSourcePackage = { + start.await() + NativeDataSourceRegistry.addToArtifacts(session, pkg) + } + }) + } + start.countDown() + val distributed = results.map(_.get(1, TimeUnit.MINUTES)).distinct + assert(distributed.length == 1) + assert(session.artifactManager.hasFileArtifact(distributed.head.distributedFileName.get)) + } finally { + pool.shutdownNow() + } + } + + nativeTest("load a package only with the checksum of the driver") { + // The package at the path of the driver is valid, but it is not the package that the driver + // read, for example because it was replaced since, and neither is any distributed copy. + val pkg = NativeDataSourcePackage.read(defaultPackage).copy(sha256 = "0" * 64) + checkError( + exception = intercept[SparkRuntimeException](NativeLibraries.get(pkg)), + condition = "NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND", + parameters = Map("fileName" -> defaultPackage.getName, "checksum" -> pkg.sha256)) + } + test("distribute a package without an active session") { val pkg = NativeDataSourcePackage( "/path/to/x.sparkpkg", "checksum", NativeDataSourceManifest(1, Seq("x"), Map.empty)) @@ -902,6 +1266,36 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { parameters = Map("provider" -> "installed_native_range")) } + nativeTest("reject an installed library replaced in place") { + val dir = Utils.createTempDir() + val installed = installLibrary(installedLibrary, "installed_native_range", dir) + withLibraryPath(dir) { + def read(): DataFrame = { + newSession().read.format("installed_native_range").option("end", "2").load().select("id") + } + checkRows(read(), Seq(Row(0L), Row(1L))) + val loaded = NativeDataSourcePackage.sha256(installed) + + // An operator installs another build at the same path, as `make install` does. The JVM + // cannot load it from the same path, so the driver and the executors that loaded the + // previous build reject it, rather than running a build that does not match the state + // that the driver serializes. + Files.copy(installedNegatingLibrary.toPath, installed.toPath, REPLACE_EXISTING) + val replaced = NativeDataSourcePackage.sha256(installed) + checkError( + exception = intercept[SparkRuntimeException](read()), + condition = "INVALID_NATIVE_DATA_SOURCE_LIBRARY.REPLACED", + parameters = Map( + "path" -> installed.getCanonicalPath, + "loadedChecksum" -> loaded, + "checksum" -> replaced)) + + // The library works again once the build that was loaded is installed again. + Files.copy(installedLibrary.toPath, installed.toPath, REPLACE_EXISTING) + checkRows(read(), Seq(Row(0L), Row(1L))) + } + } + nativeTest("find packages in directories") { val dir = Utils.createTempDir() createPackage(library, dataSources(""), dir = dir, fileName = "a.sparkpkg") @@ -1049,6 +1443,23 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } } + nativeTest("log the packages of a version that Spark does not support") { + // An operator left an outdated package in a configured directory. + val dir = Utils.createTempDir() + val outdated = createPackage(library, Seq("other_native_range"), dir = dir, + fileName = "outdated.sparkpkg", abiVersion = 2) + val logAppender = new LogAppender("packages of an unsupported version") + withLogAppender(logAppender, level = Some(org.apache.logging.log4j.Level.WARN)) { + checkRows(range(newSession(dir, defaultPackage), "end" -> "2"), + (0L until 2L).map(expectedRow)) + } + assert(logAppender.loggingEvents.exists { event => + val message = event.getMessage.getFormattedMessage + message.contains(outdated.getCanonicalPath) && + message.contains("implements version 2 of the native data source interface") + }) + } + nativeTest("forget the packages of a session that is cleaned up") { val session = spark.newSession() session.addArtifact(defaultPackage.getPath) From 4f032d5f1b9ce1086b2ce47b0126a7ff1d0eabd0 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Wed, 14 Oct 2026 10:38:00 +0900 Subject: [PATCH 21/27] Ship spark_datasource.h in the include directory of a Spark distribution The header is published like the headers of other C interfaces: a binary distribution contains it in include/, as a JDK contains include/jni.h, and libraries can copy it into their source tree, since it only depends on the C standard library. The spark-sql JAR still contains it, and the docs describe both ways to get it. Co-authored-by: Isaac --- dev/make-distribution.sh | 5 ++++ docs/sql-native-data-sources.md | 23 +++++++++++++------ .../spark/sql/datasource/spark_datasource.h | 6 +++++ 3 files changed, 27 insertions(+), 7 deletions(-) diff --git a/dev/make-distribution.sh b/dev/make-distribution.sh index 4a2bd09910129..fc73cbc9ac058 100755 --- a/dev/make-distribution.sh +++ b/dev/make-distribution.sh @@ -223,6 +223,11 @@ if [ -f "$DISTDIR"/jars/datasketches-memory-3.0.2.jar ]; then zip -d "$DISTDIR"/jars/datasketches-memory-3.0.2.jar org/apache/datasketches/memory/internal/ResourceImpl.class fi +# Copy the C header that the libraries of native data sources implement +mkdir -p "$DISTDIR/include" +cp "$SPARK_HOME/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h" \ + "$DISTDIR/include/" + # Only create the yarn directory if the yarn artifacts were built. if [ -f "$SPARK_HOME"/common/network-yarn/target/scala*/spark-*-yarn-shuffle.jar ]; then mkdir "$DISTDIR/yarn" diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 8554b6cb6be22..82d63c0202e53 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -52,6 +52,19 @@ A native data source library exports the functions of `spark_datasource.h`, such functions up in each library separately, so any number of native data sources can be loaded side by side, even though they export the same functions. +A Spark distribution contains the header in its `include` directory, as +`$SPARK_HOME/include/spark_datasource.h`, and the `spark-sql` JAR contains it as +`org/apache/spark/sql/datasource/spark_datasource.h`. The header only depends on the C standard +library, so a library can also copy it into its source tree, to build without Spark, for example +from the source of a release: + +```bash +curl -LO {{site.SPARK_GITHUB_URL}}/raw/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +``` + +A library built against the header works with all the releases of Spark that support its version, +`SPARK_DS_ABI_VERSION`: see [Compatibility](#compatibility). + Spark calls the functions of the library as follows: - **Data sources**: for each schema inference, scan and write, Spark creates a data source on the @@ -474,11 +487,7 @@ EOF The same data source in C++, with [Arrow C++](https://arrow.apache.org/docs/cpp/), which exports schemas and streams with `arrow::ExportSchema` and `arrow::ExportRecordBatchReader`. It includes -`spark_datasource.h`, which is also in the `spark-sql` JAR of a Spark distribution: - -```bash -unzip -j "$SPARK_HOME"/jars/spark-sql_*.jar org/apache/spark/sql/datasource/spark_datasource.h -``` +`spark_datasource.h`, here from the `include` directory of a Spark distribution. `cpp_range.cc`: @@ -597,8 +606,8 @@ Build the library, here on macOS with Arrow C++ 24, which requires C++20, and in package it as above: ```bash -c++ -std=c++20 -O2 -shared -fPIC -I. -I"$ARROW_HOME/include" -L"$ARROW_HOME/lib" -larrow \ - -o libspark_datasource_cpp_range.dylib cpp_range.cc +c++ -std=c++20 -O2 -shared -fPIC -I"$SPARK_HOME/include" -I"$ARROW_HOME/include" \ + -L"$ARROW_HOME/lib" -larrow -o libspark_datasource_cpp_range.dylib cpp_range.cc cp libspark_datasource_cpp_range.dylib /usr/local/lib/ ``` diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h index db08aa28e851e..ac894a4d84401 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -26,6 +26,12 @@ * source to Data Source V2, with batch and micro-batch streaming reads, batch and streaming * writes, and predicate, column and limit pushdown. * + * This header is self-contained: it only depends on the C standard library, and defines the + * structs of the Arrow C data interface unless they are already defined, as by Arrow C++. A library + * can copy it into its source tree, or use the copy in the include directory of a Spark + * distribution. A library built against it works with all the releases of Spark that support its + * SPARK_DS_ABI_VERSION. + * * Spark finds the library of a data source automatically when it is installed like other native * libraries and named after the data source: the library of my_source is * spark_datasource_my_source, such as libspark_datasource_my_source.so in /usr/local/lib on Linux. From e5352e191c14545fbd5ecb449ec2ac026632835b Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Wed, 14 Oct 2026 18:33:53 +0900 Subject: [PATCH 22/27] Do not abort a data writer that committed, and run the native data source tests on Java 25 - spark_ds_commit_data_writer releases the data writer when it returns 0, but Spark then checks the message that it passed to the sink, and failed the task with the handle still set, so the task aborted the released data writer. The handle is now cleared as soon as the library returns 0, before the message is checked. - The Java 25 Maven job of a pull request also runs NativeDataSourceSuite, NativeDataSourceClusterSuite and NativeDataSourceConnectSuite: the other jobs, on Java 17, skip the tests that load a library. Co-authored-by: Isaac --- .github/workflows/build_and_test.yml | 9 +++++ .../datasources/v2/ffi/NativeLibrary.java | 33 ++++++++++++----- .../spark/sql/datasource/spark_datasource.h | 3 +- .../datasources/v2/ffi/NativeDataSource.scala | 7 ++-- .../test_native_datasource.cc | 36 +++++++++++++++++-- .../v2/ffi/NativeDataSourceSuite.scala | 21 +++++++++++ 6 files changed, 93 insertions(+), 16 deletions(-) diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index b118c9d2774f9..8261d4e3632f8 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -952,6 +952,15 @@ jobs: export MAVEN_OPTS="-Xss64m -Xmx4g -Xms4g -XX:ReservedCodeCacheSize=128m -Dorg.slf4j.simpleLogger.defaultLogLevel=WARN" export MAVEN_CLI_OPTS="--no-transfer-progress" ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws clean install + # Native data sources are called with the FFM API, which requires Java 22 or later, so the + # other jobs, on Java 17, skip the tests that load a library. The tests build their library + # with the C++ compiler of the runner. + - name: Test native data sources + run: | + export MAVEN_OPTS="-Xss64m -Xmx4g -Xms4g -XX:ReservedCodeCacheSize=128m -Dorg.slf4j.simpleLogger.defaultLogLevel=WARN" + export MAVEN_CLI_OPTS="--no-transfer-progress" + ./build/mvn $MAVEN_CLI_OPTS -pl sql/core -Dtest=none -DwildcardSuites=org.apache.spark.sql.execution.datasources.v2.ffi.NativeDataSourceSuite,org.apache.spark.sql.execution.datasources.v2.ffi.NativeDataSourceClusterSuite test + ./build/mvn $MAVEN_CLI_OPTS -pl sql/connect/server -Dtest=none -DwildcardSuites=org.apache.spark.sql.connect.NativeDataSourceConnectSuite test # Documentation build docs: diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java index 8e0f8754bac4d..3f44d5ca106b9 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -555,9 +555,19 @@ void write(long dataWriter, long arrayAddress, long schemaAddress) throws Throwa dataWriter, arrayAddress, schemaAddress), "spark_ds_write"); } - byte[] commitDataWriter(long dataWriter) throws Throwable { - byte[] message = serialize(commitDataWriter, "spark_ds_commit_data_writer", dataWriter); - return message == null ? new byte[0] : message; + /** + * Commits a data writer, whose handle the library releases when the commit succeeds. Then + * {@code released} runs right away, before the message is checked: if the message is invalid, + * the commit fails, but the handle must not be used again. + */ + byte[] commitDataWriter(long dataWriter, Runnable released) throws Throwable { + MethodHandle f = require(commitDataWriter, "spark_ds_commit_data_writer"); + try (Sink sink = new Sink()) { + check((int) f.invokeExact(dataWriter, Ffm.SINK, sink.id), "spark_ds_commit_data_writer"); + released.run(); + byte[] message = single(sink, "spark_ds_commit_data_writer"); + return message == null ? new byte[0] : message; + } } void abortDataWriter(long dataWriter) throws Throwable { @@ -573,13 +583,18 @@ private byte[] serialize(MethodHandle function, String name, long handle) throws MethodHandle f = require(function, name); try (Sink sink = new Sink()) { check((int) f.invokeExact(handle, Ffm.SINK, sink.id), name); - byte[][] values = sink.all(); - if (values.length > 1) { - throw new IllegalStateException( - name + " returned " + values.length + " values instead of one."); - } - return values.length == 0 ? null : values[0]; + return single(sink, name); + } + } + + /** Returns the single byte array that a function passed to the sink, or null if none. */ + private static byte[] single(Sink sink, String name) throws Throwable { + byte[][] values = sink.all(); + if (values.length > 1) { + throw new IllegalStateException( + name + " returned " + values.length + " values instead of one."); } + return values.length == 0 ? null : values[0]; } // --------------------------------------------------------------------------------------------- diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h index ac894a4d84401..8dffaefc0493f 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -532,7 +532,8 @@ int spark_ds_write(void* data_writer, struct ArrowArray* array, struct ArrowSche /* * Commits the data written by the task and releases the data writer handle. If it fails, the * handle is not released: Spark then calls spark_ds_abort_data_writer with it, to abort the data - * written by the task and release the handle. + * written by the task and release the handle. If it succeeds, Spark does not use the handle again, + * even if it then fails the task, for example because set_message was called more than once. * * set_message: called once with the message to pass to spark_ds_commit on the driver, or not at * all for an empty message diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index 797ca86bd8789..55212ea4d7c08 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -669,7 +669,7 @@ class NativeDataWriter( private val library = NativeLibraries.get(location) // Zero once the data writer was committed or aborted, which releases it. It stays set if the - // commit fails, so that Spark aborts the data writer. + // library fails to commit, so that Spark aborts the data writer. private var handle: Long = call(name, "write to") { library.createDataWriter(writerState, partitionId, taskId, epochId) } @@ -707,8 +707,9 @@ class NativeDataWriter( if (handle == 0) { throw SparkException.internalError("The data writer was already committed or aborted.") } - val message = call(name, "write to")(library.commitDataWriter(handle)) - handle = 0 + // The library releases the data writer once it commits, even if its message is then rejected, + // so Spark must not abort it afterwards. + val message = call(name, "write to")(library.commitDataWriter(handle, () => handle = 0)) NativeWriterCommitMessage(bytesOrEmpty(Some(message))) } diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index d1f6ae82bd187..a3fb0db870398 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -32,7 +32,9 @@ // to, although its buffers hold them all: "id" shortens the column id, // "entries" the entries of the map and their fields, and "values" the values // of the entries. -// native_sink Writes the rows as CSV files under the directory in the option "path". +// native_sink Writes the rows as CSV files under the directory in the option "path". With +// the option "commit_messages", a data writer passes its commit message that +// many times. // native_counter Streams the rows [0, max_offset) with a single column "id", making "step" // more rows available each time Spark asks for the latest offset. With the // option "bad_offset", the latest offset is "none", "multiline" JSON or @@ -63,6 +65,7 @@ #include #include #include +#include #include #include #include @@ -727,6 +730,12 @@ struct DataWriter { Options state; }; +// The data writers that a commit released, with the path of their write, so that aborting one of +// them afterwards, which the header forbids, is recorded in the file _ABORTED_RELEASED_DATA_WRITER +// rather than crashing. +std::mutex released_data_writers_lock; +[[maybe_unused]] std::map released_data_writers; + [[maybe_unused]] void makeDirs(const std::string& path) { for (size_t pos = path.find('/', 1); pos != std::string::npos; pos = path.find('/', pos + 1)) { mkdir(path.substr(0, pos).c_str(), 0755); @@ -1082,6 +1091,7 @@ int createAnyWriter(const char* function, void* handle, struct ArrowSchema* sche writer->state["overwrite"] = overwrite ? "true" : "false"; writer->state["fail"] = get(source->options, "fail", ""); writer->state["zero"] = get(source->options, "zero", ""); + writer->state["commit_messages"] = get(source->options, "commit_messages", "1"); // Records that a writer was created, which Spark then commits or aborts, and the ID of the // streaming query. makeDirs(path); @@ -1175,7 +1185,13 @@ int spark_ds_create_data_writer(const uint8_t* writer_state, int64_t writer_stat std::to_string(partition_id) + "-" + std::to_string(task_id) + ".csv"; FILE* file = std::fopen(temp_path.c_str(), "w"); if (file == nullptr) throw std::runtime_error("cannot write " + temp_path); - *out = new DataWriter{temp_path, file, state}; + auto* writer = new DataWriter{temp_path, file, state}; + { + // The allocator may reuse the address of a released data writer. + std::lock_guard guard(released_data_writers_lock); + released_data_writers.erase(writer); + } + *out = writer; }); } @@ -1206,13 +1222,27 @@ int spark_ds_commit_data_writer(void* handle, SparkDsSink set_message, void* sin // If the commit fails, the data writer stays valid: Spark aborts it with // spark_ds_abort_data_writer. maybeFail(writer->state, "commitDataWriter"); - sendBytes(set_message, sink, writer->temp_path); + // With the option "commit_messages", the message is passed that many times. + int64_t messages = getLong(writer->state, "commit_messages", 1); + for (int64_t i = 0; i < messages; i++) sendBytes(set_message, sink, writer->temp_path); std::fclose(writer->file); + { + std::lock_guard guard(released_data_writers_lock); + released_data_writers[handle] = writer->state["path"]; + } delete writer; }); } int spark_ds_abort_data_writer(void* handle) { + { + std::lock_guard guard(released_data_writers_lock); + auto released = released_data_writers.find(handle); + if (released != released_data_writers.end()) { + writeFile(released->second + "/_ABORTED_RELEASED_DATA_WRITER", ""); + return 0; + } + } std::unique_ptr writer(fromHandle(handle)); std::fclose(writer->file); unlink(writer->temp_path.c_str()); diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 0797fe1212254..b544b43dd04fa 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -697,6 +697,27 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) } + nativeTest("do not abort a data writer that committed with an invalid message") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + // The library commits and releases the data writer, but passes two commit messages. + val e = intercept[Exception] { + session.range(10).repartition(1).write.format("native_sink").option("path", dir.getPath) + .option("commit_messages", "2").mode("append").save() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "write to", + "name" -> "native_sink", + "msg" -> "spark_ds_commit_data_writer returned 2 values instead of one.")) + // The task failed without aborting the released data writer, and the driver aborted the write. + assert(!new File(dir, "_ABORTED_RELEASED_DATA_WRITER").exists()) + assert(new File(dir, "_ABORTED").isFile) + assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) + } + nativeTest("abort a micro-batch whose write failed") { val session = newSession(defaultPackage) val dir = Utils.createTempDir() From d2df236bd234f5a1188f4c98249029ef0cc4a5dc Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Thu, 15 Oct 2026 08:07:42 +0900 Subject: [PATCH 23/27] Give the C interface a function table, structs from Spark, and capabilities Follows the review of the shape of version 1, so that it can grow without new exported functions or signatures: - A library exports only spark_ds_init, which gets SparkDsHost (the newest ABI version that Spark supports, the version of Spark, and a log function) and returns its function table, SparkDsLibrary. A library leaves the functions it does not implement NULL. - The table and the structs that Spark passes start with their size, and only get fields appended within an ABI version: Spark only reads the functions that fit in the size of a table, and a library only reads the fields that fit in the size of a struct. - read and create_data_writer get SparkDsTaskContext, with the IDs of the task and is_interrupted, valid until the stream is released or the data writer is closed. create_writer, which replaces create_writer and create_stream_writer, gets SparkDsWriteInfo, with the write mode (SPARK_DS_WRITE_APPEND or SPARK_DS_WRITE_TRUNCATE) and the query ID. - create returns the capabilities of each data source. Spark reports an operation that a data source does not declare as unsupported without calling the library, only passes SPARK_DS_WRITE_TRUNCATE with SPARK_DS_TRUNCATE, and fails on the driver if a declared capability needs a NULL function. - push_predicates answers SPARK_DS_PREDICATE_UNUSED, PARTIAL or COMPLETE for each predicate. Spark shows the used predicates as pushed, and still evaluates all of them. - Each handle is released by the close function of its kind only: commit_data_writer no longer releases the data writer, and Spark calls the new close_data_writer after the commit or abort. The docs and their Rust and C++ examples use the table. New tests check the layout of the header against the bindings, the capabilities, a table that ends early, the host, the task context, a read that only ends once is_interrupted is true, and libraries that spark_ds_init does not initialize. Co-authored-by: Isaac --- docs/sql-native-data-sources.md | 349 ++++--- .../datasources/v2/ffi/NativeLibrary.java | 731 ++++++++++----- .../spark/sql/datasource/predicate-cases.json | 2 +- .../spark/sql/datasource/spark_datasource.h | 871 ++++++++++-------- .../v2/columnar/ColumnarDataSource.scala | 7 +- .../v2/columnar/ColumnarTable.scala | 3 +- .../datasources/v2/ffi/NativeDataSource.scala | 210 +++-- .../v2/ffi/NativeDataSourcePackage.scala | 24 +- .../datasources/v2/ffi/NativePredicates.scala | 4 +- .../test_native_datasource.cc | 434 ++++++--- .../v2/columnar/ColumnarDataSourceSuite.scala | 2 +- .../v2/ffi/NativeDataSourceSuite.scala | 347 +++++-- .../v2/ffi/NativeDataSourceTestUtils.scala | 39 +- 13 files changed, 2002 insertions(+), 1021 deletions(-) diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 82d63c0202e53..4f07d96a283f8 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -25,7 +25,7 @@ license: | ## Overview A native data source is a data source implemented in native code, such as Rust or C++, without -any JVM code. It is a shared library that exports the C functions declared in the header +any JVM code. It is a shared library that implements the C interface declared in the header [`spark_datasource.h`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h). Spark finds the library automatically when it is installed like other native libraries, or when a session adds it in a package, and loads it when a query uses the data source. Spark adapts it to @@ -47,10 +47,11 @@ Java, a query that uses a native data source fails with ## How It Works -A native data source library exports the functions of `spark_datasource.h`, such as -`spark_ds_create`, as plain C functions: it needs no JVM headers or bindings. Spark looks the -functions up in each library separately, so any number of native data sources can be loaded side -by side, even though they export the same functions. +A native data source library exports a single C function, `spark_ds_init`, which returns a table +of function pointers, a `SparkDsLibrary`, that implements the data source: it needs no JVM headers +or bindings. Each library has its own table, so any number of native data sources can be loaded +side by side. When Spark loads a library, it calls `spark_ds_init` with a `SparkDsHost`, which +gives the library the version of Spark and a function to log messages to the log of Spark. A Spark distribution contains the header in its `include` directory, as `$SPARK_HOME/include/spark_datasource.h`, and the `spark-sql` JAR contains it as @@ -65,34 +66,35 @@ curl -LO {{site.SPARK_GITHUB_URL}}/raw/v{{site.SPARK_VERSION_SHORT}}/sql/core/sr A library built against the header works with all the releases of Spark that support its version, `SPARK_DS_ABI_VERSION`: see [Compatibility](#compatibility). -Spark calls the functions of the library as follows: +Spark calls the functions of the table as follows: - **Data sources**: for each schema inference, scan and write, Spark creates a data source on the - driver with `spark_ds_create`, which receives the options of the user, and closes it right after - with `spark_ds_close`. When the user does not specify a schema, Spark asks for it with - `spark_ds_schema`. -- **Batch reads**: on the driver, Spark creates a reader with `spark_ds_create_reader`, pushes - operations down with `spark_ds_push_predicates`, `spark_ds_push_limit` and - `spark_ds_prune_columns`, plans the partitions with `spark_ds_partitions`, and serializes the - state of the reader with `spark_ds_serialize_reader`. Then it calls `spark_ds_read` on the - executors, in one task for each partition, and `spark_ds_read` exports an `ArrowArrayStream`. -- **Streaming reads**: on the driver, Spark creates a stream reader with - `spark_ds_create_stream_reader`. For each micro-batch, it asks for the `spark_ds_latest_offset`, - plans the partitions with `spark_ds_stream_partitions`, and calls `spark_ds_read` on the - executors. Offsets are defined by the library, as compact JSON on a single line: Spark stores - each offset as a line of the offset log of the checkpoint, and rejects an offset that is not - valid JSON or contains a line break, such as pretty-printed JSON. + driver with `create`, which receives the options of the user and returns the capabilities of the + data source, such as `SPARK_DS_BATCH_READ` or `SPARK_DS_STREAMING_WRITE`, and closes it right + after with `close`. Spark reports the operations that a data source does not declare as + unsupported. When the user does not specify a schema, Spark asks for it with `schema`. +- **Batch reads**: on the driver, Spark creates a reader with `create_reader`, pushes operations + down with `push_predicates`, `push_limit` and `prune_columns`, plans the partitions with + `partitions`, and serializes the state of the reader with `serialize_reader`. Then it calls + `read` on the executors, in one task for each partition, and `read` exports an + `ArrowArrayStream`. +- **Streaming reads**: on the driver, Spark creates a stream reader with `create_stream_reader`. + For each micro-batch, it asks for the `latest_offset`, plans the partitions with + `stream_partitions`, and calls `read` on the executors. Offsets are defined by the library, as + compact JSON on a single line: Spark stores each offset as a line of the offset log of the + checkpoint, and rejects an offset that is not valid JSON or contains a line break, such as + pretty-printed JSON. - **Writes**: on the driver, right before the tasks of the write (or of a micro-batch) start, Spark - creates a writer with `spark_ds_create_writer` or `spark_ds_create_stream_writer`. With adaptive - query execution, it first runs the shuffle stages of the data to write, and creates no writer if - they fail, but the last stage runs in the tasks of the write, so a writer can still be aborted - because computing the data failed: do irreversible work, such as removing the existing data for - an overwrite, in `spark_ds_commit`. On the executors, each task creates a data writer with - `spark_ds_create_data_writer`, writes the batches of its data with `spark_ds_write`, and calls - `spark_ds_commit_data_writer`, or `spark_ds_abort_data_writer` if the task or the commit fails. - Then the driver calls `spark_ds_commit` with the messages of the tasks, or `spark_ds_abort` if - the write failed. A streaming write commits each micro-batch with its epoch ID, and - `spark_ds_create_stream_writer` gets the ID of the query, which stays the same when the query + creates a writer with `create_writer`, which gets a `SparkDsWriteInfo` with the mode of the write, + `SPARK_DS_WRITE_APPEND` or `SPARK_DS_WRITE_TRUNCATE`, and the ID of the query of a streaming + write. With adaptive query execution, it first runs the shuffle stages of the data to write, and + creates no writer if they fail, but the last stage runs in the tasks of the write, so a writer can + still be aborted because computing the data failed: do irreversible work, such as removing the + existing data for an overwrite, in `commit`. On the executors, each task creates a data writer + with `create_data_writer`, writes the batches of its data with `write`, and calls + `commit_data_writer`, or `abort_data_writer` if the task or the commit fails. Then the driver + calls `commit` with the messages of the tasks, or `abort` if the write failed. A streaming write + commits each micro-batch with its epoch ID, and the ID of the query stays the same when the query restarts from its checkpoint: together, they identify the output of a micro-batch, so that a library can skip an epoch that it already committed. A micro-batch streaming query creates a new stream writer for each micro-batch, and Spark closes the previous one when it creates the next @@ -103,31 +105,37 @@ The library defines the format of the partitions, of the state of its readers an the commit messages, which Spark sends to the executors and back as byte arrays. The library returns them, and the offsets, by passing them to a `SparkDsSink` callback, which copies them. A function reports an error by returning a non-zero value, with the message returned by -`spark_ds_last_error`, and Spark reports it as a `NATIVE_DATA_SOURCE_ERROR`. A C++ exception or a -Rust panic must not cross the interface. The close functions cannot fail, and they can run after -the query, so do the work that can fail, such as flushing data, in `spark_ds_commit`. Strings are -UTF-8, and the JSON of the predicates is ASCII, with the other characters escaped. +`last_error`, and Spark reports it as a `NATIVE_DATA_SOURCE_ERROR`. A C++ exception or a Rust panic +must not cross the interface. Each handle that a `create` function returns is released exactly +once, by the `close` function of its kind. The close functions cannot fail, and they can run after +the query, so do the work that can fail, such as flushing data, in `commit`. Strings are UTF-8, and +the JSON of the predicates is ASCII, with the other characters escaped. -A library only exports the functions of the operations it supports. If a function is missing, -Spark treats the operation as unsupported, except for the optional functions: +A library only sets the functions of the operations it supports, and leaves the others `NULL`. A +data source must not declare a capability whose functions are `NULL`: -| Operation | Required functions | Optional functions | +| Capability | Required functions | Optional functions | |-----------|--------------------|--------------------| -| All | `spark_ds_abi_version`, `spark_ds_create`, `spark_ds_close` | `spark_ds_last_error`, `spark_ds_schema` | -| Batch read | `spark_ds_create_reader`, `spark_ds_partitions`, `spark_ds_read`, `spark_ds_close_reader` | `spark_ds_push_predicates`, `spark_ds_push_limit`, `spark_ds_prune_columns`, `spark_ds_serialize_reader` | -| Streaming read | `spark_ds_create_stream_reader`, `spark_ds_initial_offset`, `spark_ds_latest_offset`, `spark_ds_stream_partitions`, `spark_ds_read`, `spark_ds_close_stream_reader` | `spark_ds_serialize_stream_reader`, `spark_ds_commit_offset` | -| Batch write | `spark_ds_create_writer`, `spark_ds_commit`, `spark_ds_abort`, `spark_ds_close_writer`, `spark_ds_create_data_writer`, `spark_ds_write`, `spark_ds_commit_data_writer`, `spark_ds_abort_data_writer` | `spark_ds_serialize_writer` | -| Streaming write | `spark_ds_create_stream_writer`, `spark_ds_commit`, `spark_ds_abort`, `spark_ds_close_writer`, `spark_ds_create_data_writer`, `spark_ds_write`, `spark_ds_commit_data_writer`, `spark_ds_abort_data_writer` | `spark_ds_serialize_writer` | +| All | `create`, `close` | `last_error`, `schema` | +| `SPARK_DS_BATCH_READ` | `create_reader`, `partitions`, `read`, `close_reader` | `push_predicates`, `push_limit`, `prune_columns`, `serialize_reader` | +| `SPARK_DS_MICRO_BATCH_READ` | `create_stream_reader`, `initial_offset`, `latest_offset`, `stream_partitions`, `read`, `close_stream_reader` | `serialize_stream_reader`, `commit_offset` | +| `SPARK_DS_BATCH_WRITE`, `SPARK_DS_STREAMING_WRITE` | `create_writer`, `commit`, `abort`, `close_writer`, `create_data_writer`, `write`, `commit_data_writer`, `abort_data_writer`, `close_data_writer` | `serialize_writer` | + +`SPARK_DS_TRUNCATE` declares that a writer supports `SPARK_DS_WRITE_TRUNCATE`: the save mode +`overwrite` of a batch write, and the output mode `complete` of a streaming write. The comments of [`spark_datasource.h`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h) specify each function, including the ownership of the Arrow structs and the JSON format of the -pushed predicates. Spark evaluates all the pushed predicates again after reading, so a library can -use them to skip data, for example files or row groups whose statistics show that no row matches, -but it must never skip a row for which the predicates are true under the semantics of Spark SQL: -NaN values are equal to each other and greater than any other value, `-0.0` is equal to `0.0`, -comparisons with null are null, and strings compare by their UTF-8 bytes. The comparisons of IEEE -754 and the total orders of Arrow libraries, such as the comparison kernels of arrow-rs, differ for -some of these values. The jar of Spark SQL contains test cases for a library, +pushed predicates. `push_predicates` tells Spark how the reader uses each predicate: not at all, +only to skip data (`SPARK_DS_PREDICATE_PARTIAL`), or to only return the rows for which it is true +(`SPARK_DS_PREDICATE_COMPLETE`). Spark shows the predicates that the reader uses as pushed in the +plan of the query, and evaluates all the pushed predicates again after reading either way, so a +library can use them to skip data, for example files or row groups whose statistics show that no +row matches, but it must never skip a row for which the predicates are true under the semantics of +Spark SQL: NaN values are equal to each other and greater than any other value, `-0.0` is equal to +`0.0`, comparisons with null are null, and strings compare by their UTF-8 bytes. The comparisons of +IEEE 754 and the total orders of Arrow libraries, such as the comparison kernels of arrow-rs, +differ for some of these values. The jar of Spark SQL contains test cases for a library, `org/apache/spark/sql/datasource/predicate-cases.json`: predicates as Spark pushes them, with the rows of a table for which they are true: @@ -141,28 +149,39 @@ has at least as many values as its parent refers to, but it cannot check that th arrays hold those values, because the C data interface does not carry the sizes of buffers: a library that exports a buffer that is too short makes Spark read past its end. -The `get_next` function of the stream that `spark_ds_read` exports may block, for example on -network I/O, but Spark cannot interrupt native code: when a task is killed, Spark stops reading -only once `get_next` returns. Use timeouts for blocking operations, or enable the task reaper with -`spark.task.reaper.enabled` and `spark.task.reaper.killTimeout`, which kills an executor whose -killed task does not end. +`read` and `create_data_writer` get a `SparkDsTaskContext`, with the IDs of the task, and its +partition and attempt, and an `is_interrupted` function that returns whether the task was killed, +for example because its query was canceled. Spark cannot interrupt native code: when a task is +killed, Spark stops reading only once the `get_next` function of the stream returns. So a read or +a write that runs for long, or blocks, for example on network I/O, should check `is_interrupted` +regularly and use timeouts. The task reaper, `spark.task.reaper.enabled` and +`spark.task.reaper.killTimeout`, kills an executor whose killed task does not end. ### Compatibility -Spark binds the functions of a library by their names, so the names and the signatures of the -functions of `spark_datasource.h` are the interface. Its version, 1, is `SPARK_DS_ABI_VERSION`, -which `spark_ds_abi_version` returns and the manifest of a package declares. Within a version, the -signature and the semantics of an existing function never change, and new functions are only -added as optional functions, so a library keeps working with later releases of Spark. New values -may appear where the header says so: for example, later releases may push down predicates with -other names, which a library must ignore if it does not know them. An incompatible change comes -with a new version, with new names. Spark supports version 1 at least in all the releases of Spark -4 from 4.4.0. +The function table, `SparkDsLibrary`, and the structs that Spark passes to the library, +`SparkDsHost`, `SparkDsTaskContext` and `SparkDsWriteInfo`, start with their size, so that they can +grow. The version of the interface, 1, is `SPARK_DS_ABI_VERSION`, which the table and the manifest +of a package declare. Within a version: + +- fields are only appended to these structs, and the existing fields and the semantics of the + existing functions never change; +- Spark only reads the functions that fit in the size of the table of a library, so a library built + against an earlier header keeps working with later releases of Spark, with the functions that it + does not know as `NULL`. A library must only read the fields that fit in the size of a struct + that Spark passes, so that it also runs on earlier releases of Spark that support its version; +- new values may appear where the header says so: for example, later releases may push down + predicates with other names, which a library must ignore if it does not know them. Spark only + uses a new capability or write mode with a data source that declares it. + +An incompatible change comes with a new version: `spark_ds_init` gets the newest version that Spark +supports, in `SparkDsHost`, and returns a table of a version that the library implements. Spark +supports version 1 at least in all the releases of Spark 4 from 4.4.0. ### Limitations Version 1 of the interface does not cover these features of Data Source V2 yet. They can be added -later as optional functions: +later as functions or fields appended to the structs of the interface: - statistics: Spark estimates the size of a native scan with `spark.sql.defaultSizeInBytes`, so it does not plan a native scan as the broadcast side of a join, unless the query has a broadcast @@ -172,7 +191,6 @@ later as optional functions: - admission control for streaming reads, such as a maximum number of rows for each micro-batch, and continuous streaming; - the pushdown of aggregates, of the sort order and of the partitioning of the data; -- a function to cancel a read in progress; - memory accounting: see below. ### Memory @@ -202,8 +220,9 @@ like other native libraries is found automatically, without any configuration. T named after the data source, in lower case, with the prefix `spark_datasource_`: the library of `rust_range` is `libspark_datasource_rust_range.so` on Linux, `libspark_datasource_rust_range.dylib` on macOS and `spark_datasource_rust_range.dll` on Windows. -Spark passes the name of the data source to `spark_ds_create`. A library that implements several -data sources is installed under each of their names, for example with symbolic links. +Spark passes the name of the data source to the `create` function of the library. A library that +implements several data sources is installed under each of their names, for example with symbolic +links. Spark looks for the library in the native library path, which consists of, in order: @@ -260,7 +279,7 @@ is a zip file with the extension `.sparkpkg`, which contains a manifest named - `abiVersion` is the version of `spark_datasource.h` that the library implements, currently 1. - `dataSources` lists the names of the data sources that the library implements. The names are - case-insensitive, and Spark passes the name to `spark_ds_create`. + case-insensitive, and Spark passes the name to `create`. - `libraries` maps each platform, `-` with an `os` of `linux`, `osx` or `windows` and an `arch` of `x86_64` or `aarch64`, to the path of the library in the package. @@ -288,9 +307,12 @@ installed library provides it. This data source reads the numbers `[0, end)` as a column `id`, in two partitions. It only depends on the `arrow-array` and `arrow-schema` crates of [arrow-rs](https://github.com/apache/arrow-rs), -which export schemas and streams with `FFI_ArrowSchema` and `FFI_ArrowArrayStream`, and declares -the functions of `spark_datasource.h` that it implements as `extern "C"` functions. `Cargo.toml`, -where the name of the library is the name that Spark finds when it is installed: +which export schemas and streams with `FFI_ArrowSchema` and `FFI_ArrowArrayStream`. It declares +`SparkDsLibrary` as a `#[repr(C)]` struct up to `read`, the last function that it implements, since +Spark only reads the functions that fit in the size of the table, and leaves the functions that it +does not implement `None`. The fields must follow the order of `spark_datasource.h`, which +[bindgen](https://github.com/rust-lang/rust-bindgen) can also generate from the header. +`Cargo.toml`, where the name of the library is the name that Spark finds when it is installed: ```toml [package] @@ -327,7 +349,93 @@ use arrow_schema::{DataType, Field, Schema, SchemaRef}; type Result = std::result::Result>; /// `SparkDsSink`: receives a byte array returned to Spark, and copies it. -type Sink = extern "C" fn(sink: *mut c_void, data: *const u8, length: i64); +type Sink = unsafe extern "C" fn(sink: *mut c_void, data: *const u8, length: i64); + +/// A function of `SparkDsLibrary` that this data source does not implement, left NULL. +type Unused = Option; + +/// `SparkDsLibrary` of `spark_datasource.h`, up to `read`, the last function that this data +/// source implements: Spark only reads the functions that fit in `size`. +#[repr(C)] +pub struct SparkDsLibrary { + size: i64, + abi_version: i32, + last_error: unsafe extern "C" fn() -> *const c_char, + create: unsafe extern "C" fn( + name: *const c_char, + option_keys: *const *const c_char, + option_values: *const *const c_char, + num_options: i64, + capabilities: *mut u64, + source: *mut *mut c_void, + ) -> c_int, + schema: unsafe extern "C" fn(source: *mut c_void, out: *mut FFI_ArrowSchema) -> c_int, + create_reader: unsafe extern "C" fn( + source: *mut c_void, + schema: *mut FFI_ArrowSchema, + reader: *mut *mut c_void, + ) -> c_int, + create_stream_reader: Unused, + create_writer: Unused, + close: unsafe extern "C" fn(source: *mut c_void), + push_predicates: Unused, + push_limit: Unused, + prune_columns: Unused, + partitions: unsafe extern "C" fn(reader: *mut c_void, add: Sink, sink: *mut c_void) -> c_int, + serialize_reader: Unused, + close_reader: unsafe extern "C" fn(reader: *mut c_void), + initial_offset: Unused, + latest_offset: Unused, + stream_partitions: Unused, + serialize_stream_reader: Unused, + commit_offset: Unused, + close_stream_reader: Unused, + read: unsafe extern "C" fn( + task: *const c_void, + reader_state: *const u8, + reader_state_length: i64, + partition: *const u8, + partition_length: i64, + out: *mut FFI_ArrowArrayStream, + ) -> c_int, +} + +const SPARK_DS_BATCH_READ: u64 = 1; + +static LIBRARY: SparkDsLibrary = SparkDsLibrary { + size: std::mem::size_of::() as i64, + abi_version: 1, + last_error, + create, + schema, + create_reader, + create_stream_reader: None, + create_writer: None, + close, + push_predicates: None, + push_limit: None, + prune_columns: None, + partitions, + serialize_reader: None, + close_reader, + initial_offset: None, + latest_offset: None, + stream_partitions: None, + serialize_stream_reader: None, + commit_offset: None, + close_stream_reader: None, + read, +}; + +/// The function that Spark calls when it loads the library. +#[no_mangle] +pub unsafe extern "C" fn spark_ds_init( + _host: *const c_void, + library: *mut *const SparkDsLibrary, +) -> c_int { + *library = &LIBRARY; + 0 +} struct RangeSource { end: i64, @@ -341,7 +449,7 @@ thread_local! { static LAST_ERROR: RefCell = RefCell::new(CString::default()); } -fn schema() -> SchemaRef { +fn arrow_schema() -> SchemaRef { Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])) } @@ -356,22 +464,16 @@ fn call(f: impl FnOnce() -> Result<()>) -> c_int { 1 } -#[no_mangle] -pub extern "C" fn spark_ds_last_error() -> *const c_char { +unsafe extern "C" fn last_error() -> *const c_char { LAST_ERROR.with(|error| error.borrow().as_ptr()) } -#[no_mangle] -pub extern "C" fn spark_ds_abi_version() -> i32 { - 1 -} - -#[no_mangle] -pub unsafe extern "C" fn spark_ds_create( +unsafe extern "C" fn create( _name: *const c_char, option_keys: *const *const c_char, option_values: *const *const c_char, num_options: i64, + capabilities: *mut u64, source: *mut *mut c_void, ) -> c_int { call(|| { @@ -381,21 +483,20 @@ pub unsafe extern "C" fn spark_ds_create( range.end = CStr::from_ptr(*option_values.add(i)).to_str()?.parse()?; } } + *capabilities = SPARK_DS_BATCH_READ; *source = Box::into_raw(Box::new(range)) as *mut c_void; Ok(()) }) } -#[no_mangle] -pub unsafe extern "C" fn spark_ds_schema(_source: *mut c_void, out: *mut FFI_ArrowSchema) -> c_int { +unsafe extern "C" fn schema(_source: *mut c_void, out: *mut FFI_ArrowSchema) -> c_int { call(|| { - std::ptr::write(out, FFI_ArrowSchema::try_from(schema().as_ref())?); + std::ptr::write(out, FFI_ArrowSchema::try_from(arrow_schema().as_ref())?); Ok(()) }) } -#[no_mangle] -pub unsafe extern "C" fn spark_ds_create_reader( +unsafe extern "C" fn create_reader( source: *mut c_void, _schema: *mut FFI_ArrowSchema, reader: *mut *mut c_void, @@ -407,34 +508,27 @@ pub unsafe extern "C" fn spark_ds_create_reader( }) } -#[no_mangle] -pub unsafe extern "C" fn spark_ds_close(source: *mut c_void) { +unsafe extern "C" fn close(source: *mut c_void) { drop(Box::from_raw(source as *mut RangeSource)); } -#[no_mangle] -pub unsafe extern "C" fn spark_ds_partitions( - reader: *mut c_void, - add_partition: Sink, - sink: *mut c_void, -) -> c_int { +unsafe extern "C" fn partitions(reader: *mut c_void, add: Sink, sink: *mut c_void) -> c_int { call(|| { let end = (*(reader as *const RangeReader)).end; for (start, end) in [(0, end / 2), (end / 2, end)] { let bytes = [start.to_le_bytes(), end.to_le_bytes()].concat(); - add_partition(sink, bytes.as_ptr(), bytes.len() as i64); + add(sink, bytes.as_ptr(), bytes.len() as i64); } Ok(()) }) } -#[no_mangle] -pub unsafe extern "C" fn spark_ds_close_reader(reader: *mut c_void) { +unsafe extern "C" fn close_reader(reader: *mut c_void) { drop(Box::from_raw(reader as *mut RangeReader)); } -#[no_mangle] -pub unsafe extern "C" fn spark_ds_read( +unsafe extern "C" fn read( + _task: *const c_void, _reader_state: *const u8, _reader_state_length: i64, partition: *const u8, @@ -446,8 +540,8 @@ pub unsafe extern "C" fn spark_ds_read( let start = i64::from_le_bytes(bytes[0..8].try_into()?); let end = i64::from_le_bytes(bytes[8..16].try_into()?); let ids = Int64Array::from_iter_values(start..end); - let batch = RecordBatch::try_new(schema(), vec![Arc::new(ids)])?; - let batches = RecordBatchIterator::new(vec![Ok(batch)], schema()); + let batch = RecordBatch::try_new(arrow_schema(), vec![Arc::new(ids)])?; + let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema()); std::ptr::write(out, FFI_ArrowArrayStream::new(Box::new(batches))); Ok(()) }) @@ -487,7 +581,8 @@ EOF The same data source in C++, with [Arrow C++](https://arrow.apache.org/docs/cpp/), which exports schemas and streams with `arrow::ExportSchema` and `arrow::ExportRecordBatchReader`. It includes -`spark_datasource.h`, here from the `include` directory of a Spark distribution. +`spark_datasource.h`, here from the `include` directory of a Spark distribution, and fills +`SparkDsLibrary` with designated initializers, which leave the other functions `NULL`. `cpp_range.cc`: @@ -542,36 +637,31 @@ void Check(const arrow::Status& status) { if (!status.ok()) throw std::runtime_error(status.ToString()); } -} // namespace - -extern "C" { +const char* LastError() { return last_error.c_str(); } -const char* spark_ds_last_error(void) { return last_error.c_str(); } - -int32_t spark_ds_abi_version(void) { return SPARK_DS_ABI_VERSION; } - -int spark_ds_create(const char*, const char* const* option_keys, - const char* const* option_values, int64_t num_options, void** source) { +int Create(const char*, const char* const* option_keys, const char* const* option_values, + int64_t num_options, uint64_t* capabilities, void** source) { return Call([&] { auto range = std::make_unique(); for (int64_t i = 0; i < num_options; i++) { if (std::string(option_keys[i]) == "end") range->end = std::stoll(option_values[i]); } + *capabilities = SPARK_DS_BATCH_READ; *source = range.release(); }); } -int spark_ds_schema(void*, struct ArrowSchema* out) { +int GetSchema(void*, struct ArrowSchema* out) { return Call([&] { Check(arrow::ExportSchema(*Schema(), out)); }); } -int spark_ds_create_reader(void* source, struct ArrowSchema*, void** reader) { +int CreateReader(void* source, struct ArrowSchema*, void** reader) { return Call([&] { *reader = new RangeReader{static_cast(source)->end}; }); } -void spark_ds_close(void* source) { delete static_cast(source); } +void Close(void* source) { delete static_cast(source); } -int spark_ds_partitions(void* reader, SparkDsSink add_partition, void* sink) { +int Partitions(void* reader, SparkDsSink add_partition, void* sink) { return Call([&] { int64_t end = static_cast(reader)->end; int64_t bounds[2][2] = {{0, end / 2}, {end / 2, end}}; @@ -581,10 +671,10 @@ int spark_ds_partitions(void* reader, SparkDsSink add_partition, void* sink) { }); } -void spark_ds_close_reader(void* reader) { delete static_cast(reader); } +void CloseReader(void* reader) { delete static_cast(reader); } -int spark_ds_read(const uint8_t*, int64_t, const uint8_t* partition, int64_t, - struct ArrowArrayStream* out) { +int Read(const SparkDsTaskContext*, const uint8_t*, int64_t, const uint8_t* partition, int64_t, + struct ArrowArrayStream* out) { return Call([&] { int64_t bounds[2]; std::memcpy(bounds, partition, sizeof(bounds)); @@ -598,7 +688,26 @@ int spark_ds_read(const uint8_t*, int64_t, const uint8_t* partition, int64_t, }); } -} // extern "C" +// The functions of the data source. The ones that it does not implement are NULL. +const SparkDsLibrary kLibrary = { + .size = sizeof(SparkDsLibrary), + .abi_version = SPARK_DS_ABI_VERSION, + .last_error = LastError, + .create = Create, + .schema = GetSchema, + .create_reader = CreateReader, + .close = Close, + .partitions = Partitions, + .close_reader = CloseReader, + .read = Read, +}; + +} // namespace + +extern "C" SPARK_DS_EXPORT int spark_ds_init(const SparkDsHost*, const SparkDsLibrary** library) { + *library = &kLibrary; + return 0; +} ``` {% endraw %} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java index 3f44d5ca106b9..6c2368ea7c224 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -24,11 +24,18 @@ import java.nio.charset.StandardCharsets; import java.nio.file.Path; import java.util.ArrayList; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicLong; +import java.util.function.BooleanSupplier; +import org.apache.spark.internal.LogKeys; +import org.apache.spark.internal.MDC; +import org.apache.spark.internal.SparkLogger; +import org.apache.spark.internal.SparkLoggerFactory; import org.apache.spark.unsafe.Platform; /** @@ -37,18 +44,21 @@ * and Memory (FFM) API of Java 22 and later. *

    * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached - * through method handles, looked up once when the first library is loaded. Each function of a - * library is bound to a method handle that takes and returns primitive values only, with - * pointers as {@code long} addresses, so the calls themselves use no reflection. The memory that - * Spark passes to the library is allocated with {@link Platform}. + * through method handles, looked up once when the first library is loaded. The library exports + * {@code spark_ds_init}, which returns its table of functions, {@code SparkDsLibrary}. Each + * function of the table is bound to a method handle that takes and returns primitive values only, + * with pointers as {@code long} addresses, so the calls themselves use no reflection. The memory + * that Spark passes to the library is allocated with {@link Platform}, at the offsets of the + * structs of the header, which {@link #layout()} lists for the tests. *

    - * Each library is looked up separately, so all the libraries can export the same functions. - * Calling a function that the library does not export throws {@link MissingFunctionException}. A - * function that fails throws a {@link RuntimeException} with the message of - * {@code spark_ds_last_error}. + * Calling a function that the library leaves NULL in its table throws + * {@link MissingFunctionException}. A function that fails throws a {@link RuntimeException} with + * the message of {@code last_error}. */ final class NativeLibrary { + private static final SparkLogger LOGGER = SparkLoggerFactory.getLogger(NativeLibrary.class); + /** The version of the interface that Spark supports, {@code SPARK_DS_ABI_VERSION}. */ static final int ABI_VERSION = 1; @@ -60,86 +70,158 @@ static boolean isSupported() { return Runtime.version().feature() >= MINIMUM_JAVA_VERSION; } + // The capabilities of a data source, SPARK_DS_BATCH_READ and the others. + static final long BATCH_READ = 1; + static final long MICRO_BATCH_READ = 2; + static final long BATCH_WRITE = 4; + static final long STREAMING_WRITE = 8; + static final long TRUNCATE = 16; + + // The modes of a write, SPARK_DS_WRITE_APPEND and SPARK_DS_WRITE_TRUNCATE. + static final int WRITE_APPEND = 0; + static final int WRITE_TRUNCATE = 1; + + /** A predicate that the reader does not use, SPARK_DS_PREDICATE_UNUSED. */ + static final int PREDICATE_UNUSED = 0; + // The C types of the parameters and the results of the functions. - private enum C { VOID, BOOL, INT32, INT64, POINTER } - - private final MethodHandle lastError; - private final MethodHandle abiVersion; - private final MethodHandle create; - private final MethodHandle schema; - private final MethodHandle createReader; - private final MethodHandle createStreamReader; - private final MethodHandle createWriter; - private final MethodHandle createStreamWriter; - private final MethodHandle close; - private final MethodHandle pushPredicates; - private final MethodHandle pushLimit; - private final MethodHandle pruneColumns; - private final MethodHandle partitions; - private final MethodHandle serializeReader; - private final MethodHandle closeReader; - private final MethodHandle initialOffset; - private final MethodHandle latestOffset; - private final MethodHandle streamPartitions; - private final MethodHandle serializeStreamReader; - private final MethodHandle commitOffset; - private final MethodHandle closeStreamReader; - private final MethodHandle read; - private final MethodHandle serializeWriter; - private final MethodHandle commit; - private final MethodHandle abort; - private final MethodHandle closeWriter; - private final MethodHandle createDataWriter; - private final MethodHandle write; - private final MethodHandle commitDataWriter; - private final MethodHandle abortDataWriter; - - private NativeLibrary(Object lookup) throws Throwable { - C i32 = C.INT32; - C i64 = C.INT64; - C ptr = C.POINTER; - lastError = bind(lookup, "spark_ds_last_error", ptr); - abiVersion = bind(lookup, "spark_ds_abi_version", i32); - create = bind(lookup, "spark_ds_create", i32, ptr, ptr, ptr, i64, ptr); - schema = bind(lookup, "spark_ds_schema", i32, ptr, ptr); - createReader = bind(lookup, "spark_ds_create_reader", i32, ptr, ptr, ptr); - createStreamReader = bind(lookup, "spark_ds_create_stream_reader", i32, ptr, ptr, ptr); - createWriter = bind(lookup, "spark_ds_create_writer", i32, ptr, ptr, C.BOOL, ptr); - createStreamWriter = - bind(lookup, "spark_ds_create_stream_writer", i32, ptr, ptr, C.BOOL, ptr, ptr); - close = bind(lookup, "spark_ds_close", C.VOID, ptr); - pushPredicates = bind(lookup, "spark_ds_push_predicates", i32, ptr, ptr, i64, ptr); - pushLimit = bind(lookup, "spark_ds_push_limit", i32, ptr, i32, ptr); - pruneColumns = bind(lookup, "spark_ds_prune_columns", i32, ptr, ptr, i64, ptr); - partitions = bind(lookup, "spark_ds_partitions", i32, ptr, ptr, ptr); - serializeReader = bind(lookup, "spark_ds_serialize_reader", i32, ptr, ptr, ptr); - closeReader = bind(lookup, "spark_ds_close_reader", C.VOID, ptr); - initialOffset = bind(lookup, "spark_ds_initial_offset", i32, ptr, ptr, ptr); - latestOffset = bind(lookup, "spark_ds_latest_offset", i32, ptr, ptr, ptr); - streamPartitions = - bind(lookup, "spark_ds_stream_partitions", i32, ptr, ptr, ptr, ptr, ptr); - serializeStreamReader = - bind(lookup, "spark_ds_serialize_stream_reader", i32, ptr, ptr, ptr); - commitOffset = bind(lookup, "spark_ds_commit_offset", i32, ptr, ptr); - closeStreamReader = bind(lookup, "spark_ds_close_stream_reader", C.VOID, ptr); - read = bind(lookup, "spark_ds_read", i32, ptr, i64, ptr, i64, ptr); - serializeWriter = bind(lookup, "spark_ds_serialize_writer", i32, ptr, ptr, ptr); - commit = bind(lookup, "spark_ds_commit", i32, ptr, i64, ptr, ptr, i64); - abort = bind(lookup, "spark_ds_abort", i32, ptr, i64, ptr, ptr, i64); - closeWriter = bind(lookup, "spark_ds_close_writer", C.VOID, ptr); - createDataWriter = - bind(lookup, "spark_ds_create_data_writer", i32, ptr, i64, i32, i64, i64, ptr); - write = bind(lookup, "spark_ds_write", i32, ptr, ptr, ptr); - commitDataWriter = bind(lookup, "spark_ds_commit_data_writer", i32, ptr, ptr, ptr); - abortDataWriter = bind(lookup, "spark_ds_abort_data_writer", i32, ptr); + private enum C { VOID, BOOL, I32, I64, PTR } + + /** The functions of the table {@code SparkDsLibrary}, in their order in the table. */ + enum Function { + LAST_ERROR("last_error", C.PTR), + CREATE("create", C.I32, C.PTR, C.PTR, C.PTR, C.I64, C.PTR, C.PTR), + SCHEMA("schema", C.I32, C.PTR, C.PTR), + CREATE_READER("create_reader", C.I32, C.PTR, C.PTR, C.PTR), + CREATE_STREAM_READER("create_stream_reader", C.I32, C.PTR, C.PTR, C.PTR), + CREATE_WRITER("create_writer", C.I32, C.PTR, C.PTR, C.PTR, C.PTR), + CLOSE("close", C.VOID, C.PTR), + PUSH_PREDICATES("push_predicates", C.I32, C.PTR, C.PTR, C.I64, C.PTR), + PUSH_LIMIT("push_limit", C.I32, C.PTR, C.I32, C.PTR), + PRUNE_COLUMNS("prune_columns", C.I32, C.PTR, C.PTR, C.I64, C.PTR), + PARTITIONS("partitions", C.I32, C.PTR, C.PTR, C.PTR), + SERIALIZE_READER("serialize_reader", C.I32, C.PTR, C.PTR, C.PTR), + CLOSE_READER("close_reader", C.VOID, C.PTR), + INITIAL_OFFSET("initial_offset", C.I32, C.PTR, C.PTR, C.PTR), + LATEST_OFFSET("latest_offset", C.I32, C.PTR, C.PTR, C.PTR), + STREAM_PARTITIONS("stream_partitions", C.I32, C.PTR, C.PTR, C.PTR, C.PTR, C.PTR), + SERIALIZE_STREAM_READER("serialize_stream_reader", C.I32, C.PTR, C.PTR, C.PTR), + COMMIT_OFFSET("commit_offset", C.I32, C.PTR, C.PTR), + CLOSE_STREAM_READER("close_stream_reader", C.VOID, C.PTR), + READ("read", C.I32, C.PTR, C.PTR, C.I64, C.PTR, C.I64, C.PTR), + SERIALIZE_WRITER("serialize_writer", C.I32, C.PTR, C.PTR, C.PTR), + COMMIT("commit", C.I32, C.PTR, C.I64, C.PTR, C.PTR, C.I64), + ABORT("abort", C.I32, C.PTR, C.I64, C.PTR, C.PTR, C.I64), + CLOSE_WRITER("close_writer", C.VOID, C.PTR), + CREATE_DATA_WRITER("create_data_writer", C.I32, C.PTR, C.PTR, C.I64, C.I64, C.PTR), + WRITE("write", C.I32, C.PTR, C.PTR, C.PTR), + COMMIT_DATA_WRITER("commit_data_writer", C.I32, C.PTR, C.PTR, C.PTR), + ABORT_DATA_WRITER("abort_data_writer", C.I32, C.PTR), + CLOSE_DATA_WRITER("close_data_writer", C.VOID, C.PTR); + + /** The name of the field of the function in the table. */ + final String fieldName; + private final C result; + private final C[] parameters; + + Function(String fieldName, C result, C... parameters) { + this.fieldName = fieldName; + this.result = result; + this.parameters = parameters; + } + + /** The offset of the function in the table, after the size and the version of the table. */ + long offset() { + return TABLE_FUNCTIONS + 8L * ordinal(); + } + } + + // The layouts of the structs of the header, on the 64-bit platforms that Spark supports. Each + // struct starts with its size. + private static final long TABLE_ABI_VERSION = 8; + private static final long TABLE_FUNCTIONS = 16; + private static final long HOST_SIZE = 32; + private static final long HOST_ABI_VERSION = 8; + private static final long HOST_SPARK_VERSION = 16; + private static final long HOST_LOG = 24; + private static final long TASK_SIZE = 40; + private static final long TASK_STAGE_ID = 8; + private static final long TASK_STAGE_ATTEMPT_NUMBER = 12; + private static final long TASK_PARTITION_ID = 16; + private static final long TASK_ATTEMPT_NUMBER = 20; + private static final long TASK_ATTEMPT_ID = 24; + private static final long TASK_IS_INTERRUPTED = 32; + private static final long WRITE_INFO_SIZE = 24; + private static final long WRITE_INFO_MODE = 8; + private static final long WRITE_INFO_QUERY_ID = 16; + + /** + * The offsets of the fields of the structs that Spark reads or writes, by {@code struct.field}, + * and the sizes of the structs, by {@code struct}, for the tests to compare with the header. + */ + static Map layout() { + Map layout = new LinkedHashMap<>(); + layout.put("SparkDsLibrary", TABLE_FUNCTIONS + 8L * Function.values().length); + layout.put("SparkDsLibrary.size", 0L); + layout.put("SparkDsLibrary.abi_version", TABLE_ABI_VERSION); + for (Function function : Function.values()) { + layout.put("SparkDsLibrary." + function.fieldName, function.offset()); + } + layout.put("SparkDsHost", HOST_SIZE); + layout.put("SparkDsHost.size", 0L); + layout.put("SparkDsHost.abi_version", HOST_ABI_VERSION); + layout.put("SparkDsHost.spark_version", HOST_SPARK_VERSION); + layout.put("SparkDsHost.log", HOST_LOG); + layout.put("SparkDsTaskContext", TASK_SIZE); + layout.put("SparkDsTaskContext.size", 0L); + layout.put("SparkDsTaskContext.stage_id", TASK_STAGE_ID); + layout.put("SparkDsTaskContext.stage_attempt_number", TASK_STAGE_ATTEMPT_NUMBER); + layout.put("SparkDsTaskContext.partition_id", TASK_PARTITION_ID); + layout.put("SparkDsTaskContext.attempt_number", TASK_ATTEMPT_NUMBER); + layout.put("SparkDsTaskContext.task_attempt_id", TASK_ATTEMPT_ID); + layout.put("SparkDsTaskContext.is_interrupted", TASK_IS_INTERRUPTED); + layout.put("SparkDsWriteInfo", WRITE_INFO_SIZE); + layout.put("SparkDsWriteInfo.size", 0L); + layout.put("SparkDsWriteInfo.mode", WRITE_INFO_MODE); + layout.put("SparkDsWriteInfo.query_id", WRITE_INFO_QUERY_ID); + return layout; + } + + private final int abiVersion; + // The functions of the table, by ordinal, null for the functions that the library leaves NULL. + private final MethodHandle[] functions = new MethodHandle[Function.values().length]; + + private NativeLibrary(long table) throws Throwable { + long size = Platform.getLong(null, table); + if (size < TABLE_FUNCTIONS) { + throw new LoadException( + "spark_ds_init returned a function table with a size of " + size + " bytes.", null, true); + } + abiVersion = Platform.getInt(null, table + TABLE_ABI_VERSION); + // The table of another version of the interface has another layout. Only the functions that + // fit in the size of the table are read: the table of a library built against an earlier + // header ends before the functions that were added later. + if (abiVersion == ABI_VERSION) { + for (Function function : Function.values()) { + if (function.offset() + 8 <= size) { + long address = Platform.getLong(null, table + function.offset()); + if (address != 0) { + functions[function.ordinal()] = Ffm.downcall( + Ffm.OF_ADDRESS.invoke(address), function.result, function.parameters); + } + } + } + } } /** - * Loads the library at the given path. The library stays loaded until the JVM exits. + * Loads the library at the given path, and gets its table with {@code spark_ds_init}. The + * library stays loaded until the JVM exits. * + * @param sparkVersion the version of Spark, which the library gets in {@code SparkDsHost} * @throws LoadException if the library cannot be loaded */ - static NativeLibrary load(String path) throws Throwable { + static NativeLibrary load(String path, String sparkVersion) throws Throwable { if (!isSupported()) { throw new UnsupportedOperationException( "Native data sources require Java " + MINIMUM_JAVA_VERSION + " or later."); @@ -156,36 +238,30 @@ static NativeLibrary load(String path) throws Throwable { throw new LoadException( "Cannot use the Foreign Function and Memory API: " + rootCause(e), e); } - return new NativeLibrary(lookup); - } - - /** - * Binds a function of the library to a method handle whose parameters and result have the - * Java types of the given C types, with {@code long} for pointers. Returns null if the library - * does not export the function. - */ - private static MethodHandle bind(Object lookup, String name, C result, C... parameters) - throws Throwable { - Optional symbol = (Optional) Ffm.FIND.invoke(lookup, name); - if (symbol.isEmpty()) { - return null; - } - Object layouts = Ffm.layouts(parameters); - Object descriptor = result == C.VOID - ? Ffm.FUNCTION_OF_VOID.invoke(layouts) - : Ffm.FUNCTION_OF.invoke(Ffm.layout(result), layouts); - MethodHandle handle = (MethodHandle) Ffm.DOWNCALL.invoke( - Ffm.LINKER, symbol.get(), descriptor, Ffm.NO_OPTIONS); - // Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses. - for (int i = 0; i < parameters.length; i++) { - if (parameters[i] == C.POINTER) { - handle = MethodHandles.filterArguments(handle, i, Ffm.OF_ADDRESS); - } + Optional init = (Optional) Ffm.FIND.invoke(lookup, "spark_ds_init"); + if (init.isEmpty()) { + throw new LoadException( + "The library does not export the function spark_ds_init.", null, true); } - if (result == C.POINTER) { - handle = MethodHandles.filterReturnValue(handle, Ffm.ADDRESS_OF); + MethodHandle initFunction = Ffm.downcall(init.get(), C.I32, C.PTR, C.PTR); + // The library may keep the host until the JVM exits, even if spark_ds_init fails. + long host = Ffm.host(path, sparkVersion); + long out = Platform.allocateMemory(8); + try { + Platform.putLong(null, out, 0); + int status = (int) initFunction.invokeExact(host, out); + if (status != 0) { + throw new LoadException( + "spark_ds_init failed with the error code " + status + ".", null, true); + } + long table = Platform.getLong(null, out); + if (table == 0) { + throw new LoadException("spark_ds_init returned no function table.", null, true); + } + return new NativeLibrary(table); + } finally { + Platform.freeMemory(out); } - return handle; } /** The method handles of the FFM API, initialized when the first library is loaded. */ @@ -197,6 +273,7 @@ private static final class Ffm { static final MethodHandle LIBRARY_LOOKUP; static final MethodHandle FIND; static final MethodHandle DOWNCALL; + static final MethodHandle UPCALL; static final MethodHandle FUNCTION_OF; static final MethodHandle FUNCTION_OF_VOID; // (long) -> MemorySegment and (MemorySegment) -> long. @@ -206,8 +283,13 @@ private static final class Ffm { static final Object INT32; static final Object INT64; static final Object POINTER; - // The address of the C function SparkDsSink, which calls NativeLibrary.sink. + // NativeLibrary.log, which SparkDsHost.log calls with the path of its library. + static final MethodHandle LOG; + // The addresses of the C functions that call Java, which live as long as the JVM: + // SparkDsSink, which calls NativeLibrary.sink, and SparkDsTaskContext.is_interrupted, which + // calls NativeLibrary.isInterrupted. static final long SINK; + static final long IS_INTERRUPTED; static { try { @@ -237,6 +319,9 @@ private static final class Ffm { DOWNCALL = lookup.findVirtual(linker, "downcallHandle", MethodType.methodType(MethodHandle.class, segment, descriptor, options)) .asFixedArity(); + UPCALL = lookup.findVirtual(linker, "upcallStub", + MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options)) + .asFixedArity(); FUNCTION_OF = lookup.findStatic(descriptor, "of", MethodType.methodType(descriptor, layout, layoutArray)).asFixedArity(); FUNCTION_OF_VOID = lookup.findStatic(descriptor, "ofVoid", @@ -249,17 +334,16 @@ private static final class Ffm { INT64 = valueLayout.getField("JAVA_LONG").get(null); POINTER = valueLayout.getField("ADDRESS").get(null); - // The upcall stub of the sink, which lives as long as the JVM. - MethodHandle target = lookup.findStatic(NativeLibrary.class, "sink", - MethodType.methodType(void.class, long.class, long.class, long.class)); - target = MethodHandles.filterArguments(target, 0, ADDRESS_OF, ADDRESS_OF); - Object sinkDescriptor = - FUNCTION_OF_VOID.invoke(layouts(C.POINTER, C.POINTER, C.INT64)); - Object stub = lookup.findVirtual(linker, "upcallStub", - MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options)) - .asFixedArity() - .invoke(LINKER, target, sinkDescriptor, GLOBAL_ARENA, NO_OPTIONS); - SINK = (long) ADDRESS_OF.invoke(stub); + LOG = lookup.findStatic(NativeLibrary.class, "log", + MethodType.methodType(void.class, String.class, int.class, long.class)); + SINK = upcall( + lookup.findStatic(NativeLibrary.class, "sink", + MethodType.methodType(void.class, long.class, long.class, long.class)), + C.VOID, C.PTR, C.PTR, C.I64); + IS_INTERRUPTED = upcall( + lookup.findStatic(NativeLibrary.class, "isInterrupted", + MethodType.methodType(boolean.class, long.class)), + C.BOOL, C.PTR); } catch (Throwable e) { // The message names the cause, because the JVM only keeps this message for the later // uses of the class. @@ -267,24 +351,82 @@ private static final class Ffm { } } - /** Returns a MemoryLayout[] with the layouts of the given types. */ - static Object layouts(C... types) { - Object layouts = Array.newInstance(LAYOUT_CLASS, types.length); - for (int i = 0; i < types.length; i++) { - Array.set(layouts, i, layout(types[i])); + /** Returns the FunctionDescriptor of a C function with the given types. */ + static Object descriptor(C result, C... parameters) throws Throwable { + Object layouts = Array.newInstance(LAYOUT_CLASS, parameters.length); + for (int i = 0; i < parameters.length; i++) { + Array.set(layouts, i, layout(parameters[i])); } - return layouts; + return result == C.VOID + ? FUNCTION_OF_VOID.invoke(layouts) + : FUNCTION_OF.invoke(layout(result), layouts); } static Object layout(C type) { return switch (type) { case BOOL -> BOOL; - case INT32 -> INT32; - case INT64 -> INT64; - case POINTER -> POINTER; + case I32 -> INT32; + case I64 -> INT64; + case PTR -> POINTER; case VOID -> throw new IllegalArgumentException("void has no layout"); }; } + + /** + * Binds the C function at the given address, a MemorySegment, to a method handle whose + * parameters and result have the Java types of the given C types, with {@code long} for + * pointers. + */ + static MethodHandle downcall(Object address, C result, C... parameters) throws Throwable { + MethodHandle handle = (MethodHandle) DOWNCALL.invoke( + LINKER, address, descriptor(result, parameters), NO_OPTIONS); + // Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses. + for (int i = 0; i < parameters.length; i++) { + if (parameters[i] == C.PTR) { + handle = MethodHandles.filterArguments(handle, i, OF_ADDRESS); + } + } + if (result == C.PTR) { + handle = MethodHandles.filterReturnValue(handle, ADDRESS_OF); + } + return handle; + } + + /** + * Returns the address of a C function with the given types that calls `target`, whose + * parameters have the Java types of the C types, with {@code long} for pointers. The + * function lives as long as the JVM. + */ + static long upcall(MethodHandle target, C result, C... parameters) throws Throwable { + for (int i = 0; i < parameters.length; i++) { + if (parameters[i] == C.PTR) { + target = MethodHandles.filterArguments(target, i, ADDRESS_OF); + } + } + Object stub = UPCALL.invoke( + LINKER, target, descriptor(result, parameters), GLOBAL_ARENA, NO_OPTIONS); + return (long) ADDRESS_OF.invoke(stub); + } + + /** + * Allocates the SparkDsHost of a library. The library may keep it until the JVM exits, so it + * is never freed. + */ + static long host(String path, String sparkVersion) throws Throwable { + long log = upcall(MethodHandles.insertArguments(LOG, 0, path), C.VOID, C.I32, C.PTR); + byte[] version = sparkVersion.getBytes(StandardCharsets.UTF_8); + long versionAddress = Platform.allocateMemory(version.length + 1L); + Platform.copyMemory( + version, Platform.BYTE_ARRAY_OFFSET, null, versionAddress, version.length); + Platform.putByte(null, versionAddress + version.length, (byte) 0); + long host = Platform.allocateMemory(HOST_SIZE); + Platform.setMemory(host, (byte) 0, HOST_SIZE); + Platform.putLong(null, host, HOST_SIZE); + Platform.putInt(null, host + HOST_ABI_VERSION, ABI_VERSION); + Platform.putLong(null, host + HOST_SPARK_VERSION, versionAddress); + Platform.putLong(null, host + HOST_LOG, log); + return host; + } } private static Throwable rootCause(Throwable e) { @@ -295,232 +437,284 @@ private static Throwable rootCause(Throwable e) { return cause; } - /** Thrown when a library cannot be loaded. */ + /** Thrown when a library cannot be loaded, or cannot be used. */ static final class LoadException extends RuntimeException { + /** + * Whether the library was loaded, in which case it stays loaded: Spark found that it cannot + * use it afterwards, for example because {@code spark_ds_init} failed. + */ + final boolean loaded; + LoadException(String message, Throwable cause) { + this(message, cause, false); + } + + LoadException(String message, Throwable cause, boolean loaded) { super(message, cause); + this.loaded = loaded; } } - /** Thrown when a function is called that the library does not export, with its name. */ + /** Thrown when a function is called that the library leaves NULL in its table, with its name. */ static final class MissingFunctionException extends RuntimeException { MissingFunctionException(String name) { super(name); } } + /** + * Called by the library, through {@code SparkDsHost.log}, with the path of the library. It must + * not throw: an exception thrown by an upcall crashes the JVM. + */ + private static void log(String path, int level, long message) { + try { + String text = message == 0 ? "" : new String(readCString(message), StandardCharsets.UTF_8); + switch (level) { + case 1 -> LOGGER.error("Native data source library {}: {}", + MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text)); + case 2 -> LOGGER.warn("Native data source library {}: {}", + MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text)); + case 4 -> LOGGER.debug("Native data source library {}: {}", path, text); + case 5 -> LOGGER.trace("Native data source library {}: {}", path, text); + default -> LOGGER.info("Native data source library {}: {}", + MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text)); + } + } catch (Throwable e) { + // The message is lost. + } + } + // --------------------------------------------------------------------------------------------- // Calls // --------------------------------------------------------------------------------------------- - private static MethodHandle require(MethodHandle function, String name) { - if (function == null) { - throw new MissingFunctionException(name); + private MethodHandle function(Function function) { + MethodHandle handle = functions[function.ordinal()]; + if (handle == null) { + throw new MissingFunctionException(function.fieldName); } - return function; + return handle; } /** Throws the error that the library reported if a function failed. */ - private void check(int status, String name) throws Throwable { + private void check(int status, Function function) throws Throwable { if (status == 0) { return; } + MethodHandle lastError = functions[Function.LAST_ERROR.ordinal()]; long address = lastError == null ? 0 : (long) lastError.invokeExact(); String message = address == 0 ? "" : new String(readCString(address), StandardCharsets.UTF_8); throw new RuntimeException(message.isEmpty() - ? name + " failed with the error code " + status + "." + ? function.fieldName + " failed with the error code " + status + "." : message); } - private static long readHandle(long out, String name) { + private static long readHandle(long out, Function function) { long handle = Platform.getLong(null, out); if (handle == 0) { - throw new IllegalStateException(name + " returned a NULL handle."); + throw new IllegalStateException(function.fieldName + " returned a NULL handle."); } return handle; } - int abiVersion() throws Throwable { - return (int) require(abiVersion, "spark_ds_abi_version").invokeExact(); + /** The version of the interface that the table of the library implements. */ + int abiVersion() { + return abiVersion; + } + + /** Whether the library implements the function, which it does not leave NULL in its table. */ + boolean hasFunction(Function function) { + return functions[function.ordinal()] != null; } - long createDataSource(String name, String[] optionKeys, String[] optionValues) + /** A data source handle, with the capabilities of the data source. */ + static final class DataSourceHandle { + final long handle; + final long capabilities; + + DataSourceHandle(long handle, long capabilities) { + this.handle = handle; + this.capabilities = capabilities; + } + } + + DataSourceHandle createDataSource(String name, String[] optionKeys, String[] optionValues) throws Throwable { - MethodHandle f = require(create, "spark_ds_create"); + MethodHandle f = function(Function.CREATE); try (Memory memory = new Memory()) { + long capabilities = memory.zeroed(8); long out = memory.zeroed(8); check((int) f.invokeExact(memory.string(name), memory.strings(optionKeys), - memory.strings(optionValues), (long) optionKeys.length, out), "spark_ds_create"); - return readHandle(out, "spark_ds_create"); + memory.strings(optionValues), (long) optionKeys.length, capabilities, out), + Function.CREATE); + return new DataSourceHandle( + readHandle(out, Function.CREATE), Platform.getLong(null, capabilities)); } } void schema(long dataSource, long schemaAddress) throws Throwable { - check((int) require(schema, "spark_ds_schema").invokeExact(dataSource, schemaAddress), - "spark_ds_schema"); + check((int) function(Function.SCHEMA).invokeExact(dataSource, schemaAddress), + Function.SCHEMA); } long createReader(long dataSource, long schemaAddress) throws Throwable { - return createWithSchema(createReader, "spark_ds_create_reader", dataSource, schemaAddress); + return createWithSchema(Function.CREATE_READER, dataSource, schemaAddress); } long createStreamReader(long dataSource, long schemaAddress) throws Throwable { - return createWithSchema( - createStreamReader, "spark_ds_create_stream_reader", dataSource, schemaAddress); - } - - private long createWithSchema( - MethodHandle function, String name, long dataSource, long schemaAddress) throws Throwable { - MethodHandle f = require(function, name); - try (Memory memory = new Memory()) { - long out = memory.zeroed(8); - check((int) f.invokeExact(dataSource, schemaAddress, out), name); - return readHandle(out, name); - } + return createWithSchema(Function.CREATE_STREAM_READER, dataSource, schemaAddress); } - long createWriter(long dataSource, long schemaAddress, boolean overwrite) throws Throwable { - MethodHandle f = require(createWriter, "spark_ds_create_writer"); + private long createWithSchema(Function function, long dataSource, long schemaAddress) + throws Throwable { + MethodHandle f = function(function); try (Memory memory = new Memory()) { long out = memory.zeroed(8); - check((int) f.invokeExact(dataSource, schemaAddress, overwrite, out), - "spark_ds_create_writer"); - return readHandle(out, "spark_ds_create_writer"); + check((int) f.invokeExact(dataSource, schemaAddress, out), function); + return readHandle(out, function); } } - long createStreamWriter(long dataSource, long schemaAddress, boolean overwrite, String queryId) + /** + * Creates a writer. + * + * @param mode {@link #WRITE_APPEND} or {@link #WRITE_TRUNCATE} + * @param queryId the ID of the streaming query of a streaming write, or null for a batch write + */ + long createWriter(long dataSource, long schemaAddress, int mode, String queryId) throws Throwable { - MethodHandle f = require(createStreamWriter, "spark_ds_create_stream_writer"); + MethodHandle f = function(Function.CREATE_WRITER); try (Memory memory = new Memory()) { + long info = memory.zeroed(WRITE_INFO_SIZE); + Platform.putLong(null, info, WRITE_INFO_SIZE); + Platform.putInt(null, info + WRITE_INFO_MODE, mode); + Platform.putLong( + null, info + WRITE_INFO_QUERY_ID, queryId == null ? 0 : memory.string(queryId)); long out = memory.zeroed(8); - check((int) f.invokeExact(dataSource, schemaAddress, overwrite, memory.string(queryId), out), - "spark_ds_create_stream_writer"); - return readHandle(out, "spark_ds_create_stream_writer"); + check((int) f.invokeExact(dataSource, schemaAddress, info, out), Function.CREATE_WRITER); + return readHandle(out, Function.CREATE_WRITER); } } void closeDataSource(long dataSource) throws Throwable { - require(close, "spark_ds_close").invokeExact(dataSource); + function(Function.CLOSE).invokeExact(dataSource); } - boolean[] pushPredicates(long reader, String[] predicates) throws Throwable { - MethodHandle f = require(pushPredicates, "spark_ds_push_predicates"); + /** Returns how the reader uses each predicate, as SPARK_DS_PREDICATE_* values. */ + int[] pushPredicates(long reader, String[] predicates) throws Throwable { + MethodHandle f = function(Function.PUSH_PREDICATES); try (Memory memory = new Memory()) { - long evaluated = memory.zeroed(Math.max(1, predicates.length)); + long uses = memory.zeroed(4L * Math.max(1, predicates.length)); check((int) f.invokeExact(reader, memory.strings(predicates), (long) predicates.length, - evaluated), "spark_ds_push_predicates"); - boolean[] result = new boolean[predicates.length]; + uses), Function.PUSH_PREDICATES); + int[] result = new int[predicates.length]; for (int i = 0; i < result.length; i++) { - result[i] = Platform.getByte(null, evaluated + i) != 0; + result[i] = Platform.getInt(null, uses + 4L * i); } return result; } } boolean pushLimit(long reader, int limit) throws Throwable { - MethodHandle f = require(pushLimit, "spark_ds_push_limit"); + MethodHandle f = function(Function.PUSH_LIMIT); try (Memory memory = new Memory()) { long applied = memory.zeroed(1); - check((int) f.invokeExact(reader, limit, applied), "spark_ds_push_limit"); + check((int) f.invokeExact(reader, limit, applied), Function.PUSH_LIMIT); return Platform.getByte(null, applied) != 0; } } boolean pruneColumns(long reader, String[] columnNames) throws Throwable { - MethodHandle f = require(pruneColumns, "spark_ds_prune_columns"); + MethodHandle f = function(Function.PRUNE_COLUMNS); try (Memory memory = new Memory()) { long pruned = memory.zeroed(1); check((int) f.invokeExact(reader, memory.strings(columnNames), (long) columnNames.length, - pruned), "spark_ds_prune_columns"); + pruned), Function.PRUNE_COLUMNS); return Platform.getByte(null, pruned) != 0; } } byte[][] partitions(long reader) throws Throwable { - MethodHandle f = require(partitions, "spark_ds_partitions"); + MethodHandle f = function(Function.PARTITIONS); try (Sink sink = new Sink()) { - check((int) f.invokeExact(reader, Ffm.SINK, sink.id), "spark_ds_partitions"); + check((int) f.invokeExact(reader, Ffm.SINK, sink.id), Function.PARTITIONS); return sink.all(); } } byte[] serializeReader(long reader) throws Throwable { - return serialize(serializeReader, "spark_ds_serialize_reader", reader); + return serialize(Function.SERIALIZE_READER, reader); } void closeReader(long reader) throws Throwable { - require(closeReader, "spark_ds_close_reader").invokeExact(reader); + function(Function.CLOSE_READER).invokeExact(reader); } String initialOffset(long streamReader) throws Throwable { - return offset(initialOffset, "spark_ds_initial_offset", streamReader); + return offset(Function.INITIAL_OFFSET, streamReader); } String latestOffset(long streamReader) throws Throwable { - return offset(latestOffset, "spark_ds_latest_offset", streamReader); + return offset(Function.LATEST_OFFSET, streamReader); } - private String offset(MethodHandle function, String name, long streamReader) - throws Throwable { - byte[] offset = serialize(function, name, streamReader); + private String offset(Function function, long streamReader) throws Throwable { + byte[] offset = serialize(function, streamReader); if (offset == null) { - throw new IllegalStateException(name + " returned no offset."); + throw new IllegalStateException(function.fieldName + " returned no offset."); } return new String(offset, StandardCharsets.UTF_8); } byte[][] streamPartitions(long streamReader, String start, String end) throws Throwable { - MethodHandle f = require(streamPartitions, "spark_ds_stream_partitions"); + MethodHandle f = function(Function.STREAM_PARTITIONS); try (Memory memory = new Memory(); Sink sink = new Sink()) { check((int) f.invokeExact(streamReader, memory.string(start), memory.string(end), - Ffm.SINK, sink.id), "spark_ds_stream_partitions"); + Ffm.SINK, sink.id), Function.STREAM_PARTITIONS); return sink.all(); } } byte[] serializeStreamReader(long streamReader) throws Throwable { - return serialize(serializeStreamReader, "spark_ds_serialize_stream_reader", streamReader); + return serialize(Function.SERIALIZE_STREAM_READER, streamReader); } void commitOffset(long streamReader, String end) throws Throwable { - MethodHandle f = require(commitOffset, "spark_ds_commit_offset"); + MethodHandle f = function(Function.COMMIT_OFFSET); try (Memory memory = new Memory()) { - check((int) f.invokeExact(streamReader, memory.string(end)), "spark_ds_commit_offset"); + check((int) f.invokeExact(streamReader, memory.string(end)), Function.COMMIT_OFFSET); } } void closeStreamReader(long streamReader) throws Throwable { - require(closeStreamReader, "spark_ds_close_stream_reader").invokeExact(streamReader); + function(Function.CLOSE_STREAM_READER).invokeExact(streamReader); } - void read(byte[] readerState, byte[] partition, long streamAddress) throws Throwable { - MethodHandle f = require(read, "spark_ds_read"); + void read(Task task, byte[] readerState, byte[] partition, long streamAddress) + throws Throwable { + MethodHandle f = function(Function.READ); try (Memory memory = new Memory()) { - check((int) f.invokeExact(memory.bytes(readerState), (long) readerState.length, - memory.bytes(partition), (long) partition.length, streamAddress), "spark_ds_read"); + check((int) f.invokeExact(task.address, memory.bytes(readerState), + (long) readerState.length, memory.bytes(partition), (long) partition.length, + streamAddress), Function.READ); } } byte[] serializeWriter(long writer) throws Throwable { - return serialize(serializeWriter, "spark_ds_serialize_writer", writer); + return serialize(Function.SERIALIZE_WRITER, writer); } void commit(long writer, long epochId, byte[][] messages) throws Throwable { - finishWrite(commit, "spark_ds_commit", writer, epochId, messages); + finishWrite(Function.COMMIT, writer, epochId, messages); } void abort(long writer, long epochId, byte[][] messages) throws Throwable { - finishWrite(abort, "spark_ds_abort", writer, epochId, messages); + finishWrite(Function.ABORT, writer, epochId, messages); } - private void finishWrite( - MethodHandle function, - String name, - long writer, - long epochId, - byte[][] messages) throws Throwable { - MethodHandle f = require(function, name); + private void finishWrite(Function function, long writer, long epochId, byte[][] messages) + throws Throwable { + MethodHandle f = function(function); try (Memory memory = new Memory()) { long pointers = memory.zeroed(8L * Math.max(1, messages.length)); long lengths = memory.zeroed(8L * Math.max(1, messages.length)); @@ -531,68 +725,70 @@ private void finishWrite( } } check((int) f.invokeExact(writer, epochId, pointers, lengths, (long) messages.length), - name); + function); } } void closeWriter(long writer) throws Throwable { - require(closeWriter, "spark_ds_close_writer").invokeExact(writer); + function(Function.CLOSE_WRITER).invokeExact(writer); } - long createDataWriter(byte[] writerState, int partitionId, long taskId, long epochId) - throws Throwable { - MethodHandle f = require(createDataWriter, "spark_ds_create_data_writer"); + long createDataWriter(Task task, byte[] writerState, long epochId) throws Throwable { + MethodHandle f = function(Function.CREATE_DATA_WRITER); try (Memory memory = new Memory()) { long out = memory.zeroed(8); - check((int) f.invokeExact(memory.bytes(writerState), (long) writerState.length, - partitionId, taskId, epochId, out), "spark_ds_create_data_writer"); - return readHandle(out, "spark_ds_create_data_writer"); + check((int) f.invokeExact(task.address, memory.bytes(writerState), + (long) writerState.length, epochId, out), Function.CREATE_DATA_WRITER); + return readHandle(out, Function.CREATE_DATA_WRITER); } } void write(long dataWriter, long arrayAddress, long schemaAddress) throws Throwable { - check((int) require(write, "spark_ds_write").invokeExact( - dataWriter, arrayAddress, schemaAddress), "spark_ds_write"); + check((int) function(Function.WRITE).invokeExact(dataWriter, arrayAddress, schemaAddress), + Function.WRITE); } /** - * Commits a data writer, whose handle the library releases when the commit succeeds. Then - * {@code released} runs right away, before the message is checked: if the message is invalid, - * the commit fails, but the handle must not be used again. + * Commits a data writer. Then {@code committed} runs right away, before the message is checked: + * if the message is invalid, the commit fails, but the data writer must not be aborted. */ - byte[] commitDataWriter(long dataWriter, Runnable released) throws Throwable { - MethodHandle f = require(commitDataWriter, "spark_ds_commit_data_writer"); + byte[] commitDataWriter(long dataWriter, Runnable committed) throws Throwable { + MethodHandle f = function(Function.COMMIT_DATA_WRITER); try (Sink sink = new Sink()) { - check((int) f.invokeExact(dataWriter, Ffm.SINK, sink.id), "spark_ds_commit_data_writer"); - released.run(); - byte[] message = single(sink, "spark_ds_commit_data_writer"); + check((int) f.invokeExact(dataWriter, Ffm.SINK, sink.id), Function.COMMIT_DATA_WRITER); + committed.run(); + byte[] message = single(sink, Function.COMMIT_DATA_WRITER); return message == null ? new byte[0] : message; } } void abortDataWriter(long dataWriter) throws Throwable { - check((int) require(abortDataWriter, "spark_ds_abort_data_writer").invokeExact(dataWriter), - "spark_ds_abort_data_writer"); + check((int) function(Function.ABORT_DATA_WRITER).invokeExact(dataWriter), + Function.ABORT_DATA_WRITER); + } + + void closeDataWriter(long dataWriter) throws Throwable { + function(Function.CLOSE_DATA_WRITER).invokeExact(dataWriter); } /** * Calls a function that returns a single byte array through a sink. Returns null if the * function does not call the sink. */ - private byte[] serialize(MethodHandle function, String name, long handle) throws Throwable { - MethodHandle f = require(function, name); + private byte[] serialize(Function function, long handle) throws Throwable { + MethodHandle f = function(function); try (Sink sink = new Sink()) { - check((int) f.invokeExact(handle, Ffm.SINK, sink.id), name); - return single(sink, name); + check((int) f.invokeExact(handle, Ffm.SINK, sink.id), function); + return single(sink, function); } } /** Returns the single byte array that a function passed to the sink, or null if none. */ - private static byte[] single(Sink sink, String name) throws Throwable { + private static byte[] single(Sink sink, Function function) throws Throwable { byte[][] values = sink.all(); if (values.length > 1) { throw new IllegalStateException( - name + " returned " + values.length + " values instead of one."); + function.fieldName + " returned " + values.length + " values instead of one."); } return values.length == 0 ? null : values[0]; } @@ -652,6 +848,65 @@ private static byte[] readCString(long address) { return bytes; } + // --------------------------------------------------------------------------------------------- + // Task contexts + // --------------------------------------------------------------------------------------------- + + // Whether the task of each SparkDsTaskContext was killed, by the address of the struct. + private static final ConcurrentHashMap tasks = new ConcurrentHashMap<>(); + + /** + * The {@code SparkDsTaskContext} of a task that reads a partition or writes data. The reader or + * the data writer closes it once the library no longer uses it, which frees it. + */ + static final class Task implements AutoCloseable { + final long address; + private boolean closed; + + Task( + int stageId, + int stageAttemptNumber, + int partitionId, + int attemptNumber, + long taskAttemptId, + BooleanSupplier interrupted) { + address = Platform.allocateMemory(TASK_SIZE); + Platform.setMemory(address, (byte) 0, TASK_SIZE); + Platform.putLong(null, address, TASK_SIZE); + Platform.putInt(null, address + TASK_STAGE_ID, stageId); + Platform.putInt(null, address + TASK_STAGE_ATTEMPT_NUMBER, stageAttemptNumber); + Platform.putInt(null, address + TASK_PARTITION_ID, partitionId); + Platform.putInt(null, address + TASK_ATTEMPT_NUMBER, attemptNumber); + Platform.putLong(null, address + TASK_ATTEMPT_ID, taskAttemptId); + Platform.putLong(null, address + TASK_IS_INTERRUPTED, Ffm.IS_INTERRUPTED); + tasks.put(address, interrupted); + } + + @Override + public void close() { + if (!closed) { + closed = true; + // A call to is_interrupted that is still running only looks the address up. + tasks.remove(address); + Platform.freeMemory(address); + } + } + } + + /** + * Called by the library, through {@code SparkDsTaskContext.is_interrupted}, with the address of + * the struct. It must not throw: an exception thrown by an upcall crashes the JVM. + */ + private static boolean isInterrupted(long task) { + try { + BooleanSupplier interrupted = tasks.get(task); + // A freed task context belongs to a read or a write that is over. + return interrupted == null || interrupted.getAsBoolean(); + } catch (Throwable e) { + return false; + } + } + // --------------------------------------------------------------------------------------------- // Sink // --------------------------------------------------------------------------------------------- diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json b/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json index b6d74c3b6303a..a2be573c93029 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json @@ -1,5 +1,5 @@ { - "description": "Predicates as Spark passes them to spark_ds_push_predicates, with the ids of the rows of the table for which they are true. A case lists the predicates that Spark pushes for the filter in its SQL, which are combined with AND. A row has a value for each column of the schema; NaN and infinite doubles are strings.", + "description": "Predicates as Spark passes them to push_predicates of spark_datasource.h, with the ids of the rows of the table for which they are true. A case lists the predicates that Spark pushes for the filter in its SQL, which are combined with AND. A row has a value for each column of the schema; NaN and infinite doubles are strings.", "schema": "id INT,i INT,d DOUBLE,s STRING", "rows": [ [0, 1, 0.0, "a"], diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h index 8dffaefc0493f..dd55f6795aa4d 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -19,12 +19,13 @@ * spark_datasource.h: the binary interface between Spark and data sources implemented in native * code, such as Rust or C++. Since Spark 4.4.0. * - * A native data source is a shared library that exports the C functions declared here. Spark - * calls them with the Foreign Function and Memory API of Java, so native data sources require - * Java 22 or later. Spark looks the functions up in each library separately, so any number of - * libraries that export the same functions can be loaded side by side. Spark adapts a native data - * source to Data Source V2, with batch and micro-batch streaming reads, batch and streaming - * writes, and predicate, column and limit pushdown. + * A native data source is a shared library that exports a single C function, spark_ds_init, + * which returns the table of the functions that implement the data source, a SparkDsLibrary. + * Spark calls them with the Foreign Function and Memory API of Java, so native data sources + * require Java 22 or later. Each library has its own table, so any number of libraries can be + * loaded side by side. Spark adapts a native data source to Data Source V2, with batch and + * micro-batch streaming reads, batch and streaming writes, and predicate, column and limit + * pushdown. * * This header is self-contained: it only depends on the C standard library, and defines the * structs of the Arrow C data interface unless they are already defined, as by Arrow C++. A library @@ -45,78 +46,76 @@ * * Lifecycle * --------- - * Spark creates a data source with spark_ds_create on the driver for each schema inference, scan - * and write, and closes it right after. - * - * - Batch reads. On the driver, Spark creates a reader with spark_ds_create_reader, pushes - * operations down by calling, each at most once and in this order, spark_ds_push_predicates, - * spark_ds_push_limit and spark_ds_prune_columns, plans the partitions with - * spark_ds_partitions, serializes the state of the reader with spark_ds_serialize_reader, and - * closes it. Then it calls spark_ds_read on the executors, in one task for each partition. If - * Spark never plans the partitions, for example to only explain the plan of the query, it - * closes the reader once the reader is no longer referenced. - * - Streaming reads. On the driver, Spark creates a stream reader with - * spark_ds_create_stream_reader and keeps it until the streaming query stops. For each - * micro-batch, it asks for the spark_ds_latest_offset, plans the partitions with - * spark_ds_stream_partitions and serializes the state with spark_ds_serialize_stream_reader, - * then calls spark_ds_read on the executors. Offsets are defined by the library, as compact - * JSON on a single line: Spark stores each offset as a line of the offset log of the + * Spark loads a library once in each JVM, and calls spark_ds_init before any other function. It + * creates a data source with create on the driver for each schema inference, scan and write, and + * closes it right after. create also returns the capabilities of the data source: Spark reports + * the operations that a data source does not declare as unsupported, without calling the library. + * + * - Batch reads. On the driver, Spark creates a reader with create_reader, pushes operations down + * by calling, each at most once and in this order, push_predicates, push_limit and + * prune_columns, plans the partitions with partitions, serializes the state of the reader with + * serialize_reader, and closes it. Then it calls read on the executors, in one task for each + * partition. If Spark never plans the partitions, for example to only explain the plan of the + * query, it closes the reader once the reader is no longer referenced. + * - Streaming reads. On the driver, Spark creates a stream reader with create_stream_reader and + * keeps it until the streaming query stops. For each micro-batch, it asks for the latest_offset, + * plans the partitions with stream_partitions and serializes the state with + * serialize_stream_reader, then calls read on the executors. Offsets are defined by the library, + * as compact JSON on a single line: Spark stores each offset as a line of the offset log of the * checkpoint, so it rejects an offset that is not valid JSON or contains a line break. When a * streaming query restarts, from its checkpoint, Spark creates a new stream reader and may call - * spark_ds_stream_partitions again for the last range that it planned before, so the partitions - * of a range must read the same data each time. + * stream_partitions again for the last range that it planned before, so the partitions of a + * range must read the same data each time. * - Writes. On the driver, right before it starts the tasks of the write (or of a micro-batch), - * Spark creates a writer with spark_ds_create_writer or spark_ds_create_stream_writer, and - * serializes its state with spark_ds_serialize_writer. With adaptive query execution, Spark - * first runs the stages of the data to write that end with a shuffle, and creates no writer if - * they fail. The rest of the data, such as its last stage, is computed in the tasks of the - * write, so a writer may still be aborted because computing the data failed: do irreversible - * work, such as removing the existing data for an overwrite, in spark_ds_commit. On the - * executors, each task creates a data writer with spark_ds_create_data_writer, calls - * spark_ds_write for each batch of at most - * spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends once it holds - * spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls spark_ds_commit_data_writer, or - * spark_ds_abort_data_writer if the task or the commit fails. Then the driver calls - * spark_ds_commit with the messages of the tasks if they all succeeded, or spark_ds_abort - * otherwise, including when spark_ds_commit fails, and closes the writer. Several attempts of - * a task may write the same partition, after a failure or with speculative execution, each - * with a data writer of its own. Spark lets the attempts of a partition commit one at a time, - * and passes to spark_ds_commit the message of one committed attempt for each partition, so - * the write consists of the data of these attempts only. It aborts the data writers of the - * other attempts, except on an executor that is lost, where a data writer is neither committed - * nor aborted, and may even have committed before its message reached the driver. A streaming - * write commits or aborts each micro-batch (epoch). Spark may commit the same epoch more than - * once, for example when a streaming query restarts after a failure, so committing an epoch - * that is already committed must have no effect, for exactly-once output. The ID of the query, - * which spark_ds_create_stream_writer gets, stays the same when the query restarts from its - * checkpoint, so the ID of the query and the epoch ID identify the output of a micro-batch, - * also when other queries write to the same destination. A micro-batch streaming query creates - * a new stream writer for each micro-batch, and Spark closes the previous stream writer of the - * query when the query creates a new one. Spark is not told when a streaming query stops, so it - * closes the last stream writer of the query once the writer is no longer referenced, or when - * the query is restarted. + * Spark creates a writer with create_writer, and serializes its state with serialize_writer. + * With adaptive query execution, Spark first runs the stages of the data to write that end with + * a shuffle, and creates no writer if they fail. The rest of the data, such as its last stage, + * is computed in the tasks of the write, so a writer may still be aborted because computing the + * data failed: do irreversible work, such as removing the existing data for an overwrite, in + * commit. On the executors, each task creates a data writer with create_data_writer, calls write + * for each batch of at most spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends + * once it holds spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls commit_data_writer, + * or abort_data_writer if the task or the commit fails, before it closes the data writer. Then + * the driver calls commit with the messages of the tasks if they all succeeded, or abort + * otherwise, including when commit fails, and closes the writer. Several attempts of a task may + * write the same partition, after a failure or with speculative execution, each with a data + * writer of its own. Spark lets the attempts of a partition commit one at a time, and passes to + * commit the message of one committed attempt for each partition, so the write consists of the + * data of these attempts only. It aborts the data writers of the other attempts, except on an + * executor that is lost, where a data writer is neither committed nor aborted, and may even have + * committed before its message reached the driver. A streaming write commits or aborts each + * micro-batch (epoch). Spark may commit the same epoch more than once, for example when a + * streaming query restarts after a failure, so committing an epoch that is already committed + * must have no effect, for exactly-once output. The ID of the query, which create_writer gets for + * a streaming write, stays the same when the query restarts from its checkpoint, so the ID of the + * query and the epoch ID identify the output of a micro-batch, also when other queries write to + * the same destination. A micro-batch streaming query creates a new stream writer for each + * micro-batch, and Spark closes the previous stream writer of the query when the query creates a + * new one. Spark is not told when a streaming query stops, so it closes the last stream writer of + * the query once the writer is no longer referenced, or when the query is restarted. * * Conventions * ----------- * - Errors. The functions that return an int return 0 on success, or any other value on failure. - * They then set the message that spark_ds_last_error returns, and Spark reports it as a + * They then set the message that last_error returns, and Spark reports it as a * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail, and they may run after the query, - * from another thread: do the work that can fail a query, such as flushing data, in - * spark_ds_commit. A C++ exception or a Rust panic must not cross the interface: catch it and - * return an error instead. - * - Handles. The create functions return an opaque, non-NULL handle in their last argument, - * typically a pointer: Spark reports a NULL handle as an error. Spark releases each handle - * exactly once with the matching close function, or with spark_ds_commit_data_writer or - * spark_ds_abort_data_writer for data writers, though not always right after its last use: a + * from another thread: do the work that can fail a query, such as flushing data, in commit. A + * C++ exception or a Rust panic must not cross the interface: catch it and return an error + * instead. + * - Handles. create, create_reader, create_stream_reader, create_writer and create_data_writer + * return an opaque, non-NULL handle in their last argument, typically a pointer: Spark reports a + * NULL handle as an error. Spark releases each handle exactly once, with the close function of + * its kind: close, close_reader, close_stream_reader, close_writer or close_data_writer, and no + * other function releases a handle. A handle is not always released right after its last use: a * batch reader whose partitions are never planned and the last stream writer of a query are * released once they are no longer referenced (see Lifecycle above). A reader or writer handle - * must not depend on the data - * source handle it was created from, because Spark closes the data source handle right after - * creating it. - * - Memory. What Spark passes to a function, such as strings and byte arrays, is only valid - * during the call. Strings are UTF-8 and NUL-terminated. The JSON of the predicates is ASCII: - * it escapes the other characters, which any JSON parser decodes. The library returns strings - * and byte arrays by passing them to a SparkDsSink, which copies them. + * must not depend on the data source handle it was created from, because Spark closes the data + * source handle right after creating it. + * - Memory. What Spark passes to a function, such as strings, byte arrays and SparkDsWriteInfo, is + * only valid during the call, except for SparkDsHost and SparkDsTaskContext, which document how + * long they live. Strings are UTF-8 and NUL-terminated. The JSON of the predicates is ASCII: it + * escapes the other characters, which any JSON parser decodes. The library returns strings and + * byte arrays by passing them to a SparkDsSink, which copies them. * - State. Partitions, the state of readers and writers, and commit messages are byte arrays in a * format defined by the library. Spark sends them to the executors and back. * - Arrow data. Schemas and data are exchanged with the Arrow C data interface @@ -138,30 +137,36 @@ * large lists, fixed-size lists and list views, and dictionary encoding. * - Options. Option keys are passed in lower case, because Spark matches them * case-insensitively. - * - Optional functions. A library only needs to export the functions for the operations it - * supports. If a function is missing, Spark treats the operation as unsupported, except for the - * pushdown functions and the serialize functions, which then push nothing down and return no - * state, spark_ds_commit_offset, which then does nothing, and spark_ds_last_error. - * spark_ds_abi_version is required. + * - Optional functions. A library only sets the functions of the operations it supports in its + * table, and leaves the others NULL. A data source must not declare a capability whose + * functions are NULL: Spark then fails the operation. The pushdown functions and the serialize + * functions are optional: without them, Spark pushes nothing down, and the reader or writer has + * no state. So are schema, without which the user must specify the schema, commit_offset and + * last_error. * - Threads. Spark calls the library from multiple threads, and calls on different handles may * run concurrently. Calls on the same handle never run concurrently. A handle may be released * from a thread other than the threads that used it, such as a JVM Cleaner thread when Spark * releases a handle once it is no longer referenced, for example the last stream writer of a * streaming query. A SparkDsSink must be called on the thread of the call that it was passed - * to, during that call. + * to, during that call. The functions of SparkDsHost and SparkDsTaskContext can be called from + * any thread. * * Compatibility * ------------- - * Spark binds the functions of a library by their names, so the names of the functions and their - * signatures are the interface. Within a SPARK_DS_ABI_VERSION: - * - the signature and the semantics of an existing function never change; - * - new functions may be added, only as optional functions, so that a library that does not - * export them keeps working as before; + * The function table and the structs that Spark passes to the library start with their size, so + * that they can grow. Within a SPARK_DS_ABI_VERSION: + * - fields are only appended to the structs, including SparkDsLibrary, and the existing fields + * never change, nor do the semantics of the existing functions; + * - Spark only reads the functions that fit in the size of the table of a library, so a library + * built against an earlier header keeps working, with the functions that it does not know as + * NULL. A library must only read the fields that fit in the size of a struct that Spark passes; * - new values may appear where this documentation says so, such as the names of pushed - * predicates. - * An incompatible change uses a new SPARK_DS_ABI_VERSION, with new names. A library reports the - * version it implements with spark_ds_abi_version, and Spark rejects a library of a version it - * does not support. Spark supports version 1 at least in all the releases of Spark 4 from 4.4.0. + * predicates. Spark only uses a new capability or write mode with a data source that declares + * it. + * An incompatible change uses a new SPARK_DS_ABI_VERSION: spark_ds_init gets the newest version + * that Spark supports, the library returns a table of a version that it implements, and Spark + * rejects a table of a version that it does not support. Spark supports version 1 at least in all + * the releases of Spark 4 from 4.4.0. */ #ifndef SPARK_DATASOURCE_H #define SPARK_DATASOURCE_H @@ -176,6 +181,13 @@ extern "C" { /* The version of this interface. */ #define SPARK_DS_ABI_VERSION 1 +/* Exports spark_ds_init from the library, also when the library hides its other symbols. */ +#if defined(_WIN32) +#define SPARK_DS_EXPORT __declspec(dllexport) +#else +#define SPARK_DS_EXPORT __attribute__((visibility("default"))) +#endif + /* The Arrow C data interface: https://arrow.apache.org/docs/format/CDataInterface.html */ #ifndef ARROW_C_DATA_INTERFACE #define ARROW_C_DATA_INTERFACE @@ -230,318 +242,425 @@ struct ArrowArrayStream { */ typedef void (*SparkDsSink)(void* sink, const uint8_t* data, int64_t length); -/* - * Returns the message of the last error on the calling thread, as a NUL-terminated UTF-8 string - * that stays valid until the next call on that thread. Optional: without it, Spark only reports - * the value that the function returned. - */ -const char* spark_ds_last_error(void); - -/* Returns the version of this interface that the library implements, SPARK_DS_ABI_VERSION. */ -int32_t spark_ds_abi_version(void); - /* --------------------------------------------------------------------------------------------- - * Data source (driver) + * What Spark passes to the library * --------------------------------------------------------------------------------------------- */ -/* - * Creates a data source with the options given by the user, such as - * spark.read.option("key", "value"). - * - * name: the name of the data source, as listed in the package manifest, or in lower case for an - * installed library. A library can implement several data sources. - * option_keys, option_values: the num_options options, with the keys in lower case - * source: set to the data source handle - */ -int spark_ds_create(const char* name, const char* const* option_keys, - const char* const* option_values, int64_t num_options, void** source); - -/* - * Exports the schema of the data. Spark calls it only when the user does not specify a schema. - * - * out: initialized with a struct type whose fields are the columns - */ -int spark_ds_schema(void* source, struct ArrowSchema* out); - -/* - * Creates a reader for a batch scan. - * - * schema: the schema to read: the one the user specified, or the one returned by spark_ds_schema - * reader: set to the reader handle - */ -int spark_ds_create_reader(void* source, struct ArrowSchema* schema, void** reader); - -/* - * Creates a reader for a micro-batch streaming scan. - * - * schema: the schema to read: the one the user specified, or the one returned by spark_ds_schema - * reader: set to the stream reader handle - */ -int spark_ds_create_stream_reader(void* source, struct ArrowSchema* schema, void** reader); - -/* - * Creates a writer for a batch write. - * - * schema: the schema of the data to write - * overwrite: whether to replace the existing data (save mode "overwrite") instead of appending to - * it - * writer: set to the writer handle - */ -int spark_ds_create_writer(void* source, struct ArrowSchema* schema, bool overwrite, - void** writer); +/* The levels of the messages that a library logs with SparkDsHost.log. */ +#define SPARK_DS_LOG_ERROR 1 +#define SPARK_DS_LOG_WARN 2 +#define SPARK_DS_LOG_INFO 3 +#define SPARK_DS_LOG_DEBUG 4 +#define SPARK_DS_LOG_TRACE 5 + +/* Spark, as spark_ds_init gets it. It stays valid until the JVM exits. */ +typedef struct SparkDsHost { + /* + * The size of this struct in the header that Spark was built with. A later version of Spark may + * append fields, so a library must only read the fields that fit in it. + */ + int64_t size; + /* The newest version of this interface that Spark supports. */ + int32_t abi_version; + /* The version of Spark, such as "4.4.0". */ + const char* spark_version; + /* + * Logs a message, a NUL-terminated UTF-8 string, at one of the SPARK_DS_LOG_* levels: in the log + * of the driver or of the executor, with the path of the library. + */ + void (*log)(int32_t level, const char* message); +} SparkDsHost; /* - * Creates a writer for a streaming write. - * - * schema: the schema of the data to write - * overwrite: whether to replace the existing data in each micro-batch (output mode "complete") - * instead of appending to it - * query_id: the ID of the streaming query. It stays the same when the query restarts from its - * checkpoint, so with the epoch IDs, it identifies the commits of the query: see Writes - * above. - * writer: set to the writer handle + * The task that calls read or create_data_writer, on an executor. It stays valid until the stream + * that read returns is released, or until the data writer is closed. */ -int spark_ds_create_stream_writer(void* source, struct ArrowSchema* schema, bool overwrite, - const char* query_id, void** writer); - -/* Releases a data source handle. */ -void spark_ds_close(void* source); +typedef struct SparkDsTaskContext { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* The ID of the stage of the task, and the attempt number of the stage. */ + int32_t stage_id; + int32_t stage_attempt_number; + /* The partition that the task reads or writes, and the attempt number of the task. */ + int32_t partition_id; + int32_t attempt_number; + /* The ID of the task attempt, unique in the Spark application. */ + int64_t task_attempt_id; + /* + * Returns whether the task was killed, for example because its query was canceled. Spark cannot + * interrupt native code, so a read or a write that runs for long, or blocks, should check it + * regularly, and fail once it is true. + */ + bool (*is_interrupted)(const struct SparkDsTaskContext* task); +} SparkDsTaskContext; + +/* The modes of a write. */ +#define SPARK_DS_WRITE_APPEND 0 +#define SPARK_DS_WRITE_TRUNCATE 1 + +/* A write, as create_writer gets it. */ +typedef struct SparkDsWriteInfo { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* + * SPARK_DS_WRITE_APPEND to add the data to the existing data, or SPARK_DS_WRITE_TRUNCATE to + * replace the existing data: with the save mode "overwrite" of a batch write, and in each + * micro-batch with the output mode "complete" of a streaming write. Spark only passes + * SPARK_DS_WRITE_TRUNCATE to a data source with the capability SPARK_DS_TRUNCATE. + */ + int32_t mode; + /* + * The ID of the streaming query of a streaming write, or NULL for a batch write. It stays the + * same when the query restarts from its checkpoint, so with the epoch IDs, it identifies the + * commits of the query: see Writes above. + */ + const char* query_id; +} SparkDsWriteInfo; /* --------------------------------------------------------------------------------------------- - * Batch reader (driver) + * What the library returns to Spark * --------------------------------------------------------------------------------------------- */ -/* - * Pushes predicates down to the reader. The predicates are combined with AND. - * - * Each predicate is a JSON expression tree. An expression is one of: - * - {"type": "column", "name": ["a", "b"]}: a column. A nested field has a name part for each - * level. - * - {"type": "literal", "dataType": "bigint", "value": 1}: a literal. The data type is a Spark SQL - * type name such as int, string or decimal(10,2). The value is null, a JSON boolean or number - * for boolean and numeric types, a string for string, a base64 string for binary, a decimal - * string for decimal, the days since the epoch for date, and the microseconds since the epoch - * for timestamp and timestamp_ntz. - * - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in - * org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=> (equal, - * with null equal to null), <, <=, > and >= with two children, IN with the value to look for - * and then the values to look in, IS_NULL and IS_NOT_NULL with one child, STARTS_WITH, - * ENDS_WITH and CONTAINS with the string and the string to look for, AND and OR with two - * predicates, and NOT with one. The children of the other functions are columns, literals or - * predicates. - * Spark does not pass predicates that it cannot express this way, such as predicates that use - * arithmetic or other functions, nor predicates on columns that contain strings with a - * non-binary collation, such as UTF8_LCASE: strings are compared by their UTF-8 bytes. A later - * version of Spark may pass predicates with other names without changing SPARK_DS_ABI_VERSION, - * so a reader must ignore a predicate that contains a name that it does not know, or a data type - * that it does not support. - * - * Spark evaluates all the predicates again after reading, so a reader may use them only to skip - * data, for example the files whose statistics show that no row matches. But it must never drop - * or skip a row for which the predicates are true under the semantics of Spark SQL, which Spark - * cannot detect. In particular: - * - For floating-point columns, NaN is equal to NaN and greater than any other value, including - * positive infinity, and -0.0 is equal to 0.0. IEEE 754 comparisons treat NaN differently, and - * total orders, such as f64::total_cmp in Rust or the comparison kernels of arrow-rs, also - * order -0.0 before 0.0. For example, d > 1.0 is true for a NaN value of d, and d = 0.0 for - * -0.0. - * - Comparisons with null are null, and a row only matches if the predicates are true. For - * example, NOT (i IN (1, NULL)) matches no row. - * - Strings are compared by their UTF-8 bytes, which order the characters outside the Basic - * Multilingual Plane after all the others, unlike comparisons of UTF-16 strings. - * - To use only a part of a predicate, a reader may take the parts that it does not support as - * true, but not under a NOT, where that would drop rows. - * The file org/apache/spark/sql/datasource/predicate-cases.json in the jar of Spark SQL lists - * predicates as Spark passes them to this function, with the rows of a table for which they are - * true, to test a reader with. - * - * predicates: the count predicates, as JSON - * evaluated: an array of count values, all false, to set to whether the reader evaluates each - * predicate completely. Spark shows these predicates as pushed in the plan of the - * query, and evaluates all the predicates again after reading anyway. - */ -int spark_ds_push_predicates(void* reader, const char* const* predicates, int64_t count, - bool* evaluated); - -/* - * Pushes a LIMIT down to the reader. Spark only pushes a limit when the query does not filter the - * rows, since it evaluates the predicates after reading. - * - * applied: set to whether the reader returns at most `limit` rows in total. Spark applies the - * limit again after reading either way. - */ -int spark_ds_push_limit(void* reader, int32_t limit, bool* applied); - -/* - * Prunes the columns to read. - * - * columns: the names of the count top-level columns that the query needs, which include the - * columns of the pushed predicates. There may be none, for example for - * SELECT count(*): then the batches only have a number of rows. - * pruned: set to whether spark_ds_read returns exactly these columns, in this order. If not, it - * returns all the columns of the schema to read. - */ -int spark_ds_prune_columns(void* reader, const char* const* columns, int64_t count, - bool* pruned); +/* The capabilities of a data source, which create returns: the operations that it supports. */ +#define SPARK_DS_BATCH_READ 1 /* batch reads, with create_reader */ +#define SPARK_DS_MICRO_BATCH_READ 2 /* micro-batch streaming reads, with create_stream_reader */ +#define SPARK_DS_BATCH_WRITE 4 /* batch writes, with create_writer */ +#define SPARK_DS_STREAMING_WRITE 8 /* streaming writes, with create_writer */ +#define SPARK_DS_TRUNCATE 16 /* writes with the mode SPARK_DS_WRITE_TRUNCATE */ -/* - * Plans the partitions of the scan. Spark reads each partition in a separate task. - * - * add_partition: called with each partition, serialized in a format defined by the library - */ -int spark_ds_partitions(void* reader, SparkDsSink add_partition, void* sink); +/* How a reader uses a predicate pushed with push_predicates. */ +#define SPARK_DS_PREDICATE_UNUSED 0 /* not at all */ +/* only to skip data, so it may return rows for which the predicate is not true */ +#define SPARK_DS_PREDICATE_PARTIAL 1 +/* to only return the rows for which the predicate is true */ +#define SPARK_DS_PREDICATE_COMPLETE 2 /* - * Serializes the state of the reader that spark_ds_read needs on the executors, such as the - * pushed predicates. Spark calls it after spark_ds_partitions. - * - * set_state: called once with the state, or not at all if there is none + * The functions of a library, which spark_ds_init returns. A library sets the functions of the + * operations it supports, and leaves the others NULL: see "Optional functions" above. */ -int spark_ds_serialize_reader(void* reader, SparkDsSink set_state, void* sink); - -/* Releases a batch reader handle. */ -void spark_ds_close_reader(void* reader); - -/* --------------------------------------------------------------------------------------------- - * Stream reader (driver) - * --------------------------------------------------------------------------------------------- */ +typedef struct SparkDsLibrary { + /* + * The size of this struct in the header that the library was built with. Spark only reads the + * functions that fit in it. + */ + int64_t size; + /* The version of this interface that the table implements, SPARK_DS_ABI_VERSION. */ + int32_t abi_version; + + /* + * Returns the message of the last error on the calling thread, as a NUL-terminated UTF-8 string + * that stays valid until the next call on that thread. Optional: without it, Spark only reports + * the value that the function returned. + */ + const char* (*last_error)(void); + + /* ------------------------------------------------------------------------------------------- + * Data source (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Creates a data source with the options given by the user, such as + * spark.read.option("key", "value"). + * + * name: the name of the data source, as listed in the package manifest, or in lower case for an + * installed library. A library can implement several data sources. + * option_keys, option_values: the num_options options, with the keys in lower case + * capabilities: set to the capabilities of the data source, SPARK_DS_* flags combined with |, + * which may depend on its name and its options. It is 0 when create is called. + * source: set to the data source handle + */ + int (*create)(const char* name, const char* const* option_keys, + const char* const* option_values, int64_t num_options, uint64_t* capabilities, + void** source); + + /* + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. + * + * out: initialized with a struct type whose fields are the columns + */ + int (*schema)(void* source, struct ArrowSchema* out); + + /* + * Creates a reader for a batch scan. + * + * schema: the schema to read: the one the user specified, or the one returned by schema + * reader: set to the reader handle + */ + int (*create_reader)(void* source, struct ArrowSchema* schema, void** reader); + + /* + * Creates a reader for a micro-batch streaming scan. + * + * schema: the schema to read: the one the user specified, or the one returned by schema + * reader: set to the stream reader handle + */ + int (*create_stream_reader)(void* source, struct ArrowSchema* schema, void** reader); + + /* + * Creates a writer for a batch write, or for a micro-batch of a streaming write. + * + * schema: the schema of the data to write + * info: the mode of the write, and the ID of the query of a streaming write + * writer: set to the writer handle + */ + int (*create_writer)(void* source, struct ArrowSchema* schema, const SparkDsWriteInfo* info, + void** writer); + + /* Releases a data source handle. */ + void (*close)(void* source); + + /* ------------------------------------------------------------------------------------------- + * Batch reader (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Pushes predicates down to the reader. The predicates are combined with AND. + * + * Each predicate is a JSON expression tree. An expression is one of: + * - {"type": "column", "name": ["a", "b"]}: a column. A nested field has a name part for each + * level. + * - {"type": "literal", "dataType": "bigint", "value": 1}: a literal. The data type is a Spark + * SQL type name such as int, string or decimal(10,2). The value is null, a JSON boolean or + * number for boolean and numeric types, a string for string, a base64 string for binary, a + * decimal string for decimal, the days since the epoch for date, and the microseconds since + * the epoch for timestamp and timestamp_ntz. + * - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in + * org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=> + * (equal, with null equal to null), <, <=, > and >= with two children, IN with the value to + * look for and then the values to look in, IS_NULL and IS_NOT_NULL with one child, + * STARTS_WITH, ENDS_WITH and CONTAINS with the string and the string to look for, AND and OR + * with two predicates, and NOT with one. The children of the other functions are columns, + * literals or predicates. + * Spark does not pass predicates that it cannot express this way, such as predicates that use + * arithmetic or other functions, nor predicates on columns that contain strings with a + * non-binary collation, such as UTF8_LCASE: strings are compared by their UTF-8 bytes. A later + * version of Spark may pass predicates with other names without changing SPARK_DS_ABI_VERSION, + * so a reader must ignore a predicate that contains a name that it does not know, or a data + * type that it does not support. + * + * Spark evaluates all the predicates again after reading, so a reader may use them only to skip + * data, for example the files whose statistics show that no row matches. But it must never drop + * or skip a row for which the predicates are true under the semantics of Spark SQL, which Spark + * cannot detect. In particular: + * - For floating-point columns, NaN is equal to NaN and greater than any other value, including + * positive infinity, and -0.0 is equal to 0.0. IEEE 754 comparisons treat NaN differently, and + * total orders, such as f64::total_cmp in Rust or the comparison kernels of arrow-rs, also + * order -0.0 before 0.0. For example, d > 1.0 is true for a NaN value of d, and d = 0.0 for + * -0.0. + * - Comparisons with null are null, and a row only matches if the predicates are true. For + * example, NOT (i IN (1, NULL)) matches no row. + * - Strings are compared by their UTF-8 bytes, which order the characters outside the Basic + * Multilingual Plane after all the others, unlike comparisons of UTF-16 strings. + * - To use only a part of a predicate, a reader may take the parts that it does not support as + * true, but not under a NOT, where that would drop rows. + * The file org/apache/spark/sql/datasource/predicate-cases.json in the jar of Spark SQL lists + * predicates as Spark passes them to this function, with the rows of a table for which they are + * true, to test a reader with. + * + * predicates: the count predicates, as JSON + * uses: an array of count values, all SPARK_DS_PREDICATE_UNUSED, to set to how the reader uses + * each predicate, a SPARK_DS_PREDICATE_* value. Spark shows the predicates that the reader + * uses as pushed in the plan of the query, and evaluates all the predicates again after + * reading anyway. It takes other values as SPARK_DS_PREDICATE_PARTIAL. + */ + int (*push_predicates)(void* reader, const char* const* predicates, int64_t count, + int32_t* uses); + + /* + * Pushes a LIMIT down to the reader. Spark only pushes a limit when the query does not filter + * the rows, since it evaluates the predicates after reading. + * + * applied: set to whether the reader returns at most `limit` rows in total. Spark applies the + * limit again after reading either way. + */ + int (*push_limit)(void* reader, int32_t limit, bool* applied); + + /* + * Prunes the columns to read. + * + * columns: the names of the count top-level columns that the query needs, which include the + * columns of the pushed predicates. There may be none, for example for + * SELECT count(*): then the batches only have a number of rows. + * pruned: set to whether read returns exactly these columns, in this order. If not, it returns + * all the columns of the schema to read. + */ + int (*prune_columns)(void* reader, const char* const* columns, int64_t count, bool* pruned); + + /* + * Plans the partitions of the scan. Spark reads each partition in a separate task. + * + * add_partition: called with each partition, serialized in a format defined by the library + */ + int (*partitions)(void* reader, SparkDsSink add_partition, void* sink); + + /* + * Serializes the state of the reader that read needs on the executors, such as the pushed + * predicates. Spark calls it after partitions. + * + * set_state: called once with the state, or not at all if there is none + */ + int (*serialize_reader)(void* reader, SparkDsSink set_state, void* sink); + + /* Releases a batch reader handle. */ + void (*close_reader)(void* reader); + + /* ------------------------------------------------------------------------------------------- + * Stream reader (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Returns the offset to start from when the streaming query has no checkpoint. + * + * set_offset: called once with the offset, as compact JSON on a single line + */ + int (*initial_offset)(void* reader, SparkDsSink set_offset, void* sink); + + /* + * Returns the most recent offset available. + * + * set_offset: called once with the offset, as compact JSON on a single line + */ + int (*latest_offset)(void* reader, SparkDsSink set_offset, void* sink); + + /* + * Plans the partitions that read the data after `start` up to and including `end`, two offsets + * as JSON. + * + * add_partition: called with each partition, serialized in a format defined by the library + */ + int (*stream_partitions)(void* reader, const char* start, const char* end, + SparkDsSink add_partition, void* sink); + + /* + * Serializes the state of the stream reader that read needs on the executors. Spark calls it + * after each call to stream_partitions. + * + * set_state: called once with the state, or not at all if there is none + */ + int (*serialize_stream_reader)(void* reader, SparkDsSink set_state, void* sink); + + /* + * Tells the library that Spark has processed all the data up to and including `end`, an offset + * as JSON, and will not ask for it again. + */ + int (*commit_offset)(void* reader, const char* end); + + /* Releases a stream reader handle, when the streaming query stops. */ + void (*close_stream_reader)(void* reader); + + /* ------------------------------------------------------------------------------------------- + * Partition reader (executor) + * ------------------------------------------------------------------------------------------- */ + + /* + * Reads a partition of a batch or streaming scan. + * + * task: the task that reads the partition, valid until the stream is released + * reader_state: the state returned by serialize_reader or serialize_stream_reader, empty if + * there is none + * partition: a partition returned by partitions or stream_partitions + * out: the stream to initialize. Each array of the stream is a struct array whose fields are the + * columns to read: the columns of the schema to read, or the pruned columns. Spark calls the + * stream from the thread of the task that reads the partition. get_next may block, for + * example on network I/O, but Spark cannot interrupt native code: when the task is killed, + * Spark only stops reading once get_next returns, and then releases the stream. So a + * get_next that blocks forever keeps the task running, unless the task reaper + * (spark.task.reaper.enabled and spark.task.reaper.killTimeout) kills the executor. Use + * timeouts for blocking operations, and check task->is_interrupted. + */ + int (*read)(const SparkDsTaskContext* task, const uint8_t* reader_state, + int64_t reader_state_length, const uint8_t* partition, int64_t partition_length, + struct ArrowArrayStream* out); + + /* ------------------------------------------------------------------------------------------- + * Writer (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Serializes the state of the writer that create_data_writer needs on the executors. Spark + * calls it once, after creating the writer. + * + * set_state: called once with the state, or not at all if there is none + */ + int (*serialize_writer)(void* writer, SparkDsSink set_state, void* sink); + + /* + * Commits a write, when all the tasks succeeded. + * + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * messages, message_lengths: the num_messages messages returned by commit_data_writer of each + * task + */ + int (*commit)(void* writer, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages); + + /* + * Aborts a write, when it failed. + * + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * messages, message_lengths: the num_messages messages returned by commit_data_writer of each + * task. A message is NULL if its task did not commit. + */ + int (*abort)(void* writer, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages); + + /* Releases a writer handle. */ + void (*close_writer)(void* writer); + + /* ------------------------------------------------------------------------------------------- + * Data writer (executor) + * ------------------------------------------------------------------------------------------- */ + + /* + * Creates the writer for one task. + * + * task: the task, with the partition that it writes, valid until the data writer is closed + * writer_state: the state returned by serialize_writer, empty if there is none + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * data_writer: set to the data writer handle + */ + int (*create_data_writer)(const SparkDsTaskContext* task, const uint8_t* writer_state, + int64_t writer_state_length, int64_t epoch_id, void** data_writer); + + /* + * Writes a batch of data. + * + * array: a struct array whose fields are the columns + * schema: the schema of the array + */ + int (*write)(void* data_writer, struct ArrowArray* array, struct ArrowSchema* schema); + + /* + * Commits the data written by the task. If it fails, Spark calls abort_data_writer. If it + * succeeds, Spark does not call abort_data_writer, even if it then fails the task, for example + * because set_message was called more than once. + * + * set_message: called once with the message to pass to commit on the driver, or not at all for + * an empty message + */ + int (*commit_data_writer)(void* data_writer, SparkDsSink set_message, void* sink); + + /* Aborts the data written by the task. */ + int (*abort_data_writer)(void* data_writer); + + /* Releases a data writer handle, after its commit or abort. */ + void (*close_data_writer)(void* data_writer); +} SparkDsLibrary; /* - * Returns the offset to start from when the streaming query has no checkpoint. + * The function that a library exports. Spark calls it once in each JVM, when it loads the library, + * before any other function. * - * set_offset: called once with the offset, as compact JSON on a single line - */ -int spark_ds_initial_offset(void* reader, SparkDsSink set_offset, void* sink); - -/* - * Returns the most recent offset available. + * host: Spark, valid until the JVM exits + * library: set to the function table of the library, for the version of this interface that it + * implements, or the newest one up to host->abi_version if it implements several. Spark + * reads the table after the call, so it is typically the address of a static struct. * - * set_offset: called once with the offset, as compact JSON on a single line + * Returns 0 on success, or any other value if the library cannot run, after it logs why with + * host->log. */ -int spark_ds_latest_offset(void* reader, SparkDsSink set_offset, void* sink); - -/* - * Plans the partitions that read the data after `start` up to and including `end`, two offsets - * as JSON. - * - * add_partition: called with each partition, serialized in a format defined by the library - */ -int spark_ds_stream_partitions(void* reader, const char* start, const char* end, - SparkDsSink add_partition, void* sink); - -/* - * Serializes the state of the stream reader that spark_ds_read needs on the executors. Spark - * calls it after each call to spark_ds_stream_partitions. - * - * set_state: called once with the state, or not at all if there is none - */ -int spark_ds_serialize_stream_reader(void* reader, SparkDsSink set_state, void* sink); - -/* - * Tells the library that Spark has processed all the data up to and including `end`, an offset - * as JSON, and will not ask for it again. - */ -int spark_ds_commit_offset(void* reader, const char* end); - -/* Releases a stream reader handle, when the streaming query stops. */ -void spark_ds_close_stream_reader(void* reader); - -/* --------------------------------------------------------------------------------------------- - * Partition reader (executor) - * --------------------------------------------------------------------------------------------- */ - -/* - * Reads a partition of a batch or streaming scan. - * - * reader_state: the state returned by spark_ds_serialize_reader or - * spark_ds_serialize_stream_reader, empty if there is none - * partition: a partition returned by spark_ds_partitions or spark_ds_stream_partitions - * out: the stream to initialize. Each array of the stream is a struct array whose fields are the - * columns to read: the columns of the schema to read, or the pruned columns. Spark calls the - * stream from the thread of the task that reads the partition. get_next may block, for - * example on network I/O, but Spark cannot interrupt native code: when the task is killed, - * Spark only stops reading once get_next returns, and then releases the stream. So a - * get_next that blocks forever keeps the task running, unless the task reaper - * (spark.task.reaper.enabled and spark.task.reaper.killTimeout) kills the executor. Use - * timeouts for blocking operations. - */ -int spark_ds_read(const uint8_t* reader_state, int64_t reader_state_length, - const uint8_t* partition, int64_t partition_length, - struct ArrowArrayStream* out); - -/* --------------------------------------------------------------------------------------------- - * Writer (driver) - * --------------------------------------------------------------------------------------------- */ - -/* - * Serializes the state of the writer that spark_ds_create_data_writer needs on the executors. - * Spark calls it once, after creating the writer. - * - * set_state: called once with the state, or not at all if there is none - */ -int spark_ds_serialize_writer(void* writer, SparkDsSink set_state, void* sink); - -/* - * Commits a write, when all the tasks succeeded. - * - * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write - * messages, message_lengths: the num_messages messages returned by spark_ds_commit_data_writer - * of each task - */ -int spark_ds_commit(void* writer, int64_t epoch_id, const uint8_t* const* messages, - const int64_t* message_lengths, int64_t num_messages); - -/* - * Aborts a write, when it failed. - * - * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write - * messages, message_lengths: the num_messages messages returned by spark_ds_commit_data_writer - * of each task. A message is NULL if its task did not commit. - */ -int spark_ds_abort(void* writer, int64_t epoch_id, const uint8_t* const* messages, - const int64_t* message_lengths, int64_t num_messages); - -/* Releases a writer handle. */ -void spark_ds_close_writer(void* writer); - -/* --------------------------------------------------------------------------------------------- - * Data writer (executor) - * --------------------------------------------------------------------------------------------- */ - -/* - * Creates the writer for one task. - * - * writer_state: the state returned by spark_ds_serialize_writer, empty if there is none - * partition_id: the partition that the task writes - * task_id: the ID of the task attempt - * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write - * data_writer: set to the data writer handle - */ -int spark_ds_create_data_writer(const uint8_t* writer_state, int64_t writer_state_length, - int32_t partition_id, int64_t task_id, int64_t epoch_id, - void** data_writer); - -/* - * Writes a batch of data. - * - * array: a struct array whose fields are the columns - * schema: the schema of the array - */ -int spark_ds_write(void* data_writer, struct ArrowArray* array, struct ArrowSchema* schema); - -/* - * Commits the data written by the task and releases the data writer handle. If it fails, the - * handle is not released: Spark then calls spark_ds_abort_data_writer with it, to abort the data - * written by the task and release the handle. If it succeeds, Spark does not use the handle again, - * even if it then fails the task, for example because set_message was called more than once. - * - * set_message: called once with the message to pass to spark_ds_commit on the driver, or not at - * all for an empty message - */ -int spark_ds_commit_data_writer(void* data_writer, SparkDsSink set_message, void* sink); - -/* Aborts the data written by the task and releases the data writer handle. */ -int spark_ds_abort_data_writer(void* data_writer); +SPARK_DS_EXPORT int spark_ds_init(const SparkDsHost* host, const SparkDsLibrary** library); #ifdef __cplusplus } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala index d1bb7705c33f2..f739fdfe3e3d0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala @@ -68,10 +68,11 @@ trait ColumnarDataSource { trait ColumnarReader extends Serializable { /** - * Returns the predicates that the reader does not evaluate completely. Spark evaluates all the - * predicates after reading anyway, and shows the others as pushed in the plan. + * Returns the predicates that the reader uses, to skip data or to only return the rows for which + * they are true. Spark evaluates all the predicates after reading anyway, and shows the returned + * ones as pushed in the plan. */ - def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = predicates + def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = Array.empty /** Returns whether the reader returns at most `limit` rows. Spark applies it again anyway. */ def pushLimit(limit: Int): Boolean = false diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala index 808025d01b7fd..1393713a78f11 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala @@ -103,8 +103,7 @@ class ColumnarScanBuilder( } override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { - val remaining = batchReader.pushPredicates(predicates) - pushed = predicates.filterNot(remaining.contains) + pushed = batchReader.pushPredicates(predicates) // Spark evaluates all the predicates again after reading, even the ones that the reader // evaluates completely: Spark SQL compares values in ways that a data source can easily get // wrong, such as NaN with other values. The pushed predicates are shown in the plan. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index 55212ea4d7c08..3edb430366e8e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -31,7 +31,7 @@ import org.apache.arrow.vector.{FieldVector, ValueVector, VectorSchemaRoot} import org.apache.arrow.vector.complex.MapVector import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} -import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException, TaskContext} import org.apache.spark.internal.Logging import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} @@ -65,7 +65,7 @@ class NativeDataSource( val library = NativeLibraries.get(location) withDataSource(library, "infer the schema of") { dataSource => call(name, "infer the schema of") { - optional(NativeArrow.importSchema(address => library.schema(dataSource, address))) + optional(NativeArrow.importSchema(address => library.schema(dataSource.handle, address))) }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) } } @@ -73,9 +73,12 @@ class NativeDataSource( override def reader(schema: StructType): ColumnarReader = { val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) - val handle = create(library, "plan a scan of", "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { + val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ, + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) { dataSource => - NativeArrow.exportSchema(schema)(address => library.createReader(dataSource, address)) + NativeArrow.exportSchema(schema) { address => + library.createReader(dataSource.handle, address) + } } new NativeDataSourceReader( executorLocation, name, schema, new NativeHandle(handle, library.closeReader)) @@ -84,9 +87,12 @@ class NativeDataSource( override def streamReader(schema: StructType): ColumnarStreamReader = { val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) - val handle = create(library, "plan a scan of", "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { + val handle = create(library, "plan a scan of", NativeLibrary.MICRO_BATCH_READ, + "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) { dataSource => - NativeArrow.exportSchema(schema)(address => library.createStreamReader(dataSource, address)) + NativeArrow.exportSchema(schema) { address => + library.createStreamReader(dataSource.handle, address) + } } new NativeDataSourceStreamReader( executorLocation, name, schema, new NativeHandle(handle, library.closeStreamReader)) @@ -96,12 +102,7 @@ class NativeDataSource( val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) val handle = new NativeHandle( - create(library, "plan a write to", "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource => - NativeArrow.exportSchema(schema) { address => - library.createWriter(dataSource, address, overwrite) - } - }, - library.closeWriter) + createWriter(library, schema, overwrite, queryId = null), library.closeWriter) new NativeDataSourceWriter(executorLocation, name, handle, serializeWriter(library, handle)) } @@ -111,14 +112,8 @@ class NativeDataSource( queryId: String): ColumnarStreamWriter = { val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) - val handle = new NativeHandle( - create(library, "plan a write to", "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { - dataSource => - NativeArrow.exportSchema(schema) { address => - library.createStreamWriter(dataSource, address, overwrite, queryId) - } - }, - library.closeWriter) + val handle = + new NativeHandle(createWriter(library, schema, overwrite, queryId), library.closeWriter) val writer = new NativeDataSourceStreamWriter( executorLocation, name, handle, serializeWriter(library, handle)) NativeStreamWriters.register(queryId, handle) @@ -126,24 +121,72 @@ class NativeDataSource( } /** Runs `f` with a new native data source, which is closed afterwards. */ - private def withDataSource[T](library: NativeLibrary, action: String)(f: Long => T): T = { + private def withDataSource[T](library: NativeLibrary, action: String)( + f: NativeLibrary.DataSourceHandle => T): T = { + // Checked first, so that a data source is never created without a way to release it. + requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE)) // The keys of a case-insensitive map are in lower case. val entries = options.entrySet().asScala.toArray val dataSource = call(name, "create") { library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue)) } Utils.tryWithSafeFinally(f(dataSource)) { - call(name, action)(library.closeDataSource(dataSource)) + call(name, action)(library.closeDataSource(dataSource.handle)) } } - /** Creates a reader or a writer with `f`, which fails if the library does not implement it. */ - private def create(library: NativeLibrary, action: String, unsupportedError: String)( - f: Long => Long): Long = { + /** + * Creates a reader or a writer with `f`, which fails if the data source does not declare the + * capability, or if the library does not implement the functions that it requires. + */ + private def create( + library: NativeLibrary, + action: String, + capability: Long, + unsupportedError: String, + functions: Seq[NativeLibrary.Function])( + f: NativeLibrary.DataSourceHandle => Long): Long = { withDataSource(library, action) { dataSource => - call(name, action)(optional(f(dataSource))).getOrElse { + if ((dataSource.capabilities & capability) == 0) { throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) } + requireFunctions(library, action, functions) + call(name, action)(f(dataSource)) + } + } + + /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */ + private def createWriter( + library: NativeLibrary, + schema: StructType, + overwrite: Boolean, + queryId: String): Long = { + val action = "plan a write to" + val (capability, unsupportedError, overwriteMode) = if (queryId == null) { + (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", + "the save mode overwrite") + } else { + (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", + "the output mode complete") + } + create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) { + dataSource => + if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) { + throw error(name, action, s"The data source does not support $overwriteMode.") + } + val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND + NativeArrow.exportSchema(schema) { address => + library.createWriter(dataSource.handle, address, mode, queryId) + } + } + } + + private def requireFunctions( + library: NativeLibrary, + action: String, + functions: Seq[NativeLibrary.Function]): Unit = { + functions.find(!library.hasFunction(_)).foreach { function => + throw error(name, action, s"The library does not implement ${function.fieldName}.") } } @@ -158,6 +201,18 @@ class NativeDataSource( } } +object NativeDataSource { + import NativeLibrary.Function._ + + // The functions that the operations of a data source require, besides create and close. The + // others are optional. They are checked on the driver, before the operation starts. + private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ) + private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET, + STREAM_PARTITIONS, CLOSE_STREAM_READER, READ) + private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER, + WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER) +} + /** Reads a native data source for a batch scan. */ class NativeDataSourceReader( location: NativeLibraryLocation, @@ -182,21 +237,20 @@ class NativeDataSourceReader( val encoded = predicates.zipWithIndex.flatMap { case (predicate, i) => NativePredicates.toJson(predicate, schema).map(_ -> i) } - val pushed = if (encoded.isEmpty) { + val used = if (encoded.isEmpty) { Set.empty[Int] } else { + // How the reader uses each predicate: whether it uses it completely or only to skip data + // does not matter, since Spark evaluates all the predicates again after reading. call(name, "plan a scan of") { optional(library.pushPredicates(handle.get, encoded.map(_._1))) - } match { - case Some(flags) if flags != null && flags.length == encoded.length => - encoded.map(_._2).zip(flags).collect { case (i, true) => i }.toSet - case Some(flags) => - throw error(name, "plan a scan of", s"pushPredicates returned " + - s"${Option(flags).map(_.length).getOrElse(0)} values for ${encoded.length} predicates.") - case None => Set.empty[Int] - } + }.map { uses => + encoded.map(_._2).zip(uses).collect { + case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i + }.toSet + }.getOrElse(Set.empty[Int]) } - predicates.zipWithIndex.collect { case (predicate, i) if !pushed.contains(i) => predicate } + predicates.zipWithIndex.collect { case (predicate, i) if used.contains(i) => predicate } } override def pushLimit(limit: Int): Boolean = { @@ -257,13 +311,11 @@ class NativeDataSourceStreamReader( @transient private lazy val library = NativeLibraries.get(location) override def initialOffset(): String = { - checkOffset("spark_ds_initial_offset", - call(name, "plan a scan of")(library.initialOffset(handle.get))) + checkOffset("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get))) } override def latestOffset(): String = { - checkOffset("spark_ds_latest_offset", - call(name, "plan a scan of")(library.latestOffset(handle.get))) + checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get))) } /** @@ -422,6 +474,8 @@ class NativePartitionReader( // Read here rather than with `Data.importArrayStream`, to check each array before importing it. private val stream = ArrowArrayStream.allocateNew(allocator) private val dictionaries = new CDataDictionaryProvider() + // The task that reads the partition, which the library may use until the stream is released. + private var task: NativeLibrary.Task = _ private var root: VectorSchemaRoot = _ private var batch: ColumnarBatch = _ @@ -438,11 +492,8 @@ class NativePartitionReader( private def openStream(): Unit = { val library = NativeLibraries.get(location) - call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) - .getOrElse { - throw new SparkUnsupportedOperationException( - "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) - } + task = NativeTask.create() + call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress())) } private def importSchema(): Schema = { @@ -641,7 +692,12 @@ class NativePartitionReader( // Releases the stream of the library, unless it was not exported. stream.release() stream.close() - }(Utils.tryLogNonFatalError(allocator.close())) + } { + if (task != null) { + task.close() + } + Utils.tryLogNonFatalError(allocator.close()) + } } } @@ -668,14 +724,23 @@ class NativeDataWriter( import NativeCalls._ private val library = NativeLibraries.get(location) - // Zero once the data writer was committed or aborted, which releases it. It stays set if the - // library fails to commit, so that Spark aborts the data writer. - private var handle: Long = call(name, "write to") { - library.createDataWriter(writerState, partitionId, taskId, epochId) + // The task, which the library may use until the data writer is closed. + private val task = NativeTask.create(Some(partitionId), Some(taskId)) + private val handle: Long = try { + call(name, "write to")(library.createDataWriter(task, writerState, epochId)) + } catch { + case e: Throwable => + task.close() + throw e } + // Whether the data writer was committed or aborted. It stays false if the library fails to + // commit, so that Spark aborts the data writer, but becomes true as soon as the library commits, + // even if Spark then fails the commit, for example because of an invalid message. + private var finished = false + private var closed = false override def write(batch: ColumnarBatch): Unit = { - if (handle == 0) { + if (finished || closed) { throw SparkException.internalError("The data writer was already committed or aborted.") } val vectors = (0 until batch.numCols()).map { i => @@ -704,24 +769,49 @@ class NativeDataWriter( } override def commit(): WriterCommitMessage = { - if (handle == 0) { + if (finished || closed) { throw SparkException.internalError("The data writer was already committed or aborted.") } - // The library releases the data writer once it commits, even if its message is then rejected, - // so Spark must not abort it afterwards. - val message = call(name, "write to")(library.commitDataWriter(handle, () => handle = 0)) + val message = call(name, "write to")(library.commitDataWriter(handle, () => finished = true)) NativeWriterCommitMessage(bytesOrEmpty(Some(message))) } override def abort(): Unit = { - if (handle != 0) { - val dataWriter = handle - handle = 0 - call(name, "write to")(library.abortDataWriter(dataWriter)) + if (!finished && !closed) { + finished = true + call(name, "write to")(library.abortDataWriter(handle)) + } + } + + override def close(): Unit = { + if (!closed) { + closed = true + Utils.tryWithSafeFinally { + call(name, "write to")(library.closeDataWriter(handle)) + }(task.close()) } } +} - override def close(): Unit = abort() +/** Creates the `SparkDsTaskContext` of the running task. */ +private[ffi] object NativeTask { + /** + * @param partitionId the partition that the task reads or writes, by default the partition of + * the task + * @param taskAttemptId the ID of the task attempt, by default the one of the task + */ + def create( + partitionId: Option[Int] = None, + taskAttemptId: Option[Long] = None): NativeLibrary.Task = { + val context = Option(TaskContext.get()) + new NativeLibrary.Task( + context.map(_.stageId()).getOrElse(-1), + context.map(_.stageAttemptNumber()).getOrElse(-1), + partitionId.orElse(context.map(_.partitionId())).getOrElse(-1), + context.map(_.attemptNumber()).getOrElse(0), + taskAttemptId.orElse(context.map(_.taskAttemptId())).getOrElse(-1L), + () => context.exists(_.isInterrupted())) + } } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala index 938369ad6dd89..2e39c67122f04 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -26,11 +26,12 @@ import java.util.concurrent.ConcurrentHashMap import java.util.zip.ZipFile import scala.jdk.CollectionConverters._ -import scala.util.Try +import scala.util.{Failure, Try} import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} import com.google.common.cache.{Cache, CacheBuilder} +import org.apache.spark.SPARK_VERSION import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.{PATH, VERSION_NUM} import org.apache.spark.sql.Artifact @@ -385,7 +386,8 @@ object NativeLibraries extends Logging { /** * Loads a library, and checks the version of the interface that it implements. Throws the error - * if the library cannot be loaded, and returns it if the library is loaded but rejected. + * if the library cannot be loaded, and returns it if the library is loaded but rejected, so that + * it is kept: the library stays loaded, and Spark only initializes it once. * * @param invalidLibrary returns the error for an invalid library, with the reason * @param unsupportedAbiVersion returns the error for an unsupported version of the interface, @@ -397,20 +399,18 @@ object NativeLibraries extends Logging { unsupportedAbiVersion: Map[String, String] => Throwable): Try[NativeLibrary] = { logInfo(log"Loading the native data source library ${MDC(PATH, file)}.") val library = try { - NativeLibrary.load(file.getPath) + NativeLibrary.load(file.getPath, SPARK_VERSION) } catch { + // A library that was loaded stays loaded, and Spark must not initialize it again. + case e: NativeLibrary.LoadException if e.loaded => + return Failure(invalidLibrary(e.getMessage, e)) case e: NativeLibrary.LoadException => throw invalidLibrary(e.getMessage, e) } Try { - val abiVersion = try { - library.abiVersion() - } catch { - case e: NativeLibrary.MissingFunctionException => - throw invalidLibrary("It does not implement spark_ds_abi_version.", e) - } - if (abiVersion != NativeLibrary.ABI_VERSION) { - throw unsupportedAbiVersion( - Map("version" -> abiVersion.toString, "supported" -> NativeLibrary.ABI_VERSION.toString)) + if (library.abiVersion() != NativeLibrary.ABI_VERSION) { + throw unsupportedAbiVersion(Map( + "version" -> library.abiVersion().toString, + "supported" -> NativeLibrary.ABI_VERSION.toString)) } library } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala index 18167548c03e5..3e49dd2dae37d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.types._ /** * Encodes Data Source V2 predicates as the JSON expression trees passed to - * `spark_ds_push_predicates`. + * `push_predicates` of `spark_datasource.h`. */ object NativePredicates { // Writes ASCII only, escaping the other characters, which any JSON parser decodes. @@ -36,7 +36,7 @@ object NativePredicates { /** * The names of the predicates that are passed to the libraries, as listed in - * `spark_ds_push_predicates`. They are part of the interface, so a predicate with any other + * `push_predicates`. They are part of the interface, so a predicate with any other * name, such as an arithmetic function, is not passed: its semantics, for example on overflow, * would be part of the interface too, and the names of Data Source V2 grow with new releases. */ diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index a3fb0db870398..6ba8fa78ff5b0 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -22,43 +22,63 @@ // native_range Reads rows [0, end) with columns of several types. Supports predicate, // column and limit pushdown. With the option "claim_predicates", it reports // all the predicates as evaluated completely, but ignores them. With the -// option "extra", the schema has one more column: "map" adds a map whose -// entry fields are named "keys" and "values", as arrow-rs names them, "struct" -// adds a struct, "decimal256" adds a decimal256, "listview" adds a list view, +// option "predicate_use", it reports the predicates that it evaluates with +// another value: "partial" for SPARK_DS_PREDICATE_PARTIAL, as a reader that +// reads the blocks of 10 rows that may match, or a number. With the option +// "extra", the schema has one more column: "map" adds a map whose entry fields +// are named "keys" and "values", as arrow-rs names them, "struct" adds a +// struct, "decimal256" adds a decimal256, "listview" adds a list view, // "largelist" adds a large list, and "dictionary" adds a dictionary-encoded // string column, whose data Spark rejects before reading it. With the option // "slice", the columns of the batches have an offset of 1. With the option // "short", an array of the batches has one value fewer than its parent refers // to, although its buffers hold them all: "id" shortens the column id, // "entries" the entries of the map and their fields, and "values" the values -// of the entries. +// of the entries. With the option "wait_for_interrupt", the stream of a +// partition creates the file .waiting, waits until its task +// is killed, and then creates the file and fails. // native_sink Writes the rows as CSV files under the directory in the option "path". With // the option "commit_messages", a data writer passes its commit message that -// many times. +// many times. Creating and closing a data writer create a file named after +// it in _created_data_writers and _closed_data_writers. // native_counter Streams the rows [0, max_offset) with a single column "id", making "step" // more rows available each time Spark asks for the latest offset. With the // option "bad_offset", the latest offset is "none", "multiline" JSON or // "invalid" JSON. With the option "closed_path", closing the stream reader // creates that file. // +// native_range supports batch reads, native_counter micro-batch streaming reads and native_sink +// batch and streaming writes, which may replace the existing data. With the option +// "capabilities", a data source declares these capabilities instead, a number. With the option +// "task_path", the reads and the data writers of native_range and native_sink record their task +// in a file named after its partition, in that directory. With the option "log", creating a data +// source logs its value as a warning, and with the option "spark_version", it fails unless Spark +// has that version. +// // With the option "fail", a step fails with an error. With the option "zero", a function that -// creates a handle succeeds without setting it: "spark_ds_create", "spark_ds_create_reader", -// "spark_ds_create_stream_reader", "spark_ds_create_writer", "spark_ds_create_stream_writer" or -// "spark_ds_create_data_writer". +// creates a handle succeeds without setting it: "create", "create_reader", +// "create_stream_reader", "create_writer" or "create_data_writer". // // The suite compiles it with these macros to get several distinct libraries: // // NAME_PREFIX A prefix for the names of the data sources, "" by default. // ID_SIGN Multiplies the ids produced by native_range, 1 by default. // READ_ONLY Leaves out the functions for writes, streaming reads and pushdown. -// ABI_VERSION The version returned by spark_ds_abi_version, 1 by default. -// NO_LAST_ERROR Leaves out spark_ds_last_error. +// ABI_VERSION The version of the table, 1 by default. +// NO_LAST_ERROR Leaves out last_error. +// TABLE_SIZE The size of the table, sizeof(SparkDsLibrary) by default: Spark must not read +// the functions after it. +// INIT_FAILURE 1 to make spark_ds_init fail, after logging why, 2 to make it succeed without +// a table. +// NO_INIT Leaves out spark_ds_init, which makes the library invalid. #include "spark_datasource.h" #include #include #include +#include +#include #include #include #include @@ -69,6 +89,7 @@ #include #include #include +#include #include #include @@ -85,6 +106,12 @@ #ifndef ABI_VERSION #define ABI_VERSION 1 #endif +#ifndef TABLE_SIZE +#define TABLE_SIZE sizeof(SparkDsLibrary) +#endif +#ifndef INIT_FAILURE +#define INIT_FAILURE 0 +#endif namespace { @@ -599,6 +626,17 @@ std::vector split(const std::string& text) { thread_local std::string last_error; +// Spark, as spark_ds_init got it. +const SparkDsHost* host = nullptr; + +// Logs with the log of Spark. A library must only use the fields that fit in the size of the +// struct that Spark passes, which may be of another version of the header. +void hostLog(int32_t level, const std::string& message) { + if (host->size >= static_cast(offsetof(SparkDsHost, log) + sizeof(host->log))) { + host->log(level, message.c_str()); + } +} + // Runs `f`, and turns a C++ exception into an error: none may cross the boundary. template int guarded(F&& f) { @@ -663,6 +701,33 @@ bool returnsZero(const Options& options, const std::string& function) { return get(options, "zero", "") == function; } +void writeFile(const std::string& path, const std::string& content) { + FILE* file = std::fopen(path.c_str(), "w"); + if (file == nullptr) throw std::runtime_error("cannot write " + path); + std::fputs(content.c_str(), file); + std::fclose(file); +} + +void makeDirs(const std::string& path) { + for (size_t pos = path.find('/', 1); pos != std::string::npos; pos = path.find('/', pos + 1)) { + mkdir(path.substr(0, pos).c_str(), 0755); + } + mkdir(path.c_str(), 0755); +} + +// With the option "task_path", records the task of a read or a data writer in a file named after +// its partition: the size of the struct and its fields. +void recordTask(const Options& options, const SparkDsTaskContext* task) { + std::string dir = get(options, "task_path", ""); + if (dir.empty()) return; + makeDirs(dir); + writeFile(dir + "/" + std::to_string(task->partition_id), + std::to_string(task->size) + "," + std::to_string(task->stage_id) + "," + + std::to_string(task->stage_attempt_number) + "," + + std::to_string(task->partition_id) + "," + std::to_string(task->attempt_number) + + "," + std::to_string(task->task_attempt_id)); +} + // The driver-side reader of native_range and native_counter. struct Reader { Options state; // Serialized to the executors. @@ -677,6 +742,8 @@ struct Stream { bool fail; bool slice; std::string short_array; + const SparkDsTaskContext* task; + std::string wait_for_interrupt; std::string last_error; }; @@ -692,6 +759,17 @@ int streamGetNext(ArrowArrayStream* stream, ArrowArray* out) { state->last_error = "injected failure in get_next"; return EIO; } + if (!state->wait_for_interrupt.empty()) { + // Waits, as for data that does not come, until the task is killed, for at most a minute. + writeFile(state->wait_for_interrupt + ".waiting", ""); + for (int i = 0; i < 6000 && !state->task->is_interrupted(state->task); i++) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + bool interrupted = state->task->is_interrupted(state->task); + writeFile(state->wait_for_interrupt, interrupted ? "interrupted" : "timed out"); + state->last_error = "the task was killed"; + return ECANCELED; + } if (state->next >= state->end) { out->release = nullptr; // The end of the stream. return 0; @@ -728,20 +806,15 @@ struct DataWriter { std::string temp_path; FILE* file; Options state; + bool committed; }; -// The data writers that a commit released, with the path of their write, so that aborting one of -// them afterwards, which the header forbids, is recorded in the file _ABORTED_RELEASED_DATA_WRITER +// The data writers that Spark closed, with the path of their write, so that using one of them +// afterwards, which the header forbids, is recorded in the file _DATA_WRITER_USED_AFTER_CLOSE // rather than crashing. -std::mutex released_data_writers_lock; -[[maybe_unused]] std::map released_data_writers; +std::mutex closed_data_writers_lock; +[[maybe_unused]] std::map closed_data_writers; -[[maybe_unused]] void makeDirs(const std::string& path) { - for (size_t pos = path.find('/', 1); pos != std::string::npos; pos = path.find('/', pos + 1)) { - mkdir(path.substr(0, pos).c_str(), 0755); - } - mkdir(path.c_str(), 0755); -} [[maybe_unused]] bool isValid(const ArrowArray* array, int64_t i) { if (array->null_count == 0 || array->buffers[0] == nullptr) return true; @@ -788,30 +861,31 @@ std::vector listFiles(const std::string& dir, const std::string& pr return files; } -[[maybe_unused]] void writeFile(const std::string& path, const std::string& content) { - FILE* file = std::fopen(path.c_str(), "w"); - if (file == nullptr) throw std::runtime_error("cannot write " + path); - std::fputs(content.c_str(), file); - std::fclose(file); -} - -} // namespace - - // --------------------------------------------------------------------------------------------- -// spark_datasource.h +// The functions of the table // --------------------------------------------------------------------------------------------- -extern "C" { - #ifndef NO_LAST_ERROR -const char* spark_ds_last_error(void) { return last_error.c_str(); } +const char* lastError(void) { return last_error.c_str(); } #endif -int32_t spark_ds_abi_version(void) { return ABI_VERSION; } +// The capabilities of a data source, unless the option "capabilities" declares others. +uint64_t capabilitiesOf(const DataSource& source) { + std::string declared = get(source.options, "capabilities", ""); + if (!declared.empty()) return std::stoull(declared); + if (source.name == name("native_range")) return SPARK_DS_BATCH_READ; +#ifndef READ_ONLY + if (source.name == name("native_counter")) return SPARK_DS_MICRO_BATCH_READ; + if (source.name == name("native_sink")) { + return SPARK_DS_BATCH_WRITE | SPARK_DS_STREAMING_WRITE | SPARK_DS_TRUNCATE; + } +#endif + return 0; +} -int spark_ds_create(const char* name_value, const char* const* option_keys, - const char* const* option_values, int64_t num_options, void** handle) { +int createSource(const char* name_value, const char* const* option_keys, + const char* const* option_values, int64_t num_options, uint64_t* capabilities, + void** handle) { return guarded([&] { auto source = std::make_unique(DataSource{name_value, {}}); auto keys = toStrings(option_keys, num_options); @@ -822,12 +896,21 @@ int spark_ds_create(const char* name_value, const char* const* option_keys, source->name != name("native_counter")) { throw std::runtime_error("unknown data source " + source->name); } - if (returnsZero(source->options, "spark_ds_create")) return; + std::string message = get(source->options, "log", ""); + if (!message.empty()) hostLog(SPARK_DS_LOG_WARN, message); + std::string version = get(source->options, "spark_version", ""); + if (!version.empty() && version != host->spark_version) { + throw std::runtime_error("the version of Spark is " + std::string(host->spark_version) + + ", which supports version " + std::to_string(host->abi_version) + + " of the interface"); + } + *capabilities = capabilitiesOf(*source); + if (returnsZero(source->options, "create")) return; *handle = source.release(); }); } -int spark_ds_schema(void* handle, struct ArrowSchema* out) { +int sourceSchema(void* handle, struct ArrowSchema* out) { return guarded([&] { auto* source = fromHandle(handle); maybeFail(source->options, "schema"); @@ -839,7 +922,7 @@ int spark_ds_schema(void* handle, struct ArrowSchema* out) { }); } -int spark_ds_create_reader(void* handle, struct ArrowSchema* schema, void** out) { +int createReader(void* handle, struct ArrowSchema* schema, void** out) { return guarded([&] { auto* source = fromHandle(handle); if (source->name != name("native_range")) { @@ -847,31 +930,33 @@ int spark_ds_create_reader(void* handle, struct ArrowSchema* schema, void** out) } auto columns = importColumns(schema); for (const auto& column : columns) findColumn(column); - if (returnsZero(source->options, "spark_ds_create_reader")) return; + if (returnsZero(source->options, "create_reader")) return; auto* reader = new Reader(); reader->state["columns"] = join(columns); reader->state["lo"] = "0"; reader->state["hi"] = get(source->options, "end", "10"); reader->state["partitions"] = get(source->options, "partitions", "2"); reader->state["batch_size"] = get(source->options, "batch_size", "1000"); - reader->state["fail"] = get(source->options, "fail", ""); - reader->state["slice"] = get(source->options, "slice", ""); - reader->state["short"] = get(source->options, "short", ""); - reader->state["claim_predicates"] = get(source->options, "claim_predicates", ""); + for (const char* option : {"fail", "slice", "short", "claim_predicates", "predicate_use", + "wait_for_interrupt", "task_path"}) { + reader->state[option] = get(source->options, option, ""); + } *out = reader; }); } +void closeSource(void* handle) { delete fromHandle(handle); } + #ifndef READ_ONLY -int spark_ds_create_stream_reader(void* handle, struct ArrowSchema* schema, void** out) { +int createStreamReader(void* handle, struct ArrowSchema* schema, void** out) { return guarded([&] { auto* source = fromHandle(handle); if (source->name != name("native_counter")) { throw std::runtime_error(source->name + " does not support streaming reads"); } auto columns = join(importColumns(schema)); - if (returnsZero(source->options, "spark_ds_create_stream_reader")) return; + if (returnsZero(source->options, "create_stream_reader")) return; auto* reader = new Reader(); reader->state["columns"] = columns; reader->state["max_offset"] = get(source->options, "max_offset", "10"); @@ -885,25 +970,24 @@ int spark_ds_create_stream_reader(void* handle, struct ArrowSchema* schema, void }); } -#endif // READ_ONLY - -void spark_ds_close(void* handle) { delete fromHandle(handle); } - -#ifndef READ_ONLY - -int spark_ds_push_predicates(void* handle, const char* const* predicates, int64_t count, - bool* evaluated) { +int pushPredicates(void* handle, const char* const* predicates, int64_t count, int32_t* uses) { return guarded([&] { auto* reader = fromHandle(handle); auto json = toStrings(predicates, count); + int64_t end = getLong(reader->state, "hi", 0); int64_t lo = getLong(reader->state, "lo", 0); - int64_t hi = getLong(reader->state, "hi", 0); + int64_t hi = end; // With the option "claim_predicates", reports every predicate as evaluated completely, but // evaluates none of them. if (reader->state["claim_predicates"] == "true") { - for (size_t i = 0; i < json.size(); i++) evaluated[i] = true; + for (size_t i = 0; i < json.size(); i++) uses[i] = SPARK_DS_PREDICATE_COMPLETE; return; } + std::string predicate_use = reader->state["predicate_use"]; + bool partial = predicate_use == "partial"; + int32_t use = predicate_use.empty() ? SPARK_DS_PREDICATE_COMPLETE + : partial ? SPARK_DS_PREDICATE_PARTIAL + : std::stoi(predicate_use); for (size_t i = 0; i < json.size(); i++) { // Accepts "id " for the comparison operators. Json predicate = JsonParser(json[i]).parse(); @@ -927,7 +1011,12 @@ int spark_ds_push_predicates(void* handle, const char* const* predicates, int64_ hi = std::min(hi, value + 1); } } - evaluated[i] = ok; + uses[i] = ok ? use : SPARK_DS_PREDICATE_UNUSED; + } + if (partial) { + // Reads the blocks of 10 rows that contain matching rows, with all their rows. + lo = lo / 10 * 10; + hi = std::min(end, (hi + 9) / 10 * 10); } reader->state["lo"] = std::to_string(lo); reader->state["hi"] = std::to_string(std::max(lo, hi)); @@ -935,18 +1024,17 @@ int spark_ds_push_predicates(void* handle, const char* const* predicates, int64_ }); } -int spark_ds_push_limit(void* handle, int32_t limit, bool* applied) { +int pushLimit(void* handle, int32_t limit, bool* applied) { return guarded([&] { auto* reader = fromHandle(handle); - // Spark calls spark_ds_push_limit at most once. + // Spark calls push_limit at most once. if (reader->state.count("limit") > 0) throw std::runtime_error("pushLimit was called twice"); reader->state["limit"] = std::to_string(limit); *applied = true; }); } -int spark_ds_prune_columns(void* handle, const char* const* columns, int64_t count, - bool* pruned) { +int pruneColumns(void* handle, const char* const* columns, int64_t count, bool* pruned) { return guarded([&] { auto names = toStrings(columns, count); for (const auto& column : names) findColumn(column); @@ -957,7 +1045,7 @@ int spark_ds_prune_columns(void* handle, const char* const* columns, int64_t cou #endif // READ_ONLY -int spark_ds_partitions(void* handle, SparkDsSink add_partition, void* sink) { +int planPartitions(void* handle, SparkDsSink add_partition, void* sink) { return guarded([&] { auto* reader = fromHandle(handle); maybeFail(reader->state, "plan"); @@ -977,20 +1065,20 @@ int spark_ds_partitions(void* handle, SparkDsSink add_partition, void* sink) { }); } -int spark_ds_serialize_reader(void* handle, SparkDsSink set_state, void* sink) { +int serializeReader(void* handle, SparkDsSink set_state, void* sink) { return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); } -void spark_ds_close_reader(void* handle) { delete fromHandle(handle); } +void closeReader(void* handle) { delete fromHandle(handle); } #ifndef READ_ONLY -int spark_ds_initial_offset(void*, SparkDsSink set_offset, void* sink) { +int initialOffset(void*, SparkDsSink set_offset, void* sink) { return guarded([&] { sendBytes(set_offset, sink, "{\"offset\":0}"); }); } // Each call makes `step` more rows available, up to `max_offset`. -int spark_ds_latest_offset(void* handle, SparkDsSink set_offset, void* sink) { +int latestOffset(void* handle, SparkDsSink set_offset, void* sink) { return guarded([&] { Options& state = fromHandle(handle)->state; int64_t latest = std::min(getLong(state, "max_offset", 10), @@ -1009,8 +1097,8 @@ int spark_ds_latest_offset(void* handle, SparkDsSink set_offset, void* sink) { }); } -int spark_ds_stream_partitions(void*, const char* start, const char* end, - SparkDsSink add_partition, void* sink) { +int streamPartitions(void*, const char* start, const char* end, SparkDsSink add_partition, + void* sink) { return guarded([&] { auto lo = static_cast(JsonParser(start).parse()["offset"].number); auto hi = static_cast(JsonParser(end).parse()["offset"].number); @@ -1018,18 +1106,18 @@ int spark_ds_stream_partitions(void*, const char* start, const char* end, }); } -int spark_ds_serialize_stream_reader(void* handle, SparkDsSink set_state, void* sink) { +int serializeStreamReader(void* handle, SparkDsSink set_state, void* sink) { return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); } -int spark_ds_commit_offset(void* handle, const char* end) { +int commitOffset(void* handle, const char* end) { return guarded([&] { std::string path = get(fromHandle(handle)->state, "commit_path", ""); if (!path.empty()) writeFile(path, end); }); } -void spark_ds_close_stream_reader(void* handle) { +void closeStreamReader(void* handle) { std::unique_ptr reader(fromHandle(handle)); guarded([&] { std::string path = get(reader->state, "closed_path", ""); @@ -1039,9 +1127,9 @@ void spark_ds_close_stream_reader(void* handle) { #endif // READ_ONLY -int spark_ds_read(const uint8_t* reader_state, int64_t reader_state_length, - const uint8_t* partition, int64_t partition_length, - struct ArrowArrayStream* stream) { +int readPartition(const SparkDsTaskContext* task, const uint8_t* reader_state, + int64_t reader_state_length, const uint8_t* partition, + int64_t partition_length, struct ArrowArrayStream* stream) { return guarded([&] { Options state = deserialize(toBytes(reader_state, reader_state_length)); maybeFail(state, "open"); @@ -1049,6 +1137,7 @@ int spark_ds_read(const uint8_t* reader_state, int64_t reader_state_length, if (get(state, "fail", "") == "uninitialized") return; auto bounds = split(toBytes(partition, partition_length)); if (bounds.size() != 2) throw std::runtime_error("invalid partition"); + recordTask(state, task); stream->get_schema = streamGetSchema; stream->get_next = streamGetNext; stream->get_last_error = streamGetLastError; @@ -1057,24 +1146,22 @@ int spark_ds_read(const uint8_t* reader_state, int64_t reader_state_length, std::stoll(bounds[1]), getLong(state, "batch_size", 1000), get(state, "fail", "") == "read", get(state, "slice", "") == "true", get(state, "short", ""), - ""}; + task, get(state, "wait_for_interrupt", ""), ""}; }); } #ifndef READ_ONLY -} // extern "C" - -namespace { - -// Creates the writer of a batch write, or of a streaming write with the ID of its query. -int createAnyWriter(const char* function, void* handle, struct ArrowSchema* schema, - bool overwrite, const char* query_id, void** out) { +int createWriter(void* handle, struct ArrowSchema* schema, const SparkDsWriteInfo* info, + void** out) { return guarded([&] { auto* source = fromHandle(handle); if (source->name != name("native_sink")) { throw std::runtime_error(source->name + " does not support writes"); } + if (info->mode != SPARK_DS_WRITE_APPEND && info->mode != SPARK_DS_WRITE_TRUNCATE) { + throw std::runtime_error("unknown mode " + std::to_string(info->mode)); + } // Spark releases the schema after the call, since it is not moved here. if (schema->n_children == 0) throw std::runtime_error("no columns to write"); // The location of a table in a catalog is passed as a URI. @@ -1085,43 +1172,29 @@ int createAnyWriter(const char* function, void* handle, struct ArrowSchema* sche path = path.substr(5); } if (path.empty()) throw std::runtime_error("the option 'path' is required"); - if (returnsZero(source->options, function)) return; + if (returnsZero(source->options, "create_writer")) return; auto* writer = new Writer(); writer->state["path"] = path; - writer->state["overwrite"] = overwrite ? "true" : "false"; + writer->state["overwrite"] = info->mode == SPARK_DS_WRITE_TRUNCATE ? "true" : "false"; writer->state["fail"] = get(source->options, "fail", ""); writer->state["zero"] = get(source->options, "zero", ""); writer->state["commit_messages"] = get(source->options, "commit_messages", "1"); + writer->state["task_path"] = get(source->options, "task_path", ""); // Records that a writer was created, which Spark then commits or aborts, and the ID of the // streaming query. makeDirs(path); writeFile(path + "/_CREATED", ""); - if (query_id != nullptr) writeFile(path + "/_QUERY_ID", query_id); + if (info->query_id != nullptr) writeFile(path + "/_QUERY_ID", info->query_id); *out = writer; }); } -} // namespace - -extern "C" { - -int spark_ds_create_writer(void* handle, struct ArrowSchema* schema, bool overwrite, - void** writer) { - return createAnyWriter("spark_ds_create_writer", handle, schema, overwrite, nullptr, writer); -} - -int spark_ds_create_stream_writer(void* handle, struct ArrowSchema* schema, bool overwrite, - const char* query_id, void** writer) { - return createAnyWriter("spark_ds_create_stream_writer", handle, schema, overwrite, query_id, - writer); -} - -int spark_ds_serialize_writer(void* handle, SparkDsSink set_state, void* sink) { +int serializeWriter(void* handle, SparkDsSink set_state, void* sink) { return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); } -int spark_ds_commit(void* handle, int64_t epoch_id, const uint8_t* const* messages, - const int64_t* message_lengths, int64_t num_messages) { +int commitWrite(void* handle, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages) { return guarded([&] { auto* writer = fromHandle(handle); std::string path = writer->state["path"]; @@ -1144,8 +1217,8 @@ int spark_ds_commit(void* handle, int64_t epoch_id, const uint8_t* const* messag // Records the aborted write in _ABORTED for a batch write, or _ABORTED_ for a streaming // one, with the number of tasks that did not commit, whose messages are NULL. -int spark_ds_abort(void* handle, int64_t epoch_id, const uint8_t* const* messages, - const int64_t* message_lengths, int64_t num_messages) { +int abortWrite(void* handle, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages) { return guarded([&] { auto* writer = fromHandle(handle); writer->closed_marker = closedMarker(epoch_id); @@ -1164,7 +1237,7 @@ int spark_ds_abort(void* handle, int64_t epoch_id, const uint8_t* const* message }); } -void spark_ds_close_writer(void* handle) { +void closeWriter(void* handle) { std::unique_ptr writer(fromHandle(handle)); guarded([&] { if (!writer->closed_marker.empty()) { @@ -1173,29 +1246,41 @@ void spark_ds_close_writer(void* handle) { }); } -int spark_ds_create_data_writer(const uint8_t* writer_state, int64_t writer_state_length, - int32_t partition_id, int64_t task_id, int64_t epoch_id, - void** out) { +// Whether Spark uses a data writer that it closed, which the header forbids. +bool usedAfterClose(void* handle) { + std::lock_guard guard(closed_data_writers_lock); + auto closed = closed_data_writers.find(handle); + if (closed == closed_data_writers.end()) return false; + writeFile(closed->second + "/_DATA_WRITER_USED_AFTER_CLOSE", ""); + return true; +} + +int createDataWriter(const SparkDsTaskContext* task, const uint8_t* writer_state, + int64_t writer_state_length, int64_t epoch_id, void** out) { return guarded([&] { Options state = deserialize(toBytes(writer_state, writer_state_length)); - if (returnsZero(state, "spark_ds_create_data_writer")) return; + if (returnsZero(state, "create_data_writer")) return; + recordTask(state, task); std::string dir = state["path"] + "/_temporary"; makeDirs(dir); std::string temp_path = dir + "/" + std::to_string(epoch_id) + "-" + - std::to_string(partition_id) + "-" + std::to_string(task_id) + ".csv"; + std::to_string(task->partition_id) + "-" + + std::to_string(task->task_attempt_id) + ".csv"; FILE* file = std::fopen(temp_path.c_str(), "w"); if (file == nullptr) throw std::runtime_error("cannot write " + temp_path); - auto* writer = new DataWriter{temp_path, file, state}; + makeDirs(state["path"] + "/_created_data_writers"); + writeFile(state["path"] + "/_created_data_writers/" + temp_path.substr(dir.size() + 1), ""); + auto* writer = new DataWriter{temp_path, file, state, false}; { - // The allocator may reuse the address of a released data writer. - std::lock_guard guard(released_data_writers_lock); - released_data_writers.erase(writer); + // The allocator may reuse the address of a closed data writer. + std::lock_guard guard(closed_data_writers_lock); + closed_data_writers.erase(writer); } *out = writer; }); } -int spark_ds_write(void* handle, struct ArrowArray* array, struct ArrowSchema* schema) { +int writeBatch(void* handle, struct ArrowArray* array, struct ArrowSchema* schema) { return guarded([&] { auto* writer = fromHandle(handle); // The batch is released here, so Spark does not release it after the call. @@ -1216,39 +1301,114 @@ int spark_ds_write(void* handle, struct ArrowArray* array, struct ArrowSchema* s }); } -int spark_ds_commit_data_writer(void* handle, SparkDsSink set_message, void* sink) { +int commitDataWriter(void* handle, SparkDsSink set_message, void* sink) { + if (usedAfterClose(handle)) return 0; return guarded([&] { auto* writer = fromHandle(handle); - // If the commit fails, the data writer stays valid: Spark aborts it with - // spark_ds_abort_data_writer. + // If the commit fails, Spark aborts the data writer. maybeFail(writer->state, "commitDataWriter"); // With the option "commit_messages", the message is passed that many times. int64_t messages = getLong(writer->state, "commit_messages", 1); for (int64_t i = 0; i < messages; i++) sendBytes(set_message, sink, writer->temp_path); std::fclose(writer->file); - { - std::lock_guard guard(released_data_writers_lock); - released_data_writers[handle] = writer->state["path"]; - } - delete writer; + writer->file = nullptr; + writer->committed = true; }); } -int spark_ds_abort_data_writer(void* handle) { - { - std::lock_guard guard(released_data_writers_lock); - auto released = released_data_writers.find(handle); - if (released != released_data_writers.end()) { - writeFile(released->second + "/_ABORTED_RELEASED_DATA_WRITER", ""); - return 0; +int abortDataWriter(void* handle) { + if (usedAfterClose(handle)) return 0; + return guarded([&] { + auto* writer = fromHandle(handle); + // Spark does not abort a data writer that committed: records it if it does. + if (writer->committed) { + writeFile(writer->state["path"] + "/_ABORTED_COMMITTED_DATA_WRITER", ""); + return; } - } + if (writer->file != nullptr) { + std::fclose(writer->file); + writer->file = nullptr; + } + unlink(writer->temp_path.c_str()); + }); +} + +void closeDataWriter(void* handle) { + if (usedAfterClose(handle)) return; std::unique_ptr writer(fromHandle(handle)); - std::fclose(writer->file); - unlink(writer->temp_path.c_str()); - return 0; + guarded([&] { + if (writer->file != nullptr) std::fclose(writer->file); + std::string path = writer->state["path"]; + std::string dir = path + "/_closed_data_writers"; + makeDirs(dir); + writeFile(dir + "/" + writer->temp_path.substr(writer->temp_path.rfind('/') + 1), ""); + std::lock_guard guard(closed_data_writers_lock); + closed_data_writers[handle] = path; + }); } #endif // READ_ONLY -} // extern "C" +// The table, which NO_INIT leaves unused. +[[maybe_unused]] SparkDsLibrary makeTable() { + SparkDsLibrary table{}; + table.size = TABLE_SIZE; + table.abi_version = ABI_VERSION; +#ifndef NO_LAST_ERROR + table.last_error = lastError; +#endif + table.create = createSource; + table.schema = sourceSchema; + table.create_reader = createReader; + table.close = closeSource; + table.partitions = planPartitions; + table.serialize_reader = serializeReader; + table.close_reader = closeReader; + table.read = readPartition; +#ifndef READ_ONLY + table.create_stream_reader = createStreamReader; + table.create_writer = createWriter; + table.push_predicates = pushPredicates; + table.push_limit = pushLimit; + table.prune_columns = pruneColumns; + table.initial_offset = initialOffset; + table.latest_offset = latestOffset; + table.stream_partitions = streamPartitions; + table.serialize_stream_reader = serializeStreamReader; + table.commit_offset = commitOffset; + table.close_stream_reader = closeStreamReader; + table.serialize_writer = serializeWriter; + table.commit = commitWrite; + table.abort = abortWrite; + table.close_writer = closeWriter; + table.create_data_writer = createDataWriter; + table.write = writeBatch; + table.commit_data_writer = commitDataWriter; + table.abort_data_writer = abortDataWriter; + table.close_data_writer = closeDataWriter; +#endif + return table; +} + +} // namespace + +// --------------------------------------------------------------------------------------------- +// spark_datasource.h +// --------------------------------------------------------------------------------------------- + +#ifndef NO_INIT + +int spark_ds_init(const SparkDsHost* spark, const SparkDsLibrary** library) { + host = spark; + if (INIT_FAILURE == 1) { + hostLog(SPARK_DS_LOG_ERROR, "injected failure in spark_ds_init"); + return 1; + } + // Succeeds without returning a table, which Spark must reject. + if (INIT_FAILURE == 2) return 0; + static const SparkDsLibrary table = makeTable(); + *library = &table; + return 0; +} + +#endif // NO_INIT diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala index bc94db65cf0ab..898e1796abe03 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala @@ -248,7 +248,7 @@ class TestRangeReader(end: Long, schema: StructType) extends ColumnarReader { private var readSchema = schema override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { - predicates.filterNot { predicate => + predicates.filter { predicate => (predicate.name(), predicate.children()) match { case (">", Array(column: NamedReference, literal: Literal[_])) if column.fieldNames().sameElements(Array("id")) => diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index b544b43dd04fa..3776ddf633f8c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -24,6 +24,7 @@ import java.sql.{Date, Timestamp} import java.time.{Instant, LocalDate} import java.util.UUID import java.util.concurrent.{Callable, CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicReference import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ @@ -31,8 +32,9 @@ import scala.jdk.CollectionConverters._ import com.fasterxml.jackson.core.json.JsonWriteFeature import com.fasterxml.jackson.databind.ObjectMapper import com.fasterxml.jackson.databind.json.JsonMapper +import org.scalatest.time.SpanSugar._ -import org.apache.spark.{SparkClassNotFoundException, SparkFiles, SparkRuntimeException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.{SPARK_VERSION, SparkClassNotFoundException, SparkFiles, SparkRuntimeException, SparkThrowable, SparkUnsupportedOperationException} import org.apache.spark.api.python.{PythonBroadcast, SimplePythonFunction} import org.apache.spark.broadcast.Broadcast import org.apache.spark.sql.{AnalysisException, QueryTest, Row} @@ -76,6 +78,14 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { compile("test_native_datasource_v2", "-DABI_VERSION=2", "-DNAME_PREFIX=\"v2_\"") private lazy val noLastErrorLibrary = compile("test_native_datasource_ne", "-DNO_LAST_ERROR", "-DNAME_PREFIX=\"ne_\"") + private lazy val smallTableLibrary = compile("test_native_datasource_small", + "-DTABLE_SIZE=offsetof(SparkDsLibrary,read)", "-DNAME_PREFIX=\"small_\"") + private lazy val failingInitLibrary = + compile("test_native_datasource_if", "-DINIT_FAILURE=1", "-DNAME_PREFIX=\"if_\"") + private lazy val noTableLibrary = + compile("test_native_datasource_nt", "-DINIT_FAILURE=2", "-DNAME_PREFIX=\"nt_\"") + private lazy val noInitLibrary = + compile("test_native_datasource_ni", "-DNO_INIT", "-DNAME_PREFIX=\"ni_\"") private lazy val installedLibrary = compile("test_native_datasource_installed", "-DNAME_PREFIX=\"installed_\"") private lazy val installedNegatingLibrary = compile( @@ -131,6 +141,19 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { .getOrElse(throw e) } + /** + * Checks that Spark closed each data writer that native_sink created under `dir`, and did not + * use it afterwards. When a write fails, the tasks that Spark kills may still close their data + * writers after the write returns. + */ + private def checkDataWritersClosed(dir: File): Unit = eventually(timeout(1.minute)) { + def names(subdir: String): Set[String] = + Option(new File(dir, subdir).list()).map(_.toSet).getOrElse(Set.empty) + assert(names("_created_data_writers").nonEmpty) + assert(names("_closed_data_writers") == names("_created_data_writers")) + assert(!new File(dir, "_DATA_WRITER_USED_AFTER_CLOSE").exists()) + } + // Formats a double as the test library does. private def format(d: Double): String = if (d == d.toLong) d.toLong.toString else d.toString @@ -195,7 +218,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { """"children":[{"type":"column","name":["id"]},""" + """{"type":"literal","dataType":"bigint","value":1},""" + """{"type":"literal","dataType":"bigint","value":2}]}]}""")) - // Only the functions listed in spark_ds_push_predicates are passed, so not arithmetic. + // Only the functions listed in push_predicates are passed, so not arithmetic. NativePredicates.FUNCTION_NAMES.foreach { name => assert(toJson(predicate(name, id, id)).nonEmpty, name) } @@ -255,7 +278,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { /** * Generates `org/apache/spark/sql/datasource/predicate-cases.json`, which lists predicates as - * Spark passes them to `spark_ds_push_predicates`, with the rows of a table for which Spark + * Spark passes them to `push_predicates`, with the rows of a table for which Spark * finds them true, for libraries to test their pushdown with. Regenerate it with * `SPARK_GENERATE_GOLDEN_FILES=1`. */ @@ -311,8 +334,8 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } val generated = s"""{ - | "description": ${json("Predicates as Spark passes them to spark_ds_push_" + - "predicates, with the ids of the rows of the table for which they are true. A " + + | "description": ${json("Predicates as Spark passes them to push_predicates of " + + "spark_datasource.h, with the ids of the rows of the table for which they are true. A " + "case lists the predicates that Spark pushes for the filter in its SQL, which are " + "combined with AND. A row has a value for each column of the schema; NaN and infinite " + "doubles are strings.")}, @@ -341,14 +364,14 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } test("tell missing functions from link errors") { - // A function that the library does not export is an unsupported operation. - val missing = new NativeLibrary.MissingFunctionException("spark_ds_f") + // A function that the library leaves NULL in its table is an unsupported operation. + val missing = new NativeLibrary.MissingFunctionException("push_limit") assert(NativeCalls.optional(throw missing).isEmpty) checkError( exception = intercept[SparkRuntimeException](NativeCalls.call("x", "read")(throw missing)), condition = "NATIVE_DATA_SOURCE_ERROR", parameters = Map( - "action" -> "read", "name" -> "x", "msg" -> "The library does not implement spark_ds_f.")) + "action" -> "read", "name" -> "x", "msg" -> "The library does not implement push_limit.")) // Link errors are reported as they are, such as when Arrow, which still calls the Arrow C // structs through JNI, cannot load its JNI library, which it extracts to java.io.tmpdir, for // example if the directory does not allow it. @@ -391,6 +414,33 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } } + test("the bindings match the layout of spark_datasource.h") { + assume(compiler.isDefined, "A C++ compiler is required.") + val layout = NativeLibrary.layout().asScala.map { case (key, value) => key -> value.longValue } + val prints = layout.keys.map { key => + val value = key.split('.') match { + case Array(struct) => s"sizeof($struct)" + case Array(struct, field) => s"offsetof($struct, $field)" + } + s""" std::printf("%s %zu\\n", "$key", $value);""" + } + val output = runProgram("layout", + s"""#include + |#include + |#include "spark_datasource.h" + | + |int main() { + |${prints.mkString("\n")} + | return 0; + |} + |""".stripMargin) + val actual = output.linesIterator.map { line => + val Array(key, value) = line.split(" ") + key -> value.toLong + }.toMap + assert(actual == layout.toMap) + } + nativeTest("disable native data sources") { val conf = new SQLConf conf.setConf(SQLConf.NATIVE_DATA_SOURCE_PATHS, Seq(defaultPackage.getPath)) @@ -464,6 +514,17 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { checkRows(claimed, (16L until 20L).map(expectedRow)) assert(scan(claimed).getMetaData()("PushedPredicates").contains("id > 15")) + // A library that only uses the predicates to skip data, here by reading the blocks of 10 + // rows that may match, with all their rows. Spark shows these predicates as pushed too, and + // so it does with a value that it does not know. + Seq("partial", "7").foreach { use => + val skipping = range(newSession(defaultPackage), "end" -> "45", "predicate_use" -> use) + .filter("id > 13 AND id < 27") + checkRows(skipping, (14L until 27L).map(expectedRow)) + val pushed = scan(skipping).getMetaData()("PushedPredicates") + assert(pushed.contains("id > 13") && pushed.contains("id < 27"), pushed) + } + // The answers of the library are for the predicates that Spark passes to it: here, not the // first one, which uses another function. val mapped = range(newSession(defaultPackage), "end" -> "20") @@ -584,8 +645,9 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { write("append", 0 until 4) assert(written() == (0 until 4).map(i => s"$i,v$i").toSet) assert(new File(dir, "_SUCCESS").isFile) - // The writer is closed after the commit. + // The writer is closed after the commit, and so are the data writers of the tasks. assert(new File(dir, "_CLOSED").isFile) + checkDataWritersClosed(dir) write("append", 4 until 6) assert(written() == (0 until 6).map(i => s"$i,v$i").toSet) write("overwrite", 10 until 12) @@ -691,16 +753,18 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { "action" -> "write to", "name" -> "native_sink", "msg" -> "injected failure in commitDataWriter")) - // abortDataWriter removed the temporary files of the tasks. + // abort_data_writer removed the temporary files of the tasks, whose data writers were then + // closed. assert(new File(dir, "_ABORTED").isFile) assert(new File(dir, "_temporary").listFiles().isEmpty) assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) + checkDataWritersClosed(dir) } nativeTest("do not abort a data writer that committed with an invalid message") { val session = newSession(defaultPackage) val dir = Utils.createTempDir() - // The library commits and releases the data writer, but passes two commit messages. + // The library commits the data writer, but passes two commit messages. val e = intercept[Exception] { session.range(10).repartition(1).write.format("native_sink").option("path", dir.getPath) .option("commit_messages", "2").mode("append").save() @@ -711,9 +775,11 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { parameters = Map( "action" -> "write to", "name" -> "native_sink", - "msg" -> "spark_ds_commit_data_writer returned 2 values instead of one.")) - // The task failed without aborting the released data writer, and the driver aborted the write. - assert(!new File(dir, "_ABORTED_RELEASED_DATA_WRITER").exists()) + "msg" -> "commit_data_writer returned 2 values instead of one.")) + // The task failed without aborting the committed data writer, which it closed, and the driver + // aborted the write. + assert(!new File(dir, "_ABORTED_COMMITTED_DATA_WRITER").exists()) + checkDataWritersClosed(dir) assert(new File(dir, "_ABORTED").isFile) assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) } @@ -804,10 +870,10 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { condition = "NATIVE_DATA_SOURCE_ERROR", parameters = Map("action" -> "plan a scan of", "name" -> "native_counter", "msg" -> msg)) } - check("none", "spark_ds_latest_offset returned no offset.") - check("multiline", "spark_ds_latest_offset returned an offset with a line break, but an " + + check("none", "latest_offset returned no offset.") + check("multiline", "latest_offset returned an offset with a line break, but an " + "offset must be JSON on a single line.") - check("invalid", "spark_ds_latest_offset returned an offset that is not valid JSON: offset 10") + check("invalid", "latest_offset returned an offset that is not valid JSON: offset 10") } nativeTest("pass the ID of a streaming query to its stream writers") { @@ -865,11 +931,13 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { "NULL.") { session => range(session, "fail" -> "uninitialized").collect() } - // An error of the library is not an unsupported operation. + // An error of the library is not an unsupported operation: here, of a data source that + // declares batch reads that it does not support. check("plan a scan of", "native_sink does not support batch reads", "native_sink") { - session => session.read.schema("id BIGINT").format("native_sink").load().collect() + session => session.read.schema("id BIGINT").format("native_sink") + .option("capabilities", NativeLibrary.BATCH_READ.toString).load().collect() } - // Without spark_ds_last_error, only the value that the function returned is reported. + // Without last_error, only the value that the function returned is reported. val e = intercept[Exception] { newSession(createPackage(noLastErrorLibrary, dataSources("ne_"))) .read.format("ne_native_range").option("fail", "create").load() @@ -880,7 +948,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { parameters = Map( "action" -> "create", "name" -> "ne_native_range", - "msg" -> "spark_ds_create failed with the error code 1.")) + "msg" -> "create failed with the error code 1.")) } nativeTest("a scan whose planning failed fails again with the same error") { @@ -919,71 +987,204 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { val query = start try query.processAllAvailable() finally query.stop() } - check("spark_ds_create", "create", "native_range") { session => - range(session, "zero" -> "spark_ds_create") + check("create", "create", "native_range") { session => + range(session, "zero" -> "create") } - check("spark_ds_create_reader", "plan a scan of", "native_range") { session => - range(session, "zero" -> "spark_ds_create_reader").collect() + check("create_reader", "plan a scan of", "native_range") { session => + range(session, "zero" -> "create_reader").collect() } - check("spark_ds_create_stream_reader", "plan a scan of", "native_counter") { session => + check("create_stream_reader", "plan a scan of", "native_counter") { session => runStream { session.readStream.format("native_counter") - .option("zero", "spark_ds_create_stream_reader").load() + .option("zero", "create_stream_reader").load() .writeStream.format("noop").start() } } - Seq( - "spark_ds_create_writer" -> "plan a write to", - "spark_ds_create_data_writer" -> "write to" - ).foreach { case (function, action) => - check(function, action, "native_sink") { session => - session.range(2).write.format("native_sink") - .option("path", Utils.createTempDir().getPath).option("zero", function) - .mode("append").save() - } + Seq("create_writer" -> "plan a write to", "create_data_writer" -> "write to").foreach { + case (function, action) => + check(function, action, "native_sink") { session => + session.range(2).write.format("native_sink") + .option("path", Utils.createTempDir().getPath).option("zero", function) + .mode("append").save() + } } - check("spark_ds_create_stream_writer", "plan a write to", "native_sink") { session => + check("create_writer", "plan a write to", "native_sink") { session => runStream { session.readStream.format("native_counter").load() .writeStream .format("native_sink") .option("path", Utils.createTempDir().getPath) - .option("zero", "spark_ds_create_stream_writer") + .option("zero", "create_writer") .option("checkpointLocation", Utils.createTempDir().getPath) .start() } } } - nativeTest("operations that the library does not implement") { - val session = newSession(createPackage(readOnlyLibrary, dataSources("ro_"))) + nativeTest("operations that a data source does not support") { + val session = newSession(defaultPackage, createPackage(readOnlyLibrary, dataSources("ro_"))) // The library does not push anything down, so Spark evaluates the filter. val df = session.read.format("ro_native_range").option("end", "10").load().filter("id > 6") checkRows(df, (7L until 10L).map(expectedRow)) assert(scan(df).getMetaData()("PushedPredicates") == "[]") + def check(condition: String, parameters: Map[String, String])(f: => Unit): Unit = { + val e = intercept[Exception](f) + checkError(exception = findError(e, condition), condition = condition, + parameters = parameters) + } + def runStream(start: => StreamingQuery): Unit = { + val query = start + try query.processAllAvailable() finally query.stop() + } + // A data source declares the operations that it supports with its capabilities. Spark + // reports the others as unsupported, without calling the library, which here does not + // implement them either. + check("DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", Map("description" -> "ro_native_sink")) { + session.range(1).write.format("ro_native_sink").mode("append").save() + } + check("DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> "ro_native_counter")) { + runStream { + session.readStream.format("ro_native_counter").load().writeStream.format("noop").start() + } + } + check("DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", Map("description" -> "ro_native_sink")) { + runStream { + session.readStream.format("rate").load().writeStream.format("ro_native_sink") + .option("checkpointLocation", Utils.createTempDir().getPath).start() + } + } + // The capabilities are those of each data source: native_counter only supports streaming + // reads, although its library implements batch reads. + check("DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> "native_counter")) { + session.read.format("native_counter").load().collect() + } + // A data source that declares a capability must implement its functions. + check("NATIVE_DATA_SOURCE_ERROR", Map("action" -> "plan a write to", + "name" -> "ro_native_sink", "msg" -> "The library does not implement create_writer.")) { + session.range(1).write.format("ro_native_sink") + .option("capabilities", NativeLibrary.BATCH_WRITE.toString).mode("append").save() + } + // Replacing the existing data requires the capability SPARK_DS_TRUNCATE. + val dir = Utils.createTempDir() + val noTruncate = (NativeLibrary.BATCH_WRITE | NativeLibrary.STREAMING_WRITE).toString + check("NATIVE_DATA_SOURCE_ERROR", Map("action" -> "plan a write to", "name" -> "native_sink", + "msg" -> "The data source does not support the save mode overwrite.")) { + session.range(1).write.format("native_sink").option("path", dir.getPath) + .option("capabilities", noTruncate).mode("overwrite").save() + } + check("NATIVE_DATA_SOURCE_ERROR", Map("action" -> "plan a write to", "name" -> "native_sink", + "msg" -> "The data source does not support the output mode complete.")) { + runStream { + session.readStream.format("rate").load().groupBy().count().writeStream + .outputMode("complete") + .format("native_sink") + .option("path", dir.getPath) + .option("capabilities", noTruncate) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + } + } + assert(!new File(dir, "_CREATED").exists()) + } + + nativeTest("read only the functions that fit in the table of a library") { + // The table of the library ends before read, as the table of a library built against an + // earlier header ends before the functions added later: Spark does not read past its end. + val session = newSession(createPackage(smallTableLibrary, dataSources("small_"))) + val df = session.read.format("small_native_range").load() + assert(df.schema.length == 7) checkError( - exception = intercept[SparkUnsupportedOperationException] { - session.range(1).write.format("ro_native_sink").mode("append").save() - }, - condition = "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", - parameters = Map("description" -> "ro_native_sink")) - def checkStreaming(condition: String, name: String)(start: => StreamingQuery): Unit = { - val e = intercept[Exception] { - val query = start - try query.processAllAvailable() finally query.stop() + exception = findError(intercept[Exception](df.collect()), "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "plan a scan of", + "name" -> "small_native_range", + "msg" -> "The library does not implement read.")) + } + + nativeTest("pass the version of Spark, and log with the log of Spark") { + val session = newSession(defaultPackage) + checkRows(range(session, "end" -> "2", "spark_version" -> SPARK_VERSION).select("id"), + Seq(Row(0L), Row(1L))) + checkError( + exception = findError( + intercept[Exception](range(session, "spark_version" -> "1.0.0")), + "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "create", + "name" -> "native_range", + "msg" -> (s"the version of Spark is $SPARK_VERSION, which supports version 1 of the " + + "interface"))) + + val logAppender = new LogAppender("native data source logs") + withLogAppender(logAppender, loggerNames = Seq(classOf[NativeLibrary].getName)) { + range(session, "log" -> "a message from native code") + } + assert(logAppender.loggingEvents.exists { event => + val message = event.getMessage.getFormattedMessage + event.getLevel == org.apache.logging.log4j.Level.WARN && + message.contains("a message from native code") && + message.contains(System.mapLibraryName("test_native_datasource")) + }) + } + + nativeTest("pass the task to reads and data writers") { + val session = newSession(defaultPackage) + // Checks the SparkDsTaskContext that the library got in each task, by partition. + def check(dir: File): Unit = { + val tasks = dir.listFiles().map { file => + file.getName.toInt -> Files.readString(file.toPath).split(",").map(_.toLong).toSeq + }.toMap + assert(tasks.keySet == Set(0, 1)) + tasks.foreach { + case (partition, Seq(size, stageId, stageAttempt, partitionId, attempt, taskAttemptId)) => + assert(size == NativeLibrary.layout().get("SparkDsTaskContext").longValue) + assert(partitionId == partition) + assert(stageId >= 0 && stageAttempt == 0 && attempt == 0 && taskAttemptId >= 0) } - checkError( - exception = findError(e, condition), - condition = condition, - parameters = Map("description" -> name)) + assert(tasks.values.map(_(5)).toSet.size == 2) } - checkStreaming("DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", "ro_native_counter") { - session.readStream.format("ro_native_counter").load().writeStream.format("noop").start() + val readTasks = Utils.createTempDir() + checkRows(range(session, "end" -> "6", "task_path" -> readTasks.getPath).select("id"), + (0L until 6L).map(Row(_))) + check(readTasks) + val writeTasks = Utils.createTempDir() + session.range(6).repartition(2).write.format("native_sink") + .option("path", Utils.createTempDir().getPath) + .option("task_path", writeTasks.getPath) + .mode("append") + .save() + check(writeTasks) + } + + nativeTest("tell a read that its task was killed") { + val session = newSession(defaultPackage) + val marker = new File(Utils.createTempDir(), "interrupted") + val df = range(session, "end" -> "10", "partitions" -> "1", + "wait_for_interrupt" -> marker.getPath) + val jobGroup = s"native-data-source-${UUID.randomUUID()}" + val error = new AtomicReference[Throwable]() + val thread = new Thread(() => { + session.sparkContext.setJobGroup(jobGroup, "a read that waits", interruptOnCancel = true) + try df.collect() catch { + case e: Throwable => error.set(e) + } + }) + thread.start() + // The read waits until its task is killed, which it only learns from is_interrupted: Spark + // cannot interrupt native code. + eventually(timeout(1.minute)) { + assert(new File(marker.getPath + ".waiting").isFile) } - checkStreaming("DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", "ro_native_sink") { - session.readStream.format("rate").load().writeStream.format("ro_native_sink") - .option("checkpointLocation", Utils.createTempDir().getPath).start() + session.sparkContext.cancelJobGroup(jobGroup) + thread.join(TimeUnit.MINUTES.toMillis(1)) + assert(!thread.isAlive) + assert(error.get() != null) + // The job fails right away, and its task ends once the library sees that it was killed. + eventually(timeout(1.minute)) { + assert(marker.isFile && Files.readString(marker.toPath) == "interrupted") } } @@ -1522,4 +1723,34 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } } } + + nativeTest("reject the libraries that spark_ds_init does not initialize") { + def check(library: File, prefix: String, reason: String): Unit = { + val pkg = createPackage(library, dataSources(prefix)) + // The library stays loaded, and the error is the same when it is used again: Spark + // initializes a library only once. + (1 to 2).foreach { _ => + checkError( + exception = intercept[SparkRuntimeException] { + newSession(pkg).read.format(prefix + "native_range").load() + }, + condition = "INVALID_NATIVE_DATA_SOURCE_PACKAGE.INVALID_LIBRARY", + parameters = Map( + "path" -> pkg.getCanonicalPath, + "library" -> s"${NativePlatform.current}/${library.getName}", + "reason" -> reason)) + } + } + val logAppender = new LogAppender("spark_ds_init") + withLogAppender(logAppender, loggerNames = Seq(classOf[NativeLibrary].getName)) { + check(failingInitLibrary, "if_", "spark_ds_init failed with the error code 1.") + } + // The library logged why it failed, once. + assert(logAppender.loggingEvents.count { event => + event.getLevel == org.apache.logging.log4j.Level.ERROR && + event.getMessage.getFormattedMessage.contains("injected failure in spark_ds_init") + } == 1) + check(noTableLibrary, "nt_", "spark_ds_init returned no function table.") + check(noInitLibrary, "ni_", "The library does not export the function spark_ds_init.") + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala index cc2ec905b06ad..938f838c6f6bb 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala @@ -41,13 +41,9 @@ object NativeDataSourceTestUtils { // each package has its own name. private val packageCount = new AtomicInteger() - /** - * The C++ compiler, if available on a version of Java that can load native data sources, and - * the include directory with `spark_datasource.h`. - */ - lazy val toolchain: Option[(String, Seq[String])] = { - val compiler = Seq("c++", "g++", "clang++").find(Utils.checkCommandAvailable) - compiler.filter(_ => NativeLibrary.isSupported).map { compiler => + /** The C++ compiler, if available, and the include directory with `spark_datasource.h`. */ + lazy val compiler: Option[(String, Seq[String])] = { + Seq("c++", "g++", "clang++").find(Utils.checkCommandAvailable).map { compiler => val header = new File(buildDir, "spark_datasource.h") val in = getClass.getClassLoader .getResourceAsStream("org/apache/spark/sql/datasource/spark_datasource.h") @@ -56,6 +52,10 @@ object NativeDataSourceTestUtils { } } + /** [[compiler]], if available on a version of Java that can load native data sources. */ + lazy val toolchain: Option[(String, Seq[String])] = + compiler.filter(_ => NativeLibrary.isSupported) + /** * Compiles the test library with the given preprocessor definitions, once for each name. The * library implements the data sources [[dataSources]], with the prefix `NAME_PREFIX`. @@ -73,12 +73,29 @@ object NativeDataSourceTestUtils { Utils.tryWithResource(in)(Files.copy(_, source.toPath)) } val library = new File(buildDir, System.mapLibraryName(name)) - val command = Seq(compiler, "-std=c++17", "-O1", "-shared", "-fPIC") ++ - includes.map("-I" + _) ++ defines ++ Seq("-o", library.getPath, source.getPath) + // The symbols are hidden, as libraries are often built, so only spark_ds_init is exported. + run(Seq(compiler, "-std=c++17", "-O1", "-shared", "-fPIC", "-fvisibility=hidden") ++ + includes.map("-I" + _) ++ defines ++ Seq("-o", library.getPath, source.getPath)) + library + } + + /** Compiles a C++ program that includes `spark_datasource.h`, runs it and returns its output. */ + def runProgram(name: String, source: String): String = { + val (compiler, includes) = this.compiler.get + val sourceFile = new File(buildDir, s"$name.cc") + Files.writeString(sourceFile.toPath, source) + val program = new File(buildDir, name) + run(Seq(compiler, "-std=c++17") ++ includes.map("-I" + _) ++ + Seq("-o", program.getPath, sourceFile.getPath)) + run(Seq(program.getPath)) + } + + /** Runs a command, checks that it succeeds and returns its output. */ + private def run(command: Seq[String]): String = { val process = new ProcessBuilder(command: _*).redirectErrorStream(true).start() val output = new String(process.getInputStream.readAllBytes(), UTF_8) - assert(process.waitFor() == 0, s"Failed to compile the test library: $output") - library + assert(process.waitFor() == 0, s"Failed to run ${command.mkString(" ")}: $output") + output } /** The data sources of the test library compiled with the given `NAME_PREFIX`. */ From b36ec5e73e27c1a5d721926a0d0c414adc176b2c Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Thu, 15 Oct 2026 12:43:35 +0900 Subject: [PATCH 24/27] Pass the name and the options of a data source to create in a size-prefixed struct create was the one entry point whose inputs could only grow with a second function, since it took the name and the options as separate arguments. It now takes a SparkDsCreateInfo, which starts with its size like the other structs that Spark passes, so that later releases can append fields, for example references to credentials instead of the credentials themselves in the options. The test library and the examples in the docs read the struct, and the layout test covers it. Co-authored-by: Isaac --- docs/sql-native-data-sources.md | 51 +++++++++++-------- .../datasources/v2/ffi/NativeLibrary.java | 23 +++++++-- .../spark/sql/datasource/spark_datasource.h | 40 ++++++++++----- .../test_native_datasource.cc | 16 +++--- 4 files changed, 85 insertions(+), 45 deletions(-) diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 4f07d96a283f8..25965a2a2a67c 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -69,10 +69,11 @@ A library built against the header works with all the releases of Spark that sup Spark calls the functions of the table as follows: - **Data sources**: for each schema inference, scan and write, Spark creates a data source on the - driver with `create`, which receives the options of the user and returns the capabilities of the - data source, such as `SPARK_DS_BATCH_READ` or `SPARK_DS_STREAMING_WRITE`, and closes it right - after with `close`. Spark reports the operations that a data source does not declare as - unsupported. When the user does not specify a schema, Spark asks for it with `schema`. + driver with `create`, which receives the name of the data source and the options of the user in + a `SparkDsCreateInfo`, and returns the capabilities of the data source, such as + `SPARK_DS_BATCH_READ` or `SPARK_DS_STREAMING_WRITE`, and closes it right after with `close`. + Spark reports the operations that a data source does not declare as unsupported. When the user + does not specify a schema, Spark asks for it with `schema`. - **Batch reads**: on the driver, Spark creates a reader with `create_reader`, pushes operations down with `push_predicates`, `push_limit` and `prune_columns`, plans the partitions with `partitions`, and serializes the state of the reader with `serialize_reader`. Then it calls @@ -160,9 +161,9 @@ regularly and use timeouts. The task reaper, `spark.task.reaper.enabled` and ### Compatibility The function table, `SparkDsLibrary`, and the structs that Spark passes to the library, -`SparkDsHost`, `SparkDsTaskContext` and `SparkDsWriteInfo`, start with their size, so that they can -grow. The version of the interface, 1, is `SPARK_DS_ABI_VERSION`, which the table and the manifest -of a package declare. Within a version: +`SparkDsHost`, `SparkDsCreateInfo`, `SparkDsTaskContext` and `SparkDsWriteInfo`, start with their +size, so that they can grow. The version of the interface, 1, is `SPARK_DS_ABI_VERSION`, which the +table and the manifest of a package declare. Within a version: - fields are only appended to these structs, and the existing fields and the semantics of the existing functions never change; @@ -354,6 +355,16 @@ type Sink = unsafe extern "C" fn(sink: *mut c_void, data: *const u8, length: i64 /// A function of `SparkDsLibrary` that this data source does not implement, left NULL. type Unused = Option; +/// `SparkDsCreateInfo` of `spark_datasource.h`: the name and the options of a data source. +#[repr(C)] +pub struct SparkDsCreateInfo { + size: i64, + name: *const c_char, + option_keys: *const *const c_char, + option_values: *const *const c_char, + num_options: i64, +} + /// `SparkDsLibrary` of `spark_datasource.h`, up to `read`, the last function that this data /// source implements: Spark only reads the functions that fit in `size`. #[repr(C)] @@ -362,10 +373,7 @@ pub struct SparkDsLibrary { abi_version: i32, last_error: unsafe extern "C" fn() -> *const c_char, create: unsafe extern "C" fn( - name: *const c_char, - option_keys: *const *const c_char, - option_values: *const *const c_char, - num_options: i64, + info: *const SparkDsCreateInfo, capabilities: *mut u64, source: *mut *mut c_void, ) -> c_int, @@ -469,18 +477,16 @@ unsafe extern "C" fn last_error() -> *const c_char { } unsafe extern "C" fn create( - _name: *const c_char, - option_keys: *const *const c_char, - option_values: *const *const c_char, - num_options: i64, + info: *const SparkDsCreateInfo, capabilities: *mut u64, source: *mut *mut c_void, ) -> c_int { call(|| { + let info = &*info; let mut range = RangeSource { end: 10 }; - for i in 0..num_options as usize { - if CStr::from_ptr(*option_keys.add(i)).to_str()? == "end" { - range.end = CStr::from_ptr(*option_values.add(i)).to_str()?.parse()?; + for i in 0..info.num_options as usize { + if CStr::from_ptr(*info.option_keys.add(i)).to_str()? == "end" { + range.end = CStr::from_ptr(*info.option_values.add(i)).to_str()?.parse()?; } } *capabilities = SPARK_DS_BATCH_READ; @@ -639,12 +645,13 @@ void Check(const arrow::Status& status) { const char* LastError() { return last_error.c_str(); } -int Create(const char*, const char* const* option_keys, const char* const* option_values, - int64_t num_options, uint64_t* capabilities, void** source) { +int Create(const SparkDsCreateInfo* info, uint64_t* capabilities, void** source) { return Call([&] { auto range = std::make_unique(); - for (int64_t i = 0; i < num_options; i++) { - if (std::string(option_keys[i]) == "end") range->end = std::stoll(option_values[i]); + for (int64_t i = 0; i < info->num_options; i++) { + if (std::string(info->option_keys[i]) == "end") { + range->end = std::stoll(info->option_values[i]); + } } *capabilities = SPARK_DS_BATCH_READ; *source = range.release(); diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java index 6c2368ea7c224..66c1a46b7beea 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -90,7 +90,7 @@ private enum C { VOID, BOOL, I32, I64, PTR } /** The functions of the table {@code SparkDsLibrary}, in their order in the table. */ enum Function { LAST_ERROR("last_error", C.PTR), - CREATE("create", C.I32, C.PTR, C.PTR, C.PTR, C.I64, C.PTR, C.PTR), + CREATE("create", C.I32, C.PTR, C.PTR, C.PTR), SCHEMA("schema", C.I32, C.PTR, C.PTR), CREATE_READER("create_reader", C.I32, C.PTR, C.PTR, C.PTR), CREATE_STREAM_READER("create_stream_reader", C.I32, C.PTR, C.PTR, C.PTR), @@ -144,6 +144,11 @@ long offset() { private static final long HOST_ABI_VERSION = 8; private static final long HOST_SPARK_VERSION = 16; private static final long HOST_LOG = 24; + private static final long CREATE_INFO_SIZE = 40; + private static final long CREATE_INFO_NAME = 8; + private static final long CREATE_INFO_OPTION_KEYS = 16; + private static final long CREATE_INFO_OPTION_VALUES = 24; + private static final long CREATE_INFO_NUM_OPTIONS = 32; private static final long TASK_SIZE = 40; private static final long TASK_STAGE_ID = 8; private static final long TASK_STAGE_ATTEMPT_NUMBER = 12; @@ -172,6 +177,12 @@ static Map layout() { layout.put("SparkDsHost.abi_version", HOST_ABI_VERSION); layout.put("SparkDsHost.spark_version", HOST_SPARK_VERSION); layout.put("SparkDsHost.log", HOST_LOG); + layout.put("SparkDsCreateInfo", CREATE_INFO_SIZE); + layout.put("SparkDsCreateInfo.size", 0L); + layout.put("SparkDsCreateInfo.name", CREATE_INFO_NAME); + layout.put("SparkDsCreateInfo.option_keys", CREATE_INFO_OPTION_KEYS); + layout.put("SparkDsCreateInfo.option_values", CREATE_INFO_OPTION_VALUES); + layout.put("SparkDsCreateInfo.num_options", CREATE_INFO_NUM_OPTIONS); layout.put("SparkDsTaskContext", TASK_SIZE); layout.put("SparkDsTaskContext.size", 0L); layout.put("SparkDsTaskContext.stage_id", TASK_STAGE_ID); @@ -542,11 +553,15 @@ DataSourceHandle createDataSource(String name, String[] optionKeys, String[] opt throws Throwable { MethodHandle f = function(Function.CREATE); try (Memory memory = new Memory()) { + long info = memory.zeroed(CREATE_INFO_SIZE); + Platform.putLong(null, info, CREATE_INFO_SIZE); + Platform.putLong(null, info + CREATE_INFO_NAME, memory.string(name)); + Platform.putLong(null, info + CREATE_INFO_OPTION_KEYS, memory.strings(optionKeys)); + Platform.putLong(null, info + CREATE_INFO_OPTION_VALUES, memory.strings(optionValues)); + Platform.putLong(null, info + CREATE_INFO_NUM_OPTIONS, optionKeys.length); long capabilities = memory.zeroed(8); long out = memory.zeroed(8); - check((int) f.invokeExact(memory.string(name), memory.strings(optionKeys), - memory.strings(optionValues), (long) optionKeys.length, capabilities, out), - Function.CREATE); + check((int) f.invokeExact(info, capabilities, out), Function.CREATE); return new DataSourceHandle( readHandle(out, Function.CREATE), Platform.getLong(null, capabilities)); } diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h index dd55f6795aa4d..3ff46d2871560 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -111,11 +111,12 @@ * released once they are no longer referenced (see Lifecycle above). A reader or writer handle * must not depend on the data source handle it was created from, because Spark closes the data * source handle right after creating it. - * - Memory. What Spark passes to a function, such as strings, byte arrays and SparkDsWriteInfo, is - * only valid during the call, except for SparkDsHost and SparkDsTaskContext, which document how - * long they live. Strings are UTF-8 and NUL-terminated. The JSON of the predicates is ASCII: it - * escapes the other characters, which any JSON parser decodes. The library returns strings and - * byte arrays by passing them to a SparkDsSink, which copies them. + * - Memory. What Spark passes to a function, such as strings, byte arrays, SparkDsCreateInfo and + * SparkDsWriteInfo, is only valid during the call, except for SparkDsHost and + * SparkDsTaskContext, which document how long they live. Strings are UTF-8 and NUL-terminated. + * The JSON of the predicates is ASCII: it escapes the other characters, which any JSON parser + * decodes. The library returns strings and byte arrays by passing them to a SparkDsSink, which + * copies them. * - State. Partitions, the state of readers and writers, and commit messages are byte arrays in a * format defined by the library. Spark sends them to the executors and back. * - Arrow data. Schemas and data are exchanged with the Arrow C data interface @@ -271,6 +272,24 @@ typedef struct SparkDsHost { void (*log)(int32_t level, const char* message); } SparkDsHost; +/* A data source, as create gets it. */ +typedef struct SparkDsCreateInfo { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* + * The name of the data source, as listed in the package manifest, or in lower case for an + * installed library. A library can implement several data sources. + */ + const char* name; + /* + * The num_options options given by the user, such as spark.read.option("key", "value"), with + * the keys in lower case. + */ + const char* const* option_keys; + const char* const* option_values; + int64_t num_options; +} SparkDsCreateInfo; + /* * The task that calls read or create_data_writer, on an executor. It stays valid until the stream * that read returns is released, or until the data writer is closed. @@ -360,19 +379,14 @@ typedef struct SparkDsLibrary { * ------------------------------------------------------------------------------------------- */ /* - * Creates a data source with the options given by the user, such as - * spark.read.option("key", "value"). + * Creates a data source. * - * name: the name of the data source, as listed in the package manifest, or in lower case for an - * installed library. A library can implement several data sources. - * option_keys, option_values: the num_options options, with the keys in lower case + * info: the name of the data source, and the options given by the user * capabilities: set to the capabilities of the data source, SPARK_DS_* flags combined with |, * which may depend on its name and its options. It is 0 when create is called. * source: set to the data source handle */ - int (*create)(const char* name, const char* const* option_keys, - const char* const* option_values, int64_t num_options, uint64_t* capabilities, - void** source); + int (*create)(const SparkDsCreateInfo* info, uint64_t* capabilities, void** source); /* * Exports the schema of the data. Spark calls it only when the user does not specify a schema. diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index 6ba8fa78ff5b0..dc193afba889b 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -883,13 +883,17 @@ uint64_t capabilitiesOf(const DataSource& source) { return 0; } -int createSource(const char* name_value, const char* const* option_keys, - const char* const* option_values, int64_t num_options, uint64_t* capabilities, - void** handle) { +int createSource(const SparkDsCreateInfo* info, uint64_t* capabilities, void** handle) { return guarded([&] { - auto source = std::make_unique(DataSource{name_value, {}}); - auto keys = toStrings(option_keys, num_options); - auto values = toStrings(option_values, num_options); + // Like the host, the struct may be of another version of the header: only the fields that fit + // in its size are read. + if (info->size < static_cast(offsetof(SparkDsCreateInfo, num_options) + + sizeof(info->num_options))) { + throw std::runtime_error("SparkDsCreateInfo has a size of " + std::to_string(info->size)); + } + auto source = std::make_unique(DataSource{info->name, {}}); + auto keys = toStrings(info->option_keys, info->num_options); + auto values = toStrings(info->option_values, info->num_options); for (size_t i = 0; i < keys.size(); i++) source->options[keys[i]] = values[i]; maybeFail(source->options, "create"); if (source->name != name("native_range") && source->name != name("native_sink") && From 51af65c63a9369d793aab1a198ddde4c90c8f85b Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Thu, 15 Oct 2026 14:02:24 +0900 Subject: [PATCH 25/27] Say which sources native data sources fit, and what is out of their scope A native data source fits a source whose client exists only in native code, and that Spark needs to read and write. The docs now lead with that in a new "When to Use a Native Data Source" section, which also says that extensions outside of Data Source V2, such as table functions and SQL commands, are out of the scope of the interface. The Limitations section also lists top-N pushdown, row-level operations and metadata columns. Co-authored-by: Isaac --- docs/sql-native-data-sources.md | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 25965a2a2a67c..7f69eb455e9f6 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -45,6 +45,21 @@ Spark calls native data sources with the Java, a query that uses a native data source fails with `NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION`. +### When to Use a Native Data Source + +A native data source fits a source whose client exists only in native code, such as an in-house +storage system written in C++ or the C SDK of a vendor, and that Spark needs to read and write. +Such a source then needs no JVM connector: the same library works with all the releases of Spark +that support its version of the interface, and with every Spark API, including Spark Connect +clients in any language. + +A native data source provides what Data Source V2 offers to data sources, except for the features +that version 1 of the interface does not cover yet, such as row-level operations: see +[Limitations](#limitations). Extensions of Spark outside of Data Source V2, such as table +functions, SQL commands and optimizer rules, are out of the scope of this interface. A source that +needs them is better served by a JVM connector, which can still read and write through a native +library with the Arrow C stream interface. + ## How It Works A native data source library exports a single C function, `spark_ds_init`, which returns a table @@ -191,9 +206,16 @@ later as functions or fields appended to the structs of the interface: - custom metrics, shown in the Spark UI; - admission control for streaming reads, such as a maximum number of rows for each micro-batch, and continuous streaming; -- the pushdown of aggregates, of the sort order and of the partitioning of the data; +- the pushdown of aggregates and of top-N queries (`ORDER BY` with `LIMIT`), and reporting the + partitioning and the sort order of the data, with which Spark avoids shuffles and sorts, as in + storage-partitioned joins; +- row-level operations, `UPDATE`, `DELETE` and `MERGE INTO`, and metadata columns; - memory accounting: see below. +Extensions of Spark outside of Data Source V2, such as table functions, SQL commands and optimizer +rules, are out of the scope of this interface: see +[When to Use a Native Data Source](#when-to-use-a-native-data-source). + ### Memory Spark does not manage the memory of native data sources. The memory that a library allocates, and From 2e674a22a2748ffa5ffb74b9fc957a2a35e4eed3 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Thu, 15 Oct 2026 15:20:38 +0900 Subject: [PATCH 26/27] Call native data sources with a frozen C core and versioned messages A draft for comparison with the function table of #59424, following the proposal on it: a library exports six functions that never change, spark_ds_abi_version, spark_ds_last_error, spark_ds_open, spark_ds_call, spark_ds_stream and spark_ds_close, and everything that grows is in the Protocol Buffers messages of spark_datasource.proto: - OpenRequest opens a data source, a batch or stream reader, a writer, a partition reader or a data writer, with everything that it needs. The OpenReply of a data source declares its capabilities. - CallRequest and CallReply carry the control operations. An empty reply is valid for every request. - StreamRequest exchanges the data as Arrow C streams. A data writer gets the data of its task as one stream, from DataWriter.writeAll, which the library reads on the thread of the task. - SparkDsHost has the log and is_interrupted callbacks, and schemas are Spark SQL types. The messages use the lite runtime: the full C++ runtime registers them in a pool of the process, where a second copy of a library aborts the process. The C++ test library implements the messages with a minimal codec, and the docs and their Rust and C++ examples use prost and protobuf-lite. Co-authored-by: Isaac --- dev/make-distribution.sh | 4 +- docs/sql-native-data-sources.md | 677 ++++---- .../datasources/v2/ffi/NativeLibrary.java | 708 ++------ .../sql/datasource/spark_datasource.proto | 487 ++++++ .../spark/sql/datasource/predicate-cases.json | 2 +- .../spark/sql/datasource/spark_datasource.h | 675 ++------ .../v2/columnar/ColumnarWrite.scala | 86 + .../datasources/v2/ffi/NativeDataSource.scala | 612 ++++--- .../v2/ffi/NativeDataSourcePackage.scala | 7 +- .../datasources/v2/ffi/NativePredicates.scala | 4 +- .../test_native_datasource.cc | 1510 +++++++++++------ .../v2/ffi/NativeDataSourceSuite.scala | 188 +- .../v2/ffi/NativeDataSourceTestUtils.scala | 3 +- 13 files changed, 2734 insertions(+), 2229 deletions(-) create mode 100644 sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto diff --git a/dev/make-distribution.sh b/dev/make-distribution.sh index fc73cbc9ac058..f7c29cc2eae81 100755 --- a/dev/make-distribution.sh +++ b/dev/make-distribution.sh @@ -223,10 +223,12 @@ if [ -f "$DISTDIR"/jars/datasketches-memory-3.0.2.jar ]; then zip -d "$DISTDIR"/jars/datasketches-memory-3.0.2.jar org/apache/datasketches/memory/internal/ResourceImpl.class fi -# Copy the C header that the libraries of native data sources implement +# Copy the C header that the libraries of native data sources implement, and its messages mkdir -p "$DISTDIR/include" cp "$SPARK_HOME/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h" \ "$DISTDIR/include/" +cp "$SPARK_HOME/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto" \ + "$DISTDIR/include/" # Only create the yarn directory if the yarn artifacts were built. if [ -f "$SPARK_HOME"/common/network-yarn/target/scala*/spark-*-yarn-shuffle.jar ]; then diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md index 7f69eb455e9f6..138645e859979 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -26,7 +26,8 @@ license: | A native data source is a data source implemented in native code, such as Rust or C++, without any JVM code. It is a shared library that implements the C interface declared in the header -[`spark_datasource.h`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h). Spark finds the library +[`spark_datasource.h`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h), with the +messages of `spark_datasource.proto`. Spark finds the library automatically when it is installed like other native libraries, or when a session adds it in a package, and loads it when a query uses the data source. Spark adapts it to [Data Source V2](sql-v2-data-sources.html), so it supports: @@ -62,98 +63,100 @@ library with the Arrow C stream interface. ## How It Works -A native data source library exports a single C function, `spark_ds_init`, which returns a table -of function pointers, a `SparkDsLibrary`, that implements the data source: it needs no JVM headers -or bindings. Each library has its own table, so any number of native data sources can be loaded -side by side. When Spark loads a library, it calls `spark_ds_init` with a `SparkDsHost`, which -gives the library the version of Spark and a function to log messages to the log of Spark. - -A Spark distribution contains the header in its `include` directory, as -`$SPARK_HOME/include/spark_datasource.h`, and the `spark-sql` JAR contains it as -`org/apache/spark/sql/datasource/spark_datasource.h`. The header only depends on the C standard -library, so a library can also copy it into its source tree, to build without Spark, for example -from the source of a release: +A native data source library exports the six C functions of `spark_datasource.h`, which never +change: it needs no JVM headers or bindings. Everything else is in the messages of +[`spark_datasource.proto`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto), +which Spark and the library exchange as [Protocol Buffers](https://protobuf.dev) bytes: the +library generates their code with a Protocol Buffers compiler, such as `protoc` or `prost`. The +data itself is exchanged without copying it, as Arrow C streams. Spark looks the functions up in +each library separately, so any number of native data sources can be loaded side by side. + +| Function | What it does | +|----------|--------------| +| `spark_ds_open` | Opens an object, with an `OpenRequest` that says which object and carries everything that it needs, and replies with an `OpenReply` | +| `spark_ds_call` | Runs a control operation on an object, with a `CallRequest`, and replies with a `CallReply` | +| `spark_ds_stream` | Exchanges the data of an object, as Arrow C streams, with a `StreamRequest` | +| `spark_ds_close` | Releases an object | +| `spark_ds_abi_version` | Returns the version of the interface, `SPARK_DS_ABI_VERSION` | +| `spark_ds_last_error` | Returns the message of the last error on the calling thread | + +A Spark distribution contains the header and the messages in its `include` directory, as +`$SPARK_HOME/include/spark_datasource.h` and `$SPARK_HOME/include/spark_datasource.proto`, and the +`spark-sql` JAR contains the header as `org/apache/spark/sql/datasource/spark_datasource.h`. They +only depend on the C standard library and on Protocol Buffers, so a library can also copy them +into its source tree, to build without Spark, for example from the source of a release: ```bash curl -LO {{site.SPARK_GITHUB_URL}}/raw/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +curl -LO {{site.SPARK_GITHUB_URL}}/raw/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto ``` -A library built against the header works with all the releases of Spark that support its version, +A library built against them works with all the releases of Spark that support its version, `SPARK_DS_ABI_VERSION`: see [Compatibility](#compatibility). -Spark calls the functions of the table as follows: - -- **Data sources**: for each schema inference, scan and write, Spark creates a data source on the - driver with `create`, which receives the name of the data source and the options of the user in - a `SparkDsCreateInfo`, and returns the capabilities of the data source, such as - `SPARK_DS_BATCH_READ` or `SPARK_DS_STREAMING_WRITE`, and closes it right after with `close`. - Spark reports the operations that a data source does not declare as unsupported. When the user - does not specify a schema, Spark asks for it with `schema`. -- **Batch reads**: on the driver, Spark creates a reader with `create_reader`, pushes operations - down with `push_predicates`, `push_limit` and `prune_columns`, plans the partitions with - `partitions`, and serializes the state of the reader with `serialize_reader`. Then it calls - `read` on the executors, in one task for each partition, and `read` exports an - `ArrowArrayStream`. -- **Streaming reads**: on the driver, Spark creates a stream reader with `create_stream_reader`. - For each micro-batch, it asks for the `latest_offset`, plans the partitions with - `stream_partitions`, and calls `read` on the executors. Offsets are defined by the library, as - compact JSON on a single line: Spark stores each offset as a line of the offset log of the - checkpoint, and rejects an offset that is not valid JSON or contains a line break, such as - pretty-printed JSON. +Spark opens and uses the objects of a library as follows: + +- **Data sources**: for each schema inference, scan and write, Spark opens a data source on the + driver, with the name of the data source and the options of the user. Its `OpenReply` declares + the capabilities of the data source, such as `CAPABILITY_BATCH_READ` or + `CAPABILITY_STREAMING_WRITE`, and Spark reports the operations that a data source does not + declare as unsupported. When the user does not specify a schema, Spark asks for it with + `InferSchema`. Spark closes the data source right after. +- **Batch reads**: on the driver, Spark opens a batch reader, with the schema to read, pushes + operations down with `PushPredicates`, `PushLimit` and `PruneColumns`, plans the partitions with + `PlanPartitions`, and serializes the state of the reader with `Serialize`. Then, on the + executors, it opens a partition reader for each partition, in a task, and `Read` initializes an + `ArrowArrayStream` with the batches of the partition. +- **Streaming reads**: on the driver, Spark opens a stream reader. For each micro-batch, it asks + for the `LatestOffset`, plans the partitions with `PlanStreamPartitions`, and reads them on the + executors. Offsets are defined by the library, as compact JSON on a single line: Spark stores + each offset as a line of the offset log of the checkpoint, and rejects an offset that is not + valid JSON or contains a line break, such as pretty-printed JSON. - **Writes**: on the driver, right before the tasks of the write (or of a micro-batch) start, Spark - creates a writer with `create_writer`, which gets a `SparkDsWriteInfo` with the mode of the write, - `SPARK_DS_WRITE_APPEND` or `SPARK_DS_WRITE_TRUNCATE`, and the ID of the query of a streaming - write. With adaptive query execution, it first runs the shuffle stages of the data to write, and - creates no writer if they fail, but the last stage runs in the tasks of the write, so a writer can - still be aborted because computing the data failed: do irreversible work, such as removing the - existing data for an overwrite, in `commit`. On the executors, each task creates a data writer - with `create_data_writer`, writes the batches of its data with `write`, and calls - `commit_data_writer`, or `abort_data_writer` if the task or the commit fails. Then the driver - calls `commit` with the messages of the tasks, or `abort` if the write failed. A streaming write + opens a writer, with the mode of the write, `WRITE_MODE_APPEND` or `WRITE_MODE_TRUNCATE`, and the + ID of the query of a streaming write. With adaptive query execution, it first runs the shuffle + stages of the data to write, and opens no writer if they fail, but the last stage runs in the + tasks of the write, so a writer can still be aborted because computing the data failed: do + irreversible work, such as removing the existing data for an overwrite, in `Commit`. On the + executors, each task opens a data writer, and `Write` passes it an `ArrowArrayStream` with the + data of the task, which the library reads during the call. Then the task calls + `CommitDataWriter`, or `AbortDataWriter` if the task or the commit fails, and the driver calls + `Commit` with the messages of the tasks, or `Abort` if the write failed. A streaming write commits each micro-batch with its epoch ID, and the ID of the query stays the same when the query restarts from its checkpoint: together, they identify the output of a micro-batch, so that a - library can skip an epoch that it already committed. A micro-batch streaming query creates a new - stream writer for each micro-batch, and Spark closes the previous one when it creates the next - one. Because Spark is not told when a streaming query stops, the last stream writer of a query is - closed once it is no longer referenced, or when the query is restarted. + library can skip an epoch that it already committed. A micro-batch streaming query opens a new + writer for each micro-batch, and Spark closes the previous one when it opens the next one. + Because Spark is not told when a streaming query stops, the last writer of a query is closed once + it is no longer referenced, or when the query is restarted. The library defines the format of the partitions, of the state of its readers and writers, and of -the commit messages, which Spark sends to the executors and back as byte arrays. The library -returns them, and the offsets, by passing them to a `SparkDsSink` callback, which copies them. A -function reports an error by returning a non-zero value, with the message returned by -`last_error`, and Spark reports it as a `NATIVE_DATA_SOURCE_ERROR`. A C++ exception or a Rust panic -must not cross the interface. Each handle that a `create` function returns is released exactly -once, by the `close` function of its kind. The close functions cannot fail, and they can run after -the query, so do the work that can fail, such as flushing data, in `commit`. Strings are UTF-8, and -the JSON of the predicates is ASCII, with the other characters escaped. - -A library only sets the functions of the operations it supports, and leaves the others `NULL`. A -data source must not declare a capability whose functions are `NULL`: - -| Capability | Required functions | Optional functions | -|-----------|--------------------|--------------------| -| All | `create`, `close` | `last_error`, `schema` | -| `SPARK_DS_BATCH_READ` | `create_reader`, `partitions`, `read`, `close_reader` | `push_predicates`, `push_limit`, `prune_columns`, `serialize_reader` | -| `SPARK_DS_MICRO_BATCH_READ` | `create_stream_reader`, `initial_offset`, `latest_offset`, `stream_partitions`, `read`, `close_stream_reader` | `serialize_stream_reader`, `commit_offset` | -| `SPARK_DS_BATCH_WRITE`, `SPARK_DS_STREAMING_WRITE` | `create_writer`, `commit`, `abort`, `close_writer`, `create_data_writer`, `write`, `commit_data_writer`, `abort_data_writer`, `close_data_writer` | `serialize_writer` | - -`SPARK_DS_TRUNCATE` declares that a writer supports `SPARK_DS_WRITE_TRUNCATE`: the save mode +the commit messages, which Spark sends to the executors and back as bytes. A library replies to a +function by passing the encoded reply to a `SparkDsSink` callback, which copies it, or not at all +for an empty reply, which is valid for every request: for example, a reader that does not support +pushdown replies to `PushPredicates` with nothing, and Spark pushes nothing down. A function reports +an error by returning a non-zero value, with the message returned by `spark_ds_last_error`, and +Spark reports it as a `NATIVE_DATA_SOURCE_ERROR`. A C++ exception or a Rust panic must not cross the +interface. Each object is released exactly once, with `spark_ds_close`, which cannot fail and can +run after the query, so do the work that can fail, such as flushing data, in a commit. Strings are +UTF-8, and the JSON of the predicates is ASCII, with the other characters escaped. + +`CAPABILITY_TRUNCATE` declares that a writer supports `WRITE_MODE_TRUNCATE`: the save mode `overwrite` of a batch write, and the output mode `complete` of a streaming write. The comments of [`spark_datasource.h`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h) -specify each function, including the ownership of the Arrow structs and the JSON format of the -pushed predicates. `push_predicates` tells Spark how the reader uses each predicate: not at all, -only to skip data (`SPARK_DS_PREDICATE_PARTIAL`), or to only return the rows for which it is true -(`SPARK_DS_PREDICATE_COMPLETE`). Spark shows the predicates that the reader uses as pushed in the -plan of the query, and evaluates all the pushed predicates again after reading either way, so a -library can use them to skip data, for example files or row groups whose statistics show that no -row matches, but it must never skip a row for which the predicates are true under the semantics of -Spark SQL: NaN values are equal to each other and greater than any other value, `-0.0` is equal to -`0.0`, comparisons with null are null, and strings compare by their UTF-8 bytes. The comparisons of -IEEE 754 and the total orders of Arrow libraries, such as the comparison kernels of arrow-rs, -differ for some of these values. The jar of Spark SQL contains test cases for a library, -`org/apache/spark/sql/datasource/predicate-cases.json`: predicates as Spark pushes them, with the -rows of a table for which they are true: +and `spark_datasource.proto` specify each function and message, including the ownership of the +Arrow structs and the JSON format of the pushed predicates. `PushPredicatesReply` tells Spark how +the reader uses each predicate: not at all, only to skip data (`PREDICATE_USE_PARTIAL`), or to +only return the rows for which it is true (`PREDICATE_USE_COMPLETE`). Spark shows the predicates +that the reader uses as pushed in the plan of the query, and evaluates all the pushed predicates +again after reading either way, so a library can use them to skip data, for example files or row +groups whose statistics show that no row matches, but it must never skip a row for which the +predicates are true under the semantics of Spark SQL: NaN values are equal to each other and +greater than any other value, `-0.0` is equal to `0.0`, comparisons with null are null, and strings +compare by their UTF-8 bytes. The comparisons of IEEE 754 and the total orders of Arrow libraries, +such as the comparison kernels of arrow-rs, differ for some of these values. The jar of Spark SQL +contains test cases for a library, `org/apache/spark/sql/datasource/predicate-cases.json`: +predicates as Spark pushes them, with the rows of a table for which they are true: ```bash unzip -p "$SPARK_HOME"/jars/spark-sql_*.jar org/apache/spark/sql/datasource/predicate-cases.json @@ -165,39 +168,40 @@ has at least as many values as its parent refers to, but it cannot check that th arrays hold those values, because the C data interface does not carry the sizes of buffers: a library that exports a buffer that is too short makes Spark read past its end. -`read` and `create_data_writer` get a `SparkDsTaskContext`, with the IDs of the task, and its -partition and attempt, and an `is_interrupted` function that returns whether the task was killed, -for example because its query was canceled. Spark cannot interrupt native code: when a task is -killed, Spark stops reading only once the `get_next` function of the stream returns. So a read or -a write that runs for long, or blocks, for example on network I/O, should check `is_interrupted` -regularly and use timeouts. The task reaper, `spark.task.reaper.enabled` and -`spark.task.reaper.killTimeout`, kills an executor whose killed task does not end. +Partition readers and data writers get the `TaskContext` of their task, with its IDs, and the +`is_interrupted` function of `SparkDsHost` returns whether a task was killed, for example because +its query was canceled. Spark cannot interrupt native code: when a task is killed, Spark stops +reading only once the `get_next` function of the stream returns. So a read or a write that runs +for long, or blocks, for example on network I/O, should check `is_interrupted` regularly and use +timeouts. The task reaper, `spark.task.reaper.enabled` and `spark.task.reaper.killTimeout`, kills +an executor whose killed task does not end. `SparkDsHost` also has a `log` function, which writes +to the log of the driver or of the executor. ### Compatibility -The function table, `SparkDsLibrary`, and the structs that Spark passes to the library, -`SparkDsHost`, `SparkDsCreateInfo`, `SparkDsTaskContext` and `SparkDsWriteInfo`, start with their -size, so that they can grow. The version of the interface, 1, is `SPARK_DS_ABI_VERSION`, which the -table and the manifest of a package declare. Within a version: - -- fields are only appended to these structs, and the existing fields and the semantics of the - existing functions never change; -- Spark only reads the functions that fit in the size of the table of a library, so a library built - against an earlier header keeps working with later releases of Spark, with the functions that it - does not know as `NULL`. A library must only read the fields that fit in the size of a struct - that Spark passes, so that it also runs on earlier releases of Spark that support its version; -- new values may appear where the header says so: for example, later releases may push down - predicates with other names, which a library must ignore if it does not know them. Spark only - uses a new capability or write mode with a data source that declares it. - -An incompatible change comes with a new version: `spark_ds_init` gets the newest version that Spark -supports, in `SparkDsHost`, and returns a table of a version that the library implements. Spark -supports version 1 at least in all the releases of Spark 4 from 4.4.0. +The six functions of `spark_datasource.h` never change, nor do the existing fields of +`SparkDsHost`, which starts with its size and only gets fields appended: a library must only read +the fields that fit in the size of the struct that Spark passes. Everything else grows in the +messages, as Protocol Buffers messages do: + +- a later release of Spark may add fields, messages and values of enums, but never changes or + reuses the number of a field. Both sides ignore the fields that they do not know, so a library + built against an earlier version of the messages keeps working with later releases of Spark, and + a library built against a later version works with earlier releases, without the features that + they do not know; +- Spark only sends a new kind of request to an object whose `OpenReply` declares that it supports + it, so a new request never reaches an older library; +- new values may appear where the messages say so: for example, later releases may push down + predicates with other names, which a library must ignore if it does not know them. + +The version of the interface, 1, is `SPARK_DS_ABI_VERSION`, which `spark_ds_abi_version` returns +and the manifest of a package declares. Spark supports version 1 at least in all the releases of +Spark 4 from 4.4.0. ### Limitations Version 1 of the interface does not cover these features of Data Source V2 yet. They can be added -later as functions or fields appended to the structs of the interface: +later as fields and messages of `spark_datasource.proto`: - statistics: Spark estimates the size of a native scan with `spark.sql.defaultSizeInBytes`, so it does not plan a native scan as the broadcast side of a join, unless the query has a broadcast @@ -243,7 +247,7 @@ like other native libraries is found automatically, without any configuration. T named after the data source, in lower case, with the prefix `spark_datasource_`: the library of `rust_range` is `libspark_datasource_rust_range.so` on Linux, `libspark_datasource_rust_range.dylib` on macOS and `spark_datasource_rust_range.dll` on Windows. -Spark passes the name of the data source to the `create` function of the library. A library that +Spark passes the name of the data source in the requests of the library. A library that implements several data sources is installed under each of their names, for example with symbolic links. @@ -302,7 +306,7 @@ is a zip file with the extension `.sparkpkg`, which contains a manifest named - `abiVersion` is the version of `spark_datasource.h` that the library implements, currently 1. - `dataSources` lists the names of the data sources that the library implements. The names are - case-insensitive, and Spark passes the name to `create`. + case-insensitive, and Spark passes the name in the requests of the library. - `libraries` maps each platform, `-` with an `os` of `linux`, `osx` or `windows` and an `arch` of `x86_64` or `aarch64`, to the path of the library in the package. @@ -328,14 +332,13 @@ installed library provides it. ## Example in Rust -This data source reads the numbers `[0, end)` as a column `id`, in two partitions. It only depends -on the `arrow-array` and `arrow-schema` crates of [arrow-rs](https://github.com/apache/arrow-rs), -which export schemas and streams with `FFI_ArrowSchema` and `FFI_ArrowArrayStream`. It declares -`SparkDsLibrary` as a `#[repr(C)]` struct up to `read`, the last function that it implements, since -Spark only reads the functions that fit in the size of the table, and leaves the functions that it -does not implement `None`. The fields must follow the order of `spark_datasource.h`, which -[bindgen](https://github.com/rust-lang/rust-bindgen) can also generate from the header. -`Cargo.toml`, where the name of the library is the name that Spark finds when it is installed: +This data source reads the numbers `[0, end)` as a column `id`, in two partitions. It depends on +the `arrow-array` and `arrow-schema` crates of [arrow-rs](https://github.com/apache/arrow-rs), +which export streams with `FFI_ArrowArrayStream`, and on [prost](https://github.com/tokio-rs/prost), +whose `prost-build` generates the code of the messages from `spark_datasource.proto` when the +library is built, with `protoc`. It only opens data sources, batch readers and partition readers, +and replies to the requests that it does not support with empty replies. `Cargo.toml`, where the +name of the library is the name that Spark finds when it is installed: ```toml [package] @@ -350,6 +353,21 @@ crate-type = ["cdylib"] [dependencies] arrow-array = { version = "58", features = ["ffi"] } arrow-schema = { version = "58", features = ["ffi"] } +prost = "0.14" + +[build-dependencies] +prost-build = "0.14" +``` + +`build.rs`, which generates the messages from the `include` directory of a Spark distribution: + +```rust +// Generates the code of the messages from spark_datasource.proto, in the include directory of a +// Spark distribution. prost-build runs protoc, which must be installed. +fn main() -> std::io::Result<()> { + let include = format!("{}/include", std::env::var("SPARK_HOME").expect("SPARK_HOME")); + prost_build::compile_protos(&[format!("{include}/spark_datasource.proto")], &[include]) +} ``` `src/lib.rs`: @@ -360,119 +378,32 @@ arrow-schema = { version = "58", features = ["ffi"] } //! column `id`, in two partitions. The option `end` defaults to 10. use std::cell::RefCell; -use std::ffi::{c_char, c_int, c_void, CStr, CString}; +use std::ffi::{c_char, c_int, c_void, CString}; use std::panic::{catch_unwind, AssertUnwindSafe}; use std::sync::Arc; use arrow_array::ffi_stream::FFI_ArrowArrayStream; use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator}; -use arrow_schema::ffi::FFI_ArrowSchema; use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use prost::Message; -type Result = std::result::Result>; - -/// `SparkDsSink`: receives a byte array returned to Spark, and copies it. -type Sink = unsafe extern "C" fn(sink: *mut c_void, data: *const u8, length: i64); - -/// A function of `SparkDsLibrary` that this data source does not implement, left NULL. -type Unused = Option; - -/// `SparkDsCreateInfo` of `spark_datasource.h`: the name and the options of a data source. -#[repr(C)] -pub struct SparkDsCreateInfo { - size: i64, - name: *const c_char, - option_keys: *const *const c_char, - option_values: *const *const c_char, - num_options: i64, +/// The messages of `spark_datasource.proto`. +mod proto { + include!(concat!(env!("OUT_DIR"), "/spark.datasource.rs")); } -/// `SparkDsLibrary` of `spark_datasource.h`, up to `read`, the last function that this data -/// source implements: Spark only reads the functions that fit in `size`. -#[repr(C)] -pub struct SparkDsLibrary { - size: i64, - abi_version: i32, - last_error: unsafe extern "C" fn() -> *const c_char, - create: unsafe extern "C" fn( - info: *const SparkDsCreateInfo, - capabilities: *mut u64, - source: *mut *mut c_void, - ) -> c_int, - schema: unsafe extern "C" fn(source: *mut c_void, out: *mut FFI_ArrowSchema) -> c_int, - create_reader: unsafe extern "C" fn( - source: *mut c_void, - schema: *mut FFI_ArrowSchema, - reader: *mut *mut c_void, - ) -> c_int, - create_stream_reader: Unused, - create_writer: Unused, - close: unsafe extern "C" fn(source: *mut c_void), - push_predicates: Unused, - push_limit: Unused, - prune_columns: Unused, - partitions: unsafe extern "C" fn(reader: *mut c_void, add: Sink, sink: *mut c_void) -> c_int, - serialize_reader: Unused, - close_reader: unsafe extern "C" fn(reader: *mut c_void), - initial_offset: Unused, - latest_offset: Unused, - stream_partitions: Unused, - serialize_stream_reader: Unused, - commit_offset: Unused, - close_stream_reader: Unused, - read: unsafe extern "C" fn( - task: *const c_void, - reader_state: *const u8, - reader_state_length: i64, - partition: *const u8, - partition_length: i64, - out: *mut FFI_ArrowArrayStream, - ) -> c_int, -} +use proto::{call_reply, call_request, data_type, open_request, stream_request}; -const SPARK_DS_BATCH_READ: u64 = 1; - -static LIBRARY: SparkDsLibrary = SparkDsLibrary { - size: std::mem::size_of::() as i64, - abi_version: 1, - last_error, - create, - schema, - create_reader, - create_stream_reader: None, - create_writer: None, - close, - push_predicates: None, - push_limit: None, - prune_columns: None, - partitions, - serialize_reader: None, - close_reader, - initial_offset: None, - latest_offset: None, - stream_partitions: None, - serialize_stream_reader: None, - commit_offset: None, - close_stream_reader: None, - read, -}; - -/// The function that Spark calls when it loads the library. -#[no_mangle] -pub unsafe extern "C" fn spark_ds_init( - _host: *const c_void, - library: *mut *const SparkDsLibrary, -) -> c_int { - *library = &LIBRARY; - 0 -} +type Result = std::result::Result>; -struct RangeSource { - end: i64, -} +/// `SparkDsSink`: receives the reply of a function, and copies it. +type Sink = unsafe extern "C" fn(sink: *mut c_void, data: *const u8, length: i64); -struct RangeReader { - end: i64, +/// What a handle refers to. +enum Object { + DataSource, + Reader { end: i64 }, + Partition { start: i64, end: i64 }, } thread_local! { @@ -483,6 +414,18 @@ fn arrow_schema() -> SchemaRef { Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])) } +/// The schema of the data, as the Spark SQL types of the messages. +fn spark_schema() -> proto::StructType { + let long = data_type::Kind::Long(data_type::Long {}); + proto::StructType { + fields: vec![proto::StructField { + name: "id".to_string(), + data_type: Some(proto::DataType { kind: Some(long) }), + nullable: false, + }], + } +} + /// Runs `f`, and reports its error or panic to Spark: neither may cross into the JVM. fn call(f: impl FnOnce() -> Result<()>) -> c_int { let message = match catch_unwind(AssertUnwindSafe(f)) { @@ -494,86 +437,117 @@ fn call(f: impl FnOnce() -> Result<()>) -> c_int { 1 } -unsafe extern "C" fn last_error() -> *const c_char { - LAST_ERROR.with(|error| error.borrow().as_ptr()) +unsafe fn decode(request: *const u8, length: i64) -> Result { + Ok(M::decode(std::slice::from_raw_parts(request, length as usize))?) } -unsafe extern "C" fn create( - info: *const SparkDsCreateInfo, - capabilities: *mut u64, - source: *mut *mut c_void, -) -> c_int { - call(|| { - let info = &*info; - let mut range = RangeSource { end: 10 }; - for i in 0..info.num_options as usize { - if CStr::from_ptr(*info.option_keys.add(i)).to_str()? == "end" { - range.end = CStr::from_ptr(*info.option_values.add(i)).to_str()?.parse()?; - } - } - *capabilities = SPARK_DS_BATCH_READ; - *source = Box::into_raw(Box::new(range)) as *mut c_void; - Ok(()) - }) +unsafe fn reply(reply: Sink, sink: *mut c_void, message: &impl Message) { + let bytes = message.encode_to_vec(); + reply(sink, bytes.as_ptr(), bytes.len() as i64); } -unsafe extern "C" fn schema(_source: *mut c_void, out: *mut FFI_ArrowSchema) -> c_int { - call(|| { - std::ptr::write(out, FFI_ArrowSchema::try_from(arrow_schema().as_ref())?); - Ok(()) - }) +#[no_mangle] +pub extern "C" fn spark_ds_abi_version() -> i32 { + 1 } -unsafe extern "C" fn create_reader( - source: *mut c_void, - _schema: *mut FFI_ArrowSchema, - reader: *mut *mut c_void, +#[no_mangle] +pub extern "C" fn spark_ds_last_error() -> *const c_char { + LAST_ERROR.with(|error| error.borrow().as_ptr()) +} + +#[no_mangle] +pub unsafe extern "C" fn spark_ds_open( + _host: *const c_void, + request: *const u8, + length: i64, + sink_function: Sink, + sink: *mut c_void, + handle: *mut *mut c_void, ) -> c_int { call(|| { - let end = (*(source as *const RangeSource)).end; - *reader = Box::into_raw(Box::new(RangeReader { end })) as *mut c_void; + let object = match decode::(request, length)?.object { + Some(open_request::Object::DataSource(_)) => { + let capabilities = vec![proto::Capability::BatchRead as i32]; + reply(sink_function, sink, &proto::OpenReply { capabilities }); + Object::DataSource + } + Some(open_request::Object::BatchReader(reader)) => { + let options = reader.data_source.unwrap_or_default().options; + let end = options.get("end").map(|end| end.parse()).transpose()?; + Object::Reader { end: end.unwrap_or(10) } + } + Some(open_request::Object::PartitionReader(reader)) => { + let bytes = reader.partition; + let start = i64::from_le_bytes(bytes[0..8].try_into()?); + let end = i64::from_le_bytes(bytes[8..16].try_into()?); + Object::Partition { start, end } + } + _ => return Err("this data source only supports batch reads".into()), + }; + *handle = Box::into_raw(Box::new(object)) as *mut c_void; Ok(()) }) } -unsafe extern "C" fn close(source: *mut c_void) { - drop(Box::from_raw(source as *mut RangeSource)); -} - -unsafe extern "C" fn partitions(reader: *mut c_void, add: Sink, sink: *mut c_void) -> c_int { +#[no_mangle] +pub unsafe extern "C" fn spark_ds_call( + handle: *mut c_void, + request: *const u8, + length: i64, + sink_function: Sink, + sink: *mut c_void, +) -> c_int { call(|| { - let end = (*(reader as *const RangeReader)).end; - for (start, end) in [(0, end / 2), (end / 2, end)] { - let bytes = [start.to_le_bytes(), end.to_le_bytes()].concat(); - add(sink, bytes.as_ptr(), bytes.len() as i64); - } + let object = &*(handle as *const Object); + let result = match (decode::(request, length)?.call, object) { + (Some(call_request::Call::InferSchema(_)), _) => { + let schema = Some(spark_schema()); + call_reply::Reply::InferSchema(proto::InferSchemaReply { schema }) + } + (Some(call_request::Call::PlanPartitions(_)), Object::Reader { end }) => { + let partitions = [(0, end / 2), (end / 2, *end)] + .iter() + .map(|(start, end)| [start.to_le_bytes(), end.to_le_bytes()].concat()) + .collect(); + call_reply::Reply::PlanPartitions(proto::PlanPartitionsReply { partitions }) + } + // The other requests, such as the pushdown, have empty replies. + _ => return Ok(()), + }; + reply(sink_function, sink, &proto::CallReply { reply: Some(result) }); Ok(()) }) } -unsafe extern "C" fn close_reader(reader: *mut c_void) { - drop(Box::from_raw(reader as *mut RangeReader)); -} - -unsafe extern "C" fn read( - _task: *const c_void, - _reader_state: *const u8, - _reader_state_length: i64, - partition: *const u8, - partition_length: i64, - out: *mut FFI_ArrowArrayStream, +#[no_mangle] +pub unsafe extern "C" fn spark_ds_stream( + handle: *mut c_void, + request: *const u8, + length: i64, + _input: *mut FFI_ArrowArrayStream, + output: *mut FFI_ArrowArrayStream, ) -> c_int { call(|| { - let bytes = std::slice::from_raw_parts(partition, partition_length as usize); - let start = i64::from_le_bytes(bytes[0..8].try_into()?); - let end = i64::from_le_bytes(bytes[8..16].try_into()?); - let ids = Int64Array::from_iter_values(start..end); + let Object::Partition { start, end } = &*(handle as *const Object) else { + return Err("not a partition reader".into()); + }; + let stream = decode::(request, length)?.stream; + let Some(stream_request::Stream::Read(_)) = stream else { + return Err("this data source only supports reads".into()); + }; + let ids = Int64Array::from_iter_values(*start..*end); let batch = RecordBatch::try_new(arrow_schema(), vec![Arc::new(ids)])?; let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema()); - std::ptr::write(out, FFI_ArrowArrayStream::new(Box::new(batches))); + std::ptr::write(output, FFI_ArrowArrayStream::new(Box::new(batches))); Ok(()) }) } + +#[no_mangle] +pub unsafe extern "C" fn spark_ds_close(handle: *mut c_void) { + drop(Box::from_raw(handle as *mut Object)); +} ``` {% endraw %} @@ -608,37 +582,44 @@ EOF ## Example in C++ The same data source in C++, with [Arrow C++](https://arrow.apache.org/docs/cpp/), which exports -schemas and streams with `arrow::ExportSchema` and `arrow::ExportRecordBatchReader`. It includes -`spark_datasource.h`, here from the `include` directory of a Spark distribution, and fills -`SparkDsLibrary` with designated initializers, which leave the other functions `NULL`. +streams with `arrow::ExportRecordBatchReader`, and with the code that `protoc` generates from +`spark_datasource.proto`. It includes `spark_datasource.h` and the generated messages, here from +the `include` directory of a Spark distribution. `cpp_range.cc`: {% raw %} ```cpp -// A native data source for Apache Spark, written in C++ with Arrow C++. It reads the numbers -// [0, end) as a column `id`, in two partitions. The option `end` defaults to 10. +// A native data source for Apache Spark, written in C++ with Arrow C++ and Protocol Buffers. It +// reads the numbers [0, end) as a column `id`, in two partitions. The option `end` defaults to 10. #include #include #include #include #include +#include #include #include #include "spark_datasource.h" +#include "spark_datasource.pb.h" namespace { -struct RangeSource { - int64_t end = 10; -}; +namespace sds = spark::datasource; -struct RangeReader { +// What a handle refers to. +struct DataSource {}; +struct Reader { int64_t end; }; +struct Partition { + int64_t start; + int64_t end; +}; +using Object = std::variant; thread_local std::string last_error; @@ -665,50 +646,86 @@ void Check(const arrow::Status& status) { if (!status.ok()) throw std::runtime_error(status.ToString()); } -const char* LastError() { return last_error.c_str(); } - -int Create(const SparkDsCreateInfo* info, uint64_t* capabilities, void** source) { - return Call([&] { - auto range = std::make_unique(); - for (int64_t i = 0; i < info->num_options; i++) { - if (std::string(info->option_keys[i]) == "end") { - range->end = std::stoll(info->option_values[i]); - } - } - *capabilities = SPARK_DS_BATCH_READ; - *source = range.release(); - }); +template +M Decode(const uint8_t* request, int64_t length) { + M message; + if (!message.ParseFromArray(request, static_cast(length))) { + throw std::runtime_error("invalid request"); + } + return message; } -int GetSchema(void*, struct ArrowSchema* out) { - return Call([&] { Check(arrow::ExportSchema(*Schema(), out)); }); +void Reply(SparkDsSink reply, void* sink, const google::protobuf::MessageLite& message) { + std::string bytes = message.SerializeAsString(); + reply(sink, reinterpret_cast(bytes.data()), static_cast(bytes.size())); } -int CreateReader(void* source, struct ArrowSchema*, void** reader) { - return Call([&] { *reader = new RangeReader{static_cast(source)->end}; }); -} +} // namespace + +extern "C" { -void Close(void* source) { delete static_cast(source); } +int32_t spark_ds_abi_version(void) { return SPARK_DS_ABI_VERSION; } -int Partitions(void* reader, SparkDsSink add_partition, void* sink) { +const char* spark_ds_last_error(void) { return last_error.c_str(); } + +int spark_ds_open(const SparkDsHost*, const uint8_t* request, int64_t length, SparkDsSink reply, + void* sink, void** handle) { return Call([&] { - int64_t end = static_cast(reader)->end; - int64_t bounds[2][2] = {{0, end / 2}, {end / 2, end}}; - for (const auto& partition : bounds) { - add_partition(sink, reinterpret_cast(partition), sizeof(partition)); + auto open = Decode(request, length); + std::unique_ptr object; + if (open.has_data_source()) { + sds::OpenReply capabilities; + capabilities.add_capabilities(sds::CAPABILITY_BATCH_READ); + Reply(reply, sink, capabilities); + object = std::make_unique(DataSource{}); + } else if (open.has_batch_reader()) { + const auto& options = open.batch_reader().data_source().options(); + auto end = options.find("end"); + object = std::make_unique( + Reader{end == options.end() ? 10 : std::stoll(end->second)}); + } else if (open.has_partition_reader()) { + int64_t bounds[2]; + std::memcpy(bounds, open.partition_reader().partition().data(), sizeof(bounds)); + object = std::make_unique(Partition{bounds[0], bounds[1]}); + } else { + throw std::runtime_error("this data source only supports batch reads"); } + *handle = object.release(); }); } -void CloseReader(void* reader) { delete static_cast(reader); } +int spark_ds_call(void* handle, const uint8_t* request, int64_t length, SparkDsSink reply, + void* sink) { + return Call([&] { + auto call = Decode(request, length); + sds::CallReply result; + if (call.has_infer_schema()) { + auto* field = result.mutable_infer_schema()->mutable_schema()->add_fields(); + field->set_name("id"); + field->mutable_data_type()->mutable_long_(); + field->set_nullable(false); + } else if (call.has_plan_partitions()) { + int64_t end = std::get(*static_cast(handle)).end; + int64_t bounds[2][2] = {{0, end / 2}, {end / 2, end}}; + for (const auto& partition : bounds) { + result.mutable_plan_partitions()->add_partitions(partition, sizeof(partition)); + } + } else { + return; // The other requests, such as the pushdown, have empty replies. + } + Reply(reply, sink, result); + }); +} -int Read(const SparkDsTaskContext*, const uint8_t*, int64_t, const uint8_t* partition, int64_t, - struct ArrowArrayStream* out) { +int spark_ds_stream(void* handle, const uint8_t* request, int64_t length, struct ArrowArrayStream*, + struct ArrowArrayStream* out) { return Call([&] { - int64_t bounds[2]; - std::memcpy(bounds, partition, sizeof(bounds)); + if (!Decode(request, length).has_read()) { + throw std::runtime_error("this data source only supports reads"); + } + const auto& partition = std::get(*static_cast(handle)); arrow::Int64Builder ids; - for (int64_t id = bounds[0]; id < bounds[1]; id++) Check(ids.Append(id)); + for (int64_t id = partition.start; id < partition.end; id++) Check(ids.Append(id)); std::shared_ptr array; Check(ids.Finish(&array)); auto batch = arrow::RecordBatch::Make(Schema(), array->length(), {array}); @@ -717,40 +734,30 @@ int Read(const SparkDsTaskContext*, const uint8_t*, int64_t, const uint8_t* part }); } -// The functions of the data source. The ones that it does not implement are NULL. -const SparkDsLibrary kLibrary = { - .size = sizeof(SparkDsLibrary), - .abi_version = SPARK_DS_ABI_VERSION, - .last_error = LastError, - .create = Create, - .schema = GetSchema, - .create_reader = CreateReader, - .close = Close, - .partitions = Partitions, - .close_reader = CloseReader, - .read = Read, -}; - -} // namespace +void spark_ds_close(void* handle) { delete static_cast(handle); } -extern "C" SPARK_DS_EXPORT int spark_ds_init(const SparkDsHost*, const SparkDsLibrary** library) { - *library = &kLibrary; - return 0; -} +} // extern "C" ``` {% endraw %} -Build the library, here on macOS with Arrow C++ 24, which requires C++20, and install it, or -package it as above: +Generate the messages and build the library, here on macOS with Arrow C++ 24, which requires +C++20, and install it, or package it as above: ```bash -c++ -std=c++20 -O2 -shared -fPIC -I"$SPARK_HOME/include" -I"$ARROW_HOME/include" \ - -L"$ARROW_HOME/lib" -larrow -o libspark_datasource_cpp_range.dylib cpp_range.cc +protoc --cpp_out=. -I"$SPARK_HOME/include" "$SPARK_HOME/include/spark_datasource.proto" +c++ -std=c++20 -O2 -shared -fPIC -I. -I"$SPARK_HOME/include" -I"$ARROW_HOME/include" \ + cpp_range.cc spark_datasource.pb.cc $(pkg-config --cflags --libs protobuf-lite) \ + -L"$ARROW_HOME/lib" -larrow -o libspark_datasource_cpp_range.dylib cp libspark_datasource_cpp_range.dylib /usr/local/lib/ ``` -A library loaded by Spark must find its dependencies, such as `libarrow`, on every node: link them -statically, or install them on the nodes. +A library loaded by Spark must find its dependencies, such as `libarrow` and `libprotobuf-lite`, on +every node: link them statically, or install them on the nodes. Several native data source +libraries, or several copies of one, such as an installed library and the library of a package, can +share a process. The messages use the lite runtime of Protocol Buffers, which, unlike the full one, +does not register them in a pool of the process, where a second copy would abort the process. But +libraries built against different versions of Protocol Buffers conflict if they link it +dynamically, so link it statically, with hidden symbols. ## Configuration diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java index 66c1a46b7beea..c9edcf0b0b6b2 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -32,28 +32,34 @@ import java.util.concurrent.atomic.AtomicLong; import java.util.function.BooleanSupplier; +import com.google.protobuf.InvalidProtocolBufferException; + import org.apache.spark.internal.LogKeys; import org.apache.spark.internal.MDC; import org.apache.spark.internal.SparkLogger; import org.apache.spark.internal.SparkLoggerFactory; +import org.apache.spark.sql.execution.datasources.v2.ffi.proto.CallReply; +import org.apache.spark.sql.execution.datasources.v2.ffi.proto.CallRequest; +import org.apache.spark.sql.execution.datasources.v2.ffi.proto.OpenReply; +import org.apache.spark.sql.execution.datasources.v2.ffi.proto.OpenRequest; +import org.apache.spark.sql.execution.datasources.v2.ffi.proto.StreamRequest; import org.apache.spark.unsafe.Platform; /** * A native data source library, which implements the C interface in * {@code org/apache/spark/sql/datasource/spark_datasource.h}, called through the Foreign Function - * and Memory (FFM) API of Java 22 and later. + * and Memory (FFM) API of Java 22 and later, with the messages of + * {@code org/apache/spark/sql/datasource/spark_datasource.proto}. *

    * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached - * through method handles, looked up once when the first library is loaded. The library exports - * {@code spark_ds_init}, which returns its table of functions, {@code SparkDsLibrary}. Each - * function of the table is bound to a method handle that takes and returns primitive values only, - * with pointers as {@code long} addresses, so the calls themselves use no reflection. The memory - * that Spark passes to the library is allocated with {@link Platform}, at the offsets of the - * structs of the header, which {@link #layout()} lists for the tests. + * through method handles, looked up once when the first library is loaded. Each of the six + * functions that the library exports is bound to a method handle that takes and returns primitive + * values only, with pointers as {@code long} addresses, so the calls themselves use no + * reflection. The memory that Spark passes to the library is allocated with {@link Platform}, at + * the offsets of the structs of the header, which {@link #layout()} lists for the tests. *

    - * Calling a function that the library leaves NULL in its table throws - * {@link MissingFunctionException}. A function that fails throws a {@link RuntimeException} with - * the message of {@code last_error}. + * A function that fails throws a {@link RuntimeException} with the message of + * {@code spark_ds_last_error}. */ final class NativeLibrary { @@ -62,6 +68,9 @@ final class NativeLibrary { /** The version of the interface that Spark supports, {@code SPARK_DS_ABI_VERSION}. */ static final int ABI_VERSION = 1; + /** The newest version of the messages that Spark supports. */ + static final int PROTOCOL_VERSION = 1; + /** The first version of Java with the final FFM API. */ static final int MINIMUM_JAVA_VERSION = 22; @@ -70,169 +79,77 @@ static boolean isSupported() { return Runtime.version().feature() >= MINIMUM_JAVA_VERSION; } - // The capabilities of a data source, SPARK_DS_BATCH_READ and the others. - static final long BATCH_READ = 1; - static final long MICRO_BATCH_READ = 2; - static final long BATCH_WRITE = 4; - static final long STREAMING_WRITE = 8; - static final long TRUNCATE = 16; - - // The modes of a write, SPARK_DS_WRITE_APPEND and SPARK_DS_WRITE_TRUNCATE. - static final int WRITE_APPEND = 0; - static final int WRITE_TRUNCATE = 1; - - /** A predicate that the reader does not use, SPARK_DS_PREDICATE_UNUSED. */ - static final int PREDICATE_UNUSED = 0; - // The C types of the parameters and the results of the functions. private enum C { VOID, BOOL, I32, I64, PTR } - /** The functions of the table {@code SparkDsLibrary}, in their order in the table. */ - enum Function { - LAST_ERROR("last_error", C.PTR), - CREATE("create", C.I32, C.PTR, C.PTR, C.PTR), - SCHEMA("schema", C.I32, C.PTR, C.PTR), - CREATE_READER("create_reader", C.I32, C.PTR, C.PTR, C.PTR), - CREATE_STREAM_READER("create_stream_reader", C.I32, C.PTR, C.PTR, C.PTR), - CREATE_WRITER("create_writer", C.I32, C.PTR, C.PTR, C.PTR, C.PTR), - CLOSE("close", C.VOID, C.PTR), - PUSH_PREDICATES("push_predicates", C.I32, C.PTR, C.PTR, C.I64, C.PTR), - PUSH_LIMIT("push_limit", C.I32, C.PTR, C.I32, C.PTR), - PRUNE_COLUMNS("prune_columns", C.I32, C.PTR, C.PTR, C.I64, C.PTR), - PARTITIONS("partitions", C.I32, C.PTR, C.PTR, C.PTR), - SERIALIZE_READER("serialize_reader", C.I32, C.PTR, C.PTR, C.PTR), - CLOSE_READER("close_reader", C.VOID, C.PTR), - INITIAL_OFFSET("initial_offset", C.I32, C.PTR, C.PTR, C.PTR), - LATEST_OFFSET("latest_offset", C.I32, C.PTR, C.PTR, C.PTR), - STREAM_PARTITIONS("stream_partitions", C.I32, C.PTR, C.PTR, C.PTR, C.PTR, C.PTR), - SERIALIZE_STREAM_READER("serialize_stream_reader", C.I32, C.PTR, C.PTR, C.PTR), - COMMIT_OFFSET("commit_offset", C.I32, C.PTR, C.PTR), - CLOSE_STREAM_READER("close_stream_reader", C.VOID, C.PTR), - READ("read", C.I32, C.PTR, C.PTR, C.I64, C.PTR, C.I64, C.PTR), - SERIALIZE_WRITER("serialize_writer", C.I32, C.PTR, C.PTR, C.PTR), - COMMIT("commit", C.I32, C.PTR, C.I64, C.PTR, C.PTR, C.I64), - ABORT("abort", C.I32, C.PTR, C.I64, C.PTR, C.PTR, C.I64), - CLOSE_WRITER("close_writer", C.VOID, C.PTR), - CREATE_DATA_WRITER("create_data_writer", C.I32, C.PTR, C.PTR, C.I64, C.I64, C.PTR), - WRITE("write", C.I32, C.PTR, C.PTR, C.PTR), - COMMIT_DATA_WRITER("commit_data_writer", C.I32, C.PTR, C.PTR, C.PTR), - ABORT_DATA_WRITER("abort_data_writer", C.I32, C.PTR), - CLOSE_DATA_WRITER("close_data_writer", C.VOID, C.PTR); - - /** The name of the field of the function in the table. */ - final String fieldName; - private final C result; - private final C[] parameters; - - Function(String fieldName, C result, C... parameters) { - this.fieldName = fieldName; - this.result = result; - this.parameters = parameters; - } - - /** The offset of the function in the table, after the size and the version of the table. */ - long offset() { - return TABLE_FUNCTIONS + 8L * ordinal(); - } - } - - // The layouts of the structs of the header, on the 64-bit platforms that Spark supports. Each - // struct starts with its size. - private static final long TABLE_ABI_VERSION = 8; - private static final long TABLE_FUNCTIONS = 16; - private static final long HOST_SIZE = 32; - private static final long HOST_ABI_VERSION = 8; - private static final long HOST_SPARK_VERSION = 16; - private static final long HOST_LOG = 24; - private static final long CREATE_INFO_SIZE = 40; - private static final long CREATE_INFO_NAME = 8; - private static final long CREATE_INFO_OPTION_KEYS = 16; - private static final long CREATE_INFO_OPTION_VALUES = 24; - private static final long CREATE_INFO_NUM_OPTIONS = 32; - private static final long TASK_SIZE = 40; - private static final long TASK_STAGE_ID = 8; - private static final long TASK_STAGE_ATTEMPT_NUMBER = 12; - private static final long TASK_PARTITION_ID = 16; - private static final long TASK_ATTEMPT_NUMBER = 20; - private static final long TASK_ATTEMPT_ID = 24; - private static final long TASK_IS_INTERRUPTED = 32; - private static final long WRITE_INFO_SIZE = 24; - private static final long WRITE_INFO_MODE = 8; - private static final long WRITE_INFO_QUERY_ID = 16; + // The layout of SparkDsHost on the 64-bit platforms that Spark supports. + private static final long HOST_SIZE = 24; + private static final long HOST_LOG = 8; + private static final long HOST_IS_INTERRUPTED = 16; /** - * The offsets of the fields of the structs that Spark reads or writes, by {@code struct.field}, - * and the sizes of the structs, by {@code struct}, for the tests to compare with the header. + * The offsets of the fields of the structs that Spark writes, by {@code struct.field}, and the + * sizes of the structs, by {@code struct}, for the tests to compare with the header. */ static Map layout() { Map layout = new LinkedHashMap<>(); - layout.put("SparkDsLibrary", TABLE_FUNCTIONS + 8L * Function.values().length); - layout.put("SparkDsLibrary.size", 0L); - layout.put("SparkDsLibrary.abi_version", TABLE_ABI_VERSION); - for (Function function : Function.values()) { - layout.put("SparkDsLibrary." + function.fieldName, function.offset()); - } layout.put("SparkDsHost", HOST_SIZE); layout.put("SparkDsHost.size", 0L); - layout.put("SparkDsHost.abi_version", HOST_ABI_VERSION); - layout.put("SparkDsHost.spark_version", HOST_SPARK_VERSION); layout.put("SparkDsHost.log", HOST_LOG); - layout.put("SparkDsCreateInfo", CREATE_INFO_SIZE); - layout.put("SparkDsCreateInfo.size", 0L); - layout.put("SparkDsCreateInfo.name", CREATE_INFO_NAME); - layout.put("SparkDsCreateInfo.option_keys", CREATE_INFO_OPTION_KEYS); - layout.put("SparkDsCreateInfo.option_values", CREATE_INFO_OPTION_VALUES); - layout.put("SparkDsCreateInfo.num_options", CREATE_INFO_NUM_OPTIONS); - layout.put("SparkDsTaskContext", TASK_SIZE); - layout.put("SparkDsTaskContext.size", 0L); - layout.put("SparkDsTaskContext.stage_id", TASK_STAGE_ID); - layout.put("SparkDsTaskContext.stage_attempt_number", TASK_STAGE_ATTEMPT_NUMBER); - layout.put("SparkDsTaskContext.partition_id", TASK_PARTITION_ID); - layout.put("SparkDsTaskContext.attempt_number", TASK_ATTEMPT_NUMBER); - layout.put("SparkDsTaskContext.task_attempt_id", TASK_ATTEMPT_ID); - layout.put("SparkDsTaskContext.is_interrupted", TASK_IS_INTERRUPTED); - layout.put("SparkDsWriteInfo", WRITE_INFO_SIZE); - layout.put("SparkDsWriteInfo.size", 0L); - layout.put("SparkDsWriteInfo.mode", WRITE_INFO_MODE); - layout.put("SparkDsWriteInfo.query_id", WRITE_INFO_QUERY_ID); + layout.put("SparkDsHost.is_interrupted", HOST_IS_INTERRUPTED); return layout; } + private final long host; private final int abiVersion; - // The functions of the table, by ordinal, null for the functions that the library leaves NULL. - private final MethodHandle[] functions = new MethodHandle[Function.values().length]; - - private NativeLibrary(long table) throws Throwable { - long size = Platform.getLong(null, table); - if (size < TABLE_FUNCTIONS) { + private final MethodHandle lastError; + private final MethodHandle open; + private final MethodHandle call; + private final MethodHandle stream; + private final MethodHandle close; + + private NativeLibrary(Object lookup, String path) throws Throwable { + abiVersion = (int) function(lookup, "spark_ds_abi_version", C.I32).invokeExact(); + // The functions of another version of the interface may differ: the caller rejects the + // library before any of them is called. + boolean supported = abiVersion == ABI_VERSION; + lastError = supported ? optionalFunction(lookup, "spark_ds_last_error", C.PTR) : null; + open = supported + ? function(lookup, "spark_ds_open", C.I32, C.PTR, C.PTR, C.I64, C.PTR, C.PTR, C.PTR) + : null; + call = supported + ? function(lookup, "spark_ds_call", C.I32, C.PTR, C.PTR, C.I64, C.PTR, C.PTR) + : null; + stream = supported + ? function(lookup, "spark_ds_stream", C.I32, C.PTR, C.PTR, C.I64, C.PTR, C.PTR) + : null; + close = supported ? function(lookup, "spark_ds_close", C.VOID, C.PTR) : null; + // The library may keep the host until the JVM exits. + host = supported ? Ffm.host(path) : 0; + } + + private static MethodHandle function(Object lookup, String name, C result, C... parameters) + throws Throwable { + MethodHandle function = optionalFunction(lookup, name, result, parameters); + if (function == null) { throw new LoadException( - "spark_ds_init returned a function table with a size of " + size + " bytes.", null, true); - } - abiVersion = Platform.getInt(null, table + TABLE_ABI_VERSION); - // The table of another version of the interface has another layout. Only the functions that - // fit in the size of the table are read: the table of a library built against an earlier - // header ends before the functions that were added later. - if (abiVersion == ABI_VERSION) { - for (Function function : Function.values()) { - if (function.offset() + 8 <= size) { - long address = Platform.getLong(null, table + function.offset()); - if (address != 0) { - functions[function.ordinal()] = Ffm.downcall( - Ffm.OF_ADDRESS.invoke(address), function.result, function.parameters); - } - } - } + "The library does not export the function " + name + ".", null, true); } + return function; + } + + private static MethodHandle optionalFunction( + Object lookup, String name, C result, C... parameters) throws Throwable { + Optional address = (Optional) Ffm.FIND.invoke(lookup, name); + return address.isEmpty() ? null : Ffm.downcall(address.get(), result, parameters); } /** - * Loads the library at the given path, and gets its table with {@code spark_ds_init}. The - * library stays loaded until the JVM exits. + * Loads the library at the given path. The library stays loaded until the JVM exits. * - * @param sparkVersion the version of Spark, which the library gets in {@code SparkDsHost} * @throws LoadException if the library cannot be loaded */ - static NativeLibrary load(String path, String sparkVersion) throws Throwable { + static NativeLibrary load(String path) throws Throwable { if (!isSupported()) { throw new UnsupportedOperationException( "Native data sources require Java " + MINIMUM_JAVA_VERSION + " or later."); @@ -249,30 +166,7 @@ static NativeLibrary load(String path, String sparkVersion) throws Throwable { throw new LoadException( "Cannot use the Foreign Function and Memory API: " + rootCause(e), e); } - Optional init = (Optional) Ffm.FIND.invoke(lookup, "spark_ds_init"); - if (init.isEmpty()) { - throw new LoadException( - "The library does not export the function spark_ds_init.", null, true); - } - MethodHandle initFunction = Ffm.downcall(init.get(), C.I32, C.PTR, C.PTR); - // The library may keep the host until the JVM exits, even if spark_ds_init fails. - long host = Ffm.host(path, sparkVersion); - long out = Platform.allocateMemory(8); - try { - Platform.putLong(null, out, 0); - int status = (int) initFunction.invokeExact(host, out); - if (status != 0) { - throw new LoadException( - "spark_ds_init failed with the error code " + status + ".", null, true); - } - long table = Platform.getLong(null, out); - if (table == 0) { - throw new LoadException("spark_ds_init returned no function table.", null, true); - } - return new NativeLibrary(table); - } finally { - Platform.freeMemory(out); - } + return new NativeLibrary(lookup, path); } /** The method handles of the FFM API, initialized when the first library is loaded. */ @@ -297,8 +191,8 @@ private static final class Ffm { // NativeLibrary.log, which SparkDsHost.log calls with the path of its library. static final MethodHandle LOG; // The addresses of the C functions that call Java, which live as long as the JVM: - // SparkDsSink, which calls NativeLibrary.sink, and SparkDsTaskContext.is_interrupted, which - // calls NativeLibrary.isInterrupted. + // SparkDsSink, which calls NativeLibrary.sink, and SparkDsHost.is_interrupted, which calls + // NativeLibrary.isInterrupted. static final long SINK; static final long IS_INTERRUPTED; @@ -354,7 +248,7 @@ private static final class Ffm { IS_INTERRUPTED = upcall( lookup.findStatic(NativeLibrary.class, "isInterrupted", MethodType.methodType(boolean.class, long.class)), - C.BOOL, C.PTR); + C.BOOL, C.I64); } catch (Throwable e) { // The message names the cause, because the JVM only keeps this message for the later // uses of the class. @@ -423,19 +317,13 @@ static long upcall(MethodHandle target, C result, C... parameters) throws Throwa * Allocates the SparkDsHost of a library. The library may keep it until the JVM exits, so it * is never freed. */ - static long host(String path, String sparkVersion) throws Throwable { + static long host(String path) throws Throwable { long log = upcall(MethodHandles.insertArguments(LOG, 0, path), C.VOID, C.I32, C.PTR); - byte[] version = sparkVersion.getBytes(StandardCharsets.UTF_8); - long versionAddress = Platform.allocateMemory(version.length + 1L); - Platform.copyMemory( - version, Platform.BYTE_ARRAY_OFFSET, null, versionAddress, version.length); - Platform.putByte(null, versionAddress + version.length, (byte) 0); long host = Platform.allocateMemory(HOST_SIZE); Platform.setMemory(host, (byte) 0, HOST_SIZE); Platform.putLong(null, host, HOST_SIZE); - Platform.putInt(null, host + HOST_ABI_VERSION, ABI_VERSION); - Platform.putLong(null, host + HOST_SPARK_VERSION, versionAddress); Platform.putLong(null, host + HOST_LOG, log); + Platform.putLong(null, host + HOST_IS_INTERRUPTED, IS_INTERRUPTED); return host; } } @@ -452,7 +340,7 @@ private static Throwable rootCause(Throwable e) { static final class LoadException extends RuntimeException { /** * Whether the library was loaded, in which case it stays loaded: Spark found that it cannot - * use it afterwards, for example because {@code spark_ds_init} failed. + * use it afterwards, for example because it does not export a function. */ final boolean loaded; @@ -466,13 +354,6 @@ static final class LoadException extends RuntimeException { } } - /** Thrown when a function is called that the library leaves NULL in its table, with its name. */ - static final class MissingFunctionException extends RuntimeException { - MissingFunctionException(String name) { - super(name); - } - } - /** * Called by the library, through {@code SparkDsHost.log}, with the path of the library. It must * not throw: an exception thrown by an upcall crashes the JVM. @@ -499,313 +380,101 @@ private static void log(String path, int level, long message) { // Calls // --------------------------------------------------------------------------------------------- - private MethodHandle function(Function function) { - MethodHandle handle = functions[function.ordinal()]; - if (handle == null) { - throw new MissingFunctionException(function.fieldName); - } - return handle; + /** The version of the interface that the library implements. */ + int abiVersion() { + return abiVersion; } /** Throws the error that the library reported if a function failed. */ - private void check(int status, Function function) throws Throwable { + private void check(int status, String function) throws Throwable { if (status == 0) { return; } - MethodHandle lastError = functions[Function.LAST_ERROR.ordinal()]; long address = lastError == null ? 0 : (long) lastError.invokeExact(); String message = address == 0 ? "" : new String(readCString(address), StandardCharsets.UTF_8); throw new RuntimeException(message.isEmpty() - ? function.fieldName + " failed with the error code " + status + "." + ? function + " failed with the error code " + status + "." : message); } - private static long readHandle(long out, Function function) { - long handle = Platform.getLong(null, out); - if (handle == 0) { - throw new IllegalStateException(function.fieldName + " returned a NULL handle."); - } - return handle; - } - - /** The version of the interface that the table of the library implements. */ - int abiVersion() { - return abiVersion; - } - - /** Whether the library implements the function, which it does not leave NULL in its table. */ - boolean hasFunction(Function function) { - return functions[function.ordinal()] != null; - } - - /** A data source handle, with the capabilities of the data source. */ - static final class DataSourceHandle { + /** An object that the library opened, with its handle and the reply of the library. */ + static final class Opened { final long handle; - final long capabilities; + final OpenReply reply; - DataSourceHandle(long handle, long capabilities) { + Opened(long handle, OpenReply reply) { this.handle = handle; - this.capabilities = capabilities; - } - } - - DataSourceHandle createDataSource(String name, String[] optionKeys, String[] optionValues) - throws Throwable { - MethodHandle f = function(Function.CREATE); - try (Memory memory = new Memory()) { - long info = memory.zeroed(CREATE_INFO_SIZE); - Platform.putLong(null, info, CREATE_INFO_SIZE); - Platform.putLong(null, info + CREATE_INFO_NAME, memory.string(name)); - Platform.putLong(null, info + CREATE_INFO_OPTION_KEYS, memory.strings(optionKeys)); - Platform.putLong(null, info + CREATE_INFO_OPTION_VALUES, memory.strings(optionValues)); - Platform.putLong(null, info + CREATE_INFO_NUM_OPTIONS, optionKeys.length); - long capabilities = memory.zeroed(8); - long out = memory.zeroed(8); - check((int) f.invokeExact(info, capabilities, out), Function.CREATE); - return new DataSourceHandle( - readHandle(out, Function.CREATE), Platform.getLong(null, capabilities)); - } - } - - void schema(long dataSource, long schemaAddress) throws Throwable { - check((int) function(Function.SCHEMA).invokeExact(dataSource, schemaAddress), - Function.SCHEMA); - } - - long createReader(long dataSource, long schemaAddress) throws Throwable { - return createWithSchema(Function.CREATE_READER, dataSource, schemaAddress); - } - - long createStreamReader(long dataSource, long schemaAddress) throws Throwable { - return createWithSchema(Function.CREATE_STREAM_READER, dataSource, schemaAddress); - } - - private long createWithSchema(Function function, long dataSource, long schemaAddress) - throws Throwable { - MethodHandle f = function(function); - try (Memory memory = new Memory()) { - long out = memory.zeroed(8); - check((int) f.invokeExact(dataSource, schemaAddress, out), function); - return readHandle(out, function); + this.reply = reply; } } - /** - * Creates a writer. - * - * @param mode {@link #WRITE_APPEND} or {@link #WRITE_TRUNCATE} - * @param queryId the ID of the streaming query of a streaming write, or null for a batch write - */ - long createWriter(long dataSource, long schemaAddress, int mode, String queryId) - throws Throwable { - MethodHandle f = function(Function.CREATE_WRITER); - try (Memory memory = new Memory()) { - long info = memory.zeroed(WRITE_INFO_SIZE); - Platform.putLong(null, info, WRITE_INFO_SIZE); - Platform.putInt(null, info + WRITE_INFO_MODE, mode); - Platform.putLong( - null, info + WRITE_INFO_QUERY_ID, queryId == null ? 0 : memory.string(queryId)); + /** Opens an object. */ + Opened open(OpenRequest request) throws Throwable { + try (Memory memory = new Memory(); Sink sink = new Sink()) { + byte[] bytes = request.toByteArray(); long out = memory.zeroed(8); - check((int) f.invokeExact(dataSource, schemaAddress, info, out), Function.CREATE_WRITER); - return readHandle(out, Function.CREATE_WRITER); - } - } - - void closeDataSource(long dataSource) throws Throwable { - function(Function.CLOSE).invokeExact(dataSource); - } - - /** Returns how the reader uses each predicate, as SPARK_DS_PREDICATE_* values. */ - int[] pushPredicates(long reader, String[] predicates) throws Throwable { - MethodHandle f = function(Function.PUSH_PREDICATES); - try (Memory memory = new Memory()) { - long uses = memory.zeroed(4L * Math.max(1, predicates.length)); - check((int) f.invokeExact(reader, memory.strings(predicates), (long) predicates.length, - uses), Function.PUSH_PREDICATES); - int[] result = new int[predicates.length]; - for (int i = 0; i < result.length; i++) { - result[i] = Platform.getInt(null, uses + 4L * i); + check((int) open.invokeExact(host, memory.bytes(bytes), (long) bytes.length, Ffm.SINK, + sink.id, out), "spark_ds_open"); + long handle = Platform.getLong(null, out); + if (handle == 0) { + throw new IllegalStateException("spark_ds_open returned a NULL handle."); } - return result; - } - } - - boolean pushLimit(long reader, int limit) throws Throwable { - MethodHandle f = function(Function.PUSH_LIMIT); - try (Memory memory = new Memory()) { - long applied = memory.zeroed(1); - check((int) f.invokeExact(reader, limit, applied), Function.PUSH_LIMIT); - return Platform.getByte(null, applied) != 0; - } - } - - boolean pruneColumns(long reader, String[] columnNames) throws Throwable { - MethodHandle f = function(Function.PRUNE_COLUMNS); - try (Memory memory = new Memory()) { - long pruned = memory.zeroed(1); - check((int) f.invokeExact(reader, memory.strings(columnNames), (long) columnNames.length, - pruned), Function.PRUNE_COLUMNS); - return Platform.getByte(null, pruned) != 0; - } - } - - byte[][] partitions(long reader) throws Throwable { - MethodHandle f = function(Function.PARTITIONS); - try (Sink sink = new Sink()) { - check((int) f.invokeExact(reader, Ffm.SINK, sink.id), Function.PARTITIONS); - return sink.all(); - } - } - - byte[] serializeReader(long reader) throws Throwable { - return serialize(Function.SERIALIZE_READER, reader); - } - - void closeReader(long reader) throws Throwable { - function(Function.CLOSE_READER).invokeExact(reader); - } - - String initialOffset(long streamReader) throws Throwable { - return offset(Function.INITIAL_OFFSET, streamReader); - } - - String latestOffset(long streamReader) throws Throwable { - return offset(Function.LATEST_OFFSET, streamReader); - } - - private String offset(Function function, long streamReader) throws Throwable { - byte[] offset = serialize(function, streamReader); - if (offset == null) { - throw new IllegalStateException(function.fieldName + " returned no offset."); - } - return new String(offset, StandardCharsets.UTF_8); - } - - byte[][] streamPartitions(long streamReader, String start, String end) throws Throwable { - MethodHandle f = function(Function.STREAM_PARTITIONS); - try (Memory memory = new Memory(); Sink sink = new Sink()) { - check((int) f.invokeExact(streamReader, memory.string(start), memory.string(end), - Ffm.SINK, sink.id), Function.STREAM_PARTITIONS); - return sink.all(); - } - } - - byte[] serializeStreamReader(long streamReader) throws Throwable { - return serialize(Function.SERIALIZE_STREAM_READER, streamReader); - } - - void commitOffset(long streamReader, String end) throws Throwable { - MethodHandle f = function(Function.COMMIT_OFFSET); - try (Memory memory = new Memory()) { - check((int) f.invokeExact(streamReader, memory.string(end)), Function.COMMIT_OFFSET); - } - } - - void closeStreamReader(long streamReader) throws Throwable { - function(Function.CLOSE_STREAM_READER).invokeExact(streamReader); - } - - void read(Task task, byte[] readerState, byte[] partition, long streamAddress) - throws Throwable { - MethodHandle f = function(Function.READ); - try (Memory memory = new Memory()) { - check((int) f.invokeExact(task.address, memory.bytes(readerState), - (long) readerState.length, memory.bytes(partition), (long) partition.length, - streamAddress), Function.READ); - } - } - - byte[] serializeWriter(long writer) throws Throwable { - return serialize(Function.SERIALIZE_WRITER, writer); - } - - void commit(long writer, long epochId, byte[][] messages) throws Throwable { - finishWrite(Function.COMMIT, writer, epochId, messages); - } - - void abort(long writer, long epochId, byte[][] messages) throws Throwable { - finishWrite(Function.ABORT, writer, epochId, messages); - } - - private void finishWrite(Function function, long writer, long epochId, byte[][] messages) - throws Throwable { - MethodHandle f = function(function); - try (Memory memory = new Memory()) { - long pointers = memory.zeroed(8L * Math.max(1, messages.length)); - long lengths = memory.zeroed(8L * Math.max(1, messages.length)); - for (int i = 0; i < messages.length; i++) { - if (messages[i] != null) { - Platform.putLong(null, pointers + 8L * i, memory.bytes(messages[i])); - Platform.putLong(null, lengths + 8L * i, messages[i].length); - } + try { + return new Opened(handle, OpenReply.parseFrom(sink.single("spark_ds_open"))); + } catch (InvalidProtocolBufferException e) { + // The library opened the object, but its reply is invalid. + close(handle); + throw new IllegalStateException(invalidReply("spark_ds_open", e), e); + } catch (Throwable e) { + close(handle); + throw e; } - check((int) f.invokeExact(writer, epochId, pointers, lengths, (long) messages.length), - function); - } - } - - void closeWriter(long writer) throws Throwable { - function(Function.CLOSE_WRITER).invokeExact(writer); - } - - long createDataWriter(Task task, byte[] writerState, long epochId) throws Throwable { - MethodHandle f = function(Function.CREATE_DATA_WRITER); - try (Memory memory = new Memory()) { - long out = memory.zeroed(8); - check((int) f.invokeExact(task.address, memory.bytes(writerState), - (long) writerState.length, epochId, out), Function.CREATE_DATA_WRITER); - return readHandle(out, Function.CREATE_DATA_WRITER); } } - void write(long dataWriter, long arrayAddress, long schemaAddress) throws Throwable { - check((int) function(Function.WRITE).invokeExact(dataWriter, arrayAddress, schemaAddress), - Function.WRITE); + /** Runs a control operation on an object. */ + CallReply call(long handle, CallRequest request) throws Throwable { + return call(handle, request, null); } /** - * Commits a data writer. Then {@code committed} runs right away, before the message is checked: - * if the message is invalid, the commit fails, but the data writer must not be aborted. + * Runs a control operation on an object. Then {@code succeeded} runs right away, before the + * reply is checked: for example, a data writer that committed must not be aborted, even if its + * reply is invalid. */ - byte[] commitDataWriter(long dataWriter, Runnable committed) throws Throwable { - MethodHandle f = function(Function.COMMIT_DATA_WRITER); - try (Sink sink = new Sink()) { - check((int) f.invokeExact(dataWriter, Ffm.SINK, sink.id), Function.COMMIT_DATA_WRITER); - committed.run(); - byte[] message = single(sink, Function.COMMIT_DATA_WRITER); - return message == null ? new byte[0] : message; + CallReply call(long handle, CallRequest request, Runnable succeeded) throws Throwable { + try (Memory memory = new Memory(); Sink sink = new Sink()) { + byte[] bytes = request.toByteArray(); + check((int) call.invokeExact(handle, memory.bytes(bytes), (long) bytes.length, Ffm.SINK, + sink.id), "spark_ds_call"); + if (succeeded != null) { + succeeded.run(); + } + try { + return CallReply.parseFrom(sink.single("spark_ds_call")); + } catch (InvalidProtocolBufferException e) { + throw new IllegalStateException(invalidReply("spark_ds_call", e), e); + } } } - void abortDataWriter(long dataWriter) throws Throwable { - check((int) function(Function.ABORT_DATA_WRITER).invokeExact(dataWriter), - Function.ABORT_DATA_WRITER); - } - - void closeDataWriter(long dataWriter) throws Throwable { - function(Function.CLOSE_DATA_WRITER).invokeExact(dataWriter); - } - /** - * Calls a function that returns a single byte array through a sink. Returns null if the - * function does not call the sink. + * Exchanges the data of an object. + * + * @param in the address of the input stream, or 0 + * @param out the address of the output stream, or 0 */ - private byte[] serialize(Function function, long handle) throws Throwable { - MethodHandle f = function(function); - try (Sink sink = new Sink()) { - check((int) f.invokeExact(handle, Ffm.SINK, sink.id), function); - return single(sink, function); + void stream(long handle, StreamRequest request, long in, long out) throws Throwable { + try (Memory memory = new Memory()) { + byte[] bytes = request.toByteArray(); + check((int) stream.invokeExact(handle, memory.bytes(bytes), (long) bytes.length, in, out), + "spark_ds_stream"); } } - /** Returns the single byte array that a function passed to the sink, or null if none. */ - private static byte[] single(Sink sink, Function function) throws Throwable { - byte[][] values = sink.all(); - if (values.length > 1) { - throw new IllegalStateException( - function.fieldName + " returned " + values.length + " values instead of one."); - } - return values.length == 0 ? null : values[0]; + /** Releases the handle of an object. */ + void close(long handle) throws Throwable { + close.invokeExact(handle); } // --------------------------------------------------------------------------------------------- @@ -830,23 +499,6 @@ long bytes(byte[] bytes) { return address; } - /** Copies the string as a NUL-terminated UTF-8 string. */ - long string(String value) { - byte[] bytes = value.getBytes(StandardCharsets.UTF_8); - long address = zeroed(bytes.length + 1L); - Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address, bytes.length); - return address; - } - - /** Copies the strings, and returns the address of an array of pointers to them. */ - long strings(String[] values) { - long array = zeroed(8L * Math.max(1, values.length)); - for (int i = 0; i < values.length; i++) { - Platform.putLong(null, array + 8L * i, string(values[i])); - } - return array; - } - @Override public void close() { allocations.forEach(Platform::freeMemory); @@ -864,59 +516,45 @@ private static byte[] readCString(long address) { } // --------------------------------------------------------------------------------------------- - // Task contexts + // Tasks // --------------------------------------------------------------------------------------------- - // Whether the task of each SparkDsTaskContext was killed, by the address of the struct. - private static final ConcurrentHashMap tasks = new ConcurrentHashMap<>(); + /** A task that uses the library, and the number of objects that it opened. */ + private static final class Task { + final BooleanSupplier interrupted; + int objects; - /** - * The {@code SparkDsTaskContext} of a task that reads a partition or writes data. The reader or - * the data writer closes it once the library no longer uses it, which frees it. - */ - static final class Task implements AutoCloseable { - final long address; - private boolean closed; - - Task( - int stageId, - int stageAttemptNumber, - int partitionId, - int attemptNumber, - long taskAttemptId, - BooleanSupplier interrupted) { - address = Platform.allocateMemory(TASK_SIZE); - Platform.setMemory(address, (byte) 0, TASK_SIZE); - Platform.putLong(null, address, TASK_SIZE); - Platform.putInt(null, address + TASK_STAGE_ID, stageId); - Platform.putInt(null, address + TASK_STAGE_ATTEMPT_NUMBER, stageAttemptNumber); - Platform.putInt(null, address + TASK_PARTITION_ID, partitionId); - Platform.putInt(null, address + TASK_ATTEMPT_NUMBER, attemptNumber); - Platform.putLong(null, address + TASK_ATTEMPT_ID, taskAttemptId); - Platform.putLong(null, address + TASK_IS_INTERRUPTED, Ffm.IS_INTERRUPTED); - tasks.put(address, interrupted); + Task(BooleanSupplier interrupted) { + this.interrupted = interrupted; } + } - @Override - public void close() { - if (!closed) { - closed = true; - // A call to is_interrupted that is still running only looks the address up. - tasks.remove(address); - Platform.freeMemory(address); - } - } + // The tasks that use a library, by the ID of their task attempt, for SparkDsHost.is_interrupted. + private static final ConcurrentHashMap tasks = new ConcurrentHashMap<>(); + + /** + * Registers a task for {@code SparkDsHost.is_interrupted}, until the returned + * {@link AutoCloseable} is closed, once the library no longer uses the object of the task. + */ + static AutoCloseable registerTask(long taskAttemptId, BooleanSupplier interrupted) { + tasks.compute(taskAttemptId, (id, task) -> { + Task registered = task == null ? new Task(interrupted) : task; + registered.objects++; + return registered; + }); + return () -> tasks.computeIfPresent(taskAttemptId, (id, task) -> + --task.objects == 0 ? null : task); } /** - * Called by the library, through {@code SparkDsTaskContext.is_interrupted}, with the address of - * the struct. It must not throw: an exception thrown by an upcall crashes the JVM. + * Called by the library, through {@code SparkDsHost.is_interrupted}, with the ID of a task + * attempt. It must not throw: an exception thrown by an upcall crashes the JVM. */ - private static boolean isInterrupted(long task) { + private static boolean isInterrupted(long taskAttemptId) { try { - BooleanSupplier interrupted = tasks.get(task); - // A freed task context belongs to a read or a write that is over. - return interrupted == null || interrupted.getAsBoolean(); + Task task = tasks.get(taskAttemptId); + // A task that is not registered has ended, or is not a task of this executor. + return task == null || task.interrupted.getAsBoolean(); } catch (Throwable e) { return false; } @@ -930,7 +568,7 @@ private static boolean isInterrupted(long task) { private static final ConcurrentHashMap sinks = new ConcurrentHashMap<>(); /** - * Collects the byte arrays that a library returns during a call, through the C function + * Collects the reply that a library returns during a call, through the C function * {@link Ffm#SINK}. The library passes {@link #id} back to identify it. */ private static final class Sink implements AutoCloseable { @@ -942,11 +580,16 @@ private static final class Sink implements AutoCloseable { sinks.put(id, this); } - byte[][] all() throws Throwable { + /** Returns the reply, or no bytes, an empty message, if the library did not reply. */ + byte[] single(String function) throws Throwable { if (error != null) { throw error; } - return values.toArray(new byte[0][]); + if (values.size() > 1) { + throw new IllegalStateException( + function + " replied " + values.size() + " times instead of once."); + } + return values.isEmpty() ? new byte[0] : values.get(0); } @Override @@ -956,8 +599,8 @@ public void close() { } /** - * Called by the library, through the upcall stub {@link Ffm#SINK}, with a byte array to copy. - * It must not throw: an exception thrown by an upcall crashes the JVM. + * Called by the library, through the upcall stub {@link Ffm#SINK}, with a reply to copy. It + * must not throw: an exception thrown by an upcall crashes the JVM. */ private static void sink(long id, long data, long length) { Sink sink = null; @@ -967,7 +610,7 @@ private static void sink(long id, long data, long length) { return; } if (length < 0 || length > Integer.MAX_VALUE - 8 || (data == 0 && length > 0)) { - throw new IllegalStateException("The library returned an invalid byte array, of " + + throw new IllegalStateException("The library returned an invalid reply, of " + length + " bytes at the address " + data + "."); } byte[] bytes = new byte[(int) length]; @@ -979,4 +622,9 @@ private static void sink(long id, long data, long length) { } } } + + /** Describes an invalid reply of the library. */ + private static String invalidReply(String function, InvalidProtocolBufferException e) { + return function + " returned an invalid reply: " + e.getMessage(); + } } diff --git a/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto b/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto new file mode 100644 index 0000000000000..a296f7d4f5e28 --- /dev/null +++ b/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto @@ -0,0 +1,487 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// The messages that Spark and a native data source exchange through the functions of +// spark_datasource.h, which describes how they are used. Since Spark 4.4.0. +// +// The messages only grow: a later release of Spark may add fields, messages and values of enums, +// but it never changes or reuses the number of a field, so a library built against an earlier +// version of this file keeps working, and ignores the fields that it does not know. Spark only +// sends a new kind of request to an object whose reply to spark_ds_open declares that it supports +// it, so a new request never reaches an older library. +syntax = "proto3"; + +package spark.datasource; + +option java_multiple_files = true; +option java_package = "org.apache.spark.sql.execution.datasources.v2.ffi.proto"; +// Lite messages have no descriptors, which the full C++ runtime registers in a pool of the +// process: two libraries that link the same libprotobuf, or two copies of one library, such as an +// installed one and the one of a package, would register spark_datasource.proto twice, which +// aborts the process. A library should still link its Protocol Buffers runtime statically, with +// hidden symbols, since libraries built against different versions of it can share a process. +option optimize_for = LITE_RUNTIME; + +// --------------------------------------------------------------------------------------------- +// spark_ds_open +// --------------------------------------------------------------------------------------------- + +// The request of spark_ds_open: the object to open, with what it needs. +message OpenRequest { + // The version of Spark, such as "4.4.0". + string spark_version = 1; + // The newest version of these messages that Spark supports: 1. + int32 protocol_version = 2; + + oneof object { + // A data source, on the driver, for each schema inference, scan and write. Spark closes it + // right after. + DataSource data_source = 10; + // A reader for a batch scan, on the driver. + BatchReader batch_reader = 11; + // A reader for a micro-batch streaming scan, on the driver. Spark keeps it until the + // streaming query stops. + StreamReader stream_reader = 12; + // A writer for a batch write, or for a micro-batch of a streaming write, on the driver. + Writer writer = 13; + // The reader of one partition of a batch or streaming scan, on an executor. + PartitionReader partition_reader = 14; + // The writer of one task of a write, on an executor. + DataWriter data_writer = 15; + } +} + +// The reply of spark_ds_open. +message OpenReply { + // For a data source: the operations that it supports, which may depend on its name and its + // options. Spark reports the operations that a data source does not declare as unsupported, + // without opening a reader or a writer. + repeated Capability capabilities = 1; +} + +// An operation that a data source supports. +enum Capability { + CAPABILITY_UNSPECIFIED = 0; + // Batch reads, with batch readers. + CAPABILITY_BATCH_READ = 1; + // Micro-batch streaming reads, with stream readers. + CAPABILITY_MICRO_BATCH_READ = 2; + // Batch writes, with writers. + CAPABILITY_BATCH_WRITE = 3; + // Streaming writes, with writers. + CAPABILITY_STREAMING_WRITE = 4; + // Writes with the mode WRITE_MODE_TRUNCATE. + CAPABILITY_TRUNCATE = 5; +} + +// A data source, with the options given by the user, such as spark.read.option("key", "value"). +message DataSource { + // The name of the data source, as listed in the package manifest, or in lower case for an + // installed library. A library can implement several data sources. + string name = 1; + // The options, with the keys in lower case, because Spark matches them case-insensitively. + map options = 2; +} + +message BatchReader { + DataSource data_source = 1; + // The schema to read: the one that the user specified, or the one that the data source infers. + StructType schema = 2; +} + +message StreamReader { + DataSource data_source = 1; + // The schema to read: the one that the user specified, or the one that the data source infers. + StructType schema = 2; +} + +message Writer { + DataSource data_source = 1; + // The schema of the data to write. + StructType schema = 2; + // Whether the write adds to the existing data, or replaces it. + WriteMode mode = 3; + // The ID of the streaming query of a streaming write, and empty for a batch write. It stays the + // same when the query restarts from its checkpoint, so with the epoch IDs, it identifies the + // commits of the query. + string query_id = 4; +} + +enum WriteMode { + // Adds the data to the existing data. + WRITE_MODE_APPEND = 0; + // Replaces the existing data: with the save mode "overwrite" of a batch write, and in each + // micro-batch with the output mode "complete" of a streaming write. Spark only uses it with a + // data source that declares CAPABILITY_TRUNCATE. + WRITE_MODE_TRUNCATE = 1; +} + +message PartitionReader { + // The task that reads the partition. + TaskContext task = 1; + // The state that the batch reader or the stream reader returned in its SerializeReply. + bytes reader_state = 2; + // A partition that the batch reader or the stream reader planned. + bytes partition = 3; +} + +message DataWriter { + // The task that writes, with the partition that it writes. + TaskContext task = 1; + // The state that the writer returned in its SerializeReply. + bytes writer_state = 2; + // The ID of the micro-batch for a streaming write, or -1 for a batch write. + int64 epoch_id = 3; +} + +// A task of Spark, on an executor. +message TaskContext { + // The ID of the stage of the task, and the attempt number of the stage. + int32 stage_id = 1; + int32 stage_attempt_number = 2; + // The partition that the task reads or writes, and the attempt number of the task. + int32 partition_id = 3; + int32 attempt_number = 4; + // The ID of the task attempt, unique in the Spark application. SparkDsHost.is_interrupted takes + // it. + int64 task_attempt_id = 5; +} + +// --------------------------------------------------------------------------------------------- +// spark_ds_call +// --------------------------------------------------------------------------------------------- + +// The request of spark_ds_call: a control operation on an object. The comment of each operation +// says on which object Spark calls it, and the reply has the field of the same name. +message CallRequest { + oneof call { + // On a data source, when the user does not specify a schema. + InferSchema infer_schema = 1; + // On a batch reader, at most once and in this order: push_predicates, push_limit and + // prune_columns. + PushPredicates push_predicates = 2; + PushLimit push_limit = 3; + PruneColumns prune_columns = 4; + // On a batch reader, after the pushdown. + PlanPartitions plan_partitions = 5; + // On a batch reader after plan_partitions, on a stream reader after each + // plan_stream_partitions, and on a writer once, after opening it. + Serialize serialize = 6; + // On a stream reader. + InitialOffset initial_offset = 7; + LatestOffset latest_offset = 8; + PlanStreamPartitions plan_stream_partitions = 9; + CommitOffset commit_offset = 10; + // On a writer. + Commit commit = 11; + Abort abort = 12; + // On a data writer, after the data was written. + CommitDataWriter commit_data_writer = 13; + AbortDataWriter abort_data_writer = 14; + } +} + +// The reply of spark_ds_call. An empty reply is valid for every request: it means that the +// operation has no result, such as a reader that does not use the pushed predicates. +message CallReply { + oneof reply { + InferSchemaReply infer_schema = 1; + PushPredicatesReply push_predicates = 2; + PushLimitReply push_limit = 3; + PruneColumnsReply prune_columns = 4; + PlanPartitionsReply plan_partitions = 5; + SerializeReply serialize = 6; + OffsetReply initial_offset = 7; + OffsetReply latest_offset = 8; + PlanPartitionsReply plan_stream_partitions = 9; + CommitDataWriterReply commit_data_writer = 13; + } +} + +message InferSchema {} + +message InferSchemaReply { + // The schema of the data. Without it, the user must specify the schema. + StructType schema = 1; +} + +// Pushes predicates down to a batch reader. The predicates are combined with AND. +// +// Each predicate is a JSON expression tree. An expression is one of: +// - {"type": "column", "name": ["a", "b"]}: a column. A nested field has a name part for each +// level. +// - {"type": "literal", "dataType": "bigint", "value": 1}: a literal. The data type is a Spark +// SQL type name such as int, string or decimal(10,2). The value is null, a JSON boolean or +// number for boolean and numeric types, a string for string, a base64 string for binary, a +// decimal string for decimal, the days since the epoch for date, and the microseconds since +// the epoch for timestamp and timestamp_ntz. +// - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in +// org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=> +// (equal, with null equal to null), <, <=, > and >= with two children, IN with the value to +// look for and then the values to look in, IS_NULL and IS_NOT_NULL with one child, +// STARTS_WITH, ENDS_WITH and CONTAINS with the string and the string to look for, AND and OR +// with two predicates, and NOT with one. The children of the other functions are columns, +// literals or predicates. +// Spark does not pass predicates that it cannot express this way, such as predicates that use +// arithmetic or other functions, nor predicates on columns that contain strings with a +// non-binary collation, such as UTF8_LCASE: strings are compared by their UTF-8 bytes. The JSON +// is ASCII: it escapes the other characters, which any JSON parser decodes. +// +// Spark evaluates all the predicates again after reading, so a reader may use them only to skip +// data, for example the files whose statistics show that no row matches. But it must never drop +// or skip a row for which the predicates are true under the semantics of Spark SQL, which Spark +// cannot detect: see spark_datasource.h. The file +// org/apache/spark/sql/datasource/predicate-cases.json in the jar of Spark SQL lists predicates as +// Spark passes them, with the rows of a table for which they are true, to test a reader with. +message PushPredicates { + repeated string predicates_json = 1; +} + +message PushPredicatesReply { + // How the reader uses each predicate, in the order of the request. Spark shows the predicates + // that the reader uses as pushed in the plan of the query, and evaluates all of them again after + // reading anyway. Missing answers are PREDICATE_USE_UNUSED. + repeated PredicateUse uses = 1; +} + +enum PredicateUse { + // Not at all. + PREDICATE_USE_UNUSED = 0; + // Only to skip data, so the reader may return rows for which the predicate is not true. + PREDICATE_USE_PARTIAL = 1; + // To only return the rows for which the predicate is true. + PREDICATE_USE_COMPLETE = 2; +} + +// Pushes a LIMIT down to a batch reader. Spark only pushes a limit when the query does not filter +// the rows, since it evaluates the predicates after reading. +message PushLimit { + int32 limit = 1; +} + +message PushLimitReply { + // Whether the reader returns at most `limit` rows in total. Spark applies the limit again + // after reading either way. + bool applied = 1; +} + +// Prunes the columns that a batch reader reads. +message PruneColumns { + // The names of the top-level columns that the query needs, which include the columns of the + // pushed predicates. There may be none, for example for SELECT count(*): then the batches only + // have a number of rows. + repeated string columns = 1; +} + +message PruneColumnsReply { + // Whether the partition readers return exactly these columns, in this order. If not, they return + // all the columns of the schema to read. + bool pruned = 1; +} + +// Plans the partitions of a batch scan. Spark reads each partition in a separate task. +message PlanPartitions {} + +message PlanPartitionsReply { + // The partitions, serialized in a format defined by the library. + repeated bytes partitions = 1; +} + +// Serializes the state of a reader or a writer that its partition readers or data writers need +// on the executors, such as the pushed predicates. +message Serialize {} + +message SerializeReply { + // The state, in a format defined by the library, or empty if there is none. + bytes state = 1; +} + +// Returns the offset to start from when the streaming query has no checkpoint. +message InitialOffset {} + +// Returns the most recent offset available. +message LatestOffset {} + +message OffsetReply { + // The offset, as compact JSON on a single line: Spark stores each offset as a line of the offset + // log of the checkpoint, so it rejects an offset that is not valid JSON or contains a line break. + string offset_json = 1; +} + +// Plans the partitions that read the data after `start` up to and including `end`. When a +// streaming query restarts, from its checkpoint, Spark may plan the last range that it planned +// before again, so the partitions of a range must read the same data each time. +message PlanStreamPartitions { + string start_json = 1; + string end_json = 2; +} + +// Tells the library that Spark has processed all the data up to and including `end`, and will not +// ask for it again. +message CommitOffset { + string end_json = 1; +} + +// Commits a write, when all its tasks succeeded. Spark may commit the same epoch of a streaming +// query more than once, for example when the query restarts after a failure, so committing an +// epoch that is already committed must have no effect, for exactly-once output. +message Commit { + // The ID of the micro-batch for a streaming write, or -1 for a batch write. + int64 epoch_id = 1; + // The messages that the data writers of the tasks returned, one for each partition. + repeated TaskMessage messages = 2; +} + +// Aborts a write, when it failed, including when its commit failed. +message Abort { + // The ID of the micro-batch for a streaming write, or -1 for a batch write. + int64 epoch_id = 1; + // The messages that the data writers of the tasks returned. A task that did not commit has no + // message. + repeated TaskMessage messages = 2; +} + +message TaskMessage { + // Whether the task committed, and so whether `message` is set. + bool committed = 1; + bytes message = 2; +} + +// Commits the data that a data writer wrote. If it fails, Spark aborts the data writer. If it +// succeeds, Spark does not abort it, even if it then fails the task, for example because the +// reply is invalid. +message CommitDataWriter {} + +message CommitDataWriterReply { + // The message to pass to the commit of the writer on the driver. + bytes message = 1; +} + +// Aborts the data that a data writer wrote. +message AbortDataWriter {} + +// --------------------------------------------------------------------------------------------- +// spark_ds_stream +// --------------------------------------------------------------------------------------------- + +// The request of spark_ds_stream: a data operation on an object, through Arrow C streams. +message StreamRequest { + oneof stream { + // On a partition reader: initializes the output stream with the batches of the partition. + Read read = 1; + // On a data writer: reads the batches of the input stream and writes them. + Write write = 2; + } +} + +message Read {} + +message Write {} + +// --------------------------------------------------------------------------------------------- +// Schemas +// --------------------------------------------------------------------------------------------- + +// A schema: the columns of the data, as Spark SQL types. The data itself is Arrow, with the types +// of spark.sql.execution.arrow: for example, bigint is int64, string is utf8, timestamp is +// timestamp[us] with a time zone and timestamp_ntz is timestamp[us] without one. +message StructType { + repeated StructField fields = 1; +} + +message StructField { + string name = 1; + DataType data_type = 2; + bool nullable = 3; +} + +// A Spark SQL type. +message DataType { + oneof kind { + Null null = 1; + Boolean boolean = 2; + Byte byte = 3; + Short short = 4; + Integer integer = 5; + Long long = 6; + Float float = 7; + Double double = 8; + Decimal decimal = 9; + String string = 10; + Binary binary = 11; + Date date = 12; + Timestamp timestamp = 13; + TimestampNtz timestamp_ntz = 14; + YearMonthInterval year_month_interval = 15; + DayTimeInterval day_time_interval = 16; + CalendarInterval calendar_interval = 17; + Array array = 18; + Map map = 19; + StructType struct = 20; + } + + message Null {} + message Boolean {} + message Byte {} + message Short {} + message Integer {} + message Long {} + message Float {} + message Double {} + + message Decimal { + int32 precision = 1; + int32 scale = 2; + } + + message String { + // The collation of the strings, such as "UTF8_LCASE", or empty for UTF8_BINARY. + string collation = 1; + } + + message Binary {} + message Date {} + message Timestamp {} + message TimestampNtz {} + + // The fields are those of org.apache.spark.sql.types.YearMonthIntervalType: 0 for YEAR and 1 for + // MONTH. + message YearMonthInterval { + int32 start_field = 1; + int32 end_field = 2; + } + + // The fields are those of org.apache.spark.sql.types.DayTimeIntervalType: 0 for DAY, 1 for + // HOUR, 2 for MINUTE and 3 for SECOND. + message DayTimeInterval { + int32 start_field = 1; + int32 end_field = 2; + } + + message CalendarInterval {} + + message Array { + DataType element_type = 1; + bool contains_null = 2; + } + + message Map { + DataType key_type = 1; + DataType value_type = 2; + bool value_contains_null = 3; + } +} diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json b/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json index a2be573c93029..332e36d25273f 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json @@ -1,5 +1,5 @@ { - "description": "Predicates as Spark passes them to push_predicates of spark_datasource.h, with the ids of the rows of the table for which they are true. A case lists the predicates that Spark pushes for the filter in its SQL, which are combined with AND. A row has a value for each column of the schema; NaN and infinite doubles are strings.", + "description": "Predicates as Spark passes them in PushPredicates of spark_datasource.proto, with the ids of the rows of the table for which they are true. A case lists the predicates that Spark pushes for the filter in its SQL, which are combined with AND. A row has a value for each column of the schema; NaN and infinite doubles are strings.", "schema": "id INT,i INT,d DOUBLE,s STRING", "rows": [ [0, 1, 0.0, "a"], diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h index 3ff46d2871560..1a29b7a192c9c 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -19,19 +19,22 @@ * spark_datasource.h: the binary interface between Spark and data sources implemented in native * code, such as Rust or C++. Since Spark 4.4.0. * - * A native data source is a shared library that exports a single C function, spark_ds_init, - * which returns the table of the functions that implement the data source, a SparkDsLibrary. - * Spark calls them with the Foreign Function and Memory API of Java, so native data sources - * require Java 22 or later. Each library has its own table, so any number of libraries can be - * loaded side by side. Spark adapts a native data source to Data Source V2, with batch and - * micro-batch streaming reads, batch and streaming writes, and predicate, column and limit - * pushdown. + * A native data source is a shared library that exports the six functions declared below. They + * never change. Everything that the interface grows with is in the messages of + * spark_datasource.proto, which Spark and the library exchange as Protocol Buffers bytes: which + * object to open, and every control operation on it, such as the pushdown of predicates, the + * planning of partitions, offsets and commits. The data itself is exchanged without copying, as + * Arrow C streams. Spark calls the functions with the Foreign Function and Memory API of Java, so + * native data sources require Java 22 or later. Spark adapts a native data source to Data Source + * V2, with batch and micro-batch streaming reads, batch and streaming writes, and predicate, + * column and limit pushdown. * * This header is self-contained: it only depends on the C standard library, and defines the * structs of the Arrow C data interface unless they are already defined, as by Arrow C++. A library - * can copy it into its source tree, or use the copy in the include directory of a Spark - * distribution. A library built against it works with all the releases of Spark that support its - * SPARK_DS_ABI_VERSION. + * can copy it and spark_datasource.proto into its source tree, or use the copies in the include + * directory of a Spark distribution, and generate the code of the messages with any Protocol + * Buffers compiler, such as protoc or prost. A library built against them works with all the + * releases of Spark that support its SPARK_DS_ABI_VERSION. * * Spark finds the library of a data source automatically when it is installed like other native * libraries and named after the data source: the library of my_source is @@ -46,128 +49,110 @@ * * Lifecycle * --------- - * Spark loads a library once in each JVM, and calls spark_ds_init before any other function. It - * creates a data source with create on the driver for each schema inference, scan and write, and - * closes it right after. create also returns the capabilities of the data source: Spark reports - * the operations that a data source does not declare as unsupported, without calling the library. + * Spark loads a library once in each JVM. It opens objects with spark_ds_open, whose OpenRequest + * says which object and carries everything that it needs, calls control operations on them with + * spark_ds_call, exchanges their data with spark_ds_stream, and releases them with spark_ds_close. + * The messages are those of spark_datasource.proto. * - * - Batch reads. On the driver, Spark creates a reader with create_reader, pushes operations down - * by calling, each at most once and in this order, push_predicates, push_limit and - * prune_columns, plans the partitions with partitions, serializes the state of the reader with - * serialize_reader, and closes it. Then it calls read on the executors, in one task for each - * partition. If Spark never plans the partitions, for example to only explain the plan of the - * query, it closes the reader once the reader is no longer referenced. - * - Streaming reads. On the driver, Spark creates a stream reader with create_stream_reader and - * keeps it until the streaming query stops. For each micro-batch, it asks for the latest_offset, - * plans the partitions with stream_partitions and serializes the state with - * serialize_stream_reader, then calls read on the executors. Offsets are defined by the library, - * as compact JSON on a single line: Spark stores each offset as a line of the offset log of the - * checkpoint, so it rejects an offset that is not valid JSON or contains a line break. When a - * streaming query restarts, from its checkpoint, Spark creates a new stream reader and may call - * stream_partitions again for the last range that it planned before, so the partitions of a - * range must read the same data each time. + * - Data sources. On the driver, for each schema inference, scan and write, Spark opens a data + * source, whose OpenReply declares its capabilities: Spark reports the operations that a data + * source does not declare as unsupported, without opening a reader or a writer. When the user + * does not specify a schema, Spark calls InferSchema. Spark closes the data source right after. + * - Batch reads. On the driver, Spark opens a batch reader, calls, each at most once and in this + * order, PushPredicates, PushLimit and PruneColumns, then PlanPartitions and Serialize, and + * closes it. Then, on the executors, it opens a partition reader for each partition, in a task, + * initializes its stream with a Read, reads the stream, releases it and closes the partition + * reader. If Spark never plans the partitions, for example to only explain the plan of the query, + * it closes the batch reader once the reader is no longer referenced. + * - Streaming reads. On the driver, Spark opens a stream reader and keeps it until the streaming + * query stops. For each micro-batch, it calls LatestOffset, PlanStreamPartitions and Serialize, + * then reads the partitions on the executors as for a batch read, and calls CommitOffset once it + * has processed the micro-batch. Offsets are defined by the library, as compact JSON on a single + * line. When a streaming query restarts, from its checkpoint, Spark opens a new stream reader and + * may plan the last range that it planned before again, so the partitions of a range must read + * the same data each time. * - Writes. On the driver, right before it starts the tasks of the write (or of a micro-batch), - * Spark creates a writer with create_writer, and serializes its state with serialize_writer. - * With adaptive query execution, Spark first runs the stages of the data to write that end with - * a shuffle, and creates no writer if they fail. The rest of the data, such as its last stage, - * is computed in the tasks of the write, so a writer may still be aborted because computing the - * data failed: do irreversible work, such as removing the existing data for an overwrite, in - * commit. On the executors, each task creates a data writer with create_data_writer, calls write - * for each batch of at most spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends - * once it holds spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls commit_data_writer, - * or abort_data_writer if the task or the commit fails, before it closes the data writer. Then - * the driver calls commit with the messages of the tasks if they all succeeded, or abort - * otherwise, including when commit fails, and closes the writer. Several attempts of a task may - * write the same partition, after a failure or with speculative execution, each with a data - * writer of its own. Spark lets the attempts of a partition commit one at a time, and passes to - * commit the message of one committed attempt for each partition, so the write consists of the - * data of these attempts only. It aborts the data writers of the other attempts, except on an - * executor that is lost, where a data writer is neither committed nor aborted, and may even have + * Spark opens a writer and calls Serialize. With adaptive query execution, Spark first runs the + * stages of the data to write that end with a shuffle, and opens no writer if they fail. The + * rest of the data, such as its last stage, is computed in the tasks of the write, so a writer + * may still be aborted because computing the data failed: do irreversible work, such as removing + * the existing data for an overwrite, in Commit. On the executors, each task opens a data writer, + * passes its data in one or more calls of Write, each with a stream of batches of at most + * spark.sql.execution.arrow.maxRecordsPerBatch rows, which also end once they hold + * spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls CommitDataWriter, or + * AbortDataWriter if the task or the commit fails, before it closes the data writer. Then the + * driver calls Commit with the messages of the tasks if they all succeeded, or Abort otherwise, + * including when Commit fails, and closes the writer. Several attempts of a task may write the + * same partition, after a failure or with speculative execution, each with a data writer of its + * own. Spark lets the attempts of a partition commit one at a time, and passes to Commit the + * message of one committed attempt for each partition, so the write consists of the data of + * these attempts only. It aborts the data writers of the other attempts, except on an executor + * that is lost, where a data writer is neither committed nor aborted, and may even have * committed before its message reached the driver. A streaming write commits or aborts each - * micro-batch (epoch). Spark may commit the same epoch more than once, for example when a - * streaming query restarts after a failure, so committing an epoch that is already committed - * must have no effect, for exactly-once output. The ID of the query, which create_writer gets for - * a streaming write, stays the same when the query restarts from its checkpoint, so the ID of the - * query and the epoch ID identify the output of a micro-batch, also when other queries write to - * the same destination. A micro-batch streaming query creates a new stream writer for each - * micro-batch, and Spark closes the previous stream writer of the query when the query creates a - * new one. Spark is not told when a streaming query stops, so it closes the last stream writer of - * the query once the writer is no longer referenced, or when the query is restarted. + * micro-batch (epoch), which the query ID of the writer and the epoch ID identify, also when the + * query restarts from its checkpoint and when other queries write to the same destination. A + * micro-batch streaming query opens a new writer for each micro-batch, and Spark closes the + * previous writer of the query when the query opens a new one. Spark is not told when a + * streaming query stops, so it closes the last writer of the query once the writer is no longer + * referenced, or when the query is restarted. * * Conventions * ----------- * - Errors. The functions that return an int return 0 on success, or any other value on failure. - * They then set the message that last_error returns, and Spark reports it as a - * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail, and they may run after the query, - * from another thread: do the work that can fail a query, such as flushing data, in commit. A - * C++ exception or a Rust panic must not cross the interface: catch it and return an error - * instead. - * - Handles. create, create_reader, create_stream_reader, create_writer and create_data_writer - * return an opaque, non-NULL handle in their last argument, typically a pointer: Spark reports a - * NULL handle as an error. Spark releases each handle exactly once, with the close function of - * its kind: close, close_reader, close_stream_reader, close_writer or close_data_writer, and no - * other function releases a handle. A handle is not always released right after its last use: a - * batch reader whose partitions are never planned and the last stream writer of a query are - * released once they are no longer referenced (see Lifecycle above). A reader or writer handle - * must not depend on the data source handle it was created from, because Spark closes the data - * source handle right after creating it. - * - Memory. What Spark passes to a function, such as strings, byte arrays, SparkDsCreateInfo and - * SparkDsWriteInfo, is only valid during the call, except for SparkDsHost and - * SparkDsTaskContext, which document how long they live. Strings are UTF-8 and NUL-terminated. - * The JSON of the predicates is ASCII: it escapes the other characters, which any JSON parser - * decodes. The library returns strings and byte arrays by passing them to a SparkDsSink, which - * copies them. - * - State. Partitions, the state of readers and writers, and commit messages are byte arrays in a - * format defined by the library. Spark sends them to the executors and back. - * - Arrow data. Schemas and data are exchanged with the Arrow C data interface - * (https://arrow.apache.org/docs/format/CDataInterface.html), as structs allocated by Spark. - * For an output struct, the library initializes it and Spark takes ownership. For an input - * struct, the library may move it to take ownership; otherwise, Spark releases it after the - * call returns, so the library must not use it afterwards. The arrays that the library exports - * must have an offset of 0, at every level: Arrow Java ignores the offsets of the arrays it - * imports, so Spark rejects an array with another offset, such as a zero-copy slice. Copy a - * sliced array before exporting it. Spark also rejects an array with fewer values than its - * parent refers to, such as a column shorter than its batch. The buffers of an array must hold - * all of its values, as the C data interface requires: it does not carry the sizes of the - * buffers, so Spark cannot check them, and reads past the end of a buffer that is too short. - * - Schemas. Spark types map to Arrow types as in spark.sql.execution.arrow: for example, bigint - * is int64, string is utf8, timestamp is timestamp[us] with a time zone and timestamp_ntz is - * timestamp[us] without one. Data read from a library may also use large_utf8 and utf8_view for - * strings, large_binary and binary_view for binary data, and any names for the fields of the - * entries of a map. Spark rejects the other Arrow types, such as unsigned integers, decimal256, - * large lists, fixed-size lists and list views, and dictionary encoding. - * - Options. Option keys are passed in lower case, because Spark matches them - * case-insensitively. - * - Optional functions. A library only sets the functions of the operations it supports in its - * table, and leaves the others NULL. A data source must not declare a capability whose - * functions are NULL: Spark then fails the operation. The pushdown functions and the serialize - * functions are optional: without them, Spark pushes nothing down, and the reader or writer has - * no state. So are schema, without which the user must specify the schema, commit_offset and - * last_error. + * They then set the message that spark_ds_last_error returns, and Spark reports it as a + * NATIVE_DATA_SOURCE_ERROR. spark_ds_close cannot fail, and it may run after the query, from + * another thread: do the work that can fail a query, such as flushing data, in a commit. A C++ + * exception or a Rust panic must not cross the interface: catch it and return an error instead. + * - Handles. spark_ds_open returns an opaque, non-NULL handle in its last argument, typically a + * pointer: Spark reports a NULL handle as an error. Spark releases each handle exactly once, + * with spark_ds_close, and no other function releases a handle. A handle is not always released + * right after its last use: a batch reader whose partitions are never planned and the last + * writer of a streaming query are released once they are no longer referenced (see Lifecycle + * above). A handle does not depend on another one: each OpenRequest carries everything that the + * object needs. + * - Messages. A request is valid only during the call. The library replies by calling the + * SparkDsSink that it gets once, with the encoded reply, on the thread of the call and during the + * call. Not calling it is an empty reply, which is valid for every request. Both sides ignore the + * fields that they do not know. + * - Memory. What Spark passes to a function is only valid during the call, except for SparkDsHost, + * which stays valid until the JVM exits. + * - Arrow data. Data is exchanged with the Arrow C stream interface + * (https://arrow.apache.org/docs/format/CStreamInterface.html), as structs allocated by Spark. + * For an output stream, the library initializes it and Spark takes ownership: Spark reads it + * from the thread of the task, and releases it before it closes the object. For an input stream, + * the library reads it during the call, on the thread of the call, and may move it, or the arrays + * that it returns, to take ownership; otherwise, Spark releases it after the call returns. The + * arrays that the library exports must have an offset of 0, at every level: Arrow Java ignores + * the offsets of the arrays it imports, so Spark rejects an array with another offset, such as a + * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with + * fewer values than its parent refers to, such as a column shorter than its batch. The buffers of + * an array must hold all of its values, as the C data interface requires: it does not carry the + * sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that is + * too short. + * - Schemas. The messages describe schemas with Spark SQL types, which map to Arrow types as in + * spark.sql.execution.arrow: for example, bigint is int64, string is utf8, timestamp is + * timestamp[us] with a time zone and timestamp_ntz is timestamp[us] without one. Data read from a + * library may also use large_utf8 and utf8_view for strings, large_binary and binary_view for + * binary data, and any names for the fields of the entries of a map. Spark rejects the other + * Arrow types, such as unsigned integers, decimal256, large lists, fixed-size lists and list + * views, and dictionary encoding. * - Threads. Spark calls the library from multiple threads, and calls on different handles may * run concurrently. Calls on the same handle never run concurrently. A handle may be released * from a thread other than the threads that used it, such as a JVM Cleaner thread when Spark - * releases a handle once it is no longer referenced, for example the last stream writer of a - * streaming query. A SparkDsSink must be called on the thread of the call that it was passed - * to, during that call. The functions of SparkDsHost and SparkDsTaskContext can be called from - * any thread. + * releases a handle once it is no longer referenced, for example the last writer of a streaming + * query. The functions of SparkDsHost can be called from any thread. * * Compatibility * ------------- - * The function table and the structs that Spark passes to the library start with their size, so - * that they can grow. Within a SPARK_DS_ABI_VERSION: - * - fields are only appended to the structs, including SparkDsLibrary, and the existing fields - * never change, nor do the semantics of the existing functions; - * - Spark only reads the functions that fit in the size of the table of a library, so a library - * built against an earlier header keeps working, with the functions that it does not know as - * NULL. A library must only read the fields that fit in the size of a struct that Spark passes; - * - new values may appear where this documentation says so, such as the names of pushed - * predicates. Spark only uses a new capability or write mode with a data source that declares - * it. - * An incompatible change uses a new SPARK_DS_ABI_VERSION: spark_ds_init gets the newest version - * that Spark supports, the library returns a table of a version that it implements, and Spark - * rejects a table of a version that it does not support. Spark supports version 1 at least in all - * the releases of Spark 4 from 4.4.0. + * The six functions of this header and the fields of SparkDsHost never change. SparkDsHost starts + * with its size, and only gets fields appended: a library must only read the fields that fit in + * the size of the struct that Spark passes. The messages only grow, as spark_datasource.proto + * describes, and Spark only sends a new kind of request to an object whose OpenReply declares that + * it supports it. So a library built against an earlier version of these files keeps working with + * later releases of Spark, and a library built against a later version works with earlier + * releases, without the features that they do not know. SPARK_DS_ABI_VERSION is the version of + * the functions of this header, and Spark rejects a library of a version that it does not support. + * Spark supports version 1 at least in all the releases of Spark 4 from 4.4.0. */ #ifndef SPARK_DATASOURCE_H #define SPARK_DATASOURCE_H @@ -179,10 +164,10 @@ extern "C" { #endif -/* The version of this interface. */ +/* The version of this interface, which spark_ds_abi_version returns. */ #define SPARK_DS_ABI_VERSION 1 -/* Exports spark_ds_init from the library, also when the library hides its other symbols. */ +/* Exports the functions of this header, also when the library hides its other symbols. */ #if defined(_WIN32) #define SPARK_DS_EXPORT __declspec(dllexport) #else @@ -238,15 +223,11 @@ struct ArrowArrayStream { #endif /* ARROW_C_STREAM_INTERFACE */ /* - * Receives a string or a byte array that the library returns, and copies it. `sink` is the - * argument passed with the sink to the function that returns it. + * Receives the reply of a function, encoded as a message of spark_datasource.proto, and copies it. + * `sink` is the argument passed with the sink to the function. */ typedef void (*SparkDsSink)(void* sink, const uint8_t* data, int64_t length); -/* --------------------------------------------------------------------------------------------- - * What Spark passes to the library - * --------------------------------------------------------------------------------------------- */ - /* The levels of the messages that a library logs with SparkDsHost.log. */ #define SPARK_DS_LOG_ERROR 1 #define SPARK_DS_LOG_WARN 2 @@ -254,427 +235,73 @@ typedef void (*SparkDsSink)(void* sink, const uint8_t* data, int64_t length); #define SPARK_DS_LOG_DEBUG 4 #define SPARK_DS_LOG_TRACE 5 -/* Spark, as spark_ds_init gets it. It stays valid until the JVM exits. */ +/* Spark, as spark_ds_open gets it. It stays valid until the JVM exits. */ typedef struct SparkDsHost { /* * The size of this struct in the header that Spark was built with. A later version of Spark may * append fields, so a library must only read the fields that fit in it. */ int64_t size; - /* The newest version of this interface that Spark supports. */ - int32_t abi_version; - /* The version of Spark, such as "4.4.0". */ - const char* spark_version; /* * Logs a message, a NUL-terminated UTF-8 string, at one of the SPARK_DS_LOG_* levels: in the log * of the driver or of the executor, with the path of the library. */ void (*log)(int32_t level, const char* message); -} SparkDsHost; - -/* A data source, as create gets it. */ -typedef struct SparkDsCreateInfo { - /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ - int64_t size; - /* - * The name of the data source, as listed in the package manifest, or in lower case for an - * installed library. A library can implement several data sources. - */ - const char* name; /* - * The num_options options given by the user, such as spark.read.option("key", "value"), with - * the keys in lower case. + * Returns whether the task with the given task_attempt_id of its TaskContext was killed, for + * example because its query was canceled, or has ended. Spark cannot interrupt native code, so + * a read or a write that runs for long, or blocks, should check it regularly, and fail once it is + * true. When a task is killed, Spark only stops reading once the get_next function of the stream + * returns: a get_next that blocks forever keeps the task running, unless the task reaper + * (spark.task.reaper.enabled and spark.task.reaper.killTimeout) kills the executor. */ - const char* const* option_keys; - const char* const* option_values; - int64_t num_options; -} SparkDsCreateInfo; + bool (*is_interrupted)(int64_t task_attempt_id); +} SparkDsHost; + +/* Returns SPARK_DS_ABI_VERSION, the version of this interface that the library implements. */ +SPARK_DS_EXPORT int32_t spark_ds_abi_version(void); /* - * The task that calls read or create_data_writer, on an executor. It stays valid until the stream - * that read returns is released, or until the data writer is closed. + * Returns the message of the last error on the calling thread, as a NUL-terminated UTF-8 string + * that stays valid until the next call on that thread. */ -typedef struct SparkDsTaskContext { - /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ - int64_t size; - /* The ID of the stage of the task, and the attempt number of the stage. */ - int32_t stage_id; - int32_t stage_attempt_number; - /* The partition that the task reads or writes, and the attempt number of the task. */ - int32_t partition_id; - int32_t attempt_number; - /* The ID of the task attempt, unique in the Spark application. */ - int64_t task_attempt_id; - /* - * Returns whether the task was killed, for example because its query was canceled. Spark cannot - * interrupt native code, so a read or a write that runs for long, or blocks, should check it - * regularly, and fail once it is true. - */ - bool (*is_interrupted)(const struct SparkDsTaskContext* task); -} SparkDsTaskContext; - -/* The modes of a write. */ -#define SPARK_DS_WRITE_APPEND 0 -#define SPARK_DS_WRITE_TRUNCATE 1 - -/* A write, as create_writer gets it. */ -typedef struct SparkDsWriteInfo { - /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ - int64_t size; - /* - * SPARK_DS_WRITE_APPEND to add the data to the existing data, or SPARK_DS_WRITE_TRUNCATE to - * replace the existing data: with the save mode "overwrite" of a batch write, and in each - * micro-batch with the output mode "complete" of a streaming write. Spark only passes - * SPARK_DS_WRITE_TRUNCATE to a data source with the capability SPARK_DS_TRUNCATE. - */ - int32_t mode; - /* - * The ID of the streaming query of a streaming write, or NULL for a batch write. It stays the - * same when the query restarts from its checkpoint, so with the epoch IDs, it identifies the - * commits of the query: see Writes above. - */ - const char* query_id; -} SparkDsWriteInfo; - -/* --------------------------------------------------------------------------------------------- - * What the library returns to Spark - * --------------------------------------------------------------------------------------------- */ - -/* The capabilities of a data source, which create returns: the operations that it supports. */ -#define SPARK_DS_BATCH_READ 1 /* batch reads, with create_reader */ -#define SPARK_DS_MICRO_BATCH_READ 2 /* micro-batch streaming reads, with create_stream_reader */ -#define SPARK_DS_BATCH_WRITE 4 /* batch writes, with create_writer */ -#define SPARK_DS_STREAMING_WRITE 8 /* streaming writes, with create_writer */ -#define SPARK_DS_TRUNCATE 16 /* writes with the mode SPARK_DS_WRITE_TRUNCATE */ - -/* How a reader uses a predicate pushed with push_predicates. */ -#define SPARK_DS_PREDICATE_UNUSED 0 /* not at all */ -/* only to skip data, so it may return rows for which the predicate is not true */ -#define SPARK_DS_PREDICATE_PARTIAL 1 -/* to only return the rows for which the predicate is true */ -#define SPARK_DS_PREDICATE_COMPLETE 2 +SPARK_DS_EXPORT const char* spark_ds_last_error(void); /* - * The functions of a library, which spark_ds_init returns. A library sets the functions of the - * operations it supports, and leaves the others NULL: see "Optional functions" above. + * Opens an object: a data source, a batch reader, a stream reader, a writer, a partition reader or + * a data writer. + * + * host: Spark, the same struct for every call in a JVM + * request, length: an OpenRequest + * reply: called with an OpenReply, or not at all for an empty one + * handle: set to the handle of the object */ -typedef struct SparkDsLibrary { - /* - * The size of this struct in the header that the library was built with. Spark only reads the - * functions that fit in it. - */ - int64_t size; - /* The version of this interface that the table implements, SPARK_DS_ABI_VERSION. */ - int32_t abi_version; - - /* - * Returns the message of the last error on the calling thread, as a NUL-terminated UTF-8 string - * that stays valid until the next call on that thread. Optional: without it, Spark only reports - * the value that the function returned. - */ - const char* (*last_error)(void); - - /* ------------------------------------------------------------------------------------------- - * Data source (driver) - * ------------------------------------------------------------------------------------------- */ - - /* - * Creates a data source. - * - * info: the name of the data source, and the options given by the user - * capabilities: set to the capabilities of the data source, SPARK_DS_* flags combined with |, - * which may depend on its name and its options. It is 0 when create is called. - * source: set to the data source handle - */ - int (*create)(const SparkDsCreateInfo* info, uint64_t* capabilities, void** source); - - /* - * Exports the schema of the data. Spark calls it only when the user does not specify a schema. - * - * out: initialized with a struct type whose fields are the columns - */ - int (*schema)(void* source, struct ArrowSchema* out); - - /* - * Creates a reader for a batch scan. - * - * schema: the schema to read: the one the user specified, or the one returned by schema - * reader: set to the reader handle - */ - int (*create_reader)(void* source, struct ArrowSchema* schema, void** reader); - - /* - * Creates a reader for a micro-batch streaming scan. - * - * schema: the schema to read: the one the user specified, or the one returned by schema - * reader: set to the stream reader handle - */ - int (*create_stream_reader)(void* source, struct ArrowSchema* schema, void** reader); - - /* - * Creates a writer for a batch write, or for a micro-batch of a streaming write. - * - * schema: the schema of the data to write - * info: the mode of the write, and the ID of the query of a streaming write - * writer: set to the writer handle - */ - int (*create_writer)(void* source, struct ArrowSchema* schema, const SparkDsWriteInfo* info, - void** writer); - - /* Releases a data source handle. */ - void (*close)(void* source); - - /* ------------------------------------------------------------------------------------------- - * Batch reader (driver) - * ------------------------------------------------------------------------------------------- */ - - /* - * Pushes predicates down to the reader. The predicates are combined with AND. - * - * Each predicate is a JSON expression tree. An expression is one of: - * - {"type": "column", "name": ["a", "b"]}: a column. A nested field has a name part for each - * level. - * - {"type": "literal", "dataType": "bigint", "value": 1}: a literal. The data type is a Spark - * SQL type name such as int, string or decimal(10,2). The value is null, a JSON boolean or - * number for boolean and numeric types, a string for string, a base64 string for binary, a - * decimal string for decimal, the days since the epoch for date, and the microseconds since - * the epoch for timestamp and timestamp_ntz. - * - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in - * org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=> - * (equal, with null equal to null), <, <=, > and >= with two children, IN with the value to - * look for and then the values to look in, IS_NULL and IS_NOT_NULL with one child, - * STARTS_WITH, ENDS_WITH and CONTAINS with the string and the string to look for, AND and OR - * with two predicates, and NOT with one. The children of the other functions are columns, - * literals or predicates. - * Spark does not pass predicates that it cannot express this way, such as predicates that use - * arithmetic or other functions, nor predicates on columns that contain strings with a - * non-binary collation, such as UTF8_LCASE: strings are compared by their UTF-8 bytes. A later - * version of Spark may pass predicates with other names without changing SPARK_DS_ABI_VERSION, - * so a reader must ignore a predicate that contains a name that it does not know, or a data - * type that it does not support. - * - * Spark evaluates all the predicates again after reading, so a reader may use them only to skip - * data, for example the files whose statistics show that no row matches. But it must never drop - * or skip a row for which the predicates are true under the semantics of Spark SQL, which Spark - * cannot detect. In particular: - * - For floating-point columns, NaN is equal to NaN and greater than any other value, including - * positive infinity, and -0.0 is equal to 0.0. IEEE 754 comparisons treat NaN differently, and - * total orders, such as f64::total_cmp in Rust or the comparison kernels of arrow-rs, also - * order -0.0 before 0.0. For example, d > 1.0 is true for a NaN value of d, and d = 0.0 for - * -0.0. - * - Comparisons with null are null, and a row only matches if the predicates are true. For - * example, NOT (i IN (1, NULL)) matches no row. - * - Strings are compared by their UTF-8 bytes, which order the characters outside the Basic - * Multilingual Plane after all the others, unlike comparisons of UTF-16 strings. - * - To use only a part of a predicate, a reader may take the parts that it does not support as - * true, but not under a NOT, where that would drop rows. - * The file org/apache/spark/sql/datasource/predicate-cases.json in the jar of Spark SQL lists - * predicates as Spark passes them to this function, with the rows of a table for which they are - * true, to test a reader with. - * - * predicates: the count predicates, as JSON - * uses: an array of count values, all SPARK_DS_PREDICATE_UNUSED, to set to how the reader uses - * each predicate, a SPARK_DS_PREDICATE_* value. Spark shows the predicates that the reader - * uses as pushed in the plan of the query, and evaluates all the predicates again after - * reading anyway. It takes other values as SPARK_DS_PREDICATE_PARTIAL. - */ - int (*push_predicates)(void* reader, const char* const* predicates, int64_t count, - int32_t* uses); - - /* - * Pushes a LIMIT down to the reader. Spark only pushes a limit when the query does not filter - * the rows, since it evaluates the predicates after reading. - * - * applied: set to whether the reader returns at most `limit` rows in total. Spark applies the - * limit again after reading either way. - */ - int (*push_limit)(void* reader, int32_t limit, bool* applied); - - /* - * Prunes the columns to read. - * - * columns: the names of the count top-level columns that the query needs, which include the - * columns of the pushed predicates. There may be none, for example for - * SELECT count(*): then the batches only have a number of rows. - * pruned: set to whether read returns exactly these columns, in this order. If not, it returns - * all the columns of the schema to read. - */ - int (*prune_columns)(void* reader, const char* const* columns, int64_t count, bool* pruned); - - /* - * Plans the partitions of the scan. Spark reads each partition in a separate task. - * - * add_partition: called with each partition, serialized in a format defined by the library - */ - int (*partitions)(void* reader, SparkDsSink add_partition, void* sink); - - /* - * Serializes the state of the reader that read needs on the executors, such as the pushed - * predicates. Spark calls it after partitions. - * - * set_state: called once with the state, or not at all if there is none - */ - int (*serialize_reader)(void* reader, SparkDsSink set_state, void* sink); - - /* Releases a batch reader handle. */ - void (*close_reader)(void* reader); - - /* ------------------------------------------------------------------------------------------- - * Stream reader (driver) - * ------------------------------------------------------------------------------------------- */ - - /* - * Returns the offset to start from when the streaming query has no checkpoint. - * - * set_offset: called once with the offset, as compact JSON on a single line - */ - int (*initial_offset)(void* reader, SparkDsSink set_offset, void* sink); - - /* - * Returns the most recent offset available. - * - * set_offset: called once with the offset, as compact JSON on a single line - */ - int (*latest_offset)(void* reader, SparkDsSink set_offset, void* sink); - - /* - * Plans the partitions that read the data after `start` up to and including `end`, two offsets - * as JSON. - * - * add_partition: called with each partition, serialized in a format defined by the library - */ - int (*stream_partitions)(void* reader, const char* start, const char* end, - SparkDsSink add_partition, void* sink); - - /* - * Serializes the state of the stream reader that read needs on the executors. Spark calls it - * after each call to stream_partitions. - * - * set_state: called once with the state, or not at all if there is none - */ - int (*serialize_stream_reader)(void* reader, SparkDsSink set_state, void* sink); - - /* - * Tells the library that Spark has processed all the data up to and including `end`, an offset - * as JSON, and will not ask for it again. - */ - int (*commit_offset)(void* reader, const char* end); - - /* Releases a stream reader handle, when the streaming query stops. */ - void (*close_stream_reader)(void* reader); - - /* ------------------------------------------------------------------------------------------- - * Partition reader (executor) - * ------------------------------------------------------------------------------------------- */ - - /* - * Reads a partition of a batch or streaming scan. - * - * task: the task that reads the partition, valid until the stream is released - * reader_state: the state returned by serialize_reader or serialize_stream_reader, empty if - * there is none - * partition: a partition returned by partitions or stream_partitions - * out: the stream to initialize. Each array of the stream is a struct array whose fields are the - * columns to read: the columns of the schema to read, or the pruned columns. Spark calls the - * stream from the thread of the task that reads the partition. get_next may block, for - * example on network I/O, but Spark cannot interrupt native code: when the task is killed, - * Spark only stops reading once get_next returns, and then releases the stream. So a - * get_next that blocks forever keeps the task running, unless the task reaper - * (spark.task.reaper.enabled and spark.task.reaper.killTimeout) kills the executor. Use - * timeouts for blocking operations, and check task->is_interrupted. - */ - int (*read)(const SparkDsTaskContext* task, const uint8_t* reader_state, - int64_t reader_state_length, const uint8_t* partition, int64_t partition_length, - struct ArrowArrayStream* out); - - /* ------------------------------------------------------------------------------------------- - * Writer (driver) - * ------------------------------------------------------------------------------------------- */ - - /* - * Serializes the state of the writer that create_data_writer needs on the executors. Spark - * calls it once, after creating the writer. - * - * set_state: called once with the state, or not at all if there is none - */ - int (*serialize_writer)(void* writer, SparkDsSink set_state, void* sink); - - /* - * Commits a write, when all the tasks succeeded. - * - * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write - * messages, message_lengths: the num_messages messages returned by commit_data_writer of each - * task - */ - int (*commit)(void* writer, int64_t epoch_id, const uint8_t* const* messages, - const int64_t* message_lengths, int64_t num_messages); - - /* - * Aborts a write, when it failed. - * - * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write - * messages, message_lengths: the num_messages messages returned by commit_data_writer of each - * task. A message is NULL if its task did not commit. - */ - int (*abort)(void* writer, int64_t epoch_id, const uint8_t* const* messages, - const int64_t* message_lengths, int64_t num_messages); - - /* Releases a writer handle. */ - void (*close_writer)(void* writer); - - /* ------------------------------------------------------------------------------------------- - * Data writer (executor) - * ------------------------------------------------------------------------------------------- */ - - /* - * Creates the writer for one task. - * - * task: the task, with the partition that it writes, valid until the data writer is closed - * writer_state: the state returned by serialize_writer, empty if there is none - * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write - * data_writer: set to the data writer handle - */ - int (*create_data_writer)(const SparkDsTaskContext* task, const uint8_t* writer_state, - int64_t writer_state_length, int64_t epoch_id, void** data_writer); - - /* - * Writes a batch of data. - * - * array: a struct array whose fields are the columns - * schema: the schema of the array - */ - int (*write)(void* data_writer, struct ArrowArray* array, struct ArrowSchema* schema); - - /* - * Commits the data written by the task. If it fails, Spark calls abort_data_writer. If it - * succeeds, Spark does not call abort_data_writer, even if it then fails the task, for example - * because set_message was called more than once. - * - * set_message: called once with the message to pass to commit on the driver, or not at all for - * an empty message - */ - int (*commit_data_writer)(void* data_writer, SparkDsSink set_message, void* sink); - - /* Aborts the data written by the task. */ - int (*abort_data_writer)(void* data_writer); - - /* Releases a data writer handle, after its commit or abort. */ - void (*close_data_writer)(void* data_writer); -} SparkDsLibrary; +SPARK_DS_EXPORT int spark_ds_open(const SparkDsHost* host, const uint8_t* request, int64_t length, + SparkDsSink reply, void* sink, void** handle); /* - * The function that a library exports. Spark calls it once in each JVM, when it loads the library, - * before any other function. + * Runs a control operation on an object. * - * host: Spark, valid until the JVM exits - * library: set to the function table of the library, for the version of this interface that it - * implements, or the newest one up to host->abi_version if it implements several. Spark - * reads the table after the call, so it is typically the address of a static struct. + * request, length: a CallRequest + * reply: called with a CallReply, or not at all for an empty one + */ +SPARK_DS_EXPORT int spark_ds_call(void* handle, const uint8_t* request, int64_t length, + SparkDsSink reply, void* sink); + +/* + * Exchanges the data of an object, as Arrow C streams. * - * Returns 0 on success, or any other value if the library cannot run, after it logs why with - * host->log. + * request, length: a StreamRequest + * in: for a Write, the stream of the batches to write. It is NULL otherwise. + * out: for a Read, the stream to initialize with the batches of the partition: each array of the + * stream is a struct array whose fields are the columns to read, those of the schema to read + * or the pruned columns. It is NULL otherwise. */ -SPARK_DS_EXPORT int spark_ds_init(const SparkDsHost* host, const SparkDsLibrary** library); +SPARK_DS_EXPORT int spark_ds_stream(void* handle, const uint8_t* request, int64_t length, + struct ArrowArrayStream* in, struct ArrowArrayStream* out); + +/* Releases the handle of an object. */ +SPARK_DS_EXPORT void spark_ds_close(void* handle); #ifdef __cplusplus } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala index 8759c2894ead9..8721301852c73 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala @@ -18,7 +18,10 @@ package org.apache.spark.sql.execution.datasources.v2.columnar import scala.jdk.CollectionConverters._ +import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.arrow.vector.types.pojo.Schema import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.connector.metric.CustomTaskMetric @@ -185,6 +188,15 @@ class ArrowBatchingDataWriter(delegate: DataWriter[ColumnarBatch], options: Arro private var arrowWriter: ArrowWriter = _ private var numRows = 0 + override def writeAll(records: java.util.Iterator[InternalRow]): Unit = delegate match { + case streaming: ArrowStreamDataWriter => + // The rows that `write` got before, if any, are written first. + flush() + streaming.writeStream(new ArrowRowReader(records.asScala, arrowSchema, allocator, options)) + case _ => + super.writeAll(records) + } + override def write(record: InternalRow): Unit = { if (root == null) { root = VectorSchemaRoot.create(arrowSchema, allocator) @@ -241,3 +253,77 @@ class ArrowBatchingDataWriter(delegate: DataWriter[ColumnarBatch], options: Arro override def currentMetricsValues(): Array[CustomTaskMetric] = delegate.currentMetricsValues() } + +/** + * A data writer of a [[ColumnarDataSource]] that can also write all the data of a task from an + * Arrow reader, which it reads on the calling thread, during the call. + */ +trait ArrowStreamDataWriter extends DataWriter[ColumnarBatch] { + /** Writes all the batches of `reader`. */ + def writeStream(reader: ArrowReader): Unit +} + +/** + * Reads rows as Arrow batches, which end with the limits of [[ArrowBatchingOptions]]. Each batch + * is written to new buffers, so a native writer that took ownership of the arrays of a batch can + * keep them after the next batch. The first error of the rows is kept, see + * [[ArrowRowReader.failure]]. + */ +class ArrowRowReader( + rows: Iterator[InternalRow], + schema: Schema, + allocator: BufferAllocator, + options: ArrowBatchingOptions) + extends ArrowReader(allocator) { + + @volatile private var error: Throwable = _ + + override def loadNextBatch(): Boolean = { + try { + val root = getVectorSchemaRoot + // Releases the buffers of the previous batch, which its exported arrays may still hold. + root.clear() + if (!rows.hasNext) { + false + } else { + val writer = ArrowWriter.create(root) + var numRows = 0 + var full = false + while (!full && rows.hasNext) { + writer.write(rows.next()) + numRows += 1 + // The size is not computed when the limit is the largest possible one, as in Python + // UDFs. + full = numRows >= options.maxRecordsPerBatch || + (options.maxBytesPerBatch < Int.MaxValue && + writer.sizeInBytes() >= options.maxBytesPerBatch) + } + writer.finish() + true + } + } catch { + case e: Throwable => + if (error == null) { + error = e + } + throw e + } + } + + override def bytesRead(): Long = 0L + + override protected def closeReadSource(): Unit = {} + + override protected def readSchema(): Schema = schema +} + +object ArrowRowReader { + /** + * The error of the rows of `reader`, if it is an [[ArrowRowReader]] that failed. An Arrow C + * stream that exports the reader only reports the message of the error. + */ + def failure(reader: ArrowReader): Option[Throwable] = reader match { + case rows: ArrowRowReader => Option(rows.error) + case _ => None + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala index 3edb430366e8e..1d48f409f181c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -26,20 +26,22 @@ import scala.util.control.NonFatal import com.fasterxml.jackson.core.JsonProcessingException import com.fasterxml.jackson.databind.DeserializationFeature import com.fasterxml.jackson.databind.json.JsonMapper +import com.google.protobuf.ByteString import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, CDataDictionaryProvider, Data} import org.apache.arrow.vector.{FieldVector, ValueVector, VectorSchemaRoot} import org.apache.arrow.vector.complex.MapVector +import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} -import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException, TaskContext} +import org.apache.spark.{SPARK_VERSION, SparkException, SparkThrowable, SparkUnsupportedOperationException, TaskContext} import org.apache.spark.internal.Logging +import org.apache.spark.sql.catalyst.util.CollationFactory import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} -import org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarDataSource, ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} -import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{DataType, StructField, StructType} +import org.apache.spark.sql.execution.datasources.v2.columnar.{ArrowRowReader, ArrowStreamDataWriter, ColumnarDataSource, ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} +import org.apache.spark.sql.types._ import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} import org.apache.spark.unsafe.Platform @@ -47,7 +49,8 @@ import org.apache.spark.util.Utils /** * A [[ColumnarDataSource]] implemented by a native library, through the C interface in - * `org/apache/spark/sql/datasource/spark_datasource.h`. + * `org/apache/spark/sql/datasource/spark_datasource.h` and the messages of + * `org/apache/spark/sql/datasource/spark_datasource.proto`. * * @param location where the library is on the driver. The readers and the writers distribute it * to the executors of the active session when they are created. @@ -61,48 +64,58 @@ class NativeDataSource( extends ColumnarDataSource { import NativeCalls._ + // The data source, as the requests of spark_ds_open carry it. The keys of a case-insensitive + // map are in lower case. + private def dataSource: proto.DataSource = { + val builder = proto.DataSource.newBuilder().setName(name) + options.entrySet().asScala.foreach(e => builder.putOptions(e.getKey, e.getValue)) + builder.build() + } + override def schema(): StructType = { val library = NativeLibraries.get(location) - withDataSource(library, "infer the schema of") { dataSource => - call(name, "infer the schema of") { - optional(NativeArrow.importSchema(address => library.schema(dataSource.handle, address))) - }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) + withDataSource(library, "infer the schema of") { (handle, _) => + val reply = call(name, "infer the schema of") { + library.call(handle, NativeMessages.InferSchema) + } + if (!reply.getInferSchema.hasSchema) { + throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name) + } + call(name, "infer the schema of")(NativeSchemas.fromProto(reply.getInferSchema.getSchema)) } } override def reader(schema: StructType): ColumnarReader = { val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) - val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ, - "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) { - dataSource => - NativeArrow.exportSchema(schema) { address => - library.createReader(dataSource.handle, address) - } + val handle = open(library, "plan a scan of", proto.Capability.CAPABILITY_BATCH_READ, + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { + _.setBatchReader(proto.BatchReader.newBuilder() + .setDataSource(dataSource) + .setSchema(NativeSchemas.toProto(schema))) } new NativeDataSourceReader( - executorLocation, name, schema, new NativeHandle(handle, library.closeReader)) + executorLocation, name, schema, new NativeHandle(handle, library.close)) } override def streamReader(schema: StructType): ColumnarStreamReader = { val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) - val handle = create(library, "plan a scan of", NativeLibrary.MICRO_BATCH_READ, - "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) { - dataSource => - NativeArrow.exportSchema(schema) { address => - library.createStreamReader(dataSource.handle, address) - } + val handle = open(library, "plan a scan of", proto.Capability.CAPABILITY_MICRO_BATCH_READ, + "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { + _.setStreamReader(proto.StreamReader.newBuilder() + .setDataSource(dataSource) + .setSchema(NativeSchemas.toProto(schema))) } new NativeDataSourceStreamReader( - executorLocation, name, schema, new NativeHandle(handle, library.closeStreamReader)) + executorLocation, name, schema, new NativeHandle(handle, library.close)) } override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) - val handle = new NativeHandle( - createWriter(library, schema, overwrite, queryId = null), library.closeWriter) + val handle = + new NativeHandle(openWriter(library, schema, overwrite, queryId = null), library.close) new NativeDataSourceWriter(executorLocation, name, handle, serializeWriter(library, handle)) } @@ -112,87 +125,87 @@ class NativeDataSource( queryId: String): ColumnarStreamWriter = { val executorLocation = NativeDataSourceRegistry.distribute(location) val library = NativeLibraries.get(location) - val handle = - new NativeHandle(createWriter(library, schema, overwrite, queryId), library.closeWriter) + val handle = new NativeHandle(openWriter(library, schema, overwrite, queryId), library.close) val writer = new NativeDataSourceStreamWriter( executorLocation, name, handle, serializeWriter(library, handle)) NativeStreamWriters.register(queryId, handle) writer } - /** Runs `f` with a new native data source, which is closed afterwards. */ + /** + * Runs `f` with a new native data source and its capabilities, and closes the data source + * afterwards. + */ private def withDataSource[T](library: NativeLibrary, action: String)( - f: NativeLibrary.DataSourceHandle => T): T = { - // Checked first, so that a data source is never created without a way to release it. - requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE)) - // The keys of a case-insensitive map are in lower case. - val entries = options.entrySet().asScala.toArray - val dataSource = call(name, "create") { - library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue)) + f: (Long, Set[proto.Capability]) => T): T = { + val opened = call(name, "create") { + library.open(NativeMessages.open().setDataSource(dataSource).build()) } - Utils.tryWithSafeFinally(f(dataSource)) { - call(name, action)(library.closeDataSource(dataSource.handle)) + Utils.tryWithSafeFinally { + f(opened.handle, opened.reply.getCapabilitiesList.asScala.toSet) + } { + call(name, action)(library.close(opened.handle)) } } - /** - * Creates a reader or a writer with `f`, which fails if the data source does not declare the - * capability, or if the library does not implement the functions that it requires. - */ - private def create( + /** Returns the capabilities of the data source, and fails if it does not declare `capability`. */ + private def capabilities( library: NativeLibrary, action: String, - capability: Long, - unsupportedError: String, - functions: Seq[NativeLibrary.Function])( - f: NativeLibrary.DataSourceHandle => Long): Long = { - withDataSource(library, action) { dataSource => - if ((dataSource.capabilities & capability) == 0) { - throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) - } - requireFunctions(library, action, functions) - call(name, action)(f(dataSource)) + capability: proto.Capability, + unsupportedError: String): Set[proto.Capability] = { + val capabilities = withDataSource(library, action)((_, capabilities) => capabilities) + if (!capabilities.contains(capability)) { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) } + capabilities } - /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */ - private def createWriter( + /** Opens a reader with the request that `f` builds, if the data source declares `capability`. */ + private def open( + library: NativeLibrary, + action: String, + capability: proto.Capability, + unsupportedError: String)( + f: proto.OpenRequest.Builder => proto.OpenRequest.Builder): Long = { + capabilities(library, action, capability, unsupportedError) + call(name, action)(library.open(f(NativeMessages.open()).build()).handle) + } + + /** Opens a writer for a batch write, or for a streaming write with the ID of its query. */ + private def openWriter( library: NativeLibrary, schema: StructType, overwrite: Boolean, queryId: String): Long = { val action = "plan a write to" val (capability, unsupportedError, overwriteMode) = if (queryId == null) { - (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", + (proto.Capability.CAPABILITY_BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", "the save mode overwrite") } else { - (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", + (proto.Capability.CAPABILITY_STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", "the output mode complete") } - create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) { - dataSource => - if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) { - throw error(name, action, s"The data source does not support $overwriteMode.") - } - val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND - NativeArrow.exportSchema(schema) { address => - library.createWriter(dataSource.handle, address, mode, queryId) - } + val declared = capabilities(library, action, capability, unsupportedError) + if (overwrite && !declared.contains(proto.Capability.CAPABILITY_TRUNCATE)) { + throw error(name, action, s"The data source does not support $overwriteMode.") } - } - - private def requireFunctions( - library: NativeLibrary, - action: String, - functions: Seq[NativeLibrary.Function]): Unit = { - functions.find(!library.hasFunction(_)).foreach { function => - throw error(name, action, s"The library does not implement ${function.fieldName}.") + val writer = proto.Writer.newBuilder() + .setDataSource(dataSource) + .setSchema(NativeSchemas.toProto(schema)) + .setMode( + if (overwrite) proto.WriteMode.WRITE_MODE_TRUNCATE else proto.WriteMode.WRITE_MODE_APPEND) + if (queryId != null) { + writer.setQueryId(queryId) } + call(name, action)(library.open(NativeMessages.open().setWriter(writer).build()).handle) } private def serializeWriter(library: NativeLibrary, handle: NativeHandle): Array[Byte] = { try { - bytesOrEmpty(call(name, "plan a write to")(optional(library.serializeWriter(handle.get)))) + call(name, "plan a write to") { + library.call(handle.get, NativeMessages.Serialize).getSerialize.getState.toByteArray + } } catch { case e: Throwable => handle.close() @@ -201,18 +214,6 @@ class NativeDataSource( } } -object NativeDataSource { - import NativeLibrary.Function._ - - // The functions that the operations of a data source require, besides create and close. The - // others are optional. They are checked on the driver, before the operation starts. - private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ) - private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET, - STREAM_PARTITIONS, CLOSE_STREAM_READER, READ) - private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER, - WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER) -} - /** Reads a native data source for a batch scan. */ class NativeDataSourceReader( location: NativeLibraryLocation, @@ -240,27 +241,40 @@ class NativeDataSourceReader( val used = if (encoded.isEmpty) { Set.empty[Int] } else { + val request = proto.CallRequest.newBuilder().setPushPredicates( + proto.PushPredicates.newBuilder().addAllPredicatesJson(encoded.map(_._1).toSeq.asJava)) + val uses = call(name, "plan a scan of") { + library.call(handle.get, request.build()).getPushPredicates.getUsesValueList.asScala + } + if (uses.length > encoded.length) { + throw error(name, "plan a scan of", s"PushPredicates returned ${uses.length} answers " + + s"for ${encoded.length} predicates.") + } // How the reader uses each predicate: whether it uses it completely or only to skip data - // does not matter, since Spark evaluates all the predicates again after reading. - call(name, "plan a scan of") { - optional(library.pushPredicates(handle.get, encoded.map(_._1))) - }.map { uses => - encoded.map(_._2).zip(uses).collect { - case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i - }.toSet - }.getOrElse(Set.empty[Int]) + // does not matter, since Spark evaluates all the predicates again after reading. The + // missing answers are PREDICATE_USE_UNUSED, and Spark takes the values that it does not + // know as PREDICATE_USE_PARTIAL. + encoded.map(_._2).zip(uses).collect { + case (i, use) if use != proto.PredicateUse.PREDICATE_USE_UNUSED_VALUE => i + }.toSet } predicates.zipWithIndex.collect { case (predicate, i) if used.contains(i) => predicate } } override def pushLimit(limit: Int): Boolean = { - call(name, "plan a scan of")(optional(library.pushLimit(handle.get, limit))).getOrElse(false) + val request = proto.CallRequest.newBuilder() + .setPushLimit(proto.PushLimit.newBuilder().setLimit(limit)) + call(name, "plan a scan of") { + library.call(handle.get, request.build()).getPushLimit.getApplied + } } override def pruneColumns(requiredSchema: StructType): Boolean = { + val request = proto.CallRequest.newBuilder().setPruneColumns( + proto.PruneColumns.newBuilder().addAllColumns(requiredSchema.fieldNames.toSeq.asJava)) val pruned = call(name, "plan a scan of") { - optional(library.pruneColumns(handle.get, requiredSchema.fieldNames)) - }.getOrElse(false) + library.call(handle.get, request.build()).getPruneColumns.getPruned + } if (pruned) { readSchema = requiredSchema } @@ -273,9 +287,13 @@ class NativeDataSourceReader( throw planningError } try { - val partitions = call(name, "plan a scan of")(library.partitions(handle.get)) - readerState = bytesOrEmpty( - call(name, "plan a scan of")(optional(library.serializeReader(handle.get)))) + val partitions = call(name, "plan a scan of") { + library.call(handle.get, NativeMessages.PlanPartitions).getPlanPartitions + .getPartitionsList.asScala.map(_.toByteArray).toArray + } + readerState = call(name, "plan a scan of") { + library.call(handle.get, NativeMessages.Serialize).getSerialize.getState.toByteArray + } plannedPartitions = NativeInputPartition.from(partitions) } catch { case e: Throwable => @@ -311,20 +329,27 @@ class NativeDataSourceStreamReader( @transient private lazy val library = NativeLibraries.get(location) override def initialOffset(): String = { - checkOffset("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get))) + checkOffset("InitialOffset", call(name, "plan a scan of") { + library.call(handle.get, NativeMessages.InitialOffset).getInitialOffset.getOffsetJson + }) } override def latestOffset(): String = { - checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get))) + checkOffset("LatestOffset", call(name, "plan a scan of") { + library.call(handle.get, NativeMessages.LatestOffset).getLatestOffset.getOffsetJson + }) } /** - * Checks an offset returned by the given function of the library: it must be JSON on a single - * line, because the offset log of the checkpoint stores each offset as a line. + * Checks an offset returned by the given request: it must be JSON on a single line, because the + * offset log of the checkpoint stores each offset as a line. */ - private def checkOffset(function: String, offset: String): String = { + private def checkOffset(request: String, offset: String): String = { + if (offset.isEmpty) { + throw error(name, "plan a scan of", s"$request returned no offset.") + } if (offset.exists(c => c == '\n' || c == '\r')) { - throw error(name, "plan a scan of", s"$function returned an offset with a line break, " + + throw error(name, "plan a scan of", s"$request returned an offset with a line break, " + "but an offset must be JSON on a single line.") } val valid = try { @@ -334,18 +359,22 @@ class NativeDataSourceStreamReader( case _: JsonProcessingException => false } if (!valid) { - throw error(name, "plan a scan of", s"$function returned an offset that is not valid " + + throw error(name, "plan a scan of", s"$request returned an offset that is not valid " + s"JSON: $offset") } offset } override def partitions(start: String, end: String): Array[InputPartition] = { + val request = proto.CallRequest.newBuilder().setPlanStreamPartitions( + proto.PlanStreamPartitions.newBuilder().setStartJson(start).setEndJson(end)) val partitions = call(name, "plan a scan of") { - library.streamPartitions(handle.get, start, end) + library.call(handle.get, request.build()).getPlanStreamPartitions.getPartitionsList.asScala + .map(_.toByteArray).toArray + } + readerState = call(name, "plan a scan of") { + library.call(handle.get, NativeMessages.Serialize).getSerialize.getState.toByteArray } - readerState = bytesOrEmpty( - call(name, "plan a scan of")(optional(library.serializeStreamReader(handle.get)))) NativeInputPartition.from(partitions) } @@ -355,7 +384,9 @@ class NativeDataSourceStreamReader( } override def commit(end: String): Unit = { - call(name, "commit an offset of")(optional(library.commitOffset(handle.get, end))) + val request = proto.CallRequest.newBuilder() + .setCommitOffset(proto.CommitOffset.newBuilder().setEndJson(end)) + call(name, "commit an offset of")(library.call(handle.get, request.build())) } override def stop(): Unit = handle.close() @@ -384,7 +415,7 @@ class NativeDataSourceWriter( override def commit(messages: Array[WriterCommitMessage]): Unit = { call(name, "commit a write to") { - library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + library.call(handle.get, NativeMessages.commit(-1L, messages)) } // Not closed if the commit fails: Spark then aborts the write, which closes it. handle.close() @@ -393,7 +424,7 @@ class NativeDataSourceWriter( override def abort(messages: Array[WriterCommitMessage]): Unit = { Utils.tryWithSafeFinally { call(name, "abort a write to") { - library.abort(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + library.call(handle.get, NativeMessages.abort(-1L, messages)) } }(handle.close()) } @@ -422,13 +453,13 @@ class NativeDataSourceStreamWriter( override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { call(name, "commit a write to") { - library.commit(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + library.call(handle.get, NativeMessages.commit(epochId, messages)) } } override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { call(name, "abort a write to") { - library.abort(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + library.call(handle.get, NativeMessages.abort(epochId, messages)) } } } @@ -448,11 +479,18 @@ object NativeInputPartition { case class NativeWriterCommitMessage(bytes: Array[Byte]) extends WriterCommitMessage object NativeWriterCommitMessage { - def toBytes(messages: Array[WriterCommitMessage]): Array[Array[Byte]] = messages.map { - case message: NativeWriterCommitMessage => message.bytes - case null => null - case other => - throw SparkException.internalError(s"Unexpected commit message: ${other.getClass.getName}") + /** The messages of the tasks, as the requests of a commit or an abort carry them. */ + def toTaskMessages(messages: Array[WriterCommitMessage]): Seq[proto.TaskMessage] = { + messages.toSeq.map { + case message: NativeWriterCommitMessage => + proto.TaskMessage.newBuilder() + .setCommitted(true) + .setMessage(ByteString.copyFrom(message.bytes)) + .build() + case null => proto.TaskMessage.getDefaultInstance + case other => + throw SparkException.internalError(s"Unexpected commit message: ${other.getClass.getName}") + } } } @@ -474,8 +512,11 @@ class NativePartitionReader( // Read here rather than with `Data.importArrayStream`, to check each array before importing it. private val stream = ArrowArrayStream.allocateNew(allocator) private val dictionaries = new CDataDictionaryProvider() - // The task that reads the partition, which the library may use until the stream is released. - private var task: NativeLibrary.Task = _ + private var library: NativeLibrary = _ + // The partition reader of the library, which is closed after its stream is released, and the + // registration of the task for SparkDsHost.is_interrupted, until then. + private var partitionReader: Long = 0L + private var task: AutoCloseable = _ private var root: VectorSchemaRoot = _ private var batch: ColumnarBatch = _ @@ -491,9 +532,19 @@ class NativePartitionReader( } private def openStream(): Unit = { - val library = NativeLibraries.get(location) - task = NativeTask.create() - call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress())) + library = NativeLibraries.get(location) + val (taskContext, registration) = NativeTask.create() + task = registration + val request = proto.PartitionReader.newBuilder() + .setTask(taskContext) + .setReaderState(ByteString.copyFrom(readerState)) + .setPartition(ByteString.copyFrom(partition)) + partitionReader = call(name, "read") { + library.open(NativeMessages.open().setPartitionReader(request).build()).handle + } + call(name, "read") { + library.stream(partitionReader, NativeMessages.Read, 0L, stream.memoryAddress()) + } } private def importSchema(): Schema = { @@ -693,10 +744,16 @@ class NativePartitionReader( stream.release() stream.close() } { - if (task != null) { - task.close() + Utils.tryWithSafeFinally { + if (partitionReader != 0L) { + call(name, "read")(library.close(partitionReader)) + } + } { + if (task != null) { + task.close() + } + Utils.tryLogNonFatalError(allocator.close()) } - Utils.tryLogNonFatalError(allocator.close()) } } } @@ -712,7 +769,7 @@ private class ReusedArrowColumnVector(vector: ValueVector) extends ArrowColumnVe override def closeIfFreeable(): Unit = {} } -/** Writes the batches of a task to the data writer of a native library. */ +/** Writes the data of a task to the data writer of a native library. */ class NativeDataWriter( location: NativeLibraryLocation, name: String, @@ -720,14 +777,20 @@ class NativeDataWriter( partitionId: Int, taskId: Long, epochId: Long) - extends DataWriter[ColumnarBatch] { + extends ArrowStreamDataWriter { import NativeCalls._ private val library = NativeLibraries.get(location) - // The task, which the library may use until the data writer is closed. - private val task = NativeTask.create(Some(partitionId), Some(taskId)) + // The task, registered for SparkDsHost.is_interrupted until the data writer is closed. + private val (taskContext, task) = NativeTask.create(Some(partitionId), Some(taskId)) private val handle: Long = try { - call(name, "write to")(library.createDataWriter(task, writerState, epochId)) + val request = proto.DataWriter.newBuilder() + .setTask(taskContext) + .setWriterState(ByteString.copyFrom(writerState)) + .setEpochId(epochId) + call(name, "write to") { + library.open(NativeMessages.open().setDataWriter(request).build()).handle + } } catch { case e: Throwable => task.close() @@ -735,14 +798,18 @@ class NativeDataWriter( } // Whether the data writer was committed or aborted. It stays false if the library fails to // commit, so that Spark aborts the data writer, but becomes true as soon as the library commits, - // even if Spark then fails the commit, for example because of an invalid message. + // even if Spark then fails the commit, for example because of an invalid reply. private var finished = false private var closed = false - override def write(batch: ColumnarBatch): Unit = { + private def checkOpen(): Unit = { if (finished || closed) { throw SparkException.internalError("The data writer was already committed or aborted.") } + } + + override def write(batch: ColumnarBatch): Unit = { + checkOpen() val vectors = (0 until batch.numCols()).map { i => batch.column(i) match { case vector: ArrowColumnVector => vector.getValueVector.asInstanceOf[FieldVector] @@ -750,36 +817,43 @@ class NativeDataWriter( s"Expected an Arrow column vector, but got ${vector.getClass.getName}.") } } - val root = new VectorSchemaRoot( - vectors.map(_.getField).asJava, vectors.asJava, batch.numRows()) - val array = ArrowArray.allocateNew(ArrowUtils.rootAllocator) - val arrowSchema = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) + writeStream(new SingleBatchArrowReader( + new VectorSchemaRoot(vectors.map(_.getField).asJava, vectors.asJava, batch.numRows()))) + } + + override def writeStream(reader: ArrowReader): Unit = { + checkOpen() + val stream = ArrowArrayStream.allocateNew(ArrowUtils.rootAllocator) try { - Data.exportVectorSchemaRoot(ArrowUtils.rootAllocator, root, null, array, arrowSchema) - call(name, "write to") { - library.write(handle, array.memoryAddress(), arrowSchema.memoryAddress()) + Data.exportArrayStream(ArrowUtils.rootAllocator, reader, stream) + try { + call(name, "write to") { + library.stream(handle, NativeMessages.Write, stream.memoryAddress(), 0L) + } + } catch { + // The library reads the batches through the stream, which only reports the message of + // an error of the rows, such as a failure of an upstream operator: that error is thrown. + case e: Throwable => throw ArrowRowReader.failure(reader).getOrElse(e) } } finally { - // Releases the batch unless the library moved it to take ownership. - array.release() - array.close() - arrowSchema.release() - arrowSchema.close() + // Releases the stream, and so the reader, unless the library moved it to take ownership. + stream.release() + stream.close() } } override def commit(): WriterCommitMessage = { - if (finished || closed) { - throw SparkException.internalError("The data writer was already committed or aborted.") + checkOpen() + val reply = call(name, "write to") { + library.call(handle, NativeMessages.CommitDataWriter, () => finished = true) } - val message = call(name, "write to")(library.commitDataWriter(handle, () => finished = true)) - NativeWriterCommitMessage(bytesOrEmpty(Some(message))) + NativeWriterCommitMessage(reply.getCommitDataWriter.getMessage.toByteArray) } override def abort(): Unit = { if (!finished && !closed) { finished = true - call(name, "write to")(library.abortDataWriter(handle)) + call(name, "write to")(library.call(handle, NativeMessages.AbortDataWriter)) } } @@ -787,30 +861,64 @@ class NativeDataWriter( if (!closed) { closed = true Utils.tryWithSafeFinally { - call(name, "write to")(library.closeDataWriter(handle)) + call(name, "write to")(library.close(handle)) }(task.close()) } } } -/** Creates the `SparkDsTaskContext` of the running task. */ +/** + * Exports one batch of a [[VectorSchemaRoot]] as an Arrow reader. The vectors of the batch are + * transferred to the reader, which releases them when it is closed. + */ +private class SingleBatchArrowReader(batch: VectorSchemaRoot) + extends ArrowReader(ArrowUtils.rootAllocator) { + private var loaded = false + + override def loadNextBatch(): Boolean = { + if (loaded) { + false + } else { + loaded = true + val root = getVectorSchemaRoot + batch.getFieldVectors.asScala.zip(root.getFieldVectors.asScala).foreach { + case (from, to) => from.makeTransferPair(to).transfer() + } + root.setRowCount(batch.getRowCount) + true + } + } + + override def bytesRead(): Long = 0L + + override protected def closeReadSource(): Unit = {} + + override protected def readSchema(): Schema = batch.getSchema +} + +/** Creates the `TaskContext` message of the running task. */ private[ffi] object NativeTask { /** + * Returns the message, and the registration of the task for `SparkDsHost.is_interrupted`, which + * the caller closes once the library no longer uses the object of the task. + * * @param partitionId the partition that the task reads or writes, by default the partition of * the task * @param taskAttemptId the ID of the task attempt, by default the one of the task */ def create( partitionId: Option[Int] = None, - taskAttemptId: Option[Long] = None): NativeLibrary.Task = { + taskAttemptId: Option[Long] = None): (proto.TaskContext, AutoCloseable) = { val context = Option(TaskContext.get()) - new NativeLibrary.Task( - context.map(_.stageId()).getOrElse(-1), - context.map(_.stageAttemptNumber()).getOrElse(-1), - partitionId.orElse(context.map(_.partitionId())).getOrElse(-1), - context.map(_.attemptNumber()).getOrElse(0), - taskAttemptId.orElse(context.map(_.taskAttemptId())).getOrElse(-1L), - () => context.exists(_.isInterrupted())) + val attemptId = taskAttemptId.orElse(context.map(_.taskAttemptId())).getOrElse(-1L) + val message = proto.TaskContext.newBuilder() + .setStageId(context.map(_.stageId()).getOrElse(-1)) + .setStageAttemptNumber(context.map(_.stageAttemptNumber()).getOrElse(-1)) + .setPartitionId(partitionId.orElse(context.map(_.partitionId())).getOrElse(-1)) + .setAttemptNumber(context.map(_.attemptNumber()).getOrElse(0)) + .setTaskAttemptId(attemptId) + .build() + (message, NativeLibrary.registerTask(attemptId, () => context.exists(_.isInterrupted()))) } } @@ -885,9 +993,6 @@ object NativeCalls { f } catch { case e: Throwable with SparkThrowable => throw e - case e: NativeLibrary.MissingFunctionException => - throw QueryExecutionErrors.nativeDataSourceError( - action, name, s"The library does not implement ${e.getMessage}.", e) case e @ (NonFatal(_) | _: LinkageError) => throw QueryExecutionErrors.nativeDataSourceError(action, name, message(e), e) } @@ -903,57 +1008,154 @@ object NativeCalls { } } - /** Calls a function that the library may not implement. Returns None if it does not. */ - def optional[T](f: => T): Option[T] = { - try { - Some(f) - } catch { - case _: NativeLibrary.MissingFunctionException => None - } - } - /** Returns a NATIVE_DATA_SOURCE_ERROR for an error detected by Spark. */ def error(name: String, action: String, msg: String): Throwable = { QueryExecutionErrors.nativeDataSourceError(action, name, msg, null) } +} + +/** The messages of `spark_datasource.proto` that Spark sends. */ +private[ffi] object NativeMessages { + /** An `OpenRequest`, with the versions of Spark and of the messages. */ + def open(): proto.OpenRequest.Builder = proto.OpenRequest.newBuilder() + .setSparkVersion(SPARK_VERSION) + .setProtocolVersion(NativeLibrary.PROTOCOL_VERSION) + + val InferSchema: proto.CallRequest = proto.CallRequest.newBuilder() + .setInferSchema(proto.InferSchema.getDefaultInstance).build() + val PlanPartitions: proto.CallRequest = proto.CallRequest.newBuilder() + .setPlanPartitions(proto.PlanPartitions.getDefaultInstance).build() + val Serialize: proto.CallRequest = proto.CallRequest.newBuilder() + .setSerialize(proto.Serialize.getDefaultInstance).build() + val InitialOffset: proto.CallRequest = proto.CallRequest.newBuilder() + .setInitialOffset(proto.InitialOffset.getDefaultInstance).build() + val LatestOffset: proto.CallRequest = proto.CallRequest.newBuilder() + .setLatestOffset(proto.LatestOffset.getDefaultInstance).build() + val CommitDataWriter: proto.CallRequest = proto.CallRequest.newBuilder() + .setCommitDataWriter(proto.CommitDataWriter.getDefaultInstance).build() + val AbortDataWriter: proto.CallRequest = proto.CallRequest.newBuilder() + .setAbortDataWriter(proto.AbortDataWriter.getDefaultInstance).build() + val Read: proto.StreamRequest = + proto.StreamRequest.newBuilder().setRead(proto.Read.getDefaultInstance).build() + val Write: proto.StreamRequest = + proto.StreamRequest.newBuilder().setWrite(proto.Write.getDefaultInstance).build() + + def commit(epochId: Long, messages: Array[WriterCommitMessage]): proto.CallRequest = { + proto.CallRequest.newBuilder().setCommit(proto.Commit.newBuilder() + .setEpochId(epochId) + .addAllMessages(NativeWriterCommitMessage.toTaskMessages(messages).asJava)).build() + } - def bytesOrEmpty(bytes: Option[Array[Byte]]): Array[Byte] = { - bytes.flatMap(Option(_)).getOrElse(Array.emptyByteArray) + def abort(epochId: Long, messages: Array[WriterCommitMessage]): proto.CallRequest = { + proto.CallRequest.newBuilder().setAbort(proto.Abort.newBuilder() + .setEpochId(epochId) + .addAllMessages(NativeWriterCommitMessage.toTaskMessages(messages).asJava)).build() } } -/** Exchanges schemas with native libraries through the Arrow C data interface. */ -object NativeArrow { - /** Exports the schema as an `ArrowSchema`, and calls `f` with its address. */ - def exportSchema[T](schema: StructType)(f: Long => T): T = { - val arrowSchema = ArrowUtils.toArrowSchema( - schema, SQLConf.get.sessionLocalTimeZone, errorOnDuplicatedFieldNames = true, - largeVarTypes = false) - val struct = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) - try { - Data.exportSchema(ArrowUtils.rootAllocator, arrowSchema, null, struct) - f(struct.memoryAddress()) - } finally { - // Releases the schema unless the library moved it to take ownership. - struct.release() - struct.close() +/** Converts schemas to and from the Spark SQL types of `spark_datasource.proto`. */ +private[ffi] object NativeSchemas { + import proto.{DataType => P} + + def toProto(schema: StructType): proto.StructType = { + val builder = proto.StructType.newBuilder() + schema.fields.foreach { field => + builder.addFields(proto.StructField.newBuilder() + .setName(field.name) + .setDataType(toProto(field.dataType)) + .setNullable(field.nullable)) } + builder.build() } - /** Calls `f` with the address of an `ArrowSchema` to initialize, and imports it. */ - def importSchema(f: Long => Unit): StructType = { - val struct = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) - var imported = false - try { - f(struct.memoryAddress()) - imported = true - // Releases and closes the struct. - ArrowUtils.fromArrowSchema(Data.importSchema(ArrowUtils.rootAllocator, struct, null)) - } finally { - if (!imported) { - struct.release() - struct.close() - } + def toProto(dataType: DataType): proto.DataType = { + val builder = proto.DataType.newBuilder() + dataType match { + case NullType => builder.setNull(P.Null.getDefaultInstance) + case BooleanType => builder.setBoolean(P.Boolean.getDefaultInstance) + case ByteType => builder.setByte(P.Byte.getDefaultInstance) + case ShortType => builder.setShort(P.Short.getDefaultInstance) + case IntegerType => builder.setInteger(P.Integer.getDefaultInstance) + case LongType => builder.setLong(P.Long.getDefaultInstance) + case FloatType => builder.setFloat(P.Float.getDefaultInstance) + case DoubleType => builder.setDouble(P.Double.getDefaultInstance) + case d: DecimalType => + builder.setDecimal(P.Decimal.newBuilder().setPrecision(d.precision).setScale(d.scale)) + case s: StringType => + val collation = if (s.isUTF8BinaryCollation) { + "" + } else { + CollationFactory.fetchCollation(s.collationId).collationName + } + builder.setString(P.String.newBuilder().setCollation(collation)) + case BinaryType => builder.setBinary(P.Binary.getDefaultInstance) + case DateType => builder.setDate(P.Date.getDefaultInstance) + case TimestampType => builder.setTimestamp(P.Timestamp.getDefaultInstance) + case TimestampNTZType => builder.setTimestampNtz(P.TimestampNtz.getDefaultInstance) + case y: YearMonthIntervalType => + builder.setYearMonthInterval( + P.YearMonthInterval.newBuilder().setStartField(y.startField).setEndField(y.endField)) + case d: DayTimeIntervalType => + builder.setDayTimeInterval( + P.DayTimeInterval.newBuilder().setStartField(d.startField).setEndField(d.endField)) + case CalendarIntervalType => + builder.setCalendarInterval(P.CalendarInterval.getDefaultInstance) + case a: ArrayType => + builder.setArray(P.Array.newBuilder() + .setElementType(toProto(a.elementType)) + .setContainsNull(a.containsNull)) + case m: MapType => + builder.setMap(P.Map.newBuilder() + .setKeyType(toProto(m.keyType)) + .setValueType(toProto(m.valueType)) + .setValueContainsNull(m.valueContainsNull)) + case s: StructType => builder.setStruct(toProto(s)) + case udt: UserDefinedType[_] => return toProto(udt.sqlType) + case other => + throw new IllegalArgumentException( + s"Native data sources do not support the type ${other.catalogString}.") } + builder.build() + } + + def fromProto(schema: proto.StructType): StructType = { + StructType(schema.getFieldsList.asScala.map { field => + StructField(field.getName, fromProto(field.getDataType), field.getNullable) + }.toSeq) + } + + def fromProto(dataType: proto.DataType): DataType = dataType.getKindCase match { + case P.KindCase.NULL => NullType + case P.KindCase.BOOLEAN => BooleanType + case P.KindCase.BYTE => ByteType + case P.KindCase.SHORT => ShortType + case P.KindCase.INTEGER => IntegerType + case P.KindCase.LONG => LongType + case P.KindCase.FLOAT => FloatType + case P.KindCase.DOUBLE => DoubleType + case P.KindCase.DECIMAL => + DecimalType(dataType.getDecimal.getPrecision, dataType.getDecimal.getScale) + case P.KindCase.STRING => + val collation = dataType.getString.getCollation + if (collation.isEmpty) StringType else StringType(collation) + case P.KindCase.BINARY => BinaryType + case P.KindCase.DATE => DateType + case P.KindCase.TIMESTAMP => TimestampType + case P.KindCase.TIMESTAMP_NTZ => TimestampNTZType + case P.KindCase.YEAR_MONTH_INTERVAL => + val interval = dataType.getYearMonthInterval + YearMonthIntervalType(interval.getStartField.toByte, interval.getEndField.toByte) + case P.KindCase.DAY_TIME_INTERVAL => + val interval = dataType.getDayTimeInterval + DayTimeIntervalType(interval.getStartField.toByte, interval.getEndField.toByte) + case P.KindCase.CALENDAR_INTERVAL => CalendarIntervalType + case P.KindCase.ARRAY => + ArrayType(fromProto(dataType.getArray.getElementType), dataType.getArray.getContainsNull) + case P.KindCase.MAP => + val map = dataType.getMap + MapType(fromProto(map.getKeyType), fromProto(map.getValueType), map.getValueContainsNull) + case P.KindCase.STRUCT => fromProto(dataType.getStruct) + case P.KindCase.KIND_NOT_SET => + throw new IllegalStateException("The library returned a data type without a kind.") } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala index 2e39c67122f04..817ce79ed26df 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -31,7 +31,6 @@ import scala.util.{Failure, Try} import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} import com.google.common.cache.{Cache, CacheBuilder} -import org.apache.spark.SPARK_VERSION import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.{PATH, VERSION_NUM} import org.apache.spark.sql.Artifact @@ -387,7 +386,7 @@ object NativeLibraries extends Logging { /** * Loads a library, and checks the version of the interface that it implements. Throws the error * if the library cannot be loaded, and returns it if the library is loaded but rejected, so that - * it is kept: the library stays loaded, and Spark only initializes it once. + * it is kept: the library stays loaded, and gets the same error when it is used again. * * @param invalidLibrary returns the error for an invalid library, with the reason * @param unsupportedAbiVersion returns the error for an unsupported version of the interface, @@ -399,9 +398,9 @@ object NativeLibraries extends Logging { unsupportedAbiVersion: Map[String, String] => Throwable): Try[NativeLibrary] = { logInfo(log"Loading the native data source library ${MDC(PATH, file)}.") val library = try { - NativeLibrary.load(file.getPath, SPARK_VERSION) + NativeLibrary.load(file.getPath) } catch { - // A library that was loaded stays loaded, and Spark must not initialize it again. + // A library that was loaded stays loaded, and gets the same error when it is used again. case e: NativeLibrary.LoadException if e.loaded => return Failure(invalidLibrary(e.getMessage, e)) case e: NativeLibrary.LoadException => throw invalidLibrary(e.getMessage, e) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala index 3e49dd2dae37d..a84fce136986e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.types._ /** * Encodes Data Source V2 predicates as the JSON expression trees passed to - * `push_predicates` of `spark_datasource.h`. + * `PushPredicates` of `spark_datasource.proto`. */ object NativePredicates { // Writes ASCII only, escaping the other characters, which any JSON parser decodes. @@ -36,7 +36,7 @@ object NativePredicates { /** * The names of the predicates that are passed to the libraries, as listed in - * `push_predicates`. They are part of the interface, so a predicate with any other + * `PushPredicates`. They are part of the interface, so a predicate with any other * name, such as an arithmetic function, is not passed: its semantics, for example on overflow, * would be part of the interface too, and the names of Data Source V2 grow with new releases. */ diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc index dc193afba889b..a74ef7ae2a7bd 100644 --- a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -16,14 +16,14 @@ */ // A native data source library used by NativeDataSourceSuite. It implements the C interface in -// spark_datasource.h without any other dependency, building the Arrow C data interface structs by -// hand. It provides three data sources: +// spark_datasource.h without any other dependency, building the Arrow C data interface structs and +// the Protocol Buffers messages of spark_datasource.proto by hand. It provides three data sources: // // native_range Reads rows [0, end) with columns of several types. Supports predicate, // column and limit pushdown. With the option "claim_predicates", it reports // all the predicates as evaluated completely, but ignores them. With the // option "predicate_use", it reports the predicates that it evaluates with -// another value: "partial" for SPARK_DS_PREDICATE_PARTIAL, as a reader that +// another value: "partial" for PREDICATE_USE_PARTIAL, as a reader that // reads the blocks of 10 rows that may match, or a number. With the option // "extra", the schema has one more column: "map" adds a map whose entry fields // are named "keys" and "values", as arrow-rs names them, "struct" adds a @@ -49,28 +49,26 @@ // // native_range supports batch reads, native_counter micro-batch streaming reads and native_sink // batch and streaming writes, which may replace the existing data. With the option -// "capabilities", a data source declares these capabilities instead, a number. With the option -// "task_path", the reads and the data writers of native_range and native_sink record their task -// in a file named after its partition, in that directory. With the option "log", creating a data -// source logs its value as a warning, and with the option "spark_version", it fails unless Spark -// has that version. +// "capabilities", a data source declares these capabilities instead, as the numbers of the enum +// Capability separated by commas. With the option "task_path", the reads and the data writers of +// native_range and native_sink record their task in a file named after its partition, in that +// directory. With the option "log", opening a data source logs its value as a warning, and with +// the option "spark_version", it fails unless Spark has that version. With the option +// "unknown_fields", every reply has a field that Spark does not know. // -// With the option "fail", a step fails with an error. With the option "zero", a function that -// creates a handle succeeds without setting it: "create", "create_reader", +// With the option "fail", a step fails with an error. With the option "zero", opening an object +// succeeds without setting its handle: "create" for a data source, "create_reader", // "create_stream_reader", "create_writer" or "create_data_writer". // // The suite compiles it with these macros to get several distinct libraries: // // NAME_PREFIX A prefix for the names of the data sources, "" by default. // ID_SIGN Multiplies the ids produced by native_range, 1 by default. -// READ_ONLY Leaves out the functions for writes, streaming reads and pushdown. -// ABI_VERSION The version of the table, 1 by default. -// NO_LAST_ERROR Leaves out last_error. -// TABLE_SIZE The size of the table, sizeof(SparkDsLibrary) by default: Spark must not read -// the functions after it. -// INIT_FAILURE 1 to make spark_ds_init fail, after logging why, 2 to make it succeed without -// a table. -// NO_INIT Leaves out spark_ds_init, which makes the library invalid. +// READ_ONLY Leaves out writes, streaming reads and pushdown: the replies to the pushdown +// requests are empty. +// ABI_VERSION The value of spark_ds_abi_version, 1 by default. +// NO_LAST_ERROR Leaves out spark_ds_last_error. +// NO_OPEN Leaves out spark_ds_open, which makes the library invalid. #include "spark_datasource.h" @@ -106,12 +104,6 @@ #ifndef ABI_VERSION #define ABI_VERSION 1 #endif -#ifndef TABLE_SIZE -#define TABLE_SIZE sizeof(SparkDsLibrary) -#endif -#ifndef INIT_FAILURE -#define INIT_FAILURE 0 -#endif namespace { @@ -241,16 +233,6 @@ void exportSchema(const std::vector& columns, ArrowSchema* out) { } } -// Reads the column names of a struct schema exported by Spark, and releases it. -std::vector importColumns(ArrowSchema* schema) { - std::vector columns; - for (int64_t i = 0; i < schema->n_children; i++) { - columns.push_back(schema->children[i]->name); - } - schema->release(schema); - return columns; -} - // --------------------------------------------------------------------------------------------- // Exporting arrays // --------------------------------------------------------------------------------------------- @@ -626,7 +608,7 @@ std::vector split(const std::string& text) { thread_local std::string last_error; -// Spark, as spark_ds_init got it. +// Spark, as spark_ds_open got it. const SparkDsHost* host = nullptr; // Logs with the log of Spark. A library must only use the fields that fit in the size of the @@ -637,6 +619,13 @@ void hostLog(int32_t level, const std::string& message) { } } +// Whether the task with the given ID was killed. +bool isInterrupted(int64_t task_attempt_id) { + return host->size >= static_cast(offsetof(SparkDsHost, is_interrupted) + + sizeof(host->is_interrupted)) && + host->is_interrupted(task_attempt_id); +} + // Runs `f`, and turns a C++ exception into an error: none may cross the boundary. template int guarded(F&& f) { @@ -649,10 +638,6 @@ int guarded(F&& f) { } } -std::vector toStrings(const char* const* values, int64_t count) { - return std::vector(values, values + count); -} - std::string toBytes(const uint8_t* data, int64_t length) { return std::string(reinterpret_cast(data), static_cast(length)); } @@ -662,36 +647,12 @@ void sendBytes(SparkDsSink sink_function, void* sink, const std::string& bytes) static_cast(bytes.size())); } -[[maybe_unused]] -std::vector toByteStrings(const uint8_t* const* messages, const int64_t* lengths, - int64_t count, bool* has_null) { - std::vector result; - for (int64_t i = 0; i < count; i++) { - if (messages[i] == nullptr) { - *has_null = true; - } else { - result.push_back(toBytes(messages[i], lengths[i])); - } - } - return result; -} - -template -T* fromHandle(void* handle) { - return static_cast(handle); -} - // --------------------------------------------------------------------------------------------- // Data sources // --------------------------------------------------------------------------------------------- std::string name(const std::string& base) { return std::string(NAME_PREFIX) + base; } -struct DataSource { - std::string name; - Options options; -}; - void maybeFail(const Options& options, const std::string& step) { if (get(options, "fail", "") == step) throw std::runtime_error("injected failure in " + step); } @@ -715,24 +676,6 @@ void makeDirs(const std::string& path) { mkdir(path.c_str(), 0755); } -// With the option "task_path", records the task of a read or a data writer in a file named after -// its partition: the size of the struct and its fields. -void recordTask(const Options& options, const SparkDsTaskContext* task) { - std::string dir = get(options, "task_path", ""); - if (dir.empty()) return; - makeDirs(dir); - writeFile(dir + "/" + std::to_string(task->partition_id), - std::to_string(task->size) + "," + std::to_string(task->stage_id) + "," + - std::to_string(task->stage_attempt_number) + "," + - std::to_string(task->partition_id) + "," + std::to_string(task->attempt_number) + - "," + std::to_string(task->task_attempt_id)); -} - -// The driver-side reader of native_range and native_counter. -struct Reader { - Options state; // Serialized to the executors. -}; - // Reads the ids [lo, hi) of the columns in the state, in batches. struct Stream { std::vector columns; @@ -742,7 +685,7 @@ struct Stream { bool fail; bool slice; std::string short_array; - const SparkDsTaskContext* task; + int64_t task_attempt_id; std::string wait_for_interrupt; std::string last_error; }; @@ -762,10 +705,10 @@ int streamGetNext(ArrowArrayStream* stream, ArrowArray* out) { if (!state->wait_for_interrupt.empty()) { // Waits, as for data that does not come, until the task is killed, for at most a minute. writeFile(state->wait_for_interrupt + ".waiting", ""); - for (int i = 0; i < 6000 && !state->task->is_interrupted(state->task); i++) { + for (int i = 0; i < 6000 && !isInterrupted(state->task_attempt_id); i++) { std::this_thread::sleep_for(std::chrono::milliseconds(10)); } - bool interrupted = state->task->is_interrupted(state->task); + bool interrupted = isInterrupted(state->task_attempt_id); writeFile(state->wait_for_interrupt, interrupted ? "interrupted" : "timed out"); state->last_error = "the task was killed"; return ECANCELED; @@ -789,33 +732,6 @@ void streamRelease(ArrowArrayStream* stream) { stream->release = nullptr; } -// The driver-side writer of native_sink. -struct Writer { - Options state; // path, overwrite, fail. - // The file that closeWriter creates, so that the tests can check that the writer of a committed - // or aborted write is closed: _CLOSED for a batch write, _CLOSED_ for a streaming one. - std::string closed_marker; -}; - -[[maybe_unused]] std::string closedMarker(int64_t epoch_id) { - return epoch_id < 0 ? "/_CLOSED" : "/_CLOSED_" + std::to_string(epoch_id); -} - -// Writes the rows of a task as CSV lines to a temporary file, renamed on commit. -struct DataWriter { - std::string temp_path; - FILE* file; - Options state; - bool committed; -}; - -// The data writers that Spark closed, with the path of their write, so that using one of them -// afterwards, which the header forbids, is recorded in the file _DATA_WRITER_USED_AFTER_CLOSE -// rather than crashing. -std::mutex closed_data_writers_lock; -[[maybe_unused]] std::map closed_data_writers; - - [[maybe_unused]] bool isValid(const ArrowArray* array, int64_t i) { if (array->null_count == 0 || array->buffers[0] == nullptr) return true; int64_t bit = array->offset + i; @@ -862,135 +778,609 @@ std::vector listFiles(const std::string& dir, const std::string& pr } // --------------------------------------------------------------------------------------------- -// The functions of the table +// Messages // --------------------------------------------------------------------------------------------- -#ifndef NO_LAST_ERROR -const char* lastError(void) { return last_error.c_str(); } -#endif +// A minimal Protocol Buffers codec for the messages of spark_datasource.proto, so that the test +// library needs nothing but a C++ compiler. A library would rather use the code that a Protocol +// Buffers compiler generates from spark_datasource.proto. +class ProtoReader { + public: + // Keeps the bytes of the message, which are often a temporary string. + explicit ProtoReader(std::string bytes) + : bytes_(std::move(bytes)), + p_(reinterpret_cast(bytes_.data())), + end_(p_ + bytes_.size()) {} + + // Moves to the next field, and returns false at the end of the message. + bool next() { + if (p_ == end_) return false; + uint64_t key = varint(); + field_ = static_cast(key >> 3); + wire_type_ = static_cast(key & 7); + return true; + } -// The capabilities of a data source, unless the option "capabilities" declares others. -uint64_t capabilitiesOf(const DataSource& source) { - std::string declared = get(source.options, "capabilities", ""); - if (!declared.empty()) return std::stoull(declared); - if (source.name == name("native_range")) return SPARK_DS_BATCH_READ; -#ifndef READ_ONLY - if (source.name == name("native_counter")) return SPARK_DS_MICRO_BATCH_READ; - if (source.name == name("native_sink")) { - return SPARK_DS_BATCH_WRITE | SPARK_DS_STREAMING_WRITE | SPARK_DS_TRUNCATE; + int field() const { return field_; } + + uint64_t varint() { + uint64_t value = 0; + for (int shift = 0; shift < 64; shift += 7) { + if (p_ == end_) throw std::runtime_error("truncated message"); + uint8_t byte = *p_++; + value |= static_cast(byte & 0x7f) << shift; + if ((byte & 0x80) == 0) return value; + } + throw std::runtime_error("invalid varint"); } -#endif - return 0; -} -int createSource(const SparkDsCreateInfo* info, uint64_t* capabilities, void** handle) { - return guarded([&] { - // Like the host, the struct may be of another version of the header: only the fields that fit - // in its size are read. - if (info->size < static_cast(offsetof(SparkDsCreateInfo, num_options) + - sizeof(info->num_options))) { - throw std::runtime_error("SparkDsCreateInfo has a size of " + std::to_string(info->size)); + int64_t int64() { return static_cast(varint()); } + + // A string, bytes or a message. + std::string bytes() { + uint64_t length = varint(); + if (length > static_cast(end_ - p_)) throw std::runtime_error("truncated message"); + std::string value(reinterpret_cast(p_), static_cast(length)); + p_ += length; + return value; + } + + // Skips a field that the library does not know, as Protocol Buffers requires. + void skip() { + switch (wire_type_) { + case 0: varint(); break; + case 1: advance(8); break; + case 2: bytes(); break; + case 5: advance(4); break; + default: throw std::runtime_error("unsupported wire type " + std::to_string(wire_type_)); } - auto source = std::make_unique(DataSource{info->name, {}}); - auto keys = toStrings(info->option_keys, info->num_options); - auto values = toStrings(info->option_values, info->num_options); - for (size_t i = 0; i < keys.size(); i++) source->options[keys[i]] = values[i]; - maybeFail(source->options, "create"); - if (source->name != name("native_range") && source->name != name("native_sink") && - source->name != name("native_counter")) { - throw std::runtime_error("unknown data source " + source->name); + } + + private: + void advance(size_t n) { + if (n > static_cast(end_ - p_)) throw std::runtime_error("truncated message"); + p_ += n; + } + + std::string bytes_; + const uint8_t* p_; + const uint8_t* end_; + int field_ = 0; + int wire_type_ = 0; +}; + +class ProtoWriter { + public: + ProtoWriter& varint(int field, uint64_t value) { + key(field, 0); + raw(value); + return *this; + } + + ProtoWriter& bytes(int field, const std::string& value) { + key(field, 2); + raw(value.size()); + out_ += value; + return *this; + } + + ProtoWriter& message(int field, const ProtoWriter& value) { return bytes(field, value.out_); } + + // A packed repeated field of enums. + ProtoWriter& packed(int field, const std::vector& values) { + ProtoWriter packed; + for (int32_t value : values) packed.raw(static_cast(static_cast(value))); + return bytes(field, packed.out_); + } + + const std::string& str() const { return out_; } + + private: + void key(int field, int wire_type) { + raw((static_cast(field) << 3) | static_cast(wire_type)); + } + + void raw(uint64_t value) { + while (value >= 0x80) { + out_ += static_cast((value & 0x7f) | 0x80); + value >>= 7; } - std::string message = get(source->options, "log", ""); - if (!message.empty()) hostLog(SPARK_DS_LOG_WARN, message); - std::string version = get(source->options, "spark_version", ""); - if (!version.empty() && version != host->spark_version) { - throw std::runtime_error("the version of Spark is " + std::string(host->spark_version) + - ", which supports version " + std::to_string(host->abi_version) + - " of the interface"); + out_ += static_cast(value); + } + + std::string out_; +}; + +// The fields of the messages, by their number in spark_datasource.proto. +enum OpenField { kSparkVersion = 1, kDataSource = 10, kBatchReader, kStreamReader, kWriter, + kPartitionReader, kDataWriter }; +enum CallField { kInferSchema = 1, kPushPredicates, kPushLimit, kPruneColumns, kPlanPartitions, + kSerialize, kInitialOffset, kLatestOffset, kPlanStreamPartitions, kCommitOffset, + kCommit, kAbort, kCommitDataWriter, kAbortDataWriter }; +enum StreamField { kRead = 1, kWrite }; +enum Capability { kBatchRead = 1, kMicroBatchRead, kBatchWrite, kStreamingWrite, kTruncate }; +enum PredicateUse { kUnused = 0, kPartial, kComplete }; +enum WriteMode { kAppend = 0, kTruncateMode }; + +// With the option "unknown_fields", every reply has a field that Spark does not know, which it +// must ignore. +std::string withUnknownField(const Options& options, const ProtoWriter& reply) { + ProtoWriter result = reply; + if (get(options, "unknown_fields", "") == "true") result.bytes(1000, "unknown"); + return result.str(); +} + +void sendReply(SparkDsSink reply, void* sink, const Options& options, const ProtoWriter& message) { + sendBytes(reply, sink, withUnknownField(options, message)); +} + +// The fields of a message, by number: the last value of each field. +std::map parseMessage(const std::string& bytes) { + std::map fields; + ProtoReader reader(bytes); + while (reader.next()) { + // The fields that the library reads are strings, bytes and messages, or varints, which are + // returned as decimal numbers. + int field = reader.field(); + try { + fields[field] = reader.bytes(); + } catch (const std::exception&) { + throw std::runtime_error("unexpected field " + std::to_string(field)); } - *capabilities = capabilitiesOf(*source); - if (returnsZero(source->options, "create")) return; - *handle = source.release(); - }); + } + return fields; } -int sourceSchema(void* handle, struct ArrowSchema* out) { - return guarded([&] { - auto* source = fromHandle(handle); - maybeFail(source->options, "schema"); - if (source->name == name("native_sink")) { - throw std::runtime_error("native_sink has no schema; it can only be written"); - } - bool counter = source->name == name("native_counter"); - exportSchema(allColumns(counter, get(source->options, "extra", "")), out); - }); +// The varints of a message, by number. +std::map parseVarints(const std::string& bytes) { + std::map fields; + ProtoReader reader(bytes); + while (reader.next()) { + int field = reader.field(); + fields[field] = reader.varint(); + } + return fields; } -int createReader(void* handle, struct ArrowSchema* schema, void** out) { - return guarded([&] { - auto* source = fromHandle(handle); - if (source->name != name("native_range")) { - throw std::runtime_error(source->name + " does not support batch reads"); +// A TaskContext message. +struct TaskInfo { + int64_t stage_id = 0; + int64_t stage_attempt_number = 0; + int64_t partition_id = 0; + int64_t attempt_number = 0; + int64_t task_attempt_id = 0; +}; + +TaskInfo parseTask(const std::string& bytes) { + auto fields = parseVarints(bytes); + auto field = [&](int number) { + return static_cast(static_cast(fields.count(number) ? fields[number] : 0)); + }; + TaskInfo task; + task.stage_id = field(1); + task.stage_attempt_number = field(2); + task.partition_id = field(3); + task.attempt_number = field(4); + task.task_attempt_id = fields.count(5) ? static_cast(fields[5]) : 0; + return task; +} + +// The names of the columns of a StructType message. +std::vector parseColumns(const std::string& struct_type) { + std::vector columns; + ProtoReader fields(struct_type); + while (fields.next()) { + if (fields.field() != 1) { + fields.skip(); + continue; } - auto columns = importColumns(schema); - for (const auto& column : columns) findColumn(column); - if (returnsZero(source->options, "create_reader")) return; - auto* reader = new Reader(); - reader->state["columns"] = join(columns); - reader->state["lo"] = "0"; - reader->state["hi"] = get(source->options, "end", "10"); - reader->state["partitions"] = get(source->options, "partitions", "2"); - reader->state["batch_size"] = get(source->options, "batch_size", "1000"); - for (const char* option : {"fail", "slice", "short", "claim_predicates", "predicate_use", - "wait_for_interrupt", "task_path"}) { - reader->state[option] = get(source->options, option, ""); + ProtoReader field(fields.bytes()); + std::string name; + while (field.next()) { + if (field.field() == 1) { + name = field.bytes(); + } else { + field.skip(); + } } - *out = reader; - }); + columns.push_back(name); + } + return columns; } -void closeSource(void* handle) { delete fromHandle(handle); } +// A DataType message of a type without parameters, by the number of its kind. +ProtoWriter simpleType(int kind) { return ProtoWriter().message(kind, ProtoWriter()); } + +// The DataType message of a column, as Spark infers it from the Arrow type: Spark cannot express +// some of the Arrow types of the extra columns, whose data it rejects. +ProtoWriter columnType(const std::string& name) { + enum { kInteger = 5, kLong, kDouble = 8, kDecimal, kString, kBoolean = 2, kDate = 12, + kTimestamp, kArray = 18, kMap, kStruct }; + auto field = [](const std::string& field_name, const ProtoWriter& type) { + return ProtoWriter().bytes(1, field_name).message(2, type).varint(3, 1); + }; + std::string format = findColumn(name)->format; + if (name == "m") { + return ProtoWriter().message(kMap, ProtoWriter() + .message(1, ProtoWriter().message(kString, ProtoWriter())) + .message(2, simpleType(kLong)) + .varint(3, 1)); + } else if (name == "st") { + return ProtoWriter().message(kStruct, ProtoWriter() + .message(1, field("a", simpleType(kLong))) + .message(1, field("b", ProtoWriter().message(kString, ProtoWriter())))); + } else if (name == "lv" || name == "ll") { + return ProtoWriter().message(kArray, + ProtoWriter().message(1, simpleType(kLong)).varint(2, 1)); + } else if (name == "dict" || format == "u") { + return ProtoWriter().message(kString, ProtoWriter()); + } else if (format == "l") { + return simpleType(kLong); + } else if (format == "i") { + return simpleType(kInteger); + } else if (format == "g") { + return simpleType(kDouble); + } else if (format == "b") { + return simpleType(kBoolean); + } else if (format == "tdD") { + return simpleType(kDate); + } else if (format.rfind("ts", 0) == 0) { + return simpleType(kTimestamp); + } else if (format.rfind("d:", 0) == 0) { + return ProtoWriter().message(kDecimal, ProtoWriter().varint(1, 10).varint(2, 2)); + } + throw std::runtime_error("no Spark type for the column " + name); +} + +// The StructType message of the given columns. +ProtoWriter structType(const std::vector& columns) { + ProtoWriter schema; + for (const auto& column : columns) { + schema.message(1, ProtoWriter().bytes(1, column).message(2, columnType(column)).varint(3, 1)); + } + return schema; +} + +// --------------------------------------------------------------------------------------------- +// Objects +// --------------------------------------------------------------------------------------------- + +// What a handle refers to. +struct Object { + virtual ~Object() = default; +}; + +struct DataSource : Object { + std::string name; + Options options; +}; + +// The batch or stream reader of native_range or native_counter. +struct Reader : Object { + Options state; // Serialized to the executors. +}; +// The reader of one partition, on an executor. +struct PartitionReader : Object { + Options state; + std::string partition; + TaskInfo task; +}; + +// The writer of native_sink. +struct Writer : Object { + Options state; // path, overwrite, fail. + // The file that closing the writer creates, so that the tests can check that the writer of a + // committed or aborted write is closed: _CLOSED for a batch write, _CLOSED_ for a + // streaming one. + std::string closed_marker; +}; + +[[maybe_unused]] std::string closedMarker(int64_t epoch_id) { + return epoch_id < 0 ? "/_CLOSED" : "/_CLOSED_" + std::to_string(epoch_id); +} + +// Writes the rows of a task as CSV lines to a temporary file, renamed on commit. The commit records +// in _streams the number of streams and of batches that the data writer got. +struct DataWriter : Object { + std::string temp_path; + FILE* file = nullptr; + Options state; + bool committed = false; + int64_t streams = 0; + int64_t batches = 0; +}; + +// The data writers that Spark closed, with the path of their write, so that using one of them +// afterwards, which the header forbids, is recorded in the file _DATA_WRITER_USED_AFTER_CLOSE +// rather than crashing. +std::mutex closed_data_writers_lock; +std::map closed_data_writers; + +// Whether Spark uses a data writer that it closed, which the header forbids. +bool usedAfterClose(void* handle) { + std::lock_guard guard(closed_data_writers_lock); + auto closed = closed_data_writers.find(handle); + if (closed == closed_data_writers.end()) return false; + writeFile(closed->second + "/_DATA_WRITER_USED_AFTER_CLOSE", ""); + return true; +} + +// Returns a new object as a handle. The allocator may reuse the address of a closed data writer. +void* newHandle(Object* object) { + std::lock_guard guard(closed_data_writers_lock); + closed_data_writers.erase(object); + return object; +} + +template +T* as(void* handle, const char* request) { + auto* object = dynamic_cast(static_cast(handle)); + if (object == nullptr) throw std::runtime_error(std::string(request) + " on another object"); + return object; +} + +// With the option "task_path", records the task of a read or a data writer in a file named after +// its partition: its fields. +void recordTask(const Options& options, const TaskInfo& task) { + std::string dir = get(options, "task_path", ""); + if (dir.empty()) return; + makeDirs(dir); + writeFile(dir + "/" + std::to_string(task.partition_id), + std::to_string(task.stage_id) + "," + std::to_string(task.stage_attempt_number) + + "," + std::to_string(task.partition_id) + "," + + std::to_string(task.attempt_number) + "," + + std::to_string(task.task_attempt_id)); +} + +DataSource parseDataSource(const std::string& bytes) { + DataSource source; + ProtoReader reader(bytes); + while (reader.next()) { + if (reader.field() == 1) { + source.name = reader.bytes(); + } else if (reader.field() == 2) { + // A map is a repeated message of entries, with the key and the value. + ProtoReader entry(reader.bytes()); + std::string key; + std::string value; + while (entry.next()) { + if (entry.field() == 1) { + key = entry.bytes(); + } else if (entry.field() == 2) { + value = entry.bytes(); + } else { + entry.skip(); + } + } + source.options[key] = value; + } else { + reader.skip(); + } + } + return source; +} + +// The capabilities of a data source, unless the option "capabilities" declares others: their +// numbers, separated by commas. +std::vector capabilitiesOf(const DataSource& source) { + auto declared = source.options.find("capabilities"); + if (declared != source.options.end()) { + std::vector capabilities; + for (const auto& value : split(declared->second)) capabilities.push_back(std::stoi(value)); + return capabilities; + } + if (source.name == name("native_range")) return {kBatchRead}; #ifndef READ_ONLY + if (source.name == name("native_counter")) return {kMicroBatchRead}; + if (source.name == name("native_sink")) return {kBatchWrite, kStreamingWrite, kTruncate}; +#endif + return {}; +} -int createStreamReader(void* handle, struct ArrowSchema* schema, void** out) { - return guarded([&] { - auto* source = fromHandle(handle); - if (source->name != name("native_counter")) { - throw std::runtime_error(source->name + " does not support streaming reads"); +// --------------------------------------------------------------------------------------------- +// spark_ds_open +// --------------------------------------------------------------------------------------------- + +void* openDataSource(const std::string& message, const std::string& spark_version, + SparkDsSink reply, void* sink) { + auto source = std::make_unique(parseDataSource(message)); + maybeFail(source->options, "create"); + if (source->name != name("native_range") && source->name != name("native_sink") && + source->name != name("native_counter")) { + throw std::runtime_error("unknown data source " + source->name); + } + std::string log = get(source->options, "log", ""); + if (!log.empty()) hostLog(SPARK_DS_LOG_WARN, log); + std::string version = get(source->options, "spark_version", ""); + if (!version.empty() && version != spark_version) { + throw std::runtime_error("the version of Spark is " + spark_version); + } + sendReply(reply, sink, source->options, ProtoWriter().packed(1, capabilitiesOf(*source))); + if (returnsZero(source->options, "create")) return nullptr; + return newHandle(source.release()); +} + +void* openBatchReader(const std::string& message) { + auto fields = parseMessage(message); + DataSource source = parseDataSource(fields[1]); + if (source.name != name("native_range")) { + throw std::runtime_error(source.name + " does not support batch reads"); + } + auto columns = parseColumns(fields[2]); + for (const auto& column : columns) findColumn(column); + if (returnsZero(source.options, "create_reader")) return nullptr; + auto reader = std::make_unique(); + reader->state["columns"] = join(columns); + reader->state["lo"] = "0"; + reader->state["hi"] = get(source.options, "end", "10"); + reader->state["partitions"] = get(source.options, "partitions", "2"); + reader->state["batch_size"] = get(source.options, "batch_size", "1000"); + for (const char* option : {"fail", "slice", "short", "claim_predicates", "predicate_use", + "wait_for_interrupt", "task_path", "unknown_fields"}) { + reader->state[option] = get(source.options, option, ""); + } + return newHandle(reader.release()); +} + +#ifndef READ_ONLY + +void* openStreamReader(const std::string& message) { + auto fields = parseMessage(message); + DataSource source = parseDataSource(fields[1]); + if (source.name != name("native_counter")) { + throw std::runtime_error(source.name + " does not support streaming reads"); + } + auto columns = join(parseColumns(fields[2])); + if (returnsZero(source.options, "create_stream_reader")) return nullptr; + auto reader = std::make_unique(); + reader->state["columns"] = columns; + reader->state["max_offset"] = get(source.options, "max_offset", "10"); + reader->state["step"] = get(source.options, "step", reader->state["max_offset"]); + reader->state["current"] = "0"; + reader->state["batch_size"] = get(source.options, "batch_size", "1000"); + reader->state["commit_path"] = get(source.options, "commit_path", ""); + reader->state["bad_offset"] = get(source.options, "bad_offset", ""); + reader->state["closed_path"] = get(source.options, "closed_path", ""); + reader->state["stream"] = "true"; + return newHandle(reader.release()); +} + +void* openWriter(const std::string& message) { + ProtoReader reader(message); + DataSource source; + std::vector columns; + uint64_t mode = kAppend; + std::string query_id; + while (reader.next()) { + switch (reader.field()) { + case 1: source = parseDataSource(reader.bytes()); break; + case 2: columns = parseColumns(reader.bytes()); break; + case 3: mode = reader.varint(); break; + case 4: query_id = reader.bytes(); break; + default: reader.skip(); } - auto columns = join(importColumns(schema)); - if (returnsZero(source->options, "create_stream_reader")) return; - auto* reader = new Reader(); - reader->state["columns"] = columns; - reader->state["max_offset"] = get(source->options, "max_offset", "10"); - reader->state["step"] = get(source->options, "step", reader->state["max_offset"]); - reader->state["current"] = "0"; - reader->state["batch_size"] = get(source->options, "batch_size", "1000"); - reader->state["commit_path"] = get(source->options, "commit_path", ""); - reader->state["bad_offset"] = get(source->options, "bad_offset", ""); - reader->state["closed_path"] = get(source->options, "closed_path", ""); - *out = reader; - }); + } + if (source.name != name("native_sink")) { + throw std::runtime_error(source.name + " does not support writes"); + } + if (mode != kAppend && mode != kTruncateMode) { + throw std::runtime_error("unknown mode " + std::to_string(mode)); + } + if (columns.empty()) throw std::runtime_error("no columns to write"); + // The location of a table in a catalog is passed as a URI. + std::string path = get(source.options, "path", ""); + if (path.rfind("file://", 0) == 0) { + path = path.substr(7); + } else if (path.rfind("file:", 0) == 0) { + path = path.substr(5); + } + if (path.empty()) throw std::runtime_error("the option 'path' is required"); + if (returnsZero(source.options, "create_writer")) return nullptr; + auto writer = std::make_unique(); + writer->state["path"] = path; + writer->state["overwrite"] = mode == kTruncateMode ? "true" : "false"; + for (const char* option : {"fail", "zero", "task_path", "unknown_fields"}) { + writer->state[option] = get(source.options, option, ""); + } + writer->state["commit_messages"] = get(source.options, "commit_messages", "1"); + // Records that a writer was created, which Spark then commits or aborts, and the ID of the + // streaming query. + makeDirs(path); + writeFile(path + "/_CREATED", ""); + if (!query_id.empty()) writeFile(path + "/_QUERY_ID", query_id); + return newHandle(writer.release()); +} + +void* openDataWriter(const std::string& message) { + ProtoReader reader(message); + TaskInfo task; + Options state; + int64_t epoch_id = 0; + while (reader.next()) { + switch (reader.field()) { + case 1: task = parseTask(reader.bytes()); break; + case 2: state = deserialize(reader.bytes()); break; + case 3: epoch_id = reader.int64(); break; + default: reader.skip(); + } + } + if (returnsZero(state, "create_data_writer")) return nullptr; + recordTask(state, task); + std::string dir = state["path"] + "/_temporary"; + makeDirs(dir); + std::string temp_path = dir + "/" + std::to_string(epoch_id) + "-" + + std::to_string(task.partition_id) + "-" + + std::to_string(task.task_attempt_id) + ".csv"; + FILE* file = std::fopen(temp_path.c_str(), "w"); + if (file == nullptr) throw std::runtime_error("cannot write " + temp_path); + makeDirs(state["path"] + "/_created_data_writers"); + writeFile(state["path"] + "/_created_data_writers/" + temp_path.substr(dir.size() + 1), ""); + auto writer = std::make_unique(); + writer->temp_path = temp_path; + writer->file = file; + writer->state = state; + return newHandle(writer.release()); } -int pushPredicates(void* handle, const char* const* predicates, int64_t count, int32_t* uses) { - return guarded([&] { - auto* reader = fromHandle(handle); - auto json = toStrings(predicates, count); - int64_t end = getLong(reader->state, "hi", 0); - int64_t lo = getLong(reader->state, "lo", 0); - int64_t hi = end; - // With the option "claim_predicates", reports every predicate as evaluated completely, but - // evaluates none of them. - if (reader->state["claim_predicates"] == "true") { - for (size_t i = 0; i < json.size(); i++) uses[i] = SPARK_DS_PREDICATE_COMPLETE; - return; +#endif // READ_ONLY + +void* openPartitionReader(const std::string& message) { + ProtoReader reader(message); + auto partition = std::make_unique(); + while (reader.next()) { + switch (reader.field()) { + case 1: partition->task = parseTask(reader.bytes()); break; + case 2: partition->state = deserialize(reader.bytes()); break; + case 3: partition->partition = reader.bytes(); break; + default: reader.skip(); + } + } + maybeFail(partition->state, "open"); + recordTask(partition->state, partition->task); + return newHandle(partition.release()); +} + +// --------------------------------------------------------------------------------------------- +// spark_ds_call +// --------------------------------------------------------------------------------------------- + +void inferSchema(DataSource* source, SparkDsSink reply, void* sink) { + maybeFail(source->options, "schema"); + if (source->name == name("native_sink")) { + throw std::runtime_error("native_sink has no schema; it can only be written"); + } + bool counter = source->name == name("native_counter"); + auto columns = allColumns(counter, get(source->options, "extra", "")); + sendReply(reply, sink, source->options, + ProtoWriter().message(kInferSchema, ProtoWriter().message(1, structType(columns)))); +} + +#ifndef READ_ONLY + +void pushPredicates(Reader* reader, const std::string& message, SparkDsSink reply, void* sink) { + std::vector json; + ProtoReader request(message); + while (request.next()) { + if (request.field() == 1) { + json.push_back(request.bytes()); + } else { + request.skip(); } + } + int64_t end = getLong(reader->state, "hi", 0); + int64_t lo = getLong(reader->state, "lo", 0); + int64_t hi = end; + std::vector uses(json.size(), kUnused); + // With the option "claim_predicates", reports every predicate as evaluated completely, but + // evaluates none of them. + if (reader->state["claim_predicates"] == "true") { + std::fill(uses.begin(), uses.end(), kComplete); + } else { std::string predicate_use = reader->state["predicate_use"]; bool partial = predicate_use == "partial"; - int32_t use = predicate_use.empty() ? SPARK_DS_PREDICATE_COMPLETE - : partial ? SPARK_DS_PREDICATE_PARTIAL + int32_t use = predicate_use.empty() ? kComplete + : partial ? kPartial : std::stoi(predicate_use); for (size_t i = 0; i < json.size(); i++) { // Accepts "id " for the comparison operators. @@ -1015,7 +1405,7 @@ int pushPredicates(void* handle, const char* const* predicates, int64_t count, i hi = std::min(hi, value + 1); } } - uses[i] = ok ? use : SPARK_DS_PREDICATE_UNUSED; + uses[i] = ok ? use : kUnused; } if (partial) { // Reads the blocks of 10 rows that contain matching rows, with all their rows. @@ -1024,374 +1414,380 @@ int pushPredicates(void* handle, const char* const* predicates, int64_t count, i } reader->state["lo"] = std::to_string(lo); reader->state["hi"] = std::to_string(std::max(lo, hi)); - reader->state["predicates"] = std::to_string(json.size()); - }); + } + reader->state["predicates"] = std::to_string(json.size()); + sendReply(reply, sink, reader->state, + ProtoWriter().message(kPushPredicates, ProtoWriter().packed(1, uses))); } -int pushLimit(void* handle, int32_t limit, bool* applied) { - return guarded([&] { - auto* reader = fromHandle(handle); - // Spark calls push_limit at most once. - if (reader->state.count("limit") > 0) throw std::runtime_error("pushLimit was called twice"); - reader->state["limit"] = std::to_string(limit); - *applied = true; - }); +void pushLimit(Reader* reader, const std::string& message, SparkDsSink reply, void* sink) { + // Spark pushes a limit at most once. + if (reader->state.count("limit") > 0) throw std::runtime_error("PushLimit was called twice"); + auto fields = parseVarints(message); + reader->state["limit"] = std::to_string(static_cast(fields[1])); + sendReply(reply, sink, reader->state, + ProtoWriter().message(kPushLimit, ProtoWriter().varint(1, 1))); } -int pruneColumns(void* handle, const char* const* columns, int64_t count, bool* pruned) { - return guarded([&] { - auto names = toStrings(columns, count); - for (const auto& column : names) findColumn(column); - fromHandle(handle)->state["columns"] = join(names); - *pruned = true; - }); +void pruneColumns(Reader* reader, const std::string& message, SparkDsSink reply, void* sink) { + std::vector names; + ProtoReader request(message); + while (request.next()) { + if (request.field() == 1) { + names.push_back(request.bytes()); + } else { + request.skip(); + } + } + for (const auto& column : names) findColumn(column); + reader->state["columns"] = join(names); + sendReply(reply, sink, reader->state, + ProtoWriter().message(kPruneColumns, ProtoWriter().varint(1, 1))); } #endif // READ_ONLY -int planPartitions(void* handle, SparkDsSink add_partition, void* sink) { - return guarded([&] { - auto* reader = fromHandle(handle); - maybeFail(reader->state, "plan"); - int64_t lo = getLong(reader->state, "lo", 0); - int64_t hi = getLong(reader->state, "hi", 0); - int64_t count = getLong(reader->state, "partitions", 1); - if (reader->state.count("limit") > 0) { - // A pushed limit is applied by reading a single partition with at most `limit` rows. - hi = std::min(hi, lo + getLong(reader->state, "limit", 0)); - count = 1; - } - int64_t step = std::max(1, (hi - lo + count - 1) / count); - for (int64_t start = lo; start < hi; start += step) { - int64_t stop = std::min(hi, start + step); - sendBytes(add_partition, sink, std::to_string(start) + "," + std::to_string(stop)); - } - }); +void planPartitions(Reader* reader, SparkDsSink reply, void* sink) { + maybeFail(reader->state, "plan"); + int64_t lo = getLong(reader->state, "lo", 0); + int64_t hi = getLong(reader->state, "hi", 0); + int64_t count = getLong(reader->state, "partitions", 1); + if (reader->state.count("limit") > 0) { + // A pushed limit is applied by reading a single partition with at most `limit` rows. + hi = std::min(hi, lo + getLong(reader->state, "limit", 0)); + count = 1; + } + ProtoWriter partitions; + int64_t step = std::max(1, (hi - lo + count - 1) / count); + for (int64_t start = lo; start < hi; start += step) { + int64_t stop = std::min(hi, start + step); + partitions.bytes(1, std::to_string(start) + "," + std::to_string(stop)); + } + sendReply(reply, sink, reader->state, ProtoWriter().message(kPlanPartitions, partitions)); } -int serializeReader(void* handle, SparkDsSink set_state, void* sink) { - return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); +void serializeState(const Options& state, SparkDsSink reply, void* sink) { + sendReply(reply, sink, state, + ProtoWriter().message(kSerialize, ProtoWriter().bytes(1, serialize(state)))); } -void closeReader(void* handle) { delete fromHandle(handle); } - #ifndef READ_ONLY -int initialOffset(void*, SparkDsSink set_offset, void* sink) { - return guarded([&] { sendBytes(set_offset, sink, "{\"offset\":0}"); }); +void initialOffset(Reader* reader, SparkDsSink reply, void* sink) { + sendReply(reply, sink, reader->state, + ProtoWriter().message(kInitialOffset, ProtoWriter().bytes(1, "{\"offset\":0}"))); } // Each call makes `step` more rows available, up to `max_offset`. -int latestOffset(void* handle, SparkDsSink set_offset, void* sink) { - return guarded([&] { - Options& state = fromHandle(handle)->state; - int64_t latest = std::min(getLong(state, "max_offset", 10), - getLong(state, "current", 0) + getLong(state, "step", 10)); - state["current"] = std::to_string(latest); - std::string bad_offset = get(state, "bad_offset", ""); - if (bad_offset == "none") return; - if (bad_offset == "multiline") { - // Pretty-printed JSON, as serde_json::to_string_pretty writes it. - sendBytes(set_offset, sink, "{\n \"offset\": " + std::to_string(latest) + "\n}"); - } else if (bad_offset == "invalid") { - sendBytes(set_offset, sink, "offset " + std::to_string(latest)); +void latestOffset(Reader* reader, SparkDsSink reply, void* sink) { + Options& state = reader->state; + int64_t latest = std::min(getLong(state, "max_offset", 10), + getLong(state, "current", 0) + getLong(state, "step", 10)); + state["current"] = std::to_string(latest); + std::string bad_offset = get(state, "bad_offset", ""); + if (bad_offset == "none") return; + std::string offset = "{\"offset\":" + std::to_string(latest) + "}"; + if (bad_offset == "multiline") { + // Pretty-printed JSON, as serde_json::to_string_pretty writes it. + offset = "{\n \"offset\": " + std::to_string(latest) + "\n}"; + } else if (bad_offset == "invalid") { + offset = "offset " + std::to_string(latest); + } + sendReply(reply, sink, state, + ProtoWriter().message(kLatestOffset, ProtoWriter().bytes(1, offset))); +} + +void planStreamPartitions(Reader* reader, const std::string& message, SparkDsSink reply, + void* sink) { + auto fields = parseMessage(message); + auto lo = static_cast(JsonParser(fields[1]).parse()["offset"].number); + auto hi = static_cast(JsonParser(fields[2]).parse()["offset"].number); + ProtoWriter partitions; + if (hi > lo) partitions.bytes(1, std::to_string(lo) + "," + std::to_string(hi)); + sendReply(reply, sink, reader->state, + ProtoWriter().message(kPlanStreamPartitions, partitions)); +} + +void commitOffset(Reader* reader, const std::string& message) { + std::string path = get(reader->state, "commit_path", ""); + if (!path.empty()) writeFile(path, parseMessage(message)[1]); +} + +// The epoch ID and the messages of a Commit or an Abort, with whether a task did not commit. +int64_t parseTaskMessages(const std::string& message, std::vector* messages, + int64_t* uncommitted) { + int64_t epoch_id = 0; + ProtoReader request(message); + while (request.next()) { + if (request.field() == 1) { + epoch_id = request.int64(); + } else if (request.field() == 2) { + ProtoReader task(request.bytes()); + bool committed = false; + std::string bytes; + while (task.next()) { + if (task.field() == 1) { + committed = task.varint() != 0; + } else if (task.field() == 2) { + bytes = task.bytes(); + } else { + task.skip(); + } + } + if (committed) { + messages->push_back(bytes); + } else { + (*uncommitted)++; + } } else { - sendBytes(set_offset, sink, "{\"offset\":" + std::to_string(latest) + "}"); + request.skip(); } - }); -} - -int streamPartitions(void*, const char* start, const char* end, SparkDsSink add_partition, - void* sink) { - return guarded([&] { - auto lo = static_cast(JsonParser(start).parse()["offset"].number); - auto hi = static_cast(JsonParser(end).parse()["offset"].number); - if (hi > lo) sendBytes(add_partition, sink, std::to_string(lo) + "," + std::to_string(hi)); - }); -} - -int serializeStreamReader(void* handle, SparkDsSink set_state, void* sink) { - return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); + } + return epoch_id; +} + +void commitWrite(Writer* writer, const std::string& message) { + std::string path = writer->state["path"]; + maybeFail(writer->state, "commit"); + std::vector temp_files; + int64_t uncommitted = 0; + int64_t epoch_id = parseTaskMessages(message, &temp_files, &uncommitted); + if (uncommitted > 0) throw std::runtime_error("commit got a task that did not commit"); + if (writer->state["overwrite"] == "true") { + for (const auto& file : listFiles(path, "part-")) unlink((path + "/" + file).c_str()); + } + for (const auto& temp_file : temp_files) { + std::string file_name = temp_file.substr(temp_file.rfind('/') + 1); + std::rename(temp_file.c_str(), (path + "/part-" + file_name).c_str()); + } + writeFile(path + (epoch_id < 0 ? "/_SUCCESS" : "/_EPOCH_" + std::to_string(epoch_id)), + std::to_string(temp_files.size())); + writer->closed_marker = closedMarker(epoch_id); } -int commitOffset(void* handle, const char* end) { - return guarded([&] { - std::string path = get(fromHandle(handle)->state, "commit_path", ""); - if (!path.empty()) writeFile(path, end); - }); +// Records the aborted write in _ABORTED for a batch write, or _ABORTED_ for a streaming +// one, with the number of tasks that did not commit. +void abortWrite(Writer* writer, const std::string& message) { + std::vector temp_files; + int64_t uncommitted = 0; + int64_t epoch_id = parseTaskMessages(message, &temp_files, &uncommitted); + writer->closed_marker = closedMarker(epoch_id); + for (const auto& temp_file : temp_files) unlink(temp_file.c_str()); + writeFile(writer->state["path"] + + (epoch_id < 0 ? "/_ABORTED" : "/_ABORTED_" + std::to_string(epoch_id)), + std::to_string(uncommitted)); +} + +void commitDataWriter(DataWriter* writer, SparkDsSink reply, void* sink) { + // If the commit fails, Spark aborts the data writer. + maybeFail(writer->state, "commitDataWriter"); + // With the option "commit_messages", the reply is passed that many times. + int64_t messages = getLong(writer->state, "commit_messages", 1); + for (int64_t i = 0; i < messages; i++) { + sendReply(reply, sink, writer->state, ProtoWriter().message( + kCommitDataWriter, ProtoWriter().bytes(1, writer->temp_path))); + } + std::fclose(writer->file); + writer->file = nullptr; + writer->committed = true; + std::string dir = writer->state["path"] + "/_streams"; + makeDirs(dir); + writeFile(dir + "/" + writer->temp_path.substr(writer->temp_path.rfind('/') + 1), + std::to_string(writer->streams) + "," + std::to_string(writer->batches)); } -void closeStreamReader(void* handle) { - std::unique_ptr reader(fromHandle(handle)); - guarded([&] { - std::string path = get(reader->state, "closed_path", ""); - if (!path.empty()) writeFile(path, ""); - }); +void abortDataWriter(DataWriter* writer) { + // Spark does not abort a data writer that committed: records it if it does. + if (writer->committed) { + writeFile(writer->state["path"] + "/_ABORTED_COMMITTED_DATA_WRITER", ""); + return; + } + if (writer->file != nullptr) { + std::fclose(writer->file); + writer->file = nullptr; + } + unlink(writer->temp_path.c_str()); } #endif // READ_ONLY -int readPartition(const SparkDsTaskContext* task, const uint8_t* reader_state, - int64_t reader_state_length, const uint8_t* partition, - int64_t partition_length, struct ArrowArrayStream* stream) { - return guarded([&] { - Options state = deserialize(toBytes(reader_state, reader_state_length)); - maybeFail(state, "open"); - // Returns success without initializing the stream, which Spark must reject. - if (get(state, "fail", "") == "uninitialized") return; - auto bounds = split(toBytes(partition, partition_length)); - if (bounds.size() != 2) throw std::runtime_error("invalid partition"); - recordTask(state, task); - stream->get_schema = streamGetSchema; - stream->get_next = streamGetNext; - stream->get_last_error = streamGetLastError; - stream->release = streamRelease; - stream->private_data = new Stream{split(get(state, "columns", "")), std::stoll(bounds[0]), - std::stoll(bounds[1]), getLong(state, "batch_size", 1000), - get(state, "fail", "") == "read", - get(state, "slice", "") == "true", get(state, "short", ""), - task, get(state, "wait_for_interrupt", ""), ""}; - }); -} - -#ifndef READ_ONLY - -int createWriter(void* handle, struct ArrowSchema* schema, const SparkDsWriteInfo* info, - void** out) { - return guarded([&] { - auto* source = fromHandle(handle); - if (source->name != name("native_sink")) { - throw std::runtime_error(source->name + " does not support writes"); - } - if (info->mode != SPARK_DS_WRITE_APPEND && info->mode != SPARK_DS_WRITE_TRUNCATE) { - throw std::runtime_error("unknown mode " + std::to_string(info->mode)); - } - // Spark releases the schema after the call, since it is not moved here. - if (schema->n_children == 0) throw std::runtime_error("no columns to write"); - // The location of a table in a catalog is passed as a URI. - std::string path = get(source->options, "path", ""); - if (path.rfind("file://", 0) == 0) { - path = path.substr(7); - } else if (path.rfind("file:", 0) == 0) { - path = path.substr(5); - } - if (path.empty()) throw std::runtime_error("the option 'path' is required"); - if (returnsZero(source->options, "create_writer")) return; - auto* writer = new Writer(); - writer->state["path"] = path; - writer->state["overwrite"] = info->mode == SPARK_DS_WRITE_TRUNCATE ? "true" : "false"; - writer->state["fail"] = get(source->options, "fail", ""); - writer->state["zero"] = get(source->options, "zero", ""); - writer->state["commit_messages"] = get(source->options, "commit_messages", "1"); - writer->state["task_path"] = get(source->options, "task_path", ""); - // Records that a writer was created, which Spark then commits or aborts, and the ID of the - // streaming query. - makeDirs(path); - writeFile(path + "/_CREATED", ""); - if (info->query_id != nullptr) writeFile(path + "/_QUERY_ID", info->query_id); - *out = writer; - }); -} - -int serializeWriter(void* handle, SparkDsSink set_state, void* sink) { - return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); -} - -int commitWrite(void* handle, int64_t epoch_id, const uint8_t* const* messages, - const int64_t* message_lengths, int64_t num_messages) { - return guarded([&] { - auto* writer = fromHandle(handle); - std::string path = writer->state["path"]; - maybeFail(writer->state, "commit"); - bool has_null = false; - auto temp_files = toByteStrings(messages, message_lengths, num_messages, &has_null); - if (has_null) throw std::runtime_error("commit got a null message"); - if (writer->state["overwrite"] == "true") { - for (const auto& file : listFiles(path, "part-")) unlink((path + "/" + file).c_str()); +void call(void* handle, const std::string& request, SparkDsSink reply, void* sink) { + ProtoReader message(request); + while (message.next()) { + int field = message.field(); + if (field < kInferSchema || field > kAbortDataWriter) { + // A request that the library does not know, which Spark does not send to it. + message.skip(); + continue; } - for (const auto& temp_file : temp_files) { - std::string file_name = temp_file.substr(temp_file.rfind('/') + 1); - std::rename(temp_file.c_str(), (path + "/part-" + file_name).c_str()); + std::string body = message.bytes(); + switch (field) { + case kInferSchema: return inferSchema(as(handle, "InferSchema"), reply, sink); + case kPlanPartitions: + return planPartitions(as(handle, "PlanPartitions"), reply, sink); + case kSerialize: { + auto* object = static_cast(handle); + if (auto* reader = dynamic_cast(object)) { + return serializeState(reader->state, reply, sink); + } + return serializeState(as(handle, "Serialize")->state, reply, sink); + } +#ifdef READ_ONLY + // A library without pushdown replies to it with empty replies: Spark pushes nothing. + case kPushPredicates: + case kPushLimit: + case kPruneColumns: + return; + default: + throw std::runtime_error("unsupported request " + std::to_string(field)); +#else + case kPushPredicates: + return pushPredicates(as(handle, "PushPredicates"), body, reply, sink); + case kPushLimit: return pushLimit(as(handle, "PushLimit"), body, reply, sink); + case kPruneColumns: + return pruneColumns(as(handle, "PruneColumns"), body, reply, sink); + case kInitialOffset: + return initialOffset(as(handle, "InitialOffset"), reply, sink); + case kLatestOffset: return latestOffset(as(handle, "LatestOffset"), reply, sink); + case kPlanStreamPartitions: + return planStreamPartitions(as(handle, "PlanStreamPartitions"), body, reply, sink); + case kCommitOffset: return commitOffset(as(handle, "CommitOffset"), body); + case kCommit: return commitWrite(as(handle, "Commit"), body); + case kAbort: return abortWrite(as(handle, "Abort"), body); + case kCommitDataWriter: + return commitDataWriter(as(handle, "CommitDataWriter"), reply, sink); + case kAbortDataWriter: + return abortDataWriter(as(handle, "AbortDataWriter")); +#endif } - writeFile(path + (epoch_id < 0 ? "/_SUCCESS" : "/_EPOCH_" + std::to_string(epoch_id)), - std::to_string(temp_files.size())); - writer->closed_marker = closedMarker(epoch_id); - }); + } } -// Records the aborted write in _ABORTED for a batch write, or _ABORTED_ for a streaming -// one, with the number of tasks that did not commit, whose messages are NULL. -int abortWrite(void* handle, int64_t epoch_id, const uint8_t* const* messages, - const int64_t* message_lengths, int64_t num_messages) { - return guarded([&] { - auto* writer = fromHandle(handle); - writer->closed_marker = closedMarker(epoch_id); - bool has_null = false; - for (const auto& temp_file : - toByteStrings(messages, message_lengths, num_messages, &has_null)) { - unlink(temp_file.c_str()); - } - int64_t null_messages = 0; - for (int64_t i = 0; i < num_messages; i++) { - if (messages[i] == nullptr) null_messages++; - } - writeFile(writer->state["path"] + - (epoch_id < 0 ? "/_ABORTED" : "/_ABORTED_" + std::to_string(epoch_id)), - std::to_string(null_messages)); - }); -} +// --------------------------------------------------------------------------------------------- +// spark_ds_stream +// --------------------------------------------------------------------------------------------- -void closeWriter(void* handle) { - std::unique_ptr writer(fromHandle(handle)); - guarded([&] { - if (!writer->closed_marker.empty()) { - writeFile(writer->state["path"] + writer->closed_marker, ""); - } - }); +void read(PartitionReader* partition, struct ArrowArrayStream* stream) { + const Options& state = partition->state; + // Returns success without initializing the stream, which Spark must reject. + if (get(state, "fail", "") == "uninitialized") return; + auto bounds = split(partition->partition); + if (bounds.size() != 2) throw std::runtime_error("invalid partition"); + stream->get_schema = streamGetSchema; + stream->get_next = streamGetNext; + stream->get_last_error = streamGetLastError; + stream->release = streamRelease; + stream->private_data = new Stream{split(get(state, "columns", "")), std::stoll(bounds[0]), + std::stoll(bounds[1]), getLong(state, "batch_size", 1000), + get(state, "fail", "") == "read", + get(state, "slice", "") == "true", get(state, "short", ""), + partition->task.task_attempt_id, + get(state, "wait_for_interrupt", ""), ""}; } -// Whether Spark uses a data writer that it closed, which the header forbids. -bool usedAfterClose(void* handle) { - std::lock_guard guard(closed_data_writers_lock); - auto closed = closed_data_writers.find(handle); - if (closed == closed_data_writers.end()) return false; - writeFile(closed->second + "/_DATA_WRITER_USED_AFTER_CLOSE", ""); - return true; -} +#ifndef READ_ONLY -int createDataWriter(const SparkDsTaskContext* task, const uint8_t* writer_state, - int64_t writer_state_length, int64_t epoch_id, void** out) { - return guarded([&] { - Options state = deserialize(toBytes(writer_state, writer_state_length)); - if (returnsZero(state, "create_data_writer")) return; - recordTask(state, task); - std::string dir = state["path"] + "/_temporary"; - makeDirs(dir); - std::string temp_path = dir + "/" + std::to_string(epoch_id) + "-" + - std::to_string(task->partition_id) + "-" + - std::to_string(task->task_attempt_id) + ".csv"; - FILE* file = std::fopen(temp_path.c_str(), "w"); - if (file == nullptr) throw std::runtime_error("cannot write " + temp_path); - makeDirs(state["path"] + "/_created_data_writers"); - writeFile(state["path"] + "/_created_data_writers/" + temp_path.substr(dir.size() + 1), ""); - auto* writer = new DataWriter{temp_path, file, state, false}; - { - // The allocator may reuse the address of a closed data writer. - std::lock_guard guard(closed_data_writers_lock); - closed_data_writers.erase(writer); +// Writes the batches of the stream. Spark releases the stream after the call, since it is not +// moved here, but the schema and the arrays that the stream returns are the library's to release. +void write(DataWriter* writer, struct ArrowArrayStream* stream) { + maybeFail(writer->state, "write"); + writer->streams++; + ArrowSchema schema; + if (stream->get_schema(stream, &schema) != 0) { + throw std::runtime_error(std::string("get_schema failed: ") + stream->get_last_error(stream)); + } + std::unique_ptr schema_owner( + &schema, [](ArrowSchema* s) { s->release(s); }); + while (true) { + ArrowArray array; + if (stream->get_next(stream, &array) != 0) { + throw std::runtime_error(std::string("get_next failed: ") + stream->get_last_error(stream)); } - *out = writer; - }); -} - -int writeBatch(void* handle, struct ArrowArray* array, struct ArrowSchema* schema) { - return guarded([&] { - auto* writer = fromHandle(handle); - // The batch is released here, so Spark does not release it after the call. + if (array.release == nullptr) break; // The end of the stream. std::unique_ptr array_owner( - array, [](ArrowArray* a) { a->release(a); }); - std::unique_ptr schema_owner( - schema, [](ArrowSchema* s) { s->release(s); }); - maybeFail(writer->state, "write"); - for (int64_t i = 0; i < array->length; i++) { + &array, [](ArrowArray* a) { a->release(a); }); + writer->batches++; + for (int64_t i = 0; i < array.length; i++) { std::string line; - for (int64_t c = 0; c < array->n_children; c++) { + for (int64_t c = 0; c < array.n_children; c++) { if (c > 0) line += ","; - line += formatValue(schema->children[c], array->children[c], array->offset + i); + line += formatValue(schema.children[c], array.children[c], array.offset + i); } line += "\n"; std::fputs(line.c_str(), writer->file); } - }); + } } -int commitDataWriter(void* handle, SparkDsSink set_message, void* sink) { - if (usedAfterClose(handle)) return 0; - return guarded([&] { - auto* writer = fromHandle(handle); - // If the commit fails, Spark aborts the data writer. - maybeFail(writer->state, "commitDataWriter"); - // With the option "commit_messages", the message is passed that many times. - int64_t messages = getLong(writer->state, "commit_messages", 1); - for (int64_t i = 0; i < messages; i++) sendBytes(set_message, sink, writer->temp_path); - std::fclose(writer->file); - writer->file = nullptr; - writer->committed = true; - }); -} +#endif // READ_ONLY -int abortDataWriter(void* handle) { - if (usedAfterClose(handle)) return 0; - return guarded([&] { - auto* writer = fromHandle(handle); - // Spark does not abort a data writer that committed: records it if it does. - if (writer->committed) { - writeFile(writer->state["path"] + "/_ABORTED_COMMITTED_DATA_WRITER", ""); - return; - } - if (writer->file != nullptr) { - std::fclose(writer->file); - writer->file = nullptr; - } - unlink(writer->temp_path.c_str()); - }); -} +// --------------------------------------------------------------------------------------------- +// spark_ds_close +// --------------------------------------------------------------------------------------------- -void closeDataWriter(void* handle) { - if (usedAfterClose(handle)) return; - std::unique_ptr writer(fromHandle(handle)); +void close(void* handle) { + std::unique_ptr object(static_cast(handle)); guarded([&] { - if (writer->file != nullptr) std::fclose(writer->file); - std::string path = writer->state["path"]; - std::string dir = path + "/_closed_data_writers"; - makeDirs(dir); - writeFile(dir + "/" + writer->temp_path.substr(writer->temp_path.rfind('/') + 1), ""); - std::lock_guard guard(closed_data_writers_lock); - closed_data_writers[handle] = path; +#ifndef READ_ONLY + if (auto* reader = dynamic_cast(object.get())) { + std::string path = get(reader->state, "closed_path", ""); + if (!path.empty()) writeFile(path, ""); + } else if (auto* writer = dynamic_cast(object.get())) { + if (!writer->closed_marker.empty()) { + writeFile(writer->state["path"] + writer->closed_marker, ""); + } + } else if (auto* data_writer = dynamic_cast(object.get())) { + if (data_writer->file != nullptr) std::fclose(data_writer->file); + std::string path = data_writer->state["path"]; + std::string dir = path + "/_closed_data_writers"; + makeDirs(dir); + std::string temp_path = data_writer->temp_path; + writeFile(dir + "/" + temp_path.substr(temp_path.rfind('/') + 1), ""); + std::lock_guard guard(closed_data_writers_lock); + closed_data_writers[handle] = path; + } +#endif }); } -#endif // READ_ONLY - -// The table, which NO_INIT leaves unused. -[[maybe_unused]] SparkDsLibrary makeTable() { - SparkDsLibrary table{}; - table.size = TABLE_SIZE; - table.abi_version = ABI_VERSION; -#ifndef NO_LAST_ERROR - table.last_error = lastError; -#endif - table.create = createSource; - table.schema = sourceSchema; - table.create_reader = createReader; - table.close = closeSource; - table.partitions = planPartitions; - table.serialize_reader = serializeReader; - table.close_reader = closeReader; - table.read = readPartition; +// spark_ds_open, which NO_OPEN does not export. +[[maybe_unused]] int open(const SparkDsHost* spark, const uint8_t* request, int64_t length, + SparkDsSink reply, void* sink, void** handle) { + return guarded([&] { + host = spark; + // The fields may come in any order, so they are all read before the object is opened. + std::string spark_version; + int object = 0; + std::string message; + ProtoReader reader(toBytes(request, length)); + while (reader.next()) { + if (reader.field() == kSparkVersion) { + spark_version = reader.bytes(); + } else if (reader.field() >= kDataSource && reader.field() <= kDataWriter) { + object = reader.field(); + message = reader.bytes(); + } else { + reader.skip(); + } + } + switch (object) { + case kDataSource: *handle = openDataSource(message, spark_version, reply, sink); break; + case kBatchReader: *handle = openBatchReader(message); break; + case kPartitionReader: *handle = openPartitionReader(message); break; #ifndef READ_ONLY - table.create_stream_reader = createStreamReader; - table.create_writer = createWriter; - table.push_predicates = pushPredicates; - table.push_limit = pushLimit; - table.prune_columns = pruneColumns; - table.initial_offset = initialOffset; - table.latest_offset = latestOffset; - table.stream_partitions = streamPartitions; - table.serialize_stream_reader = serializeStreamReader; - table.commit_offset = commitOffset; - table.close_stream_reader = closeStreamReader; - table.serialize_writer = serializeWriter; - table.commit = commitWrite; - table.abort = abortWrite; - table.close_writer = closeWriter; - table.create_data_writer = createDataWriter; - table.write = writeBatch; - table.commit_data_writer = commitDataWriter; - table.abort_data_writer = abortDataWriter; - table.close_data_writer = closeDataWriter; + case kStreamReader: *handle = openStreamReader(message); break; + case kWriter: *handle = openWriter(message); break; + case kDataWriter: *handle = openDataWriter(message); break; #endif - return table; + default: { + const char* kinds[] = {"data source", "batch reader", "stream reader", "writer", + "partition reader", "data writer"}; + bool known = object >= kDataSource && object <= kDataWriter; + throw std::runtime_error(std::string("the library cannot open a ") + + (known ? kinds[object - kDataSource] : "new kind of object")); + } + } + }); } } // namespace @@ -1400,19 +1796,45 @@ void closeDataWriter(void* handle) { // spark_datasource.h // --------------------------------------------------------------------------------------------- -#ifndef NO_INIT +int32_t spark_ds_abi_version(void) { return ABI_VERSION; } -int spark_ds_init(const SparkDsHost* spark, const SparkDsLibrary** library) { - host = spark; - if (INIT_FAILURE == 1) { - hostLog(SPARK_DS_LOG_ERROR, "injected failure in spark_ds_init"); - return 1; - } - // Succeeds without returning a table, which Spark must reject. - if (INIT_FAILURE == 2) return 0; - static const SparkDsLibrary table = makeTable(); - *library = &table; - return 0; +#ifndef NO_LAST_ERROR +const char* spark_ds_last_error(void) { return last_error.c_str(); } +#endif + +#ifndef NO_OPEN + +int spark_ds_open(const SparkDsHost* spark, const uint8_t* request, int64_t length, + SparkDsSink reply, void* sink, void** handle) { + return open(spark, request, length, reply, sink, handle); +} + +#endif // NO_OPEN + +int spark_ds_call(void* handle, const uint8_t* request, int64_t length, SparkDsSink reply, + void* sink) { + if (usedAfterClose(handle)) return 0; + return guarded([&] { call(handle, toBytes(request, length), reply, sink); }); } -#endif // NO_INIT +int spark_ds_stream(void* handle, const uint8_t* request, int64_t length, + [[maybe_unused]] struct ArrowArrayStream* in, struct ArrowArrayStream* out) { + if (usedAfterClose(handle)) return 0; + return guarded([&] { + ProtoReader message(toBytes(request, length)); + while (message.next()) { + int field = message.field(); + message.skip(); + if (field == kRead) return read(as(handle, "Read"), out); +#ifndef READ_ONLY + if (field == kWrite) return write(as(handle, "Write"), in); +#endif + } + throw std::runtime_error("unsupported stream request"); + }); +} + +void spark_ds_close(void* handle) { + if (usedAfterClose(handle)) return; + close(handle); +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala index 3776ddf633f8c..133d04946b796 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -48,6 +48,7 @@ import org.apache.spark.sql.execution.datasources.DataSource import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2ScanRelation, DataSourceV2Strategy} import org.apache.spark.sql.execution.datasources.v2.columnar.ColumnarScan import org.apache.spark.sql.execution.datasources.v2.python.{PythonDataSourceV2, UserDefinedPythonDataSource} +import org.apache.spark.sql.functions.{col, udf} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.streaming.StreamingQuery import org.apache.spark.sql.test.SharedSparkSession @@ -78,14 +79,8 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { compile("test_native_datasource_v2", "-DABI_VERSION=2", "-DNAME_PREFIX=\"v2_\"") private lazy val noLastErrorLibrary = compile("test_native_datasource_ne", "-DNO_LAST_ERROR", "-DNAME_PREFIX=\"ne_\"") - private lazy val smallTableLibrary = compile("test_native_datasource_small", - "-DTABLE_SIZE=offsetof(SparkDsLibrary,read)", "-DNAME_PREFIX=\"small_\"") - private lazy val failingInitLibrary = - compile("test_native_datasource_if", "-DINIT_FAILURE=1", "-DNAME_PREFIX=\"if_\"") - private lazy val noTableLibrary = - compile("test_native_datasource_nt", "-DINIT_FAILURE=2", "-DNAME_PREFIX=\"nt_\"") - private lazy val noInitLibrary = - compile("test_native_datasource_ni", "-DNO_INIT", "-DNAME_PREFIX=\"ni_\"") + private lazy val noOpenLibrary = + compile("test_native_datasource_no", "-DNO_OPEN", "-DNAME_PREFIX=\"no_\"") private lazy val installedLibrary = compile("test_native_datasource_installed", "-DNAME_PREFIX=\"installed_\"") private lazy val installedNegatingLibrary = compile( @@ -218,7 +213,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { """"children":[{"type":"column","name":["id"]},""" + """{"type":"literal","dataType":"bigint","value":1},""" + """{"type":"literal","dataType":"bigint","value":2}]}]}""")) - // Only the functions listed in push_predicates are passed, so not arithmetic. + // Only the functions listed in PushPredicates are passed, so not arithmetic. NativePredicates.FUNCTION_NAMES.foreach { name => assert(toJson(predicate(name, id, id)).nonEmpty, name) } @@ -278,7 +273,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { /** * Generates `org/apache/spark/sql/datasource/predicate-cases.json`, which lists predicates as - * Spark passes them to `push_predicates`, with the rows of a table for which Spark + * Spark passes them in `PushPredicates`, with the rows of a table for which Spark * finds them true, for libraries to test their pushdown with. Regenerate it with * `SPARK_GENERATE_GOLDEN_FILES=1`. */ @@ -334,9 +329,9 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } val generated = s"""{ - | "description": ${json("Predicates as Spark passes them to push_predicates of " + - "spark_datasource.h, with the ids of the rows of the table for which they are true. A " + - "case lists the predicates that Spark pushes for the filter in its SQL, which are " + + | "description": ${json("Predicates as Spark passes them in PushPredicates of " + + "spark_datasource.proto, with the ids of the rows of the table for which they are true. " + + "A case lists the predicates that Spark pushes for the filter in its SQL, which are " + "combined with AND. A row has a value for each column of the schema; NaN and infinite " + "doubles are strings.")}, | "schema": ${json(schema.toDDL)}, @@ -363,20 +358,11 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { "org/apache/spark/sql/datasource/predicate-cases.json") != null) } - test("tell missing functions from link errors") { - // A function that the library leaves NULL in its table is an unsupported operation. - val missing = new NativeLibrary.MissingFunctionException("push_limit") - assert(NativeCalls.optional(throw missing).isEmpty) - checkError( - exception = intercept[SparkRuntimeException](NativeCalls.call("x", "read")(throw missing)), - condition = "NATIVE_DATA_SOURCE_ERROR", - parameters = Map( - "action" -> "read", "name" -> "x", "msg" -> "The library does not implement push_limit.")) + test("report link errors") { // Link errors are reported as they are, such as when Arrow, which still calls the Arrow C // structs through JNI, cannot load its JNI library, which it extracts to java.io.tmpdir, for // example if the directory does not allow it. val linkError = new UnsatisfiedLinkError("/tmp/libarrow_cdata_jni.so: failed to map segment") - assert(intercept[UnsatisfiedLinkError](NativeCalls.optional(throw linkError)) eq linkError) checkError( exception = intercept[SparkRuntimeException](NativeCalls.call("x", "read")(throw linkError)), condition = "NATIVE_DATA_SOURCE_ERROR", @@ -775,7 +761,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { parameters = Map( "action" -> "write to", "name" -> "native_sink", - "msg" -> "commit_data_writer returned 2 values instead of one.")) + "msg" -> "spark_ds_call replied 2 times instead of once.")) // The task failed without aborting the committed data writer, which it closed, and the driver // aborted the write. assert(!new File(dir, "_ABORTED_COMMITTED_DATA_WRITER").exists()) @@ -784,6 +770,42 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) } + nativeTest("write the data of a task in one stream") { + val session = newSession(defaultPackage) + session.conf.set(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key, "3") + val dir = Utils.createTempDir() + session.range(10).repartition(1).write.format("native_sink").option("path", dir.getPath) + .mode("append").save() + // The library reads the 10 rows of the task in one stream, in batches of at most 3 rows. + val streams = new File(dir, "_streams").listFiles() + assert(streams.length == 1) + assert(Files.readString(streams.head.toPath) == "1,4") + } + + nativeTest("report the errors of the data that a library writes") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val failing = udf { (id: Long) => + if (id == 5) throw new IllegalStateException("injected failure in the data") + id + } + val e = intercept[Exception] { + session.range(10).repartition(1).select(failing(col("id")).as("id")).write + .format("native_sink").option("path", dir.getPath).mode("append").save() + } + // The library reads the data through the stream, which only gets the message of the error: + // the error itself is reported, rather than the error that the library returns. + val causes = Iterator.iterate[Throwable](e)(_.getCause).takeWhile(_ != null).toSeq + assert(causes.exists(c => c.isInstanceOf[IllegalStateException] && + c.getMessage == "injected failure in the data"), e) + assert(!causes.exists { + case t: SparkThrowable => t.getCondition == "NATIVE_DATA_SOURCE_ERROR" + case _ => false + }, e) + checkDataWritersClosed(dir) + assert(new File(dir, "_ABORTED").isFile) + } + nativeTest("abort a micro-batch whose write failed") { val session = newSession(defaultPackage) val dir = Utils.createTempDir() @@ -870,10 +892,10 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { condition = "NATIVE_DATA_SOURCE_ERROR", parameters = Map("action" -> "plan a scan of", "name" -> "native_counter", "msg" -> msg)) } - check("none", "latest_offset returned no offset.") - check("multiline", "latest_offset returned an offset with a line break, but an " + + check("none", "LatestOffset returned no offset.") + check("multiline", "LatestOffset returned an offset with a line break, but an " + "offset must be JSON on a single line.") - check("invalid", "latest_offset returned an offset that is not valid JSON: offset 10") + check("invalid", "LatestOffset returned an offset that is not valid JSON: offset 10") } nativeTest("pass the ID of a streaming query to its stream writers") { @@ -935,9 +957,10 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { // declares batch reads that it does not support. check("plan a scan of", "native_sink does not support batch reads", "native_sink") { session => session.read.schema("id BIGINT").format("native_sink") - .option("capabilities", NativeLibrary.BATCH_READ.toString).load().collect() + .option("capabilities", proto.Capability.CAPABILITY_BATCH_READ.getNumber.toString) + .load().collect() } - // Without last_error, only the value that the function returned is reported. + // Without spark_ds_last_error, only the value that the function returned is reported. val e = intercept[Exception] { newSession(createPackage(noLastErrorLibrary, dataSources("ne_"))) .read.format("ne_native_range").option("fail", "create").load() @@ -948,7 +971,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { parameters = Map( "action" -> "create", "name" -> "ne_native_range", - "msg" -> "create failed with the error code 1.")) + "msg" -> "spark_ds_open failed with the error code 1.")) } nativeTest("a scan whose planning failed fails again with the same error") { @@ -981,7 +1004,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { parameters = Map( "action" -> action, "name" -> name, - "msg" -> s"$function returned a NULL handle.")) + "msg" -> "spark_ds_open returned a NULL handle.")) } def runStream(start: => StreamingQuery): Unit = { val query = start @@ -1059,15 +1082,19 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { check("DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> "native_counter")) { session.read.format("native_counter").load().collect() } - // A data source that declares a capability must implement its functions. + // A data source that declares a capability must support its operations: here, the library + // cannot open a writer. check("NATIVE_DATA_SOURCE_ERROR", Map("action" -> "plan a write to", - "name" -> "ro_native_sink", "msg" -> "The library does not implement create_writer.")) { + "name" -> "ro_native_sink", "msg" -> "the library cannot open a writer")) { session.range(1).write.format("ro_native_sink") - .option("capabilities", NativeLibrary.BATCH_WRITE.toString).mode("append").save() + .option("capabilities", proto.Capability.CAPABILITY_BATCH_WRITE.getNumber.toString) + .mode("append").save() } - // Replacing the existing data requires the capability SPARK_DS_TRUNCATE. + // Replacing the existing data requires the capability CAPABILITY_TRUNCATE. val dir = Utils.createTempDir() - val noTruncate = (NativeLibrary.BATCH_WRITE | NativeLibrary.STREAMING_WRITE).toString + val noTruncate = + Seq(proto.Capability.CAPABILITY_BATCH_WRITE, proto.Capability.CAPABILITY_STREAMING_WRITE) + .map(_.getNumber).mkString(",") check("NATIVE_DATA_SOURCE_ERROR", Map("action" -> "plan a write to", "name" -> "native_sink", "msg" -> "The data source does not support the save mode overwrite.")) { session.range(1).write.format("native_sink").option("path", dir.getPath) @@ -1088,19 +1115,32 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { assert(!new File(dir, "_CREATED").exists()) } - nativeTest("read only the functions that fit in the table of a library") { - // The table of the library ends before read, as the table of a library built against an - // earlier header ends before the functions added later: Spark does not read past its end. - val session = newSession(createPackage(smallTableLibrary, dataSources("small_"))) - val df = session.read.format("small_native_range").load() - assert(df.schema.length == 7) - checkError( - exception = findError(intercept[Exception](df.collect()), "NATIVE_DATA_SOURCE_ERROR"), - condition = "NATIVE_DATA_SOURCE_ERROR", - parameters = Map( - "action" -> "plan a scan of", - "name" -> "small_native_range", - "msg" -> "The library does not implement read.")) + nativeTest("ignore the fields of messages that they do not know") { + // A library built against a later version of the messages may reply with fields that Spark + // does not know: here, every reply has one. + val session = newSession(defaultPackage) + checkRows(range(session, "end" -> "6", "unknown_fields" -> "true").filter("id > 2"), + (3L until 6L).map(expectedRow)) + // And a library built against an earlier version gets requests with fields that it does not + // know: here, a request of a later version, with a field that the library does not know. + val (name, location) = + NativeDataSourceRegistry.lookup("native_range", session.sessionState.conf).get + val library = NativeLibraries.get(location) + val known = NativeMessages.open() + .setDataSource(proto.DataSource.newBuilder().setName(name)) + .build() + // A message with a field of a later version: field 1000, with a string. Spark keeps the fields + // that it does not know of a message that it parses. + val unknown = Array[Byte](0xc2.toByte, 0x3e, 3) ++ "new".getBytes(UTF_8) + val request = proto.OpenRequest.parseFrom(known.toByteArray ++ unknown) + assert(request.toByteArray.length == known.toByteArray.length + unknown.length) + val opened = library.open(request) + try { + assert(opened.reply.getCapabilitiesList.asScala == + Seq(proto.Capability.CAPABILITY_BATCH_READ)) + } finally { + library.close(opened.handle) + } } nativeTest("pass the version of Spark, and log with the log of Spark") { @@ -1115,8 +1155,7 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { parameters = Map( "action" -> "create", "name" -> "native_range", - "msg" -> (s"the version of Spark is $SPARK_VERSION, which supports version 1 of the " + - "interface"))) + "msg" -> s"the version of Spark is $SPARK_VERSION")) val logAppender = new LogAppender("native data source logs") withLogAppender(logAppender, loggerNames = Seq(classOf[NativeLibrary].getName)) { @@ -1132,19 +1171,18 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { nativeTest("pass the task to reads and data writers") { val session = newSession(defaultPackage) - // Checks the SparkDsTaskContext that the library got in each task, by partition. + // Checks the TaskContext that the library got in each task, by partition. def check(dir: File): Unit = { val tasks = dir.listFiles().map { file => file.getName.toInt -> Files.readString(file.toPath).split(",").map(_.toLong).toSeq }.toMap assert(tasks.keySet == Set(0, 1)) tasks.foreach { - case (partition, Seq(size, stageId, stageAttempt, partitionId, attempt, taskAttemptId)) => - assert(size == NativeLibrary.layout().get("SparkDsTaskContext").longValue) + case (partition, Seq(stageId, stageAttempt, partitionId, attempt, taskAttemptId)) => assert(partitionId == partition) assert(stageId >= 0 && stageAttempt == 0 && attempt == 0 && taskAttemptId >= 0) } - assert(tasks.values.map(_(5)).toSet.size == 2) + assert(tasks.values.map(_(4)).toSet.size == 2) } val readTasks = Utils.createTempDir() checkRows(range(session, "end" -> "6", "task_path" -> readTasks.getPath).select("id"), @@ -1724,33 +1762,19 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } } - nativeTest("reject the libraries that spark_ds_init does not initialize") { - def check(library: File, prefix: String, reason: String): Unit = { - val pkg = createPackage(library, dataSources(prefix)) - // The library stays loaded, and the error is the same when it is used again: Spark - // initializes a library only once. - (1 to 2).foreach { _ => - checkError( - exception = intercept[SparkRuntimeException] { - newSession(pkg).read.format(prefix + "native_range").load() - }, - condition = "INVALID_NATIVE_DATA_SOURCE_PACKAGE.INVALID_LIBRARY", - parameters = Map( - "path" -> pkg.getCanonicalPath, - "library" -> s"${NativePlatform.current}/${library.getName}", - "reason" -> reason)) - } + nativeTest("reject the libraries that do not export the functions of the interface") { + val pkg = createPackage(noOpenLibrary, dataSources("no_")) + // The library stays loaded, and the error is the same when it is used again. + (1 to 2).foreach { _ => + checkError( + exception = intercept[SparkRuntimeException] { + newSession(pkg).read.format("no_native_range").load() + }, + condition = "INVALID_NATIVE_DATA_SOURCE_PACKAGE.INVALID_LIBRARY", + parameters = Map( + "path" -> pkg.getCanonicalPath, + "library" -> s"${NativePlatform.current}/${noOpenLibrary.getName}", + "reason" -> "The library does not export the function spark_ds_open.")) } - val logAppender = new LogAppender("spark_ds_init") - withLogAppender(logAppender, loggerNames = Seq(classOf[NativeLibrary].getName)) { - check(failingInitLibrary, "if_", "spark_ds_init failed with the error code 1.") - } - // The library logged why it failed, once. - assert(logAppender.loggingEvents.count { event => - event.getLevel == org.apache.logging.log4j.Level.ERROR && - event.getMessage.getFormattedMessage.contains("injected failure in spark_ds_init") - } == 1) - check(noTableLibrary, "nt_", "spark_ds_init returned no function table.") - check(noInitLibrary, "ni_", "The library does not export the function spark_ds_init.") } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala index 938f838c6f6bb..4b7e3cf55a49c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala @@ -73,7 +73,8 @@ object NativeDataSourceTestUtils { Utils.tryWithResource(in)(Files.copy(_, source.toPath)) } val library = new File(buildDir, System.mapLibraryName(name)) - // The symbols are hidden, as libraries are often built, so only spark_ds_init is exported. + // The symbols are hidden, as libraries are often built, so only the functions of the header + // are exported. run(Seq(compiler, "-std=c++17", "-O1", "-shared", "-fPIC", "-fvisibility=hidden") ++ includes.map("-I" + _) ++ defines ++ Seq("-o", library.getPath, source.getPath)) library From 85cc900328498be9e8f61096f47f3d524f4071f6 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Thu, 15 Oct 2026 15:28:00 +0900 Subject: [PATCH 27/27] Exclude the messages of native data sources from the Python code of the streaming protos dev/streaming-gen-protos.sh generates the Python code of every message of sql/core with buf, and the CI checks that it matches python/pyspark/sql/streaming/proto. PySpark does not use the messages of native data sources, so the buf module of sql/core excludes them. Co-authored-by: Isaac --- sql/core/src/main/protobuf/buf.yaml | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 sql/core/src/main/protobuf/buf.yaml diff --git a/sql/core/src/main/protobuf/buf.yaml b/sql/core/src/main/protobuf/buf.yaml new file mode 100644 index 0000000000000..986dbaa5d43a2 --- /dev/null +++ b/sql/core/src/main/protobuf/buf.yaml @@ -0,0 +1,22 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +version: v1 +build: + excludes: + # The messages of native data sources, which PySpark does not use: dev/streaming-gen-protos.sh + # only generates the Python code of the messages of the Python streaming APIs. + - org/apache/spark/sql/datasource