[SPARK-60101][SQL] Support native data sources with a frozen C core and versioned messages

[HyukjinKwon]

What changes were proposed in this pull request?

This is a draft for comparison with #59424, following @tdcmeehan's proposal on how the interface of native data sources grows. It is #59424 up to 51af65c, with its C interface, a table of typed functions, replaced by a frozen C core and versioned Protocol Buffers messages. Everything else is the same as in #59424: finding the libraries, the packages, the Data Source V2 adapter, the Arrow data path, the checks of the data, the docs and the tests.

The C core: spark_datasource.h declares six functions, which never change:

int32_t spark_ds_abi_version(void);
const char* spark_ds_last_error(void);
int spark_ds_open(const SparkDsHost* host, const uint8_t* request, int64_t length,
                  SparkDsSink reply, void* sink, void** handle);
int spark_ds_call(void* handle, const uint8_t* request, int64_t length,
                  SparkDsSink reply, void* sink);
int spark_ds_stream(void* handle, const uint8_t* request, int64_t length,
                    struct ArrowArrayStream* in, struct ArrowArrayStream* out);
void spark_ds_close(void* handle);

and SparkDsHost, which starts with its size and only gets fields appended, with the callbacks of Spark: log, which writes to the log of the driver or the executor, and is_interrupted(task_attempt_id), which tells a read or a write that its task was killed.

The messages: everything else is in spark_datasource.proto (sql/core/src/main/protobuf/org/apache/spark/sql/datasource/, also copied to the include directory of a distribution):

• OpenRequest opens a data source, a batch reader, a stream reader, a writer, a partition reader or a data writer, with everything that it needs: the name and the options of the data source, the schema, the write mode and the ID of a streaming query, the TaskContext, or the serialized state and partition. It also carries the version of Spark and of the messages. The OpenReply of a data source declares its capabilities, and Spark reports the operations that a data source does not declare as unsupported, without opening a reader or a writer.
• CallRequest and CallReply carry every control operation: InferSchema, PushPredicates (with the answer of the reader for each predicate, PredicateUse), PushLimit, PruneColumns, PlanPartitions, Serialize, InitialOffset, LatestOffset, PlanStreamPartitions, CommitOffset, Commit, Abort, CommitDataWriter and AbortDataWriter. An empty reply is valid for every request: a reader without pushdown replies to PushPredicates with nothing.
• StreamRequest exchanges the data, as Arrow C streams, without copying: Read initializes the output stream of a partition reader, and Write gives a data writer the data of its task as an input stream, which the library reads during the call, on the thread of the task. Spark exports the rows of the task as one stream, from DataWriter.writeAll, in batches that end at spark.sql.execution.arrow.maxRecordsPerBatch rows or spark.sql.execution.arrow.maxBytesPerBatch bytes.
• Schemas are Spark SQL types, as in Spark Connect, so that every control message can be forwarded unchanged to a library that runs in another process, for example to isolate it, while the data is converted between Arrow C streams and Arrow IPC.
• The messages use optimize_for = LITE_RUNTIME. With the full C++ runtime, the generated code registers spark_datasource.proto in a pool of the process: the C++ example of the docs, linked to the libprotobuf of Homebrew, aborted the JVM with File already exists in database: spark_datasource.proto when Spark loaded a second copy of it, the library of a package after the installed one. Lite messages register nothing. Libraries built against different versions of Protocol Buffers still conflict if they link it dynamically, so the docs ask for static linking, with hidden symbols.

Compatibility: the six functions and the fields of SparkDsHost never change. The messages only grow: a later release may add fields, messages and values of enums, but never changes or reuses the number of a field, both sides ignore the fields that they do not know, and Spark only sends a new kind of request to an object whose OpenReply declares that it supports it. SPARK_DS_ABI_VERSION stays 1.

How Spark calls the libraries: NativeLibrary.java binds the six functions by name with the FFM API, as #59424 does, checks the ABI version before it binds the others, allocates one SparkDsHost for each library, and never throws from the upcalls (SparkDsSink, log and is_interrupted). The tasks that use a library are registered by their attempt ID for is_interrupted. NativeDataSource.scala builds the requests and reads the replies, and NativeSchemas converts the schemas. The columnar adapter gets ArrowStreamDataWriter and ArrowRowReader, so that a native data writer gets the input of a task as one Arrow stream; an error of the rows, which the stream only reports as a message, is rethrown as it is.

	#59424 (function table)	This PR (messages)
