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..3dfcff8ed728f 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -4792,6 +4792,57 @@ ], "sqlState" : "42000" }, + "INVALID_NATIVE_DATA_SOURCE_LIBRARY" : { + "message" : [ + "Invalid native data source library :" + ], + "subClass" : { + "CANNOT_LOAD" : { + "message" : [ + "" + ] + }, + "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 ." + ] + } + }, + "sqlState" : "22023" + }, + "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 +6313,30 @@ ], "sqlState" : "0A000" }, + "NATIVE_DATA_SOURCE_ERROR" : { + "message" : [ + "Failed to native data source : " + ], + "sqlState" : "38000" + }, + "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND" : { + "message" : [ + "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" + }, + "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..3112614546113 100644 --- a/docs/_data/menu-sql.yaml +++ b/docs/_data/menu-sql.yaml @@ -75,6 +75,19 @@ url: sql-v2-data-sources.html#streaming - text: Further Reading url: sql-v2-data-sources.html#further-reading +- text: Native Data Sources + url: sql-native-data-sources.html + subitems: + - text: Overview + 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-native-data-sources.md b/docs/sql-native-data-sources.md new file mode 100644 index 0000000000000..4572905c2de64 --- /dev/null +++ b/docs/sql-native-data-sources.md @@ -0,0 +1,654 @@ +--- +layout: global +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 + 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 + +A native data source is a data source implemented in native code, such as Rust or C++, without +any JVM code. It is a shared library that implements the +[`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) binary interface. +Spark finds the library automatically when it is installed like other native libraries, or when a +session adds it in a package, and loads it when a query uses the data source. Spark adapts it to +[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. + +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. + +## How It Works + +A native data source library implements the `native` methods of `NativeBridge` as JNI functions, +named after the class, such as `Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource`. +Spark loads each library into its own copy of the class, so any number of native data sources can +be loaded side by side. + +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 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`. 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 functions, the `serialize` +functions and `commitOffset`, 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`](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. 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 +`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 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. + +### 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 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; +- 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 + +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, 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. + +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 +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/libspark_datasource_rust_range.so", + "linux-aarch64": "linux-aarch64/libspark_datasource_rust_range.so", + "osx-aarch64": "osx-aarch64/libspark_datasource_rust_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. + +A session finds the packages in two places, and the executors get them automatically: + +- **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, 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. 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 + +This data source reads the numbers `[0, end)` as a column `id`, in two partitions. It uses the +[`jni`](https://crates.io/crates/jni) crate, and the `arrow-array` and `arrow-schema` crates of +[arrow-rs](https://github.com/apache/arrow-rs), which export schemas and streams with +`FFI_ArrowSchema` and `FFI_ArrowArrayStream`. `Cargo.toml`, where the name of the library is the +name that Spark finds when it is installed: + +```toml +[package] +name = "rust_range" +version = "0.1.0" +edition = "2021" + +[lib] +name = "spark_datasource_rust_range" +crate-type = ["cdylib"] + +[dependencies] +arrow-array = { version = "58", features = ["ffi"] } +arrow-schema = { version = "58", features = ["ffi"] } +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. + +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(()) + }) +} +``` +{% endraw %} + +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/libspark_datasource_rust_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()); +} + +// 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); + 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))); + }); +} +``` +{% endraw %} + +Build the library, here on macOS with Arrow C++ 24, which requires C++20, and install it, or +package it as above: + +```bash +c++ -std=c++20 -O2 -shared -fPIC \ + -I"$JAVA_HOME/include" -I"$JAVA_HOME/include/darwin" \ + -I"$ARROW_HOME/include" -L"$ARROW_HOME/lib" -larrow \ + -o libspark_datasource_cpp_range.dylib cpp_range.cc +cp libspark_datasource_cpp_range.dylib /usr/local/lib/ +``` + +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. They take precedence over the installed libraries. | + +## Security + +A native data source runs native code in the driver and executor processes, with their +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/docs/sql-v2-data-sources.md b/docs/sql-v2-data-sources.md index 4b4d0f35cb4e4..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. | +| **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,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. +- [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/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..916d4a1640ea6 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -41,6 +41,7 @@ no_type_check, overload, ) +from urllib.parse import urlparse from pyspark.conf import SparkConf from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError @@ -2267,10 +2268,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 +2297,18 @@ def addArtifacts( "condition": "True together", }, ) + # Native data source packages are artifacts of the session in the JVM, whatever the + # 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 == "file": + 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) 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..73f18abc9676e 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 @@ -105,6 +106,47 @@ 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 packages do not affect the other tests. + spark = self.spark.newSession() + with tempfile.TemporaryDirectory(prefix="test_add_native_data_source_package") as d: + + 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")) + # 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() + self.assertEqual( + sorted(packages.apply(i).getName() for i in range(packages.size())), + [ + "my_file_source.sparkpkg", + "my_other_source.sparkpkg", + "my_relative_source.sparkpkg", + "my_source.sparkpkg", + "my_source:1.sparkpkg", + "my_uri_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/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/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..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 @@ -802,6 +802,50 @@ 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 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, checksum: String, path: String): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND", + messageParameters = Map("fileName" -> fileName, "checksum" -> checksum, "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 d3211aba4986f..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 @@ -7207,6 +7207,19 @@ 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. These " + + "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 + .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..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 @@ -464,6 +464,22 @@ 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++. 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") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(true) + val KAFKA_DISALLOWED_OPTIONS = buildStaticConf("spark.sql.kafka.disallowedOptions") .internal() 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.*"), 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..efef2b08240a7 --- /dev/null +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala @@ -0,0 +1,105 @@ +/* + * 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/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/NativeBridge.java b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java new file mode 100644 index 0000000000000..1500e1c6c4344 --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java @@ -0,0 +1,494 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.datasource; + +import org.apache.spark.annotation.Evolving; + +/** + * The binary interface between Spark and data sources implemented in native code, such as Rust + * or C++. + *

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

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

