[HyukjinKwon]

What changes were proposed in this pull request?

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.

Public API: a single class, org.apache.spark.sql.datasource.NativeBridge (@Evolving, since 4.4.0). A library implements its native methods as JNI functions, such as Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource, and its Javadoc is the specification of the interface: the lifecycle, the handles, the ownership of the Arrow structs, errors (Java exceptions, reported as NATIVE_DATA_SOURCE_ERROR), the JSON format of the pushed predicates (a fixed set of comparison, IN, null, string and boolean predicates), which Spark evaluates again after reading, so that a library only uses them to skip data and must never skip a row for which they are true under the semantics of Spark SQL, for example for NaN and -0.0 (the jar of Spark SQL contains test cases for libraries, org/apache/spark/sql/datasource/predicate-cases.json), the encoding of strings (Java strings, which JNI converts to modified UTF-8; the JSON of the predicates is ASCII), and the compatibility rules: within an ABI version, functions are only added as optional functions, and the signature and semantics of an existing function never change. Spark loads each library into its own copy of the class, with a separate class loader, so any number of libraries exporting the same functions can be loaded side by side. Spark itself ships no native code for native data sources, besides the JNI library that the new arrow-c-data dependency bundles. The interface uses JNI rather than the FFM API, which is final only from Java 22 while Spark supports Java 17 and 21; the interface is versioned, so that a later version can be a plain C ABI called through FFM.

Operation	Required functions	Optional functions
All	abiVersion, createDataSource, closeDataSource	schema
Batch read	createReader, partitions, read, closeReader	pushPredicates, pushLimit, pruneColumns, serializeReader
Streaming read	createStreamReader, initialOffset, latestOffset, streamPartitions, read, closeStreamReader	serializeStreamReader, commitOffset
Batch write	createWriter, commit, abort, closeWriter, createDataWriter, write, commitDataWriter, abortDataWriter	serializeWriter
Streaming write	createStreamWriter, commit, abort, closeWriter, createDataWriter, write, commitDataWriter, abortDataWriter	serializeWriter

A library only exports the functions of the operations it supports. Spark reports a missing required function as DATA_SOURCE_*_NOT_SUPPORTED.

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. A JVM loads a library from a path only once, so after a new build is installed at the same path, the queries of an application that loaded the previous build fail with INVALID_NATIVE_DATA_SOURCE_LIBRARY.REPLACED, which asks to restart it, instead of running mixed builds.

• 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).

Internals:

• execution.datasources.v2.columnar: an internal columnar data source interface, adapted to Data Source V2. Writes convert rows to Arrow batches that end at spark.sql.execution.arrow.maxRecordsPerBatch rows or once they hold spark.sql.execution.arrow.maxBytesPerBatch bytes, whichever is reached first, as for Python UDFs. The writer of a write is created right before its tasks start, after the shuffle stages that adaptive query execution runs first, so that each writer is committed or aborted: the last stage of the data runs in the tasks of the write, so a failure there aborts the writer. Spark evaluates all the pushed predicates again after reading, and only shows the predicates that a library evaluates completely in the plan. A micro-batch streaming write gets a new stream writer for each micro-batch, and the previous one of the query is closed when the next one is created.
• execution.datasources.v2.ffi: the discovery of packages and installed libraries, package parsing, the distribution of packages to the executors, library loading, the implementation of the internal interface on top of NativeBridge with arrow-c-data, which reads the Arrow streams array by array to reject arrays with a non-zero offset (Arrow Java ignores offsets) or with fewer values than their parent refers to (Arrow Java pads them with nulls), the JSON encoding of pushed predicates, limited to the predicates listed in NativeBridge, and the checks of what a library returns: non-zero handles and offsets that are JSON on a single line. The reader imports the batches of a partition into the same Arrow vectors, which are not freeable between batches.

Other changes:

• Configurations: spark.sql.dataSource.native.enabled (static, default true) and spark.sql.dataSource.native.paths.
• Error conditions: NATIVE_DATA_SOURCE_ERROR, INVALID_NATIVE_DATA_SOURCE_LIBRARY (CANNOT_LOAD, REPLACED, UNSUPPORTED_ABI_VERSION), INVALID_NATIVE_DATA_SOURCE_PACKAGE (INVALID_MANIFEST, INVALID_LIBRARY, UNSUPPORTED_ABI_VERSION, UNSUPPORTED_PLATFORM), NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND, NATIVE_DATA_SOURCE_PACKAGE_CONFLICT and NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND.
• New dependency: org.apache.arrow:arrow-c-data, with the same version as arrow-vector. It bundles its own JNI library for Linux, macOS and Windows.
• Documentation: a new page, "Native Data Sources", with the examples below, linked from the Data Source V2 page, and the Scaladoc of SparkSession.addArtifact in sql/api, which now lists .sparkpkg files.
• Resource: org/apache/spark/sql/datasource/predicate-cases.json in sql/core, the predicate test cases for libraries, generated by NativeDataSourceSuite from the evaluation of Spark.
• 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 (passing file: URIs as they are and resolving the other paths in Python), and with Spark Connect as file artifacts.
  • CheckConnectJvmClientCompatibility excludes the new org.apache.spark.sql.datasource package, which the Spark Connect client does not have.

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 8 functions needed for batch reads.

Rust (with the jni crate and arrow-rs)

Cargo.toml:

[package]
name = "rust_range"
version = "0.1.0"
edition = "2021"

[lib]
name = "spark_datasource_rust_range"
crate-type = ["cdylib"]