What a library exports	spark_ds_init, which returns a table of 29 functions	6 functions, which never change
What grows	functions appended to the table, fields appended to its size-prefixed structs	fields and messages of spark_datasource.proto
What checks it	the C, C++ or Rust compiler, for each function	the generated code of the messages; the compiler checks the 6 functions
What a library needs	the header	the header, the .proto, and a Protocol Buffers compiler and runtime, such as prost or protobuf C++
Libraries in one process	share nothing but the header	each has its own Protocol Buffers runtime: lite messages, linked statically with hidden symbols for C++
A host in another process	maps each function to a message of its own	forwards the control messages unchanged
The binding of Spark	a method handle for each function, and the structs	6 method handles; the requests and replies in Scala
The C++ test library, without dependencies	fills the table	a minimal Protocol Buffers codec written by hand, of about 200 lines

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 open data sources, batch readers and partition readers, and reply to the other requests, such as the pushdown, with empty replies. They are the examples of the docs.

Rust (with arrow-rs and prost)

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"] }
prost = "0.14"

[build-dependencies]
prost-build = "0.14"

build.rs, which generates the code of the messages from the include directory of a Spark distribution with prost-build, which runs protoc:

// Generates the code of the messages from spark_datasource.proto, in the include directory of a
// Spark distribution. prost-build runs protoc, which must be installed.
fn main() -> std::io::Result<()> {
    let include = format!("{}/include", std::env::var("SPARK_HOME").expect("SPARK_HOME"));
    prost_build::compile_protos(&[format!("{include}/spark_datasource.proto")], &[include])
}

src/lib.rs:

//! 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, 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::{DataType, Field, Schema, SchemaRef};
use prost::Message;

/// The messages of `spark_datasource.proto`.
mod proto {
    include!(concat!(env!("OUT_DIR"), "/spark.datasource.rs"));
}

use proto::{call_reply, call_request, data_type, open_request, stream_request};

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

/// `SparkDsSink`: receives the reply of a function, and copies it.
type Sink = unsafe extern "C" fn(sink: *mut c_void, data: *const u8, length: i64);

/// What a handle refers to.
enum Object {
    DataSource,
    Reader { end: i64 },
    Partition { start: i64, end: i64 },
}

thread_local! {
    static LAST_ERROR: RefCell<CString> = RefCell::new(CString::default());
}

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

/// The schema of the data, as the Spark SQL types of the messages.
fn spark_schema() -> proto::StructType {
    let long = data_type::Kind::Long(data_type::Long {});
    proto::StructType {
        fields: vec![proto::StructField {
            name: "id".to_string(),
            data_type: Some(proto::DataType { kind: Some(long) }),
            nullable: false,
        }],
    }
}