Lifecycle

+ * Spark creates a data source with {@link #createDataSource} on the driver for each schema + * inference, scan and write, and closes it right after. + *
    + *
  • Batch reads. On the driver, Spark creates a reader with {@link #createReader}, + * pushes operations down by calling, each at most once and in this order, + * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions + * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, + * and closes it. Then it calls {@link #read} on the executors, in one task for each partition. + * 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 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 + * {@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. 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 + * query is restarted.
    • + *
+ * + *

Conventions

+ *
    + *
  • Handles. The {@code create} methods return an opaque, non-zero handle, typically a + * 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}, 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 + * allocated by Spark. For an output struct, the library initializes it and Spark takes + * ownership. For an input struct, the library may move it to take ownership; otherwise, Spark + * releases it after the call returns, so the library must not use it afterwards. The arrays + * that the library exports must have an offset of 0, at every level: Arrow Java ignores the + * offsets of the arrays it imports, so Spark rejects an array with another offset, such as a + * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with + * fewer values than its parent refers to, such as a column shorter than its batch. The buffers + * of an array must hold all of its values, as the C data interface requires: it does not carry + * the sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that + * is too short.
    • + *
  • Schemas. Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: + * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, + * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is + * {@code timestamp[us]} without one. Data read from a library may also use {@code large_utf8} + * 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, 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 + * {@code Cleaner} thread when Spark releases a handle once it is no longer referenced, for + * example the last stream writer of a streaming query.
    • + *
+ * + *

Compatibility

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

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

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

+ * Spark evaluates all the predicates again after reading, so a reader may use them only to skip + * data, for example the files whose statistics show that no row matches. But it must never drop + * or skip a row for which the predicates are true under the semantics of Spark SQL, which Spark + * cannot detect. In particular: + *

    + *
  • For floating-point columns, NaN is equal to NaN and greater than any other value, + * including positive infinity, and {@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. 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. 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. + */ + public static native boolean pushLimit(long reader, int limit); + + /** + * Prunes the columns to read. + * + * @param columnNames the names of the top-level columns that the query needs, 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. + */ + public static native boolean pruneColumns(long reader, String[] columnNames); + + /** + * Plans the partitions of the scan. Spark reads each partition in a separate task. + * + * @return the partitions, each serialized in a format defined by the library + */ + public static native byte[][] partitions(long reader); + + /** + * Serializes the state of the reader that {@link #read} needs on the executors, such as the + * pushed predicates. Spark calls it after {@link #partitions}. + */ + public static native byte[] serializeReader(long reader); + + /** Releases a batch reader handle. */ + public static native void closeReader(long reader); + + // --------------------------------------------------------------------------------------------- + // Stream reader (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Returns the offset to start from when the streaming query has no checkpoint, as compact JSON + * on a single line. + */ + public static native String initialOffset(long streamReader); + + /** Returns the most recent offset available, as compact JSON on a single line. */ + public static native String latestOffset(long streamReader); + + /** + * Plans the partitions that read the data after {@code start} up to and including {@code end}. + * + * @return the partitions, each serialized in a format defined by the library + */ + public static native byte[][] streamPartitions(long streamReader, String start, String end); + + /** + * Serializes the state of the stream reader that {@link #read} needs on the executors. Spark + * calls it after each call to {@link #streamPartitions}. + */ + public static native byte[] serializeStreamReader(long streamReader); + + /** + * Tells the library that Spark has processed all the data up to and including {@code end}, + * and will not ask for it again. + */ + public static native void commitOffset(long streamReader, String end); + + /** Releases a stream reader handle, when the streaming query stops. */ + public static native void closeStreamReader(long streamReader); + + // --------------------------------------------------------------------------------------------- + // Partition reader (executor) + // --------------------------------------------------------------------------------------------- + + /** + * Reads a partition of a batch or streaming scan. + * + * @param readerState the state returned by {@link #serializeReader} or + * {@link #serializeStreamReader}, or an empty array if there is none + * @param partition a partition returned by {@link #partitions} or {@link #streamPartitions} + * @param streamAddress the address of the {@code ArrowArrayStream} to initialize. Each array + * of the stream is a struct array whose fields are the columns to read: + * the columns of the schema to read, or the pruned columns. Spark calls + * the stream from the thread of the task that reads the partition. + * {@code get_next} may block, for example on network I/O, but Spark + * cannot interrupt native code: when the task is killed, Spark only + * stops reading once {@code get_next} returns, and then releases the + * stream. So a {@code get_next} that blocks forever keeps the task + * running, unless the task reaper ({@code spark.task.reaper.enabled} + * and {@code spark.task.reaper.killTimeout}) kills the executor. Use + * timeouts for blocking operations. + */ + public static native void read(byte[] readerState, byte[] partition, long streamAddress); + + // --------------------------------------------------------------------------------------------- + // Writer (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Serializes the state of the writer that {@link #createDataWriter} needs on the executors. + * Spark calls it once, after creating the writer. + */ + public static native byte[] serializeWriter(long writer); + + /** + * Commits a write, when all the tasks succeeded. + * + * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write + * @param messages the messages returned by {@link #commitDataWriter} of each task + */ + public static native void commit(long writer, long epochId, byte[][] messages); + + /** + * Aborts a write, when it failed. + * + * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write + * @param messages the messages returned by {@link #commitDataWriter} of each task. An element + * is null if its task did not commit. + */ + public static native void abort(long writer, long epochId, byte[][] messages); + + /** Releases a writer handle. */ + public static native void closeWriter(long writer); + + // --------------------------------------------------------------------------------------------- + // Data writer (executor) + // --------------------------------------------------------------------------------------------- + + /** + * Creates the writer for one task. + * + * @param writerState the state returned by {@link #serializeWriter}, or an empty array if there + * is none + * @param partitionId the partition that the task writes + * @param taskId the ID of the task attempt + * @param epochId the ID of the micro-batch for a streaming write, or -1 for a batch write + * @return the data writer handle + */ + public static native long createDataWriter( + byte[] writerState, int partitionId, long taskId, long epochId); + + /** + * Writes a batch of data. + * + * @param arrayAddress the address of the {@code ArrowArray}: a struct array whose fields are + * the columns + * @param schemaAddress the address of the {@code ArrowSchema} of the array + */ + public static native void write(long dataWriter, long arrayAddress, long schemaAddress); + + /** + * Commits the data written by the task and releases the data writer handle. If it fails, the + * handle is not released: Spark then calls {@link #abortDataWriter} with it, to abort the data + * written by the task and release the handle. + * + * @return the message to pass to {@link #commit} on the driver + */ + public static native byte[] commitDataWriter(long dataWriter); + + /** Aborts the data written by the task and releases the data writer handle. */ + public static native void abortDataWriter(long dataWriter); +} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java new file mode 100644 index 0000000000000..c073f41e9a9a9 --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -0,0 +1,324 @@ +/* + * 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 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; + 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, 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); + 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 { + 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))); + } + + /** + * 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); + // 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 { + 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, String queryId) + throws Throwable { + return (long) createStreamWriter.invokeExact(dataSource, schemaAddress, overwrite, queryId); + } + + 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/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/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/artifact/ArtifactManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala index 2a03a58502c60..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 @@ -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} @@ -153,6 +154,28 @@ 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)) + } + + /** 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) } @@ -618,6 +641,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/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/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 1301985e30158..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 @@ -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._ @@ -254,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: PythonDataSourceV2 => 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/DataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala index 52b9e803a64ed..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 @@ -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 @@ -666,6 +667,21 @@ 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[_] = { + // 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" => @@ -682,6 +698,9 @@ object DataSource extends Logging { val serviceLoader = ServiceLoader.load(classOf[DataSourceRegister], loader) lazy val isUserDefinedDataSource = SparkSession.getActiveSession.exists( _.sessionState.dataSourceManager.dataSourceExists(provider)) + // 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 { @@ -703,6 +722,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 +739,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 +804,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/ColumnarDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala new file mode 100644 index 0000000000000..d1bb7705c33f2 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala @@ -0,0 +1,134 @@ +/* + * 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. 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, queryId: String): 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 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. */ + def pushLimit(limit: Int): Boolean = false + + /** + * Returns whether `read` returns exactly the columns of `requiredSchema`, the top-level columns + * 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 + + 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 new file mode 100644 index 0000000000000..808025d01b7fd --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala @@ -0,0 +1,242 @@ +/* + * 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.Locale + +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.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 [[ColumnarDataSource]]. It creates a new data source for each + * 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 { + + 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( + shortName, tableSchema, () => createDataSource(withTableOptions(options))) + } + + override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = { + 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) + } +} + +class ColumnarScanBuilder( + shortName: String, + schema: StructType, + createDataSource: () => ColumnarDataSource) + 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: ColumnarReader = _ + + 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 = { + 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) + // 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 + + 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. + 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( + shortName, schema, readSchema, () => batchReader, () => dataSource, pushed, pushedLimit) + } +} + +class ColumnarScan( + shortName: String, + fullSchema: StructType, + prunedSchema: StructType, + batchReader: () => ColumnarReader, + dataSource: () => ColumnarDataSource, + 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 + + override def description(): String = shortName + + 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: ColumnarReader) 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: ColumnarStreamReader) 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() +} + +/** + * A [[ColumnarDataSource]] only produces 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: ColumnarReader) + extends ColumnarOnlyPartitionReaderFactory { + override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = { + reader.read(partition) + } +} + +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 new file mode 100644 index 0000000000000..8759c2894ead9 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala @@ -0,0 +1,243 @@ +/* + * 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.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: () => ColumnarDataSource) + 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: () => ColumnarDataSource) + 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.queryId()), + info.schema()) + } +} + +/** + * 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 { + + 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 = { + if (writer != null) { + writer.abort(messages) + } + } +} + +/** + * A streaming write. As for a [[ColumnarBatchWrite]], the writer of the data source is created + * 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 { + + private var writer: ColumnarStreamWriter = _ + + override def createStreamingWriterFactory( + info: PhysicalWriteInfo): StreamingDataWriterFactory = { + if (writer == null) { + writer = createWriter() + } + 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 = { + if (writer != null) { + writer.abort(epochId, messages) + } + } +} + +class ColumnarDataWriterFactory(writer: ColumnarWriter, options: ArrowBatchingOptions) + extends DataWriterFactory { + override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = { + new ArrowBatchingDataWriter(writer.createWriter(partitionId, taskId), options) + } +} + +class ColumnarStreamingDataWriterFactory( + writer: ColumnarStreamWriter, + 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. 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, + 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, + conf.arrowMaxBytesPerBatch) + } +} + +/** + * Converts the rows of a write task to Arrow-backed [[ColumnarBatch]]es for the [[DataWriter]] + * 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) + 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 + // 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() + } + } + + 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..d748e0cb4287b --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -0,0 +1,881 @@ +/* + * 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, WeakReference} +import java.util.concurrent.ConcurrentHashMap +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, ValueVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.MapVector +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} + +import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.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.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarDataSource, ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataType, StructField, StructType} +import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.unsafe.Platform +import org.apache.spark.util.Utils + +/** + * A [[ColumnarDataSource]] implemented by a native library, through + * `org.apache.spark.sql.datasource.NativeBridge`. + * + * @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 + */ +class NativeDataSource( + location: NativeLibraryLocation, + name: String, + options: CaseInsensitiveStringMap) + extends ColumnarDataSource { + import NativeCalls._ + + override def schema(): StructType = { + val library = NativeLibraries.get(location) + withDataSource(library, "infer the schema of") { dataSource => + call(name, "infer the schema of") { + optional(NativeArrow.importSchema(address => library.schema(dataSource, address))) + }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) + } + } + + override def reader(schema: StructType): ColumnarReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + 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)) + } + + override def streamReader(schema: StructType): ColumnarStreamReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + 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)) + } + + 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", "createWriter", + "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource => + NativeArrow.exportSchema(schema) { address => + library.createWriter(dataSource, address, overwrite) + } + }, + library.closeWriter) + new NativeDataSourceWriter(executorLocation, name, handle, serializeWriter(library, handle)) + } + + 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", "createStreamWriter", + "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { dataSource => + NativeArrow.exportSchema(schema) { address => + library.createStreamWriter(dataSource, address, overwrite, queryId) + } + }, + library.closeWriter) + 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. */ + 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 = 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 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 => + val handle = call(name, action)(optional(f(dataSource))).getOrElse { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) + } + checkHandle(name, action, function)(handle) + } + } + + 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( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarReader { + 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] = _ + // 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) + + 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, schema).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) { + 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() + } + } + plannedPartitions + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, readSchema) + } +} + +/** Reads a native data source for a micro-batch streaming scan. */ +class NativeDataSourceStreamReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarStreamReader { + 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(location) + + override def initialOffset(): String = { + checkOffset("initialOffset", call(name, "plan a scan of")(library.initialOffset(handle.get))) + } + + override def latestOffset(): String = { + 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] = { + 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( + location, 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() +} + +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, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, 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)) + } + // Not closed if the commit fails: Spark then aborts the write, which closes it. + 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. Its handle is closed when the streaming + * query creates its next stream writer, see [[NativeStreamWriters]]. + */ +class NativeDataSourceStreamWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarStreamWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, 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: 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, + 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) + // Read here rather than with `Data.importArrayStream`, to check each array before importing it. + private val stream = ArrowArrayStream.allocateNew(allocator) + private val dictionaries = new CDataDictionaryProvider() + private var root: VectorSchemaRoot = _ + private var batch: ColumnarBatch = _ + + try { + openStream() + root = VectorSchemaRoot.create(importSchema(), allocator) + checkSchema() + batch = new ColumnarBatch(root.getFieldVectors.asScala.map(columnVector).toArray) + } catch { + case e: Throwable => + close() + throw e + } + + private def openStream(): Unit = { + val library = NativeLibraries.get(location) + call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) + .getOrElse { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) + } + } + + private def importSchema(): Schema = { + val arrowSchema = ArrowSchema.allocateNew(allocator) + try { + 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 { + 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) + } + root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + 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 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 => + 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 = 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. + 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 ${expected.toDDL}.") + } + } + + private def columnVector(vector: FieldVector): ColumnVector = { + try { + new ReusedArrowColumnVector(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 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 { + checkArrays(snapshot) + } catch { + case e: Throwable => + array.release() + throw e + } + call(name, "read")(Data.importIntoVectorSchemaRoot(allocator, array, root, dictionaries)) + batch.setNumRows(root.getRowCount) + true + } + } finally { + array.close() + } + } + + /** + * 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 checkArrays(batch: ArrowArray.Snapshot): Unit = { + def children(array: ArrowArray.Snapshot): Seq[ArrowArray.Snapshot] = { + (0 until array.n_children.toInt).map { i => + ArrowArray.wrap(Platform.getLong(null, array.children + 8L * i)).snapshot() + } + } + // The last offset of a list or a map, in its second buffer. Arrow Java rejects the array if + // the buffer is missing. + def lastOffset(array: 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) + // The other types that Spark reads have no children. List views are rejected before. + case _ => 0 + } + def checkColumn(array: ArrowArray.Snapshot, field: Field, column: String): Unit = { + if (array.offset != 0) { + throw error(name, "read", s"The column $column has an offset of ${array.offset}, but " + + "Spark only reads arrays with an offset of 0. Copy sliced arrays before exporting them.") + } + 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 (batch.offset != 0) { + throw error(name, "read", s"The batch has an offset of ${batch.offset}, but Spark only " + + "reads arrays with an offset of 0. Copy sliced arrays before exporting them.") + } + children(batch).zip(root.getSchema.getFields.asScala).foreach { case (column, field) => + if (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) + } + } + + override def get(): ColumnarBatch = batch + + override def close(): Unit = { + Utils.tryWithSafeFinally { + 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())) + } +} + +/** + * 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, + name: String, + writerState: Array[Byte], + partitionId: Int, + taskId: Long, + epochId: Long) + extends DataWriter[ColumnarBatch] { + import NativeCalls._ + + private val library = NativeLibraries.get(location) + // Zero once the data writer was committed or aborted, which releases it. It stays set if the + // commit fails, so that Spark aborts the data writer. + 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 = { + 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 = { + 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(Some(message))) + } + + 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) + } + } + } + } +} + +/** + * 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, 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: NativeLibrary.MissingFunctionException => + throw QueryExecutionErrors.nativeDataSourceError( + action, name, s"The library does not implement ${e.getMessage}.", e) + case e @ (NonFatal(_) | _: LinkageError) => + throw QueryExecutionErrors.nativeDataSourceError(action, name, message(e), e) + } + } + + /** 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 + + private[ffi] def message(e: Throwable): String = { + Option(e.getMessage).getOrElse(e.toString) match { + case StreamError(message) => message + case message => message + } + } + + /** Calls a function that the library may not implement. Returns None if it does not. */ + def optional[T](f: => T): Option[T] = { + try { + Some(f) + } catch { + case _: NativeLibrary.MissingFunctionException => None + } + } + + /** Returns a NATIVE_DATA_SOURCE_ERROR for an error detected by Spark. */ + def error(name: String, action: String, msg: String): Throwable = { + QueryExecutionErrors.nativeDataSourceError(action, name, msg, null) + } + + 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..7710479407073 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -0,0 +1,450 @@ +/* + * 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.nio.file.attribute.BasicFileAttributes +import java.security.MessageDigest +import java.util.{HexFormat, Locale} +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} +import com.google.common.cache.{Cache, CacheBuilder} + +import org.apache.spark.internal.Logging +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 +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]) + +/** + * 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 + * 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. + * @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, + distributedFileName: Option[String] = None) extends NativeLibraryLocation { + def fileName: String = new File(path).getName +} + +/** + * 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]) + +/** + * 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[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[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. */ + 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, 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 version = try { + FileVersion(file.getCanonicalFile) + } catch { + case e: IOException => + return Left(InvalidNativeDataSourcePackage( + invalidManifestError(file, s"The package cannot be read: $e", e), Nil)) + } + Option(packages.getIfPresent(version)).getOrElse { + val canonical = new File(version.path) + val result = parse(canonical) + 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(version, 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) = { + 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)) + } + 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) => + 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) + } + } + } + } + manifest.map(NativeDataSourcePackage(file.getPath, sha256(file), _)) + } catch { + case e: IOException => invalid(s"The package is not a valid zip file: $e", cause = e) + } + } + + 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[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) + var read = in.read(buffer) + while (read != -1) { + digest.update(buffer, 0, read) + read = in.read(buffer) + } + } + HexFormat.of().formatHex(digest.digest()) + } +} + +/** + * 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. 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, sha256: String) extends NativeLibraryLocation + +object InstalledNativeLibrary { + // Bounded, like the cache of the packages, although few libraries are installed. + 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 version = FileVersion(file.getCanonicalFile) + Option(checksums.getIfPresent(version)).getOrElse { + val checksum = NativeDataSourcePackage.sha256(new File(version.path)) + checksums.put(version, checksum) + checksum + } + } +} + +/** 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 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 { + // The libraries of the packages, by the checksum of the package, and the installed libraries, + // 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, (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)).get + case installed: InstalledNativeLibrary => + // The file has the checksum of the build of the driver. + val file = locate(installed).getCanonicalFile + 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[T](libraries: ConcurrentHashMap[String, T], key: String)( + load: => T): T = { + val library = libraries.get(key) + if (library != null) library else libraries.computeIfAbsent(key, _ => load) + } + + private def loadPackaged(pkg: NativeDataSourcePackage): Try[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) + } + 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)) + } + + 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(file.getPath) + } catch { + case e: UnsatisfiedLinkError => throw invalidLibrary(e.getMessage, e) + } + Try { + val abiVersion = try { + library.abiVersion() + } catch { + case e: NativeLibrary.MissingFunctionException => + 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 + } + } + + /** 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.tryRead(file).exists(_.sha256 == pkg.sha256) + } + .getOrElse { + throw QueryExecutionErrors.nativeDataSourcePackageNotFoundError(pkg.fileName, pkg.sha256) + } + } + + /** + * 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) + (file +: NativeDataSourceRegistry.findInstalledLibraries(file.getName)) + .find(f => f.isFile && InstalledNativeLibrary.sha256(f) == installed.sha256) + .getOrElse { + throw QueryExecutionErrors.nativeDataSourceLibraryNotFoundError( + file.getName, installed.sha256, installed.path) + } + } + + 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..29567623688a5 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -0,0 +1,280 @@ +/* + * 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 java.util.Locale + +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 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 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, 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]]. Each + // lookup clears it first, so a provider never takes the result of an earlier lookup. + private val lastLookup = new ThreadLocal[(String, (String, NativeLibraryLocation))] + + /** + * 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, NativeLibraryLocation)] = { + if (!conf.getConf(StaticSQLConf.NATIVE_DATA_SOURCE_ENABLED)) { + return None + } + 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 + } + } + + /** + * 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 the last [[lookupForProvider]] on this thread, if it was for the given + * name, and forgets it. + */ + def takeLookup(name: String): Option[(String, NativeLibraryLocation)] = { + 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, + 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) => Some(result) + case _ => + throw QueryCompilationErrors.nativeDataSourcePackageConflictError( + name, found.map(_._2.path).sorted) + } + } + + 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) + findInstalledLibraries(libraryFileName(dataSource)).headOption + .map(file => dataSource -> InstalledNativeLibrary(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)) + .filter(_.isFile) + } + + /** + * 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 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, 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 => + 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 => 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 + } + + /** + * 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 + if (added.exists(_.getCanonicalPath == pkg.path)) { + pkg.copy(artifactUUID = artifactUUID) + } else { + val copyName = s"${pkg.fileName}.${pkg.sha256}" + // 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)) + } + } + + /** + * 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 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(_, fileName)) + (pkg.artifactUUID.map(uuid => new File(new File(root, uuid), fileName)).toSeq ++ + Seq(new File(root, fileName)) ++ sessionCopies).distinct + } + } + + 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) { + packageFilesIn(file) + } else if (file.isFile) { + Seq(file) + } else { + logWarning(log"The native data source path ${MDC(PATH, path)} does not exist.") + Nil + } + } + } + + private def packageFilesIn(dir: File): Seq[File] = { + 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] = { + 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..93876788182f2 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala @@ -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.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 = _ + // 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 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) + lookup.orElse(NativeDataSourceRegistry.lookup(shortName, SQLConf.get)).getOrElse { + throw QueryCompilationErrors.dataSourceDoesNotExist(shortName) + } + } + + override def inferSchema(options: CaseInsensitiveStringMap): StructType = { + 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, location) = resolved + new ColumnarTable( + shortName, schema, properties, options => new NativeDataSource(location, 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..66e1be754718b --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala @@ -0,0 +1,138 @@ +/* + * 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.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} +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 { + // 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 + * `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 with the functions in [[FUNCTION_NAMES]]. + */ + def toJson(predicate: Predicate, schema: StructType): Option[String] = { + encode(predicate, schema).map(mapper.writeValueAsString) + } + + 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 if FUNCTION_NAMES.contains(function.name()) => + 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) + } + } + + /** 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] = { + 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..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,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.execution.datasources.v2.columnar.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..923011017fba2 --- /dev/null +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -0,0 +1,1336 @@ +/* + * 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. 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. 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: +// +// 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 +}; + +// The columns that native_range adds with the option "extra", by its value. +const std::map kExtraColumns = { + {"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] + {"largelist", {"ll", "+L"}}, // large_list, without data + {"dictionary", {"dict", "i"}} // dictionary, without data +}; + +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, 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; +} + +// --------------------------------------------------------------------------------------------- +// Exporting schemas +// --------------------------------------------------------------------------------------------- + +struct SchemaData { + std::string format; + std::string name; + std::vector> children; + std::vector child_pointers; + std::unique_ptr dictionary; +}; + +void releaseSchema(ArrowSchema* schema) { + auto* data = static_cast(schema->private_data); + 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, {}, {}, nullptr}; + 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; +} + +// 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); + for (const auto& name : columns) { + 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); + } 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(); + } + } +} + +// 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)); +} + +// 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; + 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 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 == "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; + 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 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 { + // Including "ll" and "dict", which Spark rejects from the schema, before reading data. + throw std::runtime_error("no data for the column " + name); + } +} + +// 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. `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, + const std::string& short_array, ArrowArray* out) { + initArray(out, hi - lo, 0, {std::vector()}); + for (const auto& name : columns) { + 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. + } + 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--; + } + } +} + +// --------------------------------------------------------------------------------------------- +// 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); +} + +// 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. +}; + +// 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; + bool slice; + std::string short_array; + 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, state->slice, state->short_array, 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. + // 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; + 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); + } + if (returnsZero(source->options, "createDataSource")) return jlong{0}; + 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"); + } + bool counter = source->name == name("native_counter"); + exportSchema(allColumns(counter, get(source->options, "extra", "")), + 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); + if (returnsZero(source->options, "createReader")) return jlong{0}; + auto* reader = new Reader(); + reader->state["columns"] = join(columns); + reader->state["lo"] = "0"; + reader->state["hi"] = get(source->options, "end", "10"); + reader->state["partitions"] = get(source->options, "partitions", "2"); + reader->state["batch_size"] = get(source->options, "batch_size", "1000"); + reader->state["fail"] = get(source->options, "fail", ""); + reader->state["slice"] = get(source->options, "slice", ""); + reader->state["short"] = get(source->options, "short", ""); + reader->state["claim_predicates"] = get(source->options, "claim_predicates", ""); + 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 columns = join(importColumns(reinterpret_cast(schema_address))); + if (returnsZero(source->options, "createStreamReader")) return jlong{0}; + auto* reader = new Reader(); + reader->state["columns"] = columns; + reader->state["max_offset"] = get(source->options, "max_offset", "10"); + reader->state["step"] = get(source->options, "step", reader->state["max_offset"]); + reader->state["current"] = "0"; + reader->state["batch_size"] = get(source->options, "batch_size", "1000"); + reader->state["commit_path"] = get(source->options, "commit_path", ""); + reader->state["bad_offset"] = get(source->options, "bad_offset", ""); + reader->state["closed_path"] = get(source->options, "closed_path", ""); + 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); + // 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(); + 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, [&] { + 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; + }); +} + +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); + 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()); + }); +} + +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* 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 + +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", + get(state, "slice", "") == "true", get(state, "short", ""), + ""}; + }); +} + +#ifndef READ_ONLY + +namespace { + +// 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")) { + 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"); + 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", ""); + 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", ""); + 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, 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); + }); +} + +} // namespace + +SPARK_JNI(jlong, createWriter)(JNIEnv* env, jclass, jlong handle, jlong schema_address, + jboolean overwrite) { + return createAnyWriter(env, "createWriter", handle, schema_address, overwrite, nullptr); +} + +SPARK_JNI(jlong, createStreamWriter)(JNIEnv* env, jclass, jlong handle, jlong schema_address, + jboolean overwrite, jstring query_id) { + return createAnyWriter(env, "createStreamWriter", handle, schema_address, overwrite, query_id); +} + +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())); + writer->closed_marker = closedMarker(epoch_id); + }); +} + +// Records the aborted write in _ABORTED for a batch write, or _ABORTED_ for a streaming +// one, with the number of tasks that did not commit, whose messages are null. +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); + 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"] + + (epoch_id < 0 ? "/_ABORTED" : "/_ABORTED_" + std::to_string(epoch_id)), + std::to_string(null_messages)); + }); +} + +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) { + 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) + "-" + + 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, [&] { + 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); + delete writer; + return message; + }); +} + +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/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala new file mode 100644 index 0000000000000..bc94db65cf0ab --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala @@ -0,0 +1,402 @@ +/* + * 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.nio.charset.StandardCharsets.UTF_8 +import java.util +import java.util.concurrent.ConcurrentHashMap + +import scala.collection.mutable.ArrayBuffer + +import org.apache.spark.SparkException +import org.apache.spark.sql.{QueryTest, Row} +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.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 +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._ + + 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 + } + + 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("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" + (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("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("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] { + 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("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]]. */ +class TestRangeProvider extends TableProvider with DataSourceRegister { + override def shortName(): String = "test_columnar_range" + + 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, properties, new TestRangeDataSource(_)) + } + + override def supportsExternalMetadata(): Boolean = true +} + +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): ColumnarReader = new TestRangeReader(end, schema) + + override def streamReader(schema: StructType): ColumnarStreamReader = { + new TestRangeStreamReader(end, schema) + } + + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + TestColumnarSink.writers.add(table) + new TestSinkWriter(table, overwrite, options.getBoolean("fail", false)) + } + + override def streamWriter( + schema: StructType, + overwrite: Boolean, + queryId: String): ColumnarStreamWriter = { + TestColumnarSink.streamWriterQueryIds.compute( + table, (_, old) => Option(old).getOrElse(Nil) :+ queryId) + new TestSinkStreamWriter(table) + } +} + +case class TestRangePartition(lo: Long, hi: Long) extends InputPartition + +/** Accepts the predicates `id > `. */ +class TestRangeReader(end: Long, schema: StructType) extends ColumnarReader { + 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 ColumnarStreamReader { + 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]() + // 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]]() + + 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 ColumnarWriter { + 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 ColumnarStreamWriter { + 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) + } +} 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..d27d67931cabf --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.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 +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 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)) + 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 new file mode 100644 index 0000000000000..93b9baf3340f6 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -0,0 +1,1465 @@ +/* + * 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.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, 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} +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.{ThreadUtils, Utils} + +/** + * Tests native data sources with the library in `native-datasource/test_native_datasource.cc`, + * compiled with the C++ compiler of the machine by [[NativeDataSourceTestUtils]]. The tests that + * load a library are skipped when there is no C++ compiler or no JNI headers. + */ +class NativeDataSourceSuite extends QueryTest with SharedSparkSession { + import NativeDataSourceTestUtils._ + + // 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 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 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 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("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) + } + val id = FieldReference(Seq("id")) + 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(toJson(new And( + predicate("IS_NULL", FieldReference(Seq("a", "b"))), + 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":"IN",""" + + """"children":[{"type":"column","name":["id"]},""" + + """{"type":"literal","dataType":"bigint","value":1},""" + + """{"type":"literal","dataType":"bigint","value":2}]}]}""")) + // Only the functions listed in NativeBridge.pushPredicates are passed, so not arithmetic. + 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))) + .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"""")) + // 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)) + .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) + + // 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) + } + + /** + * 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 += _) + 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") { + 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"))) + + // 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() + } + 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) + // 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") + .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") { + 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")) + + // 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) + } + + 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") + + // 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") + + // 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") { + // 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("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] { + 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"), + 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") + // 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") + // 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") + // 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") { + 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) + // 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) + 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("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() + 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(new File(dir, "_CLOSED").isFile) + 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() + 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() + 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("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") + .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)) + // 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) + // 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") + .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("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 = { + 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("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. + 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 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) + assert(NativeDataSourceRegistry.distributedCopies(distributed).exists(_.isFile)) + // Adding the same package again does nothing. + assert(NativeDataSourceRegistry.addToArtifacts(other, configuredPackage) == distributed) + + // 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) + 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. + 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("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)) + 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 + 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") + + // 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") { + 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") { + 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( + session.read.format("Installed_Native_Range").option("end", "3").load().select("id"), + Seq(Row(0L), Row(1L), Row(2L))) + 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( + other.read.format("installed_native_range").option("end", "3").load().select("id"), + Seq(Row(0L), Row(-1L), Row(-2L))) + + // 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, 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, "0")) + }, + condition = "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND", + parameters = Map("fileName" -> installed.getName, "checksum" -> "0", "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("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") + // 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 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")) + } + + 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, when + // no other package or installed library provides the data source. + 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("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("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) + 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")) + 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) + } + } + } +}