[dependencies]
arrow-array = { version = "58", features = ["ffi"] }
arrow-schema = { version = "58", features = ["ffi"] }
jni = "0.21"

src/lib.rs:

//! A native data source for Apache Spark, written in Rust. It reads the numbers [0, end) as a
//! column `id`, in two partitions. The option `end` defaults to 10.

use std::panic::{catch_unwind, AssertUnwindSafe};
use std::sync::Arc;

use arrow_array::ffi_stream::FFI_ArrowArrayStream;
use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator};
use arrow_schema::ffi::FFI_ArrowSchema;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use jni::objects::{JByteArray, JClass, JObject, JObjectArray, JString};
use jni::sys::{jint, jlong, jobjectArray};
use jni::JNIEnv;

type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;

struct RangeSource {
    end: i64,
}

struct RangeReader {
    end: i64,
}

fn schema() -> SchemaRef {
    Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]))
}

/// Runs `f`, and reports its error or panic as a Java exception: neither may cross into the JVM.
fn call<T>(env: &mut JNIEnv, default: T, f: impl FnOnce(&mut JNIEnv) -> Result<T>) -> T {
    let message = match catch_unwind(AssertUnwindSafe(|| f(env))) {
        Ok(Ok(value)) => return value,
        Ok(Err(e)) => e.to_string(),
        Err(_) => "the data source panicked".to_string(),
    };
    if !env.exception_check().unwrap_or(true) {
        let _ = env.throw_new("java/lang/RuntimeException", message);
    }
    default
}

fn strings(env: &mut JNIEnv, array: &JObjectArray) -> Result<Vec<String>> {
    let mut result = Vec::new();
    for i in 0..env.get_array_length(array)? {
        let value = JString::from(env.get_object_array_element(array, i)?);
        result.push(env.get_string(&value)?.into());
    }
    Ok(result)
}

#[no_mangle]
pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_abiVersion(
    _env: JNIEnv,
    _class: JClass,
) -> jint {
    1
}

#[no_mangle]
pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource(
    mut env: JNIEnv,
    _class: JClass,
    _name: JString,
    keys: JObjectArray,
    values: JObjectArray,
) -> jlong {
    call(&mut env, 0, |env| {
        let mut source = RangeSource { end: 10 };
        for (key, value) in strings(env, &keys)?.into_iter().zip(strings(env, &values)?) {
            if key == "end" {
                source.end = value.parse()?;
            }
        }
        Ok(Box::into_raw(Box::new(source)) as jlong)
    })
}

#[no_mangle]
pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_schema(
    mut env: JNIEnv,
    _class: JClass,
    _source: jlong,
    schema_address: jlong,
) {
    call(&mut env, (), |_| {
        let schema = FFI_ArrowSchema::try_from(schema().as_ref())?;
        unsafe { std::ptr::write(schema_address as *mut FFI_ArrowSchema, schema) };
        Ok(())
    })
}

#[no_mangle]
pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_createReader(
    mut env: JNIEnv,
    _class: JClass,
    source: jlong,
    _schema_address: jlong,
) -> jlong {
    call(&mut env, 0, |_| {
        let source = unsafe { &*(source as *const RangeSource) };
        Ok(Box::into_raw(Box::new(RangeReader { end: source.end })) as jlong)
    })
}

#[no_mangle]
pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_closeDataSource(
    _env: JNIEnv,
    _class: JClass,
    source: jlong,
) {
    drop(unsafe { Box::from_raw(source as *mut RangeSource) });
}

#[no_mangle]
pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_partitions(
    mut env: JNIEnv,
    _class: JClass,
    reader: jlong,
) -> jobjectArray {
    call(&mut env, std::ptr::null_mut(), |env| {
        let end = unsafe { &*(reader as *const RangeReader) }.end;
        let bounds = [(0, end / 2), (end / 2, end)];
        let partitions = env.new_object_array(bounds.len() as i32, "[B", JObject::null())?;
        for (i, (start, end)) in bounds.into_iter().enumerate() {
            let bytes = [start.to_le_bytes(), end.to_le_bytes()].concat();
            let partition = env.byte_array_from_slice(&bytes)?;
            env.set_object_array_element(&partitions, i as i32, partition)?;
        }
        Ok(partitions.into_raw())
    })
}

#[no_mangle]
pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_closeReader(
    _env: JNIEnv,
    _class: JClass,
    reader: jlong,
) {
    drop(unsafe { Box::from_raw(reader as *mut RangeReader) });
}

#[no_mangle]
pub extern "system" fn Java_org_apache_spark_sql_datasource_NativeBridge_read(
    mut env: JNIEnv,
    _class: JClass,
    _reader_state: JByteArray,
    partition: JByteArray,
    stream_address: jlong,
) {
    call(&mut env, (), |env| {
        let bytes = env.convert_byte_array(&partition)?;
        let start = i64::from_le_bytes(bytes[0..8].try_into()?);
        let end = i64::from_le_bytes(bytes[8..16].try_into()?);
        let ids = Int64Array::from_iter_values(start..end);
        let batch = RecordBatch::try_new(schema(), vec![Arc::new(ids)])?;
        let batches = RecordBatchIterator::new(vec![Ok(batch)], schema());
        let stream = FFI_ArrowArrayStream::new(Box::new(batches));
        unsafe { std::ptr::write(stream_address as *mut FFI_ArrowArrayStream, stream) };
        Ok(())
    })
}

Build and install, here on macOS:

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:

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 .)

C++ (with Arrow C++)

cpp_range.cc:

// 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 <jni.h>