/// Runs `f`, and reports its error or panic to Spark: neither may cross into the JVM.
fn call(f: impl FnOnce() -> Result<()>) -> c_int {
    let message = match catch_unwind(AssertUnwindSafe(f)) {
        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
}

unsafe fn decode<M: Message + Default>(request: *const u8, length: i64) -> Result<M> {
    Ok(M::decode(std::slice::from_raw_parts(request, length as usize))?)
}

unsafe fn reply(reply: Sink, sink: *mut c_void, message: &impl Message) {
    let bytes = message.encode_to_vec();
    reply(sink, bytes.as_ptr(), bytes.len() as i64);
}

#[no_mangle]
pub extern "C" fn spark_ds_abi_version() -> i32 {
    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 unsafe extern "C" fn spark_ds_open(
    _host: *const c_void,
    request: *const u8,
    length: i64,
    sink_function: Sink,
    sink: *mut c_void,
    handle: *mut *mut c_void,
) -> c_int {
    call(|| {
        let object = match decode::<proto::OpenRequest>(request, length)?.object {
            Some(open_request::Object::DataSource(_)) => {
                let capabilities = vec![proto::Capability::BatchRead as i32];
                reply(sink_function, sink, &proto::OpenReply { capabilities });
                Object::DataSource
            }
            Some(open_request::Object::BatchReader(reader)) => {
                let options = reader.data_source.unwrap_or_default().options;
                let end = options.get("end").map(|end| end.parse()).transpose()?;
                Object::Reader { end: end.unwrap_or(10) }
            }
            Some(open_request::Object::PartitionReader(reader)) => {
                let bytes = reader.partition;
                let start = i64::from_le_bytes(bytes[0..8].try_into()?);
                let end = i64::from_le_bytes(bytes[8..16].try_into()?);
                Object::Partition { start, end }
            }
            _ => return Err("this data source only supports batch reads".into()),
        };
        *handle = Box::into_raw(Box::new(object)) as *mut c_void;
        Ok(())
    })
}

#[no_mangle]
pub unsafe extern "C" fn spark_ds_call(
    handle: *mut c_void,
    request: *const u8,
    length: i64,
    sink_function: Sink,
    sink: *mut c_void,
) -> c_int {
    call(|| {
        let object = &*(handle as *const Object);
        let result = match (decode::<proto::CallRequest>(request, length)?.call, object) {
            (Some(call_request::Call::InferSchema(_)), _) => {
                let schema = Some(spark_schema());
                call_reply::Reply::InferSchema(proto::InferSchemaReply { schema })
            }
            (Some(call_request::Call::PlanPartitions(_)), Object::Reader { end }) => {
                let partitions = [(0, end / 2), (end / 2, *end)]
                    .iter()
                    .map(|(start, end)| [start.to_le_bytes(), end.to_le_bytes()].concat())
                    .collect();
                call_reply::Reply::PlanPartitions(proto::PlanPartitionsReply { partitions })
            }
            // The other requests, such as the pushdown, have empty replies.
            _ => return Ok(()),
        };
        reply(sink_function, sink, &proto::CallReply { reply: Some(result) });
        Ok(())
    })
}

#[no_mangle]
pub unsafe extern "C" fn spark_ds_stream(
    handle: *mut c_void,
    request: *const u8,
    length: i64,
    _input: *mut FFI_ArrowArrayStream,
    output: *mut FFI_ArrowArrayStream,
) -> c_int {
    call(|| {
        let Object::Partition { start, end } = &*(handle as *const Object) else {
            return Err("not a partition reader".into());
        };
        let stream = decode::<proto::StreamRequest>(request, length)?.stream;
        let Some(stream_request::Stream::Read(_)) = stream else {
            return Err("this data source only supports reads".into());
        };
        let ids = Int64Array::from_iter_values(*start..*end);
        let batch = RecordBatch::try_new(arrow_schema(), vec![Arc::new(ids)])?;
        let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema());
        std::ptr::write(output, FFI_ArrowArrayStream::new(Box::new(batches)));
        Ok(())
    })
}

#[no_mangle]
pub unsafe extern "C" fn spark_ds_close(handle: *mut c_void) {
    drop(Box::from_raw(handle as *mut Object));
}

Build and install, here on macOS:

SPARK_HOME=/path/to/spark 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++ and protobuf-lite)

cpp_range.cc:

// A native data source for Apache Spark, written in C++ with Arrow C++ and Protocol Buffers. It
// reads the numbers [0, end) as a column `id`, in two partitions. The option `end` defaults to 10.

#include <cstdint>
#include <cstring>
#include <memory>
#include <stdexcept>
#include <string>
#include <variant>

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

#include "spark_datasource.h"
#include "spark_datasource.pb.h"

namespace {

namespace sds = spark::datasource;

// What a handle refers to.
struct DataSource {};
struct Reader {
  int64_t end;
};
struct Partition {
  int64_t start;
  int64_t end;
};
using Object = std::variant<DataSource, Reader, Partition>;

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());
}

template <typename M>
M Decode(const uint8_t* request, int64_t length) {
  M message;
  if (!message.ParseFromArray(request, static_cast<int>(length))) {
    throw std::runtime_error("invalid request");
  }
  return message;
}

void Reply(SparkDsSink reply, void* sink, const google::protobuf::MessageLite& message) {
  std::string bytes = message.SerializeAsString();
  reply(sink, reinterpret_cast<const uint8_t*>(bytes.data()), static_cast<int64_t>(bytes.size()));
}

}  // namespace

