From 56acc8bea7362d8618d1b3939fd2b98122c41401 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 9 Oct 2026 10:09:08 +0900 Subject: [PATCH 01/15] [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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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/15] 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 aa0f3ba4dd770c0527c0d26bec45d2a41bca6638 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Wed, 14 Oct 2026 08:03:39 +0900 Subject: [PATCH 14/15] Address review comments on predicates, strings, writes, offsets and builds of native data sources - Spark evaluates all the pushed predicates again after reading: the answers of a library only show in the plan. NativeBridge states the rule for skipping data (NaN, -0.0, nulls, UTF-8 order, parts of predicates), and the jar contains predicate-cases.json, generated by a test from Spark's own evaluation. - The JSON of the predicates is ASCII, and NativeBridge documents how strings are encoded. - createStreamWriter gets the ID of the streaming query. - Offsets must be compact JSON on a single line: null, invalid JSON and line breaks are rejected. - 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 in a session are serialized. - Handles of 0, retried planning, unsupported Arrow types and the link errors of Arrow are reported as NATIVE_DATA_SOURCE_ERROR. - The column vectors of a native read are not freeable between batches, which broke structs. - Packages of an unsupported version are logged. - PySpark addArtifact only passes file: URIs as they are. - Docs: the writer lifecycle, close errors, optional commitOffset, Arrow variants, a Memory section, broadcasts with AQE, replaced libraries, and the addArtifact Scaladoc. Co-authored-by: Isaac --- .../resources/error/error-conditions.json | 5 + docs/sql-native-data-sources.md | 98 ++++- python/pyspark/sql/session.py | 10 +- python/pyspark/sql/tests/test_artifact.py | 7 +- .../org/apache/spark/sql/SparkSession.scala | 17 +- .../spark/sql/datasource/NativeBridge.java | 134 ++++-- .../datasources/v2/ffi/NativeLibrary.java | 38 +- .../spark/sql/datasource/predicate-cases.json | 44 ++ .../v2/columnar/ColumnarDataSource.scala | 8 +- .../v2/columnar/ColumnarTable.scala | 5 +- .../v2/columnar/ColumnarWrite.scala | 11 +- .../datasources/v2/ffi/NativeDataSource.scala | 147 +++++-- .../v2/ffi/NativeDataSourcePackage.scala | 93 ++-- .../v2/ffi/NativeDataSourceRegistry.scala | 10 +- .../datasources/v2/ffi/NativePredicates.scala | 8 +- .../test_native_datasource.cc | 143 ++++-- .../v2/ffi/NativeDataSourceClusterSuite.scala | 9 +- .../v2/ffi/NativeDataSourceSuite.scala | 406 +++++++++++++++++- 18 files changed, 1009 insertions(+), 184 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 f08fcf3748be4..3dfcff8ed728f 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 8dcfa59694fd1..4572905c2de64 100644 --- a/docs/sql-native-data-sources.md +++ b/docs/sql-native-data-sources.md @@ -57,24 +57,41 @@ Spark calls the functions of the library as follows: 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. + 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 null, 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 `createWriter` or `createStreamWriter`. 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 `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 streaming write commits each micro-batch with its epoch ID, and + `createStreamWriter` 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. 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`. +reports it as a `NATIVE_DATA_SOURCE_ERROR`. Spark only logs the errors of `closeReader`, +`closeStreamReader` and `closeWriter`, which can run after the query, so do the work that can fail, +such as flushing data, in `commit`. + +The strings of the interface, such as the options, the predicates and the offsets, are Java +strings. The JNI functions `GetStringUTFChars` and `NewStringUTF` convert them from and to modified +UTF-8, which differs from UTF-8 for the character U+0000 and for the characters outside the Basic +Multilingual Plane, such as emoji: convert them where it matters. The `jni` crate of Rust converts +them to and from 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 pushdown and `serialize` functions, which -are optional: +Spark treats the operation as unsupported, except for the pushdown functions, the `serialize` +functions and `commitOffset`, which are optional: | Operation | Required functions | Optional functions | |-----------|--------------------|--------------------| @@ -86,13 +103,25 @@ 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. 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 `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 @@ -108,7 +137,7 @@ library returns and that the manifest of a package declares. Within a version, t 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 +a library must ignore if it does not know them. 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. @@ -117,13 +146,29 @@ Spark 4 from 4.4.0. 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 @@ -172,6 +217,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 @@ -489,6 +539,8 @@ void Check(const arrow::Status& status) { if (!status.ok()) throw std::runtime_error(status.ToString()); } +// Returns a Java string as modified UTF-8, which is UTF-8 except for U+0000 and the characters +// outside the Basic Multilingual Plane: enough for the option of this data source. std::string ToString(JNIEnv* env, jobject value) { const char* chars = env->GetStringUTFChars(static_cast(value), nullptr); std::string result(chars); 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/datasource/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java index ad8b7fdb06e8f..1500e1c6c4344 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 @@ -51,21 +51,28 @@ * 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. + * and closes it. Then it calls {@link #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 * {@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 + * calls {@link #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 null, 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 + * {@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, 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}. 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 {@link #commit}. 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 @@ -79,7 +86,10 @@ * 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 + * committed must have no effect, for exactly-once output. The ID of the query, which + * {@link #createStreamWriter} 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 @@ -89,16 +99,30 @@ *

    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.
      • + * pointer: Spark reports a handle of 0 as an error. 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 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. *
    • 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}.
      • + * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}, except for the errors of + * {@link #closeReader}, {@link #closeStreamReader} and {@link #closeWriter}, which Spark only + * logs: they may run after the query, from another thread. So do the work that can fail a + * query, such as flushing data, in {@link #commit} rather than when closing. + *
    • Strings. The strings of this interface, such as the name and the options of + * {@link #createDataSource}, the predicates, the column names of {@link #pruneColumns} and the + * offsets, are Java strings. The JNI functions {@code GetStringUTFChars} and + * {@code NewStringUTF} convert them from and to modified UTF-8, which is UTF-8 except for the + * character U+0000 and the characters outside the Basic Multilingual Plane, such as emoji, which + * modified UTF-8 encodes in 6 bytes rather than 4. Convert the strings where it matters, for + * example with the Java method {@code String.getBytes} and the charset UTF-8. The {@code jni} + * crate of Rust converts them to and from UTF-8. The JSON of the predicates is ASCII: it escapes + * the other characters, which any JSON parser decodes.
      • *
    • 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 @@ -115,16 +139,18 @@ *
    • 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.
      • + * {@code timestamp[us]} without one. Data read from a library may also use {@code large_utf8} + * and {@code utf8_view} for strings, {@code large_binary} and {@code 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, {@code 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 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.
      • + * the pushdown methods, the {@code serialize} methods and {@link #commitOffset}: then Spark + * pushes nothing down, sends no state and does not tell the library which offsets it + * processed. *
    • 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 @@ -179,8 +205,8 @@ static void load(String path) { * 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 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. * @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 @@ -230,10 +256,13 @@ public static native long createDataSource( * @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 + * @param queryId 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. * @return the writer handle */ public static native long createStreamWriter( - long dataSource, long schemaAddress, boolean overwrite); + long dataSource, long schemaAddress, boolean overwrite, String queryId); /** Releases a data source handle. */ public static native void closeDataSource(long dataSource); @@ -268,23 +297,41 @@ public static native long createStreamWriter( * 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. + * {@link #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. *

      - * 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. + * 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 {@code -0.0} is equal to {@code 0.0}. IEEE 754 comparisons + * treat NaN differently, and total orders, such as {@code f64::total_cmp} in Rust or the + * comparison kernels of arrow-rs, also order {@code -0.0} before {@code 0.0}. For example, + * {@code d > 1.0} is true for a NaN value of {@code d}, and {@code d = 0.0} for + * {@code -0.0}.
        • + *
      • Comparisons with null are null, and a row only matches if the predicates are true. For + * example, {@code 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, such as + * {@code String.compareTo} in Java.
        • + *
      • 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 {@code NOT}, where that would drop rows.
        • + *
      + * The file {@code org/apache/spark/sql/datasource/predicate-cases.json} in the jar of Spark SQL + * lists predicates as Spark passes them to this method, with the rows of a table for which they + * are true, to test a reader with. * - * @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. + * @return for each predicate, whether the reader evaluates it completely. Spark shows these + * predicates as pushed in the plan of the query, and evaluates all the predicates again + * after reading anyway. */ public static native boolean[] pushPredicates(long reader, String[] predicates); /** - * 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. * * @return whether the reader returns at most {@code limit} rows in total. Spark applies the * limit again after reading either way. @@ -294,7 +341,9 @@ public static native long createStreamWriter( /** * Prunes the columns to read. * - * @param columnNames the names of the top-level columns that the query needs + * @param columnNames 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 + * {@code SELECT count(*)}: then the batches only have a number of rows. * @return whether {@link #read} returns exactly these columns, in this order. If not, it * returns all the columns of the schema to read. */ @@ -320,10 +369,13 @@ public static native long createStreamWriter( // Stream reader (driver) // --------------------------------------------------------------------------------------------- - /** Returns the offset to start from when the streaming query has no checkpoint, as JSON. */ + /** + * Returns the offset to start from when the streaming query has no checkpoint, as compact JSON + * on a single line. + */ public static native String initialOffset(long streamReader); - /** Returns the most recent offset available, as JSON. */ + /** Returns the most recent offset available, as compact JSON on a single line. */ public static native String latestOffset(long streamReader); /** 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..c073f41e9a9a9 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 @@ -36,12 +36,34 @@ * 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}. + * {@link MissingFunctionException}. */ final class NativeLibrary { private static final String BRIDGE_CLASS = NativeBridge.class.getName(); + private static final MethodHandle MISSING_FUNCTION; + + static { + try { + MISSING_FUNCTION = MethodHandles.lookup().findStatic(NativeLibrary.class, "missingFunction", + MethodType.methodType(Object.class, UnsatisfiedLinkError.class)); + } catch (ReflectiveOperationException e) { + throw new IllegalStateException(e); + } + } + + /** Thrown when Spark calls a function that the library does not export. */ + static final class MissingFunctionException extends RuntimeException { + MissingFunctionException(UnsatisfiedLinkError cause) { + super(cause.getMessage(), cause); + } + } + + private static Object missingFunction(UnsatisfiedLinkError e) { + throw new MissingFunctionException(e); + } + private final MethodHandle abiVersion; private final MethodHandle createDataSource; private final MethodHandle schema; @@ -82,7 +104,8 @@ private NativeLibrary(Class bridge) throws ReflectiveOperationException { 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); + createStreamWriter = + find(bridge, "createStreamWriter", jlong, jlong, jlong, boolean.class, String.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); @@ -111,8 +134,13 @@ private NativeLibrary(Class bridge) throws ReflectiveOperationException { private static MethodHandle find( Class bridge, String name, Class returnType, Class... parameterTypes) throws ReflectiveOperationException { - return MethodHandles.publicLookup() + MethodHandle method = MethodHandles.publicLookup() .findStatic(bridge, name, MethodType.methodType(returnType, parameterTypes)); + // The JVM throws UnsatisfiedLinkError when it cannot link the native method, because the + // library does not export it. Other code also throws it, such as Arrow when it cannot load its + // own JNI library, so only the error of the method itself means that the function is missing. + return MethodHandles.catchException(method, UnsatisfiedLinkError.class, + MISSING_FUNCTION.asType(MethodType.methodType(returnType, UnsatisfiedLinkError.class))); } /** @@ -200,9 +228,9 @@ long createWriter(long dataSource, long schemaAddress, boolean overwrite) throws return (long) createWriter.invokeExact(dataSource, schemaAddress, overwrite); } - long createStreamWriter(long dataSource, long schemaAddress, boolean overwrite) + long createStreamWriter(long dataSource, long schemaAddress, boolean overwrite, String queryId) throws Throwable { - return (long) createStreamWriter.invokeExact(dataSource, schemaAddress, overwrite); + return (long) createStreamWriter.invokeExact(dataSource, schemaAddress, overwrite, queryId); } void closeDataSource(long dataSource) throws Throwable { 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..ea2de943ca89f --- /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 NativeBridge.pushPredicates, 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/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 e2307c70c3a42..d748e0cb4287b 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 @@ -70,9 +73,9 @@ 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)) + val handle = create(library, "plan a scan of", "createReader", + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { dataSource => + NativeArrow.exportSchema(schema)(address => library.createReader(dataSource, address)) } new NativeDataSourceReader( executorLocation, name, schema, new NativeHandle(handle, library.closeReader)) @@ -81,9 +84,9 @@ 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") { - dataSource => - NativeArrow.exportSchema(schema)(address => library.createStreamReader(dataSource, address)) + val handle = create(library, "plan a scan of", "createStreamReader", + "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { dataSource => + NativeArrow.exportSchema(schema)(address => library.createStreamReader(dataSource, address)) } new NativeDataSourceStreamReader( executorLocation, name, schema, new NativeHandle(handle, library.closeStreamReader)) @@ -93,7 +96,8 @@ 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 => + create(library, "plan a write to", "createWriter", + "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource => NativeArrow.exportSchema(schema) { address => library.createWriter(dataSource, address, overwrite) } @@ -109,11 +113,11 @@ class NativeDataSource( 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) - } + create(library, "plan a write to", "createStreamWriter", + "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { dataSource => + NativeArrow.exportSchema(schema) { address => + library.createStreamWriter(dataSource, address, overwrite, queryId) + } }, library.closeWriter) val writer = new NativeDataSourceStreamWriter( @@ -126,21 +130,30 @@ class NativeDataSource( 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)) + val dataSource = checkHandle(name, "create", "createDataSource") { + 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 = { + /** + * Creates a reader or a writer with `f`, which calls the given function of the library, and + * fails if the library does not implement it. + */ + private def create( + library: NativeLibrary, + action: String, + function: String, + unsupportedError: String)(f: Long => Long): Long = { withDataSource(library, action) { dataSource => - call(name, action)(optional(f(dataSource))).getOrElse { + val handle = call(name, action)(optional(f(dataSource))).getOrElse { throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) } + checkHandle(name, action, function)(handle) } } @@ -168,6 +181,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 +225,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 +267,36 @@ 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("initialOffset", 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("latestOffset", 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 == null) { + throw error(name, "plan a scan of", s"$function returned null instead of an 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, " + + "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 +320,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, @@ -454,7 +508,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. @@ -473,7 +536,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, @@ -585,6 +648,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, @@ -599,8 +673,10 @@ 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. - private var handle: Long = call(name, "write to") { - library.createDataWriter(writerState, partitionId, taskId, epochId) + private var handle: Long = checkHandle(name, "write to", "createDataWriter") { + call(name, "write to") { + library.createDataWriter(writerState, partitionId, taskId, epochId) + } } override def write(batch: ColumnarBatch): Unit = { @@ -714,20 +790,31 @@ 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 } 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) => + case e @ (NonFatal(_) | _: LinkageError) => throw QueryExecutionErrors.nativeDataSourceError(action, name, message(e), e) } } + /** Checks that the handle returned by the given function of the library is not 0. */ + def checkHandle(name: String, action: String, function: String)(handle: Long): Long = { + if (handle == 0) { + throw error(name, action, s"$function returned 0, which is not a valid handle.") + } + handle + } + // How Arrow reports the error message of an Arrow C stream. private val StreamError = """(?s)\[errno \d+\] CDataJniException\{errno=\d+, message=(.*)\}""".r @@ -743,7 +830,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..7710479407073 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.datasource.NativeBridge import org.apache.spark.sql.errors.QueryExecutionErrors @@ -95,20 +96,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. */ @@ -116,24 +130,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 != NativeBridge.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 } } @@ -244,22 +265,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 } } @@ -292,26 +310,37 @@ 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, with the checksum of the build that was loaded: 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 val installedLibraries = new ConcurrentHashMap[String, (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 { 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] = { @@ -368,7 +397,7 @@ object NativeLibraries extends Logging { val abiVersion = try { library.abiVersion() } catch { - case e: UnsatisfiedLinkError => + case e: NativeLibrary.MissingFunctionException => throw invalidLibrary("It does not implement NativeBridge.abiVersion.", e) } if (abiVersion != NativeBridge.ABI_VERSION) { 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..29567623688a5 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 @@ -208,9 +208,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 5a1c5f5892a70..66e1be754718b 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,10 @@ import org.apache.spark.sql.types._ * `NativeBridge.pushPredicates`. */ object NativePredicates { - private val mapper = new ObjectMapper() + // Writes ASCII only, escaping the other characters. The libraries often get the JSON with + // the JNI function GetStringUTFChars, which returns modified UTF-8: it differs from UTF-8 for + // the characters outside the Basic Multilingual Plane, but any JSON parser decodes the escapes. + 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 4aba24753fca4..923011017fba2 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,28 @@ // 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 -// 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 "null", "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 returns 0 without an error: "createDataSource", "createReader", +// "createStreamReader", "createWriter", "createStreamWriter" or "createDataWriter". // // The suite compiles it with these macros to get several distinct libraries: // @@ -146,9 +157,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) { @@ -182,6 +196,7 @@ struct SchemaData { std::string name; std::vector> children; std::vector child_pointers; + std::unique_ptr dictionary; }; void releaseSchema(ArrowSchema* schema) { @@ -189,13 +204,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; @@ -229,8 +247,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(); } } } @@ -369,6 +395,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; @@ -391,7 +430,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); } } @@ -715,6 +755,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 returns 0 instead of a 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. @@ -864,6 +909,7 @@ SPARK_JNI(jlong, createDataSource)(JNIEnv* env, jclass, jstring name_value, jobj source->name != name("native_counter")) { throw std::runtime_error("unknown data source " + source->name); } + if (returnsZero(source->options, "createDataSource")) return jlong{0}; return toHandle(owner.release()); }); } @@ -889,6 +935,7 @@ SPARK_JNI(jlong, createReader)(JNIEnv* env, jclass, jlong handle, jlong schema_a } auto columns = importColumns(reinterpret_cast(schema_address)); for (const auto& column : columns) findColumn(column); + if (returnsZero(source->options, "createReader")) return jlong{0}; auto* reader = new Reader(); reader->state["columns"] = join(columns); reader->state["lo"] = "0"; @@ -898,6 +945,7 @@ 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", ""); + reader->state["claim_predicates"] = get(source->options, "claim_predicates", ""); return toHandle(reader); }); } @@ -910,13 +958,17 @@ SPARK_JNI(jlong, createStreamReader)(JNIEnv* env, jclass, jlong handle, jlong sc if (source->name != name("native_counter")) { throw std::runtime_error(source->name + " does not support streaming reads"); } + auto columns = join(importColumns(reinterpret_cast(schema_address))); + if (returnsZero(source->options, "createStreamReader")) return jlong{0}; auto* reader = new Reader(); - reader->state["columns"] = join(importColumns(reinterpret_cast(schema_address))); + 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", ""); return toHandle(reader); }); } @@ -937,6 +989,14 @@ SPARK_JNI(jbooleanArray, pushPredicates)(JNIEnv* env, jclass, jlong handle, std::vector accepted; 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") { + std::vector claimed(json.size(), JNI_TRUE); + jbooleanArray result = env->NewBooleanArray(static_cast(claimed.size())); + env->SetBooleanArrayRegion(result, 0, static_cast(claimed.size()), claimed.data()); + return result; + } for (const auto& text : json) { // Accepts "id " for the comparison operators. Json predicate = JsonParser(text).parse(); @@ -1035,6 +1095,15 @@ SPARK_JNI(jstring, latestOffset)(JNIEnv* env, jclass, jlong handle) { 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 == "null") return static_cast(nullptr); + // Pretty-printed JSON, as serde_json::to_string_pretty writes it. + if (bad_offset == "multiline") { + return env->NewStringUTF(("{\n \"offset\": " + std::to_string(latest) + "\n}").c_str()); + } + if (bad_offset == "invalid") { + return env->NewStringUTF(("offset " + std::to_string(latest)).c_str()); + } return env->NewStringUTF(("{\"offset\":" + std::to_string(latest) + "}").c_str()); }); } @@ -1063,8 +1132,12 @@ SPARK_JNI(void, commitOffset)(JNIEnv* env, jclass, jlong handle, jstring end) { }); } -SPARK_JNI(void, closeStreamReader)(JNIEnv*, jclass, jlong handle) { - delete fromHandle(handle); +SPARK_JNI(void, closeStreamReader)(JNIEnv* env, jclass, jlong handle) { + std::unique_ptr reader(fromHandle(handle)); + guardedVoid(env, [&] { + std::string path = get(reader->state, "closed_path", ""); + if (!path.empty()) writeFile(path, ""); + }); } #endif // READ_ONLY @@ -1093,7 +1166,9 @@ SPARK_JNI(void, read)(JNIEnv* env, jclass, jbyteArray reader_state, jbyteArray p namespace { -jlong createAnyWriter(JNIEnv* env, jlong handle, jlong schema_address, jboolean overwrite) { +// Creates the writer of a batch write, or of a streaming write with the ID of its query. +jlong createAnyWriter(JNIEnv* env, const char* function, jlong handle, jlong schema_address, + jboolean overwrite, jstring query_id) { return guarded(env, 0, [&] { auto* source = fromHandle(handle); if (source->name != name("native_sink")) { @@ -1102,6 +1177,7 @@ jlong createAnyWriter(JNIEnv* env, jlong handle, jlong schema_address, jboolean // 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"); + if (returnsZero(source->options, function)) return jlong{0}; auto* writer = new Writer(); // The location of a table in a catalog is passed as a URI. std::string path = get(source->options, "path", ""); @@ -1113,13 +1189,16 @@ jlong createAnyWriter(JNIEnv* env, jlong handle, jlong schema_address, jboolean writer->state["path"] = path; writer->state["overwrite"] = overwrite ? "true" : "false"; writer->state["fail"] = get(source->options, "fail", ""); + writer->state["zero"] = get(source->options, "zero", ""); 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. + // 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", toString(env, query_id)); return toHandle(writer); }); } @@ -1128,12 +1207,12 @@ jlong createAnyWriter(JNIEnv* env, jlong handle, jlong schema_address, jboolean SPARK_JNI(jlong, createWriter)(JNIEnv* env, jclass, jlong handle, jlong schema_address, jboolean overwrite) { - return createAnyWriter(env, handle, schema_address, overwrite); + return createAnyWriter(env, "createWriter", handle, schema_address, overwrite, nullptr); } SPARK_JNI(jlong, createStreamWriter)(JNIEnv* env, jclass, jlong handle, jlong schema_address, - jboolean overwrite) { - return createAnyWriter(env, handle, schema_address, overwrite); + jboolean overwrite, jstring query_id) { + return createAnyWriter(env, "createStreamWriter", handle, schema_address, overwrite, query_id); } SPARK_JNI(jbyteArray, serializeWriter)(JNIEnv* env, jclass, jlong handle) { @@ -1164,16 +1243,27 @@ SPARK_JNI(void, commit)(JNIEnv* env, jclass, jlong handle, jlong 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. 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()); + int64_t null_messages = 0; + jsize n = env->GetArrayLength(messages); + for (jsize i = 0; i < n; i++) { + auto message = static_cast(env->GetObjectArrayElement(messages, i)); + if (message == nullptr) { + null_messages++; + } else { + unlink(toBytes(env, message).c_str()); + env->DeleteLocalRef(message); + } } - writeFile(writer->state["path"] + "/_ABORTED", ""); + writeFile(writer->state["path"] + + (epoch_id < 0 ? "/_ABORTED" : "/_ABORTED_" + std::to_string(epoch_id)), + std::to_string(null_messages)); }); } @@ -1190,6 +1280,7 @@ 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)); + if (returnsZero(state, "createDataWriter")) return jlong{0}; 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 d261654f7cede..d27d67931cabf 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 694c4d777db68..136338c5f1801 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`, @@ -205,6 +214,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: a library that gets it as modified UTF-8, which differs from UTF-8 for + // the characters outside the Basic Multilingual Plane, such as emoji, still decodes them. + 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)) @@ -236,6 +251,112 @@ 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 `NativeBridge.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`. + */ + 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 NativeBridge." + + "pushPredicates, 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 other link errors") { + // A function that the library does not export is an unsupported operation. + val missing = new NativeLibrary.MissingFunctionException(new UnsatisfiedLinkError("'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 'f'.")) + // Other link errors are reported as they are, such as when Arrow 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 += _) @@ -315,6 +436,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") { @@ -337,6 +472,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") { @@ -346,10 +486,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"), @@ -372,6 +534,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") { @@ -501,13 +670,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") @@ -529,6 +725,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") @@ -541,6 +739,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("null", "latestOffset returned null instead of an offset.") + check("multiline", "latestOffset returned an offset with a line break, but an offset must " + + "be JSON on a single line.") + 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") { + 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 = { @@ -571,6 +819,75 @@ class NativeDataSourceSuite extends QueryTest with SharedSparkSession { } } + 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 handles of 0") { + 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 0, which is not a valid handle.")) + } + def runStream(start: => StreamingQuery): Unit = { + val query = start + try query.processAllAvailable() finally query.stop() + } + check("createDataSource", "create", "native_range") { session => + range(session, "zero" -> "createDataSource") + } + check("createReader", "plan a scan of", "native_range") { session => + range(session, "zero" -> "createReader").collect() + } + check("createStreamReader", "plan a scan of", "native_counter") { session => + runStream { + session.readStream.format("native_counter").option("zero", "createStreamReader").load() + .writeStream.format("noop").start() + } + } + Seq("createWriter" -> "plan a write to", "createDataWriter" -> "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("createStreamWriter", "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", "createStreamWriter") + .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. @@ -672,6 +989,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)) @@ -867,6 +1220,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") @@ -1014,6 +1397,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 e3ff80cc6bba0bbdd63cf260bc3878f772401fc8 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Wed, 14 Oct 2026 08:07:01 +0900 Subject: [PATCH 15/15] Test that a native scan reads the columns of the pushed predicates Co-authored-by: Isaac --- .../datasources/v2/ffi/NativeDataSourceSuite.scala | 7 +++++++ 1 file changed, 7 insertions(+) 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 136338c5f1801..93b9baf3340f6 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 @@ -457,6 +457,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) }