#include <cstdint>
#include <memory>
#include <stdexcept>
#include <string>
#include <type_traits>

#include <arrow/api.h>
#include <arrow/c/bridge.h>

#define SPARK_JNI(ret, method) \
  extern "C" JNIEXPORT ret JNICALL Java_org_apache_spark_sql_datasource_NativeBridge_##method

namespace {

struct RangeSource {
  int64_t end = 10;
};

struct RangeReader {
  int64_t end;
};

std::shared_ptr<arrow::Schema> Schema() {
  return arrow::schema({arrow::field("id", arrow::int64(), /*nullable=*/false)});
}

// Runs `f`, and reports its error as a Java exception: no C++ exception may cross into the JVM.
template <typename F>
auto Call(JNIEnv* env, F&& f) -> decltype(f()) {
  try {
    return f();
  } catch (const std::exception& e) {
    env->ThrowNew(env->FindClass("java/lang/RuntimeException"), e.what());
    if constexpr (!std::is_void_v<decltype(f())>) return {};
  }
}

void Check(const arrow::Status& status) {
  if (!status.ok()) throw std::runtime_error(status.ToString());
}

// Returns a Java string as modified UTF-8, which is UTF-8 except for U+0000 and the characters
// outside the Basic Multilingual Plane: enough for the option of this data source.
std::string ToString(JNIEnv* env, jobject value) {
  const char* chars = env->GetStringUTFChars(static_cast<jstring>(value), nullptr);
  std::string result(chars);
  env->ReleaseStringUTFChars(static_cast<jstring>(value), chars);
  return result;
}

}  // namespace

SPARK_JNI(jint, abiVersion)(JNIEnv*, jclass) { return 1; }

SPARK_JNI(jlong, createDataSource)(JNIEnv* env, jclass, jstring, jobjectArray keys,
                                   jobjectArray values) {
  return Call(env, [&] {
    auto source = std::make_unique<RangeSource>();
    for (jsize i = 0; i < env->GetArrayLength(keys); i++) {
      if (ToString(env, env->GetObjectArrayElement(keys, i)) == "end") {
        source->end = std::stoll(ToString(env, env->GetObjectArrayElement(values, i)));
      }
    }
    return reinterpret_cast<jlong>(source.release());
  });
}

SPARK_JNI(void, schema)(JNIEnv* env, jclass, jlong, jlong schema_address) {
  Call(env, [&] {
    Check(arrow::ExportSchema(*Schema(), reinterpret_cast<ArrowSchema*>(schema_address)));
  });
}

SPARK_JNI(jlong, createReader)(JNIEnv* env, jclass, jlong source, jlong) {
  return Call(env, [&] {
    auto* reader = new RangeReader{reinterpret_cast<RangeSource*>(source)->end};
    return reinterpret_cast<jlong>(reader);
  });
}

SPARK_JNI(void, closeDataSource)(JNIEnv*, jclass, jlong source) {
  delete reinterpret_cast<RangeSource*>(source);
}

SPARK_JNI(jobjectArray, partitions)(JNIEnv* env, jclass, jlong reader) {
  return Call(env, [&] {
    int64_t end = reinterpret_cast<RangeReader*>(reader)->end;
    int64_t bounds[2][2] = {{0, end / 2}, {end / 2, end}};
    jobjectArray partitions = env->NewObjectArray(2, env->FindClass("[B"), nullptr);
    for (jsize i = 0; i < 2; i++) {
      jbyteArray partition = env->NewByteArray(sizeof(bounds[i]));
      env->SetByteArrayRegion(partition, 0, sizeof(bounds[i]),
                              reinterpret_cast<const jbyte*>(bounds[i]));
      env->SetObjectArrayElement(partitions, i, partition);
    }
    return partitions;
  });
}

SPARK_JNI(void, closeReader)(JNIEnv*, jclass, jlong reader) {
  delete reinterpret_cast<RangeReader*>(reader);
}

SPARK_JNI(void, read)(JNIEnv* env, jclass, jbyteArray, jbyteArray partition,
                      jlong stream_address) {
  Call(env, [&] {
    int64_t bounds[2];
    env->GetByteArrayRegion(partition, 0, sizeof(bounds), reinterpret_cast<jbyte*>(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,
                                         reinterpret_cast<ArrowArrayStream*>(stream_address)));
  });
}

Build (Arrow C++ 24 requires C++20) and install, or package as above:

c++ -std=c++20 -O2 -shared -fPIC \
  -I"$JAVA_HOME/include" -I"$JAVA_HOME/include/darwin" \
  -I"$ARROW_HOME/include" -L"$ARROW_HOME/lib" -larrow \
  -o libspark_datasource_cpp_range.dylib cpp_range.cc
cp libspark_datasource_cpp_range.dylib /usr/local/lib/

Usage, in classic PySpark or with Spark Connect. Once installed, the data sources are found automatically, without spark.addArtifact or any configuration:

>>> 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:

>>> 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.

Does this PR introduce any user-facing change?

Yes. It adds native data sources, with the NativeBridge interface, the configurations and the error conditions above. 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?