extern "C" {

int32_t spark_ds_abi_version(void) { return SPARK_DS_ABI_VERSION; }

const char* spark_ds_last_error(void) { return last_error.c_str(); }

int spark_ds_open(const SparkDsHost*, const uint8_t* request, int64_t length, SparkDsSink reply,
                  void* sink, void** handle) {
  return Call([&] {
    auto open = Decode<sds::OpenRequest>(request, length);
    std::unique_ptr<Object> object;
    if (open.has_data_source()) {
      sds::OpenReply capabilities;
      capabilities.add_capabilities(sds::CAPABILITY_BATCH_READ);
      Reply(reply, sink, capabilities);
      object = std::make_unique<Object>(DataSource{});
    } else if (open.has_batch_reader()) {
      const auto& options = open.batch_reader().data_source().options();
      auto end = options.find("end");
      object = std::make_unique<Object>(
          Reader{end == options.end() ? 10 : std::stoll(end->second)});
    } else if (open.has_partition_reader()) {
      int64_t bounds[2];
      std::memcpy(bounds, open.partition_reader().partition().data(), sizeof(bounds));
      object = std::make_unique<Object>(Partition{bounds[0], bounds[1]});
    } else {
      throw std::runtime_error("this data source only supports batch reads");
    }
    *handle = object.release();
  });
}

int spark_ds_call(void* handle, const uint8_t* request, int64_t length, SparkDsSink reply,
                  void* sink) {
  return Call([&] {
    auto call = Decode<sds::CallRequest>(request, length);
    sds::CallReply result;
    if (call.has_infer_schema()) {
      auto* field = result.mutable_infer_schema()->mutable_schema()->add_fields();
      field->set_name("id");
      field->mutable_data_type()->mutable_long_();
      field->set_nullable(false);
    } else if (call.has_plan_partitions()) {
      int64_t end = std::get<Reader>(*static_cast<Object*>(handle)).end;
      int64_t bounds[2][2] = {{0, end / 2}, {end / 2, end}};
      for (const auto& partition : bounds) {
        result.mutable_plan_partitions()->add_partitions(partition, sizeof(partition));
      }
    } else {
      return;  // The other requests, such as the pushdown, have empty replies.
    }
    Reply(reply, sink, result);
  });
}

int spark_ds_stream(void* handle, const uint8_t* request, int64_t length, struct ArrowArrayStream*,
                    struct ArrowArrayStream* out) {
  return Call([&] {
    if (!Decode<sds::StreamRequest>(request, length).has_read()) {
      throw std::runtime_error("this data source only supports reads");
    }
    const auto& partition = std::get<Partition>(*static_cast<Object*>(handle));
    arrow::Int64Builder ids;
    for (int64_t id = partition.start; id < partition.end; 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));
  });
}

void spark_ds_close(void* handle) { delete static_cast<Object*>(handle); }

}  // extern "C"

Generate the messages, then build (Arrow C++ 24 requires C++20) and install, or package as above:

protoc --cpp_out=. -I"$SPARK_HOME/include" "$SPARK_HOME/include/spark_datasource.proto"
c++ -std=c++20 -O2 -shared -fPIC -I. -I"$SPARK_HOME/include" -I"$ARROW_HOME/include" \
  cpp_range.cc spark_datasource.pb.cc $(pkg-config --cflags --libs protobuf-lite) \
  -L"$ARROW_HOME/lib" -larrow -o libspark_datasource_cpp_range.dylib
cp libspark_datasource_cpp_range.dylib /usr/local/lib/

Usage is the same as in #59424, 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:

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

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?

To compare the two shapes of version 1 of the interface, which library binaries get built against, before Spark 4.4.0: see the discussion in #59424. A library that runs in another process, as the UDF worker protocol (SPARK-55278) does for UDFs, could get the same control messages, while #59424 would need a message for each of its functions.

Does this PR introduce any user-facing change?

The same as #59424, with this interface: native data sources implement spark_datasource.h and spark_datasource.proto, which a Spark distribution contains in its include directory.

How was this patch tested?

All the tests of #59424, adapted to the messages, with the C++ test library implementing them with its own minimal Protocol Buffers codec, so that it still needs nothing but a C++ compiler. New tests:

• ignore the fields of messages that they do not know: Spark ignores a field of a later version in every reply of a library, and a library ignores a field of a later version in a request.
• write the data of a task in one stream: a data writer gets the 10 rows of its task in one stream of 4 batches, with spark.sql.execution.arrow.maxRecordsPerBatch at 3.
• report the errors of the data that a library writes: a UDF that fails while the library reads the rows fails the write with the error of the UDF, not with the error that the library returns, and the write is aborted.
• reject the libraries that do not export the functions of the interface: a library without spark_ds_open.

Locally on macOS (osx-aarch64), with Java 26 (Spark built with Java 17): NativeDataSourceSuite, NativeDataSourceClusterSuite (executors on Java 26 too), ColumnarDataSourceSuite and NativeDataSourceConnectSuite pass, and with Java 17, NativeDataSourceSuite passes with the tests that load a library skipped. The Rust example of the docs (arrow-rs 58, prost 0.14) and the C++ example (Arrow C++ 24, protobuf-lite 35) were built with the documented commands and read by Spark on Java 26: installed and in a package, with errors of the library reported as NATIVE_DATA_SOURCE_ERROR. The C++ test library builds without warnings with Apple clang (-Wall -Wextra) in all its variants. checkstyle and scalastyle of the changed modules pass.

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.