You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "HyukjinKwon (via GitHub)" <gi...@apache.org> on 2026/10/13 06:44:02 UTC
[PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
HyukjinKwon opened a new pull request, #59424:
URL: https://github.com/apache/spark/pull/59424
### What changes were proposed in this pull request?
This is the FFM version of #59313, as a draft for comparison: the same feature, with the binary interface as a plain C ABI called through the [Foreign Function and Memory (FFM) API](https://openjdk.org/jeps/454) instead of JNI. It includes all the changes of #59313, up to ff7b1ae, and only changes how Spark calls the libraries. Native data sources then require Java 22 or later, where the FFM API is final; with Java 17 and 21, Spark compiles and runs as before, and a query that uses a native data source fails with `NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION`.
Like #59313, this PR adds native data sources: data sources implemented in native code, such as Rust or C++, without any JVM code. A native data source is a shared library that implements a binary interface. Spark finds it automatically when it is installed like other native libraries, or in a package file added to a session, and loads it when a query uses it. Spark adapts it to Data Source V2, with batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and limit pushdown. Data is exchanged only as columnar batches, through the Arrow C data interface.
**Differences from #59313**:
| | #59313 (JNI) | This PR (FFM) |
|---|---|---|
| Interface | `org.apache.spark.sql.datasource.NativeBridge`, a Java class whose `native` methods the library implements | `spark_datasource.h`, a C header (`sql/core/src/main/resources/org/apache/spark/sql/datasource/`, also in the `spark-sql` JAR) |
| Function names | `Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource`, ... | `spark_ds_create`, ... |
| What a library needs | `jni.h`, or the `jni` crate | nothing but the header, or `extern "C"` declarations |
| Errors | throw a Java exception (`ThrowNew`) | return non-zero, with the message from `spark_ds_last_error` |
| Returned byte arrays (partitions, state, offsets, commit messages) | Java `byte[]`, `byte[][]`, `String` built by the library | passed to a `SparkDsSink` callback, which copies them |
| Side-by-side libraries | a copy of `NativeBridge` per library, in its own class loader | each library looked up separately (`SymbolLookup.libraryLookup`) |
| Java versions | 17, 21, 25 | 22 and later (25 among the supported LTS versions); a clear error otherwise |
**Public API**: a single C header, `spark_datasource.h`, whose comments are the specification of the interface, with the same content as the Javadoc of `NativeBridge` in #59313: the lifecycle, the handles, the ownership of the Arrow structs and of the memory that Spark passes, errors (non-zero return values, with the message of `spark_ds_last_error`, reported as `NATIVE_DATA_SOURCE_ERROR`), the JSON format of the pushed predicates, the Spark SQL semantics that a predicate evaluated by the library must follow, the threads, and the compatibility rules. Version 1 of the interface, `SPARK_DS_ABI_VERSION`. The `org.apache.spark.sql.datasource.NativeBridge` class of #59313 is removed: there is no JVM API for native data sources.
| Operation | Required functions | Optional functions |
|-----------|--------------------|--------------------|
| All | `spark_ds_abi_version`, `spark_ds_create`, `spark_ds_close` | `spark_ds_last_error`, `spark_ds_schema` |
| Batch read | `spark_ds_create_reader`, `spark_ds_partitions`, `spark_ds_read`, `spark_ds_close_reader` | `spark_ds_push_predicates`, `spark_ds_push_limit`, `spark_ds_prune_columns`, `spark_ds_serialize_reader` |
| Streaming read | `spark_ds_create_stream_reader`, `spark_ds_initial_offset`, `spark_ds_latest_offset`, `spark_ds_stream_partitions`, `spark_ds_read`, `spark_ds_close_stream_reader` | `spark_ds_serialize_stream_reader`, `spark_ds_commit_offset` |
| Batch write | `spark_ds_create_writer`, `spark_ds_commit`, `spark_ds_abort`, `spark_ds_close_writer`, `spark_ds_create_data_writer`, `spark_ds_write`, `spark_ds_commit_data_writer`, `spark_ds_abort_data_writer` | `spark_ds_serialize_writer` |
| Streaming write | `spark_ds_create_stream_writer`, `spark_ds_commit`, `spark_ds_abort`, `spark_ds_close_writer`, `spark_ds_create_data_writer`, `spark_ds_write`, `spark_ds_commit_data_writer`, `spark_ds_abort_data_writer` | `spark_ds_serialize_writer` |
A library only exports the functions of the operations it supports. Spark reports a missing required function as `DATA_SOURCE_*_NOT_SUPPORTED`.
**How Spark calls the libraries** (`execution.datasources.v2.ffi.NativeLibrary`): Spark still compiles for Java 17, where the FFM API does not exist, so `NativeLibrary` reaches `java.lang.foreign` only through method handles, looked up once when the first library is loaded, and only on Java 22 or later. Each function of a library is bound once to a downcall handle adapted to take and return primitive values only, with pointers as `long` addresses, so each call is a plain `invokeExact` without reflection, which the JIT compiles like a direct FFM call; a call to `spark_ds_partitions`, including the upcalls to the sink and the allocation of its arguments, takes about 175 ns on JDK 26. The memory that Spark passes to a function is allocated with `Platform` and freed after the call. A single upcall stub, `SparkDsSink`, copies the byte arrays that the library returns, and never throws, since an exception thrown by an upcall would crash the JVM. Spark already runs with `--enable-native-a
ccess=ALL-UNNAMED`, which the restricted methods of the FFM API need to avoid warnings; a library that cannot be loaded, including when native access is denied, is reported as `INVALID_NATIVE_DATA_SOURCE_LIBRARY.CANNOT_LOAD` or `INVALID_NATIVE_DATA_SOURCE_PACKAGE.INVALID_LIBRARY`, as before. Everything above `NativeLibrary` is the same as in #59313: the Scala classes keep calling it with the same methods, and only catch `NativeLibrary.MissingFunctionException` and `NativeLibrary.LoadException` instead of `UnsatisfiedLinkError`.
Spark still uses `arrow-c-data` to import the Arrow C structs, and `arrow-c-data` calls their `get_next` and `release` callbacks through the JNI library that it bundles, as in #59313. With FFM, Spark could call those callbacks itself and drop that library, which is left for later.
**Finding native data sources**: when no Java or Python data source has the requested name, Spark looks for a native data source with that name, in two places:
- **Installed libraries, automatically**: like the Python data sources installed in the Python path, a library installed like other native libraries is found without any configuration, by its file name. The library of the data source `rust_range` is `spark_datasource_rust_range`: `libspark_datasource_rust_range.so` on Linux, `libspark_datasource_rust_range.dylib` on macOS and `spark_datasource_rust_range.dll` on Windows. Spark checks whether the file exists in the native library path, which consists of:
1. the directories of `java.library.path`, where the JVM itself finds native libraries. They include `LD_LIBRARY_PATH` on Linux, `DYLD_LIBRARY_PATH` on macOS and `PATH` on Windows, and so `spark.driver.extraLibraryPath` and `spark.executor.extraLibraryPath`;
2. the `lib` directory of each installation prefix in `PATH`: `<prefix>/lib` for each `<prefix>/bin`.
So libraries are found where native libraries are usually installed: `/usr/local/lib` (`make install`, CMake, cargo-c, Homebrew on Intel macs), `/opt/homebrew/lib`, `$CONDA_PREFIX/lib` and `$VIRTUAL_ENV/lib` in an activated environment, and `/usr/lib`. A library that implements several data sources is installed under each of their names, for example with symbolic links. The executors load the library from the same path as the driver, or else find it in their own native library path, and only load the same build as the driver, which they check with the SHA-256 checksum of the file.
- **Packages of a session**: a native data source package is a zip file with the extension `.sparkpkg`, with a manifest, `spark-native-datasource.json`, that lists the ABI version, the names of the data sources, and the library for each platform such as `linux-x86_64` or `osx-aarch64`. A session adds packages with `spark.addArtifact` (Scala and Python, classic and Spark Connect: `.sparkpkg` files are added as file artifacts), or with `spark.sql.dataSource.native.paths`. They take precedence over the installed libraries, and the executors get them through the artifacts of the session, except in local mode, where they read them in place: a configured package is copied there under a name with its checksum, so it can be replaced in place. Spark ignores, and logs, a package that it cannot read or whose ABI version it does not support, unless its manifest lists the requested data source and no other package or installed library provides it.
**Short names**: no JVM class is needed for each data source. After the Java data sources (`DataSourceRegister`) and the Python ones, `DataSource.lookupDataSource` looks the name up in the packages of the session and in the native library path, case-insensitively, and resolves it to a single provider class, `NativeDataSourceV2`. The caller then gives it the requested name, the way `PythonDataSourceV2` already works: a small `NamedTableProvider` trait generalizes how `PythonDataSourceV2` gets its short name, so the native provider shares the same call sites. The provider reuses the result of the lookup that found its class, instead of looking the name up again. A Python data source can still be registered with the name of a native data source, and takes precedence over it. Lookups only read manifests and check whether files exist; Spark loads a library when a query uses it. Native data sources work with `spark.read`/`write`, `readStream`/`writeStream`, and SQL tables (`CREATE TABLE
... USING rust_range OPTIONS (...)`, `SELECT` and `INSERT INTO`).
**Other changes**, besides those of #59313:
- Error condition: `NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION` (`0A000`), raised when a query uses a native data source on Java 21 or earlier. The lookups still find the data source, so the error says why it cannot be used rather than that it does not exist.
- The test library `test_native_datasource.cc` implements the C interface, and the tests compile it with `spark_datasource.h` instead of the JNI headers.
- The `CheckConnectJvmClientCompatibility` exclusion of `org.apache.spark.sql.datasource`, added by #59313 for `NativeBridge`, is not needed.
The other changes are those of #59313: the configurations `spark.sql.dataSource.native.enabled` and `spark.sql.dataSource.native.paths`, the error conditions, the `arrow-c-data` dependency, the "Native Data Sources" documentation page (with the examples below), and:
- Existing classes:
- `DataSource.lookupDataSource` looks up the native data sources after the Python ones, and `DataSourceRegistration` skips them when it registers a Python data source.
- `ArtifactManager` lists the native data source packages of a session, and forgets the cached packages read from the artifacts of a session when it cleans them up.
- The classic `DataFrameWriter` builds its commands, and `DataStreamWriter` resolves its sink, with the session of the Dataset active, as the analysis of a query does: the lookups of the Python and native data sources use the active session.
- PySpark `addArtifacts` adds `.sparkpkg` files to the session whatever the flags: in classic mode through the JVM session, and with Spark Connect as file artifacts.
### Examples
Both examples implement a read-only data source that returns the numbers `[0, end)` as a column `id`, in two partitions, with the option `end` (10 by default). They only export the functions needed for batch reads, plus `spark_ds_last_error`.
<details>
<summary><b>Rust</b> (with arrow-rs only)</summary>
`Cargo.toml`:
```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"] }
```
`src/lib.rs`:
```rust
//! A native data source for Apache Spark, written in Rust. It reads the numbers [0, end) as a
//! column `id`, in two partitions. The option `end` defaults to 10.
use std::cell::RefCell;
use std::ffi::{c_char, c_int, c_void, CStr, CString};
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::sync::Arc;
use arrow_array::ffi_stream::FFI_ArrowArrayStream;
use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator};
use arrow_schema::ffi::FFI_ArrowSchema;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
/// `SparkDsSink`: receives a byte array returned to Spark, and copies it.
type Sink = extern "C" fn(sink: *mut c_void, data: *const u8, length: i64);
struct RangeSource {
end: i64,
}
struct RangeReader {
end: i64,
}
thread_local! {
static LAST_ERROR: RefCell<CString> = RefCell::new(CString::default());
}
fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]))
}
/// Runs `f`, and reports its error or panic to Spark: neither may cross into the JVM.
fn call(f: impl FnOnce() -> Result<()>) -> c_int {
let message = match catch_unwind(AssertUnwindSafe(f)) {
Ok(Ok(())) => return 0,
Ok(Err(e)) => e.to_string(),
Err(_) => "the data source panicked".to_string(),
};
LAST_ERROR.with(|error| *error.borrow_mut() = CString::new(message).unwrap_or_default());
1
}
#[no_mangle]
pub extern "C" fn spark_ds_last_error() -> *const c_char {
LAST_ERROR.with(|error| error.borrow().as_ptr())
}
#[no_mangle]
pub extern "C" fn spark_ds_abi_version() -> i32 {
1
}
#[no_mangle]
pub unsafe extern "C" fn spark_ds_create(
_name: *const c_char,
option_keys: *const *const c_char,
option_values: *const *const c_char,
num_options: i64,
source: *mut *mut c_void,
) -> c_int {
call(|| {
let mut range = RangeSource { end: 10 };
for i in 0..num_options as usize {
if CStr::from_ptr(*option_keys.add(i)).to_str()? == "end" {
range.end = CStr::from_ptr(*option_values.add(i)).to_str()?.parse()?;
}
}
*source = Box::into_raw(Box::new(range)) as *mut c_void;
Ok(())
})
}
#[no_mangle]
pub unsafe extern "C" fn spark_ds_schema(_source: *mut c_void, out: *mut FFI_ArrowSchema) -> c_int {
call(|| {
std::ptr::write(out, FFI_ArrowSchema::try_from(schema().as_ref())?);
Ok(())
})
}
#[no_mangle]
pub unsafe extern "C" fn spark_ds_create_reader(
source: *mut c_void,
_schema: *mut FFI_ArrowSchema,
reader: *mut *mut c_void,
) -> c_int {
call(|| {
let end = (*(source as *const RangeSource)).end;
*reader = Box::into_raw(Box::new(RangeReader { end })) as *mut c_void;
Ok(())
})
}
#[no_mangle]
pub unsafe extern "C" fn spark_ds_close(source: *mut c_void) {
drop(Box::from_raw(source as *mut RangeSource));
}
#[no_mangle]
pub unsafe extern "C" fn spark_ds_partitions(
reader: *mut c_void,
add_partition: Sink,
sink: *mut c_void,
) -> c_int {
call(|| {
let end = (*(reader as *const RangeReader)).end;
for (start, end) in [(0, end / 2), (end / 2, end)] {
let bytes = [start.to_le_bytes(), end.to_le_bytes()].concat();
add_partition(sink, bytes.as_ptr(), bytes.len() as i64);
}
Ok(())
})
}
#[no_mangle]
pub unsafe extern "C" fn spark_ds_close_reader(reader: *mut c_void) {
drop(Box::from_raw(reader as *mut RangeReader));
}
#[no_mangle]
pub unsafe extern "C" fn spark_ds_read(
_reader_state: *const u8,
_reader_state_length: i64,
partition: *const u8,
partition_length: i64,
out: *mut FFI_ArrowArrayStream,
) -> c_int {
call(|| {
let bytes = std::slice::from_raw_parts(partition, partition_length as usize);
let start = i64::from_le_bytes(bytes[0..8].try_into()?);
let end = i64::from_le_bytes(bytes[8..16].try_into()?);
let ids = Int64Array::from_iter_values(start..end);
let batch = RecordBatch::try_new(schema(), vec![Arc::new(ids)])?;
let batches = RecordBatchIterator::new(vec![Ok(batch)], schema());
std::ptr::write(out, FFI_ArrowArrayStream::new(Box::new(batches)));
Ok(())
})
}
```
Build and install, here on macOS:
```bash
cargo build --release
cp target/release/libspark_datasource_rust_range.dylib /usr/local/lib/
```
Or package, 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 <<EOF
{
"abiVersion": 1,
"dataSources": ["rust_range"],
"libraries": {"osx-aarch64": "osx-aarch64/libspark_datasource_rust_range.dylib"}
}
EOF
(cd package && zip -r ../rust_range.sparkpkg .)
```
</details>
<details>
<summary><b>C++</b> (with Arrow C++)</summary>
`cpp_range.cc`:
```cpp
// A native data source for Apache Spark, written in C++ with Arrow C++. It reads the numbers
// [0, end) as a column `id`, in two partitions. The option `end` defaults to 10.
#include <cstdint>
#include <cstring>
#include <memory>
#include <stdexcept>
#include <string>
#include <arrow/api.h>
#include <arrow/c/bridge.h>
#include "spark_datasource.h"
namespace {
struct RangeSource {
int64_t end = 10;
};
struct RangeReader {
int64_t end;
};
thread_local std::string last_error;
std::shared_ptr<arrow::Schema> Schema() {
return arrow::schema({arrow::field("id", arrow::int64(), /*nullable=*/false)});
}
// Runs `f`, and reports its error to Spark: no C++ exception may cross into the JVM.
template <typename F>
int Call(F&& f) {
try {
f();
return 0;
} catch (const std::exception& e) {
last_error = e.what();
return 1;
} catch (...) {
last_error = "unknown error";
return 1;
}
}
void Check(const arrow::Status& status) {
if (!status.ok()) throw std::runtime_error(status.ToString());
}
} // namespace
extern "C" {
const char* spark_ds_last_error(void) { return last_error.c_str(); }
int32_t spark_ds_abi_version(void) { return SPARK_DS_ABI_VERSION; }
int spark_ds_create(const char*, const char* const* option_keys,
const char* const* option_values, int64_t num_options, void** source) {
return Call([&] {
auto range = std::make_unique<RangeSource>();
for (int64_t i = 0; i < num_options; i++) {
if (std::string(option_keys[i]) == "end") range->end = std::stoll(option_values[i]);
}
*source = range.release();
});
}
int spark_ds_schema(void*, struct ArrowSchema* out) {
return Call([&] { Check(arrow::ExportSchema(*Schema(), out)); });
}
int spark_ds_create_reader(void* source, struct ArrowSchema*, void** reader) {
return Call([&] { *reader = new RangeReader{static_cast<RangeSource*>(source)->end}; });
}
void spark_ds_close(void* source) { delete static_cast<RangeSource*>(source); }
int spark_ds_partitions(void* reader, SparkDsSink add_partition, void* sink) {
return Call([&] {
int64_t end = static_cast<RangeReader*>(reader)->end;
int64_t bounds[2][2] = {{0, end / 2}, {end / 2, end}};
for (const auto& partition : bounds) {
add_partition(sink, reinterpret_cast<const uint8_t*>(partition), sizeof(partition));
}
});
}
void spark_ds_close_reader(void* reader) { delete static_cast<RangeReader*>(reader); }
int spark_ds_read(const uint8_t*, int64_t, const uint8_t* partition, int64_t,
struct ArrowArrayStream* out) {
return Call([&] {
int64_t bounds[2];
std::memcpy(bounds, partition, sizeof(bounds));
arrow::Int64Builder ids;
for (int64_t id = bounds[0]; id < bounds[1]; id++) Check(ids.Append(id));
std::shared_ptr<arrow::Array> 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, out));
});
}
} // extern "C"
```
Get `spark_datasource.h`, which is in the `spark-sql` JAR, then build (Arrow C++ 24 requires C++20) and install, or package as above:
```bash
unzip -j "$SPARK_HOME"/jars/spark-sql_*.jar org/apache/spark/sql/datasource/spark_datasource.h
c++ -std=c++20 -O2 -shared -fPIC -I. -I"$ARROW_HOME/include" -L"$ARROW_HOME/lib" -larrow \
-o libspark_datasource_cpp_range.dylib cpp_range.cc
cp libspark_datasource_cpp_range.dylib /usr/local/lib/
```
</details>
Usage is the same as in #59313, on Java 22 or later, in classic PySpark or with Spark Connect. Once installed, the data sources are found automatically, without `spark.addArtifact` or any configuration:
```python
>>> spark.read.format("rust_range").option("end", 5).load().show()
+---+
| id|
+---+
| 0|
| 1|
| 2|
| 3|
| 4|
+---+
>>> spark.read.format("cpp_range").option("end", 5).load().filter("id >= 3").collect()
[Row(id=3), Row(id=4)]
>>> spark.sql("CREATE TABLE t USING rust_range OPTIONS (end 3)")
>>> spark.sql("SELECT * FROM t").collect()
[Row(id=0), Row(id=1), Row(id=2)]
```
Or added to a session as a package:
```python
>>> spark.addArtifact("rust_range/rust_range.sparkpkg")
>>> spark.read.format("rust_range").option("end", 5).load().count()
5
```
### Why are the changes needed?
Many data systems have their best client libraries in native code, such as arrow-rs, DataFusion, delta-rs, iceberg-rust and Lance in Rust, or C++ libraries. Using them from Spark today requires writing a JVM Data Source V2 connector with its own JNI glue. With this change, a native data source implements a stable binary interface once, is installed like any other native library or distributed as a single file, and works with every Spark API, including Spark Connect clients in any language, without any JVM code.
With FFM, the interface is also a plain C ABI, so a library needs no JVM headers or bindings, and other hosts than Spark could load the same libraries.
### Does this PR introduce _any_ user-facing change?
Yes. It adds native data sources, with the `spark_datasource.h` interface, the configurations and the error conditions above, on Java 22 or later. Spark finds the native data source libraries installed in the native library path, and `spark.addArtifact` now accepts native data source packages (`.sparkpkg` files). Also, in classic mode, writing with a session that is not the active one now finds the data sources of that session, such as its Python data sources, because `DataFrameWriter` and `DataStreamWriter` resolve the data source with that session active.
### How was this patch tested?
All the tests of #59313, ported to the C interface:
- `NativeDataSourceSuite`, with the dependency-free C++ test library (`sql/core/src/test/resources/native-datasource/test_native_datasource.cc`), now implementing `spark_datasource.h`, which the tests compile with the C++ compiler of the machine against the header from the classpath. It covers inferred and user-specified schemas, partitions and batches, the supported types, predicate, column and limit pushdown, batch writes (append and overwrite), aborted writes, streaming reads and writes over several micro-batches with offset commits, SQL tables with `OPTIONS` (`SELECT` and `INSERT INTO`), errors reported by native code, missing functions, several libraries exporting the same functions side by side, `spark.addArtifact`, packages in directories, package conflicts, the precedence of Java data sources, invalid packages, installed libraries and invalid installed libraries, plus the unit tests for the manifest, the platform names, the native library path and the JSON encoding of pred
icates.
- Tests added during the review, in `NativeDataSourceSuite` unless noted:
- Packages: invalid and unsupported packages are ignored unless no other package or installed library provides the data source; a configured package replaced in place; the copies distributed to the executors are named by checksum and are not packages of the session; the cached packages of a session are forgotten when its artifacts are cleaned up; configured directories that cannot be read; distributing a package without an active session.
- Lookups: a single lookup for each resolution, which the next lookup on the thread clears; a Python data source registered with the name of an installed native one; batch and streaming writes with a session that is not the active one.
- Writes: a failed `commitDataWriter` aborts the data writer; a failed `commit` aborts the write, then closes the writer; the test library records when a writer is closed, which checks that a batch write closes its writer, and that each micro-batch of a streaming write closes the writer of the previous one; `ColumnarDataSourceSuite` checks that each micro-batch creates its stream writer with the ID of the query, and that batches end at `spark.sql.execution.arrow.maxBytesPerBatch`.
- Cluster mode: new `NativeDataSourceClusterSuite` runs in `local-cluster[2,1,1024]`, where the executors are separate JVMs: batch reads, writes and streaming with a package added with `spark.addArtifact`, and a configured package that the executors can only load from its distributed copy (named by checksum, under the UUID of the session), which they check by its checksum.
- Installed libraries: an executor does not load a library with another checksum than on the driver, at the driver's path or in its native library path.
- Predicates: only the predicates listed in `spark_datasource.h` are pushed down, not arithmetic or other functions.
- `spark.addArtifact` in classic PySpark accepts `file://` URIs and paths relative to the working directory of Python.
- Reads: predicates on columns with collated strings are not pushed down, read schemas with collations, and the scan described by the name of the data source; sliced arrays, arrays shorter than their parent (a column, a struct field or the elements of a map), `decimal256` and list views are rejected with the name of the column, maps whose entry fields are named `keys` and `values` (as arrow-rs names them) are read, and `limit(...).offset(...)` asks the library for the limit once.
- Writes whose upstream stages fail with adaptive execution create no writer (also in `ColumnarDataSourceSuite`).
- `NativeDataSourceConnectSuite` (Spark Connect), `NativeDataSourceClusterSuite` (`local-cluster`, executors in separate JVMs) and `ColumnarDataSourceSuite`.
New tests:
- `native data sources require Java 22 or later`: on Java 21 or earlier, a query that uses an installed native data source fails with `NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION`, before Spark loads the library. It runs only on Java 21 or earlier, so it runs in the default CI on Java 17.
- `report the errors of native code`: a library without `spark_ds_last_error` reports `spark_ds_create failed with the error code 1.`
The tests that load a library are skipped before Java 22, like they are skipped without a C++ compiler. So the default CI, on Java 17, only runs the tests that do not load a library, and the scheduled Java 25 builds run them all. Locally on macOS (osx-aarch64):
- Java 26 (Homebrew OpenJDK 26.0.1; Spark built with Java 17): `NativeDataSourceSuite`, `NativeDataSourceClusterSuite` (with the executors on Java 26 too), `ColumnarDataSourceSuite` and `NativeDataSourceConnectSuite` pass. With the executors of `local-cluster` on Java 17, the tasks fail with `NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION`, as expected.
- Java 17: `NativeDataSourceSuite` passes, with the tests that load a library skipped, and the Java version test run.
- The Rust and C++ examples above were built with the commands above (Rust with arrow-rs 58, C++ with Arrow C++ 24 and `-std=c++20`), installed in a directory of `java.library.path`, and read by Spark on Java 26: `spark.read.format("rust_range").option("end", 100).load()` and the same with `cpp_range` return the 100 rows in two partitions, and `option("end", "x")` fails with `NATIVE_DATA_SOURCE_ERROR` and the message of the library (`invalid digit found in string` and `stoll: no conversion`).
- The documented steps of the C++ example also work from a packaged JAR: `unzip -j` of `spark_datasource.h` from the `spark-sql` JAR built by `sbt sql/package`, then the documented `c++` command with `-I.`.
- The C++ test library builds without warnings with Apple clang (`-Wall -Wextra`), with and without `READ_ONLY`.
- `SparkThrowableSuite`, for the new error condition.
- `checkstyle` and `scalastyle` of the changed modules.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
This pull request and its description were written by Isaac.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "FelixYBW (via GitHub)" <gi...@apache.org>.
FelixYBW commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6208572316
@zhouyuan
```
read -> ArrowArrayStream -> Arrow Java -> ArrowColumnVector
-> ColumnarToRow -> RowToVeloxColumnar -> Velox operators
```
This is what @zhztheplayer's proposal [SPARK-57468](https://issues.apache.org/jira/browse/SPARK-57468) tries to solve exactly.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "gerashegalov (via GitHub)" <gi...@apache.org>.
gerashegalov commented on code in PR #59424:
URL: https://github.com/apache/spark/pull/59424#discussion_r4316541208
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,1098 @@
+/*
+ * 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
+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.{ListVector, MapVector, StructVector}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
+
+import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.util.Utils
+
+/**
+ * A [[ColumnarDataSource]] implemented by a native library, through the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @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 =>
+ val schema = call(name, "infer the schema of") {
+ optional(NativeArrow.importSchema(address => library.schema(dataSource.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ NativeArrow.checkSupported(schema, name, "infer the schema of")
+ schema
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
Review Comment:
The docs say continuous streaming is not supported, but nothing stops it: this table always reports `STREAMING_WRITE`, and Spark accepts any such sink with `Trigger.Continuous`. I ran `rate` into `native_sink` with `Trigger.Continuous("1 second")` on Java 25, and it ran:
- One stream writer was created and used for every epoch.
- When the query stopped, the data writers of the last epoch were closed without a commit or an abort. Their files were left in `_temporary/`.
The header promises a library that a data writer is always committed or aborted before it is closed, so a library can't rely on that here.
Could `ColumnarStreamingWrite.createStreamingWriterFactory` fail when the write runs in continuous mode (for example, when the local property `ContinuousExecution.EPOCH_COORDINATOR_ID_KEY` is set), with a test?
##########
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.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_<name>`, such as
+ * `libspark_datasource_my_source.so` on Linux, implements the data source `<name>`. 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 =>
+ // The version of a package is the oldest that its library implements.
+ if (pkg.manifest.abiVersion <= NativeLibrary.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" -> NativeLibrary.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(
Review Comment:
The docs suggest putting the version in the file name of a package. But if a session adds `rs-1.0.sparkpkg` and then `rs-1.1.sparkpkg`, every later query on that source fails with `NATIVE_DATA_SOURCE_PACKAGE_CONFLICT`, which says to remove all but one. A session has no way to remove an artifact, and adding `rs-1.0.sparkpkg` again with new contents fails with `ARTIFACT_ALREADY_EXISTS`. So the session is stuck until it is closed. I checked this on Java 25.
The same conflict happens when an admin's package in `spark.sql.dataSource.native.paths` and a user's added package provide the same name.
Could Spark choose one, for example the package the session added last, then the configured ones? Or else the docs could say that an upgrade needs a new session, and stop suggesting versioned file names.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala:
##########
@@ -0,0 +1,142 @@
+/*
+ * 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._
+import org.apache.spark.unsafe.types.UTF8String
+
+/**
+ * Encodes Data Source V2 predicates as the JSON expression trees passed to
+ * `push_predicates` of `spark_datasource.h`.
+ */
+object NativePredicates {
+ // Writes ASCII only, escaping the other characters, which any JSON parser decodes.
+ private val mapper = JsonMapper.builder().enable(JsonWriteFeature.ESCAPE_NON_ASCII).build()
+
+ /**
+ * The names of the predicates that are passed to the libraries, as listed in
+ * `push_predicates`. They are part of the interface, so a predicate with any other
+ * name, such as an arithmetic function, is not passed: its semantics, for example on overflow,
+ * would be part of the interface too, and the names of Data Source V2 grow with new releases.
+ */
+ 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] = {
Review Comment:
This passes a column of any type, so `s1 = s2` on two struct columns, and `a1 < a2` on two array columns, are pushed to the library as they are. I checked both with `V2ExpressionBuilder` and `toJson`.
The rules in the header only cover single values. For example, they say comparisons with null are null. But Spark treats nulls inside structs and arrays as equal to each other and smaller than other values, so `array(1, null) = array(1, null)` is true. A library that follows the header can skip those rows, and Spark can't get them back by checking the predicates again after the scan.
Could comparisons and `IN` only be pushed for columns of simple types (null checks are fine for any type)? Or else the header could describe the rules for nested values. It could also say that binary values compare as unsigned bytes.
##########
sql/core/pom.xml:
##########
@@ -298,6 +298,10 @@
<groupId>org.apache.arrow</groupId>
<artifactId>arrow-compression</artifactId>
</dependency>
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>arrow-c-data</artifactId>
Review Comment:
Spark calls the library with FFM, but it still exchanges Arrow data through `arrow-c-data`, which loads Arrow's own JNI library, `libarrow_cdata_jni`. That library is only shipped for Linux and macOS on x86_64 and aarch64, and for Windows on x86_64. On ppc64le, s390x or Alpine (musl), a correctly installed native data source fails on its first read, with an Arrow error that doesn't mention native data sources.
Could the docs list the supported platforms, and Spark fail early with a clear error elsewhere? The platform names in the manifest also can't tell glibc from musl, which will be hard to change once 4.4.0 is out.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,1098 @@
+/*
+ * 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
+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.{ListVector, MapVector, StructVector}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
+
+import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.util.Utils
+
+/**
+ * A [[ColumnarDataSource]] implemented by a native library, through the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @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 =>
+ val schema = call(name, "infer the schema of") {
+ optional(NativeArrow.importSchema(address => library.schema(dataSource.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ NativeArrow.checkSupported(schema, name, "infer the schema of")
+ schema
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ NativeArrow.checkSupported(schema, name, "plan a scan of")
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, executorLocation, name, schema, new NativeHandle(handle, library.closeReader))
+ }
+
+ override def streamReader(schema: StructType): ColumnarStreamReader = {
+ NativeArrow.checkSupported(schema, name, "plan a scan of")
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, address)
+ }
+ }
+ 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ entries.find(e => e.getKey.contains('\u0000') || e.getValue.contains('\u0000')).foreach { e =>
+ throw error(name, "create", s"The option ${NativeArrow.printable(e.getKey)} contains the " +
+ "character U+0000, which Spark cannot pass to a library.")
+ }
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ NativeArrow.checkSupported(schema, name, action)
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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)
+ }
+
+ /**
+ * Checks that the data agrees with the nullability of the schema to read: a column, a field,
+ * an element or a map value that it declares not nullable has no null, except where its parent
+ * is null. Spark relies on the nullability, for example to evaluate `isnull` of such a column to
+ * false without reading it.
+ */
+ private def checkNulls(): Unit = {
+ root.getFieldVectors.asScala.zip(schema.fields).foreach { case (vector, field) =>
+ checkNulls(vector, field.dataType, field.nullable, field.name, _ => true)
+ }
+ }
+
+ /** `referenced` tells whether a value belongs to a value of the parent that is not null. */
+ private def checkNulls(
+ vector: ValueVector,
+ dataType: DataType,
+ nullable: Boolean,
+ column: String,
+ referenced: Int => Boolean): Unit = {
+ if (!nullable && vector.getNullCount > 0 &&
+ (0 until vector.getValueCount).exists(i => vector.isNull(i) && referenced(i))) {
+ throw error(name, "read", s"The column $column has a null, but the schema to read " +
+ "declares it not nullable.")
+ }
+ val valid = (i: Int) => !vector.isNull(i) && referenced(i)
+ (dataType, vector) match {
+ case (struct: StructType, structVector: StructVector) =>
+ struct.fields.zipWithIndex.foreach { case (field, i) =>
+ checkNulls(structVector.getChildByOrdinal(i), field.dataType, field.nullable,
+ s"$column.${field.name}", valid)
+ }
+ case (map: MapType, mapVector: MapVector) =>
+ val entries = mapVector.getDataVector.asInstanceOf[StructVector]
+ val inMaps = elements(mapVector, valid)
+ // Spark does not allow null keys.
+ checkNulls(entries.getChildByOrdinal(0), map.keyType, nullable = false, s"$column.key",
+ inMaps)
+ checkNulls(entries.getChildByOrdinal(1), map.valueType, map.valueContainsNull,
+ s"$column.value", inMaps)
+ case (array: ArrayType, list: ListVector) =>
+ checkNulls(list.getDataVector, array.elementType, array.containsNull, s"$column.element",
+ elements(list, valid))
+ case _ =>
+ }
+ }
+
+ /** Whether each element of a list belongs to a list that `valid` accepts. */
+ private def elements(list: ListVector, valid: Int => Boolean): Int => Boolean = {
+ lazy val referenced = {
+ val bits = new java.util.BitSet(list.getDataVector.getValueCount)
+ (0 until list.getValueCount).foreach { i =>
+ if (valid(i)) bits.set(list.getElementStartIndex(i), list.getElementEndIndex(i))
+ }
+ bits
+ }
+ i => referenced.get(i)
+ }
+
+ 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))
+ checkNulls()
+ 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 = {
+ // The arrays of the children, one for each field of the type: Arrow Java only checks their
+ // number when an array has children, and imports the missing columns of a batch as nulls.
+ def children(
+ array: ArrowArray.Snapshot,
+ fields: Int,
+ owner: String,
+ expected: String): Seq[ArrowArray.Snapshot] = {
+ if (array.n_children != fields || (fields > 0 && array.children == 0)) {
+ throw error(name, "read", s"$owner has ${array.n_children} child arrays, but $expected.")
+ }
+ (0 until fields).map { i =>
+ val child = Platform.getLong(null, array.children + 8L * i)
+ if (child == 0) {
+ throw error(name, "read", s"$owner has a NULL child array.")
+ }
+ ArrowArray.wrap(child).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)
+ val fields = field.getChildren.size
+ children(array, fields, s"An array of the column $column", s"its type has $fields")
+ .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.")
+ }
+ val columns = root.getSchema.getFields.size
+ children(batch, columns, "The batch", s"the schema has $columns columns")
+ .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()
Review Comment:
`stream.release()` runs after `root.close()` and `dictionaries.close()` in the same block, while `task.close()` is in the `finally`. If either Arrow close throws, the library never gets `release`, but the task context is freed anyway. A library thread that still checks `task->is_interrupted` then reads freed memory. The header says the task stays valid until the stream is released.
Unlikely, but easy to rule out: could the stream be released in its own `finally`, before `task.close()`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,1098 @@
+/*
+ * 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
+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.{ListVector, MapVector, StructVector}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
+
+import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.util.Utils
+
+/**
+ * A [[ColumnarDataSource]] implemented by a native library, through the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @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 =>
+ val schema = call(name, "infer the schema of") {
+ optional(NativeArrow.importSchema(address => library.schema(dataSource.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ NativeArrow.checkSupported(schema, name, "infer the schema of")
+ schema
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ NativeArrow.checkSupported(schema, name, "plan a scan of")
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, executorLocation, name, schema, new NativeHandle(handle, library.closeReader))
+ }
+
+ override def streamReader(schema: StructType): ColumnarStreamReader = {
+ NativeArrow.checkSupported(schema, name, "plan a scan of")
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, address)
+ }
+ }
+ 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ entries.find(e => e.getKey.contains('\u0000') || e.getValue.contains('\u0000')).foreach { e =>
+ throw error(name, "create", s"The option ${NativeArrow.printable(e.getKey)} contains the " +
+ "character U+0000, which Spark cannot pass to a library.")
+ }
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ NativeArrow.checkSupported(schema, name, action)
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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)
+ }
+
+ /**
+ * Checks that the data agrees with the nullability of the schema to read: a column, a field,
+ * an element or a map value that it declares not nullable has no null, except where its parent
+ * is null. Spark relies on the nullability, for example to evaluate `isnull` of such a column to
+ * false without reading it.
+ */
+ private def checkNulls(): Unit = {
+ root.getFieldVectors.asScala.zip(schema.fields).foreach { case (vector, field) =>
+ checkNulls(vector, field.dataType, field.nullable, field.name, _ => true)
+ }
+ }
+
+ /** `referenced` tells whether a value belongs to a value of the parent that is not null. */
+ private def checkNulls(
+ vector: ValueVector,
+ dataType: DataType,
+ nullable: Boolean,
+ column: String,
+ referenced: Int => Boolean): Unit = {
+ if (!nullable && vector.getNullCount > 0 &&
+ (0 until vector.getValueCount).exists(i => vector.isNull(i) && referenced(i))) {
+ throw error(name, "read", s"The column $column has a null, but the schema to read " +
+ "declares it not nullable.")
+ }
+ val valid = (i: Int) => !vector.isNull(i) && referenced(i)
+ (dataType, vector) match {
+ case (struct: StructType, structVector: StructVector) =>
+ struct.fields.zipWithIndex.foreach { case (field, i) =>
+ checkNulls(structVector.getChildByOrdinal(i), field.dataType, field.nullable,
+ s"$column.${field.name}", valid)
+ }
+ case (map: MapType, mapVector: MapVector) =>
+ val entries = mapVector.getDataVector.asInstanceOf[StructVector]
+ val inMaps = elements(mapVector, valid)
+ // Spark does not allow null keys.
+ checkNulls(entries.getChildByOrdinal(0), map.keyType, nullable = false, s"$column.key",
+ inMaps)
+ checkNulls(entries.getChildByOrdinal(1), map.valueType, map.valueContainsNull,
+ s"$column.value", inMaps)
+ case (array: ArrayType, list: ListVector) =>
+ checkNulls(list.getDataVector, array.elementType, array.containsNull, s"$column.element",
+ elements(list, valid))
+ case _ =>
+ }
+ }
+
+ /** Whether each element of a list belongs to a list that `valid` accepts. */
+ private def elements(list: ListVector, valid: Int => Boolean): Int => Boolean = {
+ lazy val referenced = {
+ val bits = new java.util.BitSet(list.getDataVector.getValueCount)
+ (0 until list.getValueCount).foreach { i =>
+ if (valid(i)) bits.set(list.getElementStartIndex(i), list.getElementEndIndex(i))
+ }
+ bits
+ }
+ i => referenced.get(i)
+ }
+
+ 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))
+ checkNulls()
+ 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 = {
+ // The arrays of the children, one for each field of the type: Arrow Java only checks their
+ // number when an array has children, and imports the missing columns of a batch as nulls.
+ def children(
+ array: ArrowArray.Snapshot,
+ fields: Int,
+ owner: String,
+ expected: String): Seq[ArrowArray.Snapshot] = {
+ if (array.n_children != fields || (fields > 0 && array.children == 0)) {
+ throw error(name, "read", s"$owner has ${array.n_children} child arrays, but $expected.")
+ }
+ (0 until fields).map { i =>
+ val child = Platform.getLong(null, array.children + 8L * i)
+ if (child == 0) {
+ throw error(name, "read", s"$owner has a NULL child array.")
+ }
+ ArrowArray.wrap(child).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)
+ val fields = field.getChildren.size
+ children(array, fields, s"An array of the column $column", s"its type has $fields")
+ .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) {
Review Comment:
The checks here look at the offset and the children of the batch, but not its `null_count`. As far as I can tell, Arrow Java ignores the validity of the top-level struct when it imports a batch. So if a library marks some rows of the batch itself as null, Spark reads them as ordinary rows built from the child values, with no error. Arrow C++ rejects such a batch on import.
Could this also require `null_count == 0`, or `-1` with no validity buffer? It is a constant-time check.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "tdcmeehan (via GitHub)" <gi...@apache.org>.
tdcmeehan commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6173067539
Thanks for moving to the C header, and for documenting the lifecycle so carefully. I have one concern, about how the header grows. I think it's worth settling before 4.4.0: version 1 is what library binaries get built against, and this is the one change that can't come later as an optional function.
**The problem: the header grows by one function per feature.**
1. **Every feature is a new function.** The Compatibility section adds features as optional functions, and the Limitations section lists six to add that way: statistics, custom metrics, admission control, aggregate/sort/partitioning pushdown, cancellation and memory accounting. The predicate declaration and three-state answer from #59313 are planned the same way. DSv2 itself keeps growing: master has 37 `Supports*` interfaces, 7 of them added since July 2025.
2. **Each function is a separate contract that every host and SDK has to reimplement.** Each has its own argument encoding, ownership, threading and ordering rules, written in prose. There are five handle types, released by six different functions. That's where bugs land: @LuciferYang's catch in this PR, a released data writer passed to `spark_ds_abort_data_writer`, came from one of those rules, in Spark's own binding.
3. **Signatures are frozen, so a change is a second function.** For example, giving `spark_ds_read` a task context for memory accounting or cancellation means a second `read`, and a library that also runs on older releases exports both.
4. **Support is inferred from symbol names.** It's per library, not per data source or options, and a library can't declare which predicates it handles. A missing export, for example a misspelled name, isn't an error: the operation becomes unsupported, or, for the pushdown and serialize functions, silently does nothing.
5. **Nothing flows from Spark to the library.** The library reports its version, but Spark never tells it Spark's version or features, and there's no callback table. Cancellation, memory accounting and logging all need one.
**Proposal: freeze the C core, and move everything that grows into versioned messages.**
The header is already halfway there: 14 of its 30 functions pass bytes or JSON (partitions, state, commit messages, offsets, predicates). This finishes the move, with one schema. A sketch, with placeholder names:
```c
int32_t spark_ds_abi_version(void);
const char* spark_ds_last_error(void);
/* Opens a data source, reader, writer or data writer. The request says which, with what it
* needs, plus Spark's protocol version and features. The reply has the library's version and
* what this object supports. */
int spark_ds_open(const SparkDsHost* host, const uint8_t* request, int64_t length,
SparkDsSink reply, void* sink, void** handle);
/* Every control operation: schema, pushdown, partitions, offsets, commit, abort. */
int spark_ds_call(void* handle, const uint8_t* request, int64_t length,
SparkDsSink reply, void* sink);
/* Every data operation, as Arrow C streams: a read fills `out`, a write consumes `in`. */
int spark_ds_stream(void* handle, const uint8_t* request, int64_t length,
struct ArrowArrayStream* in, struct ArrowArrayStream* out);
/* Releases any handle. Nothing else does. */
void spark_ds_close(void* handle);
```
- `SparkDsHost` holds Spark's callbacks, such as memory reservation, cancellation and logging. It starts with its size, and fields are only appended.
- Requests and replies are protobuf messages, in a `.proto` next to the header. Old code ignores new fields, and field numbers are never reused.
- Spark only sends what the library declared in its open reply, so a new request never reaches an old library.
- Data stays as it is: zero-copy Arrow C streams.
How the planned additions would land:
| Addition | As optional functions | As messages |
|---|---|---|
| Declared predicates, three-state answer | new pushdown functions | a list in the open reply, an enum in the pushdown reply |
| Memory accounting, cancellation | a setup function each, or a second `read` | fields appended to `SparkDsHost` |
| Statistics, partitioning, sort order | a function each | fields in the planning reply |
| Aggregate and top-N pushdown | a function each | fields in the pushdown request |
| Admission control | a second `latest_offset` | a field in the offset request |
| Overwrite by filter | a new function, since `bool overwrite` can't carry a filter | a write mode and a predicate field |
| Substrait, as @lidavidm suggested | a new pushdown function | another encoding of the predicate field |
What it buys:
- A feature is new fields, with generated code in every language. No new C function, binding or ownership rule.
- One ownership rule: request bytes live for the call, the sink copies replies, `close` releases handles, and Arrow structs follow Arrow's rules.
- Tooling checks compatibility. Spark's CI already runs [`buf breaking`](https://github.com/apache/spark/blob/4ccf07b6d91e2b126b7001741778ee28f22c6609/.github/workflows/build_and_test.yml#L834-L838) on the Spark Connect protos, and recorded messages from each release can be replayed against current libraries.
- A host that runs a library in another process, for example to isolate it on a shared server, forwards the control messages unchanged.
- If Spark ever wants native data sources on Java 17 and 21, a JNI bridge would be six functions that never change.
What it costs:
- Library authors need a protobuf runtime. SDKs can vendor it and link it statically, with hidden symbols, since several libraries can share a process. Spark already uses protobuf for Spark Connect and the UDF worker protocol.
- The `.proto` and its generated types check the contract, not the C compiler.
- Control calls pay for encoding. They're rare, and data stays zero-copy.
- Semantics still need specifying. The NaN, null and partial-predicate rules you wrote carry over unchanged.
Precedents:
- **Arrow** splits it the same way. The C Data Interface is a few frozen structs. The IPC format is versioned messages: since 1.0, which guarantees forward and backward compatibility, it has been [extended five times](https://github.com/apache/arrow/blob/910292405bb82ccbf057b9801f5e4b922ea23788/format/Schema.fbs#L19-L26), and its [`Feature` enum](https://github.com/apache/arrow/blob/910292405bb82ccbf057b9801f5e4b922ea23788/format/Schema.fbs#L54-L81) is "A means of negotiating between a client and server what features a stream is allowed to use".
- **Comet** shows both sides. Its plan, with 24 operator, 66 expression and 27 aggregate kinds, crosses one JNI function, [`createPlan`](https://github.com/apache/datafusion-comet/blob/7d8d9b1544b734c85a93cca2cc9d5380ad70accd/spark/src/main/scala/org/apache/comet/Native.scala#L61-L79), as protobuf, and a [contrib envelope](https://github.com/apache/datafusion-comet/blob/7d8d9b1544b734c85a93cca2cc9d5380ad70accd/native/proto/src/proto/operator.proto#L49-L81) means "core's proto never has to change when a contrib is added". Its other JNI functions went from 9 to 19 in the past year, and `createPlan` gained two parameters. That's fine for Comet, whose JVM and native halves ship in one jar; under a public ABI, each would be a second function or a break. **Gluten** also passes plans as Substrait bytes, and [asks the native side](https://github.com/apache/gluten/blob/91ae3b22fcb1dbff6937ae517fc0d816c23f4501/gluten-arrow/src/main/java/org/apache/gluten/vectorized/PlanEvaluatorJniWrapper.j
ava#L46-L82) whether it supports a plan.
- **Spark's UDF worker protocol** (SPARK-55278) is protobuf, with [`Init.protocol_version`](https://github.com/apache/spark/blob/4ccf07b6d91e2b126b7001741778ee28f22c6609/udf/worker/proto/src/main/protobuf/org/apache/spark/udf/worker/udf_message.proto#L162-L167) and [`WorkerCapabilities`](https://github.com/apache/spark/blob/4ccf07b6d91e2b126b7001741778ee28f22c6609/udf/worker/proto/src/main/protobuf/org/apache/spark/udf/worker/worker_spec.proto#L113-L135) that Spark checks during planning.
- **protoc plugins:** the compiler sends its version in the request, and the plugin answers with [`supported_features`](https://github.com/protocolbuffers/protobuf/blob/cd68a2eab18e9646bf746182fa3b858c6d9caa7e/src/google/protobuf/compiler/plugin.proto#L79-L115).
- **OpenSSL 3 providers** run in process, with dispatch tables and typed key/value parameters. Result: "A FIPS provider built from any validated version may be used together with an OpenSSL library built from any supported release from OpenSSL 3.0 onwards" ([README-FIPS.md](https://github.com/openssl/openssl/blob/ed57129cd43b77e077c68fd5f7ea951976766f38/README-FIPS.md?plain=1#L36-L41)).
- **ADBC** is a wide C API that works, and shows what it takes. Drivers fill a function table through a versioned init function, and a driver manager handles the versions. The table has grown once in four years: revision 1.1.0 (2023) took it from 26 to 55 functions, 21 of them typed option getters and setters, and needed a [handshake](https://github.com/apache/arrow-adbc/blob/fe88b72a9917763a1ff2eb694341b5b6da8aa6f4/c/include/arrow-adbc/adbc.h#L1055-L1074) in which a 1.0.0 driver must fail and be retried. Since then it has added option keys and info codes, but no functions. DSv2, which this header mirrors, added 7 optional interfaces in the last 15 months.
- The other way: **DataFusion FFI** mirrors DataFusion's traits as C structs, and [recommends](https://github.com/apache/datafusion/blob/0d79252f61bd8777990b2749c216ffd07079c30d/datafusion/ffi/README.md?plain=1#L38-L40) "the same version of DataFusion by both the producer and consumer", though it's not strictly required.
Could version 1 take this shape? I'm happy to write the `.proto` for the current feature set, and the replay test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6481066910
For tracking, here are the rest of my suggestions from the Rust SDK side (#59520), which can be follow-ups, after [the three above](https://github.com/apache/spark/pull/59424#issuecomment-6480975990). Checked at `1f40484`.
1. **Give the stream reader the inputs of `toMicroBatchStream` and of admission control.**
[`ColumnarTable.toMicroBatchStream`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala#L167-L168) drops `checkpointLocation`, since [`create_stream_reader`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L419) only takes the schema, and [`latest_offset`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L492) has no place for the start offset and the `ReadLimit` of `SupportsAdmissionControl.latestOffset`, so a source cannot limit the size of a micro-batch (#59520 leaves admission control out for this reason). These functions take positional arguments, so they can only get new inputs through new functions, unless they take size-prefixed stru
cts, like `create` and `create_writer`, before 4.4.0.
2. **Prune the columns of streaming reads.**
[`toMicroBatchStream`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala#L167-L168) creates the stream reader with `fullSchema`, and [`prune_columns`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L456) only takes a batch reader, so a streaming read always reads every column. It could pass the pruned schema to `create_stream_reader`, or call a pruning function on stream readers, with a capability. The same would apply to predicates when they come back.
3. **Make the order of the calls on each handle a rule of the header.**
Spark already keeps a failed planning and stops a stream after a failure, but the header does not promise it, so the Rust SDK keeps the first failure of a stream reader and returns it from every later call ([`StreamReader::call`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/bridge.rs#L559-L568)), and tracks each writer through [`Phase`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/write.rs#L37-L42) to guard against repeated commits and aborts. Two rules would let SDKs drop most of this: after a function on a handle fails, Spark only calls its abort or close function; and Spark calls `commit_data_writer` and `abort_data_writer` at most once for each data writer, and never `write` after either. This only changes comments.
4. **Tell the library which log levels are enabled.**
[`SparkDsHost.log`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L286) does not say which levels Spark logs, so the Rust SDK [passes `Info` and more severe messages](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/host.rs#L89) unless the library sets another level: debug messages are lost when the log of Spark enables them, and formatted for nothing otherwise. A function appended to `SparkDsHost`, such as `bool (*is_log_enabled)(int32_t level)`, would let an SDK follow the level of the log of Spark, also when it changes at runtime.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #59424:
URL: https://github.com/apache/spark/pull/59424#discussion_r4299567529
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,980 @@
+/*
+ * 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
+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, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, 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 the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @param location where the library is on the driver. The readers and the writers distribute it
+ * to the executors of the active session when they are created.
+ * @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.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, 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", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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))
Review Comment:
Thanks, this is the bug that @FelixYBW reported. A fresh root for each batch, as `Data.importVectorSchemaRoot` creates, would not fix it alone: the reader would still have to release each root at some point, and nothing tells it when a consumer is done with a batch. `ColumnarBatch` lets a reader reuse its vectors, and @sarutak showed that `ArrowCachedBatchSerializer`, which creates a root for each batch and closes the previous one in `next()`, fails in the same way. So the fix is in the consumer, which kept references to the pass-through columns: SPARK-60294, #59551, copies them. This PR keeps one root, and adds a test of Arrow UDFs over a native data source in 74aa0dffc534.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,980 @@
+/*
+ * 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
+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, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, 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 the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @param location where the library is on the driver. The readers and the writers distribute it
+ * to the executors of the active session when they are created.
+ * @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.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, 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", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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) =>
Review Comment:
Thanks, fixed in 74aa0dffc534. `checkArrays` now requires the batch, and each array at every level, to have as many children as its type has fields, with non-NULL pointers, before it reads them, so it does not read past the array of pointers either. The test library can now export batches without their columns (`short=children`), and `reject the Arrow data that Spark cannot read` checks that the read fails with `The batch has 0 child arrays, but the schema has 7 columns.`
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,980 @@
+/*
+ * 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
+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, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, 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 the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @param location where the library is on the driver. The readers and the writers distribute it
+ * to the executors of the active session when they are created.
+ * @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.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, 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", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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)
Review Comment:
Thanks, fixed in 74aa0dffc534. After importing a batch, Spark now checks that the data agrees with the nullability of the schema to read, at every level: a column, a struct field, an element or a map value that the schema declares not nullable must have no null where its parent is not null, and map keys none at all. It starts from the null counts of the vectors, so only a column that has nulls is scanned further. A null fails the read with `NATIVE_DATA_SOURCE_ERROR`, such as `The column name has a null, but the schema to read declares it not nullable.`, and the header says that the data must agree with the nullability of the schema. I checked the data rather than making the inferred schema nullable, since a user can also declare a column `NOT NULL` in the schema to read, and the schema of a table in a catalog also applies to its writes. Test: `reject the data that has nulls where the schema to read has none`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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
+ * `push_predicates` of `spark_datasource.h`.
+ */
+object NativePredicates {
+ // Writes ASCII only, escaping the other characters, which any JSON parser decodes.
+ private val mapper = JsonMapper.builder().enable(JsonWriteFeature.ESCAPE_NON_ASCII).build()
+
+ /**
+ * The names of the predicates that are passed to the libraries, as listed in
+ * `push_predicates`. They are part of the interface, so a predicate with any other
+ * name, such as an arithmetic function, is not passed: its semantics, for example on overflow,
+ * would be part of the interface too, and the names of Data Source V2 grow with new releases.
+ */
+ 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))
Review Comment:
Thanks, both fixed in 74aa0dffc534:
- A float literal is now written as the double of the same value, `0.10000000149011612` for `0.1f`, which a reader that parses JSON numbers as doubles gets exactly, and which a reader that parses it as a float reads back as the same float. The header says to read a literal at the precision of its data type, and that the values of the integer types, dates and timestamps are exact 64-bit integers. `predicate-cases.json` has a float column now, with `f = CAST(0.1 AS FLOAT)` and `f > CAST(0.1 AS FLOAT)`, and writes the floats of its rows the same way.
- A string literal that is not valid UTF-8 is no longer pushed, and the header says so.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -0,0 +1,688 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*
+ * spark_datasource.h: the binary interface between Spark and data sources implemented in native
+ * code, such as Rust or C++. Since Spark 4.4.0.
+ *
+ * A native data source is a shared library that exports a single C function, spark_ds_init,
+ * which returns the table of the functions that implement the data source, a SparkDsLibrary.
+ * Spark calls them with the Foreign Function and Memory API of Java, so native data sources
+ * require Java 22 or later. Each library has its own table, so any number of libraries can be
+ * loaded side by side. Spark adapts a native data source to Data Source V2, with batch and
+ * micro-batch streaming reads, batch and streaming writes, and predicate, column and limit
+ * pushdown.
+ *
+ * This header is self-contained: it only depends on the C standard library, and defines the
+ * structs of the Arrow C data interface unless they are already defined, as by Arrow C++. A library
+ * can copy it into its source tree, or use the copy in the include directory of a Spark
+ * distribution. A library built against it works with all the releases of Spark that support its
+ * SPARK_DS_ABI_VERSION.
+ *
+ * Spark finds the library of a data source automatically when it is installed like other native
+ * libraries and named after the data source: the library of my_source is
+ * spark_datasource_my_source, such as libspark_datasource_my_source.so in /usr/local/lib on Linux.
+ * Spark looks for it in the directories of java.library.path, which include LD_LIBRARY_PATH on
+ * Linux, and in the lib directory of each installation prefix in PATH. A library can also be
+ * distributed in a native data source package: a zip file with the extension .sparkpkg that
+ * contains a manifest named spark-native-datasource.json and the library built for one or more
+ * platforms, added to a session with spark.addArtifact or under the paths in
+ * spark.sql.dataSource.native.paths. See the "Native Data Sources" page of the Spark SQL guide for
+ * details and examples.
+ *
+ * Lifecycle
+ * ---------
+ * Spark loads a library once in each JVM, and calls spark_ds_init before any other function. It
+ * creates a data source with create on the driver for each schema inference, scan and write, and
+ * closes it right after. create also returns the capabilities of the data source: Spark reports
+ * the operations that a data source does not declare as unsupported, without calling the library.
+ * Spark also creates a data source while it analyzes a query, to reject a batch write that the
+ * data source does not support before it computes the data to write. That data source has the
+ * options of the table, those of CREATE TABLE for a table in a catalog, which the options of a
+ * write can override, so Spark checks the capabilities of each operation again.
+ *
+ * - Batch reads. On the driver, Spark creates a reader with create_reader, pushes operations down
+ * by calling, each at most once and in this order, push_predicates, push_limit and
+ * prune_columns, plans the partitions with partitions, serializes the state of the reader with
+ * serialize_reader, and closes it. Then it calls read on the executors, in one task for each
+ * partition. If Spark never plans the partitions, for example to only explain the plan of the
+ * query, it closes the reader once the reader is no longer referenced.
+ * - Streaming reads. On the driver, Spark creates a stream reader with create_stream_reader and
+ * keeps it until the streaming query stops. For each micro-batch, it asks for the latest_offset,
+ * plans the partitions with stream_partitions and serializes the state with
+ * serialize_stream_reader, then calls read on the executors. Offsets are defined by the library,
+ * as compact JSON on a single line: Spark stores each offset as a line of the offset log of the
+ * checkpoint, so it rejects an offset that is not valid JSON or contains a line break. When a
+ * streaming query restarts, from its checkpoint, Spark creates a new stream reader and may call
+ * stream_partitions again for the last range that it planned before, so the partitions of a
+ * range must read the same data each time.
+ * - Writes. On the driver, right before it starts the tasks of the write (or of a micro-batch),
+ * Spark creates a writer with create_writer, and serializes its state with serialize_writer.
+ * With adaptive query execution, Spark first runs the stages of the data to write that end with
+ * a shuffle, and creates no writer if they fail. The rest of the data, such as its last stage,
+ * is computed in the tasks of the write, so a writer may still be aborted because computing the
+ * data failed: do irreversible work, such as removing the existing data for an overwrite, in
+ * commit. On the executors, each task creates a data writer with create_data_writer, calls write
+ * for each batch of at most spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends
+ * once it holds spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls commit_data_writer,
+ * or abort_data_writer if the task or the commit fails, before it closes the data writer. Then
+ * the driver calls commit with the messages of the tasks if they all succeeded, or abort
+ * otherwise, including when commit fails, and closes the writer. Several attempts of a task may
+ * write the same partition, after a failure or with speculative execution, each with a data
+ * writer of its own. Spark lets the attempts of a partition commit one at a time, and passes to
+ * commit the message of one committed attempt for each partition, so the write consists of the
+ * data of these attempts only. It aborts the data writers of the other attempts, except on an
+ * executor that is lost, where a data writer is neither committed nor aborted, and may even have
+ * committed before its message reached the driver. A streaming write commits or aborts each
+ * micro-batch (epoch). Spark may commit the same epoch more than once, for example when a
+ * streaming query restarts after a failure, so committing an epoch that is already committed
+ * must have no effect, for exactly-once output. The ID of the query, which create_writer gets for
+ * a streaming write, stays the same when the query restarts from its checkpoint, so the ID of the
+ * query and the epoch ID identify the output of a micro-batch, also when other queries write to
+ * the same destination. A micro-batch streaming query creates a new stream writer for each
+ * micro-batch, and Spark closes the previous stream writer of the query when the query creates a
+ * new one. Spark is not told when a streaming query stops, so it closes the last stream writer of
+ * the query once the writer is no longer referenced, or when the query is restarted.
+ *
+ * Conventions
+ * -----------
+ * - Errors. The functions that return an int return 0 on success, or any other value on failure.
+ * They then set the message that last_error returns, and Spark reports it as a
+ * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail, and they may run after the query,
+ * from another thread: do the work that can fail a query, such as flushing data, in commit. A
+ * C++ exception or a Rust panic must not cross the interface: catch it and return an error
+ * instead.
+ * - Handles. create, create_reader, create_stream_reader, create_writer and create_data_writer
+ * return an opaque, non-NULL handle in their last argument, typically a pointer: Spark reports a
+ * NULL handle as an error. Spark releases each handle exactly once, with the close function of
+ * its kind: close, close_reader, close_stream_reader, close_writer or close_data_writer, and no
+ * other function releases a handle. A handle is not always released right after its last use: a
+ * batch reader whose partitions are never planned and the last stream writer of a query are
+ * released once they are no longer referenced (see Lifecycle above). A reader or writer handle
+ * must not depend on the data source handle it was created from, because Spark closes the data
+ * source handle right after creating it.
+ * - Memory. What Spark passes to a function, such as strings, byte arrays, SparkDsCreateInfo and
+ * SparkDsWriteInfo, is only valid during the call, except for SparkDsHost and
+ * SparkDsTaskContext, which document how long they live. Strings are UTF-8 and NUL-terminated.
+ * The JSON of the predicates is ASCII: it escapes the other characters, which any JSON parser
+ * decodes. The library returns strings and byte arrays by passing them to a SparkDsSink, which
+ * copies them.
+ * - State. Partitions, the state of readers and writers, and commit messages are byte arrays in a
+ * format defined by the library. Spark sends them to the executors and back.
+ * - Arrow data. Schemas and data are exchanged with the Arrow C data interface
+ * (https://arrow.apache.org/docs/format/CDataInterface.html), as structs allocated by Spark.
+ * For an output struct, the library initializes it and Spark takes ownership. For an input
+ * struct, the library may move it to take ownership; otherwise, Spark releases it after the
+ * call returns, so the library must not use it afterwards. The arrays that the library exports
+ * must have an offset of 0, at every level: Arrow Java ignores the offsets of the arrays it
+ * imports, so Spark rejects an array with another offset, such as a zero-copy slice. Copy a
+ * sliced array before exporting it. Spark also rejects an array with fewer values than its
+ * parent refers to, such as a column shorter than its batch. The buffers of an array must hold
+ * all of its values, as the C data interface requires: it does not carry the sizes of the
+ * buffers, so Spark cannot check them, and reads past the end of a buffer that is too short.
+ * - Schemas. Spark types map to Arrow types as in spark.sql.execution.arrow: for example, bigint
+ * is int64, string is utf8, timestamp is timestamp[us] with a time zone and timestamp_ntz is
+ * timestamp[us] without one. Data read from a library may also use large_utf8 and utf8_view for
+ * strings, large_binary and binary_view for binary data, and any names for the fields of the
+ * entries of a map. Spark rejects the other Arrow types, such as unsigned integers, decimal256,
+ * large lists, fixed-size lists and list views, and dictionary encoding.
+ * - Options. Option keys are passed in lower case, because Spark matches them
+ * case-insensitively.
+ * - Optional functions. A library only sets the functions of the operations it supports in its
+ * table, and leaves the others NULL. A data source must not declare a capability whose
+ * functions are NULL: Spark then fails the operation. The pushdown functions and the serialize
+ * functions are optional: without them, Spark pushes nothing down, and the reader or writer has
+ * no state. So are schema, without which the user must specify the schema, commit_offset and
+ * last_error.
+ * - Threads. Spark calls the library from multiple threads, and calls on different handles may
+ * run concurrently. Calls on the same handle never run concurrently. A handle may be released
+ * from a thread other than the threads that used it, such as a JVM Cleaner thread when Spark
+ * releases a handle once it is no longer referenced, for example the last stream writer of a
+ * streaming query. A SparkDsSink must be called on the thread of the call that it was passed
+ * to, during that call. The functions of SparkDsHost and SparkDsTaskContext, and the callbacks
+ * of the stream that read initializes, can be called from any thread. The callbacks of the
+ * stream never run concurrently.
+ *
+ * Compatibility
+ * -------------
+ * The function table and the structs that Spark passes to the library start with their size, so
+ * that they can grow. Within a SPARK_DS_ABI_VERSION:
+ * - fields are only appended to the structs, including SparkDsLibrary, and the existing fields
+ * never change, nor do the semantics of the existing functions;
+ * - Spark only reads the functions that fit in the size of the table of a library, so a library
+ * built against an earlier header keeps working, with the functions that it does not know as
+ * NULL. A library must only read the fields that fit in the size of a struct that Spark passes;
+ * - new values may appear where this documentation says so, such as the names of pushed
+ * predicates. Spark only uses a new capability or write mode with a data source that declares
+ * it.
+ * An incompatible change uses a new SPARK_DS_ABI_VERSION: spark_ds_init gets the newest version
+ * that Spark supports, the library returns a table of a version that it implements, and Spark
+ * rejects a table of a version that it does not support. Spark supports version 1 at least in all
+ * the releases of Spark 4 from 4.4.0.
+ */
+#ifndef SPARK_DATASOURCE_H
+#define SPARK_DATASOURCE_H
+
+#include <stdbool.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/* The version of this interface. */
+#define SPARK_DS_ABI_VERSION 1
+
+/* Exports spark_ds_init from the library, also when the library hides its other symbols. */
+#if defined(_WIN32)
+#define SPARK_DS_EXPORT __declspec(dllexport)
+#else
+#define SPARK_DS_EXPORT __attribute__((visibility("default")))
+#endif
+
+/* The Arrow C data interface: https://arrow.apache.org/docs/format/CDataInterface.html */
+#ifndef ARROW_C_DATA_INTERFACE
+#define ARROW_C_DATA_INTERFACE
+
+#define ARROW_FLAG_DICTIONARY_ORDERED 1
+#define ARROW_FLAG_NULLABLE 2
+#define ARROW_FLAG_MAP_KEYS_SORTED 4
+
+struct ArrowSchema {
+ const char* format;
+ const char* name;
+ const char* metadata;
+ int64_t flags;
+ int64_t n_children;
+ struct ArrowSchema** children;
+ struct ArrowSchema* dictionary;
+ void (*release)(struct ArrowSchema*);
+ void* private_data;
+};
+
+struct ArrowArray {
+ int64_t length;
+ int64_t null_count;
+ int64_t offset;
+ int64_t n_buffers;
+ int64_t n_children;
+ const void** buffers;
+ struct ArrowArray** children;
+ struct ArrowArray* dictionary;
+ void (*release)(struct ArrowArray*);
+ void* private_data;
+};
+
+#endif /* ARROW_C_DATA_INTERFACE */
+
+#ifndef ARROW_C_STREAM_INTERFACE
+#define ARROW_C_STREAM_INTERFACE
+
+struct ArrowArrayStream {
+ int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out);
+ int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out);
+ const char* (*get_last_error)(struct ArrowArrayStream*);
+ void (*release)(struct ArrowArrayStream*);
+ void* private_data;
+};
+
+#endif /* ARROW_C_STREAM_INTERFACE */
+
+/*
+ * Receives a string or a byte array that the library returns, and copies it. `sink` is the
+ * argument passed with the sink to the function that returns it.
+ */
+typedef void (*SparkDsSink)(void* sink, const uint8_t* data, int64_t length);
+
+/* ---------------------------------------------------------------------------------------------
+ * What Spark passes to the library
+ * --------------------------------------------------------------------------------------------- */
+
+/* The levels of the messages that a library logs with SparkDsHost.log. */
+#define SPARK_DS_LOG_ERROR 1
+#define SPARK_DS_LOG_WARN 2
+#define SPARK_DS_LOG_INFO 3
+#define SPARK_DS_LOG_DEBUG 4
+#define SPARK_DS_LOG_TRACE 5
+
+/* Spark, as spark_ds_init gets it. It stays valid until the JVM exits. */
+typedef struct SparkDsHost {
+ /*
+ * The size of this struct in the header that Spark was built with. A later version of Spark may
+ * append fields, so a library must only read the fields that fit in it.
+ */
+ int64_t size;
+ /* The newest version of this interface that Spark supports. */
+ int32_t abi_version;
+ /* The version of Spark, such as "4.4.0". */
+ const char* spark_version;
+ /*
+ * Logs a message, a NUL-terminated UTF-8 string, at one of the SPARK_DS_LOG_* levels: in the log
+ * of the driver or of the executor, with the path of the library.
+ */
+ void (*log)(int32_t level, const char* message);
+} SparkDsHost;
+
+/* A data source, as create gets it. */
+typedef struct SparkDsCreateInfo {
+ /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */
+ int64_t size;
+ /*
+ * The name of the data source, as listed in the package manifest, or in lower case for an
+ * installed library. A library can implement several data sources.
+ */
+ const char* name;
+ /*
+ * The num_options options given by the user, such as spark.read.option("key", "value"), with
+ * the keys in lower case.
+ */
+ const char* const* option_keys;
+ const char* const* option_values;
+ int64_t num_options;
+} SparkDsCreateInfo;
+
+/*
+ * The task that calls read or create_data_writer, on an executor. It stays valid until the stream
+ * that read returns is released, or until the data writer is closed.
+ */
+typedef struct SparkDsTaskContext {
+ /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */
+ int64_t size;
+ /* The ID of the stage of the task, and the attempt number of the stage. */
+ int32_t stage_id;
+ int32_t stage_attempt_number;
+ /* The partition that the task reads or writes, and the attempt number of the task. */
+ int32_t partition_id;
+ int32_t attempt_number;
+ /* The ID of the task attempt, unique in the Spark application. */
+ int64_t task_attempt_id;
+ /*
+ * Returns whether the task was killed, for example because its query was canceled. Spark cannot
+ * interrupt native code, so a read or a write that runs for long, or blocks, should check it
+ * regularly, and fail once it is true.
+ */
+ bool (*is_interrupted)(const struct SparkDsTaskContext* task);
+} SparkDsTaskContext;
+
+/* The modes of a write. */
+#define SPARK_DS_WRITE_APPEND 0
+#define SPARK_DS_WRITE_TRUNCATE 1
+
+/* A write, as create_writer gets it. */
+typedef struct SparkDsWriteInfo {
+ /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */
+ int64_t size;
+ /*
+ * SPARK_DS_WRITE_APPEND to add the data to the existing data, or SPARK_DS_WRITE_TRUNCATE to
+ * replace the existing data: with the save mode "overwrite" of a batch write, and in each
+ * micro-batch with the output mode "complete" of a streaming write. Spark only passes
+ * SPARK_DS_WRITE_TRUNCATE to a data source with the capability SPARK_DS_TRUNCATE.
+ */
+ int32_t mode;
+ /*
+ * The ID of the streaming query of a streaming write, or NULL for a batch write. It stays the
+ * same when the query restarts from its checkpoint, so with the epoch IDs, it identifies the
+ * commits of the query: see Writes above.
+ */
+ const char* query_id;
+} SparkDsWriteInfo;
+
+/* ---------------------------------------------------------------------------------------------
+ * What the library returns to Spark
+ * --------------------------------------------------------------------------------------------- */
+
+/* The capabilities of a data source, which create returns: the operations that it supports. */
+#define SPARK_DS_BATCH_READ 1 /* batch reads, with create_reader */
+#define SPARK_DS_MICRO_BATCH_READ 2 /* micro-batch streaming reads, with create_stream_reader */
+#define SPARK_DS_BATCH_WRITE 4 /* batch writes, with create_writer */
+#define SPARK_DS_STREAMING_WRITE 8 /* streaming writes, with create_writer */
+#define SPARK_DS_TRUNCATE 16 /* writes with the mode SPARK_DS_WRITE_TRUNCATE */
+
+/* How a reader uses a predicate pushed with push_predicates. */
+#define SPARK_DS_PREDICATE_UNUSED 0 /* not at all */
+/* only to skip data, so it may return rows for which the predicate is not true */
+#define SPARK_DS_PREDICATE_PARTIAL 1
+/* to only return the rows for which the predicate is true */
+#define SPARK_DS_PREDICATE_COMPLETE 2
+
+/*
+ * The functions of a library, which spark_ds_init returns. A library sets the functions of the
+ * operations it supports, and leaves the others NULL: see "Optional functions" above.
+ */
+typedef struct SparkDsLibrary {
+ /*
+ * The size of this struct in the header that the library was built with. Spark only reads the
+ * functions that fit in it.
+ */
+ int64_t size;
+ /* The version of this interface that the table implements, SPARK_DS_ABI_VERSION. */
+ int32_t abi_version;
+
+ /*
+ * Returns the message of the last error on the calling thread, as a NUL-terminated UTF-8 string
+ * that stays valid until the next call on that thread. Optional: without it, Spark only reports
+ * the value that the function returned.
+ */
+ const char* (*last_error)(void);
+
+ /* -------------------------------------------------------------------------------------------
+ * Data source (driver)
+ * ------------------------------------------------------------------------------------------- */
+
+ /*
+ * Creates a data source.
+ *
+ * info: the name of the data source, and the options given by the user
+ * capabilities: set to the capabilities of the data source, SPARK_DS_* flags combined with |,
+ * which may depend on its name and its options. It is 0 when create is called.
+ * source: set to the data source handle
+ */
+ int (*create)(const SparkDsCreateInfo* info, uint64_t* capabilities, void** source);
+
+ /*
+ * Exports the schema of the data. Spark calls it only when the user does not specify a schema.
+ *
+ * out: initialized with a struct type whose fields are the columns
+ */
+ int (*schema)(void* source, struct ArrowSchema* out);
+
+ /*
+ * Creates a reader for a batch scan.
+ *
+ * schema: the schema to read: the one the user specified, or the one returned by schema
+ * reader: set to the reader handle
+ */
+ int (*create_reader)(void* source, struct ArrowSchema* schema, void** reader);
+
+ /*
+ * Creates a reader for a micro-batch streaming scan.
+ *
+ * schema: the schema to read: the one the user specified, or the one returned by schema
+ * reader: set to the stream reader handle
+ */
+ int (*create_stream_reader)(void* source, struct ArrowSchema* schema, void** reader);
+
+ /*
+ * Creates a writer for a batch write, or for a micro-batch of a streaming write.
+ *
+ * schema: the schema of the data to write
+ * info: the mode of the write, and the ID of the query of a streaming write
+ * writer: set to the writer handle
+ */
+ int (*create_writer)(void* source, struct ArrowSchema* schema, const SparkDsWriteInfo* info,
+ void** writer);
+
+ /* Releases a data source handle. */
+ void (*close)(void* source);
+
+ /* -------------------------------------------------------------------------------------------
+ * Batch reader (driver)
+ * ------------------------------------------------------------------------------------------- */
+
+ /*
+ * Pushes predicates down to the reader. The predicates are combined with AND.
+ *
+ * Each predicate is a JSON expression tree. An expression is one of:
+ * - {"type": "column", "name": ["a", "b"]}: a column. A nested field has a name part for each
+ * level.
+ * - {"type": "literal", "dataType": "bigint", "value": 1}: a literal. The data type is a Spark
+ * SQL type name such as int, string or decimal(10,2). The value is null, a JSON boolean or
+ * number for boolean and numeric types, a string for string, a base64 string for binary, a
+ * decimal string for decimal, the days since the epoch for date, and the microseconds since
+ * the epoch for timestamp and timestamp_ntz.
+ * - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in
+ * org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=>
+ * (equal, with null equal to null), <, <=, > and >= with two children, IN with the value to
+ * look for and then the values to look in, IS_NULL and IS_NOT_NULL with one child,
+ * STARTS_WITH, ENDS_WITH and CONTAINS with the string and the string to look for, AND and OR
+ * with two predicates, and NOT with one. The children of the other functions are columns,
+ * literals or predicates.
+ * Spark does not pass predicates that it cannot express this way, such as predicates that use
+ * arithmetic or other functions, nor predicates on columns that contain strings with a
+ * non-binary collation, such as UTF8_LCASE: strings are compared by their UTF-8 bytes. A later
+ * version of Spark may pass predicates with other names without changing SPARK_DS_ABI_VERSION,
+ * so a reader must ignore a predicate that contains a name that it does not know, or a data
+ * type that it does not support.
+ *
+ * Spark evaluates all the predicates again after reading, so a reader may use them only to skip
+ * data, for example the files whose statistics show that no row matches. But it must never drop
+ * or skip a row for which the predicates are true under the semantics of Spark SQL, which Spark
+ * cannot detect. In particular:
+ * - For floating-point columns, NaN is equal to NaN and greater than any other value, including
+ * positive infinity, and -0.0 is equal to 0.0. IEEE 754 comparisons treat NaN differently, and
+ * total orders, such as f64::total_cmp in Rust or the comparison kernels of arrow-rs, also
+ * order -0.0 before 0.0. For example, d > 1.0 is true for a NaN value of d, and d = 0.0 for
+ * -0.0.
+ * - Comparisons with null are null, and a row only matches if the predicates are true. For
+ * example, NOT (i IN (1, NULL)) matches no row.
+ * - Strings are compared by their UTF-8 bytes, which order the characters outside the Basic
+ * Multilingual Plane after all the others, unlike comparisons of UTF-16 strings.
+ * - To use only a part of a predicate, a reader may take the parts that it does not support as
+ * true, but not under a NOT, where that would drop rows.
Review Comment:
Thanks, fixed in 74aa0dffc534. The rule now says that a reader may only take a part as true when every function between that part and the root of the predicate is AND or OR, with your example, and `predicate-cases.json` has `(id > 3) = (id < 3)`, which is true for the row whose id is 3.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala:
##########
@@ -0,0 +1,279 @@
+/*
+ * 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.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_<name>`, such as
+ * `libspark_datasource_my_source.so` on Linux, implements the data source `<name>`. 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 == NativeLibrary.ABI_VERSION) Right(pkg) else Left(pkg)
Review Comment:
Thanks, settled in 74aa0dffc534 with your second option: `abiVersion` is the oldest version that the library implements. Spark uses a package whose `abiVersion` is at most the newest version that it supports, and `spark_ds_init` returns a table of a version that both support. So a library that implements 1 and 2 declares 1, and a later Spark that supports only 2 and 3 still loads it, and gets a table of version 2. The docs promise that Spark ignores the fields of the manifest that it does not know, so a later release can add one, for example for the newest version. The manifest must also have a positive `abiVersion`, and names without U+0000.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6166143391
@LuciferYang Thanks for checking the bindings against the header and the examples. On testing, good point: in e5352e191c14, the `Java 25 build with Maven` job of a pull request also runs `NativeDataSourceSuite`, `NativeDataSourceClusterSuite` and `NativeDataSourceConnectSuite` after the build, with the C++ compiler of the runner. On this commit, they ran on Linux with Java 25: 55 tests passed in `NativeDataSourceSuite` and `NativeDataSourceClusterSuite`, with only the Java version test canceled, as expected, and 2 in `NativeDataSourceConnectSuite`. So the FFM calls now run on every pull request. I also addressed the inline comment, with a test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "sarutak (via GitHub)" <gi...@apache.org>.
sarutak commented on code in PR #59424:
URL: https://github.com/apache/spark/pull/59424#discussion_r4293982862
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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
+ * `push_predicates` of `spark_datasource.h`.
+ */
+object NativePredicates {
+ // Writes ASCII only, escaping the other characters, which any JSON parser decodes.
+ private val mapper = JsonMapper.builder().enable(JsonWriteFeature.ESCAPE_NON_ASCII).build()
+
+ /**
+ * The names of the predicates that are passed to the libraries, as listed in
+ * `push_predicates`. They are part of the interface, so a predicate with any other
+ * name, such as an arithmetic function, is not passed: its semantics, for example on overflow,
+ * would be part of the interface too, and the names of Data Source V2 grow with new releases.
+ */
+ 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))
Review Comment:
Two literal encodings let a library that follows the header skip rows that match:
- A float literal is written as the shortest decimal of the float: `col("f") == lit(0.1).cast("float")` pushes `{"type":"literal","dataType":"float","value":0.1}`. A reader that parses JSON numbers as doubles, as most JSON parsers do, and widens the float column to compare it, finds `0.10000000149011612 = 0.1` false and skips the matching row. Could the header say that a literal must be read at the precision of its `dataType`, binary32 for float, and that the values of the integer types, dates and timestamps are exact 64-bit integers, which a double does not hold above 2^53? `predicate-cases.json` has no float literal, so it does not catch this.
- `s.toString` replaces invalid UTF-8 with U+FFFD: `col("s") == expr("cast(X'FF' as string)")` pushes `"value":"\uFFFD"`. Spark compares the bytes, so a reader that compares `EF BF BD` with its statistics can skip a file whose strings contain the byte `FF`. Could a string literal that is not valid UTF-8 (`!s.isValid`) make the predicate not pushed, as for the other values that cannot be expressed?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,980 @@
+/*
+ * 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
+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, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, 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 the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @param location where the library is on the driver. The readers and the writers distribute it
+ * to the executors of the active session when they are created.
+ * @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.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, 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", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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) =>
Review Comment:
`checkArrays` only checks the children that an array has, so a batch with `n_children = 0` skips the checks of all its columns. Arrow Java does not catch it either: `ArrayImporter.doImport` only compares the number of children with the vector when there is at least one child. I tried it with a library that returns a batch of 4 rows with `n_children = 0`: the scan returns 4 rows in which every column is null, without an error. These are the same silent nulls as the ones that the length check of a8f178b prevents.
Could the checks of the batch and of each column also require `n_children` to be the number of fields of the type, and the child pointers to be non-NULL, before `children` reads them? With an `n_children` larger than the array of pointers, `children` would also read past that array before Arrow gets to see it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,980 @@
+/*
+ * 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
+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, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, 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 the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @param location where the library is on the driver. The readers and the writers distribute it
+ * to the executors of the active session when they are created.
+ * @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.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, 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", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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)
Review Comment:
The nullability of the columns comes from `schema` of the library (`NativeDataSource.schema`), and this comparison ignores nullability, so nothing checks that the data agrees with the schema. I tried a library whose `schema` returns `id` as not nullable, while its stream declares `id` nullable and has nulls in 2 of 4 rows:
- `collect()` fails with `java.lang.IllegalStateException: Value at index is null` from `BigIntVector.get`, which does not say what is wrong.
- `filter(isnull(col("id"))).count()` returns 0 instead of 2, without an error, since the optimizer folds `isnull` of a non-nullable column to false.
- With `-Darrow.enable_null_check_for_get=false`, which @FelixYBW suggests on #59313 for performance, `collect()` no longer fails: it returns the values in the data buffer for the null rows (101 and 103 with my library), without an error.
Could Spark fail with `NATIVE_DATA_SOURCE_ERROR` when a column that the schema to read declares not nullable has a `null_count` other than 0 (it is in the array, so the check is cheap; a `null_count` of -1 needs a pass over the validity buffer)? Or make the inferred schema nullable. Either way, the header could say that the data must agree with the nullability of the schema.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala:
##########
@@ -0,0 +1,279 @@
+/*
+ * 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.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_<name>`, such as
+ * `libspark_datasource_my_source.so` on Linux, implements the data source `<name>`. 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 == NativeLibrary.ABI_VERSION) Right(pkg) else Left(pkg)
Review Comment:
A package declares a single `abiVersion`, and Spark only uses a package whose `abiVersion` is the version of Spark, while the header lets a library implement several versions and choose one in `spark_ds_init` (`spark_datasource.h` L676). Once there is a version 2, a library that implements both cannot say so in its manifest. With `"abiVersion": 2`, Spark 4.4 reports `UNSUPPORTED_ABI_VERSION` without calling `spark_ds_init`, although the library would return a version 1 table. With `1`, the field no longer means "the version that the library implements", as the docs say. Installed libraries have no manifest, so only packages have this problem.
The manifest of 4.4.0 is what all later packages have to work with, so could it be settled now? For example, the manifest could list the versions that the library implements, such as `"abiVersions": [1]`, and Spark would use a package if it supports one of them. Or the docs could define `abiVersion` as the oldest version that the library implements, so that a later release can add a field for the others: `parseManifest` already ignores fields that it does not know, which the docs could also promise.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -0,0 +1,688 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*
+ * spark_datasource.h: the binary interface between Spark and data sources implemented in native
+ * code, such as Rust or C++. Since Spark 4.4.0.
+ *
+ * A native data source is a shared library that exports a single C function, spark_ds_init,
+ * which returns the table of the functions that implement the data source, a SparkDsLibrary.
+ * Spark calls them with the Foreign Function and Memory API of Java, so native data sources
+ * require Java 22 or later. Each library has its own table, so any number of libraries can be
+ * loaded side by side. Spark adapts a native data source to Data Source V2, with batch and
+ * micro-batch streaming reads, batch and streaming writes, and predicate, column and limit
+ * pushdown.
+ *
+ * This header is self-contained: it only depends on the C standard library, and defines the
+ * structs of the Arrow C data interface unless they are already defined, as by Arrow C++. A library
+ * can copy it into its source tree, or use the copy in the include directory of a Spark
+ * distribution. A library built against it works with all the releases of Spark that support its
+ * SPARK_DS_ABI_VERSION.
+ *
+ * Spark finds the library of a data source automatically when it is installed like other native
+ * libraries and named after the data source: the library of my_source is
+ * spark_datasource_my_source, such as libspark_datasource_my_source.so in /usr/local/lib on Linux.
+ * Spark looks for it in the directories of java.library.path, which include LD_LIBRARY_PATH on
+ * Linux, and in the lib directory of each installation prefix in PATH. A library can also be
+ * distributed in a native data source package: a zip file with the extension .sparkpkg that
+ * contains a manifest named spark-native-datasource.json and the library built for one or more
+ * platforms, added to a session with spark.addArtifact or under the paths in
+ * spark.sql.dataSource.native.paths. See the "Native Data Sources" page of the Spark SQL guide for
+ * details and examples.
+ *
+ * Lifecycle
+ * ---------
+ * Spark loads a library once in each JVM, and calls spark_ds_init before any other function. It
+ * creates a data source with create on the driver for each schema inference, scan and write, and
+ * closes it right after. create also returns the capabilities of the data source: Spark reports
+ * the operations that a data source does not declare as unsupported, without calling the library.
+ * Spark also creates a data source while it analyzes a query, to reject a batch write that the
+ * data source does not support before it computes the data to write. That data source has the
+ * options of the table, those of CREATE TABLE for a table in a catalog, which the options of a
+ * write can override, so Spark checks the capabilities of each operation again.
+ *
+ * - Batch reads. On the driver, Spark creates a reader with create_reader, pushes operations down
+ * by calling, each at most once and in this order, push_predicates, push_limit and
+ * prune_columns, plans the partitions with partitions, serializes the state of the reader with
+ * serialize_reader, and closes it. Then it calls read on the executors, in one task for each
+ * partition. If Spark never plans the partitions, for example to only explain the plan of the
+ * query, it closes the reader once the reader is no longer referenced.
+ * - Streaming reads. On the driver, Spark creates a stream reader with create_stream_reader and
+ * keeps it until the streaming query stops. For each micro-batch, it asks for the latest_offset,
+ * plans the partitions with stream_partitions and serializes the state with
+ * serialize_stream_reader, then calls read on the executors. Offsets are defined by the library,
+ * as compact JSON on a single line: Spark stores each offset as a line of the offset log of the
+ * checkpoint, so it rejects an offset that is not valid JSON or contains a line break. When a
+ * streaming query restarts, from its checkpoint, Spark creates a new stream reader and may call
+ * stream_partitions again for the last range that it planned before, so the partitions of a
+ * range must read the same data each time.
+ * - Writes. On the driver, right before it starts the tasks of the write (or of a micro-batch),
+ * Spark creates a writer with create_writer, and serializes its state with serialize_writer.
+ * With adaptive query execution, Spark first runs the stages of the data to write that end with
+ * a shuffle, and creates no writer if they fail. The rest of the data, such as its last stage,
+ * is computed in the tasks of the write, so a writer may still be aborted because computing the
+ * data failed: do irreversible work, such as removing the existing data for an overwrite, in
+ * commit. On the executors, each task creates a data writer with create_data_writer, calls write
+ * for each batch of at most spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends
+ * once it holds spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls commit_data_writer,
+ * or abort_data_writer if the task or the commit fails, before it closes the data writer. Then
+ * the driver calls commit with the messages of the tasks if they all succeeded, or abort
+ * otherwise, including when commit fails, and closes the writer. Several attempts of a task may
+ * write the same partition, after a failure or with speculative execution, each with a data
+ * writer of its own. Spark lets the attempts of a partition commit one at a time, and passes to
+ * commit the message of one committed attempt for each partition, so the write consists of the
+ * data of these attempts only. It aborts the data writers of the other attempts, except on an
+ * executor that is lost, where a data writer is neither committed nor aborted, and may even have
+ * committed before its message reached the driver. A streaming write commits or aborts each
+ * micro-batch (epoch). Spark may commit the same epoch more than once, for example when a
+ * streaming query restarts after a failure, so committing an epoch that is already committed
+ * must have no effect, for exactly-once output. The ID of the query, which create_writer gets for
+ * a streaming write, stays the same when the query restarts from its checkpoint, so the ID of the
+ * query and the epoch ID identify the output of a micro-batch, also when other queries write to
+ * the same destination. A micro-batch streaming query creates a new stream writer for each
+ * micro-batch, and Spark closes the previous stream writer of the query when the query creates a
+ * new one. Spark is not told when a streaming query stops, so it closes the last stream writer of
+ * the query once the writer is no longer referenced, or when the query is restarted.
+ *
+ * Conventions
+ * -----------
+ * - Errors. The functions that return an int return 0 on success, or any other value on failure.
+ * They then set the message that last_error returns, and Spark reports it as a
+ * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail, and they may run after the query,
+ * from another thread: do the work that can fail a query, such as flushing data, in commit. A
+ * C++ exception or a Rust panic must not cross the interface: catch it and return an error
+ * instead.
+ * - Handles. create, create_reader, create_stream_reader, create_writer and create_data_writer
+ * return an opaque, non-NULL handle in their last argument, typically a pointer: Spark reports a
+ * NULL handle as an error. Spark releases each handle exactly once, with the close function of
+ * its kind: close, close_reader, close_stream_reader, close_writer or close_data_writer, and no
+ * other function releases a handle. A handle is not always released right after its last use: a
+ * batch reader whose partitions are never planned and the last stream writer of a query are
+ * released once they are no longer referenced (see Lifecycle above). A reader or writer handle
+ * must not depend on the data source handle it was created from, because Spark closes the data
+ * source handle right after creating it.
+ * - Memory. What Spark passes to a function, such as strings, byte arrays, SparkDsCreateInfo and
+ * SparkDsWriteInfo, is only valid during the call, except for SparkDsHost and
+ * SparkDsTaskContext, which document how long they live. Strings are UTF-8 and NUL-terminated.
+ * The JSON of the predicates is ASCII: it escapes the other characters, which any JSON parser
+ * decodes. The library returns strings and byte arrays by passing them to a SparkDsSink, which
+ * copies them.
+ * - State. Partitions, the state of readers and writers, and commit messages are byte arrays in a
+ * format defined by the library. Spark sends them to the executors and back.
+ * - Arrow data. Schemas and data are exchanged with the Arrow C data interface
+ * (https://arrow.apache.org/docs/format/CDataInterface.html), as structs allocated by Spark.
+ * For an output struct, the library initializes it and Spark takes ownership. For an input
+ * struct, the library may move it to take ownership; otherwise, Spark releases it after the
+ * call returns, so the library must not use it afterwards. The arrays that the library exports
+ * must have an offset of 0, at every level: Arrow Java ignores the offsets of the arrays it
+ * imports, so Spark rejects an array with another offset, such as a zero-copy slice. Copy a
+ * sliced array before exporting it. Spark also rejects an array with fewer values than its
+ * parent refers to, such as a column shorter than its batch. The buffers of an array must hold
+ * all of its values, as the C data interface requires: it does not carry the sizes of the
+ * buffers, so Spark cannot check them, and reads past the end of a buffer that is too short.
+ * - Schemas. Spark types map to Arrow types as in spark.sql.execution.arrow: for example, bigint
Review Comment:
The Arrow types are defined "as in spark.sql.execution.arrow", which is internal and keeps growing, and `NativeArrow.exportSchema` has no allowlist, so types that the header does not describe reach the library. With a library that logs the schema that `create_writer` gets:
- a variant column arrives as a struct `+s` of two `z` fields, `value` and `metadata`, with the metadata `variant=true` on `metadata`;
- a `time` column arrives as `ttn`, with the metadata `SPARK::time::precision=6`;
- a day-time interval arrives as `tDu`.
The header does not tell a library what these mean, and once 4.4.0 is released, these encodings are part of version 1, also when ArrowUtils changes them for Python. Could the header list the Spark types of version 1 with their Arrow types, and could Spark reject the other types on the driver with `NATIVE_DATA_SOURCE_ERROR` before it passes a schema to the library, in `create_reader` and `create_writer`? A type could then be added later, with a capability.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -0,0 +1,982 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.v2.ffi;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+import java.lang.reflect.Array;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.BooleanSupplier;
+
+import org.apache.spark.internal.LogKeys;
+import org.apache.spark.internal.MDC;
+import org.apache.spark.internal.SparkLogger;
+import org.apache.spark.internal.SparkLoggerFactory;
+import org.apache.spark.unsafe.Platform;
+
+/**
+ * A native data source library, which implements the C interface in
+ * {@code org/apache/spark/sql/datasource/spark_datasource.h}, called through the Foreign Function
+ * and Memory (FFM) API of Java 22 and later.
+ * <p>
+ * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached
+ * through method handles, looked up once when the first library is loaded. The library exports
+ * {@code spark_ds_init}, which returns its table of functions, {@code SparkDsLibrary}. Each
+ * function of the table is bound to a method handle that takes and returns primitive values only,
+ * with pointers as {@code long} addresses, so the calls themselves use no reflection. The memory
+ * that Spark passes to the library is allocated with {@link Platform}, at the offsets of the
+ * structs of the header, which {@link #layout()} lists for the tests.
+ * <p>
+ * Calling a function that the library leaves NULL in its table throws
+ * {@link MissingFunctionException}. A function that fails throws a {@link RuntimeException} with
+ * the message of {@code last_error}.
+ */
+final class NativeLibrary {
+
+ private static final SparkLogger LOGGER = SparkLoggerFactory.getLogger(NativeLibrary.class);
+
+ /** The version of the interface that Spark supports, {@code SPARK_DS_ABI_VERSION}. */
+ static final int ABI_VERSION = 1;
+
+ /** The first version of Java with the final FFM API. */
+ static final int MINIMUM_JAVA_VERSION = 22;
+
+ /** Whether this JVM can load native data source libraries. */
+ static boolean isSupported() {
+ return Runtime.version().feature() >= MINIMUM_JAVA_VERSION;
+ }
+
+ // The capabilities of a data source, SPARK_DS_BATCH_READ and the others.
+ static final long BATCH_READ = 1;
+ static final long MICRO_BATCH_READ = 2;
+ static final long BATCH_WRITE = 4;
+ static final long STREAMING_WRITE = 8;
+ static final long TRUNCATE = 16;
+
+ // The modes of a write, SPARK_DS_WRITE_APPEND and SPARK_DS_WRITE_TRUNCATE.
+ static final int WRITE_APPEND = 0;
+ static final int WRITE_TRUNCATE = 1;
+
+ /** A predicate that the reader does not use, SPARK_DS_PREDICATE_UNUSED. */
+ static final int PREDICATE_UNUSED = 0;
+
+ // The C types of the parameters and the results of the functions.
+ private enum C { VOID, BOOL, I32, I64, PTR }
+
+ /** The functions of the table {@code SparkDsLibrary}, in their order in the table. */
+ enum Function {
+ LAST_ERROR("last_error", C.PTR),
+ CREATE("create", C.I32, C.PTR, C.PTR, C.PTR),
+ SCHEMA("schema", C.I32, C.PTR, C.PTR),
+ CREATE_READER("create_reader", C.I32, C.PTR, C.PTR, C.PTR),
+ CREATE_STREAM_READER("create_stream_reader", C.I32, C.PTR, C.PTR, C.PTR),
+ CREATE_WRITER("create_writer", C.I32, C.PTR, C.PTR, C.PTR, C.PTR),
+ CLOSE("close", C.VOID, C.PTR),
+ PUSH_PREDICATES("push_predicates", C.I32, C.PTR, C.PTR, C.I64, C.PTR),
+ PUSH_LIMIT("push_limit", C.I32, C.PTR, C.I32, C.PTR),
+ PRUNE_COLUMNS("prune_columns", C.I32, C.PTR, C.PTR, C.I64, C.PTR),
+ PARTITIONS("partitions", C.I32, C.PTR, C.PTR, C.PTR),
+ SERIALIZE_READER("serialize_reader", C.I32, C.PTR, C.PTR, C.PTR),
+ CLOSE_READER("close_reader", C.VOID, C.PTR),
+ INITIAL_OFFSET("initial_offset", C.I32, C.PTR, C.PTR, C.PTR),
+ LATEST_OFFSET("latest_offset", C.I32, C.PTR, C.PTR, C.PTR),
+ STREAM_PARTITIONS("stream_partitions", C.I32, C.PTR, C.PTR, C.PTR, C.PTR, C.PTR),
+ SERIALIZE_STREAM_READER("serialize_stream_reader", C.I32, C.PTR, C.PTR, C.PTR),
+ COMMIT_OFFSET("commit_offset", C.I32, C.PTR, C.PTR),
+ CLOSE_STREAM_READER("close_stream_reader", C.VOID, C.PTR),
+ READ("read", C.I32, C.PTR, C.PTR, C.I64, C.PTR, C.I64, C.PTR),
+ SERIALIZE_WRITER("serialize_writer", C.I32, C.PTR, C.PTR, C.PTR),
+ COMMIT("commit", C.I32, C.PTR, C.I64, C.PTR, C.PTR, C.I64),
+ ABORT("abort", C.I32, C.PTR, C.I64, C.PTR, C.PTR, C.I64),
+ CLOSE_WRITER("close_writer", C.VOID, C.PTR),
+ CREATE_DATA_WRITER("create_data_writer", C.I32, C.PTR, C.PTR, C.I64, C.I64, C.PTR),
+ WRITE("write", C.I32, C.PTR, C.PTR, C.PTR),
+ COMMIT_DATA_WRITER("commit_data_writer", C.I32, C.PTR, C.PTR, C.PTR),
+ ABORT_DATA_WRITER("abort_data_writer", C.I32, C.PTR),
+ CLOSE_DATA_WRITER("close_data_writer", C.VOID, C.PTR);
+
+ /** The name of the field of the function in the table. */
+ final String fieldName;
+ private final C result;
+ private final C[] parameters;
+
+ Function(String fieldName, C result, C... parameters) {
+ this.fieldName = fieldName;
+ this.result = result;
+ this.parameters = parameters;
+ }
+
+ /** The offset of the function in the table, after the size and the version of the table. */
+ long offset() {
+ return TABLE_FUNCTIONS + 8L * ordinal();
+ }
+ }
+
+ // The layouts of the structs of the header, on the 64-bit platforms that Spark supports. Each
+ // struct starts with its size.
+ private static final long TABLE_ABI_VERSION = 8;
+ private static final long TABLE_FUNCTIONS = 16;
+ private static final long HOST_SIZE = 32;
+ private static final long HOST_ABI_VERSION = 8;
+ private static final long HOST_SPARK_VERSION = 16;
+ private static final long HOST_LOG = 24;
+ private static final long CREATE_INFO_SIZE = 40;
+ private static final long CREATE_INFO_NAME = 8;
+ private static final long CREATE_INFO_OPTION_KEYS = 16;
+ private static final long CREATE_INFO_OPTION_VALUES = 24;
+ private static final long CREATE_INFO_NUM_OPTIONS = 32;
+ private static final long TASK_SIZE = 40;
+ private static final long TASK_STAGE_ID = 8;
+ private static final long TASK_STAGE_ATTEMPT_NUMBER = 12;
+ private static final long TASK_PARTITION_ID = 16;
+ private static final long TASK_ATTEMPT_NUMBER = 20;
+ private static final long TASK_ATTEMPT_ID = 24;
+ private static final long TASK_IS_INTERRUPTED = 32;
+ private static final long WRITE_INFO_SIZE = 24;
+ private static final long WRITE_INFO_MODE = 8;
+ private static final long WRITE_INFO_QUERY_ID = 16;
+
+ /**
+ * The offsets of the fields of the structs that Spark reads or writes, by {@code struct.field},
+ * and the sizes of the structs, by {@code struct}, for the tests to compare with the header.
+ */
+ static Map<String, Long> layout() {
+ Map<String, Long> layout = new LinkedHashMap<>();
+ layout.put("SparkDsLibrary", TABLE_FUNCTIONS + 8L * Function.values().length);
+ layout.put("SparkDsLibrary.size", 0L);
+ layout.put("SparkDsLibrary.abi_version", TABLE_ABI_VERSION);
+ for (Function function : Function.values()) {
+ layout.put("SparkDsLibrary." + function.fieldName, function.offset());
+ }
+ layout.put("SparkDsHost", HOST_SIZE);
+ layout.put("SparkDsHost.size", 0L);
+ layout.put("SparkDsHost.abi_version", HOST_ABI_VERSION);
+ layout.put("SparkDsHost.spark_version", HOST_SPARK_VERSION);
+ layout.put("SparkDsHost.log", HOST_LOG);
+ layout.put("SparkDsCreateInfo", CREATE_INFO_SIZE);
+ layout.put("SparkDsCreateInfo.size", 0L);
+ layout.put("SparkDsCreateInfo.name", CREATE_INFO_NAME);
+ layout.put("SparkDsCreateInfo.option_keys", CREATE_INFO_OPTION_KEYS);
+ layout.put("SparkDsCreateInfo.option_values", CREATE_INFO_OPTION_VALUES);
+ layout.put("SparkDsCreateInfo.num_options", CREATE_INFO_NUM_OPTIONS);
+ layout.put("SparkDsTaskContext", TASK_SIZE);
+ layout.put("SparkDsTaskContext.size", 0L);
+ layout.put("SparkDsTaskContext.stage_id", TASK_STAGE_ID);
+ layout.put("SparkDsTaskContext.stage_attempt_number", TASK_STAGE_ATTEMPT_NUMBER);
+ layout.put("SparkDsTaskContext.partition_id", TASK_PARTITION_ID);
+ layout.put("SparkDsTaskContext.attempt_number", TASK_ATTEMPT_NUMBER);
+ layout.put("SparkDsTaskContext.task_attempt_id", TASK_ATTEMPT_ID);
+ layout.put("SparkDsTaskContext.is_interrupted", TASK_IS_INTERRUPTED);
+ layout.put("SparkDsWriteInfo", WRITE_INFO_SIZE);
+ layout.put("SparkDsWriteInfo.size", 0L);
+ layout.put("SparkDsWriteInfo.mode", WRITE_INFO_MODE);
+ layout.put("SparkDsWriteInfo.query_id", WRITE_INFO_QUERY_ID);
+ return layout;
+ }
+
+ private final int abiVersion;
+ // The functions of the table, by ordinal, null for the functions that the library leaves NULL.
+ private final MethodHandle[] functions = new MethodHandle[Function.values().length];
+
+ private NativeLibrary(long table) throws Throwable {
+ long size = Platform.getLong(null, table);
+ if (size < TABLE_FUNCTIONS) {
+ throw new LoadException(
+ "spark_ds_init returned a function table with a size of " + size + " bytes.", null, true);
+ }
+ abiVersion = Platform.getInt(null, table + TABLE_ABI_VERSION);
+ // The table of another version of the interface has another layout. Only the functions that
+ // fit in the size of the table are read: the table of a library built against an earlier
+ // header ends before the functions that were added later.
+ if (abiVersion == ABI_VERSION) {
+ for (Function function : Function.values()) {
+ if (function.offset() + 8 <= size) {
+ long address = Platform.getLong(null, table + function.offset());
+ if (address != 0) {
+ functions[function.ordinal()] = Ffm.downcall(
+ Ffm.OF_ADDRESS.invoke(address), function.result, function.parameters);
+ }
+ }
+ }
+ }
+ }
+
+ /**
+ * Loads the library at the given path, and gets its table with {@code spark_ds_init}. The
+ * library stays loaded until the JVM exits.
+ *
+ * @param sparkVersion the version of Spark, which the library gets in {@code SparkDsHost}
+ * @throws LoadException if the library cannot be loaded
+ */
+ static NativeLibrary load(String path, String sparkVersion) throws Throwable {
+ if (!isSupported()) {
+ throw new UnsupportedOperationException(
+ "Native data sources require Java " + MINIMUM_JAVA_VERSION + " or later.");
+ }
+ Object lookup;
+ try {
+ lookup = Ffm.LIBRARY_LOOKUP.invoke(Path.of(path), Ffm.GLOBAL_ARENA);
+ } catch (IllegalArgumentException | IllegalCallerException e) {
+ // The library is not a valid library, or native access is not enabled for Spark.
+ throw new LoadException(e.getMessage(), e);
+ } catch (LinkageError e) {
+ // The FFM API could not be initialized, see Ffm: report why, such as an
+ // IllegalCallerException when native access is not enabled for Spark.
+ throw new LoadException(
+ "Cannot use the Foreign Function and Memory API: " + rootCause(e), e);
+ }
+ Optional<?> init = (Optional<?>) Ffm.FIND.invoke(lookup, "spark_ds_init");
+ if (init.isEmpty()) {
+ throw new LoadException(
+ "The library does not export the function spark_ds_init.", null, true);
+ }
+ MethodHandle initFunction = Ffm.downcall(init.get(), C.I32, C.PTR, C.PTR);
+ // The library may keep the host until the JVM exits, even if spark_ds_init fails.
+ long host = Ffm.host(path, sparkVersion);
+ long out = Platform.allocateMemory(8);
+ try {
+ Platform.putLong(null, out, 0);
+ int status = (int) initFunction.invokeExact(host, out);
+ if (status != 0) {
+ throw new LoadException(
+ "spark_ds_init failed with the error code " + status + ".", null, true);
+ }
+ long table = Platform.getLong(null, out);
+ if (table == 0) {
+ throw new LoadException("spark_ds_init returned no function table.", null, true);
+ }
+ return new NativeLibrary(table);
+ } finally {
+ Platform.freeMemory(out);
+ }
+ }
+
+ /** The method handles of the FFM API, initialized when the first library is loaded. */
+ private static final class Ffm {
+ static final Class<?> LAYOUT_CLASS;
+ static final Object LINKER;
+ static final Object GLOBAL_ARENA;
+ static final Object NO_OPTIONS;
+ static final MethodHandle LIBRARY_LOOKUP;
+ static final MethodHandle FIND;
+ static final MethodHandle DOWNCALL;
+ static final MethodHandle UPCALL;
+ static final MethodHandle FUNCTION_OF;
+ static final MethodHandle FUNCTION_OF_VOID;
+ // (long) -> MemorySegment and (MemorySegment) -> long.
+ static final MethodHandle OF_ADDRESS;
+ static final MethodHandle ADDRESS_OF;
+ static final Object BOOL;
+ static final Object INT32;
+ static final Object INT64;
+ static final Object POINTER;
+ // NativeLibrary.log, which SparkDsHost.log calls with the path of its library.
+ static final MethodHandle LOG;
+ // The addresses of the C functions that call Java, which live as long as the JVM:
+ // SparkDsSink, which calls NativeLibrary.sink, and SparkDsTaskContext.is_interrupted, which
+ // calls NativeLibrary.isInterrupted.
+ static final long SINK;
+ static final long IS_INTERRUPTED;
+
+ static {
+ try {
+ // A full-privilege lookup, so that the restricted methods of the FFM API see this class
+ // as their caller, in the unnamed module that Spark enables native access for.
+ MethodHandles.Lookup lookup = MethodHandles.lookup();
+ Class<?> linker = Class.forName("java.lang.foreign.Linker");
+ Class<?> option = Class.forName("java.lang.foreign.Linker$Option");
+ Class<?> symbolLookup = Class.forName("java.lang.foreign.SymbolLookup");
+ Class<?> arena = Class.forName("java.lang.foreign.Arena");
+ Class<?> segment = Class.forName("java.lang.foreign.MemorySegment");
+ Class<?> layout = Class.forName("java.lang.foreign.MemoryLayout");
+ Class<?> layoutArray = Array.newInstance(layout, 0).getClass();
+ Class<?> valueLayout = Class.forName("java.lang.foreign.ValueLayout");
+ Class<?> descriptor = Class.forName("java.lang.foreign.FunctionDescriptor");
+ Class<?> options = Array.newInstance(option, 0).getClass();
+
+ LAYOUT_CLASS = layout;
+ LINKER = lookup.findStatic(linker, "nativeLinker", MethodType.methodType(linker))
+ .invoke();
+ GLOBAL_ARENA = lookup.findStatic(arena, "global", MethodType.methodType(arena)).invoke();
+ NO_OPTIONS = Array.newInstance(option, 0);
+ LIBRARY_LOOKUP = lookup.findStatic(symbolLookup, "libraryLookup",
+ MethodType.methodType(symbolLookup, Path.class, arena));
+ FIND = lookup.findVirtual(symbolLookup, "find",
+ MethodType.methodType(Optional.class, String.class));
+ DOWNCALL = lookup.findVirtual(linker, "downcallHandle",
+ MethodType.methodType(MethodHandle.class, segment, descriptor, options))
+ .asFixedArity();
+ UPCALL = lookup.findVirtual(linker, "upcallStub",
+ MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options))
+ .asFixedArity();
+ FUNCTION_OF = lookup.findStatic(descriptor, "of",
+ MethodType.methodType(descriptor, layout, layoutArray)).asFixedArity();
+ FUNCTION_OF_VOID = lookup.findStatic(descriptor, "ofVoid",
+ MethodType.methodType(descriptor, layoutArray)).asFixedArity();
+ OF_ADDRESS = lookup.findStatic(segment, "ofAddress",
+ MethodType.methodType(segment, long.class));
+ ADDRESS_OF = lookup.findVirtual(segment, "address", MethodType.methodType(long.class));
+ BOOL = valueLayout.getField("JAVA_BOOLEAN").get(null);
+ INT32 = valueLayout.getField("JAVA_INT").get(null);
+ INT64 = valueLayout.getField("JAVA_LONG").get(null);
+ POINTER = valueLayout.getField("ADDRESS").get(null);
+
+ LOG = lookup.findStatic(NativeLibrary.class, "log",
+ MethodType.methodType(void.class, String.class, int.class, long.class));
+ SINK = upcall(
+ lookup.findStatic(NativeLibrary.class, "sink",
+ MethodType.methodType(void.class, long.class, long.class, long.class)),
+ C.VOID, C.PTR, C.PTR, C.I64);
+ IS_INTERRUPTED = upcall(
+ lookup.findStatic(NativeLibrary.class, "isInterrupted",
+ MethodType.methodType(boolean.class, long.class)),
+ C.BOOL, C.PTR);
+ } catch (Throwable e) {
+ // The message names the cause, because the JVM only keeps this message for the later
+ // uses of the class.
+ throw new IllegalStateException("Cannot initialize the FFM API: " + rootCause(e), e);
+ }
+ }
+
+ /** Returns the FunctionDescriptor of a C function with the given types. */
+ static Object descriptor(C result, C... parameters) throws Throwable {
+ Object layouts = Array.newInstance(LAYOUT_CLASS, parameters.length);
+ for (int i = 0; i < parameters.length; i++) {
+ Array.set(layouts, i, layout(parameters[i]));
+ }
+ return result == C.VOID
+ ? FUNCTION_OF_VOID.invoke(layouts)
+ : FUNCTION_OF.invoke(layout(result), layouts);
+ }
+
+ static Object layout(C type) {
+ return switch (type) {
+ case BOOL -> BOOL;
+ case I32 -> INT32;
+ case I64 -> INT64;
+ case PTR -> POINTER;
+ case VOID -> throw new IllegalArgumentException("void has no layout");
+ };
+ }
+
+ /**
+ * Binds the C function at the given address, a MemorySegment, to a method handle whose
+ * parameters and result have the Java types of the given C types, with {@code long} for
+ * pointers.
+ */
+ static MethodHandle downcall(Object address, C result, C... parameters) throws Throwable {
+ MethodHandle handle = (MethodHandle) DOWNCALL.invoke(
+ LINKER, address, descriptor(result, parameters), NO_OPTIONS);
+ // Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses.
+ for (int i = 0; i < parameters.length; i++) {
+ if (parameters[i] == C.PTR) {
+ handle = MethodHandles.filterArguments(handle, i, OF_ADDRESS);
+ }
+ }
+ if (result == C.PTR) {
+ handle = MethodHandles.filterReturnValue(handle, ADDRESS_OF);
+ }
+ return handle;
+ }
+
+ /**
+ * Returns the address of a C function with the given types that calls `target`, whose
+ * parameters have the Java types of the C types, with {@code long} for pointers. The
+ * function lives as long as the JVM.
+ */
+ static long upcall(MethodHandle target, C result, C... parameters) throws Throwable {
+ for (int i = 0; i < parameters.length; i++) {
+ if (parameters[i] == C.PTR) {
+ target = MethodHandles.filterArguments(target, i, ADDRESS_OF);
+ }
+ }
+ Object stub = UPCALL.invoke(
+ LINKER, target, descriptor(result, parameters), GLOBAL_ARENA, NO_OPTIONS);
+ return (long) ADDRESS_OF.invoke(stub);
+ }
+
+ /**
+ * Allocates the SparkDsHost of a library. The library may keep it until the JVM exits, so it
+ * is never freed.
+ */
+ static long host(String path, String sparkVersion) throws Throwable {
+ long log = upcall(MethodHandles.insertArguments(LOG, 0, path), C.VOID, C.I32, C.PTR);
+ byte[] version = sparkVersion.getBytes(StandardCharsets.UTF_8);
+ long versionAddress = Platform.allocateMemory(version.length + 1L);
+ Platform.copyMemory(
+ version, Platform.BYTE_ARRAY_OFFSET, null, versionAddress, version.length);
+ Platform.putByte(null, versionAddress + version.length, (byte) 0);
+ long host = Platform.allocateMemory(HOST_SIZE);
+ Platform.setMemory(host, (byte) 0, HOST_SIZE);
+ Platform.putLong(null, host, HOST_SIZE);
+ Platform.putInt(null, host + HOST_ABI_VERSION, ABI_VERSION);
+ Platform.putLong(null, host + HOST_SPARK_VERSION, versionAddress);
+ Platform.putLong(null, host + HOST_LOG, log);
+ return host;
+ }
+ }
+
+ private static Throwable rootCause(Throwable e) {
+ Throwable cause = e;
+ while (cause.getCause() != null && cause.getCause() != cause) {
+ cause = cause.getCause();
+ }
+ return cause;
+ }
+
+ /** Thrown when a library cannot be loaded, or cannot be used. */
+ static final class LoadException extends RuntimeException {
+ /**
+ * Whether the library was loaded, in which case it stays loaded: Spark found that it cannot
+ * use it afterwards, for example because {@code spark_ds_init} failed.
+ */
+ final boolean loaded;
+
+ LoadException(String message, Throwable cause) {
+ this(message, cause, false);
+ }
+
+ LoadException(String message, Throwable cause, boolean loaded) {
+ super(message, cause);
+ this.loaded = loaded;
+ }
+ }
+
+ /** Thrown when a function is called that the library leaves NULL in its table, with its name. */
+ static final class MissingFunctionException extends RuntimeException {
+ MissingFunctionException(String name) {
+ super(name);
+ }
+ }
+
+ /**
+ * Called by the library, through {@code SparkDsHost.log}, with the path of the library. It must
+ * not throw: an exception thrown by an upcall crashes the JVM.
+ */
+ private static void log(String path, int level, long message) {
+ try {
+ String text = message == 0 ? "" : new String(readCString(message), StandardCharsets.UTF_8);
+ switch (level) {
+ case 1 -> LOGGER.error("Native data source library {}: {}",
+ MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text));
+ case 2 -> LOGGER.warn("Native data source library {}: {}",
+ MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text));
+ case 4 -> LOGGER.debug("Native data source library {}: {}", path, text);
+ case 5 -> LOGGER.trace("Native data source library {}: {}", path, text);
+ default -> LOGGER.info("Native data source library {}: {}",
+ MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text));
+ }
+ } catch (Throwable e) {
+ // The message is lost.
+ }
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Calls
+ // ---------------------------------------------------------------------------------------------
+
+ private MethodHandle function(Function function) {
+ MethodHandle handle = functions[function.ordinal()];
+ if (handle == null) {
+ throw new MissingFunctionException(function.fieldName);
+ }
+ return handle;
+ }
+
+ /** Throws the error that the library reported if a function failed. */
+ private void check(int status, Function function) throws Throwable {
+ if (status == 0) {
+ return;
+ }
+ MethodHandle lastError = functions[Function.LAST_ERROR.ordinal()];
+ long address = lastError == null ? 0 : (long) lastError.invokeExact();
+ String message = address == 0 ? "" : new String(readCString(address), StandardCharsets.UTF_8);
+ throw new RuntimeException(message.isEmpty()
+ ? function.fieldName + " failed with the error code " + status + "."
+ : message);
+ }
+
+ private static long readHandle(long out, Function function) {
+ long handle = Platform.getLong(null, out);
+ if (handle == 0) {
+ throw new IllegalStateException(function.fieldName + " returned a NULL handle.");
+ }
+ return handle;
+ }
+
+ /** The version of the interface that the table of the library implements. */
+ int abiVersion() {
+ return abiVersion;
+ }
+
+ /** Whether the library implements the function, which it does not leave NULL in its table. */
+ boolean hasFunction(Function function) {
+ return functions[function.ordinal()] != null;
+ }
+
+ /** A data source handle, with the capabilities of the data source. */
+ static final class DataSourceHandle {
+ final long handle;
+ final long capabilities;
+
+ DataSourceHandle(long handle, long capabilities) {
+ this.handle = handle;
+ this.capabilities = capabilities;
+ }
+ }
+
+ DataSourceHandle createDataSource(String name, String[] optionKeys, String[] optionValues)
+ throws Throwable {
+ MethodHandle f = function(Function.CREATE);
+ try (Memory memory = new Memory()) {
+ long info = memory.zeroed(CREATE_INFO_SIZE);
+ Platform.putLong(null, info, CREATE_INFO_SIZE);
+ Platform.putLong(null, info + CREATE_INFO_NAME, memory.string(name));
+ Platform.putLong(null, info + CREATE_INFO_OPTION_KEYS, memory.strings(optionKeys));
+ Platform.putLong(null, info + CREATE_INFO_OPTION_VALUES, memory.strings(optionValues));
+ Platform.putLong(null, info + CREATE_INFO_NUM_OPTIONS, optionKeys.length);
+ long capabilities = memory.zeroed(8);
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(info, capabilities, out), Function.CREATE);
+ return new DataSourceHandle(
+ readHandle(out, Function.CREATE), Platform.getLong(null, capabilities));
+ }
+ }
+
+ void schema(long dataSource, long schemaAddress) throws Throwable {
+ check((int) function(Function.SCHEMA).invokeExact(dataSource, schemaAddress),
+ Function.SCHEMA);
+ }
+
+ long createReader(long dataSource, long schemaAddress) throws Throwable {
+ return createWithSchema(Function.CREATE_READER, dataSource, schemaAddress);
+ }
+
+ long createStreamReader(long dataSource, long schemaAddress) throws Throwable {
+ return createWithSchema(Function.CREATE_STREAM_READER, dataSource, schemaAddress);
+ }
+
+ private long createWithSchema(Function function, long dataSource, long schemaAddress)
+ throws Throwable {
+ MethodHandle f = function(function);
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(dataSource, schemaAddress, out), function);
+ return readHandle(out, function);
+ }
+ }
+
+ /**
+ * Creates a writer.
+ *
+ * @param mode {@link #WRITE_APPEND} or {@link #WRITE_TRUNCATE}
+ * @param queryId the ID of the streaming query of a streaming write, or null for a batch write
+ */
+ long createWriter(long dataSource, long schemaAddress, int mode, String queryId)
+ throws Throwable {
+ MethodHandle f = function(Function.CREATE_WRITER);
+ try (Memory memory = new Memory()) {
+ long info = memory.zeroed(WRITE_INFO_SIZE);
+ Platform.putLong(null, info, WRITE_INFO_SIZE);
+ Platform.putInt(null, info + WRITE_INFO_MODE, mode);
+ Platform.putLong(
+ null, info + WRITE_INFO_QUERY_ID, queryId == null ? 0 : memory.string(queryId));
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(dataSource, schemaAddress, info, out), Function.CREATE_WRITER);
+ return readHandle(out, Function.CREATE_WRITER);
+ }
+ }
+
+ void closeDataSource(long dataSource) throws Throwable {
+ function(Function.CLOSE).invokeExact(dataSource);
+ }
+
+ /** Returns how the reader uses each predicate, as SPARK_DS_PREDICATE_* values. */
+ int[] pushPredicates(long reader, String[] predicates) throws Throwable {
+ MethodHandle f = function(Function.PUSH_PREDICATES);
+ try (Memory memory = new Memory()) {
+ long uses = memory.zeroed(4L * Math.max(1, predicates.length));
+ check((int) f.invokeExact(reader, memory.strings(predicates), (long) predicates.length,
+ uses), Function.PUSH_PREDICATES);
+ int[] result = new int[predicates.length];
+ for (int i = 0; i < result.length; i++) {
+ result[i] = Platform.getInt(null, uses + 4L * i);
+ }
+ return result;
+ }
+ }
+
+ boolean pushLimit(long reader, int limit) throws Throwable {
+ MethodHandle f = function(Function.PUSH_LIMIT);
+ try (Memory memory = new Memory()) {
+ long applied = memory.zeroed(1);
+ check((int) f.invokeExact(reader, limit, applied), Function.PUSH_LIMIT);
+ return Platform.getByte(null, applied) != 0;
+ }
+ }
+
+ boolean pruneColumns(long reader, String[] columnNames) throws Throwable {
+ MethodHandle f = function(Function.PRUNE_COLUMNS);
+ try (Memory memory = new Memory()) {
+ long pruned = memory.zeroed(1);
+ check((int) f.invokeExact(reader, memory.strings(columnNames), (long) columnNames.length,
+ pruned), Function.PRUNE_COLUMNS);
+ return Platform.getByte(null, pruned) != 0;
+ }
+ }
+
+ byte[][] partitions(long reader) throws Throwable {
+ MethodHandle f = function(Function.PARTITIONS);
+ try (Sink sink = new Sink()) {
+ check((int) f.invokeExact(reader, Ffm.SINK, sink.id), Function.PARTITIONS);
+ return sink.all();
+ }
+ }
+
+ byte[] serializeReader(long reader) throws Throwable {
+ return serialize(Function.SERIALIZE_READER, reader);
+ }
+
+ void closeReader(long reader) throws Throwable {
+ function(Function.CLOSE_READER).invokeExact(reader);
+ }
+
+ String initialOffset(long streamReader) throws Throwable {
+ return offset(Function.INITIAL_OFFSET, streamReader);
+ }
+
+ String latestOffset(long streamReader) throws Throwable {
+ return offset(Function.LATEST_OFFSET, streamReader);
+ }
+
+ private String offset(Function function, long streamReader) throws Throwable {
+ byte[] offset = serialize(function, streamReader);
+ if (offset == null) {
+ throw new IllegalStateException(function.fieldName + " returned no offset.");
+ }
+ return new String(offset, StandardCharsets.UTF_8);
+ }
+
+ byte[][] streamPartitions(long streamReader, String start, String end) throws Throwable {
+ MethodHandle f = function(Function.STREAM_PARTITIONS);
+ try (Memory memory = new Memory(); Sink sink = new Sink()) {
+ check((int) f.invokeExact(streamReader, memory.string(start), memory.string(end),
+ Ffm.SINK, sink.id), Function.STREAM_PARTITIONS);
+ return sink.all();
+ }
+ }
+
+ byte[] serializeStreamReader(long streamReader) throws Throwable {
+ return serialize(Function.SERIALIZE_STREAM_READER, streamReader);
+ }
+
+ void commitOffset(long streamReader, String end) throws Throwable {
+ MethodHandle f = function(Function.COMMIT_OFFSET);
+ try (Memory memory = new Memory()) {
+ check((int) f.invokeExact(streamReader, memory.string(end)), Function.COMMIT_OFFSET);
+ }
+ }
+
+ void closeStreamReader(long streamReader) throws Throwable {
+ function(Function.CLOSE_STREAM_READER).invokeExact(streamReader);
+ }
+
+ void read(Task task, byte[] readerState, byte[] partition, long streamAddress)
+ throws Throwable {
+ MethodHandle f = function(Function.READ);
+ try (Memory memory = new Memory()) {
+ check((int) f.invokeExact(task.address, memory.bytes(readerState),
+ (long) readerState.length, memory.bytes(partition), (long) partition.length,
+ streamAddress), Function.READ);
+ }
+ }
+
+ byte[] serializeWriter(long writer) throws Throwable {
+ return serialize(Function.SERIALIZE_WRITER, writer);
+ }
+
+ void commit(long writer, long epochId, byte[][] messages) throws Throwable {
+ finishWrite(Function.COMMIT, writer, epochId, messages);
+ }
+
+ void abort(long writer, long epochId, byte[][] messages) throws Throwable {
+ finishWrite(Function.ABORT, writer, epochId, messages);
+ }
+
+ private void finishWrite(Function function, long writer, long epochId, byte[][] messages)
+ throws Throwable {
+ MethodHandle f = function(function);
+ try (Memory memory = new Memory()) {
+ long pointers = memory.zeroed(8L * Math.max(1, messages.length));
+ long lengths = memory.zeroed(8L * Math.max(1, messages.length));
+ for (int i = 0; i < messages.length; i++) {
+ if (messages[i] != null) {
+ Platform.putLong(null, pointers + 8L * i, memory.bytes(messages[i]));
+ Platform.putLong(null, lengths + 8L * i, messages[i].length);
+ }
+ }
+ check((int) f.invokeExact(writer, epochId, pointers, lengths, (long) messages.length),
+ function);
+ }
+ }
+
+ void closeWriter(long writer) throws Throwable {
+ function(Function.CLOSE_WRITER).invokeExact(writer);
+ }
+
+ long createDataWriter(Task task, byte[] writerState, long epochId) throws Throwable {
+ MethodHandle f = function(Function.CREATE_DATA_WRITER);
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(task.address, memory.bytes(writerState),
+ (long) writerState.length, epochId, out), Function.CREATE_DATA_WRITER);
+ return readHandle(out, Function.CREATE_DATA_WRITER);
+ }
+ }
+
+ void write(long dataWriter, long arrayAddress, long schemaAddress) throws Throwable {
+ check((int) function(Function.WRITE).invokeExact(dataWriter, arrayAddress, schemaAddress),
+ Function.WRITE);
+ }
+
+ /**
+ * Commits a data writer. Then {@code committed} runs right away, before the message is checked:
+ * if the message is invalid, the commit fails, but the data writer must not be aborted.
+ */
+ byte[] commitDataWriter(long dataWriter, Runnable committed) throws Throwable {
+ MethodHandle f = function(Function.COMMIT_DATA_WRITER);
+ try (Sink sink = new Sink()) {
+ check((int) f.invokeExact(dataWriter, Ffm.SINK, sink.id), Function.COMMIT_DATA_WRITER);
+ committed.run();
+ byte[] message = single(sink, Function.COMMIT_DATA_WRITER);
+ return message == null ? new byte[0] : message;
+ }
+ }
+
+ void abortDataWriter(long dataWriter) throws Throwable {
+ check((int) function(Function.ABORT_DATA_WRITER).invokeExact(dataWriter),
+ Function.ABORT_DATA_WRITER);
+ }
+
+ void closeDataWriter(long dataWriter) throws Throwable {
+ function(Function.CLOSE_DATA_WRITER).invokeExact(dataWriter);
+ }
+
+ /**
+ * Calls a function that returns a single byte array through a sink. Returns null if the
+ * function does not call the sink.
+ */
+ private byte[] serialize(Function function, long handle) throws Throwable {
+ MethodHandle f = function(function);
+ try (Sink sink = new Sink()) {
+ check((int) f.invokeExact(handle, Ffm.SINK, sink.id), function);
+ return single(sink, function);
+ }
+ }
+
+ /** Returns the single byte array that a function passed to the sink, or null if none. */
+ private static byte[] single(Sink sink, Function function) throws Throwable {
+ byte[][] values = sink.all();
+ if (values.length > 1) {
+ throw new IllegalStateException(
+ function.fieldName + " returned " + values.length + " values instead of one.");
+ }
+ return values.length == 0 ? null : values[0];
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Memory
+ // ---------------------------------------------------------------------------------------------
+
+ /** Native memory for the arguments of a call, freed after the call. */
+ private static final class Memory implements AutoCloseable {
+ private final List<Long> allocations = new ArrayList<>();
+
+ long zeroed(long size) {
+ long address = Platform.allocateMemory(size);
+ allocations.add(address);
+ Platform.setMemory(address, (byte) 0, size);
+ return address;
+ }
+
+ /** Copies the bytes. Empty bytes are passed as a valid pointer too. */
+ long bytes(byte[] bytes) {
+ long address = zeroed(Math.max(1, bytes.length));
+ Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address, bytes.length);
+ return address;
+ }
+
+ /** Copies the string as a NUL-terminated UTF-8 string. */
+ long string(String value) {
Review Comment:
Nit: `string` copies `value.getBytes(UTF_8)`, so a NUL in a string ends it early for the library: an option value `"ab\u0000cd"` reaches `create` as `"ab"`, and `"\u0000"` as an empty string, without an error. Could Spark reject option keys and values that contain U+0000 with `NATIVE_DATA_SOURCE_ERROR`, or the header say what happens to them?
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -0,0 +1,688 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*
+ * spark_datasource.h: the binary interface between Spark and data sources implemented in native
+ * code, such as Rust or C++. Since Spark 4.4.0.
+ *
+ * A native data source is a shared library that exports a single C function, spark_ds_init,
+ * which returns the table of the functions that implement the data source, a SparkDsLibrary.
+ * Spark calls them with the Foreign Function and Memory API of Java, so native data sources
+ * require Java 22 or later. Each library has its own table, so any number of libraries can be
+ * loaded side by side. Spark adapts a native data source to Data Source V2, with batch and
+ * micro-batch streaming reads, batch and streaming writes, and predicate, column and limit
+ * pushdown.
+ *
+ * This header is self-contained: it only depends on the C standard library, and defines the
+ * structs of the Arrow C data interface unless they are already defined, as by Arrow C++. A library
+ * can copy it into its source tree, or use the copy in the include directory of a Spark
+ * distribution. A library built against it works with all the releases of Spark that support its
+ * SPARK_DS_ABI_VERSION.
+ *
+ * Spark finds the library of a data source automatically when it is installed like other native
+ * libraries and named after the data source: the library of my_source is
+ * spark_datasource_my_source, such as libspark_datasource_my_source.so in /usr/local/lib on Linux.
+ * Spark looks for it in the directories of java.library.path, which include LD_LIBRARY_PATH on
+ * Linux, and in the lib directory of each installation prefix in PATH. A library can also be
+ * distributed in a native data source package: a zip file with the extension .sparkpkg that
+ * contains a manifest named spark-native-datasource.json and the library built for one or more
+ * platforms, added to a session with spark.addArtifact or under the paths in
+ * spark.sql.dataSource.native.paths. See the "Native Data Sources" page of the Spark SQL guide for
+ * details and examples.
+ *
+ * Lifecycle
+ * ---------
+ * Spark loads a library once in each JVM, and calls spark_ds_init before any other function. It
+ * creates a data source with create on the driver for each schema inference, scan and write, and
+ * closes it right after. create also returns the capabilities of the data source: Spark reports
+ * the operations that a data source does not declare as unsupported, without calling the library.
+ * Spark also creates a data source while it analyzes a query, to reject a batch write that the
+ * data source does not support before it computes the data to write. That data source has the
+ * options of the table, those of CREATE TABLE for a table in a catalog, which the options of a
+ * write can override, so Spark checks the capabilities of each operation again.
+ *
+ * - Batch reads. On the driver, Spark creates a reader with create_reader, pushes operations down
+ * by calling, each at most once and in this order, push_predicates, push_limit and
+ * prune_columns, plans the partitions with partitions, serializes the state of the reader with
+ * serialize_reader, and closes it. Then it calls read on the executors, in one task for each
+ * partition. If Spark never plans the partitions, for example to only explain the plan of the
+ * query, it closes the reader once the reader is no longer referenced.
+ * - Streaming reads. On the driver, Spark creates a stream reader with create_stream_reader and
+ * keeps it until the streaming query stops. For each micro-batch, it asks for the latest_offset,
+ * plans the partitions with stream_partitions and serializes the state with
+ * serialize_stream_reader, then calls read on the executors. Offsets are defined by the library,
+ * as compact JSON on a single line: Spark stores each offset as a line of the offset log of the
+ * checkpoint, so it rejects an offset that is not valid JSON or contains a line break. When a
+ * streaming query restarts, from its checkpoint, Spark creates a new stream reader and may call
+ * stream_partitions again for the last range that it planned before, so the partitions of a
+ * range must read the same data each time.
+ * - Writes. On the driver, right before it starts the tasks of the write (or of a micro-batch),
+ * Spark creates a writer with create_writer, and serializes its state with serialize_writer.
+ * With adaptive query execution, Spark first runs the stages of the data to write that end with
+ * a shuffle, and creates no writer if they fail. The rest of the data, such as its last stage,
+ * is computed in the tasks of the write, so a writer may still be aborted because computing the
+ * data failed: do irreversible work, such as removing the existing data for an overwrite, in
+ * commit. On the executors, each task creates a data writer with create_data_writer, calls write
+ * for each batch of at most spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends
+ * once it holds spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls commit_data_writer,
+ * or abort_data_writer if the task or the commit fails, before it closes the data writer. Then
+ * the driver calls commit with the messages of the tasks if they all succeeded, or abort
+ * otherwise, including when commit fails, and closes the writer. Several attempts of a task may
+ * write the same partition, after a failure or with speculative execution, each with a data
+ * writer of its own. Spark lets the attempts of a partition commit one at a time, and passes to
+ * commit the message of one committed attempt for each partition, so the write consists of the
+ * data of these attempts only. It aborts the data writers of the other attempts, except on an
+ * executor that is lost, where a data writer is neither committed nor aborted, and may even have
+ * committed before its message reached the driver. A streaming write commits or aborts each
+ * micro-batch (epoch). Spark may commit the same epoch more than once, for example when a
+ * streaming query restarts after a failure, so committing an epoch that is already committed
+ * must have no effect, for exactly-once output. The ID of the query, which create_writer gets for
+ * a streaming write, stays the same when the query restarts from its checkpoint, so the ID of the
+ * query and the epoch ID identify the output of a micro-batch, also when other queries write to
+ * the same destination. A micro-batch streaming query creates a new stream writer for each
+ * micro-batch, and Spark closes the previous stream writer of the query when the query creates a
+ * new one. Spark is not told when a streaming query stops, so it closes the last stream writer of
+ * the query once the writer is no longer referenced, or when the query is restarted.
+ *
+ * Conventions
+ * -----------
+ * - Errors. The functions that return an int return 0 on success, or any other value on failure.
+ * They then set the message that last_error returns, and Spark reports it as a
+ * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail, and they may run after the query,
+ * from another thread: do the work that can fail a query, such as flushing data, in commit. A
+ * C++ exception or a Rust panic must not cross the interface: catch it and return an error
+ * instead.
+ * - Handles. create, create_reader, create_stream_reader, create_writer and create_data_writer
+ * return an opaque, non-NULL handle in their last argument, typically a pointer: Spark reports a
+ * NULL handle as an error. Spark releases each handle exactly once, with the close function of
+ * its kind: close, close_reader, close_stream_reader, close_writer or close_data_writer, and no
+ * other function releases a handle. A handle is not always released right after its last use: a
+ * batch reader whose partitions are never planned and the last stream writer of a query are
+ * released once they are no longer referenced (see Lifecycle above). A reader or writer handle
+ * must not depend on the data source handle it was created from, because Spark closes the data
+ * source handle right after creating it.
+ * - Memory. What Spark passes to a function, such as strings, byte arrays, SparkDsCreateInfo and
+ * SparkDsWriteInfo, is only valid during the call, except for SparkDsHost and
+ * SparkDsTaskContext, which document how long they live. Strings are UTF-8 and NUL-terminated.
+ * The JSON of the predicates is ASCII: it escapes the other characters, which any JSON parser
+ * decodes. The library returns strings and byte arrays by passing them to a SparkDsSink, which
+ * copies them.
+ * - State. Partitions, the state of readers and writers, and commit messages are byte arrays in a
+ * format defined by the library. Spark sends them to the executors and back.
+ * - Arrow data. Schemas and data are exchanged with the Arrow C data interface
+ * (https://arrow.apache.org/docs/format/CDataInterface.html), as structs allocated by Spark.
+ * For an output struct, the library initializes it and Spark takes ownership. For an input
+ * struct, the library may move it to take ownership; otherwise, Spark releases it after the
+ * call returns, so the library must not use it afterwards. The arrays that the library exports
+ * must have an offset of 0, at every level: Arrow Java ignores the offsets of the arrays it
+ * imports, so Spark rejects an array with another offset, such as a zero-copy slice. Copy a
+ * sliced array before exporting it. Spark also rejects an array with fewer values than its
+ * parent refers to, such as a column shorter than its batch. The buffers of an array must hold
+ * all of its values, as the C data interface requires: it does not carry the sizes of the
+ * buffers, so Spark cannot check them, and reads past the end of a buffer that is too short.
+ * - Schemas. Spark types map to Arrow types as in spark.sql.execution.arrow: for example, bigint
+ * is int64, string is utf8, timestamp is timestamp[us] with a time zone and timestamp_ntz is
+ * timestamp[us] without one. Data read from a library may also use large_utf8 and utf8_view for
+ * strings, large_binary and binary_view for binary data, and any names for the fields of the
+ * entries of a map. Spark rejects the other Arrow types, such as unsigned integers, decimal256,
+ * large lists, fixed-size lists and list views, and dictionary encoding.
+ * - Options. Option keys are passed in lower case, because Spark matches them
+ * case-insensitively.
+ * - Optional functions. A library only sets the functions of the operations it supports in its
+ * table, and leaves the others NULL. A data source must not declare a capability whose
+ * functions are NULL: Spark then fails the operation. The pushdown functions and the serialize
+ * functions are optional: without them, Spark pushes nothing down, and the reader or writer has
+ * no state. So are schema, without which the user must specify the schema, commit_offset and
+ * last_error.
+ * - Threads. Spark calls the library from multiple threads, and calls on different handles may
+ * run concurrently. Calls on the same handle never run concurrently. A handle may be released
+ * from a thread other than the threads that used it, such as a JVM Cleaner thread when Spark
+ * releases a handle once it is no longer referenced, for example the last stream writer of a
+ * streaming query. A SparkDsSink must be called on the thread of the call that it was passed
+ * to, during that call. The functions of SparkDsHost and SparkDsTaskContext, and the callbacks
+ * of the stream that read initializes, can be called from any thread. The callbacks of the
+ * stream never run concurrently.
+ *
+ * Compatibility
+ * -------------
+ * The function table and the structs that Spark passes to the library start with their size, so
+ * that they can grow. Within a SPARK_DS_ABI_VERSION:
+ * - fields are only appended to the structs, including SparkDsLibrary, and the existing fields
+ * never change, nor do the semantics of the existing functions;
+ * - Spark only reads the functions that fit in the size of the table of a library, so a library
+ * built against an earlier header keeps working, with the functions that it does not know as
+ * NULL. A library must only read the fields that fit in the size of a struct that Spark passes;
+ * - new values may appear where this documentation says so, such as the names of pushed
+ * predicates. Spark only uses a new capability or write mode with a data source that declares
+ * it.
+ * An incompatible change uses a new SPARK_DS_ABI_VERSION: spark_ds_init gets the newest version
+ * that Spark supports, the library returns a table of a version that it implements, and Spark
+ * rejects a table of a version that it does not support. Spark supports version 1 at least in all
+ * the releases of Spark 4 from 4.4.0.
+ */
+#ifndef SPARK_DATASOURCE_H
+#define SPARK_DATASOURCE_H
+
+#include <stdbool.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/* The version of this interface. */
+#define SPARK_DS_ABI_VERSION 1
+
+/* Exports spark_ds_init from the library, also when the library hides its other symbols. */
+#if defined(_WIN32)
+#define SPARK_DS_EXPORT __declspec(dllexport)
+#else
+#define SPARK_DS_EXPORT __attribute__((visibility("default")))
+#endif
+
+/* The Arrow C data interface: https://arrow.apache.org/docs/format/CDataInterface.html */
+#ifndef ARROW_C_DATA_INTERFACE
+#define ARROW_C_DATA_INTERFACE
+
+#define ARROW_FLAG_DICTIONARY_ORDERED 1
+#define ARROW_FLAG_NULLABLE 2
+#define ARROW_FLAG_MAP_KEYS_SORTED 4
+
+struct ArrowSchema {
+ const char* format;
+ const char* name;
+ const char* metadata;
+ int64_t flags;
+ int64_t n_children;
+ struct ArrowSchema** children;
+ struct ArrowSchema* dictionary;
+ void (*release)(struct ArrowSchema*);
+ void* private_data;
+};
+
+struct ArrowArray {
+ int64_t length;
+ int64_t null_count;
+ int64_t offset;
+ int64_t n_buffers;
+ int64_t n_children;
+ const void** buffers;
+ struct ArrowArray** children;
+ struct ArrowArray* dictionary;
+ void (*release)(struct ArrowArray*);
+ void* private_data;
+};
+
+#endif /* ARROW_C_DATA_INTERFACE */
+
+#ifndef ARROW_C_STREAM_INTERFACE
+#define ARROW_C_STREAM_INTERFACE
+
+struct ArrowArrayStream {
+ int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out);
+ int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out);
+ const char* (*get_last_error)(struct ArrowArrayStream*);
+ void (*release)(struct ArrowArrayStream*);
+ void* private_data;
+};
+
+#endif /* ARROW_C_STREAM_INTERFACE */
+
+/*
+ * Receives a string or a byte array that the library returns, and copies it. `sink` is the
+ * argument passed with the sink to the function that returns it.
+ */
+typedef void (*SparkDsSink)(void* sink, const uint8_t* data, int64_t length);
+
+/* ---------------------------------------------------------------------------------------------
+ * What Spark passes to the library
+ * --------------------------------------------------------------------------------------------- */
+
+/* The levels of the messages that a library logs with SparkDsHost.log. */
+#define SPARK_DS_LOG_ERROR 1
+#define SPARK_DS_LOG_WARN 2
+#define SPARK_DS_LOG_INFO 3
+#define SPARK_DS_LOG_DEBUG 4
+#define SPARK_DS_LOG_TRACE 5
+
+/* Spark, as spark_ds_init gets it. It stays valid until the JVM exits. */
+typedef struct SparkDsHost {
+ /*
+ * The size of this struct in the header that Spark was built with. A later version of Spark may
+ * append fields, so a library must only read the fields that fit in it.
+ */
+ int64_t size;
+ /* The newest version of this interface that Spark supports. */
+ int32_t abi_version;
+ /* The version of Spark, such as "4.4.0". */
+ const char* spark_version;
+ /*
+ * Logs a message, a NUL-terminated UTF-8 string, at one of the SPARK_DS_LOG_* levels: in the log
+ * of the driver or of the executor, with the path of the library.
+ */
+ void (*log)(int32_t level, const char* message);
+} SparkDsHost;
+
+/* A data source, as create gets it. */
+typedef struct SparkDsCreateInfo {
+ /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */
+ int64_t size;
+ /*
+ * The name of the data source, as listed in the package manifest, or in lower case for an
+ * installed library. A library can implement several data sources.
+ */
+ const char* name;
+ /*
+ * The num_options options given by the user, such as spark.read.option("key", "value"), with
+ * the keys in lower case.
+ */
+ const char* const* option_keys;
+ const char* const* option_values;
+ int64_t num_options;
+} SparkDsCreateInfo;
+
+/*
+ * The task that calls read or create_data_writer, on an executor. It stays valid until the stream
+ * that read returns is released, or until the data writer is closed.
+ */
+typedef struct SparkDsTaskContext {
+ /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */
+ int64_t size;
+ /* The ID of the stage of the task, and the attempt number of the stage. */
+ int32_t stage_id;
+ int32_t stage_attempt_number;
+ /* The partition that the task reads or writes, and the attempt number of the task. */
+ int32_t partition_id;
+ int32_t attempt_number;
+ /* The ID of the task attempt, unique in the Spark application. */
+ int64_t task_attempt_id;
+ /*
+ * Returns whether the task was killed, for example because its query was canceled. Spark cannot
+ * interrupt native code, so a read or a write that runs for long, or blocks, should check it
+ * regularly, and fail once it is true.
+ */
+ bool (*is_interrupted)(const struct SparkDsTaskContext* task);
+} SparkDsTaskContext;
+
+/* The modes of a write. */
+#define SPARK_DS_WRITE_APPEND 0
+#define SPARK_DS_WRITE_TRUNCATE 1
+
+/* A write, as create_writer gets it. */
+typedef struct SparkDsWriteInfo {
+ /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */
+ int64_t size;
+ /*
+ * SPARK_DS_WRITE_APPEND to add the data to the existing data, or SPARK_DS_WRITE_TRUNCATE to
+ * replace the existing data: with the save mode "overwrite" of a batch write, and in each
+ * micro-batch with the output mode "complete" of a streaming write. Spark only passes
+ * SPARK_DS_WRITE_TRUNCATE to a data source with the capability SPARK_DS_TRUNCATE.
+ */
+ int32_t mode;
+ /*
+ * The ID of the streaming query of a streaming write, or NULL for a batch write. It stays the
+ * same when the query restarts from its checkpoint, so with the epoch IDs, it identifies the
+ * commits of the query: see Writes above.
+ */
+ const char* query_id;
+} SparkDsWriteInfo;
+
+/* ---------------------------------------------------------------------------------------------
+ * What the library returns to Spark
+ * --------------------------------------------------------------------------------------------- */
+
+/* The capabilities of a data source, which create returns: the operations that it supports. */
+#define SPARK_DS_BATCH_READ 1 /* batch reads, with create_reader */
+#define SPARK_DS_MICRO_BATCH_READ 2 /* micro-batch streaming reads, with create_stream_reader */
+#define SPARK_DS_BATCH_WRITE 4 /* batch writes, with create_writer */
+#define SPARK_DS_STREAMING_WRITE 8 /* streaming writes, with create_writer */
+#define SPARK_DS_TRUNCATE 16 /* writes with the mode SPARK_DS_WRITE_TRUNCATE */
+
+/* How a reader uses a predicate pushed with push_predicates. */
+#define SPARK_DS_PREDICATE_UNUSED 0 /* not at all */
+/* only to skip data, so it may return rows for which the predicate is not true */
+#define SPARK_DS_PREDICATE_PARTIAL 1
+/* to only return the rows for which the predicate is true */
+#define SPARK_DS_PREDICATE_COMPLETE 2
+
+/*
+ * The functions of a library, which spark_ds_init returns. A library sets the functions of the
+ * operations it supports, and leaves the others NULL: see "Optional functions" above.
+ */
+typedef struct SparkDsLibrary {
+ /*
+ * The size of this struct in the header that the library was built with. Spark only reads the
+ * functions that fit in it.
+ */
+ int64_t size;
+ /* The version of this interface that the table implements, SPARK_DS_ABI_VERSION. */
+ int32_t abi_version;
+
+ /*
+ * Returns the message of the last error on the calling thread, as a NUL-terminated UTF-8 string
+ * that stays valid until the next call on that thread. Optional: without it, Spark only reports
+ * the value that the function returned.
+ */
+ const char* (*last_error)(void);
+
+ /* -------------------------------------------------------------------------------------------
+ * Data source (driver)
+ * ------------------------------------------------------------------------------------------- */
+
+ /*
+ * Creates a data source.
+ *
+ * info: the name of the data source, and the options given by the user
+ * capabilities: set to the capabilities of the data source, SPARK_DS_* flags combined with |,
+ * which may depend on its name and its options. It is 0 when create is called.
+ * source: set to the data source handle
+ */
+ int (*create)(const SparkDsCreateInfo* info, uint64_t* capabilities, void** source);
+
+ /*
+ * Exports the schema of the data. Spark calls it only when the user does not specify a schema.
+ *
+ * out: initialized with a struct type whose fields are the columns
+ */
+ int (*schema)(void* source, struct ArrowSchema* out);
+
+ /*
+ * Creates a reader for a batch scan.
+ *
+ * schema: the schema to read: the one the user specified, or the one returned by schema
+ * reader: set to the reader handle
+ */
+ int (*create_reader)(void* source, struct ArrowSchema* schema, void** reader);
+
+ /*
+ * Creates a reader for a micro-batch streaming scan.
+ *
+ * schema: the schema to read: the one the user specified, or the one returned by schema
+ * reader: set to the stream reader handle
+ */
+ int (*create_stream_reader)(void* source, struct ArrowSchema* schema, void** reader);
+
+ /*
+ * Creates a writer for a batch write, or for a micro-batch of a streaming write.
+ *
+ * schema: the schema of the data to write
+ * info: the mode of the write, and the ID of the query of a streaming write
+ * writer: set to the writer handle
+ */
+ int (*create_writer)(void* source, struct ArrowSchema* schema, const SparkDsWriteInfo* info,
+ void** writer);
+
+ /* Releases a data source handle. */
+ void (*close)(void* source);
+
+ /* -------------------------------------------------------------------------------------------
+ * Batch reader (driver)
+ * ------------------------------------------------------------------------------------------- */
+
+ /*
+ * Pushes predicates down to the reader. The predicates are combined with AND.
+ *
+ * Each predicate is a JSON expression tree. An expression is one of:
+ * - {"type": "column", "name": ["a", "b"]}: a column. A nested field has a name part for each
+ * level.
+ * - {"type": "literal", "dataType": "bigint", "value": 1}: a literal. The data type is a Spark
+ * SQL type name such as int, string or decimal(10,2). The value is null, a JSON boolean or
+ * number for boolean and numeric types, a string for string, a base64 string for binary, a
+ * decimal string for decimal, the days since the epoch for date, and the microseconds since
+ * the epoch for timestamp and timestamp_ntz.
+ * - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in
+ * org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=>
+ * (equal, with null equal to null), <, <=, > and >= with two children, IN with the value to
+ * look for and then the values to look in, IS_NULL and IS_NOT_NULL with one child,
+ * STARTS_WITH, ENDS_WITH and CONTAINS with the string and the string to look for, AND and OR
+ * with two predicates, and NOT with one. The children of the other functions are columns,
+ * literals or predicates.
+ * Spark does not pass predicates that it cannot express this way, such as predicates that use
+ * arithmetic or other functions, nor predicates on columns that contain strings with a
+ * non-binary collation, such as UTF8_LCASE: strings are compared by their UTF-8 bytes. A later
+ * version of Spark may pass predicates with other names without changing SPARK_DS_ABI_VERSION,
+ * so a reader must ignore a predicate that contains a name that it does not know, or a data
+ * type that it does not support.
+ *
+ * Spark evaluates all the predicates again after reading, so a reader may use them only to skip
+ * data, for example the files whose statistics show that no row matches. But it must never drop
+ * or skip a row for which the predicates are true under the semantics of Spark SQL, which Spark
+ * cannot detect. In particular:
+ * - For floating-point columns, NaN is equal to NaN and greater than any other value, including
+ * positive infinity, and -0.0 is equal to 0.0. IEEE 754 comparisons treat NaN differently, and
+ * total orders, such as f64::total_cmp in Rust or the comparison kernels of arrow-rs, also
+ * order -0.0 before 0.0. For example, d > 1.0 is true for a NaN value of d, and d = 0.0 for
+ * -0.0.
+ * - Comparisons with null are null, and a row only matches if the predicates are true. For
+ * example, NOT (i IN (1, NULL)) matches no row.
+ * - Strings are compared by their UTF-8 bytes, which order the characters outside the Basic
+ * Multilingual Plane after all the others, unlike comparisons of UTF-16 strings.
+ * - To use only a part of a predicate, a reader may take the parts that it does not support as
+ * true, but not under a NOT, where that would drop rows.
Review Comment:
The rule for using a part of a predicate is only safe under `AND` and `OR`, but Spark also pushes predicates as children of `=` and `<=>`. `(col("id") > 101) == (col("id") < 103)` is pushed as `{"type":"function","name":"=","children":[{..."name":">"...},{..."name":"<"...}]}`, and `(col("id") > 101).eqNullSafe(lit(True))` as `<=>` with a predicate child. Taking a part as true there can drop rows. For `(id > 101) = (id < 101)`, the row with `id = 101` matches (`false = false`), but a reader that takes `id < 101` as true evaluates `false = true` and skips it.
Could the rule say that a reader may only take a part as true when every function between that part and the root is `AND` or `OR`, and must otherwise not use the predicate? `predicate-cases.json` could have one such case.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "zhztheplayer (via GitHub)" <gi...@apache.org>.
zhztheplayer commented on code in PR #59424:
URL: https://github.com/apache/spark/pull/59424#discussion_r4293697066
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,980 @@
+/*
+ * 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
+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, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, 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 the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @param location where the library is on the driver. The readers and the writers distribute it
+ * to the executors of the active session when they are created.
+ * @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.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, 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", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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))
Review Comment:
I felt this may cause bug since reusing a fixed `VectorSchemaRoot` can cause data corruption when the output of `NativeDataSource` is accumulated by the downstream operators ([example](https://github.com/apache/spark/blob/ca6524e34ce487bb61ab1fda4b440a86e2c2177b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala#L191-L192)), since the new data will replace the content of the old vectors in early `ColumnarBatch` instances.
It might be worth to try the root-free `Data.importIntoVectorSchemaRoot` APIs that creates fresh roots internally?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6199181641
I opened a companion Rust SDK prototype in #59520 for this PR's C ABI, pinned to `51af65c`. It separates driver-side `BatchSource` / `ScanBuilder` from a serializable `BatchReaderFactory`, caches immutable plans after pushdown negotiation, and exposes a lifetime-bound task context with cooperative cancellation. It advertises only `SPARK_DS_BATCH_READ`, leaving unsupported function slots null, without changing the native ABI.
The SDK handles the function table, plan serialization, provider errors and unwinding panics, and Arrow ownership. It includes a Range example, packaging instructions, and passing standalone Rust/C tests. End-to-end Spark validation is still pending.
Feedback on the connector API and lifecycle is welcome!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6184834495
@LuciferYang Thank you, the comparison with lance-spark is very helpful. I agree with the positioning. In 51af65c63a93, the docs have a new "When to Use a Native Data Source" section, which leads with a source whose client exists only in native code, such as an in-house storage system in C++ or the C SDK of a vendor, and that Spark needs to read and write. The PR description leads with it too, and no longer lists Lance among the motivating examples. This work did not start from one specific source: the goal was to let a data source be written only in native code, which is why the examples in the docs read a range of numbers.
On scope, the line I would draw is Data Source V2 itself:
- Features of Data Source V2 that version 1 does not cover are in scope for later, as functions and capabilities appended to the table: aggregate pushdown, including `count(*)`, and top-N pushdown; statistics; reporting the partitioning and the sort order of the data, which storage-partitioned joins use and which is the closest to fragment-aware joins; overwrite by filter; and row-level operations with metadata columns, as `SupportsRowLevelOperations` and `SupportsDelta` do with a row ID. The Limitations section now lists top-N pushdown, row-level operations and metadata columns too.
- Extensions outside of Data Source V2, such as table functions like `VECTOR_SEARCH`, SQL commands like `CREATE INDEX`, and parser and optimizer rules, are out of the scope of this interface on purpose: they plug into Spark's extension points, such as `SparkSessionExtensions`, rather than into Data Source V2. A source that needs them is better served by a JVM connector like lance-spark, which can keep sharing the native read and write path through the Arrow C stream interface, as it does now. The new section of the docs says so.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6204557249
@zhouyuan Thank you, running a native data source under Gluten or Comet without copies is a good use of the C ABI.
2. **Threading contract**: agreed, and done in 534dfca15723. The header now says that the callbacks of the stream that `read` initializes can be called from any thread, but never concurrently, like the functions of `SparkDsHost` and `SparkDsTaskContext`, and the docs say so too. Spark itself still calls them from the thread of the task.
1. **A developer API to open a partition as an Arrow C stream**: this fits version 1 without changing the C ABI, since `read` already initializes a stream at an address that the caller allocates. I would rather add it in a follow-up, as a `@DeveloperApi` with its own review, than make it public in this PR, because a few things need deciding first: which object an engine gets it from; who owns the `SparkDsTaskContext`, and when it is freed; and what the engine must check before importing the arrays, which Spark checks today (sliced arrays, lengths shorter than the parent, the types it cannot read). After the planner split in 534dfca15723, from @gengliangwang's review, the reader factory of a batch scan only holds the location of the library, the name of the data source, the read schema and the state of the reader, which is what your `SupportsArrowStreamRead` needs.
**Columnar write entry point**: also a follow-up. The data writers of the adapter already take Arrow-backed `ColumnarBatch`es, and an engine could pass an Arrow C stream to `write` directly in the same way.
I'd be glad to have your review on both follow-ups.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #59424:
URL: https://github.com/apache/spark/pull/59424#discussion_r4320557939
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,1098 @@
+/*
+ * 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
+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.{ListVector, MapVector, StructVector}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
+
+import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.util.Utils
+
+/**
+ * A [[ColumnarDataSource]] implemented by a native library, through the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @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 =>
+ val schema = call(name, "infer the schema of") {
+ optional(NativeArrow.importSchema(address => library.schema(dataSource.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ NativeArrow.checkSupported(schema, name, "infer the schema of")
+ schema
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
Review Comment:
Thanks, fixed in 840971d8dc2f. `ColumnarStreamingWrite.createStreamingWriterFactory` now fails with the new `DATA_SOURCE_CONTINUOUS_WRITE_NOT_SUPPORTED` when the local property `ContinuousExecution.EPOCH_COORDINATOR_ID_KEY` is set, which only continuous processing sets, right before it runs the write. So Spark creates no writer, and a library is never left with data writers that are closed without a commit or an abort. Continuous reads already fail when the query starts, since a native scan does not report `CONTINUOUS_READ`. The header now says that streaming writes are micro-batch only, and the docs list continuous processing among the limitations. Test: `continuous processing fails with a native data source` writes `rate` to `native_sink` with `Trigger.Continuous`, and checks that the library created no writer.
##########
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.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_<name>`, such as
+ * `libspark_datasource_my_source.so` on Linux, implements the data source `<name>`. 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 =>
+ // The version of a package is the oldest that its library implements.
+ if (pkg.manifest.abiVersion <= NativeLibrary.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" -> NativeLibrary.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(
Review Comment:
Thanks, fixed in 840971d8dc2f with the first option, which is what Spark does for Python data sources: one registered again replaces the previous one, and a statically registered one, with a warning. The package that a session added last now takes precedence over the packages it added before and over the configured ones, and Spark logs once that it replaces them. So a session upgrades `rs-1.0.sparkpkg` by adding `rs-1.1.sparkpkg`, and a package that a user adds replaces an admin's package in `spark.sql.dataSource.native.paths` for that session. Two configured packages that provide the same data source still fail with `NATIVE_DATA_SOURCE_PACKAGE_CONFLICT`, whose message now names the configuration: they are files on the driver that can be removed, and their file names do not tell which one is newer. The docs describe the order and how to upgrade a package. Test: `the package that a session added last replaces the other packages`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala:
##########
@@ -0,0 +1,142 @@
+/*
+ * 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._
+import org.apache.spark.unsafe.types.UTF8String
+
+/**
+ * Encodes Data Source V2 predicates as the JSON expression trees passed to
+ * `push_predicates` of `spark_datasource.h`.
+ */
+object NativePredicates {
+ // Writes ASCII only, escaping the other characters, which any JSON parser decodes.
+ private val mapper = JsonMapper.builder().enable(JsonWriteFeature.ESCAPE_NON_ASCII).build()
+
+ /**
+ * The names of the predicates that are passed to the libraries, as listed in
+ * `push_predicates`. They are part of the interface, so a predicate with any other
+ * name, such as an arithmetic function, is not passed: its semantics, for example on overflow,
+ * would be part of the interface too, and the names of Data Source V2 grow with new releases.
+ */
+ 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] = {
Review Comment:
Thanks, fixed in 840971d8dc2f. Comparisons and `IN` are now only pushed when none of their operands is an array, a map or a struct, so `s1 = s2` and `a1 < a2` are no longer pushed. `IS NULL` and `IS NOT NULL` still are, for any type, and so are the comparisons of a field of a struct that has a simple type, such as `s.x = 1`. The header says both, and that binary values compare as unsigned bytes. `predicate-cases.json` now has a binary column, with cases for which unsigned and signed comparisons differ: `b > X'7F'` matches `X'80'`, `X'FF'` and `X'7F00'`. Test: `comparisons and IN on arrays and structs are not pushed`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,1098 @@
+/*
+ * 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
+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.{ListVector, MapVector, StructVector}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
+
+import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.util.Utils
+
+/**
+ * A [[ColumnarDataSource]] implemented by a native library, through the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @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 =>
+ val schema = call(name, "infer the schema of") {
+ optional(NativeArrow.importSchema(address => library.schema(dataSource.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ NativeArrow.checkSupported(schema, name, "infer the schema of")
+ schema
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ NativeArrow.checkSupported(schema, name, "plan a scan of")
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, executorLocation, name, schema, new NativeHandle(handle, library.closeReader))
+ }
+
+ override def streamReader(schema: StructType): ColumnarStreamReader = {
+ NativeArrow.checkSupported(schema, name, "plan a scan of")
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, address)
+ }
+ }
+ 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ entries.find(e => e.getKey.contains('\u0000') || e.getValue.contains('\u0000')).foreach { e =>
+ throw error(name, "create", s"The option ${NativeArrow.printable(e.getKey)} contains the " +
+ "character U+0000, which Spark cannot pass to a library.")
+ }
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ NativeArrow.checkSupported(schema, name, action)
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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)
+ }
+
+ /**
+ * Checks that the data agrees with the nullability of the schema to read: a column, a field,
+ * an element or a map value that it declares not nullable has no null, except where its parent
+ * is null. Spark relies on the nullability, for example to evaluate `isnull` of such a column to
+ * false without reading it.
+ */
+ private def checkNulls(): Unit = {
+ root.getFieldVectors.asScala.zip(schema.fields).foreach { case (vector, field) =>
+ checkNulls(vector, field.dataType, field.nullable, field.name, _ => true)
+ }
+ }
+
+ /** `referenced` tells whether a value belongs to a value of the parent that is not null. */
+ private def checkNulls(
+ vector: ValueVector,
+ dataType: DataType,
+ nullable: Boolean,
+ column: String,
+ referenced: Int => Boolean): Unit = {
+ if (!nullable && vector.getNullCount > 0 &&
+ (0 until vector.getValueCount).exists(i => vector.isNull(i) && referenced(i))) {
+ throw error(name, "read", s"The column $column has a null, but the schema to read " +
+ "declares it not nullable.")
+ }
+ val valid = (i: Int) => !vector.isNull(i) && referenced(i)
+ (dataType, vector) match {
+ case (struct: StructType, structVector: StructVector) =>
+ struct.fields.zipWithIndex.foreach { case (field, i) =>
+ checkNulls(structVector.getChildByOrdinal(i), field.dataType, field.nullable,
+ s"$column.${field.name}", valid)
+ }
+ case (map: MapType, mapVector: MapVector) =>
+ val entries = mapVector.getDataVector.asInstanceOf[StructVector]
+ val inMaps = elements(mapVector, valid)
+ // Spark does not allow null keys.
+ checkNulls(entries.getChildByOrdinal(0), map.keyType, nullable = false, s"$column.key",
+ inMaps)
+ checkNulls(entries.getChildByOrdinal(1), map.valueType, map.valueContainsNull,
+ s"$column.value", inMaps)
+ case (array: ArrayType, list: ListVector) =>
+ checkNulls(list.getDataVector, array.elementType, array.containsNull, s"$column.element",
+ elements(list, valid))
+ case _ =>
+ }
+ }
+
+ /** Whether each element of a list belongs to a list that `valid` accepts. */
+ private def elements(list: ListVector, valid: Int => Boolean): Int => Boolean = {
+ lazy val referenced = {
+ val bits = new java.util.BitSet(list.getDataVector.getValueCount)
+ (0 until list.getValueCount).foreach { i =>
+ if (valid(i)) bits.set(list.getElementStartIndex(i), list.getElementEndIndex(i))
+ }
+ bits
+ }
+ i => referenced.get(i)
+ }
+
+ 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))
+ checkNulls()
+ 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 = {
+ // The arrays of the children, one for each field of the type: Arrow Java only checks their
+ // number when an array has children, and imports the missing columns of a batch as nulls.
+ def children(
+ array: ArrowArray.Snapshot,
+ fields: Int,
+ owner: String,
+ expected: String): Seq[ArrowArray.Snapshot] = {
+ if (array.n_children != fields || (fields > 0 && array.children == 0)) {
+ throw error(name, "read", s"$owner has ${array.n_children} child arrays, but $expected.")
+ }
+ (0 until fields).map { i =>
+ val child = Platform.getLong(null, array.children + 8L * i)
+ if (child == 0) {
+ throw error(name, "read", s"$owner has a NULL child array.")
+ }
+ ArrowArray.wrap(child).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)
+ val fields = field.getChildren.size
+ children(array, fields, s"An array of the column $column", s"its type has $fields")
+ .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) {
Review Comment:
Thanks, fixed in 840971d8dc2f. `checkArrays` now rejects a batch whose struct marks rows as null, with `The struct of the batch marks 1 of its rows as null, but Spark only reads batches without null rows: mark the values of the columns as null instead.`
Two details differ from the suggestion. I checked them with `ImportRecordBatch` of Arrow C++ 24, on a struct of 2 rows:
| `null_count` | validity buffer | Arrow C++ | Spark now |
|---|---|---|---|
| -1 | no nulls | imported | read |
| -1 | first row null | `ArrowArray struct has non-zero null count, cannot be imported as RecordBatch` | rejected |
| 1 | first row null | same | rejected |
| -1 | none | imported | rejected |
- When the library did not count the nulls, with a `null_count` of -1, Spark counts them in the validity buffer, as Arrow C++ does, rather than rejecting a batch that has none. This reads `length / 8` bytes, which the buffer must hold anyway.
- A `null_count` of -1 without a validity buffer is not allowed by the C data interface: "The buffer pointers MAY be null only in two situations: 1. for the null bitmap buffer, if ArrowArray.null_count is 0". Arrow Java takes it as a mix of nulls and values, and fails to import it with `Could not load buffers for field ...: Cannot invoke "org.apache.arrow.memory.ArrowBuf.getReferenceManager()"`. So Spark now rejects it, for every array and not only the batch, with `The batch has a null_count of -1 but no validity buffer, which the C data interface only allows with a null_count of 0.`
The header now says both. The test library has a new option `null_rows`, and `reject the Arrow data that Spark cannot read` checks each row of the table, and a column with a `null_count` of -1 but no validity buffer.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,1098 @@
+/*
+ * 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
+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.{ListVector, MapVector, StructVector}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
+
+import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.util.Utils
+
+/**
+ * A [[ColumnarDataSource]] implemented by a native library, through the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @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 =>
+ val schema = call(name, "infer the schema of") {
+ optional(NativeArrow.importSchema(address => library.schema(dataSource.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ NativeArrow.checkSupported(schema, name, "infer the schema of")
+ schema
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ NativeArrow.checkSupported(schema, name, "plan a scan of")
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, executorLocation, name, schema, new NativeHandle(handle, library.closeReader))
+ }
+
+ override def streamReader(schema: StructType): ColumnarStreamReader = {
+ NativeArrow.checkSupported(schema, name, "plan a scan of")
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, address)
+ }
+ }
+ 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ entries.find(e => e.getKey.contains('\u0000') || e.getValue.contains('\u0000')).foreach { e =>
+ throw error(name, "create", s"The option ${NativeArrow.printable(e.getKey)} contains the " +
+ "character U+0000, which Spark cannot pass to a library.")
+ }
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ NativeArrow.checkSupported(schema, name, action)
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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)
+ }
+
+ /**
+ * Checks that the data agrees with the nullability of the schema to read: a column, a field,
+ * an element or a map value that it declares not nullable has no null, except where its parent
+ * is null. Spark relies on the nullability, for example to evaluate `isnull` of such a column to
+ * false without reading it.
+ */
+ private def checkNulls(): Unit = {
+ root.getFieldVectors.asScala.zip(schema.fields).foreach { case (vector, field) =>
+ checkNulls(vector, field.dataType, field.nullable, field.name, _ => true)
+ }
+ }
+
+ /** `referenced` tells whether a value belongs to a value of the parent that is not null. */
+ private def checkNulls(
+ vector: ValueVector,
+ dataType: DataType,
+ nullable: Boolean,
+ column: String,
+ referenced: Int => Boolean): Unit = {
+ if (!nullable && vector.getNullCount > 0 &&
+ (0 until vector.getValueCount).exists(i => vector.isNull(i) && referenced(i))) {
+ throw error(name, "read", s"The column $column has a null, but the schema to read " +
+ "declares it not nullable.")
+ }
+ val valid = (i: Int) => !vector.isNull(i) && referenced(i)
+ (dataType, vector) match {
+ case (struct: StructType, structVector: StructVector) =>
+ struct.fields.zipWithIndex.foreach { case (field, i) =>
+ checkNulls(structVector.getChildByOrdinal(i), field.dataType, field.nullable,
+ s"$column.${field.name}", valid)
+ }
+ case (map: MapType, mapVector: MapVector) =>
+ val entries = mapVector.getDataVector.asInstanceOf[StructVector]
+ val inMaps = elements(mapVector, valid)
+ // Spark does not allow null keys.
+ checkNulls(entries.getChildByOrdinal(0), map.keyType, nullable = false, s"$column.key",
+ inMaps)
+ checkNulls(entries.getChildByOrdinal(1), map.valueType, map.valueContainsNull,
+ s"$column.value", inMaps)
+ case (array: ArrayType, list: ListVector) =>
+ checkNulls(list.getDataVector, array.elementType, array.containsNull, s"$column.element",
+ elements(list, valid))
+ case _ =>
+ }
+ }
+
+ /** Whether each element of a list belongs to a list that `valid` accepts. */
+ private def elements(list: ListVector, valid: Int => Boolean): Int => Boolean = {
+ lazy val referenced = {
+ val bits = new java.util.BitSet(list.getDataVector.getValueCount)
+ (0 until list.getValueCount).foreach { i =>
+ if (valid(i)) bits.set(list.getElementStartIndex(i), list.getElementEndIndex(i))
+ }
+ bits
+ }
+ i => referenced.get(i)
+ }
+
+ 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))
+ checkNulls()
+ 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 = {
+ // The arrays of the children, one for each field of the type: Arrow Java only checks their
+ // number when an array has children, and imports the missing columns of a batch as nulls.
+ def children(
+ array: ArrowArray.Snapshot,
+ fields: Int,
+ owner: String,
+ expected: String): Seq[ArrowArray.Snapshot] = {
+ if (array.n_children != fields || (fields > 0 && array.children == 0)) {
+ throw error(name, "read", s"$owner has ${array.n_children} child arrays, but $expected.")
+ }
+ (0 until fields).map { i =>
+ val child = Platform.getLong(null, array.children + 8L * i)
+ if (child == 0) {
+ throw error(name, "read", s"$owner has a NULL child array.")
+ }
+ ArrowArray.wrap(child).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)
+ val fields = field.getChildren.size
+ children(array, fields, s"An array of the column $column", s"its type has $fields")
+ .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.")
+ }
+ val columns = root.getSchema.getFields.size
+ children(batch, columns, "The batch", s"the schema has $columns columns")
+ .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()
Review Comment:
Thanks, fixed in 840971d8dc2f: the stream is now released in a `finally` of closing the root and the dictionaries, and the task is closed after it, so a failure there no longer skips the release. The data writer was already in that order: `close_data_writer` is the first thing that its `close` does, before the task is closed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6199268275
Following up on [my two DSv2 adapter suggestions on #59313](https://github.com/apache/spark/pull/59313#issuecomment-6198311277), after checking this PR at `51af65c`. The planner/reader separation still applies; the capability suggestion is partly addressed and can now be narrowed to the DSv2 table layer.
1. **Separate driver-side planning from executor-side reading.**
[`ColumnarReader`](https://github.com/apache/spark/blob/51af65c63a9369d793aab1a198ddde4c90c8f85b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala#L62-L90) still combines pushdown negotiation and partition planning with executor `read(partition)`. [`NativeDataSourceReader`](https://github.com/apache/spark/blob/51af65c63a9369d793aab1a198ddde4c90c8f85b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala#L217-L295) still changes roles through a transient native handle and mutable `readerState`, and [`ColumnarScanBuilder.build()`](https://github.com/apache/spark/blob/51af65c63a9369d793aab1a198ddde4c90c8f85b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala#L142-L145) passes closures capturing the mutable builder.
Could we keep `ColumnarTable -> ColumnarScanBuilder -> ColumnarScan`, but introduce a driver-only `BatchReadPlanner` whose `plan()` returns `PlannedBatch(partitions, readerFactory)`? The serializable executor factory would contain only the library location, source name, read schema, and serialized native state. `build()` would establish stable scan metadata and explicit planner ownership, with defined behavior for repeated calls.
This can preserve the C ABI and lazy planning point: call `partitions`, then `serialize_reader`, and close the native reader in `finally`. Keep the existing caching of successful plans **and planning failures**, lazy batch-planner creation for streaming-only sources, and the Cleaner fallback for scans that never reach planning. The benefit is clearer ownership and a smaller executor object, without changing when native partitions are planned.
2. **Expose source-specific capabilities through `Table.capabilities()`.**
The new capability flags returned by native `create` and the [checks in `NativeDataSource.create`](https://github.com/apache/spark/blob/51af65c63a9369d793aab1a198ddde4c90c8f85b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala#L138-L155) address discovery and reject unsupported operations before invoking native reader/writer creation. However, [`ColumnarTable.capabilities()`](https://github.com/apache/spark/blob/51af65c63a9369d793aab1a198ddde4c90c8f85b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala#L55-L56) still advertises batch/streaming reads and writes plus truncate for every source, and [`NativeDataSourceV2.getTable()`](https://github.com/apache/spark/blob/51af65c63a9369d793aab1a198ddde4c90c8f85b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala#L59-L66) constructs that table without a capability override.
Could we propagate the existing source-specific capability information into the DSv2 table as well, so a batch-only source is represented as such during DSv2 capability checks? The discovery mechanism now exists, so there is no need to add another native ABI query. Since capabilities may depend on options and query options can override table options, the table-level declaration needs defined, stable semantics; the operation-specific checks should remain.
These are internal adapter suggestions. The Rust-side planner/factory split in #59520 complements them, but does not change these Scala classes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6180680750
@tdcmeehan Thank you for the detailed proposal, and for offering to write the `.proto`. I agree that version 1 is what libraries get built against, so its shape has to be settled now. In d2df236bd234, I kept typed C functions, but behind a single entry point and size-prefixed structs, the way ADBC's function table and JNI's `JNIEnv` grow, and moved what grows into fields and capabilities. On your five points:
1. **Every feature is a new function.** Features that are parameters now go into structs that grow: memory accounting, for example, would be functions appended to `SparkDsTaskContext`, and overwrite by filter a write mode with a field appended to `SparkDsWriteInfo`. A new operation, such as statistics, is still a new function, but a slot appended to the function table rather than a new exported symbol, and Spark only calls it for a data source that declares it.
2. **Each function is a separate contract.** Handles now follow one rule: each handle is released by the close function of its kind, and by nothing else. `commit_data_writer` no longer releases the data writer, which removes the class of bug that @LuciferYang found: Spark calls the new `close_data_writer` after the commit or the abort. What Spark passes is valid only during the call, except `SparkDsHost`, until the JVM exits, and `SparkDsTaskContext`, until the stream is released or the data writer is closed.
3. **Signatures are frozen.** `read` and `create_data_writer` get a `const SparkDsTaskContext*`, and `create_writer`, which replaces `create_writer` and `create_stream_writer`, gets a `const SparkDsWriteInfo*` with the write mode and the query ID. Both start with their size and only get fields appended, so a task context for memory accounting or cancellation does not need a second `read`.
4. **Support is inferred from symbol names.** A library exports only `spark_ds_init`, which returns its table, `SparkDsLibrary`, so in C and C++ the compiler checks the names and types of the functions. `create` returns the capabilities of each data source, which can depend on its options. Spark reports an operation that a data source does not declare as unsupported without calling the library, and if a declared capability needs a function that the table leaves NULL, it fails on the driver with the name of the function. `push_predicates` now answers, for each predicate, `SPARK_DS_PREDICATE_UNUSED`, `PARTIAL` (only to skip data) or `COMPLETE`. Spark shows the used ones as pushed and still evaluates all of them, so trusting `COMPLETE` can come later as an opt-in capability.
5. **Nothing flows from Spark to the library.** `spark_ds_init` gets `SparkDsHost`: the newest ABI version that Spark supports, the version of Spark, and a `log` function that writes to the log of the driver or the executor. `SparkDsTaskContext` has the IDs of the stage, partition, attempt and task attempt, and `is_interrupted`, so a read or a write that blocks can stop once its task is killed. The library returns a table of the version that it implements, so a new version does not need ADBC's fail-and-retry handshake.
Your table, in these terms:
| Addition | Now |
|---|---|
| Declared predicates, three-state answer | in version 1: the answers of `push_predicates` |
| Cancellation | in version 1: `SparkDsTaskContext.is_interrupted` |
| Memory accounting | functions appended to `SparkDsTaskContext` |
| Statistics, partitioning, sort order | a function appended to the table, with a capability |
| Aggregate and top-N pushdown | a function each appended to the table, with a capability |
| Admission control | a function appended to the table, with a capability |
| Overwrite by filter | a write mode, and a predicate field appended to `SparkDsWriteInfo`, with a capability |
| Substrait | a pushdown function appended to the table, with a capability |
Why I would keep typed functions rather than messages:
- The header only depends on the C standard library, like the Arrow C data interface that it builds on. Messages would add a protobuf runtime and generated code to every library, which weighs most on C and C++ libraries.
- The compilers check the table. With messages, the encoding is checked, but which fields a request carries, and who owns what, would still be prose.
- Control calls are rare, and the data is Arrow C structs either way, so neither design is faster.
Two costs that messages would avoid remain: a new operation is still a new function to specify, and a host that runs a library in another process would marshal each function instead of forwarding bytes.
Tests: `the bindings match the layout of spark_datasource.h` compiles the header and compares every offset and size with the bindings, and a library whose table ends before `read` checks that Spark does not read past the size of a table. Others cover the capabilities, the host and its logging, the task context, a read that only ends once `is_interrupted` is true, and a failing `spark_ds_init`. Once 4.4.0 is released, the C counterpart of your replay test would be to keep building the test library against the released header. I also updated the PR description.
Does this address your concern? If you still think the control calls should be messages, could you point to the addition that you think does not fit a table?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "zhztheplayer (via GitHub)" <gi...@apache.org>.
zhztheplayer commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6213737648
> This is what @zhztheplayer's proposal [SPARK-57468](https://issues.apache.org/jira/browse/SPARK-57468) tries to solve exactly.
Yes, I am working on the PR https://github.com/apache/spark/pull/59154 as a fix. It will help Spark wire legacy columnar operators and Arrow-based operators (like what this PR is adding) correctly.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6225142534
@sarutak Thank you for the thorough review, with a probe library. All the inline comments are addressed in 74aa0dffc534, with replies there.
On Arrow UDFs: I agree that the root cause is in the columnar input path, and thanks for showing it with `ArrowCachedBatchSerializer`. I opened SPARK-60294 with #59551. It copies the pass-through columns of each batch into vectors of the task with `VectorAppender`, and releases them once the output batch that holds them is no longer used. Its test reproduces both failures without native data sources, with a test source that refills the same vectors for every batch. On your two questions:
- This PR waits for #59551, rather than sending its batches to the row queue path: with the fix, the batches of a native data source take the Arrow columnar path like the batches of other sources, and the Arrow cache is fixed too.
- `Arrow UDFs on the columns of a native data source` runs a UDF over one batch, two and 2,000, and checks the pass-through columns of every row against the result of the UDF. Locally on Java 26, it fails without #59551 and passes with it. The CI of pull requests skips it, since the job that runs the native data source suites has no pandas.
@zhztheplayer Thanks for the review.
- Native library infrastructure: I agree that discovery, platform selection, distribution, checksums and loading are not specific to data sources. I would keep them internal in this PR, because a general facility needs its own API: what a plugin declares in a manifest, how it looks up its libraries, and how the libraries of the sessions are isolated. They are already in classes separate from the data source integration (`NativeDataSourcePackage`, `NativeLibraries` and `NativeDataSourceRegistry`, apart from `NativeDataSource`), so a follow-up can extract them without changing the package format, once there is a second user such as Gluten. I'd be glad to work on that with you.
- Columnar writes: yes, a desirable follow-up, which @zhouyuan asked for too. The writers of the adapter already take Arrow-backed `ColumnarBatch`es, so `V2TableWriteExec` could pass the batches of a columnar child to them directly.
- SPARK-57468 and #59154: as @sarutak said, they are about the conventions of the batches and the transitions between operators, so the lifetime of the input batches of Python UDFs needs its own fix, #59551.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "zhouyuan (via GitHub)" <gi...@apache.org>.
zhouyuan commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6196064644
Moving my comment from #59313 here, updated for the C header.
Thanks for this. The plain C ABI also makes native data sources a good fit for native execution engines running inside Spark, such as [Apache Gluten](https://github.com/apache/incubator-gluten) (Velox) and Comet (DataFusion): as the description says, other hosts can load the same libraries. The library reads the data, and the engine runs the rest of the query natively. What's missing is on the Spark side.
**Today**, a query on a native data source with Gluten enabled looks like this:
```
read -> ArrowArrayStream -> Arrow Java -> ArrowColumnVector
-> ColumnarToRow -> RowToVeloxColumnar -> Velox operators
```
That's two extra copies, so turning on a native engine can make these queries slower. With a small hook, the data could stay columnar in native memory from the source to the shuffle:
```
read -> ArrowArrayStream -> Velox (import) -> filter / agg / join / shuffle
```
A source library written once would then run on vanilla Spark and get native downstream execution under Gluten or Comet, without per-source connector work in each engine.
**Two small requests:**
1. **A developer API to open a partition as an Arrow C stream.** The library location, reader state and partitions are private to `execution.datasources.v2.ffi`, so an external engine can't call `read` for a planned scan today. Something like:
```scala
@DeveloperApi
trait SupportsArrowStreamRead { this: PartitionReaderFactory =>
/** Initializes the ArrowArrayStream at `streamAddress` with the data of `partition`.
* The caller owns the stream and validates the arrays it reads. */
def openArrowStream(partition: InputPartition, streamAddress: Long): Unit
}
```
`ColumnarPartitionReaderFactory` would implement it by calling `read` with Spark's own `SparkDsTaskContext`, which is freed when the caller releases the stream. Library loading, the checksum checks and `spark_ds_init` stay in Spark, and the vanilla path stays exactly as it is.
2. **A threading contract for the stream in the header.** The doc of `read` says Spark calls the stream "from the thread of the task that reads the partition", so library authors may rely on it. A native engine may call `get_next` from its own driver threads. Could version 1 guarantee instead that the stream may be called from any thread, though never concurrently? The functions of `SparkDsHost` and `SparkDsTaskContext` can already be called from any thread, so this would only extend that to the stream. It's free now, but hard to add once libraries depend on the current wording.
**Optional, could be a follow-up:** a columnar write entry point (the internal `DataWriter[ColumnarBatch]` that takes Arrow batches), so a native engine can write to a native sink without going through rows.
If this direction sounds reasonable, I'm happy to help review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6204554351
@gengliangwang Thank you for both suggestions, and for the Rust SDK prototype in #59520.
1. **Planning on the driver, reading on the executors**: done in 534dfca15723, as you suggested. `ColumnarReader` is replaced by a driver-only `ColumnarBatchPlanner`, which gets the pushdown, and whose `plan()` returns `ColumnarPlannedBatch(partitions, readerFactory)`. For a native data source, `NativeBatchPlanner` holds the native reader: `plan()` calls `partitions`, then `serialize_reader`, and closes the reader in `finally`. The serializable `NativePartitionReaderFactory` only has the location of the library, the name of the data source, the read schema and the state of the reader. `ColumnarScanBuilder.build()` builds its scan once, and the scan takes over the planner that the operations were pushed down to, with no closure over the builder. A scan whose operations were not pushed down only creates a planner if its batch is planned, so a streaming scan never creates one. `ColumnarBatchScan` plans once, when Spark first asks for the partitions or the reader factory, and keeps th
e result or the failure. The `Cleaner` still releases the reader of a scan that is never planned. Tests: `build and plan a batch scan once` in `ColumnarDataSourceSuite`, and `a scan whose planning failed fails again with the same error` in `NativeDataSourceSuite`, now through `ColumnarBatchScan`.
2. **Capabilities in `Table.capabilities()`**: done for the writes. `ColumnarTable.capabilities()` now asks the data source, once per table, with the options of the table: for a native data source, Spark creates a data source with them and reads its capabilities. The table reports `BATCH_WRITE` and `TRUNCATE` as the data source declares them, so Data Source V2 rejects a batch write or overwrite that the source does not support while it analyzes the query, with `UNSUPPORTED_FEATURE.TABLE_OPERATION`, before it computes the data to write. Before, with adaptive execution, the shuffle stages of the data to write ran first.
The table always reports `BATCH_READ`, `MICRO_BATCH_READ` and `STREAMING_WRITE`. When a `TableProvider`'s table lacks one of them, Data Source V2 does not reject the query: it falls back to Data Source V1, in `DataFrameReader.load`, `DataStreamReader.load` and `DataStreamWriter.start`. The query then fails with `<format> is not a valid Spark SQL Data Source.` or `Data source <format> does not support streamed reading.` (legacy errors), which do not say that the data source does not declare the operation. So those operations fail when Spark plans them instead, still on the driver and before any task, with `DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED`, `DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED` or `DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED`.
On the semantics: the table has the options of the table, which for a table in a catalog are those of `CREATE TABLE`, and the options of a write can override them. So, as you said, the checks of each operation remain, with its own options. The header and the docs say so.
The test `operations that a data source does not support` covers both: a batch write whose data would fail in a shuffle stage now fails while the query is analyzed, and a write to a table in a catalog whose options remove a capability of the table fails when it is planned.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "zhztheplayer (via GitHub)" <gi...@apache.org>.
zhztheplayer commented on code in PR #59424:
URL: https://github.com/apache/spark/pull/59424#discussion_r4293697066
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,980 @@
+/*
+ * 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
+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, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, 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 the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @param location where the library is on the driver. The readers and the writers distribute it
+ * to the executors of the active session when they are created.
+ * @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.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, 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", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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))
Review Comment:
I felt this may cause bug since reusing a fixed `VectorSchemaRoot` can cause data corruption when the output of `NativeDataSource` is accumulated by the downstream operators ([example](https://github.com/apache/spark/blob/ca6524e34ce487bb61ab1fda4b440a86e2c2177b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala#L191-L192)), since the new data will replace the content of the old vectors in early `ColumnarBatch` instances.
It might be worth to try the root-free `Data.importVectorSchemaRoot` APIs that creates fresh roots internally?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #59424:
URL: https://github.com/apache/spark/pull/59424#discussion_r4320558849
##########
sql/core/pom.xml:
##########
@@ -298,6 +298,10 @@
<groupId>org.apache.arrow</groupId>
<artifactId>arrow-compression</artifactId>
</dependency>
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>arrow-c-data</artifactId>
Review Comment:
Thanks, fixed in 840971d8dc2f:
- Before it loads a library, on the driver and on the executors, Spark now checks that the JNI library of the C data interface of Arrow Java loads, and otherwise fails with the new `NATIVE_DATA_SOURCE_UNSUPPORTED_PLATFORM`. The error names the platform and keeps Arrow's error as its cause: an `IllegalStateException` when Arrow Java has no library for the architecture, as on ppc64le or s390x, or an `UnsatisfiedLinkError` when the library does not link, as with musl.
- The docs list the platforms for which Arrow Java 19 ships the library: Linux with glibc and macOS on x86_64 and aarch64, and Windows on x86_64.
- In the manifest, `linux-<arch>` now means Linux with glibc, and Linux with musl is `linux-<arch>-musl`, as the Rust targets for musl end with `-musl`. Spark detects musl from the C library mapped in its own process (`/proc/self/maps`), rather than from files that a glibc system may also have, so it never picks a glibc build on musl, and a package can add a musl build once native data sources run there.
Tests: `fail clearly where the C data interface of Arrow Java does not load`, and the musl cases of `platform names`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #59424:
URL: https://github.com/apache/spark/pull/59424#discussion_r4269101424
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -0,0 +1,697 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.v2.ffi;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+import java.lang.reflect.Array;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.apache.spark.unsafe.Platform;
+
+/**
+ * A native data source library, which implements the C interface in
+ * {@code org/apache/spark/sql/datasource/spark_datasource.h}, called through the Foreign Function
+ * and Memory (FFM) API of Java 22 and later.
+ * <p>
+ * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached
+ * through method handles, looked up once when the first library is loaded. Each function of a
+ * library is bound to a method handle that takes and returns primitive values only, with
+ * pointers as {@code long} addresses, so the calls themselves use no reflection. The memory that
+ * Spark passes to the library is allocated with {@link Platform}.
+ * <p>
+ * Each library is looked up separately, so all the libraries can export the same functions.
+ * Calling a function that the library does not export throws {@link MissingFunctionException}. A
+ * function that fails throws a {@link RuntimeException} with the message of
+ * {@code spark_ds_last_error}.
+ */
+final class NativeLibrary {
+
+ /** The version of the interface that Spark supports, {@code SPARK_DS_ABI_VERSION}. */
+ static final int ABI_VERSION = 1;
+
+ /** The first version of Java with the final FFM API. */
+ static final int MINIMUM_JAVA_VERSION = 22;
+
+ /** Whether this JVM can load native data source libraries. */
+ static boolean isSupported() {
+ return Runtime.version().feature() >= MINIMUM_JAVA_VERSION;
+ }
+
+ // The C types of the parameters and the results of the functions.
+ private enum C { VOID, BOOL, INT32, INT64, POINTER }
+
+ private final MethodHandle lastError;
+ private final MethodHandle abiVersion;
+ private final MethodHandle create;
+ private final MethodHandle schema;
+ private final MethodHandle createReader;
+ private final MethodHandle createStreamReader;
+ private final MethodHandle createWriter;
+ private final MethodHandle createStreamWriter;
+ private final MethodHandle close;
+ private final MethodHandle pushPredicates;
+ private final MethodHandle pushLimit;
+ private final MethodHandle pruneColumns;
+ private final MethodHandle partitions;
+ private final MethodHandle serializeReader;
+ private final MethodHandle closeReader;
+ private final MethodHandle initialOffset;
+ private final MethodHandle latestOffset;
+ private final MethodHandle streamPartitions;
+ private final MethodHandle serializeStreamReader;
+ private final MethodHandle commitOffset;
+ private final MethodHandle closeStreamReader;
+ private final MethodHandle read;
+ private final MethodHandle serializeWriter;
+ private final MethodHandle commit;
+ private final MethodHandle abort;
+ private final MethodHandle closeWriter;
+ private final MethodHandle createDataWriter;
+ private final MethodHandle write;
+ private final MethodHandle commitDataWriter;
+ private final MethodHandle abortDataWriter;
+
+ private NativeLibrary(Object lookup) throws Throwable {
+ C i32 = C.INT32;
+ C i64 = C.INT64;
+ C ptr = C.POINTER;
+ lastError = bind(lookup, "spark_ds_last_error", ptr);
+ abiVersion = bind(lookup, "spark_ds_abi_version", i32);
+ create = bind(lookup, "spark_ds_create", i32, ptr, ptr, ptr, i64, ptr);
+ schema = bind(lookup, "spark_ds_schema", i32, ptr, ptr);
+ createReader = bind(lookup, "spark_ds_create_reader", i32, ptr, ptr, ptr);
+ createStreamReader = bind(lookup, "spark_ds_create_stream_reader", i32, ptr, ptr, ptr);
+ createWriter = bind(lookup, "spark_ds_create_writer", i32, ptr, ptr, C.BOOL, ptr);
+ createStreamWriter =
+ bind(lookup, "spark_ds_create_stream_writer", i32, ptr, ptr, C.BOOL, ptr, ptr);
+ close = bind(lookup, "spark_ds_close", C.VOID, ptr);
+ pushPredicates = bind(lookup, "spark_ds_push_predicates", i32, ptr, ptr, i64, ptr);
+ pushLimit = bind(lookup, "spark_ds_push_limit", i32, ptr, i32, ptr);
+ pruneColumns = bind(lookup, "spark_ds_prune_columns", i32, ptr, ptr, i64, ptr);
+ partitions = bind(lookup, "spark_ds_partitions", i32, ptr, ptr, ptr);
+ serializeReader = bind(lookup, "spark_ds_serialize_reader", i32, ptr, ptr, ptr);
+ closeReader = bind(lookup, "spark_ds_close_reader", C.VOID, ptr);
+ initialOffset = bind(lookup, "spark_ds_initial_offset", i32, ptr, ptr, ptr);
+ latestOffset = bind(lookup, "spark_ds_latest_offset", i32, ptr, ptr, ptr);
+ streamPartitions =
+ bind(lookup, "spark_ds_stream_partitions", i32, ptr, ptr, ptr, ptr, ptr);
+ serializeStreamReader =
+ bind(lookup, "spark_ds_serialize_stream_reader", i32, ptr, ptr, ptr);
+ commitOffset = bind(lookup, "spark_ds_commit_offset", i32, ptr, ptr);
+ closeStreamReader = bind(lookup, "spark_ds_close_stream_reader", C.VOID, ptr);
+ read = bind(lookup, "spark_ds_read", i32, ptr, i64, ptr, i64, ptr);
+ serializeWriter = bind(lookup, "spark_ds_serialize_writer", i32, ptr, ptr, ptr);
+ commit = bind(lookup, "spark_ds_commit", i32, ptr, i64, ptr, ptr, i64);
+ abort = bind(lookup, "spark_ds_abort", i32, ptr, i64, ptr, ptr, i64);
+ closeWriter = bind(lookup, "spark_ds_close_writer", C.VOID, ptr);
+ createDataWriter =
+ bind(lookup, "spark_ds_create_data_writer", i32, ptr, i64, i32, i64, i64, ptr);
+ write = bind(lookup, "spark_ds_write", i32, ptr, ptr, ptr);
+ commitDataWriter = bind(lookup, "spark_ds_commit_data_writer", i32, ptr, ptr, ptr);
+ abortDataWriter = bind(lookup, "spark_ds_abort_data_writer", i32, ptr);
+ }
+
+ /**
+ * Loads the library at the given path. The library stays loaded until the JVM exits.
+ *
+ * @throws LoadException if the library cannot be loaded
+ */
+ static NativeLibrary load(String path) throws Throwable {
+ if (!isSupported()) {
+ throw new UnsupportedOperationException(
+ "Native data sources require Java " + MINIMUM_JAVA_VERSION + " or later.");
+ }
+ Object lookup;
+ try {
+ lookup = Ffm.LIBRARY_LOOKUP.invoke(Path.of(path), Ffm.GLOBAL_ARENA);
+ } catch (IllegalArgumentException | IllegalCallerException e) {
+ // The library is not a valid library, or native access is not enabled for Spark.
+ throw new LoadException(e.getMessage(), e);
+ } catch (LinkageError e) {
+ // The FFM API could not be initialized, see Ffm: report why, such as an
+ // IllegalCallerException when native access is not enabled for Spark.
+ throw new LoadException(
+ "Cannot use the Foreign Function and Memory API: " + rootCause(e), e);
+ }
+ return new NativeLibrary(lookup);
+ }
+
+ /**
+ * Binds a function of the library to a method handle whose parameters and result have the
+ * Java types of the given C types, with {@code long} for pointers. Returns null if the library
+ * does not export the function.
+ */
+ private static MethodHandle bind(Object lookup, String name, C result, C... parameters)
+ throws Throwable {
+ Optional<?> symbol = (Optional<?>) Ffm.FIND.invoke(lookup, name);
+ if (symbol.isEmpty()) {
+ return null;
+ }
+ Object layouts = Ffm.layouts(parameters);
+ Object descriptor = result == C.VOID
+ ? Ffm.FUNCTION_OF_VOID.invoke(layouts)
+ : Ffm.FUNCTION_OF.invoke(Ffm.layout(result), layouts);
+ MethodHandle handle = (MethodHandle) Ffm.DOWNCALL.invoke(
+ Ffm.LINKER, symbol.get(), descriptor, Ffm.NO_OPTIONS);
+ // Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses.
+ for (int i = 0; i < parameters.length; i++) {
+ if (parameters[i] == C.POINTER) {
+ handle = MethodHandles.filterArguments(handle, i, Ffm.OF_ADDRESS);
+ }
+ }
+ if (result == C.POINTER) {
+ handle = MethodHandles.filterReturnValue(handle, Ffm.ADDRESS_OF);
+ }
+ return handle;
+ }
+
+ /** The method handles of the FFM API, initialized when the first library is loaded. */
+ private static final class Ffm {
+ static final Class<?> LAYOUT_CLASS;
+ static final Object LINKER;
+ static final Object GLOBAL_ARENA;
+ static final Object NO_OPTIONS;
+ static final MethodHandle LIBRARY_LOOKUP;
+ static final MethodHandle FIND;
+ static final MethodHandle DOWNCALL;
+ static final MethodHandle FUNCTION_OF;
+ static final MethodHandle FUNCTION_OF_VOID;
+ // (long) -> MemorySegment and (MemorySegment) -> long.
+ static final MethodHandle OF_ADDRESS;
+ static final MethodHandle ADDRESS_OF;
+ static final Object BOOL;
+ static final Object INT32;
+ static final Object INT64;
+ static final Object POINTER;
+ // The address of the C function SparkDsSink, which calls NativeLibrary.sink.
+ static final long SINK;
+
+ static {
+ try {
+ // A full-privilege lookup, so that the restricted methods of the FFM API see this class
+ // as their caller, in the unnamed module that Spark enables native access for.
+ MethodHandles.Lookup lookup = MethodHandles.lookup();
+ Class<?> linker = Class.forName("java.lang.foreign.Linker");
+ Class<?> option = Class.forName("java.lang.foreign.Linker$Option");
+ Class<?> symbolLookup = Class.forName("java.lang.foreign.SymbolLookup");
+ Class<?> arena = Class.forName("java.lang.foreign.Arena");
+ Class<?> segment = Class.forName("java.lang.foreign.MemorySegment");
+ Class<?> layout = Class.forName("java.lang.foreign.MemoryLayout");
+ Class<?> layoutArray = Array.newInstance(layout, 0).getClass();
+ Class<?> valueLayout = Class.forName("java.lang.foreign.ValueLayout");
+ Class<?> descriptor = Class.forName("java.lang.foreign.FunctionDescriptor");
+ Class<?> options = Array.newInstance(option, 0).getClass();
+
+ LAYOUT_CLASS = layout;
+ LINKER = lookup.findStatic(linker, "nativeLinker", MethodType.methodType(linker))
+ .invoke();
+ GLOBAL_ARENA = lookup.findStatic(arena, "global", MethodType.methodType(arena)).invoke();
+ NO_OPTIONS = Array.newInstance(option, 0);
+ LIBRARY_LOOKUP = lookup.findStatic(symbolLookup, "libraryLookup",
+ MethodType.methodType(symbolLookup, Path.class, arena));
+ FIND = lookup.findVirtual(symbolLookup, "find",
+ MethodType.methodType(Optional.class, String.class));
+ DOWNCALL = lookup.findVirtual(linker, "downcallHandle",
+ MethodType.methodType(MethodHandle.class, segment, descriptor, options))
+ .asFixedArity();
+ FUNCTION_OF = lookup.findStatic(descriptor, "of",
+ MethodType.methodType(descriptor, layout, layoutArray)).asFixedArity();
+ FUNCTION_OF_VOID = lookup.findStatic(descriptor, "ofVoid",
+ MethodType.methodType(descriptor, layoutArray)).asFixedArity();
+ OF_ADDRESS = lookup.findStatic(segment, "ofAddress",
+ MethodType.methodType(segment, long.class));
+ ADDRESS_OF = lookup.findVirtual(segment, "address", MethodType.methodType(long.class));
+ BOOL = valueLayout.getField("JAVA_BOOLEAN").get(null);
+ INT32 = valueLayout.getField("JAVA_INT").get(null);
+ INT64 = valueLayout.getField("JAVA_LONG").get(null);
+ POINTER = valueLayout.getField("ADDRESS").get(null);
+
+ // The upcall stub of the sink, which lives as long as the JVM.
+ MethodHandle target = lookup.findStatic(NativeLibrary.class, "sink",
+ MethodType.methodType(void.class, long.class, long.class, long.class));
+ target = MethodHandles.filterArguments(target, 0, ADDRESS_OF, ADDRESS_OF);
+ Object sinkDescriptor =
+ FUNCTION_OF_VOID.invoke(layouts(C.POINTER, C.POINTER, C.INT64));
+ Object stub = lookup.findVirtual(linker, "upcallStub",
+ MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options))
+ .asFixedArity()
+ .invoke(LINKER, target, sinkDescriptor, GLOBAL_ARENA, NO_OPTIONS);
+ SINK = (long) ADDRESS_OF.invoke(stub);
+ } catch (Throwable e) {
+ // The message names the cause, because the JVM only keeps this message for the later
+ // uses of the class.
+ throw new IllegalStateException("Cannot initialize the FFM API: " + rootCause(e), e);
+ }
+ }
+
+ /** Returns a MemoryLayout[] with the layouts of the given types. */
+ static Object layouts(C... types) {
+ Object layouts = Array.newInstance(LAYOUT_CLASS, types.length);
+ for (int i = 0; i < types.length; i++) {
+ Array.set(layouts, i, layout(types[i]));
+ }
+ return layouts;
+ }
+
+ static Object layout(C type) {
+ return switch (type) {
+ case BOOL -> BOOL;
+ case INT32 -> INT32;
+ case INT64 -> INT64;
+ case POINTER -> POINTER;
+ case VOID -> throw new IllegalArgumentException("void has no layout");
+ };
+ }
+ }
+
+ private static Throwable rootCause(Throwable e) {
+ Throwable cause = e;
+ while (cause.getCause() != null && cause.getCause() != cause) {
+ cause = cause.getCause();
+ }
+ return cause;
+ }
+
+ /** Thrown when a library cannot be loaded. */
+ static final class LoadException extends RuntimeException {
+ LoadException(String message, Throwable cause) {
+ super(message, cause);
+ }
+ }
+
+ /** Thrown when a function is called that the library does not export, with its name. */
+ static final class MissingFunctionException extends RuntimeException {
+ MissingFunctionException(String name) {
+ super(name);
+ }
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Calls
+ // ---------------------------------------------------------------------------------------------
+
+ private static MethodHandle require(MethodHandle function, String name) {
+ if (function == null) {
+ throw new MissingFunctionException(name);
+ }
+ return function;
+ }
+
+ /** Throws the error that the library reported if a function failed. */
+ private void check(int status, String name) throws Throwable {
+ if (status == 0) {
+ return;
+ }
+ long address = lastError == null ? 0 : (long) lastError.invokeExact();
+ String message = address == 0 ? "" : new String(readCString(address), StandardCharsets.UTF_8);
+ throw new RuntimeException(message.isEmpty()
+ ? name + " failed with the error code " + status + "."
+ : message);
+ }
+
+ private static long readHandle(long out, String name) {
+ long handle = Platform.getLong(null, out);
+ if (handle == 0) {
+ throw new IllegalStateException(name + " returned a NULL handle.");
+ }
+ return handle;
+ }
+
+ int abiVersion() throws Throwable {
+ return (int) require(abiVersion, "spark_ds_abi_version").invokeExact();
+ }
+
+ long createDataSource(String name, String[] optionKeys, String[] optionValues)
+ throws Throwable {
+ MethodHandle f = require(create, "spark_ds_create");
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(memory.string(name), memory.strings(optionKeys),
+ memory.strings(optionValues), (long) optionKeys.length, out), "spark_ds_create");
+ return readHandle(out, "spark_ds_create");
+ }
+ }
+
+ void schema(long dataSource, long schemaAddress) throws Throwable {
+ check((int) require(schema, "spark_ds_schema").invokeExact(dataSource, schemaAddress),
+ "spark_ds_schema");
+ }
+
+ long createReader(long dataSource, long schemaAddress) throws Throwable {
+ return createWithSchema(createReader, "spark_ds_create_reader", dataSource, schemaAddress);
+ }
+
+ long createStreamReader(long dataSource, long schemaAddress) throws Throwable {
+ return createWithSchema(
+ createStreamReader, "spark_ds_create_stream_reader", dataSource, schemaAddress);
+ }
+
+ private long createWithSchema(
+ MethodHandle function, String name, long dataSource, long schemaAddress) throws Throwable {
+ MethodHandle f = require(function, name);
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(dataSource, schemaAddress, out), name);
+ return readHandle(out, name);
+ }
+ }
+
+ long createWriter(long dataSource, long schemaAddress, boolean overwrite) throws Throwable {
+ MethodHandle f = require(createWriter, "spark_ds_create_writer");
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(dataSource, schemaAddress, overwrite, out),
+ "spark_ds_create_writer");
+ return readHandle(out, "spark_ds_create_writer");
+ }
+ }
+
+ long createStreamWriter(long dataSource, long schemaAddress, boolean overwrite, String queryId)
+ throws Throwable {
+ MethodHandle f = require(createStreamWriter, "spark_ds_create_stream_writer");
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(dataSource, schemaAddress, overwrite, memory.string(queryId), out),
+ "spark_ds_create_stream_writer");
+ return readHandle(out, "spark_ds_create_stream_writer");
+ }
+ }
+
+ void closeDataSource(long dataSource) throws Throwable {
+ require(close, "spark_ds_close").invokeExact(dataSource);
+ }
+
+ boolean[] pushPredicates(long reader, String[] predicates) throws Throwable {
+ MethodHandle f = require(pushPredicates, "spark_ds_push_predicates");
+ try (Memory memory = new Memory()) {
+ long evaluated = memory.zeroed(Math.max(1, predicates.length));
+ check((int) f.invokeExact(reader, memory.strings(predicates), (long) predicates.length,
+ evaluated), "spark_ds_push_predicates");
+ boolean[] result = new boolean[predicates.length];
+ for (int i = 0; i < result.length; i++) {
+ result[i] = Platform.getByte(null, evaluated + i) != 0;
+ }
+ return result;
+ }
+ }
+
+ boolean pushLimit(long reader, int limit) throws Throwable {
+ MethodHandle f = require(pushLimit, "spark_ds_push_limit");
+ try (Memory memory = new Memory()) {
+ long applied = memory.zeroed(1);
+ check((int) f.invokeExact(reader, limit, applied), "spark_ds_push_limit");
+ return Platform.getByte(null, applied) != 0;
+ }
+ }
+
+ boolean pruneColumns(long reader, String[] columnNames) throws Throwable {
+ MethodHandle f = require(pruneColumns, "spark_ds_prune_columns");
+ try (Memory memory = new Memory()) {
+ long pruned = memory.zeroed(1);
+ check((int) f.invokeExact(reader, memory.strings(columnNames), (long) columnNames.length,
+ pruned), "spark_ds_prune_columns");
+ return Platform.getByte(null, pruned) != 0;
+ }
+ }
+
+ byte[][] partitions(long reader) throws Throwable {
+ MethodHandle f = require(partitions, "spark_ds_partitions");
+ try (Sink sink = new Sink()) {
+ check((int) f.invokeExact(reader, Ffm.SINK, sink.id), "spark_ds_partitions");
+ return sink.all();
+ }
+ }
+
+ byte[] serializeReader(long reader) throws Throwable {
+ return serialize(serializeReader, "spark_ds_serialize_reader", reader);
+ }
+
+ void closeReader(long reader) throws Throwable {
+ require(closeReader, "spark_ds_close_reader").invokeExact(reader);
+ }
+
+ String initialOffset(long streamReader) throws Throwable {
+ return offset(initialOffset, "spark_ds_initial_offset", streamReader);
+ }
+
+ String latestOffset(long streamReader) throws Throwable {
+ return offset(latestOffset, "spark_ds_latest_offset", streamReader);
+ }
+
+ private String offset(MethodHandle function, String name, long streamReader)
+ throws Throwable {
+ byte[] offset = serialize(function, name, streamReader);
+ if (offset == null) {
+ throw new IllegalStateException(name + " returned no offset.");
+ }
+ return new String(offset, StandardCharsets.UTF_8);
+ }
+
+ byte[][] streamPartitions(long streamReader, String start, String end) throws Throwable {
+ MethodHandle f = require(streamPartitions, "spark_ds_stream_partitions");
+ try (Memory memory = new Memory(); Sink sink = new Sink()) {
+ check((int) f.invokeExact(streamReader, memory.string(start), memory.string(end),
+ Ffm.SINK, sink.id), "spark_ds_stream_partitions");
+ return sink.all();
+ }
+ }
+
+ byte[] serializeStreamReader(long streamReader) throws Throwable {
+ return serialize(serializeStreamReader, "spark_ds_serialize_stream_reader", streamReader);
+ }
+
+ void commitOffset(long streamReader, String end) throws Throwable {
+ MethodHandle f = require(commitOffset, "spark_ds_commit_offset");
+ try (Memory memory = new Memory()) {
+ check((int) f.invokeExact(streamReader, memory.string(end)), "spark_ds_commit_offset");
+ }
+ }
+
+ void closeStreamReader(long streamReader) throws Throwable {
+ require(closeStreamReader, "spark_ds_close_stream_reader").invokeExact(streamReader);
+ }
+
+ void read(byte[] readerState, byte[] partition, long streamAddress) throws Throwable {
+ MethodHandle f = require(read, "spark_ds_read");
+ try (Memory memory = new Memory()) {
+ check((int) f.invokeExact(memory.bytes(readerState), (long) readerState.length,
+ memory.bytes(partition), (long) partition.length, streamAddress), "spark_ds_read");
+ }
+ }
+
+ byte[] serializeWriter(long writer) throws Throwable {
+ return serialize(serializeWriter, "spark_ds_serialize_writer", writer);
+ }
+
+ void commit(long writer, long epochId, byte[][] messages) throws Throwable {
+ finishWrite(commit, "spark_ds_commit", writer, epochId, messages);
+ }
+
+ void abort(long writer, long epochId, byte[][] messages) throws Throwable {
+ finishWrite(abort, "spark_ds_abort", writer, epochId, messages);
+ }
+
+ private void finishWrite(
+ MethodHandle function,
+ String name,
+ long writer,
+ long epochId,
+ byte[][] messages) throws Throwable {
+ MethodHandle f = require(function, name);
+ try (Memory memory = new Memory()) {
+ long pointers = memory.zeroed(8L * Math.max(1, messages.length));
+ long lengths = memory.zeroed(8L * Math.max(1, messages.length));
+ for (int i = 0; i < messages.length; i++) {
+ if (messages[i] != null) {
+ Platform.putLong(null, pointers + 8L * i, memory.bytes(messages[i]));
+ Platform.putLong(null, lengths + 8L * i, messages[i].length);
+ }
+ }
+ check((int) f.invokeExact(writer, epochId, pointers, lengths, (long) messages.length),
+ name);
+ }
+ }
+
+ void closeWriter(long writer) throws Throwable {
+ require(closeWriter, "spark_ds_close_writer").invokeExact(writer);
+ }
+
+ long createDataWriter(byte[] writerState, int partitionId, long taskId, long epochId)
+ throws Throwable {
+ MethodHandle f = require(createDataWriter, "spark_ds_create_data_writer");
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(memory.bytes(writerState), (long) writerState.length,
+ partitionId, taskId, epochId, out), "spark_ds_create_data_writer");
+ return readHandle(out, "spark_ds_create_data_writer");
+ }
+ }
+
+ void write(long dataWriter, long arrayAddress, long schemaAddress) throws Throwable {
+ check((int) require(write, "spark_ds_write").invokeExact(
+ dataWriter, arrayAddress, schemaAddress), "spark_ds_write");
+ }
+
+ byte[] commitDataWriter(long dataWriter) throws Throwable {
+ byte[] message = serialize(commitDataWriter, "spark_ds_commit_data_writer", dataWriter);
Review Comment:
Thanks, good catch. Fixed in e5352e191c14: `NativeLibrary.commitDataWriter` runs a callback as soon as the library returns 0, and `NativeDataWriter` clears the handle there, before the message in the sink is checked. A commit that the library completed is never aborted, and a commit that failed in the library is still aborted, as for dongjoon-hyun's comment on #59313. The header now says that Spark does not use the handle again once `spark_ds_commit_data_writer` succeeds, even if it then fails the task.
New test: "do not abort a data writer that committed with an invalid message". The test library passes its message twice (`commit_messages=2`) and records it if Spark aborts a data writer that it released, instead of crashing. The write fails with `spark_ds_commit_data_writer returned 2 values instead of one.`, the released data writer is not aborted, and the driver aborts the write. With the handle cleared after the call returns, as before, the test fails.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "zhztheplayer (via GitHub)" <gi...@apache.org>.
zhztheplayer commented on code in PR #59424:
URL: https://github.com/apache/spark/pull/59424#discussion_r4293697066
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,980 @@
+/*
+ * 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
+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, TaskContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.catalog.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.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, 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 the C interface in
+ * `org/apache/spark/sql/datasource/spark_datasource.h`.
+ *
+ * @param location where the library is on the driver. The readers and the writers distribute it
+ * to the executors of the active session when they are created.
+ * @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.handle, address)))
+ }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ }
+ }
+
+ /**
+ * Returns the batch writes that the data source declares, and always the reads and the
+ * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source,
+ * which then fails with errors that do not say why. A scan or a streaming write that the data
+ * source does not declare fails when Spark plans it instead, on the driver.
+ */
+ override def capabilities(): util.Set[TableCapability] = {
+ val library = NativeLibraries.get(location)
+ val declared = withDataSource(library, "check the capabilities of")(_.capabilities)
+ val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE)
+ if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
+ capabilities.add(BATCH_WRITE)
+ }
+ if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ capabilities.add(TRUNCATE)
+ }
+ capabilities
+ }
+
+ override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ,
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createReader(dataSource.handle, address)
+ }
+ }
+ new NativeBatchPlanner(
+ library, 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", NativeLibrary.MICRO_BATCH_READ,
+ "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamReader(dataSource.handle, 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(
+ createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, 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: NativeLibrary.DataSourceHandle => T): T = {
+ // Checked first, so that a data source is never created without a way to release it.
+ requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE))
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource.handle))
+ }
+ }
+
+ /**
+ * Creates a reader or a writer with `f`, which fails if the data source does not declare the
+ * capability, or if the library does not implement the functions that it requires.
+ */
+ private def create(
+ library: NativeLibrary,
+ action: String,
+ capability: Long,
+ unsupportedError: String,
+ functions: Seq[NativeLibrary.Function])(
+ f: NativeLibrary.DataSourceHandle => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ if ((dataSource.capabilities & capability) == 0) {
+ throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name))
+ }
+ requireFunctions(library, action, functions)
+ call(name, action)(f(dataSource))
+ }
+ }
+
+ /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */
+ private def createWriter(
+ library: NativeLibrary,
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): Long = {
+ val action = "plan a write to"
+ val (capability, unsupportedError, overwriteMode) = if (queryId == null) {
+ (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED",
+ "the save mode overwrite")
+ } else {
+ (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED",
+ "the output mode complete")
+ }
+ create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) {
+ dataSource =>
+ if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) {
+ throw error(name, action, s"The data source does not support $overwriteMode.")
+ }
+ val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource.handle, address, mode, queryId)
+ }
+ }
+ }
+
+ private def requireFunctions(
+ library: NativeLibrary,
+ action: String,
+ functions: Seq[NativeLibrary.Function]): Unit = {
+ functions.find(!library.hasFunction(_)).foreach { function =>
+ throw error(name, action, s"The library does not implement ${function.fieldName}.")
+ }
+ }
+
+ 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
+ }
+ }
+}
+
+object NativeDataSource {
+ import NativeLibrary.Function._
+
+ // The functions that the operations of a data source require, besides create and close. The
+ // others are optional. They are checked on the driver, before the operation starts.
+ private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ)
+ private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET,
+ STREAM_PARTITIONS, CLOSE_STREAM_READER, READ)
+ private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER,
+ WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER)
+}
+
+/**
+ * Plans a batch scan of a native data source, on the driver, with a reader of the library. The
+ * reader is released once the scan is planned, or when the planner becomes unreachable if the
+ * scan is never planned: the executors read the partitions with the state that it serialized.
+ *
+ * @param executorLocation where the executors find the library
+ */
+class NativeBatchPlanner(
+ library: NativeLibrary,
+ executorLocation: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ handle: NativeHandle)
+ extends ColumnarBatchPlanner {
+ import NativeCalls._
+
+ private var readSchema: StructType = schema
+
+ 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 used = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ // How the reader uses each predicate: whether it uses it completely or only to skip data
+ // does not matter, since Spark evaluates all the predicates again after reading.
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ }.map { uses =>
+ encoded.map(_._2).zip(uses).collect {
+ case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i
+ }.toSet
+ }.getOrElse(Set.empty[Int])
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = {
+ try {
+ val partitions = call(name, "plan a scan of")(library.partitions(handle.get))
+ val readerState = bytesOrEmpty(
+ call(name, "plan a scan of")(optional(library.serializeReader(handle.get))))
+ ColumnarPlannedBatch(
+ NativeInputPartition.from(partitions),
+ new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema))
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+}
+
+/** Reads the partitions of a planned batch scan of a native data source, on the executors. */
+class NativePartitionReaderFactory(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ readSchema: StructType)
+ extends ColumnarOnlyPartitionReaderFactory {
+ override def createColumnarReader(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("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get)))
+ }
+
+ override def latestOffset(): String = {
+ checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get)))
+ }
+
+ /**
+ * Checks an offset returned by the given function of the library: it must be JSON on a single
+ * line, because the offset log of the checkpoint stores each offset as a line.
+ */
+ private def checkOffset(function: String, offset: String): String = {
+ if (offset.exists(c => c == '\n' || c == '\r')) {
+ throw error(name, "plan a scan of", s"$function returned an offset with a line break, " +
+ "but an offset must be JSON on a single line.")
+ }
+ val valid = try {
+ val json = NativeDataSourceStreamReader.mapper.readTree(offset)
+ json != null && !json.isMissingNode
+ } catch {
+ case _: JsonProcessingException => false
+ }
+ if (!valid) {
+ throw error(name, "plan a scan of", s"$function returned an offset that is not valid " +
+ s"JSON: $offset")
+ }
+ offset
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] = {
+ 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()
+ // The task that reads the partition, which the library may use until the stream is released.
+ private var task: NativeLibrary.Task = _
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ 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)
+ task = NativeTask.create()
+ call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress()))
+ }
+
+ private def importSchema(): Schema = {
+ // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash
+ // the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream, or a callback of " +
+ "the stream is NULL.")
+ }
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ 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))
Review Comment:
I felt this may cause bug since reusing a fixed `VectorSchemaRoot` can cause data corruption when the output of `NativeDataSource` is accumulated by the downstream operators ([example](https://github.com/apache/spark/blob/ca6524e34ce487bb61ab1fda4b440a86e2c2177b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala#L191-L192)), since the new data will replace the content of the old vectors in early `ColumnarBatch` instances.
It might be worth trying the root-free `Data.importIntoVectorSchemaRoot` APIs that creates fresh roots internally?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6184757389
Thanks for fixing the data writer so quickly, with a test, and for running the native suites on Java 25 in the CI of pull requests. The function table also makes the header easier to build against than JNI: the compiler checks every function a library provides.
Speaking as one of the Lance maintainers, I'd like to share how this looks from lance-spark, since Lance is one of the motivating examples in the description, and hear your thoughts.
For reads and writes, lance-spark already goes through the Rust core of Lance: the Lance Java bindings fill an `ArrowArrayStream` at an address that the JVM allocates (`LanceFragmentScanner.exportArrowStream`), which is the same exchange as the `read` function of `SparkDsLibrary`. So a native data source would not remove a layer or a copy for Lance.
Most of what users get from lance-spark is beyond reads and writes: the `VECTOR_SEARCH`, `SEARCH` and `HYBRID_SEARCH` table functions, `CREATE INDEX` and the other index commands, `UPDATE`, `DELETE` and `MERGE INTO` through `_rowid` and `_rowaddr`, fragment-aware joins, `count(*)` and top-N pushdown, and branches and tags. Version 1 of the interface covers reads and writes with filter, column and limit pushdown, so a native Lance data source would give users less than lance-spark does today.
This is also where our cost of supporting several Spark releases is. lance-spark has a module for each Spark version from 3.4 to 4.2, but the read and write code is shared. What differs between versions is mainly the SQL extensions (parser and optimizer rules) and the row-level operations. For example, Spark 4.2 plans an insert-only `MERGE INTO` as `InsertOnlyMerge` instead of `AppendData`, and we needed a new rule so that blob columns are still copied. A stable ABI would not help with that part.
Where native data sources look like a clear win to me is a source that only has a native client and only needs reads and writes, such as an in-house storage system in C++ or a vendor's C SDK. Could the description and the docs lead with that kind of source, and with a concrete one if this work started from one? Your reply to @tdcmeehan lists what could be appended to the table later, such as statistics, aggregate and top-N pushdown, and overwrite by filter. For a source like Lance, table functions and row-level operations matter as much. Do you see them in scope for the table later, or out of scope on purpose? Saying so in the docs would help library authors tell early whether it fits their source.
None of this blocks the implementation; it is only about how the feature is positioned.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6207026653
@dongjoon-hyun @sarutak @gerashegalov You approved #59313, which this PR replaces, so the approvals there do not carry over here. Could you take another look? Since #59313, which this PR includes up to e3ff80c, the changes are:
- The interface is a C header, `spark_datasource.h`, called through the FFM API on Java 22 or later instead of JNI (b9ad01ac93be), and shipped in the `include` directory of a distribution (4f032d5f1b9c).
- Following the review here, a library exports only `spark_ds_init`, which returns a function table, with size-prefixed structs from Spark and the capabilities of each data source (d2df236bd234, b36ec5e73e27).
- A data writer that committed is no longer aborted, and the native data source suites run in the `Java 25 build with Maven` job of pull requests (e5352e191c14). The docs say which sources native data sources fit (51af65c63a93).
- Batch scans are planned on the driver, with a serializable reader factory for the executors, and the Data Source V2 table reports the batch writes that a data source supports, so an unsupported write fails while the query is analyzed (534dfca15723).
The "Differences from #59313" table at the top of the description summarizes the interface. CI is green on 534dfca15723.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #59424:
URL: https://github.com/apache/spark/pull/59424#discussion_r4268982738
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -0,0 +1,697 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.v2.ffi;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+import java.lang.reflect.Array;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.apache.spark.unsafe.Platform;
+
+/**
+ * A native data source library, which implements the C interface in
+ * {@code org/apache/spark/sql/datasource/spark_datasource.h}, called through the Foreign Function
+ * and Memory (FFM) API of Java 22 and later.
+ * <p>
+ * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached
+ * through method handles, looked up once when the first library is loaded. Each function of a
+ * library is bound to a method handle that takes and returns primitive values only, with
+ * pointers as {@code long} addresses, so the calls themselves use no reflection. The memory that
+ * Spark passes to the library is allocated with {@link Platform}.
+ * <p>
+ * Each library is looked up separately, so all the libraries can export the same functions.
+ * Calling a function that the library does not export throws {@link MissingFunctionException}. A
+ * function that fails throws a {@link RuntimeException} with the message of
+ * {@code spark_ds_last_error}.
+ */
+final class NativeLibrary {
+
+ /** The version of the interface that Spark supports, {@code SPARK_DS_ABI_VERSION}. */
+ static final int ABI_VERSION = 1;
+
+ /** The first version of Java with the final FFM API. */
+ static final int MINIMUM_JAVA_VERSION = 22;
+
+ /** Whether this JVM can load native data source libraries. */
+ static boolean isSupported() {
+ return Runtime.version().feature() >= MINIMUM_JAVA_VERSION;
+ }
+
+ // The C types of the parameters and the results of the functions.
+ private enum C { VOID, BOOL, INT32, INT64, POINTER }
+
+ private final MethodHandle lastError;
+ private final MethodHandle abiVersion;
+ private final MethodHandle create;
+ private final MethodHandle schema;
+ private final MethodHandle createReader;
+ private final MethodHandle createStreamReader;
+ private final MethodHandle createWriter;
+ private final MethodHandle createStreamWriter;
+ private final MethodHandle close;
+ private final MethodHandle pushPredicates;
+ private final MethodHandle pushLimit;
+ private final MethodHandle pruneColumns;
+ private final MethodHandle partitions;
+ private final MethodHandle serializeReader;
+ private final MethodHandle closeReader;
+ private final MethodHandle initialOffset;
+ private final MethodHandle latestOffset;
+ private final MethodHandle streamPartitions;
+ private final MethodHandle serializeStreamReader;
+ private final MethodHandle commitOffset;
+ private final MethodHandle closeStreamReader;
+ private final MethodHandle read;
+ private final MethodHandle serializeWriter;
+ private final MethodHandle commit;
+ private final MethodHandle abort;
+ private final MethodHandle closeWriter;
+ private final MethodHandle createDataWriter;
+ private final MethodHandle write;
+ private final MethodHandle commitDataWriter;
+ private final MethodHandle abortDataWriter;
+
+ private NativeLibrary(Object lookup) throws Throwable {
+ C i32 = C.INT32;
+ C i64 = C.INT64;
+ C ptr = C.POINTER;
+ lastError = bind(lookup, "spark_ds_last_error", ptr);
+ abiVersion = bind(lookup, "spark_ds_abi_version", i32);
+ create = bind(lookup, "spark_ds_create", i32, ptr, ptr, ptr, i64, ptr);
+ schema = bind(lookup, "spark_ds_schema", i32, ptr, ptr);
+ createReader = bind(lookup, "spark_ds_create_reader", i32, ptr, ptr, ptr);
+ createStreamReader = bind(lookup, "spark_ds_create_stream_reader", i32, ptr, ptr, ptr);
+ createWriter = bind(lookup, "spark_ds_create_writer", i32, ptr, ptr, C.BOOL, ptr);
+ createStreamWriter =
+ bind(lookup, "spark_ds_create_stream_writer", i32, ptr, ptr, C.BOOL, ptr, ptr);
+ close = bind(lookup, "spark_ds_close", C.VOID, ptr);
+ pushPredicates = bind(lookup, "spark_ds_push_predicates", i32, ptr, ptr, i64, ptr);
+ pushLimit = bind(lookup, "spark_ds_push_limit", i32, ptr, i32, ptr);
+ pruneColumns = bind(lookup, "spark_ds_prune_columns", i32, ptr, ptr, i64, ptr);
+ partitions = bind(lookup, "spark_ds_partitions", i32, ptr, ptr, ptr);
+ serializeReader = bind(lookup, "spark_ds_serialize_reader", i32, ptr, ptr, ptr);
+ closeReader = bind(lookup, "spark_ds_close_reader", C.VOID, ptr);
+ initialOffset = bind(lookup, "spark_ds_initial_offset", i32, ptr, ptr, ptr);
+ latestOffset = bind(lookup, "spark_ds_latest_offset", i32, ptr, ptr, ptr);
+ streamPartitions =
+ bind(lookup, "spark_ds_stream_partitions", i32, ptr, ptr, ptr, ptr, ptr);
+ serializeStreamReader =
+ bind(lookup, "spark_ds_serialize_stream_reader", i32, ptr, ptr, ptr);
+ commitOffset = bind(lookup, "spark_ds_commit_offset", i32, ptr, ptr);
+ closeStreamReader = bind(lookup, "spark_ds_close_stream_reader", C.VOID, ptr);
+ read = bind(lookup, "spark_ds_read", i32, ptr, i64, ptr, i64, ptr);
+ serializeWriter = bind(lookup, "spark_ds_serialize_writer", i32, ptr, ptr, ptr);
+ commit = bind(lookup, "spark_ds_commit", i32, ptr, i64, ptr, ptr, i64);
+ abort = bind(lookup, "spark_ds_abort", i32, ptr, i64, ptr, ptr, i64);
+ closeWriter = bind(lookup, "spark_ds_close_writer", C.VOID, ptr);
+ createDataWriter =
+ bind(lookup, "spark_ds_create_data_writer", i32, ptr, i64, i32, i64, i64, ptr);
+ write = bind(lookup, "spark_ds_write", i32, ptr, ptr, ptr);
+ commitDataWriter = bind(lookup, "spark_ds_commit_data_writer", i32, ptr, ptr, ptr);
+ abortDataWriter = bind(lookup, "spark_ds_abort_data_writer", i32, ptr);
+ }
+
+ /**
+ * Loads the library at the given path. The library stays loaded until the JVM exits.
+ *
+ * @throws LoadException if the library cannot be loaded
+ */
+ static NativeLibrary load(String path) throws Throwable {
+ if (!isSupported()) {
+ throw new UnsupportedOperationException(
+ "Native data sources require Java " + MINIMUM_JAVA_VERSION + " or later.");
+ }
+ Object lookup;
+ try {
+ lookup = Ffm.LIBRARY_LOOKUP.invoke(Path.of(path), Ffm.GLOBAL_ARENA);
+ } catch (IllegalArgumentException | IllegalCallerException e) {
+ // The library is not a valid library, or native access is not enabled for Spark.
+ throw new LoadException(e.getMessage(), e);
+ } catch (LinkageError e) {
+ // The FFM API could not be initialized, see Ffm: report why, such as an
+ // IllegalCallerException when native access is not enabled for Spark.
+ throw new LoadException(
+ "Cannot use the Foreign Function and Memory API: " + rootCause(e), e);
+ }
+ return new NativeLibrary(lookup);
+ }
+
+ /**
+ * Binds a function of the library to a method handle whose parameters and result have the
+ * Java types of the given C types, with {@code long} for pointers. Returns null if the library
+ * does not export the function.
+ */
+ private static MethodHandle bind(Object lookup, String name, C result, C... parameters)
+ throws Throwable {
+ Optional<?> symbol = (Optional<?>) Ffm.FIND.invoke(lookup, name);
+ if (symbol.isEmpty()) {
+ return null;
+ }
+ Object layouts = Ffm.layouts(parameters);
+ Object descriptor = result == C.VOID
+ ? Ffm.FUNCTION_OF_VOID.invoke(layouts)
+ : Ffm.FUNCTION_OF.invoke(Ffm.layout(result), layouts);
+ MethodHandle handle = (MethodHandle) Ffm.DOWNCALL.invoke(
+ Ffm.LINKER, symbol.get(), descriptor, Ffm.NO_OPTIONS);
+ // Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses.
+ for (int i = 0; i < parameters.length; i++) {
+ if (parameters[i] == C.POINTER) {
+ handle = MethodHandles.filterArguments(handle, i, Ffm.OF_ADDRESS);
+ }
+ }
+ if (result == C.POINTER) {
+ handle = MethodHandles.filterReturnValue(handle, Ffm.ADDRESS_OF);
+ }
+ return handle;
+ }
+
+ /** The method handles of the FFM API, initialized when the first library is loaded. */
+ private static final class Ffm {
+ static final Class<?> LAYOUT_CLASS;
+ static final Object LINKER;
+ static final Object GLOBAL_ARENA;
+ static final Object NO_OPTIONS;
+ static final MethodHandle LIBRARY_LOOKUP;
+ static final MethodHandle FIND;
+ static final MethodHandle DOWNCALL;
+ static final MethodHandle FUNCTION_OF;
+ static final MethodHandle FUNCTION_OF_VOID;
+ // (long) -> MemorySegment and (MemorySegment) -> long.
+ static final MethodHandle OF_ADDRESS;
+ static final MethodHandle ADDRESS_OF;
+ static final Object BOOL;
+ static final Object INT32;
+ static final Object INT64;
+ static final Object POINTER;
+ // The address of the C function SparkDsSink, which calls NativeLibrary.sink.
+ static final long SINK;
+
+ static {
+ try {
+ // A full-privilege lookup, so that the restricted methods of the FFM API see this class
+ // as their caller, in the unnamed module that Spark enables native access for.
+ MethodHandles.Lookup lookup = MethodHandles.lookup();
+ Class<?> linker = Class.forName("java.lang.foreign.Linker");
+ Class<?> option = Class.forName("java.lang.foreign.Linker$Option");
+ Class<?> symbolLookup = Class.forName("java.lang.foreign.SymbolLookup");
+ Class<?> arena = Class.forName("java.lang.foreign.Arena");
+ Class<?> segment = Class.forName("java.lang.foreign.MemorySegment");
+ Class<?> layout = Class.forName("java.lang.foreign.MemoryLayout");
+ Class<?> layoutArray = Array.newInstance(layout, 0).getClass();
+ Class<?> valueLayout = Class.forName("java.lang.foreign.ValueLayout");
+ Class<?> descriptor = Class.forName("java.lang.foreign.FunctionDescriptor");
+ Class<?> options = Array.newInstance(option, 0).getClass();
+
+ LAYOUT_CLASS = layout;
+ LINKER = lookup.findStatic(linker, "nativeLinker", MethodType.methodType(linker))
+ .invoke();
+ GLOBAL_ARENA = lookup.findStatic(arena, "global", MethodType.methodType(arena)).invoke();
+ NO_OPTIONS = Array.newInstance(option, 0);
+ LIBRARY_LOOKUP = lookup.findStatic(symbolLookup, "libraryLookup",
+ MethodType.methodType(symbolLookup, Path.class, arena));
+ FIND = lookup.findVirtual(symbolLookup, "find",
+ MethodType.methodType(Optional.class, String.class));
+ DOWNCALL = lookup.findVirtual(linker, "downcallHandle",
+ MethodType.methodType(MethodHandle.class, segment, descriptor, options))
+ .asFixedArity();
+ FUNCTION_OF = lookup.findStatic(descriptor, "of",
+ MethodType.methodType(descriptor, layout, layoutArray)).asFixedArity();
+ FUNCTION_OF_VOID = lookup.findStatic(descriptor, "ofVoid",
+ MethodType.methodType(descriptor, layoutArray)).asFixedArity();
+ OF_ADDRESS = lookup.findStatic(segment, "ofAddress",
+ MethodType.methodType(segment, long.class));
+ ADDRESS_OF = lookup.findVirtual(segment, "address", MethodType.methodType(long.class));
+ BOOL = valueLayout.getField("JAVA_BOOLEAN").get(null);
+ INT32 = valueLayout.getField("JAVA_INT").get(null);
+ INT64 = valueLayout.getField("JAVA_LONG").get(null);
+ POINTER = valueLayout.getField("ADDRESS").get(null);
+
+ // The upcall stub of the sink, which lives as long as the JVM.
+ MethodHandle target = lookup.findStatic(NativeLibrary.class, "sink",
+ MethodType.methodType(void.class, long.class, long.class, long.class));
+ target = MethodHandles.filterArguments(target, 0, ADDRESS_OF, ADDRESS_OF);
+ Object sinkDescriptor =
+ FUNCTION_OF_VOID.invoke(layouts(C.POINTER, C.POINTER, C.INT64));
+ Object stub = lookup.findVirtual(linker, "upcallStub",
+ MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options))
+ .asFixedArity()
+ .invoke(LINKER, target, sinkDescriptor, GLOBAL_ARENA, NO_OPTIONS);
+ SINK = (long) ADDRESS_OF.invoke(stub);
+ } catch (Throwable e) {
+ // The message names the cause, because the JVM only keeps this message for the later
+ // uses of the class.
+ throw new IllegalStateException("Cannot initialize the FFM API: " + rootCause(e), e);
+ }
+ }
+
+ /** Returns a MemoryLayout[] with the layouts of the given types. */
+ static Object layouts(C... types) {
+ Object layouts = Array.newInstance(LAYOUT_CLASS, types.length);
+ for (int i = 0; i < types.length; i++) {
+ Array.set(layouts, i, layout(types[i]));
+ }
+ return layouts;
+ }
+
+ static Object layout(C type) {
+ return switch (type) {
+ case BOOL -> BOOL;
+ case INT32 -> INT32;
+ case INT64 -> INT64;
+ case POINTER -> POINTER;
+ case VOID -> throw new IllegalArgumentException("void has no layout");
+ };
+ }
+ }
+
+ private static Throwable rootCause(Throwable e) {
+ Throwable cause = e;
+ while (cause.getCause() != null && cause.getCause() != cause) {
+ cause = cause.getCause();
+ }
+ return cause;
+ }
+
+ /** Thrown when a library cannot be loaded. */
+ static final class LoadException extends RuntimeException {
+ LoadException(String message, Throwable cause) {
+ super(message, cause);
+ }
+ }
+
+ /** Thrown when a function is called that the library does not export, with its name. */
+ static final class MissingFunctionException extends RuntimeException {
+ MissingFunctionException(String name) {
+ super(name);
+ }
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Calls
+ // ---------------------------------------------------------------------------------------------
+
+ private static MethodHandle require(MethodHandle function, String name) {
+ if (function == null) {
+ throw new MissingFunctionException(name);
+ }
+ return function;
+ }
+
+ /** Throws the error that the library reported if a function failed. */
+ private void check(int status, String name) throws Throwable {
+ if (status == 0) {
+ return;
+ }
+ long address = lastError == null ? 0 : (long) lastError.invokeExact();
+ String message = address == 0 ? "" : new String(readCString(address), StandardCharsets.UTF_8);
+ throw new RuntimeException(message.isEmpty()
+ ? name + " failed with the error code " + status + "."
+ : message);
+ }
+
+ private static long readHandle(long out, String name) {
+ long handle = Platform.getLong(null, out);
+ if (handle == 0) {
+ throw new IllegalStateException(name + " returned a NULL handle.");
+ }
+ return handle;
+ }
+
+ int abiVersion() throws Throwable {
+ return (int) require(abiVersion, "spark_ds_abi_version").invokeExact();
+ }
+
+ long createDataSource(String name, String[] optionKeys, String[] optionValues)
+ throws Throwable {
+ MethodHandle f = require(create, "spark_ds_create");
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(memory.string(name), memory.strings(optionKeys),
+ memory.strings(optionValues), (long) optionKeys.length, out), "spark_ds_create");
+ return readHandle(out, "spark_ds_create");
+ }
+ }
+
+ void schema(long dataSource, long schemaAddress) throws Throwable {
+ check((int) require(schema, "spark_ds_schema").invokeExact(dataSource, schemaAddress),
+ "spark_ds_schema");
+ }
+
+ long createReader(long dataSource, long schemaAddress) throws Throwable {
+ return createWithSchema(createReader, "spark_ds_create_reader", dataSource, schemaAddress);
+ }
+
+ long createStreamReader(long dataSource, long schemaAddress) throws Throwable {
+ return createWithSchema(
+ createStreamReader, "spark_ds_create_stream_reader", dataSource, schemaAddress);
+ }
+
+ private long createWithSchema(
+ MethodHandle function, String name, long dataSource, long schemaAddress) throws Throwable {
+ MethodHandle f = require(function, name);
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(dataSource, schemaAddress, out), name);
+ return readHandle(out, name);
+ }
+ }
+
+ long createWriter(long dataSource, long schemaAddress, boolean overwrite) throws Throwable {
+ MethodHandle f = require(createWriter, "spark_ds_create_writer");
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(dataSource, schemaAddress, overwrite, out),
+ "spark_ds_create_writer");
+ return readHandle(out, "spark_ds_create_writer");
+ }
+ }
+
+ long createStreamWriter(long dataSource, long schemaAddress, boolean overwrite, String queryId)
+ throws Throwable {
+ MethodHandle f = require(createStreamWriter, "spark_ds_create_stream_writer");
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(dataSource, schemaAddress, overwrite, memory.string(queryId), out),
+ "spark_ds_create_stream_writer");
+ return readHandle(out, "spark_ds_create_stream_writer");
+ }
+ }
+
+ void closeDataSource(long dataSource) throws Throwable {
+ require(close, "spark_ds_close").invokeExact(dataSource);
+ }
+
+ boolean[] pushPredicates(long reader, String[] predicates) throws Throwable {
+ MethodHandle f = require(pushPredicates, "spark_ds_push_predicates");
+ try (Memory memory = new Memory()) {
+ long evaluated = memory.zeroed(Math.max(1, predicates.length));
+ check((int) f.invokeExact(reader, memory.strings(predicates), (long) predicates.length,
+ evaluated), "spark_ds_push_predicates");
+ boolean[] result = new boolean[predicates.length];
+ for (int i = 0; i < result.length; i++) {
+ result[i] = Platform.getByte(null, evaluated + i) != 0;
+ }
+ return result;
+ }
+ }
+
+ boolean pushLimit(long reader, int limit) throws Throwable {
+ MethodHandle f = require(pushLimit, "spark_ds_push_limit");
+ try (Memory memory = new Memory()) {
+ long applied = memory.zeroed(1);
+ check((int) f.invokeExact(reader, limit, applied), "spark_ds_push_limit");
+ return Platform.getByte(null, applied) != 0;
+ }
+ }
+
+ boolean pruneColumns(long reader, String[] columnNames) throws Throwable {
+ MethodHandle f = require(pruneColumns, "spark_ds_prune_columns");
+ try (Memory memory = new Memory()) {
+ long pruned = memory.zeroed(1);
+ check((int) f.invokeExact(reader, memory.strings(columnNames), (long) columnNames.length,
+ pruned), "spark_ds_prune_columns");
+ return Platform.getByte(null, pruned) != 0;
+ }
+ }
+
+ byte[][] partitions(long reader) throws Throwable {
+ MethodHandle f = require(partitions, "spark_ds_partitions");
+ try (Sink sink = new Sink()) {
+ check((int) f.invokeExact(reader, Ffm.SINK, sink.id), "spark_ds_partitions");
+ return sink.all();
+ }
+ }
+
+ byte[] serializeReader(long reader) throws Throwable {
+ return serialize(serializeReader, "spark_ds_serialize_reader", reader);
+ }
+
+ void closeReader(long reader) throws Throwable {
+ require(closeReader, "spark_ds_close_reader").invokeExact(reader);
+ }
+
+ String initialOffset(long streamReader) throws Throwable {
+ return offset(initialOffset, "spark_ds_initial_offset", streamReader);
+ }
+
+ String latestOffset(long streamReader) throws Throwable {
+ return offset(latestOffset, "spark_ds_latest_offset", streamReader);
+ }
+
+ private String offset(MethodHandle function, String name, long streamReader)
+ throws Throwable {
+ byte[] offset = serialize(function, name, streamReader);
+ if (offset == null) {
+ throw new IllegalStateException(name + " returned no offset.");
+ }
+ return new String(offset, StandardCharsets.UTF_8);
+ }
+
+ byte[][] streamPartitions(long streamReader, String start, String end) throws Throwable {
+ MethodHandle f = require(streamPartitions, "spark_ds_stream_partitions");
+ try (Memory memory = new Memory(); Sink sink = new Sink()) {
+ check((int) f.invokeExact(streamReader, memory.string(start), memory.string(end),
+ Ffm.SINK, sink.id), "spark_ds_stream_partitions");
+ return sink.all();
+ }
+ }
+
+ byte[] serializeStreamReader(long streamReader) throws Throwable {
+ return serialize(serializeStreamReader, "spark_ds_serialize_stream_reader", streamReader);
+ }
+
+ void commitOffset(long streamReader, String end) throws Throwable {
+ MethodHandle f = require(commitOffset, "spark_ds_commit_offset");
+ try (Memory memory = new Memory()) {
+ check((int) f.invokeExact(streamReader, memory.string(end)), "spark_ds_commit_offset");
+ }
+ }
+
+ void closeStreamReader(long streamReader) throws Throwable {
+ require(closeStreamReader, "spark_ds_close_stream_reader").invokeExact(streamReader);
+ }
+
+ void read(byte[] readerState, byte[] partition, long streamAddress) throws Throwable {
+ MethodHandle f = require(read, "spark_ds_read");
+ try (Memory memory = new Memory()) {
+ check((int) f.invokeExact(memory.bytes(readerState), (long) readerState.length,
+ memory.bytes(partition), (long) partition.length, streamAddress), "spark_ds_read");
+ }
+ }
+
+ byte[] serializeWriter(long writer) throws Throwable {
+ return serialize(serializeWriter, "spark_ds_serialize_writer", writer);
+ }
+
+ void commit(long writer, long epochId, byte[][] messages) throws Throwable {
+ finishWrite(commit, "spark_ds_commit", writer, epochId, messages);
+ }
+
+ void abort(long writer, long epochId, byte[][] messages) throws Throwable {
+ finishWrite(abort, "spark_ds_abort", writer, epochId, messages);
+ }
+
+ private void finishWrite(
+ MethodHandle function,
+ String name,
+ long writer,
+ long epochId,
+ byte[][] messages) throws Throwable {
+ MethodHandle f = require(function, name);
+ try (Memory memory = new Memory()) {
+ long pointers = memory.zeroed(8L * Math.max(1, messages.length));
+ long lengths = memory.zeroed(8L * Math.max(1, messages.length));
+ for (int i = 0; i < messages.length; i++) {
+ if (messages[i] != null) {
+ Platform.putLong(null, pointers + 8L * i, memory.bytes(messages[i]));
+ Platform.putLong(null, lengths + 8L * i, messages[i].length);
+ }
+ }
+ check((int) f.invokeExact(writer, epochId, pointers, lengths, (long) messages.length),
+ name);
+ }
+ }
+
+ void closeWriter(long writer) throws Throwable {
+ require(closeWriter, "spark_ds_close_writer").invokeExact(writer);
+ }
+
+ long createDataWriter(byte[] writerState, int partitionId, long taskId, long epochId)
+ throws Throwable {
+ MethodHandle f = require(createDataWriter, "spark_ds_create_data_writer");
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(memory.bytes(writerState), (long) writerState.length,
+ partitionId, taskId, epochId, out), "spark_ds_create_data_writer");
+ return readHandle(out, "spark_ds_create_data_writer");
+ }
+ }
+
+ void write(long dataWriter, long arrayAddress, long schemaAddress) throws Throwable {
+ check((int) require(write, "spark_ds_write").invokeExact(
+ dataWriter, arrayAddress, schemaAddress), "spark_ds_write");
+ }
+
+ byte[] commitDataWriter(long dataWriter) throws Throwable {
+ byte[] message = serialize(commitDataWriter, "spark_ds_commit_data_writer", dataWriter);
Review Comment:
When `spark_ds_commit_data_writer` returns 0, the library has released the handle, but `serialize` still checks the sink afterwards and throws, for example if `set_message` was called twice. `NativeDataWriter.commit` only clears `handle` once `commitDataWriter` returns without throwing, so the task's catch block passes the released handle to `spark_ds_abort_data_writer`. A library that follows the header then gets a use-after-free or a double free, which can crash the executor. On the PR head, with the test library calling `set_message` twice, the write fails with `returned 2 values instead of one.`, and abort still gets the released handle.
Could `handle` be cleared once the call returns 0, before the sink is checked? A test for it would help too. This is the other side of the fix for dongjoon-hyun's comment on #59313, which clears the handle only after `commitDataWriter` returns: a commit that failed in the library must still be aborted, so the handle can only be cleared once the library has returned 0.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #59424:
URL: https://github.com/apache/spark/pull/59424#discussion_r4299568926
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -0,0 +1,688 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*
+ * spark_datasource.h: the binary interface between Spark and data sources implemented in native
+ * code, such as Rust or C++. Since Spark 4.4.0.
+ *
+ * A native data source is a shared library that exports a single C function, spark_ds_init,
+ * which returns the table of the functions that implement the data source, a SparkDsLibrary.
+ * Spark calls them with the Foreign Function and Memory API of Java, so native data sources
+ * require Java 22 or later. Each library has its own table, so any number of libraries can be
+ * loaded side by side. Spark adapts a native data source to Data Source V2, with batch and
+ * micro-batch streaming reads, batch and streaming writes, and predicate, column and limit
+ * pushdown.
+ *
+ * This header is self-contained: it only depends on the C standard library, and defines the
+ * structs of the Arrow C data interface unless they are already defined, as by Arrow C++. A library
+ * can copy it into its source tree, or use the copy in the include directory of a Spark
+ * distribution. A library built against it works with all the releases of Spark that support its
+ * SPARK_DS_ABI_VERSION.
+ *
+ * Spark finds the library of a data source automatically when it is installed like other native
+ * libraries and named after the data source: the library of my_source is
+ * spark_datasource_my_source, such as libspark_datasource_my_source.so in /usr/local/lib on Linux.
+ * Spark looks for it in the directories of java.library.path, which include LD_LIBRARY_PATH on
+ * Linux, and in the lib directory of each installation prefix in PATH. A library can also be
+ * distributed in a native data source package: a zip file with the extension .sparkpkg that
+ * contains a manifest named spark-native-datasource.json and the library built for one or more
+ * platforms, added to a session with spark.addArtifact or under the paths in
+ * spark.sql.dataSource.native.paths. See the "Native Data Sources" page of the Spark SQL guide for
+ * details and examples.
+ *
+ * Lifecycle
+ * ---------
+ * Spark loads a library once in each JVM, and calls spark_ds_init before any other function. It
+ * creates a data source with create on the driver for each schema inference, scan and write, and
+ * closes it right after. create also returns the capabilities of the data source: Spark reports
+ * the operations that a data source does not declare as unsupported, without calling the library.
+ * Spark also creates a data source while it analyzes a query, to reject a batch write that the
+ * data source does not support before it computes the data to write. That data source has the
+ * options of the table, those of CREATE TABLE for a table in a catalog, which the options of a
+ * write can override, so Spark checks the capabilities of each operation again.
+ *
+ * - Batch reads. On the driver, Spark creates a reader with create_reader, pushes operations down
+ * by calling, each at most once and in this order, push_predicates, push_limit and
+ * prune_columns, plans the partitions with partitions, serializes the state of the reader with
+ * serialize_reader, and closes it. Then it calls read on the executors, in one task for each
+ * partition. If Spark never plans the partitions, for example to only explain the plan of the
+ * query, it closes the reader once the reader is no longer referenced.
+ * - Streaming reads. On the driver, Spark creates a stream reader with create_stream_reader and
+ * keeps it until the streaming query stops. For each micro-batch, it asks for the latest_offset,
+ * plans the partitions with stream_partitions and serializes the state with
+ * serialize_stream_reader, then calls read on the executors. Offsets are defined by the library,
+ * as compact JSON on a single line: Spark stores each offset as a line of the offset log of the
+ * checkpoint, so it rejects an offset that is not valid JSON or contains a line break. When a
+ * streaming query restarts, from its checkpoint, Spark creates a new stream reader and may call
+ * stream_partitions again for the last range that it planned before, so the partitions of a
+ * range must read the same data each time.
+ * - Writes. On the driver, right before it starts the tasks of the write (or of a micro-batch),
+ * Spark creates a writer with create_writer, and serializes its state with serialize_writer.
+ * With adaptive query execution, Spark first runs the stages of the data to write that end with
+ * a shuffle, and creates no writer if they fail. The rest of the data, such as its last stage,
+ * is computed in the tasks of the write, so a writer may still be aborted because computing the
+ * data failed: do irreversible work, such as removing the existing data for an overwrite, in
+ * commit. On the executors, each task creates a data writer with create_data_writer, calls write
+ * for each batch of at most spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends
+ * once it holds spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls commit_data_writer,
+ * or abort_data_writer if the task or the commit fails, before it closes the data writer. Then
+ * the driver calls commit with the messages of the tasks if they all succeeded, or abort
+ * otherwise, including when commit fails, and closes the writer. Several attempts of a task may
+ * write the same partition, after a failure or with speculative execution, each with a data
+ * writer of its own. Spark lets the attempts of a partition commit one at a time, and passes to
+ * commit the message of one committed attempt for each partition, so the write consists of the
+ * data of these attempts only. It aborts the data writers of the other attempts, except on an
+ * executor that is lost, where a data writer is neither committed nor aborted, and may even have
+ * committed before its message reached the driver. A streaming write commits or aborts each
+ * micro-batch (epoch). Spark may commit the same epoch more than once, for example when a
+ * streaming query restarts after a failure, so committing an epoch that is already committed
+ * must have no effect, for exactly-once output. The ID of the query, which create_writer gets for
+ * a streaming write, stays the same when the query restarts from its checkpoint, so the ID of the
+ * query and the epoch ID identify the output of a micro-batch, also when other queries write to
+ * the same destination. A micro-batch streaming query creates a new stream writer for each
+ * micro-batch, and Spark closes the previous stream writer of the query when the query creates a
+ * new one. Spark is not told when a streaming query stops, so it closes the last stream writer of
+ * the query once the writer is no longer referenced, or when the query is restarted.
+ *
+ * Conventions
+ * -----------
+ * - Errors. The functions that return an int return 0 on success, or any other value on failure.
+ * They then set the message that last_error returns, and Spark reports it as a
+ * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail, and they may run after the query,
+ * from another thread: do the work that can fail a query, such as flushing data, in commit. A
+ * C++ exception or a Rust panic must not cross the interface: catch it and return an error
+ * instead.
+ * - Handles. create, create_reader, create_stream_reader, create_writer and create_data_writer
+ * return an opaque, non-NULL handle in their last argument, typically a pointer: Spark reports a
+ * NULL handle as an error. Spark releases each handle exactly once, with the close function of
+ * its kind: close, close_reader, close_stream_reader, close_writer or close_data_writer, and no
+ * other function releases a handle. A handle is not always released right after its last use: a
+ * batch reader whose partitions are never planned and the last stream writer of a query are
+ * released once they are no longer referenced (see Lifecycle above). A reader or writer handle
+ * must not depend on the data source handle it was created from, because Spark closes the data
+ * source handle right after creating it.
+ * - Memory. What Spark passes to a function, such as strings, byte arrays, SparkDsCreateInfo and
+ * SparkDsWriteInfo, is only valid during the call, except for SparkDsHost and
+ * SparkDsTaskContext, which document how long they live. Strings are UTF-8 and NUL-terminated.
+ * The JSON of the predicates is ASCII: it escapes the other characters, which any JSON parser
+ * decodes. The library returns strings and byte arrays by passing them to a SparkDsSink, which
+ * copies them.
+ * - State. Partitions, the state of readers and writers, and commit messages are byte arrays in a
+ * format defined by the library. Spark sends them to the executors and back.
+ * - Arrow data. Schemas and data are exchanged with the Arrow C data interface
+ * (https://arrow.apache.org/docs/format/CDataInterface.html), as structs allocated by Spark.
+ * For an output struct, the library initializes it and Spark takes ownership. For an input
+ * struct, the library may move it to take ownership; otherwise, Spark releases it after the
+ * call returns, so the library must not use it afterwards. The arrays that the library exports
+ * must have an offset of 0, at every level: Arrow Java ignores the offsets of the arrays it
+ * imports, so Spark rejects an array with another offset, such as a zero-copy slice. Copy a
+ * sliced array before exporting it. Spark also rejects an array with fewer values than its
+ * parent refers to, such as a column shorter than its batch. The buffers of an array must hold
+ * all of its values, as the C data interface requires: it does not carry the sizes of the
+ * buffers, so Spark cannot check them, and reads past the end of a buffer that is too short.
+ * - Schemas. Spark types map to Arrow types as in spark.sql.execution.arrow: for example, bigint
Review Comment:
Thanks, done in 74aa0dffc534. The header now lists the Spark types of version 1, with their Arrow types and formats: boolean, the integer types, float and double, decimal(p, s) as decimal128, string, binary, date, timestamp, timestamp_ntz, and array, map and struct. Spark fails an operation whose schema has another type with `NATIVE_DATA_SOURCE_ERROR` before it passes the schema to the library: the schema to read before `create_reader` and `create_stream_reader`, the data to write before `create_writer`, and a schema that `schema` returns. So variants, times and intervals do not reach a library, and a later version can add them for the data sources that declare a capability for them. Test: `reject the types and the strings that a library cannot get`, with a day-time interval in a read, and an array of variants in a write, which fails before `create_writer`.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -0,0 +1,982 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.v2.ffi;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+import java.lang.reflect.Array;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.BooleanSupplier;
+
+import org.apache.spark.internal.LogKeys;
+import org.apache.spark.internal.MDC;
+import org.apache.spark.internal.SparkLogger;
+import org.apache.spark.internal.SparkLoggerFactory;
+import org.apache.spark.unsafe.Platform;
+
+/**
+ * A native data source library, which implements the C interface in
+ * {@code org/apache/spark/sql/datasource/spark_datasource.h}, called through the Foreign Function
+ * and Memory (FFM) API of Java 22 and later.
+ * <p>
+ * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached
+ * through method handles, looked up once when the first library is loaded. The library exports
+ * {@code spark_ds_init}, which returns its table of functions, {@code SparkDsLibrary}. Each
+ * function of the table is bound to a method handle that takes and returns primitive values only,
+ * with pointers as {@code long} addresses, so the calls themselves use no reflection. The memory
+ * that Spark passes to the library is allocated with {@link Platform}, at the offsets of the
+ * structs of the header, which {@link #layout()} lists for the tests.
+ * <p>
+ * Calling a function that the library leaves NULL in its table throws
+ * {@link MissingFunctionException}. A function that fails throws a {@link RuntimeException} with
+ * the message of {@code last_error}.
+ */
+final class NativeLibrary {
+
+ private static final SparkLogger LOGGER = SparkLoggerFactory.getLogger(NativeLibrary.class);
+
+ /** The version of the interface that Spark supports, {@code SPARK_DS_ABI_VERSION}. */
+ static final int ABI_VERSION = 1;
+
+ /** The first version of Java with the final FFM API. */
+ static final int MINIMUM_JAVA_VERSION = 22;
+
+ /** Whether this JVM can load native data source libraries. */
+ static boolean isSupported() {
+ return Runtime.version().feature() >= MINIMUM_JAVA_VERSION;
+ }
+
+ // The capabilities of a data source, SPARK_DS_BATCH_READ and the others.
+ static final long BATCH_READ = 1;
+ static final long MICRO_BATCH_READ = 2;
+ static final long BATCH_WRITE = 4;
+ static final long STREAMING_WRITE = 8;
+ static final long TRUNCATE = 16;
+
+ // The modes of a write, SPARK_DS_WRITE_APPEND and SPARK_DS_WRITE_TRUNCATE.
+ static final int WRITE_APPEND = 0;
+ static final int WRITE_TRUNCATE = 1;
+
+ /** A predicate that the reader does not use, SPARK_DS_PREDICATE_UNUSED. */
+ static final int PREDICATE_UNUSED = 0;
+
+ // The C types of the parameters and the results of the functions.
+ private enum C { VOID, BOOL, I32, I64, PTR }
+
+ /** The functions of the table {@code SparkDsLibrary}, in their order in the table. */
+ enum Function {
+ LAST_ERROR("last_error", C.PTR),
+ CREATE("create", C.I32, C.PTR, C.PTR, C.PTR),
+ SCHEMA("schema", C.I32, C.PTR, C.PTR),
+ CREATE_READER("create_reader", C.I32, C.PTR, C.PTR, C.PTR),
+ CREATE_STREAM_READER("create_stream_reader", C.I32, C.PTR, C.PTR, C.PTR),
+ CREATE_WRITER("create_writer", C.I32, C.PTR, C.PTR, C.PTR, C.PTR),
+ CLOSE("close", C.VOID, C.PTR),
+ PUSH_PREDICATES("push_predicates", C.I32, C.PTR, C.PTR, C.I64, C.PTR),
+ PUSH_LIMIT("push_limit", C.I32, C.PTR, C.I32, C.PTR),
+ PRUNE_COLUMNS("prune_columns", C.I32, C.PTR, C.PTR, C.I64, C.PTR),
+ PARTITIONS("partitions", C.I32, C.PTR, C.PTR, C.PTR),
+ SERIALIZE_READER("serialize_reader", C.I32, C.PTR, C.PTR, C.PTR),
+ CLOSE_READER("close_reader", C.VOID, C.PTR),
+ INITIAL_OFFSET("initial_offset", C.I32, C.PTR, C.PTR, C.PTR),
+ LATEST_OFFSET("latest_offset", C.I32, C.PTR, C.PTR, C.PTR),
+ STREAM_PARTITIONS("stream_partitions", C.I32, C.PTR, C.PTR, C.PTR, C.PTR, C.PTR),
+ SERIALIZE_STREAM_READER("serialize_stream_reader", C.I32, C.PTR, C.PTR, C.PTR),
+ COMMIT_OFFSET("commit_offset", C.I32, C.PTR, C.PTR),
+ CLOSE_STREAM_READER("close_stream_reader", C.VOID, C.PTR),
+ READ("read", C.I32, C.PTR, C.PTR, C.I64, C.PTR, C.I64, C.PTR),
+ SERIALIZE_WRITER("serialize_writer", C.I32, C.PTR, C.PTR, C.PTR),
+ COMMIT("commit", C.I32, C.PTR, C.I64, C.PTR, C.PTR, C.I64),
+ ABORT("abort", C.I32, C.PTR, C.I64, C.PTR, C.PTR, C.I64),
+ CLOSE_WRITER("close_writer", C.VOID, C.PTR),
+ CREATE_DATA_WRITER("create_data_writer", C.I32, C.PTR, C.PTR, C.I64, C.I64, C.PTR),
+ WRITE("write", C.I32, C.PTR, C.PTR, C.PTR),
+ COMMIT_DATA_WRITER("commit_data_writer", C.I32, C.PTR, C.PTR, C.PTR),
+ ABORT_DATA_WRITER("abort_data_writer", C.I32, C.PTR),
+ CLOSE_DATA_WRITER("close_data_writer", C.VOID, C.PTR);
+
+ /** The name of the field of the function in the table. */
+ final String fieldName;
+ private final C result;
+ private final C[] parameters;
+
+ Function(String fieldName, C result, C... parameters) {
+ this.fieldName = fieldName;
+ this.result = result;
+ this.parameters = parameters;
+ }
+
+ /** The offset of the function in the table, after the size and the version of the table. */
+ long offset() {
+ return TABLE_FUNCTIONS + 8L * ordinal();
+ }
+ }
+
+ // The layouts of the structs of the header, on the 64-bit platforms that Spark supports. Each
+ // struct starts with its size.
+ private static final long TABLE_ABI_VERSION = 8;
+ private static final long TABLE_FUNCTIONS = 16;
+ private static final long HOST_SIZE = 32;
+ private static final long HOST_ABI_VERSION = 8;
+ private static final long HOST_SPARK_VERSION = 16;
+ private static final long HOST_LOG = 24;
+ private static final long CREATE_INFO_SIZE = 40;
+ private static final long CREATE_INFO_NAME = 8;
+ private static final long CREATE_INFO_OPTION_KEYS = 16;
+ private static final long CREATE_INFO_OPTION_VALUES = 24;
+ private static final long CREATE_INFO_NUM_OPTIONS = 32;
+ private static final long TASK_SIZE = 40;
+ private static final long TASK_STAGE_ID = 8;
+ private static final long TASK_STAGE_ATTEMPT_NUMBER = 12;
+ private static final long TASK_PARTITION_ID = 16;
+ private static final long TASK_ATTEMPT_NUMBER = 20;
+ private static final long TASK_ATTEMPT_ID = 24;
+ private static final long TASK_IS_INTERRUPTED = 32;
+ private static final long WRITE_INFO_SIZE = 24;
+ private static final long WRITE_INFO_MODE = 8;
+ private static final long WRITE_INFO_QUERY_ID = 16;
+
+ /**
+ * The offsets of the fields of the structs that Spark reads or writes, by {@code struct.field},
+ * and the sizes of the structs, by {@code struct}, for the tests to compare with the header.
+ */
+ static Map<String, Long> layout() {
+ Map<String, Long> layout = new LinkedHashMap<>();
+ layout.put("SparkDsLibrary", TABLE_FUNCTIONS + 8L * Function.values().length);
+ layout.put("SparkDsLibrary.size", 0L);
+ layout.put("SparkDsLibrary.abi_version", TABLE_ABI_VERSION);
+ for (Function function : Function.values()) {
+ layout.put("SparkDsLibrary." + function.fieldName, function.offset());
+ }
+ layout.put("SparkDsHost", HOST_SIZE);
+ layout.put("SparkDsHost.size", 0L);
+ layout.put("SparkDsHost.abi_version", HOST_ABI_VERSION);
+ layout.put("SparkDsHost.spark_version", HOST_SPARK_VERSION);
+ layout.put("SparkDsHost.log", HOST_LOG);
+ layout.put("SparkDsCreateInfo", CREATE_INFO_SIZE);
+ layout.put("SparkDsCreateInfo.size", 0L);
+ layout.put("SparkDsCreateInfo.name", CREATE_INFO_NAME);
+ layout.put("SparkDsCreateInfo.option_keys", CREATE_INFO_OPTION_KEYS);
+ layout.put("SparkDsCreateInfo.option_values", CREATE_INFO_OPTION_VALUES);
+ layout.put("SparkDsCreateInfo.num_options", CREATE_INFO_NUM_OPTIONS);
+ layout.put("SparkDsTaskContext", TASK_SIZE);
+ layout.put("SparkDsTaskContext.size", 0L);
+ layout.put("SparkDsTaskContext.stage_id", TASK_STAGE_ID);
+ layout.put("SparkDsTaskContext.stage_attempt_number", TASK_STAGE_ATTEMPT_NUMBER);
+ layout.put("SparkDsTaskContext.partition_id", TASK_PARTITION_ID);
+ layout.put("SparkDsTaskContext.attempt_number", TASK_ATTEMPT_NUMBER);
+ layout.put("SparkDsTaskContext.task_attempt_id", TASK_ATTEMPT_ID);
+ layout.put("SparkDsTaskContext.is_interrupted", TASK_IS_INTERRUPTED);
+ layout.put("SparkDsWriteInfo", WRITE_INFO_SIZE);
+ layout.put("SparkDsWriteInfo.size", 0L);
+ layout.put("SparkDsWriteInfo.mode", WRITE_INFO_MODE);
+ layout.put("SparkDsWriteInfo.query_id", WRITE_INFO_QUERY_ID);
+ return layout;
+ }
+
+ private final int abiVersion;
+ // The functions of the table, by ordinal, null for the functions that the library leaves NULL.
+ private final MethodHandle[] functions = new MethodHandle[Function.values().length];
+
+ private NativeLibrary(long table) throws Throwable {
+ long size = Platform.getLong(null, table);
+ if (size < TABLE_FUNCTIONS) {
+ throw new LoadException(
+ "spark_ds_init returned a function table with a size of " + size + " bytes.", null, true);
+ }
+ abiVersion = Platform.getInt(null, table + TABLE_ABI_VERSION);
+ // The table of another version of the interface has another layout. Only the functions that
+ // fit in the size of the table are read: the table of a library built against an earlier
+ // header ends before the functions that were added later.
+ if (abiVersion == ABI_VERSION) {
+ for (Function function : Function.values()) {
+ if (function.offset() + 8 <= size) {
+ long address = Platform.getLong(null, table + function.offset());
+ if (address != 0) {
+ functions[function.ordinal()] = Ffm.downcall(
+ Ffm.OF_ADDRESS.invoke(address), function.result, function.parameters);
+ }
+ }
+ }
+ }
+ }
+
+ /**
+ * Loads the library at the given path, and gets its table with {@code spark_ds_init}. The
+ * library stays loaded until the JVM exits.
+ *
+ * @param sparkVersion the version of Spark, which the library gets in {@code SparkDsHost}
+ * @throws LoadException if the library cannot be loaded
+ */
+ static NativeLibrary load(String path, String sparkVersion) throws Throwable {
+ if (!isSupported()) {
+ throw new UnsupportedOperationException(
+ "Native data sources require Java " + MINIMUM_JAVA_VERSION + " or later.");
+ }
+ Object lookup;
+ try {
+ lookup = Ffm.LIBRARY_LOOKUP.invoke(Path.of(path), Ffm.GLOBAL_ARENA);
+ } catch (IllegalArgumentException | IllegalCallerException e) {
+ // The library is not a valid library, or native access is not enabled for Spark.
+ throw new LoadException(e.getMessage(), e);
+ } catch (LinkageError e) {
+ // The FFM API could not be initialized, see Ffm: report why, such as an
+ // IllegalCallerException when native access is not enabled for Spark.
+ throw new LoadException(
+ "Cannot use the Foreign Function and Memory API: " + rootCause(e), e);
+ }
+ Optional<?> init = (Optional<?>) Ffm.FIND.invoke(lookup, "spark_ds_init");
+ if (init.isEmpty()) {
+ throw new LoadException(
+ "The library does not export the function spark_ds_init.", null, true);
+ }
+ MethodHandle initFunction = Ffm.downcall(init.get(), C.I32, C.PTR, C.PTR);
+ // The library may keep the host until the JVM exits, even if spark_ds_init fails.
+ long host = Ffm.host(path, sparkVersion);
+ long out = Platform.allocateMemory(8);
+ try {
+ Platform.putLong(null, out, 0);
+ int status = (int) initFunction.invokeExact(host, out);
+ if (status != 0) {
+ throw new LoadException(
+ "spark_ds_init failed with the error code " + status + ".", null, true);
+ }
+ long table = Platform.getLong(null, out);
+ if (table == 0) {
+ throw new LoadException("spark_ds_init returned no function table.", null, true);
+ }
+ return new NativeLibrary(table);
+ } finally {
+ Platform.freeMemory(out);
+ }
+ }
+
+ /** The method handles of the FFM API, initialized when the first library is loaded. */
+ private static final class Ffm {
+ static final Class<?> LAYOUT_CLASS;
+ static final Object LINKER;
+ static final Object GLOBAL_ARENA;
+ static final Object NO_OPTIONS;
+ static final MethodHandle LIBRARY_LOOKUP;
+ static final MethodHandle FIND;
+ static final MethodHandle DOWNCALL;
+ static final MethodHandle UPCALL;
+ static final MethodHandle FUNCTION_OF;
+ static final MethodHandle FUNCTION_OF_VOID;
+ // (long) -> MemorySegment and (MemorySegment) -> long.
+ static final MethodHandle OF_ADDRESS;
+ static final MethodHandle ADDRESS_OF;
+ static final Object BOOL;
+ static final Object INT32;
+ static final Object INT64;
+ static final Object POINTER;
+ // NativeLibrary.log, which SparkDsHost.log calls with the path of its library.
+ static final MethodHandle LOG;
+ // The addresses of the C functions that call Java, which live as long as the JVM:
+ // SparkDsSink, which calls NativeLibrary.sink, and SparkDsTaskContext.is_interrupted, which
+ // calls NativeLibrary.isInterrupted.
+ static final long SINK;
+ static final long IS_INTERRUPTED;
+
+ static {
+ try {
+ // A full-privilege lookup, so that the restricted methods of the FFM API see this class
+ // as their caller, in the unnamed module that Spark enables native access for.
+ MethodHandles.Lookup lookup = MethodHandles.lookup();
+ Class<?> linker = Class.forName("java.lang.foreign.Linker");
+ Class<?> option = Class.forName("java.lang.foreign.Linker$Option");
+ Class<?> symbolLookup = Class.forName("java.lang.foreign.SymbolLookup");
+ Class<?> arena = Class.forName("java.lang.foreign.Arena");
+ Class<?> segment = Class.forName("java.lang.foreign.MemorySegment");
+ Class<?> layout = Class.forName("java.lang.foreign.MemoryLayout");
+ Class<?> layoutArray = Array.newInstance(layout, 0).getClass();
+ Class<?> valueLayout = Class.forName("java.lang.foreign.ValueLayout");
+ Class<?> descriptor = Class.forName("java.lang.foreign.FunctionDescriptor");
+ Class<?> options = Array.newInstance(option, 0).getClass();
+
+ LAYOUT_CLASS = layout;
+ LINKER = lookup.findStatic(linker, "nativeLinker", MethodType.methodType(linker))
+ .invoke();
+ GLOBAL_ARENA = lookup.findStatic(arena, "global", MethodType.methodType(arena)).invoke();
+ NO_OPTIONS = Array.newInstance(option, 0);
+ LIBRARY_LOOKUP = lookup.findStatic(symbolLookup, "libraryLookup",
+ MethodType.methodType(symbolLookup, Path.class, arena));
+ FIND = lookup.findVirtual(symbolLookup, "find",
+ MethodType.methodType(Optional.class, String.class));
+ DOWNCALL = lookup.findVirtual(linker, "downcallHandle",
+ MethodType.methodType(MethodHandle.class, segment, descriptor, options))
+ .asFixedArity();
+ UPCALL = lookup.findVirtual(linker, "upcallStub",
+ MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options))
+ .asFixedArity();
+ FUNCTION_OF = lookup.findStatic(descriptor, "of",
+ MethodType.methodType(descriptor, layout, layoutArray)).asFixedArity();
+ FUNCTION_OF_VOID = lookup.findStatic(descriptor, "ofVoid",
+ MethodType.methodType(descriptor, layoutArray)).asFixedArity();
+ OF_ADDRESS = lookup.findStatic(segment, "ofAddress",
+ MethodType.methodType(segment, long.class));
+ ADDRESS_OF = lookup.findVirtual(segment, "address", MethodType.methodType(long.class));
+ BOOL = valueLayout.getField("JAVA_BOOLEAN").get(null);
+ INT32 = valueLayout.getField("JAVA_INT").get(null);
+ INT64 = valueLayout.getField("JAVA_LONG").get(null);
+ POINTER = valueLayout.getField("ADDRESS").get(null);
+
+ LOG = lookup.findStatic(NativeLibrary.class, "log",
+ MethodType.methodType(void.class, String.class, int.class, long.class));
+ SINK = upcall(
+ lookup.findStatic(NativeLibrary.class, "sink",
+ MethodType.methodType(void.class, long.class, long.class, long.class)),
+ C.VOID, C.PTR, C.PTR, C.I64);
+ IS_INTERRUPTED = upcall(
+ lookup.findStatic(NativeLibrary.class, "isInterrupted",
+ MethodType.methodType(boolean.class, long.class)),
+ C.BOOL, C.PTR);
+ } catch (Throwable e) {
+ // The message names the cause, because the JVM only keeps this message for the later
+ // uses of the class.
+ throw new IllegalStateException("Cannot initialize the FFM API: " + rootCause(e), e);
+ }
+ }
+
+ /** Returns the FunctionDescriptor of a C function with the given types. */
+ static Object descriptor(C result, C... parameters) throws Throwable {
+ Object layouts = Array.newInstance(LAYOUT_CLASS, parameters.length);
+ for (int i = 0; i < parameters.length; i++) {
+ Array.set(layouts, i, layout(parameters[i]));
+ }
+ return result == C.VOID
+ ? FUNCTION_OF_VOID.invoke(layouts)
+ : FUNCTION_OF.invoke(layout(result), layouts);
+ }
+
+ static Object layout(C type) {
+ return switch (type) {
+ case BOOL -> BOOL;
+ case I32 -> INT32;
+ case I64 -> INT64;
+ case PTR -> POINTER;
+ case VOID -> throw new IllegalArgumentException("void has no layout");
+ };
+ }
+
+ /**
+ * Binds the C function at the given address, a MemorySegment, to a method handle whose
+ * parameters and result have the Java types of the given C types, with {@code long} for
+ * pointers.
+ */
+ static MethodHandle downcall(Object address, C result, C... parameters) throws Throwable {
+ MethodHandle handle = (MethodHandle) DOWNCALL.invoke(
+ LINKER, address, descriptor(result, parameters), NO_OPTIONS);
+ // Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses.
+ for (int i = 0; i < parameters.length; i++) {
+ if (parameters[i] == C.PTR) {
+ handle = MethodHandles.filterArguments(handle, i, OF_ADDRESS);
+ }
+ }
+ if (result == C.PTR) {
+ handle = MethodHandles.filterReturnValue(handle, ADDRESS_OF);
+ }
+ return handle;
+ }
+
+ /**
+ * Returns the address of a C function with the given types that calls `target`, whose
+ * parameters have the Java types of the C types, with {@code long} for pointers. The
+ * function lives as long as the JVM.
+ */
+ static long upcall(MethodHandle target, C result, C... parameters) throws Throwable {
+ for (int i = 0; i < parameters.length; i++) {
+ if (parameters[i] == C.PTR) {
+ target = MethodHandles.filterArguments(target, i, ADDRESS_OF);
+ }
+ }
+ Object stub = UPCALL.invoke(
+ LINKER, target, descriptor(result, parameters), GLOBAL_ARENA, NO_OPTIONS);
+ return (long) ADDRESS_OF.invoke(stub);
+ }
+
+ /**
+ * Allocates the SparkDsHost of a library. The library may keep it until the JVM exits, so it
+ * is never freed.
+ */
+ static long host(String path, String sparkVersion) throws Throwable {
+ long log = upcall(MethodHandles.insertArguments(LOG, 0, path), C.VOID, C.I32, C.PTR);
+ byte[] version = sparkVersion.getBytes(StandardCharsets.UTF_8);
+ long versionAddress = Platform.allocateMemory(version.length + 1L);
+ Platform.copyMemory(
+ version, Platform.BYTE_ARRAY_OFFSET, null, versionAddress, version.length);
+ Platform.putByte(null, versionAddress + version.length, (byte) 0);
+ long host = Platform.allocateMemory(HOST_SIZE);
+ Platform.setMemory(host, (byte) 0, HOST_SIZE);
+ Platform.putLong(null, host, HOST_SIZE);
+ Platform.putInt(null, host + HOST_ABI_VERSION, ABI_VERSION);
+ Platform.putLong(null, host + HOST_SPARK_VERSION, versionAddress);
+ Platform.putLong(null, host + HOST_LOG, log);
+ return host;
+ }
+ }
+
+ private static Throwable rootCause(Throwable e) {
+ Throwable cause = e;
+ while (cause.getCause() != null && cause.getCause() != cause) {
+ cause = cause.getCause();
+ }
+ return cause;
+ }
+
+ /** Thrown when a library cannot be loaded, or cannot be used. */
+ static final class LoadException extends RuntimeException {
+ /**
+ * Whether the library was loaded, in which case it stays loaded: Spark found that it cannot
+ * use it afterwards, for example because {@code spark_ds_init} failed.
+ */
+ final boolean loaded;
+
+ LoadException(String message, Throwable cause) {
+ this(message, cause, false);
+ }
+
+ LoadException(String message, Throwable cause, boolean loaded) {
+ super(message, cause);
+ this.loaded = loaded;
+ }
+ }
+
+ /** Thrown when a function is called that the library leaves NULL in its table, with its name. */
+ static final class MissingFunctionException extends RuntimeException {
+ MissingFunctionException(String name) {
+ super(name);
+ }
+ }
+
+ /**
+ * Called by the library, through {@code SparkDsHost.log}, with the path of the library. It must
+ * not throw: an exception thrown by an upcall crashes the JVM.
+ */
+ private static void log(String path, int level, long message) {
+ try {
+ String text = message == 0 ? "" : new String(readCString(message), StandardCharsets.UTF_8);
+ switch (level) {
+ case 1 -> LOGGER.error("Native data source library {}: {}",
+ MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text));
+ case 2 -> LOGGER.warn("Native data source library {}: {}",
+ MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text));
+ case 4 -> LOGGER.debug("Native data source library {}: {}", path, text);
+ case 5 -> LOGGER.trace("Native data source library {}: {}", path, text);
+ default -> LOGGER.info("Native data source library {}: {}",
+ MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text));
+ }
+ } catch (Throwable e) {
+ // The message is lost.
+ }
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Calls
+ // ---------------------------------------------------------------------------------------------
+
+ private MethodHandle function(Function function) {
+ MethodHandle handle = functions[function.ordinal()];
+ if (handle == null) {
+ throw new MissingFunctionException(function.fieldName);
+ }
+ return handle;
+ }
+
+ /** Throws the error that the library reported if a function failed. */
+ private void check(int status, Function function) throws Throwable {
+ if (status == 0) {
+ return;
+ }
+ MethodHandle lastError = functions[Function.LAST_ERROR.ordinal()];
+ long address = lastError == null ? 0 : (long) lastError.invokeExact();
+ String message = address == 0 ? "" : new String(readCString(address), StandardCharsets.UTF_8);
+ throw new RuntimeException(message.isEmpty()
+ ? function.fieldName + " failed with the error code " + status + "."
+ : message);
+ }
+
+ private static long readHandle(long out, Function function) {
+ long handle = Platform.getLong(null, out);
+ if (handle == 0) {
+ throw new IllegalStateException(function.fieldName + " returned a NULL handle.");
+ }
+ return handle;
+ }
+
+ /** The version of the interface that the table of the library implements. */
+ int abiVersion() {
+ return abiVersion;
+ }
+
+ /** Whether the library implements the function, which it does not leave NULL in its table. */
+ boolean hasFunction(Function function) {
+ return functions[function.ordinal()] != null;
+ }
+
+ /** A data source handle, with the capabilities of the data source. */
+ static final class DataSourceHandle {
+ final long handle;
+ final long capabilities;
+
+ DataSourceHandle(long handle, long capabilities) {
+ this.handle = handle;
+ this.capabilities = capabilities;
+ }
+ }
+
+ DataSourceHandle createDataSource(String name, String[] optionKeys, String[] optionValues)
+ throws Throwable {
+ MethodHandle f = function(Function.CREATE);
+ try (Memory memory = new Memory()) {
+ long info = memory.zeroed(CREATE_INFO_SIZE);
+ Platform.putLong(null, info, CREATE_INFO_SIZE);
+ Platform.putLong(null, info + CREATE_INFO_NAME, memory.string(name));
+ Platform.putLong(null, info + CREATE_INFO_OPTION_KEYS, memory.strings(optionKeys));
+ Platform.putLong(null, info + CREATE_INFO_OPTION_VALUES, memory.strings(optionValues));
+ Platform.putLong(null, info + CREATE_INFO_NUM_OPTIONS, optionKeys.length);
+ long capabilities = memory.zeroed(8);
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(info, capabilities, out), Function.CREATE);
+ return new DataSourceHandle(
+ readHandle(out, Function.CREATE), Platform.getLong(null, capabilities));
+ }
+ }
+
+ void schema(long dataSource, long schemaAddress) throws Throwable {
+ check((int) function(Function.SCHEMA).invokeExact(dataSource, schemaAddress),
+ Function.SCHEMA);
+ }
+
+ long createReader(long dataSource, long schemaAddress) throws Throwable {
+ return createWithSchema(Function.CREATE_READER, dataSource, schemaAddress);
+ }
+
+ long createStreamReader(long dataSource, long schemaAddress) throws Throwable {
+ return createWithSchema(Function.CREATE_STREAM_READER, dataSource, schemaAddress);
+ }
+
+ private long createWithSchema(Function function, long dataSource, long schemaAddress)
+ throws Throwable {
+ MethodHandle f = function(function);
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(dataSource, schemaAddress, out), function);
+ return readHandle(out, function);
+ }
+ }
+
+ /**
+ * Creates a writer.
+ *
+ * @param mode {@link #WRITE_APPEND} or {@link #WRITE_TRUNCATE}
+ * @param queryId the ID of the streaming query of a streaming write, or null for a batch write
+ */
+ long createWriter(long dataSource, long schemaAddress, int mode, String queryId)
+ throws Throwable {
+ MethodHandle f = function(Function.CREATE_WRITER);
+ try (Memory memory = new Memory()) {
+ long info = memory.zeroed(WRITE_INFO_SIZE);
+ Platform.putLong(null, info, WRITE_INFO_SIZE);
+ Platform.putInt(null, info + WRITE_INFO_MODE, mode);
+ Platform.putLong(
+ null, info + WRITE_INFO_QUERY_ID, queryId == null ? 0 : memory.string(queryId));
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(dataSource, schemaAddress, info, out), Function.CREATE_WRITER);
+ return readHandle(out, Function.CREATE_WRITER);
+ }
+ }
+
+ void closeDataSource(long dataSource) throws Throwable {
+ function(Function.CLOSE).invokeExact(dataSource);
+ }
+
+ /** Returns how the reader uses each predicate, as SPARK_DS_PREDICATE_* values. */
+ int[] pushPredicates(long reader, String[] predicates) throws Throwable {
+ MethodHandle f = function(Function.PUSH_PREDICATES);
+ try (Memory memory = new Memory()) {
+ long uses = memory.zeroed(4L * Math.max(1, predicates.length));
+ check((int) f.invokeExact(reader, memory.strings(predicates), (long) predicates.length,
+ uses), Function.PUSH_PREDICATES);
+ int[] result = new int[predicates.length];
+ for (int i = 0; i < result.length; i++) {
+ result[i] = Platform.getInt(null, uses + 4L * i);
+ }
+ return result;
+ }
+ }
+
+ boolean pushLimit(long reader, int limit) throws Throwable {
+ MethodHandle f = function(Function.PUSH_LIMIT);
+ try (Memory memory = new Memory()) {
+ long applied = memory.zeroed(1);
+ check((int) f.invokeExact(reader, limit, applied), Function.PUSH_LIMIT);
+ return Platform.getByte(null, applied) != 0;
+ }
+ }
+
+ boolean pruneColumns(long reader, String[] columnNames) throws Throwable {
+ MethodHandle f = function(Function.PRUNE_COLUMNS);
+ try (Memory memory = new Memory()) {
+ long pruned = memory.zeroed(1);
+ check((int) f.invokeExact(reader, memory.strings(columnNames), (long) columnNames.length,
+ pruned), Function.PRUNE_COLUMNS);
+ return Platform.getByte(null, pruned) != 0;
+ }
+ }
+
+ byte[][] partitions(long reader) throws Throwable {
+ MethodHandle f = function(Function.PARTITIONS);
+ try (Sink sink = new Sink()) {
+ check((int) f.invokeExact(reader, Ffm.SINK, sink.id), Function.PARTITIONS);
+ return sink.all();
+ }
+ }
+
+ byte[] serializeReader(long reader) throws Throwable {
+ return serialize(Function.SERIALIZE_READER, reader);
+ }
+
+ void closeReader(long reader) throws Throwable {
+ function(Function.CLOSE_READER).invokeExact(reader);
+ }
+
+ String initialOffset(long streamReader) throws Throwable {
+ return offset(Function.INITIAL_OFFSET, streamReader);
+ }
+
+ String latestOffset(long streamReader) throws Throwable {
+ return offset(Function.LATEST_OFFSET, streamReader);
+ }
+
+ private String offset(Function function, long streamReader) throws Throwable {
+ byte[] offset = serialize(function, streamReader);
+ if (offset == null) {
+ throw new IllegalStateException(function.fieldName + " returned no offset.");
+ }
+ return new String(offset, StandardCharsets.UTF_8);
+ }
+
+ byte[][] streamPartitions(long streamReader, String start, String end) throws Throwable {
+ MethodHandle f = function(Function.STREAM_PARTITIONS);
+ try (Memory memory = new Memory(); Sink sink = new Sink()) {
+ check((int) f.invokeExact(streamReader, memory.string(start), memory.string(end),
+ Ffm.SINK, sink.id), Function.STREAM_PARTITIONS);
+ return sink.all();
+ }
+ }
+
+ byte[] serializeStreamReader(long streamReader) throws Throwable {
+ return serialize(Function.SERIALIZE_STREAM_READER, streamReader);
+ }
+
+ void commitOffset(long streamReader, String end) throws Throwable {
+ MethodHandle f = function(Function.COMMIT_OFFSET);
+ try (Memory memory = new Memory()) {
+ check((int) f.invokeExact(streamReader, memory.string(end)), Function.COMMIT_OFFSET);
+ }
+ }
+
+ void closeStreamReader(long streamReader) throws Throwable {
+ function(Function.CLOSE_STREAM_READER).invokeExact(streamReader);
+ }
+
+ void read(Task task, byte[] readerState, byte[] partition, long streamAddress)
+ throws Throwable {
+ MethodHandle f = function(Function.READ);
+ try (Memory memory = new Memory()) {
+ check((int) f.invokeExact(task.address, memory.bytes(readerState),
+ (long) readerState.length, memory.bytes(partition), (long) partition.length,
+ streamAddress), Function.READ);
+ }
+ }
+
+ byte[] serializeWriter(long writer) throws Throwable {
+ return serialize(Function.SERIALIZE_WRITER, writer);
+ }
+
+ void commit(long writer, long epochId, byte[][] messages) throws Throwable {
+ finishWrite(Function.COMMIT, writer, epochId, messages);
+ }
+
+ void abort(long writer, long epochId, byte[][] messages) throws Throwable {
+ finishWrite(Function.ABORT, writer, epochId, messages);
+ }
+
+ private void finishWrite(Function function, long writer, long epochId, byte[][] messages)
+ throws Throwable {
+ MethodHandle f = function(function);
+ try (Memory memory = new Memory()) {
+ long pointers = memory.zeroed(8L * Math.max(1, messages.length));
+ long lengths = memory.zeroed(8L * Math.max(1, messages.length));
+ for (int i = 0; i < messages.length; i++) {
+ if (messages[i] != null) {
+ Platform.putLong(null, pointers + 8L * i, memory.bytes(messages[i]));
+ Platform.putLong(null, lengths + 8L * i, messages[i].length);
+ }
+ }
+ check((int) f.invokeExact(writer, epochId, pointers, lengths, (long) messages.length),
+ function);
+ }
+ }
+
+ void closeWriter(long writer) throws Throwable {
+ function(Function.CLOSE_WRITER).invokeExact(writer);
+ }
+
+ long createDataWriter(Task task, byte[] writerState, long epochId) throws Throwable {
+ MethodHandle f = function(Function.CREATE_DATA_WRITER);
+ try (Memory memory = new Memory()) {
+ long out = memory.zeroed(8);
+ check((int) f.invokeExact(task.address, memory.bytes(writerState),
+ (long) writerState.length, epochId, out), Function.CREATE_DATA_WRITER);
+ return readHandle(out, Function.CREATE_DATA_WRITER);
+ }
+ }
+
+ void write(long dataWriter, long arrayAddress, long schemaAddress) throws Throwable {
+ check((int) function(Function.WRITE).invokeExact(dataWriter, arrayAddress, schemaAddress),
+ Function.WRITE);
+ }
+
+ /**
+ * Commits a data writer. Then {@code committed} runs right away, before the message is checked:
+ * if the message is invalid, the commit fails, but the data writer must not be aborted.
+ */
+ byte[] commitDataWriter(long dataWriter, Runnable committed) throws Throwable {
+ MethodHandle f = function(Function.COMMIT_DATA_WRITER);
+ try (Sink sink = new Sink()) {
+ check((int) f.invokeExact(dataWriter, Ffm.SINK, sink.id), Function.COMMIT_DATA_WRITER);
+ committed.run();
+ byte[] message = single(sink, Function.COMMIT_DATA_WRITER);
+ return message == null ? new byte[0] : message;
+ }
+ }
+
+ void abortDataWriter(long dataWriter) throws Throwable {
+ check((int) function(Function.ABORT_DATA_WRITER).invokeExact(dataWriter),
+ Function.ABORT_DATA_WRITER);
+ }
+
+ void closeDataWriter(long dataWriter) throws Throwable {
+ function(Function.CLOSE_DATA_WRITER).invokeExact(dataWriter);
+ }
+
+ /**
+ * Calls a function that returns a single byte array through a sink. Returns null if the
+ * function does not call the sink.
+ */
+ private byte[] serialize(Function function, long handle) throws Throwable {
+ MethodHandle f = function(function);
+ try (Sink sink = new Sink()) {
+ check((int) f.invokeExact(handle, Ffm.SINK, sink.id), function);
+ return single(sink, function);
+ }
+ }
+
+ /** Returns the single byte array that a function passed to the sink, or null if none. */
+ private static byte[] single(Sink sink, Function function) throws Throwable {
+ byte[][] values = sink.all();
+ if (values.length > 1) {
+ throw new IllegalStateException(
+ function.fieldName + " returned " + values.length + " values instead of one.");
+ }
+ return values.length == 0 ? null : values[0];
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Memory
+ // ---------------------------------------------------------------------------------------------
+
+ /** Native memory for the arguments of a call, freed after the call. */
+ private static final class Memory implements AutoCloseable {
+ private final List<Long> allocations = new ArrayList<>();
+
+ long zeroed(long size) {
+ long address = Platform.allocateMemory(size);
+ allocations.add(address);
+ Platform.setMemory(address, (byte) 0, size);
+ return address;
+ }
+
+ /** Copies the bytes. Empty bytes are passed as a valid pointer too. */
+ long bytes(byte[] bytes) {
+ long address = zeroed(Math.max(1, bytes.length));
+ Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address, bytes.length);
+ return address;
+ }
+
+ /** Copies the string as a NUL-terminated UTF-8 string. */
+ long string(String value) {
Review Comment:
Thanks, fixed in 74aa0dffc534. `create` now fails with `NATIVE_DATA_SOURCE_ERROR` when an option key or value contains U+0000, and so does an operation whose schema has a column name with it. The data source names of a manifest cannot contain it either. The header says so.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6480975990
Following up from the Rust SDK side (#59520), after checking this PR at `1f40484`. Three suggestions for the C ABI that would remove code from every SDK without adding much to the header. They change existing signatures or the meaning of a flag, so they need to land before version 1 ships in 4.4.0, but they can be follow-ups to this PR.
1. **Pass the schema, and whether the operation is batch or streaming, to the executors.**
[`read`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L540-L542) and [`create_data_writer`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L591-L592) get neither. Spark has the schema in both places: [`NativePartitionReaderFactory`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala#L295-L304) holds `readSchema`, and a write knows its input schema. So the Rust SDK serializes an Arrow schema into its own state ([`ReaderState`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/bridge.rs#L458-L461), [`WriterState`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sq
l/native/rust/src/write.rs#L59-L62)), and tags that state with an enum ([`Kind`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/bridge.rs#L464-L469)) to tell batch from streaming, since both share these two entry points. A C++ SDK would have to do the same, with Arrow IPC or an encoding of its own.
Could `read` and `create_data_writer` take the schema (the read schema after pruning, and the input schema of the write) as a `struct ArrowSchema*`, and `read` a batch/streaming flag? For writes, `epoch_id == -1` already tells them apart. With the schema in `create_data_writer`, [`write`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L600) could also drop its `ArrowSchema` argument, or the header could say that it is always that schema: the SDK imports it and compares its fields for every batch today ([`import_batch`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/write.rs#L392-L414)).
2. **Split `SPARK_DS_TRUNCATE` into a flag for batch writes and one for streaming writes.**
[`SPARK_DS_TRUNCATE`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L362) covers both the save mode "overwrite" and the output mode "complete". A source that supports only one of them either declares it and fails late, in `create_writer`, as the Rust SDK does ([`effective`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/bridge.rs#L122-L131), [the checks of `create_writer`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/bridge.rs#L286-L305)), or does not declare it and loses the one that it supports. With two flags, such as `SPARK_DS_BATCH_TRUNCATE` and `SPARK_DS_STREAMING_TRUNCATE`, the table reports `TRUNCATE` from the batch flag while Spark analyzes a query, and Spark rejects the output mode "complete" from the streaming flag when the query starts, as it does today.
Since this changes the meaning of an existing flag rather than appending one, it cannot wait until after the release.
3. **Return the serialized state from the planning call, instead of a separate serialize call.**
Spark always calls [`partitions`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L463) and then [`serialize_reader`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L471) ([`plan`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala#L279-L291)), [`stream_partitions`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L500-L501) and then [`serialize_stream_reader`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L509), and [`create_
writer`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L428-L429) and then [`serialize_writer`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L554) ([`writer`](https://github.com/apache/spark/blob/1f4048404befa88ac1b956e982711388bf694520/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala#L127-L133)), back to back. Because they are separate calls, the header needs rules for their order, and the Rust SDK keeps a state machine and caches to return the state of the plan that it already made ([`Planner`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439ac/sql/native/rust/src/bridge.rs#L440-L450), [`DriverWriter::serialized`](https://github.com/apache/spark/blob/19351be1babb11b5ea97b2d7a72dc31374f439
ac/sql/native/rust/src/write.rs#L156-L187)). Could `partitions` and `stream_partitions` take a second sink, `set_state`, and `create_writer` a sink for the state of the writer? That removes three functions from the table and the ordering rules from the header.
I will update #59520 to whatever shape lands.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #59424:
URL: https://github.com/apache/spark/pull/59424#issuecomment-6484354011
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/99a1bf9a837c8a0bfeff3236dc5d69c626edc929
- merged into branch-4.x https://github.com/apache/spark/commit/54705ce0c77185524ebe13ce7cd0453409c1d5b6
*Posted by `merge_spark_pr.py`*
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon closed pull request #59424: [SPARK-60101][SQL] Support native data sources implemented in Rust or C++ with the FFM API
URL: https://github.com/apache/spark/pull/59424
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org