• New NativeDataSourceSuite, with a dependency-free C++ test library (sql/core/src/test/resources/native-datasource/test_native_datasource.cc) compiled by the tests with the C++ compiler of the machine; the tests that load it are skipped if there is no compiler. 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, and invalid packages (manifest, ABI version, platform and library). It also has unit tests for the manifest, the platform names, the native library path and the JSON encoding of predicates.
• NativeDataSourceSuite also covers installed libraries: found automatically in the native library path with case-insensitive names, a library installed under several names, writes, the precedence of the packages of a session, the lookup by name on executors (NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND), and invalid installed libraries (CANNOT_LOAD, and UNSUPPORTED_ABI_VERSION, with the same error when the library is used again).
• New NativeDataSourceConnectSuite, with a Spark Connect client and server: spark.addArtifact of a package, reads, writes, SQL tables, the isolation of the packages of different sessions, and installed libraries.
• New ColumnarDataSourceSuite, which tests the internal Data Source V2 adapter with a data source implemented on the JVM.
• New PySpark tests for adding .sparkpkg artifacts in classic mode, whatever the flags and in lists with other files, and with Spark Connect.
• 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 NativeBridge 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).
  • Predicates: Spark evaluates the predicates again when a library reports them all as evaluated completely without applying them; the answers of a library map back to the predicates that Spark passed when the first one cannot be encoded; no limit is pushed with a filter; the JSON is ASCII for strings with characters outside the BMP; and the predicate test cases are regenerated and compared, as a golden file.
  • Reads: a struct column read in several batches through whole-stage codegen; a nested field read with its top-level column; large lists and dictionary encoding rejected with the name of the column; the same reader planned twice after a failure.
  • Streaming: offsets that are null, multi-line or not JSON are rejected; the stream writer gets the same query ID after a restart from the checkpoint; a failed micro-batch is aborted with its epoch ID and a null message for its task; stop() closes the stream reader.
  • Handles of 0 from each of the six create functions are rejected.
  • Distribution: concurrent distributions of the same configured package; a valid package with another checksum than the driver's is not loaded; packages of an unsupported version are logged.
  • Installed libraries: a library replaced in place fails with REPLACED, and works again once the loaded build is installed again.
  • The link errors of Arrow are not taken for missing functions.
  • PySpark addArtifact with a relative file name that contains a colon.
  • The tests of the predicates, the structs, the retried planning, the offsets, the closing of stream readers, the concurrent distributions, the package checksum, the replaced library, the logging, the unsupported Arrow types and the link errors were checked to fail with their fix reverted.
• The Rust and C++ examples above were built from scratch with the commands above, and run on macOS (osx-aarch64) with classic PySpark and with Spark Connect, with the output shown above: installed in the lib directory of an installation prefix in PATH without any configuration (and not found without it), found through java.library.path, added with spark.addArtifact, and as SQL tables.
• The C++ test library also builds without warnings with GCC 15 and libstdc++ (-Wall -Wextra), and NativeDataSourceSuite runs on Linux with g++ in CI.
• Existing suites: ResolvedDataSourceSuite, ArtifactManagerSuite, PythonDataSourceSuite, PythonStreamingDataSourceSuite, DataStreamReaderWriterSuite, DataFrameReaderWriterSuite, DataFrameWriterV2Suite, DataSourceSuite, DDLSourceLoadSuite, DataSourceV2Suite, StreamingDataSourceV2Suite, SQLConfSuite, SparkThrowableSuite and SparkConfigBindingPolicySuite.

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.

[pan3793]

Just out of curiosity, was the FFM API considered instead of JNI?

[HyukjinKwon]

@pan3793 Good question. Yes, but FFM cannot be the only bridge yet because of the Java baseline. Spark compiles with --release 17 and supports Java 17, 21 and 25, while the FFM API is final only from Java 22 (JEP 454): it is a preview API in Java 21 and an incubator module in Java 17. So an FFM bridge would need a separate source set for Java 22+, and we would still need the JNI path for Java 17 and 21, with two interfaces for library authors to choose from.

FFM would fit this design well, though, and the design leaves room for it:

• With FFM, a library could export a plain C ABI, without JNIEnv or the jni crate, and report errors with return codes instead of Java exceptions. Spark could also open each library with SymbolLookup.libraryLookup instead of defining a copy of NativeBridge in a class loader of its own for each library.
• Performance does not decide it either way: the calls are coarse-grained (per scan, partition and batch), and the data crosses the boundary as Arrow C Data Interface structs, without copying.
• The interface is versioned (NativeBridge.abiVersion, and abiVersion in the manifest of a package), and only NativeBridge is public: loading and calling the libraries is internal. Once the minimum Java version of Spark allows it, a later version of the interface can be a plain C ABI called through FFM, while Spark keeps loading the libraries of version 1.
• Spark already runs with --enable-native-access=ALL-UNNAMED, without which recent JDKs warn about both JNI (JEP 472) and the restricted methods of FFM, so neither needs new flags.

[pan3793]

... FFM cannot be the only bridge yet because of the Java baseline.

I suppose we can raise our JDK baseline in the upcoming 5.0.0 (early 2027)? And this can also be a 25-specific feature before that happens, with the decoupled Connect client-server architecture, upgrading server-side JDK version should be much easier.

With FFM, a library could export a plain C ABI, without JNIEnv or the jni crate ...

Yeah, this is the biggest advantage in my mind, non-JVM developers see a clean C header file without needing to learn JNI stuff.

[HyukjinKwon]

I took a look about FFM (and did prototype). I think it's sort of good and bad. For example, we have to release header file together whereas JNI one does not require it. BTW, we will release 4.4.0 in Dec. and 4.5.0 around Mar. since we will do the quarterly release. I think we can keep JNI version, and do FFM one later around Spark 5.0.0.

[pan3793]

we have to release header file together whereas JNI one does not require it.

for JNI, we define the interface via a Java class (NativeBridge), and the developer uses javap to generate the header file; for FFM, we define a header file as the interface directly.

