HyukjinKwon commented on code in PR #59313:
URL: https://github.com/apache/spark/pull/59313#discussion_r4265011347
##########
sql/api/src/main/scala/org/apache/spark/sql/Artifact.scala:
##########
@@ -48,13 +52,15 @@ private[sql] object Artifact {
newJarArtifact(targetFilePath, storage)
case cf if cf.endsWith(".class") =>
newClassArtifact(targetFilePath, storage)
+ case pkg if pkg.endsWith(NATIVE_DATA_SOURCE_PACKAGE_EXTENSION) =>
Review Comment:
Thanks, fixed in aa0f3ba4dd77. The five Scaladocs in
`sql/api/.../SparkSession.scala` now say that since 4.4.0, `.sparkpkg` files
are native data source packages, which make the native data sources that they
contain available to the session. For the two methods with a target,
`.sparkpkg` is listed with `.jar` and `.class`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,794 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.datasources.v2.ffi
+
+import java.lang.ref.{Cleaner, WeakReference}
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicLong
+
+import scala.jdk.CollectionConverters._
+import scala.util.control.NonFatal
+
+import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema,
CDataDictionaryProvider, Data}
+import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.MapVector
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
+
+import org.apache.spark.{SparkException, SparkThrowable,
SparkUnsupportedOperationException}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.expressions.filter.Predicate
+import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader}
+import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage}
+import org.apache.spark.sql.errors.{QueryCompilationErrors,
QueryExecutionErrors}
+import
org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarDataSource,
ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{DataType, StructType}
+import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch,
ColumnVector}
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.util.Utils
+
+/**
+ * A [[ColumnarDataSource]] implemented by a native library, through
+ * `org.apache.spark.sql.datasource.NativeBridge`.
+ *
+ * @param location where the library is on the driver. The readers and the
writers distribute it
+ * to the executors of the active session when they are
created.
+ * @param name the name of the data source, as listed in the manifest of its
package or in lower
+ * case for an installed library
+ */
+class NativeDataSource(
+ location: NativeLibraryLocation,
+ name: String,
+ options: CaseInsensitiveStringMap)
+ extends ColumnarDataSource {
+ import NativeCalls._
+
+ override def schema(): StructType = {
+ val library = NativeLibraries.get(location)
+ withDataSource(library, "infer the schema of") { dataSource =>
+ call(name, "infer the schema of") {
+ optional(NativeArrow.importSchema(address =>
library.schema(dataSource, address)))
+ }.getOrElse(throw
QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ }
+ }
+
+ override def reader(schema: StructType): ColumnarReader = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of",
"DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") {
+ dataSource =>
+ NativeArrow.exportSchema(schema)(address =>
library.createReader(dataSource, address))
+ }
+ new NativeDataSourceReader(
+ executorLocation, name, schema, new NativeHandle(handle,
library.closeReader))
+ }
+
+ override def streamReader(schema: StructType): ColumnarStreamReader = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of",
"DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") {
+ dataSource =>
+ NativeArrow.exportSchema(schema)(address =>
library.createStreamReader(dataSource, address))
+ }
+ new NativeDataSourceStreamReader(
+ executorLocation, name, schema, new NativeHandle(handle,
library.closeStreamReader))
+ }
+
+ override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter
= {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = new NativeHandle(
+ create(library, "plan a write to",
"DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource, address, overwrite)
+ }
+ },
+ library.closeWriter)
+ new NativeDataSourceWriter(executorLocation, name, handle,
serializeWriter(library, handle))
+ }
+
+ override def streamWriter(
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): ColumnarStreamWriter = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = new NativeHandle(
+ create(library, "plan a write to",
"DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamWriter(dataSource, address, overwrite)
+ }
+ },
+ library.closeWriter)
+ val writer = new NativeDataSourceStreamWriter(
+ executorLocation, name, handle, serializeWriter(library, handle))
+ NativeStreamWriters.register(queryId, handle)
+ writer
+ }
+
+ /** Runs `f` with a new native data source, which is closed afterwards. */
+ private def withDataSource[T](library: NativeLibrary, action: String)(f:
Long => T): T = {
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey),
entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource))
+ }
+ }
+
+ /** Creates a reader or a writer with `f`, which fails if the library does
not implement it. */
+ private def create(library: NativeLibrary, action: String, unsupportedError:
String)(
+ f: Long => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ call(name, action)(optional(f(dataSource))).getOrElse {
Review Comment:
Thanks, fixed in aa0f3ba4dd77. When one of the six `create` functions
(`createDataSource`, `createReader`, `createStreamReader`, `createWriter`,
`createStreamWriter`, `createDataWriter`) returns 0, Spark fails with a
`NATIVE_DATA_SOURCE_ERROR` that names the function, before it uses the handle.
The Handles convention says that Spark reports a handle of 0 as an error. New
test: "reject handles of 0", with a test library option that makes each of
these functions return 0 without an error.
##########
docs/sql-native-data-sources.md:
##########
@@ -0,0 +1,602 @@
+---
+layout: global
+title: Native Data Sources
+displayTitle: Native Data Sources
+license: |
+ Licensed to the Apache Software Foundation (ASF) under one or more
+ contributor license agreements. See the NOTICE file distributed with
+ this work for additional information regarding copyright ownership.
+ The ASF licenses this file to You under the Apache License, Version 2.0
+ (the "License"); you may not use this file except in compliance with
+ the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+---
+
+* Table of contents
+{:toc}
+
+## Overview
+
+A native data source is a data source implemented in native code, such as Rust
or C++, without
+any JVM code. It is a shared library that implements the
+[`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html)
binary interface.
+Spark finds the library automatically when it is installed like other native
libraries, or when a
+session adds it in a package, and loads it when a query uses the data source.
Spark adapts it to
+[Data Source V2](sql-v2-data-sources.html), so it supports:
+
+- batch and micro-batch streaming reads,
+- batch and streaming writes, appending to or overwriting the existing data,
+- predicate, column and limit pushdown.
+
+Data is exchanged only as columnar batches, through the
+[Arrow C data
interface](https://arrow.apache.org/docs/format/CDataInterface.html), so Spark
reads
+the data that a library returns without copying it.
+
+## How It Works
+
+A native data source library implements the `native` methods of `NativeBridge`
as JNI functions,
+named after the class, such as
`Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource`.
+Spark loads each library into its own copy of the class, so any number of
native data sources can
+be loaded side by side.
+
+Spark calls the functions of the library as follows:
+
+- **Data sources**: for each schema inference, scan and write, Spark creates a
data source on the
+ driver with `createDataSource`, which receives the options of the user, and
closes it right
+ after. When the user does not specify a schema, Spark asks for it with
`schema`.
+- **Batch reads**: on the driver, Spark creates a reader with `createReader`,
pushes operations down
+ with `pushPredicates`, `pushLimit` and `pruneColumns`, plans the partitions
with `partitions`,
+ and serializes the state of the reader with `serializeReader`. Then it calls
`read` on the
+ executors, in one task for each partition, and `read` exports an
`ArrowArrayStream`.
+- **Streaming reads**: on the driver, Spark creates a stream reader with
`createStreamReader`. For
+ each micro-batch, it asks for the `latestOffset`, plans the partitions with
`streamPartitions`,
+ and calls `read` on the executors. Offsets are JSON strings defined by the
library.
+- **Writes**: on the driver, Spark creates a writer with `createWriter` or
`createStreamWriter`. On
+ the executors, each task creates a data writer with `createDataWriter`,
writes the batches of its
+ data with `write`, and calls `commitDataWriter`, or `abortDataWriter` if the
task or the commit
+ fails. Then the driver calls `commit` with the messages of the tasks, or
`abort` if the write
+ failed. A micro-batch streaming query creates a new stream writer for each
micro-batch, and Spark
+ closes the previous one when it creates the next one. Because Spark is not
told when a streaming
+ query stops, the last stream writer of a query is closed once it is no
longer referenced, or
+ when the query is restarted.
+
+The library defines the format of the partitions, of the state of its readers
and writers, and of
+the commit messages, which Spark sends to the executors and back as byte
arrays. A library reports
+an error by throwing a Java exception, for example with the JNI function
`ThrowNew`, and Spark
+reports it as a `NATIVE_DATA_SOURCE_ERROR`.
+
+A library only exports the functions of the operations it supports. If a
function is missing,
+Spark treats the operation as unsupported, except for the pushdown and
`serialize` functions, which
+are optional:
+
+| Operation | Required functions | Optional functions |
+|-----------|--------------------|--------------------|
+| All | `abiVersion`, `createDataSource`, `closeDataSource` | `schema` |
+| Batch read | `createReader`, `partitions`, `read`, `closeReader` |
`pushPredicates`, `pushLimit`, `pruneColumns`, `serializeReader` |
+| Streaming read | `createStreamReader`, `initialOffset`, `latestOffset`,
`streamPartitions`, `read`, `closeStreamReader` | `serializeStreamReader`,
`commitOffset` |
+| Batch write | `createWriter`, `commit`, `abort`, `closeWriter`,
`createDataWriter`, `write`, `commitDataWriter`, `abortDataWriter` |
`serializeWriter` |
+| Streaming write | `createStreamWriter`, `commit`, `abort`, `closeWriter`,
`createDataWriter`, `write`, `commitDataWriter`, `abortDataWriter` |
`serializeWriter` |
+
+The [Javadoc of
`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html)
+specifies each function, including the ownership of the Arrow structs and the
JSON format of the
+pushed predicates. A predicate that a library reports as evaluated completely
must follow the
+semantics of Spark SQL, for example for NaN values, which are equal to each
other and greater
+than any other value. The arrays that a library exports must have an offset of
0: copy a sliced
+array, such as the result of `RecordBatch::slice` in arrow-rs, before
exporting it. Spark also
+checks that each array has at least as many values as its parent refers to,
but it cannot check
+that the buffers of the arrays hold those values, because the C data interface
does not carry the
+sizes of buffers: a library that exports a buffer that is too short makes
Spark read past its end.
+
+The `get_next` function of the stream that `read` exports may block, for
example on network I/O,
+but Spark cannot interrupt native code: when a task is killed, Spark stops
reading only once
+`get_next` returns. Use timeouts for blocking operations, or enable the task
reaper with
+`spark.task.reaper.enabled` and `spark.task.reaper.killTimeout`, which kills
an executor whose
+killed task does not end.
+
+### Compatibility
+
+The JVM binds the functions of a library by their names only, so the names and
the signatures of
+the methods of `NativeBridge` are the interface. Its version, 1, is the
`abiVersion` that a
+library returns and that the manifest of a package declares. Within a version,
the signature and
+the semantics of an existing function never change, and new functions are only
added as optional
+functions, so a library keeps working with later releases of Spark. New values
may appear where
+the Javadoc says so: for example, later releases may push down predicates with
other names, which
+a library must not report as evaluated completely. An incompatible change
comes with a new
+version, with new names or a new class. Spark supports version 1 at least in
all the releases of
+Spark 4 from 4.4.0.
+
+### Limitations
+
+Version 1 of the interface does not cover these features of Data Source V2
yet. They can be added
+later as optional functions:
+
+- statistics: Spark estimates the size of a native scan with
`spark.sql.defaultSizeInBytes`, so a
+ native scan is never the broadcast side of a join, unless the query has a
broadcast hint;
Review Comment:
Thanks, both fixed in aa0f3ba4dd77:
1. The Limitations section now says that Spark does not plan a native scan
as the broadcast side of a join without a hint, but that adaptive query
execution can still broadcast it at run time, from the size of a shuffle of the
scan.
2. Rather than rewording the sentence, Spark now also logs a package of an
unsupported `abiVersion` when it reads it. That is once per version of the
file, since the result is cached. So the docs page and the PR description are
now accurate. New test: "log the packages of a version that Spark does not
support".
##########
sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java:
##########
@@ -0,0 +1,442 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.datasource;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * The binary interface between Spark and data sources implemented in native
code, such as Rust
+ * or C++.
+ * <p>
+ * A native data source is a shared library that implements the {@code native}
methods of this
+ * class as JNI functions, for example
+ * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}.
Spark loads each
+ * library into its own copy of this class, so any number of libraries that
export the same
+ * functions can be loaded side by side. Spark adapts a native data source to
Data Source V2, with
+ * batch and micro-batch streaming reads, batch and streaming writes, and
predicate, column and
+ * limit pushdown.
+ * <p>
+ * Spark finds the library of a data source automatically when it is installed
like other native
+ * libraries and named after the data source: the library of {@code my_source}
is
+ * {@code spark_datasource_my_source}, such as {@code
libspark_datasource_my_source.so} in
+ * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of
+ * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux,
and in the
+ * {@code lib} directory of each installation prefix in {@code PATH}. A
library can also be
+ * distributed in a native data source package: a zip file with the extension
{@code .sparkpkg}
+ * that contains a manifest named {@code spark-native-datasource.json} and the
library built for
+ * one or more platforms, added to a session with {@code spark.addArtifact} or
under the paths in
+ * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources"
page of the Spark SQL
+ * guide for details and examples.
+ *
+ * <h2>Lifecycle</h2>
+ * Spark creates a data source with {@link #createDataSource} on the driver
for each schema
+ * inference, scan and write, and closes it right after.
+ * <ul>
+ * <li><b>Batch reads.</b> On the driver, Spark creates a reader with {@link
#createReader},
+ * pushes operations down by calling, each at most once and in this order,
+ * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns},
plans the partitions
+ * with {@link #partitions}, serializes the state of the reader with {@link
#serializeReader},
+ * and closes it. Then it calls {@link #read} on the executors, in one task
for each
+ * partition.</li>
+ * <li><b>Streaming reads.</b> On the driver, Spark creates a stream reader
with
+ * {@link #createStreamReader} and keeps it until the streaming query stops.
For each
+ * micro-batch, it asks for the {@link #latestOffset}, plans the partitions
with
+ * {@link #streamPartitions} and serializes the state with {@link
#serializeStreamReader}, then
+ * calls {@link #read} on the executors. Offsets are JSON strings defined by
the library. When
+ * a streaming query restarts, from its checkpoint, Spark creates a new
stream reader and may
+ * call {@link #streamPartitions} again for the last range that it planned
before, so the
+ * partitions of a range must read the same data each time.</li>
+ * <li><b>Writes.</b> On the driver, once it computed the data to write,
right before it starts
+ * the tasks of the write (or of a micro-batch), Spark creates a writer with
{@link #createWriter}
+ * or {@link #createStreamWriter}, and serializes its state with {@link
#serializeWriter}. If
+ * computing the data fails, Spark creates no writer. On the executors, each
task creates a data
+ * writer with {@link #createDataWriter}, calls {@link #write} for each
batch of at most
+ * {@code spark.sql.execution.arrow.maxRecordsPerBatch} rows, which also
ends once it holds
+ * {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, and calls
{@link #commitDataWriter},
+ * or {@link #abortDataWriter} if the task or the commit fails. Then the
driver calls
+ * {@link #commit} with the messages of the tasks if they all succeeded, or
{@link #abort}
+ * otherwise, including when {@link #commit} fails, and closes the writer.
Several attempts of a
+ * task may write the same partition, after a failure or with speculative
execution, each with a
+ * data writer of its own. Spark lets the attempts of a partition commit one
at a time, and passes
+ * to {@link #commit} the message of one committed attempt for each
partition, so the write
+ * consists of the data of these attempts only. It aborts the data writers
of the other attempts,
+ * except on an executor that is lost, where a data writer is neither
committed nor aborted, and
+ * may even have committed before its message reached the driver. A
streaming write commits or
+ * aborts each micro-batch (epoch). Spark may commit the same epoch more
than once, for example
+ * when a streaming query restarts after a failure, so committing an epoch
that is already
+ * committed must have no effect, for exactly-once output. A micro-batch
streaming query creates a
+ * new stream writer for each micro-batch, and Spark closes the previous
stream writer of the
+ * query when the query creates a new one. Spark is not told when a
streaming query stops, so it
+ * closes the last stream writer of the query once the writer is no longer
referenced, or when the
+ * query is restarted.</li>
+ * </ul>
+ *
+ * <h2>Conventions</h2>
+ * <ul>
+ * <li><b>Handles.</b> The {@code create} methods return an opaque, non-zero
handle, typically a
+ * pointer. Spark releases each handle exactly once with the matching {@code
close} method, or
+ * with {@link #commitDataWriter} or {@link #abortDataWriter} for data
writers, though not
+ * always right after the last use of a stream writer (see Writes above). A
reader or
+ * writer handle must not depend on the data source handle it was created
from, because Spark
+ * closes the data source handle right after creating it.</li>
+ * <li><b>State.</b> 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.</li>
+ * <li><b>Errors.</b> To report an error, throw a Java exception, for
example with the JNI
+ * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and
return any value. Spark
+ * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.</li>
+ * <li><b>Arrow data.</b> Schemas and data are exchanged with the
+ * <a href="https://arrow.apache.org/docs/format/CDataInterface.html";>Arrow
C data
+ * interface</a>. A {@code long} argument named {@code ...Address} is the
address of a C struct
+ * allocated by Spark. For an output struct, the library initializes it and
Spark takes
+ * ownership. For an input struct, the library may move it to take
ownership; otherwise, Spark
+ * releases it after the call returns, so the library must not use it
afterwards. The arrays
+ * that the library exports must have an offset of 0, at every level: Arrow
Java ignores the
+ * offsets of the arrays it imports, so Spark rejects an array with another
offset, such as a
+ * zero-copy slice. Copy a sliced array before exporting it. Spark also
rejects an array with
+ * fewer values than its parent refers to, such as a column shorter than its
batch. The buffers
+ * of an array must hold all of its values, as the C data interface
requires: it does not carry
+ * the sizes of the buffers, so Spark cannot check them, and reads past the
end of a buffer that
+ * is too short.</li>
+ * <li><b>Schemas.</b> Spark types map to Arrow types as in {@code
spark.sql.execution.arrow}:
+ * for example, {@code bigint} is {@code int64}, {@code string} is {@code
utf8},
+ * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code
timestamp_ntz} is
+ * {@code timestamp[us]} without one. Data read from a library may also use
+ * {@code large_utf8}, {@code utf8_view} and the other variants of the same
type, and any names
+ * for the fields of the entries of a map, but not unsigned integers, {@code
decimal256}, list
+ * views or dictionary encoding.</li>
+ * <li><b>Options.</b> Option keys are passed in lower case, because Spark
matches them
+ * case-insensitively.</li>
+ * <li><b>Optional methods.</b> A library only needs to export the functions
for the operations
+ * it supports. If a function is missing, Spark treats the operation as
unsupported, except for
+ * the pushdown methods and the {@code serialize} methods, which then push
nothing down and
+ * return no state.</li>
+ * <li><b>Threads.</b> Spark calls the library from multiple threads, and
calls on different
+ * handles may run concurrently. Calls on the same handle never run
concurrently. A handle may
+ * be released from a thread other than the threads that used it, such as a
JVM
+ * {@code Cleaner} thread when Spark releases a handle once it is no longer
referenced, for
+ * example the last stream writer of a streaming query.</li>
+ * </ul>
+ *
+ * <h2>Compatibility</h2>
+ * The JVM binds the functions of a library by their names only, so the name
of this class, the
+ * names of its methods and their signatures are the interface. Within an
{@link #ABI_VERSION}:
+ * <ul>
+ * <li>the signature and the semantics of an existing method never
change;</li>
+ * <li>new methods may be added, only as optional methods, so that a library
that does not
+ * export them keeps working as before;</li>
+ * <li>new values may appear where this documentation says so, such as the
names of pushed
+ * predicates.</li>
+ * </ul>
+ * An incompatible change uses a new {@link #ABI_VERSION}, with new names or a
new class. A
+ * library reports the version it implements with {@link #abiVersion}, and
Spark rejects a library
+ * of a version it does not support. Spark supports version 1 at least in all
the releases of
+ * Spark 4 from 4.4.0.
+ *
+ * @since 4.4.0
+ */
+@Evolving
+public final class NativeBridge {
+
+ /**
+ * The version of this interface. It changes when the interface changes
incompatibly: see
+ * Compatibility above.
+ */
+ public static final int ABI_VERSION = 1;
+
+ private NativeBridge() {}
+
+ /**
+ * Loads the library that implements the native methods of this copy of the
class. Spark calls
+ * it on the copy of the class it creates for the library.
+ */
+ static void load(String path) {
+ System.load(path);
+ }
+
+ /** Returns the version of this interface that the library implements,
{@link #ABI_VERSION}. */
+ public static native int abiVersion();
+
+ //
---------------------------------------------------------------------------------------------
+ // Data source (driver)
+ //
---------------------------------------------------------------------------------------------
+
+ /**
+ * Creates a data source with the options given by the user, such as
+ * {@code spark.read.option("key", "value")}.
+ *
+ * @param name the name of the data source, as listed in the package
manifest. A library can
+ * implement several data sources.
+ * @param optionKeys the option keys, in lower case
+ * @param optionValues the option values, in the same order as the keys
+ * @return the data source handle
+ */
+ public static native long createDataSource(
+ String name, String[] optionKeys, String[] optionValues);
+
+ /**
+ * Exports the schema of the data. Spark calls it only when the user does
not specify a schema.
+ *
+ * @param schemaAddress the address of the {@code ArrowSchema} to initialize
with a struct type
+ * whose fields are the columns
+ */
+ public static native void schema(long dataSource, long schemaAddress);
+
+ /**
+ * Creates a reader for a batch scan.
+ *
+ * @param schemaAddress the address of the {@code ArrowSchema} of the schema
to read: the one the
+ * user specified, or the one returned by {@link
#schema}
+ * @return the reader handle
+ */
+ public static native long createReader(long dataSource, long schemaAddress);
+
+ /**
+ * Creates a reader for a micro-batch streaming scan.
+ *
+ * @param schemaAddress the address of the {@code ArrowSchema} of the schema
to read: the one the
+ * user specified, or the one returned by {@link
#schema}
+ * @return the stream reader handle
+ */
+ public static native long createStreamReader(long dataSource, long
schemaAddress);
+
+ /**
+ * Creates a writer for a batch write.
+ *
+ * @param schemaAddress the address of the {@code ArrowSchema} of the data
to write
+ * @param overwrite whether to replace the existing data (save mode
"overwrite") instead of
+ * appending to it
+ * @return the writer handle
+ */
+ public static native long createWriter(long dataSource, long schemaAddress,
boolean overwrite);
+
+ /**
+ * Creates a writer for a streaming write.
+ *
+ * @param schemaAddress the address of the {@code ArrowSchema} of the data
to write
+ * @param overwrite whether to replace the existing data in each micro-batch
(output mode
+ * "complete") instead of appending to it
+ * @return the writer handle
+ */
+ public static native long createStreamWriter(
+ long dataSource, long schemaAddress, boolean overwrite);
+
+ /** Releases a data source handle. */
+ public static native void closeDataSource(long dataSource);
+
+ //
---------------------------------------------------------------------------------------------
+ // Batch reader (driver)
+ //
---------------------------------------------------------------------------------------------
+
+ /**
+ * Pushes predicates down to the reader. The predicates are combined with
AND.
+ * <p>
+ * Each predicate is a JSON expression tree. An expression is one of:
+ * <ul>
+ * <li>{@code {"type": "column", "name": ["a", "b"]}}: a column. A nested
field has a name
+ * part for each level.</li>
+ * <li>{@code {"type": "literal", "dataType": "bigint", "value": 1}}: a
literal. The data type
+ * is a Spark SQL type name such as {@code int}, {@code string} or {@code
decimal(10,2)}. The
+ * value is {@code null}, a JSON boolean or number for boolean and numeric
types, a string
+ * for {@code string}, a base64 string for {@code binary}, a decimal
string for
+ * {@code decimal}, the days since the epoch for {@code date}, and the
microseconds since the
+ * epoch for {@code timestamp} and {@code timestamp_ntz}.</li>
+ * <li>{@code {"type": "function", "name": ">", "children": [...]}}: a
predicate, with the
+ * name it has in {@code
org.apache.spark.sql.connector.expressions.filter.Predicate}. The
+ * names are {@code =}, {@code <>}, {@code <=>} (equal, with null equal to
null), {@code <},
+ * {@code <=}, {@code >} and {@code >=} with two children, {@code IN} with
the value to look
+ * for and then the values to look in, {@code IS_NULL} and {@code
IS_NOT_NULL} with one
+ * child, {@code STARTS_WITH}, {@code ENDS_WITH} and {@code CONTAINS} with
the string and the
+ * string to look for, {@code AND} and {@code OR} with two predicates, and
{@code NOT} with
+ * one. The children of the other functions are columns, literals or
predicates.</li>
+ * </ul>
+ * Spark does not pass predicates that it cannot express this way, such as
predicates that use
+ * arithmetic or other functions, nor predicates on columns that contain
strings with a
+ * non-binary collation, such as {@code UTF8_LCASE}: strings are compared by
their UTF-8 bytes.
+ * A later version of Spark may pass predicates with other names without
changing
+ * {@link #ABI_VERSION}, so a reader must not report a predicate as
evaluated completely if it
+ * contains a name that the reader does not know, or a data type that it
does not support.
+ * <p>
+ * A predicate that the reader evaluates completely must follow the
semantics of Spark SQL. In
+ * particular, for floating-point columns, NaN is equal to NaN and greater
than any other value,
+ * including positive infinity, and {@code -0.0} is equal to {@code 0.0},
unlike IEEE 754
+ * comparisons. For example, {@code d > 1.0} is true for a NaN value of
{@code d}. Comparisons
+ * with null are null, so a row only matches if the predicate is true.
+ *
+ * @return for each predicate, whether the reader evaluates it completely,
so that Spark does
+ * not have to evaluate it again. The reader can still use the other
predicates to skip
+ * data.
Review Comment:
Thanks for the detailed write-up. On 1 to 3, all in aa0f3ba4dd77:
1. Spark now evaluates all the pushed predicates again after reading, also
when a library reports them as evaluated completely:
`ColumnarScanBuilder.pushPredicates` returns all of them as post-scan filters.
The answer of the library only shows in the plan (`PushedPredicates`). So a
library can no longer return extra rows; it can only use the predicates to skip
data. As a side effect, Spark only pushes a limit down when there is no filter,
which the Javadoc of `pushLimit` says. Trusting a library can come later as an
opt-in.
2. The Javadoc of `pushPredicates` and the docs page state the rule: a
reader must never drop or skip a row for which the predicates are true under
Spark SQL. They spell out NaN, `-0.0` (including the total orders of
`f64::total_cmp` and the arrow-rs kernels), nulls, and the UTF-8 byte order of
strings, which differs from UTF-16 order outside the BMP. They also give the
rule for using a part of a predicate: take the parts that it does not support
as true, but not under a `NOT`.
3. The jar of Spark SQL now contains
`org/apache/spark/sql/datasource/predicate-cases.json`. It has 28 predicates as
Spark pushes them, with the rows of a table for which Spark finds them true.
They cover NaN, `-0.0`, infinities, nulls under `NOT` and `IN`, `<=>` with a
value and with null, `= NULL`, `OR`, conjunctions, and strings outside the BMP,
compared by their UTF-8 bytes. A golden-file test in `NativeDataSourceSuite`
generates it from the evaluation of Spark itself, so it stays in sync with
Spark.
On the protocol in your second comment: with Spark keeping every filter in
version 1, the boolean only feeds `EXPLAIN`. So declaring the supported names
and types, splitting conjuncts in Spark, and a three-state answer can come
later as optional functions, without changing `pushPredicates`, together with
the opt-in to trust the library.
##########
sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java:
##########
@@ -0,0 +1,442 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.datasource;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * The binary interface between Spark and data sources implemented in native
code, such as Rust
+ * or C++.
+ * <p>
+ * A native data source is a shared library that implements the {@code native}
methods of this
+ * class as JNI functions, for example
+ * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}.
Spark loads each
+ * library into its own copy of this class, so any number of libraries that
export the same
+ * functions can be loaded side by side. Spark adapts a native data source to
Data Source V2, with
+ * batch and micro-batch streaming reads, batch and streaming writes, and
predicate, column and
+ * limit pushdown.
+ * <p>
+ * Spark finds the library of a data source automatically when it is installed
like other native
+ * libraries and named after the data source: the library of {@code my_source}
is
+ * {@code spark_datasource_my_source}, such as {@code
libspark_datasource_my_source.so} in
+ * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of
+ * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux,
and in the
+ * {@code lib} directory of each installation prefix in {@code PATH}. A
library can also be
+ * distributed in a native data source package: a zip file with the extension
{@code .sparkpkg}
+ * that contains a manifest named {@code spark-native-datasource.json} and the
library built for
+ * one or more platforms, added to a session with {@code spark.addArtifact} or
under the paths in
+ * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources"
page of the Spark SQL
+ * guide for details and examples.
+ *
+ * <h2>Lifecycle</h2>
+ * Spark creates a data source with {@link #createDataSource} on the driver
for each schema
+ * inference, scan and write, and closes it right after.
+ * <ul>
+ * <li><b>Batch reads.</b> On the driver, Spark creates a reader with {@link
#createReader},
+ * pushes operations down by calling, each at most once and in this order,
+ * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns},
plans the partitions
+ * with {@link #partitions}, serializes the state of the reader with {@link
#serializeReader},
+ * and closes it. Then it calls {@link #read} on the executors, in one task
for each
+ * partition.</li>
+ * <li><b>Streaming reads.</b> On the driver, Spark creates a stream reader
with
+ * {@link #createStreamReader} and keeps it until the streaming query stops.
For each
+ * micro-batch, it asks for the {@link #latestOffset}, plans the partitions
with
+ * {@link #streamPartitions} and serializes the state with {@link
#serializeStreamReader}, then
+ * calls {@link #read} on the executors. Offsets are JSON strings defined by
the library. When
+ * a streaming query restarts, from its checkpoint, Spark creates a new
stream reader and may
+ * call {@link #streamPartitions} again for the last range that it planned
before, so the
+ * partitions of a range must read the same data each time.</li>
+ * <li><b>Writes.</b> On the driver, once it computed the data to write,
right before it starts
+ * the tasks of the write (or of a micro-batch), Spark creates a writer with
{@link #createWriter}
+ * or {@link #createStreamWriter}, and serializes its state with {@link
#serializeWriter}. If
+ * computing the data fails, Spark creates no writer. On the executors, each
task creates a data
+ * writer with {@link #createDataWriter}, calls {@link #write} for each
batch of at most
+ * {@code spark.sql.execution.arrow.maxRecordsPerBatch} rows, which also
ends once it holds
+ * {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, and calls
{@link #commitDataWriter},
+ * or {@link #abortDataWriter} if the task or the commit fails. Then the
driver calls
+ * {@link #commit} with the messages of the tasks if they all succeeded, or
{@link #abort}
+ * otherwise, including when {@link #commit} fails, and closes the writer.
Several attempts of a
+ * task may write the same partition, after a failure or with speculative
execution, each with a
+ * data writer of its own. Spark lets the attempts of a partition commit one
at a time, and passes
+ * to {@link #commit} the message of one committed attempt for each
partition, so the write
+ * consists of the data of these attempts only. It aborts the data writers
of the other attempts,
+ * except on an executor that is lost, where a data writer is neither
committed nor aborted, and
+ * may even have committed before its message reached the driver. A
streaming write commits or
+ * aborts each micro-batch (epoch). Spark may commit the same epoch more
than once, for example
+ * when a streaming query restarts after a failure, so committing an epoch
that is already
+ * committed must have no effect, for exactly-once output. A micro-batch
streaming query creates a
+ * new stream writer for each micro-batch, and Spark closes the previous
stream writer of the
+ * query when the query creates a new one. Spark is not told when a
streaming query stops, so it
+ * closes the last stream writer of the query once the writer is no longer
referenced, or when the
+ * query is restarted.</li>
+ * </ul>
+ *
+ * <h2>Conventions</h2>
+ * <ul>
+ * <li><b>Handles.</b> The {@code create} methods return an opaque, non-zero
handle, typically a
+ * pointer. Spark releases each handle exactly once with the matching {@code
close} method, or
+ * with {@link #commitDataWriter} or {@link #abortDataWriter} for data
writers, though not
+ * always right after the last use of a stream writer (see Writes above). A
reader or
+ * writer handle must not depend on the data source handle it was created
from, because Spark
+ * closes the data source handle right after creating it.</li>
+ * <li><b>State.</b> 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.</li>
+ * <li><b>Errors.</b> To report an error, throw a Java exception, for
example with the JNI
+ * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and
return any value. Spark
+ * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.</li>
+ * <li><b>Arrow data.</b> Schemas and data are exchanged with the
+ * <a href="https://arrow.apache.org/docs/format/CDataInterface.html";>Arrow
C data
+ * interface</a>. A {@code long} argument named {@code ...Address} is the
address of a C struct
+ * allocated by Spark. For an output struct, the library initializes it and
Spark takes
+ * ownership. For an input struct, the library may move it to take
ownership; otherwise, Spark
+ * releases it after the call returns, so the library must not use it
afterwards. The arrays
+ * that the library exports must have an offset of 0, at every level: Arrow
Java ignores the
+ * offsets of the arrays it imports, so Spark rejects an array with another
offset, such as a
+ * zero-copy slice. Copy a sliced array before exporting it. Spark also
rejects an array with
+ * fewer values than its parent refers to, such as a column shorter than its
batch. The buffers
+ * of an array must hold all of its values, as the C data interface
requires: it does not carry
+ * the sizes of the buffers, so Spark cannot check them, and reads past the
end of a buffer that
+ * is too short.</li>
+ * <li><b>Schemas.</b> Spark types map to Arrow types as in {@code
spark.sql.execution.arrow}:
+ * for example, {@code bigint} is {@code int64}, {@code string} is {@code
utf8},
+ * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code
timestamp_ntz} is
+ * {@code timestamp[us]} without one. Data read from a library may also use
+ * {@code large_utf8}, {@code utf8_view} and the other variants of the same
type, and any names
+ * for the fields of the entries of a map, but not unsigned integers, {@code
decimal256}, list
+ * views or dictionary encoding.</li>
+ * <li><b>Options.</b> Option keys are passed in lower case, because Spark
matches them
+ * case-insensitively.</li>
+ * <li><b>Optional methods.</b> A library only needs to export the functions
for the operations
+ * it supports. If a function is missing, Spark treats the operation as
unsupported, except for
+ * the pushdown methods and the {@code serialize} methods, which then push
nothing down and
+ * return no state.</li>
+ * <li><b>Threads.</b> Spark calls the library from multiple threads, and
calls on different
+ * handles may run concurrently. Calls on the same handle never run
concurrently. A handle may
+ * be released from a thread other than the threads that used it, such as a
JVM
+ * {@code Cleaner} thread when Spark releases a handle once it is no longer
referenced, for
+ * example the last stream writer of a streaming query.</li>
+ * </ul>
+ *
+ * <h2>Compatibility</h2>
+ * The JVM binds the functions of a library by their names only, so the name
of this class, the
+ * names of its methods and their signatures are the interface. Within an
{@link #ABI_VERSION}:
+ * <ul>
+ * <li>the signature and the semantics of an existing method never
change;</li>
+ * <li>new methods may be added, only as optional methods, so that a library
that does not
+ * export them keeps working as before;</li>
+ * <li>new values may appear where this documentation says so, such as the
names of pushed
+ * predicates.</li>
+ * </ul>
+ * An incompatible change uses a new {@link #ABI_VERSION}, with new names or a
new class. A
+ * library reports the version it implements with {@link #abiVersion}, and
Spark rejects a library
+ * of a version it does not support. Spark supports version 1 at least in all
the releases of
+ * Spark 4 from 4.4.0.
+ *
+ * @since 4.4.0
+ */
+@Evolving
+public final class NativeBridge {
+
+ /**
+ * The version of this interface. It changes when the interface changes
incompatibly: see
+ * Compatibility above.
+ */
+ public static final int ABI_VERSION = 1;
+
+ private NativeBridge() {}
+
+ /**
+ * Loads the library that implements the native methods of this copy of the
class. Spark calls
+ * it on the copy of the class it creates for the library.
+ */
+ static void load(String path) {
+ System.load(path);
+ }
+
+ /** Returns the version of this interface that the library implements,
{@link #ABI_VERSION}. */
+ public static native int abiVersion();
+
+ //
---------------------------------------------------------------------------------------------
+ // Data source (driver)
+ //
---------------------------------------------------------------------------------------------
+
+ /**
+ * Creates a data source with the options given by the user, such as
+ * {@code spark.read.option("key", "value")}.
+ *
+ * @param name the name of the data source, as listed in the package
manifest. A library can
+ * implement several data sources.
+ * @param optionKeys the option keys, in lower case
+ * @param optionValues the option values, in the same order as the keys
+ * @return the data source handle
+ */
+ public static native long createDataSource(
+ String name, String[] optionKeys, String[] optionValues);
+
+ /**
+ * Exports the schema of the data. Spark calls it only when the user does
not specify a schema.
+ *
+ * @param schemaAddress the address of the {@code ArrowSchema} to initialize
with a struct type
+ * whose fields are the columns
+ */
+ public static native void schema(long dataSource, long schemaAddress);
+
+ /**
+ * Creates a reader for a batch scan.
+ *
+ * @param schemaAddress the address of the {@code ArrowSchema} of the schema
to read: the one the
+ * user specified, or the one returned by {@link
#schema}
+ * @return the reader handle
+ */
+ public static native long createReader(long dataSource, long schemaAddress);
+
+ /**
+ * Creates a reader for a micro-batch streaming scan.
+ *
+ * @param schemaAddress the address of the {@code ArrowSchema} of the schema
to read: the one the
+ * user specified, or the one returned by {@link
#schema}
+ * @return the stream reader handle
+ */
+ public static native long createStreamReader(long dataSource, long
schemaAddress);
+
+ /**
+ * Creates a writer for a batch write.
+ *
+ * @param schemaAddress the address of the {@code ArrowSchema} of the data
to write
+ * @param overwrite whether to replace the existing data (save mode
"overwrite") instead of
+ * appending to it
+ * @return the writer handle
+ */
+ public static native long createWriter(long dataSource, long schemaAddress,
boolean overwrite);
+
+ /**
+ * Creates a writer for a streaming write.
+ *
+ * @param schemaAddress the address of the {@code ArrowSchema} of the data
to write
+ * @param overwrite whether to replace the existing data in each micro-batch
(output mode
+ * "complete") instead of appending to it
+ * @return the writer handle
+ */
+ public static native long createStreamWriter(
+ long dataSource, long schemaAddress, boolean overwrite);
+
+ /** Releases a data source handle. */
+ public static native void closeDataSource(long dataSource);
+
+ //
---------------------------------------------------------------------------------------------
+ // Batch reader (driver)
+ //
---------------------------------------------------------------------------------------------
+
+ /**
+ * Pushes predicates down to the reader. The predicates are combined with
AND.
+ * <p>
+ * Each predicate is a JSON expression tree. An expression is one of:
+ * <ul>
+ * <li>{@code {"type": "column", "name": ["a", "b"]}}: a column. A nested
field has a name
+ * part for each level.</li>
+ * <li>{@code {"type": "literal", "dataType": "bigint", "value": 1}}: a
literal. The data type
+ * is a Spark SQL type name such as {@code int}, {@code string} or {@code
decimal(10,2)}. The
+ * value is {@code null}, a JSON boolean or number for boolean and numeric
types, a string
+ * for {@code string}, a base64 string for {@code binary}, a decimal
string for
+ * {@code decimal}, the days since the epoch for {@code date}, and the
microseconds since the
+ * epoch for {@code timestamp} and {@code timestamp_ntz}.</li>
+ * <li>{@code {"type": "function", "name": ">", "children": [...]}}: a
predicate, with the
+ * name it has in {@code
org.apache.spark.sql.connector.expressions.filter.Predicate}. The
+ * names are {@code =}, {@code <>}, {@code <=>} (equal, with null equal to
null), {@code <},
+ * {@code <=}, {@code >} and {@code >=} with two children, {@code IN} with
the value to look
+ * for and then the values to look in, {@code IS_NULL} and {@code
IS_NOT_NULL} with one
+ * child, {@code STARTS_WITH}, {@code ENDS_WITH} and {@code CONTAINS} with
the string and the
+ * string to look for, {@code AND} and {@code OR} with two predicates, and
{@code NOT} with
+ * one. The children of the other functions are columns, literals or
predicates.</li>
+ * </ul>
+ * Spark does not pass predicates that it cannot express this way, such as
predicates that use
+ * arithmetic or other functions, nor predicates on columns that contain
strings with a
+ * non-binary collation, such as {@code UTF8_LCASE}: strings are compared by
their UTF-8 bytes.
+ * A later version of Spark may pass predicates with other names without
changing
+ * {@link #ABI_VERSION}, so a reader must not report a predicate as
evaluated completely if it
+ * contains a name that the reader does not know, or a data type that it
does not support.
+ * <p>
+ * A predicate that the reader evaluates completely must follow the
semantics of Spark SQL. In
+ * particular, for floating-point columns, NaN is equal to NaN and greater
than any other value,
+ * including positive infinity, and {@code -0.0} is equal to {@code 0.0},
unlike IEEE 754
+ * comparisons. For example, {@code d > 1.0} is true for a NaN value of
{@code d}. Comparisons
+ * with null are null, so a row only matches if the predicate is true.
+ *
+ * @return for each predicate, whether the reader evaluates it completely,
so that Spark does
+ * not have to evaluate it again. The reader can still use the other
predicates to skip
+ * data.
+ */
+ public static native boolean[] pushPredicates(long reader, String[]
predicates);
+
+ /**
+ * Pushes a LIMIT down to the reader.
+ *
+ * @return whether the reader returns at most {@code limit} rows in total.
Spark applies the
+ * limit again after reading either way.
+ */
+ public static native boolean pushLimit(long reader, int limit);
+
+ /**
+ * Prunes the columns to read.
+ *
+ * @param columnNames the names of the top-level columns that the query needs
+ * @return whether {@link #read} returns exactly these columns, in this
order. If not, it
+ * returns all the columns of the schema to read.
+ */
+ public static native boolean pruneColumns(long reader, String[] columnNames);
+
+ /**
+ * Plans the partitions of the scan. Spark reads each partition in a
separate task.
+ *
+ * @return the partitions, each serialized in a format defined by the library
+ */
+ public static native byte[][] partitions(long reader);
+
+ /**
+ * Serializes the state of the reader that {@link #read} needs on the
executors, such as the
+ * pushed predicates. Spark calls it after {@link #partitions}.
+ */
+ public static native byte[] serializeReader(long reader);
+
+ /** Releases a batch reader handle. */
+ public static native void closeReader(long reader);
+
+ //
---------------------------------------------------------------------------------------------
+ // Stream reader (driver)
+ //
---------------------------------------------------------------------------------------------
+
+ /** Returns the offset to start from when the streaming query has no
checkpoint, as JSON. */
+ public static native String initialOffset(long streamReader);
+
+ /** Returns the most recent offset available, as JSON. */
+ public static native String latestOffset(long streamReader);
+
+ /**
+ * Plans the partitions that read the data after {@code start} up to and
including {@code end}.
+ *
+ * @return the partitions, each serialized in a format defined by the library
+ */
+ public static native byte[][] streamPartitions(long streamReader, String
start, String end);
+
+ /**
+ * Serializes the state of the stream reader that {@link #read} needs on the
executors. Spark
+ * calls it after each call to {@link #streamPartitions}.
+ */
+ public static native byte[] serializeStreamReader(long streamReader);
+
+ /**
+ * Tells the library that Spark has processed all the data up to and
including {@code end},
+ * and will not ask for it again.
+ */
+ public static native void commitOffset(long streamReader, String end);
+
+ /** Releases a stream reader handle, when the streaming query stops. */
+ public static native void closeStreamReader(long streamReader);
+
+ //
---------------------------------------------------------------------------------------------
+ // Partition reader (executor)
+ //
---------------------------------------------------------------------------------------------
+
+ /**
+ * Reads a partition of a batch or streaming scan.
+ *
+ * @param readerState the state returned by {@link #serializeReader} or
+ * {@link #serializeStreamReader}, or an empty array if
there is none
+ * @param partition a partition returned by {@link #partitions} or {@link
#streamPartitions}
+ * @param streamAddress the address of the {@code ArrowArrayStream} to
initialize. Each array
+ * of the stream is a struct array whose fields are the
columns to read:
+ * the columns of the schema to read, or the pruned
columns. Spark calls
+ * the stream from the thread of the task that reads
the partition.
+ * {@code get_next} may block, for example on network
I/O, but Spark
+ * cannot interrupt native code: when the task is
killed, Spark only
+ * stops reading once {@code get_next} returns, and
then releases the
+ * stream. So a {@code get_next} that blocks forever
keeps the task
+ * running, unless the task reaper ({@code
spark.task.reaper.enabled}
+ * and {@code spark.task.reaper.killTimeout}) kills the
executor. Use
+ * timeouts for blocking operations.
+ */
+ public static native void read(byte[] readerState, byte[] partition, long
streamAddress);
Review Comment:
Thanks. I agree that this needs a short design, and I'd like to do it
separately. Reserving from native code needs a callback into the JVM from any
thread of the library, or a function table in a C ABI such as #59424. The
budget also deserves its own discussion: the off-heap pool, or a dedicated
setting like `spark.executor.pyspark.memory`. I'd be glad to have you review it.
For this PR, aa0f3ba4dd77 adds a Memory section to the docs page. It says
what Spark does not account for: the memory that the library allocates, the
Arrow buffers that it exports, and the batches that Spark writes. It also says
that this memory comes out of `spark.driver.memoryOverhead` and
`spark.executor.memoryOverhead`, and that YARN and Kubernetes kill an executor
that exceeds it. Limitations lists memory accounting too. On putting the
context in the version 1 signatures: libraries written against version 1 would
run unaccounted. But an optional function, through which a library gets the
memory context of its task, would reach the libraries that adopt it, like any
other optional function, and a later ABI version could make it required.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala:
##########
@@ -0,0 +1,421 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.datasources.v2.ffi
+
+import java.io.{File, FileInputStream, IOException}
+import java.nio.charset.StandardCharsets
+import java.nio.file.{Files, StandardCopyOption}
+import java.security.MessageDigest
+import java.util.{HexFormat, Locale}
+import java.util.concurrent.ConcurrentHashMap
+import java.util.zip.ZipFile
+
+import scala.jdk.CollectionConverters._
+import scala.util.Try
+
+import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper}
+import com.google.common.cache.{Cache, CacheBuilder}
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.PATH
+import org.apache.spark.sql.Artifact
+import org.apache.spark.sql.datasource.NativeBridge
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.util.Utils
+
+/**
+ * The manifest of a native data source package,
`spark-native-datasource.json`:
+ * {{{
+ * {
+ * "abiVersion": 1,
+ * "dataSources": ["my_source"],
+ * "libraries": {
+ * "linux-x86_64": "linux-x86_64/libmy_source.so",
+ * "osx-aarch64": "osx-aarch64/libmy_source.dylib"
+ * }
+ * }
+ * }}}
+ *
+ * @param abiVersion the version of [[NativeBridge]] that the libraries
implement
+ * @param dataSources the names of the data sources that the libraries
implement
+ * @param libraries the path of the library in the package for each platform,
as returned by
+ * [[NativePlatform.current]]
+ */
+case class NativeDataSourceManifest(
+ abiVersion: Int,
+ dataSources: Seq[String],
+ libraries: Map[String, String])
+
+/**
+ * Where the library of a native data source is: in a
[[NativeDataSourcePackage]], or installed in
+ * the native library path, as an [[InstalledNativeLibrary]]. It is sent to
the executors, which
+ * load the library with [[NativeLibraries.get]].
+ */
+sealed trait NativeLibraryLocation extends Serializable
+
+/**
+ * A native data source package: a zip file that contains a
[[NativeDataSourceManifest]] and the
+ * libraries it lists. The SHA-256 checksum of the file identifies the
package, so that an
+ * executor can check that it loads the same package as the driver.
+ *
+ * @param path the path of the package file on the driver
+ * @param artifactUUID the UUID of the session whose artifacts contain the
package, if the session
+ * is isolated. Executors store the artifacts of such a
session separately.
+ * @param distributedFileName the name of the copy of the package in the
artifacts of the session,
+ * if the package is not itself an artifact of the
session
+ */
+case class NativeDataSourcePackage(
+ path: String,
+ sha256: String,
+ manifest: NativeDataSourceManifest,
+ artifactUUID: Option[String] = None,
+ distributedFileName: Option[String] = None) extends NativeLibraryLocation {
+ def fileName: String = new File(path).getName
+}
+
+/**
+ * A package file that is not a valid native data source package.
+ *
+ * @param error the error that explains why the package is invalid
+ * @param dataSources the data sources listed in the manifest of the package,
if it can be read
+ */
+case class InvalidNativeDataSourcePackage(error: Throwable, dataSources:
Seq[String])
+
+object NativeDataSourcePackage extends Logging {
+ val FILE_EXTENSION: String = Artifact.NATIVE_DATA_SOURCE_PACKAGE_EXTENSION
+ val MANIFEST_NAME = "spark-native-datasource.json"
+
+ private case class CacheKey(path: String, length: Long, lastModified: Long)
+
+ private[ffi] type ReadResult = Either[InvalidNativeDataSourcePackage,
NativeDataSourcePackage]
+
+ // Bounded, because each session stores its artifacts in a directory of its
own, so a
+ // long-running server where many sessions add packages reads many different
files.
+ private val MAX_CACHED_PACKAGES = 1000
+
+ private val packages: Cache[CacheKey, ReadResult] =
+ CacheBuilder.newBuilder().maximumSize(MAX_CACHED_PACKAGES).build[CacheKey,
ReadResult]()
+ private val mapper = new ObjectMapper()
+
+ /** Reads the package in the given file. The result is cached until the file
changes. */
+ def read(file: File): NativeDataSourcePackage = tryRead(file).fold(p =>
throw p.error, identity)
+
+ /**
+ * Reads the package in the given file, or returns why it is invalid. The
result is cached until
+ * the file changes, and an invalid package is logged when it is read.
+ */
+ private[ffi] def tryRead(file: File): ReadResult = {
+ val canonical = try {
+ file.getCanonicalFile
+ } catch {
+ case e: IOException =>
+ return Left(InvalidNativeDataSourcePackage(
+ invalidManifestError(file, s"The package cannot be read: $e", e),
Nil))
+ }
+ val key = CacheKey(canonical.getPath, canonical.length(),
canonical.lastModified())
+ Option(packages.getIfPresent(key)).getOrElse {
+ val result = parse(canonical)
+ result.swap.foreach { invalid =>
+ logWarning(log"Invalid native data source package ${MDC(PATH,
canonical)}.",
+ invalid.error)
+ }
+ packages.put(key, result)
+ result
+ }
+ }
+
+ /**
+ * Forgets the packages read from the files under the given directory, such
as the artifacts of
+ * a session that is cleaned up.
+ */
+ def forgetPackagesIn(dir: File): Unit = {
+ val prefix = dir.getCanonicalPath + File.separator
+ packages.asMap().keySet().removeIf(_.path.startsWith(prefix))
+ }
+
+ /** Returns whether a package read from the file at the given canonical path
is cached. */
+ private[ffi] def isCached(canonicalPath: String): Boolean = {
+ packages.asMap().keySet().asScala.exists(_.path == canonicalPath)
+ }
+
+ private def invalidManifestError(file: File, reason: String, cause:
Throwable = null) = {
+ QueryExecutionErrors.invalidNativeDataSourcePackageError(
+ file.getPath, "INVALID_MANIFEST", Map("reason" -> reason), cause)
+ }
+
+ private def parse(file: File): ReadResult = {
+ def invalid(reason: String, dataSources: Seq[String] = Nil, cause:
Throwable = null) = {
+ Left(InvalidNativeDataSourcePackage(invalidManifestError(file, reason,
cause), dataSources))
+ }
+ try {
+ val manifest = Utils.tryWithResource(new ZipFile(file)) { zip =>
+ Option(zip.getEntry(MANIFEST_NAME)) match {
+ case None => invalid(s"The package does not contain the manifest
$MANIFEST_NAME.")
+ case Some(entry) =>
+ val json = Utils.tryWithResource(zip.getInputStream(entry)) { in =>
+ new String(in.readAllBytes(), StandardCharsets.UTF_8)
+ }
+ parseManifest(json) match {
+ case Left(reason) => invalid(reason)
+ case Right(manifest) =>
+ manifest.libraries.values.find(zip.getEntry(_) == null) match {
+ case Some(library) => invalid(
+ s"The package does not contain the library $library.",
manifest.dataSources)
+ case None => Right(manifest)
+ }
+ }
+ }
+ }
+ manifest.map(NativeDataSourcePackage(file.getPath, sha256(file), _))
+ } catch {
+ case e: IOException => invalid(s"The package is not a valid zip file:
$e", cause = e)
+ }
+ }
+
+ private[ffi] def parseManifest(json: String): Either[String,
NativeDataSourceManifest] = {
+ val root: JsonNode = try {
+ mapper.readTree(json)
+ } catch {
+ case e: IOException => return Left(s"The manifest is not valid JSON:
${e.getMessage}")
+ }
+ if (root == null || !root.isObject) {
+ return Left("The manifest must be a JSON object.")
+ }
+ val abiVersion = root.get("abiVersion")
+ if (abiVersion == null || !abiVersion.isInt) {
+ return Left("The manifest must have an integer field 'abiVersion'.")
+ }
+ val dataSources = root.get("dataSources")
+ if (dataSources == null || !dataSources.isArray || dataSources.isEmpty ||
+ !dataSources.elements().asScala.forall(n => n.isTextual &&
n.asText().nonEmpty)) {
+ return Left("The manifest must have a field 'dataSources' with a
non-empty array of " +
+ "data source names.")
+ }
+ val libraries = root.get("libraries")
+ if (libraries == null || !libraries.isObject || libraries.isEmpty ||
+ !libraries.properties().asScala.forall(e => e.getValue.isTextual)) {
+ return Left("The manifest must have a field 'libraries' with an object
that maps " +
+ "platforms to the paths of the libraries in the package.")
+ }
+ Right(NativeDataSourceManifest(
+ abiVersion.asInt(),
+ dataSources.elements().asScala.map(_.asText()).toSeq,
+ libraries.properties().asScala.map(e => e.getKey ->
e.getValue.asText()).toMap))
+ }
+
+ private[ffi] def sha256(file: File): String = {
+ val digest = MessageDigest.getInstance("SHA-256")
+ Utils.tryWithResource(new FileInputStream(file)) { in =>
+ val buffer = new Array[Byte](64 * 1024)
+ var read = in.read(buffer)
+ while (read != -1) {
+ digest.update(buffer, 0, read)
+ read = in.read(buffer)
+ }
+ }
+ HexFormat.of().formatHex(digest.digest())
+ }
+}
+
+/**
+ * A native data source library installed in the native library path, where
+ * [[NativeDataSourceRegistry]] finds it by its file name. It must be
installed on every node. The
+ * SHA-256 checksum of the file identifies the build of the library, so that
an executor can check
+ * that it loads the same build as the driver: the state that the driver
serializes is only
+ * meant for the same build.
+ *
+ * @param path the path of the library on the driver
+ * @param sha256 the SHA-256 checksum of the library on the driver
+ */
+case class InstalledNativeLibrary(path: String, sha256: String) extends
NativeLibraryLocation
+
+object InstalledNativeLibrary {
+ private case class CacheKey(path: String, length: Long, lastModified: Long)
+
+ // Bounded, like the cache of the packages, although few libraries are
installed.
+ private val checksums: Cache[CacheKey, String] =
+ CacheBuilder.newBuilder().maximumSize(1000).build[CacheKey, String]()
+
+ /** The installed library in the given file. */
+ def apply(file: File): InstalledNativeLibrary =
InstalledNativeLibrary(file.getPath, sha256(file))
+
+ /** The SHA-256 checksum of the given file. The result is cached until the
file changes. */
+ private[ffi] def sha256(file: File): String = {
+ val canonical = file.getCanonicalFile
+ val key = CacheKey(canonical.getPath, canonical.length(),
canonical.lastModified())
+ Option(checksums.getIfPresent(key)).getOrElse {
+ val checksum = NativeDataSourcePackage.sha256(canonical)
+ checksums.put(key, checksum)
+ checksum
+ }
+ }
+}
+
+/** The platform names used in [[NativeDataSourceManifest.libraries]]. */
+object NativePlatform {
+ /** The platform of this JVM, such as `linux-x86_64`, `linux-aarch64` or
`osx-aarch64`. */
+ lazy val current: String = name(System.getProperty("os.name"),
System.getProperty("os.arch"))
+
+ private[ffi] def name(osName: String, osArch: String): String = {
+ val os = osName.toLowerCase(Locale.ROOT) match {
+ case n if n.startsWith("linux") => "linux"
+ case n if n.startsWith("mac") || n.startsWith("darwin") => "osx"
+ case n if n.startsWith("windows") => "windows"
+ case n => n.replaceAll("[^a-z0-9]", "")
+ }
+ val arch = osArch.toLowerCase(Locale.ROOT) match {
+ case "amd64" | "x86_64" | "x64" => "x86_64"
+ case "aarch64" | "arm64" => "aarch64"
+ case a => a
+ }
+ s"$os-$arch"
+ }
+}
+
+/**
+ * The native data source libraries loaded in this JVM. A library is loaded
once, on first use,
+ * after it is extracted from its package if needed, and stays loaded until
the JVM exits.
+ */
+object NativeLibraries extends Logging {
+ // The libraries of the packages, by the checksum of the package, and the
installed libraries,
+ // by their canonical path: the JVM loads a library file in a single class
loader. A library
+ // that is loaded but rejected stays loaded, so its error is kept: it cannot
be loaded again.
+ private val packagedLibraries = new ConcurrentHashMap[String,
Try[NativeLibrary]]()
+ private val installedLibraries = new ConcurrentHashMap[String,
Try[NativeLibrary]]()
+
+ private lazy val extractionDir: File = Utils.createTempDir(namePrefix =
"native-datasources")
+
+ /** Returns the library at the given location, and loads it if needed. */
+ private[ffi] def get(location: NativeLibraryLocation): NativeLibrary =
location match {
+ case pkg: NativeDataSourcePackage =>
+ getOrLoad(packagedLibraries, pkg.sha256)(loadPackaged(pkg))
+ case installed: InstalledNativeLibrary =>
+ val file = locate(installed).getCanonicalFile
+ getOrLoad(installedLibraries, file.getPath)(loadInstalled(file))
Review Comment:
Thanks for checking it on the PR head. Fixed in aa0f3ba4dd77 as you
suggested. Spark records the checksum of the build that it loads from each
path. When another build is requested from that path, it fails with
`INVALID_NATIVE_DATA_SOURCE_LIBRARY.REPLACED`, which names the path and the two
checksums, and asks to restart the application. That happens on the driver,
before any state is serialized, and on the executors that loaded the old build.
I agree about not loading a copy from a directory named after the checksum,
because of `$ORIGIN` and `@loader_path`. The docs page says that replacing an
installed library needs a restart. New test: "reject an installed library
replaced in place", which replaces the file with an `ID_SIGN=-1` build, as you
did, then installs the loaded build again.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -0,0 +1,794 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.datasources.v2.ffi
+
+import java.lang.ref.{Cleaner, WeakReference}
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicLong
+
+import scala.jdk.CollectionConverters._
+import scala.util.control.NonFatal
+
+import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema,
CDataDictionaryProvider, Data}
+import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.MapVector
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
+
+import org.apache.spark.{SparkException, SparkThrowable,
SparkUnsupportedOperationException}
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connector.expressions.filter.Predicate
+import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader}
+import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage}
+import org.apache.spark.sql.errors.{QueryCompilationErrors,
QueryExecutionErrors}
+import
org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarDataSource,
ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{DataType, StructType}
+import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch,
ColumnVector}
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.util.Utils
+
+/**
+ * A [[ColumnarDataSource]] implemented by a native library, through
+ * `org.apache.spark.sql.datasource.NativeBridge`.
+ *
+ * @param location where the library is on the driver. The readers and the
writers distribute it
+ * to the executors of the active session when they are
created.
+ * @param name the name of the data source, as listed in the manifest of its
package or in lower
+ * case for an installed library
+ */
+class NativeDataSource(
+ location: NativeLibraryLocation,
+ name: String,
+ options: CaseInsensitiveStringMap)
+ extends ColumnarDataSource {
+ import NativeCalls._
+
+ override def schema(): StructType = {
+ val library = NativeLibraries.get(location)
+ withDataSource(library, "infer the schema of") { dataSource =>
+ call(name, "infer the schema of") {
+ optional(NativeArrow.importSchema(address =>
library.schema(dataSource, address)))
+ }.getOrElse(throw
QueryCompilationErrors.dataSchemaNotSpecifiedError(name))
+ }
+ }
+
+ override def reader(schema: StructType): ColumnarReader = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of",
"DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") {
+ dataSource =>
+ NativeArrow.exportSchema(schema)(address =>
library.createReader(dataSource, address))
+ }
+ new NativeDataSourceReader(
+ executorLocation, name, schema, new NativeHandle(handle,
library.closeReader))
+ }
+
+ override def streamReader(schema: StructType): ColumnarStreamReader = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = create(library, "plan a scan of",
"DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") {
+ dataSource =>
+ NativeArrow.exportSchema(schema)(address =>
library.createStreamReader(dataSource, address))
+ }
+ new NativeDataSourceStreamReader(
+ executorLocation, name, schema, new NativeHandle(handle,
library.closeStreamReader))
+ }
+
+ override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter
= {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = new NativeHandle(
+ create(library, "plan a write to",
"DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createWriter(dataSource, address, overwrite)
+ }
+ },
+ library.closeWriter)
+ new NativeDataSourceWriter(executorLocation, name, handle,
serializeWriter(library, handle))
+ }
+
+ override def streamWriter(
+ schema: StructType,
+ overwrite: Boolean,
+ queryId: String): ColumnarStreamWriter = {
+ val executorLocation = NativeDataSourceRegistry.distribute(location)
+ val library = NativeLibraries.get(location)
+ val handle = new NativeHandle(
+ create(library, "plan a write to",
"DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") {
+ dataSource =>
+ NativeArrow.exportSchema(schema) { address =>
+ library.createStreamWriter(dataSource, address, overwrite)
+ }
+ },
+ library.closeWriter)
+ val writer = new NativeDataSourceStreamWriter(
+ executorLocation, name, handle, serializeWriter(library, handle))
+ NativeStreamWriters.register(queryId, handle)
+ writer
+ }
+
+ /** Runs `f` with a new native data source, which is closed afterwards. */
+ private def withDataSource[T](library: NativeLibrary, action: String)(f:
Long => T): T = {
+ // The keys of a case-insensitive map are in lower case.
+ val entries = options.entrySet().asScala.toArray
+ val dataSource = call(name, "create") {
+ library.createDataSource(name, entries.map(_.getKey),
entries.map(_.getValue))
+ }
+ Utils.tryWithSafeFinally(f(dataSource)) {
+ call(name, action)(library.closeDataSource(dataSource))
+ }
+ }
+
+ /** Creates a reader or a writer with `f`, which fails if the library does
not implement it. */
+ private def create(library: NativeLibrary, action: String, unsupportedError:
String)(
+ f: Long => Long): Long = {
+ withDataSource(library, action) { dataSource =>
+ call(name, action)(optional(f(dataSource))).getOrElse {
+ throw new SparkUnsupportedOperationException(unsupportedError,
Map("description" -> name))
+ }
+ }
+ }
+
+ private def serializeWriter(library: NativeLibrary, handle: NativeHandle):
Array[Byte] = {
+ try {
+ bytesOrEmpty(call(name, "plan a write
to")(optional(library.serializeWriter(handle.get))))
+ } catch {
+ case e: Throwable =>
+ handle.close()
+ throw e
+ }
+ }
+}
+
+/** Reads a native data source for a batch scan. */
+class NativeDataSourceReader(
+ location: NativeLibraryLocation,
+ name: String,
+ schema: StructType,
+ @transient private val handle: NativeHandle)
+ extends ColumnarReader {
+ import NativeCalls._
+
+ // Set on the driver when the partitions are planned, and serialized to the
executors.
+ private var readerState: Array[Byte] = Array.emptyByteArray
+ private var readSchema: StructType = schema
+ @transient private var plannedPartitions: Array[InputPartition] = _
+
+ @transient private lazy val library = NativeLibraries.get(location)
+
+ override def pushPredicates(predicates: Array[Predicate]): Array[Predicate]
= {
+ // Only the predicates that can be expressed in JSON are passed to the
library.
+ val encoded = predicates.zipWithIndex.flatMap { case (predicate, i) =>
+ NativePredicates.toJson(predicate, schema).map(_ -> i)
+ }
+ val pushed = if (encoded.isEmpty) {
+ Set.empty[Int]
+ } else {
+ call(name, "plan a scan of") {
+ optional(library.pushPredicates(handle.get, encoded.map(_._1)))
+ } match {
+ case Some(flags) if flags != null && flags.length == encoded.length =>
+ encoded.map(_._2).zip(flags).collect { case (i, true) => i }.toSet
+ case Some(flags) =>
+ throw error(name, "plan a scan of", s"pushPredicates returned " +
+ s"${Option(flags).map(_.length).getOrElse(0)} values for
${encoded.length} predicates.")
+ case None => Set.empty[Int]
+ }
+ }
+ predicates.zipWithIndex.collect { case (predicate, i) if
!pushed.contains(i) => predicate }
+ }
+
+ override def pushLimit(limit: Int): Boolean = {
+ call(name, "plan a scan of")(optional(library.pushLimit(handle.get,
limit))).getOrElse(false)
+ }
+
+ override def pruneColumns(requiredSchema: StructType): Boolean = {
+ val pruned = call(name, "plan a scan of") {
+ optional(library.pruneColumns(handle.get, requiredSchema.fieldNames))
+ }.getOrElse(false)
+ if (pruned) {
+ readSchema = requiredSchema
+ }
+ pruned
+ }
+
+ override def partitions(): Array[InputPartition] = synchronized {
+ if (plannedPartitions == null) {
+ try {
+ val partitions = call(name, "plan a scan
of")(library.partitions(handle.get))
+ readerState = bytesOrEmpty(
+ call(name, "plan a scan
of")(optional(library.serializeReader(handle.get))))
+ plannedPartitions = NativeInputPartition.from(partitions)
+ } finally {
+ // The executors only need the serialized state.
+ handle.close()
+ }
+ }
+ plannedPartitions
+ }
+
+ override def read(partition: InputPartition): PartitionReader[ColumnarBatch]
= {
+ new NativePartitionReader(
+ 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 = {
+ call(name, "plan a scan of")(library.initialOffset(handle.get))
+ }
+
+ override def latestOffset(): String = {
+ call(name, "plan a scan of")(library.latestOffset(handle.get))
+ }
+
+ override def partitions(start: String, end: String): Array[InputPartition] =
{
+ val partitions = call(name, "plan a scan of") {
+ library.streamPartitions(handle.get, start, end)
+ }
+ readerState = bytesOrEmpty(
+ call(name, "plan a scan
of")(optional(library.serializeStreamReader(handle.get))))
+ NativeInputPartition.from(partitions)
+ }
+
+ override def read(partition: InputPartition): PartitionReader[ColumnarBatch]
= {
+ new NativePartitionReader(
+ 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()
+}
+
+/** Writes to a native data source for a batch write. */
+class NativeDataSourceWriter(
+ location: NativeLibraryLocation,
+ name: String,
+ @transient private val handle: NativeHandle,
+ writerState: Array[Byte])
+ extends ColumnarWriter {
+ import NativeCalls._
+
+ @transient private lazy val library = NativeLibraries.get(location)
+
+ override def createWriter(partitionId: Int, taskId: Long):
DataWriter[ColumnarBatch] = {
+ new NativeDataWriter(location, name, writerState, partitionId, taskId,
epochId = -1L)
+ }
+
+ override def commit(messages: Array[WriterCommitMessage]): Unit = {
+ call(name, "commit a write to") {
+ library.commit(handle.get, -1L,
NativeWriterCommitMessage.toBytes(messages))
+ }
+ // Not closed if the commit fails: Spark then aborts the write, which
closes it.
+ handle.close()
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {
+ Utils.tryWithSafeFinally {
+ call(name, "abort a write to") {
+ library.abort(handle.get, -1L,
NativeWriterCommitMessage.toBytes(messages))
+ }
+ }(handle.close())
+ }
+}
+
+/**
+ * Writes to a native data source for a streaming write. Its handle is closed
when the streaming
+ * query creates its next stream writer, see [[NativeStreamWriters]].
+ */
+class NativeDataSourceStreamWriter(
+ location: NativeLibraryLocation,
+ name: String,
+ @transient private val handle: NativeHandle,
+ writerState: Array[Byte])
+ extends ColumnarStreamWriter {
+ import NativeCalls._
+
+ @transient private lazy val library = NativeLibraries.get(location)
+
+ override def createWriter(
+ partitionId: Int,
+ taskId: Long,
+ epochId: Long): DataWriter[ColumnarBatch] = {
+ new NativeDataWriter(location, name, writerState, partitionId, taskId,
epochId)
+ }
+
+ override def commit(epochId: Long, messages: Array[WriterCommitMessage]):
Unit = {
+ call(name, "commit a write to") {
+ library.commit(handle.get, epochId,
NativeWriterCommitMessage.toBytes(messages))
+ }
+ }
+
+ override def abort(epochId: Long, messages: Array[WriterCommitMessage]):
Unit = {
+ call(name, "abort a write to") {
+ library.abort(handle.get, epochId,
NativeWriterCommitMessage.toBytes(messages))
+ }
+ }
+}
+
+/** A partition planned by a native library. */
+case class NativeInputPartition(index: Int, bytes: Array[Byte]) extends
InputPartition
+
+object NativeInputPartition {
+ def from(partitions: Array[Array[Byte]]): Array[InputPartition] = {
+ Option(partitions).getOrElse(Array.empty).zipWithIndex.map { case (bytes,
i) =>
+ NativeInputPartition(i, Option(bytes).getOrElse(Array.emptyByteArray))
+ }
+ }
+}
+
+/** The message returned by the data writer of a task of a native library. */
+case class NativeWriterCommitMessage(bytes: Array[Byte]) extends
WriterCommitMessage
+
+object NativeWriterCommitMessage {
+ def toBytes(messages: Array[WriterCommitMessage]): Array[Array[Byte]] =
messages.map {
+ case message: NativeWriterCommitMessage => message.bytes
+ case null => null
+ case other =>
+ throw SparkException.internalError(s"Unexpected commit message:
${other.getClass.getName}")
+ }
+}
+
+/**
+ * Reads a partition of a native data source: each array of the Arrow stream
exported by the
+ * library is checked, and imported into the Arrow vectors that back the
returned batch.
+ */
+class NativePartitionReader(
+ location: NativeLibraryLocation,
+ name: String,
+ readerState: Array[Byte],
+ partition: Array[Byte],
+ schema: StructType)
+ extends PartitionReader[ColumnarBatch] {
+ import NativeCalls._
+
+ private val allocator =
+ ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0,
Long.MaxValue)
+ // Read here rather than with `Data.importArrayStream`, to check each array
before importing it.
+ private val stream = ArrowArrayStream.allocateNew(allocator)
+ private val dictionaries = new CDataDictionaryProvider()
+ private var root: VectorSchemaRoot = _
+ private var batch: ColumnarBatch = _
+
+ try {
+ openStream()
+ root = VectorSchemaRoot.create(importSchema(), allocator)
+ checkSchema()
+ batch = new
ColumnarBatch(root.getFieldVectors.asScala.map(columnVector).toArray)
+ } catch {
+ case e: Throwable =>
+ close()
+ throw e
+ }
+
+ private def openStream(): Unit = {
+ val library = NativeLibraries.get(location)
+ call(name, "read")(optional(library.read(readerState, partition,
stream.memoryAddress())))
+ .getOrElse {
+ throw new SparkUnsupportedOperationException(
+ "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name))
+ }
+ }
+
+ private def importSchema(): Schema = {
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ try {
+ call(name, "read")(stream.getSchema(arrowSchema))
+ val imported = call(name, "read")(Data.importSchema(allocator,
arrowSchema, dictionaries))
+ new Schema(
+ imported.getFields.asScala.map(withMapEntryNames).asJava,
imported.getCustomMetadata)
+ } finally {
+ arrowSchema.close()
+ }
+ }
+
+ /**
+ * Names the fields of the entries of the maps `key` and `value`, which
ArrowColumnVector looks
+ * up by name. Arrow does not fix these names: arrow-rs names them `keys`
and `values`, for
+ * example. The arrays are imported by position, so the names do not change
how they are read.
+ */
+ private def withMapEntryNames(field: Field): Field = {
+ val children = field.getChildren.asScala.map(withMapEntryNames).toSeq
+ val renamed = (field.getType, children) match {
+ case (_: ArrowType.Map, Seq(entries)) if entries.getChildren.size == 2 =>
+ val Seq(key, value) = entries.getChildren.asScala.toSeq
+ Seq(new Field(entries.getName, entries.getFieldType, Seq(
+ new Field(MapVector.KEY_NAME, key.getFieldType, key.getChildren),
+ new Field(MapVector.VALUE_NAME, value.getFieldType,
value.getChildren)).asJava))
+ case _ => children
+ }
+ new Field(field.getName, field.getFieldType, renamed.asJava)
+ }
+
+ private def checkSchema(): Unit = {
+ def isDictionaryEncoded(field: Field): Boolean = {
+ field.getDictionary != null ||
field.getChildren.asScala.exists(isDictionaryEncoded)
+ }
+ root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field
=>
+ throw error(name, "read", s"The column ${field.getName} is
dictionary-encoded, which " +
+ "is not supported.")
+ }
+ // The Arrow types that convert to Spark types, but that Spark does not
read: ArrowColumnVector
+ // cannot read decimals of other widths than 128 bits, and the lengths of
the elements of a
+ // list view cannot be checked without a pass over all its lists (see
checkArrays).
+ def unsupportedType(field: Field): Option[ArrowType] = field.getType match
{
+ case decimal: ArrowType.Decimal if decimal.getBitWidth != 128 =>
Some(decimal)
+ case listView: ArrowType.ListView => Some(listView)
+ case _ =>
field.getChildren.asScala.view.flatMap(unsupportedType).headOption
+ }
+ root.getSchema.getFields.asScala.foreach { field =>
+ unsupportedType(field).foreach { arrowType =>
+ throw error(name, "read", s"The column ${field.getName} contains the
Arrow type " +
+ s"$arrowType, which is not supported.")
+ }
+ }
+ val actual = ArrowUtils.fromArrowSchema(root.getSchema)
+ // The schema to read as the library got it: Arrow does not carry
everything, such as the
+ // collations of strings, the fields of intervals and user-defined types,
which still apply to
+ // 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 ArrowColumnVector(vector)
Review Comment:
Thanks, good catch: confirmed and fixed in aa0f3ba4dd77. The column vectors
of the reader are no longer freeable between batches: `ReusedArrowColumnVector`
overrides `closeIfFreeable` to do nothing, as `WritableColumnVector` does. The
reader closes the Arrow vectors when it is closed. New test: "read structs in
several batches" reads a struct column in batches of 3 rows through whole-stage
codegen. Without the fix, importing the second batch fails in
`ColumnarToRowExec.advanceBatch` with "should have as many children as in the
schema: found 0 expected 2", because closing the struct also closed its child
vectors. The same test also reads a nested field.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