I suppose, from the perspective of the native datasource connector developer, the latter should be more friendly (hygienic, and save one step to generate it from Java).

If we choose to define a header file as the interface directly, we can:

• for Java 17 and 21: Spark provides a JNI bridge to call it
• for Java 25 and later: Spark uses the FFM API to call it directly

and once we raise our JDK baseline to 25+, we can get rid of the JNI glue code, the ABI keeps no change from the beginning.

[HyukjinKwon]

Yeah. Let's actually do that. We can deprecate this in 4.5.0, and drop it in 5.0.0 with replacing it to FFM. I am thinking about bringing more native code into Spark 5 in fact ..

[dongjoon-hyun]

Thank you for this proposal, @HyukjinKwon. This is great! I reviewed the lookup, distribution and lifecycle of the native data sources, and left 10 inline comments, summarized below in the same order.

Main issues

1. One invalid package breaks every lookup miss (NativeDataSourceRegistry.scala L65)
2. Streaming tasks look up the package under the wrong session UUID (NativeDataSourceRegistry.scala L168)

Other issues

3. Classic addArtifact(..., file=True) doesn't register the package (session.py L2299)
4. Spark comparison semantics (NaN, collation) for pushed predicates (NativePredicates.scala L79)
5. Artifact dedup by file name can ship a different package (NativeDataSourceRegistry.scala L152)
6. The stream writer handle is never closed explicitly (NativeDataSource.scala L294)
7. A failed commitDataWriter skips the task abort (NativeDataSource.scala L484)
8. The registry lookup runs twice for each resolution (NativeDataSourceV2.scala L46)
9. Installed native libraries now block Python data source registration (DataSource.scala L686)
10. The scan description shows the class name instead of the data source name (ColumnarTable.scala L148)

The first two items seem to be the most important: one invalid package breaks unrelated lookups, and streaming queries on a cluster with isolated sessions would not find the package on the executors.

[dongjoon-hyun]

Thank you for addressing all the comments, @HyukjinKwon. I checked 9dd4f88 against each thread, and all 10 points are handled. I left 4 more minor, non-blocking comments:

1. An invalid package that lists the name hides a valid one (NativeDataSourceRegistry.scala L88)
2. The thread-local can keep a stale lookup on pooled threads (NativeDataSourceRegistry.scala L70)
3. The package cache is never evicted (NativeDataSourcePackage.scala L102)
4. No end-to-end test for closing the previous stream writer (NativeDataSource.scala L565)

[dongjoon-hyun]

+1, LGTM.

[sarutak]

I went over the lookup/distribution and the reader/writer lifecycle again on 9dd4f88. The fixes for the earlier round all check out in the code, and the overall design (per-library BridgeClassLoader for side-by-side loading, NativeHandle backed by AtomicLong plus Cleaner, Arrow struct release/close on every path, tryWithSafeFinally around the native calls) is solid.

I left four inline comments, none blocking:

1. distribute collapses "no active session" and "local mode" into one branch, so a missing active session would silently skip distribution rather than fail clearly. I could not find a planning path that actually has no active session, so this is defensive hardening.
2. The batch writer's commit closes the handle outside tryWithSafeFinally, unlike abort; if library.commit throws, the driver-side handle falls back to the Cleaner. One-line symmetry fix.
3. spark.sql.dataSource.native.enabled defaults to true: a note on whether opt-in is preferable for the first release of a feature that loads native code (nit).
4. NativeLibrary.load uses setAccessible across class loaders: a comment on why would help future readers (nit).

The last two are nits. Thanks for the very thorough design and docs on this.

[LuciferYang]

I went through 3814147 again with a local build and the test library. The fixes from the earlier rounds hold up. Five more inline comments: four reproduced on 3814147, and the batch-size one is computed from the offset limit.

[sarutak]

One small thing before merge: could you update the PR description to mention spark.sql.execution.arrow.maxBytesPerBatch? The Internals section still says

Writes convert rows to Arrow batches of at most
spark.sql.execution.arrow.maxRecordsPerBatch rows.

Since 9ac5a49, batches are also bounded by maxBytesPerBatch, so a line like "...at most maxRecordsPerBatch rows or maxBytesPerBatch bytes, whichever is reached first" would match the code. The description already names maxRecordsPerBatch explicitly, so adding the byte limit keeps the two consistent. Not blocking, just so the description stays accurate for the record.

[HyukjinKwon]

BTW, I plan to leave this open a while. I want others to take a look too :-).

[HyukjinKwon]

@sarutak Thanks, updated. The Internals section now says writes convert rows to Arrow batches "that end at spark.sql.execution.arrow.maxRecordsPerBatch rows or once they hold spark.sql.execution.arrow.maxBytesPerBatch bytes, whichever is reached first, as for Python UDFs". I worded it this way because the byte limit is checked after each row, so a batch can go slightly over it, as in PythonArrowInput.

While there, I went through the rest of the description against 9ac5a49 and updated the other parts that changed during the review:

• the per-micro-batch stream writers;
• the precedence of usable packages over invalid or unsupported ones;
• the checksum-named copies of configured packages;
• the single lookup per resolution and the Python data source precedence on registration;
• the changes to existing classes (DataFrameWriter/DataStreamWriter under withActive, ArtifactManager, DataSourceRegistration, PySpark addArtifacts, the Connect compatibility exclusion);
• the user-facing fix for writes from a non-active session;
• the tests added during the review.

[LuciferYang]

I went through 9ac5a49 again. The fixes from the earlier rounds hold up in the code. Four more inline comments, each checked on 9ac5a49 with a small test suite: Arrow Java for the first two, and a JVM ColumnarDataSource for the last two.

The first one is the one I'd like fixed before merge: when a library returns a sliced Arrow array, Spark reads the wrong rows without any error. The other three are places where checkSchema lets through a layout that ArrowColumnVector can't read, or where NativeBridge promises something Spark doesn't do (calling abort when a write fails, and calling pushLimit at most once).

[HyukjinKwon]

@cloud-fan Thanks for the thorough review. I addressed all the inline comments in aa0f3ba (and e3ff80c, a test), and replied to each. On the findings of the summary:

Six NativeBridge statements, all fixed in the Javadoc and on the docs page:

• Errors: the errors of closeReader, closeStreamReader and closeWriter are only logged, so the work that can fail a query, such as flushing data, belongs in commit.
• Optional methods: commitOffset is listed with the pushdown and serialize functions.
• createDataSource: the name is the one listed in the manifest, or the requested name in lower case for an installed library.
• Schemas: the variants that Spark reads are listed exactly: large_utf8 and utf8_view, large_binary and binary_view, and any names for the entries of a map. The rejected types now include large and fixed-size lists. Spark also rejects them with a NATIVE_DATA_SOURCE_ERROR that names the column, instead of UNSUPPORTED_ARROWTYPE.
• Handles: a batch reader whose partitions are never planned, for example to only explain the query, is released once it is no longer referenced. The Lifecycle and Handles sections say so.
• pushPredicates: the -0.0 sentence is fixed. IEEE 754 comparisons treat -0.0 as equal to 0.0, while total orders such as f64::total_cmp and the arrow-rs kernels do not, and NaN differs from IEEE 754. Spark now also evaluates all the predicates again after reading (see @tdcmeehan's thread), so the rule is that a library must never skip a row for which they are true.

Five guards without tests, all added:

• The mapping of the answers: in upper(name) = 'N1' AND id < 5, the first predicate cannot be encoded, and only id < 5 is shown as pushed. The test fails if the answers are zipped with all the predicates.
• Dictionary-encoded columns: the test library can export one (extra=dictionary).
• A failed micro-batch write: Spark aborts epoch 0 with a null message for its task. The test library now writes _ABORTED_<epoch>, with the number of null messages.
• stop() releases the stream reader: closeStreamReader of the test library writes a marker.
• A nested field: select("st.b") reads the whole st column. The test fails if the pruned struct is passed as the read schema, because the check of the schema rejects the data.

PR description: both sentences are updated. The logging one is accurate now, since Spark also logs the packages of an unsupported version.

Should version 1 be a C ABI? That is #59424, which I opened for comparison. It is the same feature with a C header as the interface: spark_datasource.h, with spark_ds_* functions on int64_t, int32_t and bool, UTF-8 strings, the Arrow C structs, status codes with spark_ds_last_error, and byte arrays through a sink callback. It is called through the FFM API, so it requires Java 22 or later. A JNI shim for Java 17 and 21 would mean building and shipping native code with Spark for each platform.

Could a failed load of Arrow's JNI library look like an unsupported operation? Yes. In Arrow 19, JniLoader loads the bundled arrow_cdata_jni with System.load in a block that only catches IOException, so an UnsatisfiedLinkError escapes. That happens for example when java.io.tmpdir does not allow mapping the library. Fixed as you suggested: NativeLibrary turns the UnsatisfiedLinkError of the method handle of a function into a MissingFunctionException, and only that exception means a missing function to optional and call. The link errors of Arrow are reported as a NATIVE_DATA_SOURCE_ERROR with their message. New test: "tell missing functions from other link errors".

[HyukjinKwon]

@LuciferYang Thanks for running the suites again. On memory, aa0f3ba adds a Memory section to the docs page. It says that the memory of a library, the Arrow buffers that it exports and the batches that Spark writes are not accounted by Spark, and that this memory comes out of spark.executor.memoryOverhead (and spark.driver.memoryOverhead on the driver). It also says that the size of a read batch is up to the library. Thanks for the pointers to Gluten's ManagedAllocationListener and ReservationListener, which I'll look at for the design of memory accounting that @tdcmeehan asked for.

@tokoko Thanks, that's an interesting direction for the C ABI of #59424. Version 1 needs a few things that ADBC does not have today: streaming reads with offsets, the protocol of distributed writes (task commits with messages, aborts, and epochs for streaming writes), and the pushdown semantics of Spark. With the ingest API that you are working on (apache/arrow-adbc#4317), building on ADBC with extensions where needed could be worth exploring for that interface, so that one driver serves several engines.

@gengliangwang @gerashegalov Thanks, all addressed in aa0f3ba, with replies inline.

[lidavidm]

Pushdown semantics has come up quite a few times (e.g. @SamArch27 evaluating how to do pushdown for DuckDB-ADBC). I think it would require something like Substrait so that semantics are consistent across engines and drivers, with each driver supporting a limited subset (e.g. table scans with filters) and translating to SQL if needed.

ADBC does do streaming reads, and some drivers can even expose a read as a set of partitions that can be distributed across workers. As for reads with offsets, that could perhaps be done via Substrait as well, if we go that direction.

[lidavidm]

Also note that @CurtHagenlocher has been thinking about how to support pluggable filesystems and file type drivers (e.g. Excel files over S3) in the ADBC API, which may be relevant here as well, as it would reuse the ADBC ABI.

[HyukjinKwon]

(let me turn this as a draft, and make #59424 non-draft. per the review comments, I think we should go that way)

[HyukjinKwon]

@lidavidm Thanks for the context. You're right that ADBC has streaming result sets and partitioned execution (AdbcStatementExecutePartitions and AdbcConnectionReadPartition), which map well to how Spark plans and reads partitions. What I meant above is Spark's micro-batch streaming, where the source tracks offsets across micro-batches and replays ranges after restarts, and the write side (task commit messages, aborts, and epochs).

Substrait for the pushdown would also address the semantics concerns that @tdcmeehan raised here, with each driver supporting a declared subset. For now, Spark evaluates every pushed predicate again after reading, so a library only uses them to skip data. The pluggable filesystems and file type drivers that @CurtHagenlocher is looking at sound relevant too, for a C ABI such as the one in #59424.

[gerashegalov]

LGTM

[zhouyuan]

Thanks for this, it's a really nice design. Since the PR is still a draft, I'd like to raise a use case from Apache Gluten, the native execution engine for Spark built on Velox. Other native engines such as Comet (DataFusion) are in the same position.

Why it matters. Native data sources and native execution engines fit together well. The library reads the data, and the engine runs the rest of the query natively. Today, a query on a native data source with Gluten enabled looks like this:

lib.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:

lib.read -> ArrowArrayStream -> Velox (import) -> filter / agg / join / shuffle

No C ABI change is needed for this. The JNI read call only initializes the ArrowArrayStream. After that, get_next and release are plain C function pointers that a native engine can call directly. 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, and each library is loaded through its own copy of NativeBridge, so an external engine can't do this today. Something like:

   @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, and the vanilla path would stay exactly as it is.

2. A threading contract for the stream in the NativeBridge docs. The docs currently say Spark calls the stream "from the thread of the task that reads the partition". Could ABI v1 instead guarantee that get_next/release don't use JNI or the JVM, may be called from a thread other than the one that called read (never concurrently), and report errors through get_last_error? 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.

[LuciferYang]

Thanks for this, it's a really nice design. Since the PR is still a draft, I'd like to raise a use case from Apache Gluten, the native execution engine for Spark built on Velox. Other native engines such as Comet (DataFusion) are in the same position.

Why it matters. Native data sources and native execution engines fit together well. The library reads the data, and the engine runs the rest of the query natively. Today, a query on a native data source with Gluten enabled looks like this:

lib.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:

lib.read -> ArrowArrayStream -> Velox (import) -> filter / agg / join / shuffle

No C ABI change is needed for this. The JNI read call only initializes the ArrowArrayStream. After that, get_next and release are plain C function pointers that a native engine can call directly. 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, and each library is loaded through its own copy of NativeBridge, so an external engine can't do this today. Something like:

   @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, and the vanilla path would stay exactly as it is.

2. A threading contract for the stream in the NativeBridge docs. The docs currently say Spark calls the stream "from the thread of the task that reads the partition". Could ABI v1 instead guarantee that get_next/release don't use JNI or the JVM, may be called from a thread other than the one that called read (never concurrently), and report errors through get_last_error? 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.

@zhouyuan should have commented on #59424.

[zhouyuan]

@LuciferYang Thanks, moved to #59424

[gengliangwang]

I have two suggestions on the DSv2 adapter:

1. Separate driver-side planning from executor-side reading.

   ColumnarReader currently combines pushdown negotiation and partition planning with serialization and read(partition) on executors. NativeDataSourceReader switches between these roles using a transient native handle and mutable readerState. Also, ColumnarScanBuilder.build() passes closures that capture the mutable builder, alongside the scan's schema and pushdown metadata.

   Could we keep ColumnarTable -> ColumnarScanBuilder -> ColumnarScan, but introduce a driver-only BatchReadPlanner whose plan() returns PlannedBatch(partitions, readerFactory)? The executor factory would contain only the library location, source name, read schema, and serialized reader state. build() could establish stable scan metadata and explicit planner ownership, with defined behavior for repeated calls.

   This can preserve the native ABI and the current lazy partition-planning point: call partitions, then serializeReader, and close the native reader in finally. The batch adapter would cache the planned result, and we should retain the existing caching of planning failures as well. Batch planner creation should remain lazy for streaming-only sources. Plans that never reach partition planning would still need the Cleaner fallback, since DSv2 has no general scan-close hook. The benefit is a clearer driver/executor boundary and resource ownership.

2. Have Table.capabilities() reflect the individual source's supported operations.

   ColumnarTable.capabilities() currently advertises batch/streaming reads and writes plus truncate for every source, including read-only libraries. Unsupported operations are then discovered when the corresponding native creation method is called.

   Could the adapter receive a source-specific capability set, so Spark can reject unsupported operations earlier? For native sources, this needs a reliable discovery contract, such as capability metadata or an explicit capability query, covering both packaged and installed libraries. That mechanism is a separate design choice: a new native query would need an ABI compatibility plan, while the planner/factory split above can remain entirely internal.

[FelixYBW]

@HyukjinKwon I'm comparing the implementation vs. what we have done in Gluten. Claude reported a bug when it work with a pyarrow UDF. FYI, @viirya

import pyarrow as pa
from pyspark.sql import SparkSession
from pyspark.sql.functions import arrow_udf, col, assert_true, concat_ws, lit
from pyspark.sql.types import LongType

@arrow_udf(LongType())
def ident(v: pa.Array) -> pa.Array:
    return v

spark = (SparkSession.builder.master("local[1]")
         .config("spark.driver.extraLibraryPath", "/tmp/nds")
         .getOrCreate())

def df(end, batch_size):
    return (spark.read.format("native_range")
            .option("end", str(end)).option("partitions", "1")
            .option("batch_size", str(batch_size)).load()
            .select("id")
            .withColumn("uid", ident(col("id"))))

# Case A: crash. Small input: the reader is closed while outputs are still pending.
try:
    df(200, 100).collect()
except Exception as e:
    print("A:", "IndexOutOfBoundsException" in str(e))

# Case B: wrong results. Large input: the reused vectors are overwritten by read-ahead.
check = df(2_000_000, 1000).select(
    assert_true(col("id") == col("uid"),
                concat_ws(" ", lit("MISMATCH id"), col("id"), lit("uid"), col("uid"))))
try:
    check.collect()          # use collect(), not count(): count() prunes the check away
except Exception as e:
    print("B:", next(l for l in str(e).splitlines() if "MISMATCH" in l))

# Control: same query on the row path, correct
spark.conf.set("spark.sql.execution.arrow.pythonUDF.columnarInput.enabled", "false")
print("control mismatches:", sum(r.id != r.uid for r in df(200, 100).collect()))

reproduce:

A: True
B: ... MISMATCH id 1383000 uid 0 ... (the exact id depends on how far the runner read ahead)
control mismatches: 0

There are two separate causes:

1. Vector reuse gives wrong results. The Python runner read about 1,400 input batches ahead before returning the first result. Every batch is imported into the reader's single set of vectors. So the pass-through columns waiting for UDF results show the latest batch, not their own: result 0 got paired with id 1,383,000.
2. Early close crashes the task. When reader.next() returns false, DataSourceRDD's PartitionIterator.hasNext closes the reader immediately. NativePartitionReader.close() then calls root.close(), which empties the vectors that pending output batches still point to. A partition with a single batch is enough to trigger it.

I'd suggest to use ColumnarBatch types in SPARK-57468 to solve the batch type conversions elegantly.

[HyukjinKwon]

@FelixYBW do you kindly mind testing it with #59424?

[FelixYBW]

@FelixYBW do you kindly mind testing it with #59424?

Just tried, the same.

[FelixYBW]

Glad to see Spark community started to adding native support. We are thrilled to contribute and improve.

I quickly implemented a velox datasource PoC with int column only: https://github.com/FelixYBW/velox/tree/spark-native-datasource-poc/velox/experimental/spark_datasource

Here are some benchmark results. Single Parquet file with zstd compression, 553 MiB file (73,420,230 rows, 5 row groups), 11 columns all are INT

Benchmark is quite simple, scan all 11 columns then sum. parquet is on local disk, single core on a m7i EC2. Most of the time is actually spent on zstd decompression.

Here is a quick benchmark:

Configuration	Time (s)	Speedup vs. Spark
Spark	6.72	1.00x
Spark + Velox datasource	5.83	1.15x
Gluten scan only	10.03	0.67x
Gluten full	5.44	1.24x

@HyukjinKwon We need to disable arrow check (-Darrow.enable_unsafe_memory_access=true -Darrow.enable_null_check_for_get=false) in spark+velox datasource, otherwise perf is too bad. See if we can disable it by default. Then the PR can get expected improvement.

Gluten scan only performance is too bad due to the VeloxC2R operator. Looks like we need to use Spark's default one instead.

Velox datasource query plan is like this:

All done with Claude

[HyukjinKwon]

@FelixYBW Thank you for the report and the reproduction. The bug is in the Arrow columnar input path of Python UDFs (SPARK-56350): it kept the pass-through columns of a batch until the result of the batch came back from Python, after the source had refilled or released their vectors. @sarutak reproduced it without native data sources too, with ArrowCachedBatchSerializer. It is fixed in SPARK-60294, #59551, and #59424, which replaces this PR, adds a test of Arrow UDFs over a native data source.

[FelixYBW]

#59424, which replaces this PR, adds a test of Arrow UDFs over a native data source.

Do you mean you'll support FFM only for the native implementation?

#59551

Looks like the PR introduced a full copy of every pass-through column of every batch. To Python UDF, the extra copy isn't a big issue. But we should be able to find way to avoid it. Reusing data is one of the main considerations we introduced with the Arrow data format. We spent a lot of effort in Gluten to avoid any extra memcpy

[HyukjinKwon]

Do you mean you'll support FFM only for the native implementation?

Yeah that's currently what I'm thinking.

[FelixYBW]

Yeah that's currently what I'm thinking.

We also have plan to move Gluten to FFM API. But it's still too early as all customers still use jdk8-17.

[peterenescu]

Thanks @HyukjinKwon, this is really exciting. Wondering if there are future plans to offload more operators to native libraries beyond the table scan, such as filters, projections, aggregations and joins, up to a native boundary at the exchange. Gluten of course helps with this, but a native solution to help run plans in Velox would be amazing.

If so, I think @zhouyuan's request for a threading contract on the stream makes sense. The docs say Spark calls the stream from the task's thread, but a native engine would want to consume it directly from its own threads.

[FelixYBW]

Wondering if there are future plans to offload more operators to native libraries beyond the table scan,

Compared to more operators, the other missing parts may be more important and key to production:

• configurations, key to native performance
• memory management, memory allocation in offheap in native library
• metrics reported from the native library
• introducing refcount-based batch management, instead of current independent batch + reusable batch
• automatic batch type detection and transition SPARK-57468
• relational expressions of pushdown filter and agg; currently, the PR uses plain SQL. Better to reuse Spark Connect, Substrate, or any other 3rd-party protocol.
• task cancellation and cleanup of native library

and many others.