diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index b118c9d2774f..8261d4e3632f 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -952,6 +952,15 @@ jobs: export MAVEN_OPTS="-Xss64m -Xmx4g -Xms4g -XX:ReservedCodeCacheSize=128m -Dorg.slf4j.simpleLogger.defaultLogLevel=WARN" export MAVEN_CLI_OPTS="--no-transfer-progress" ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws clean install + # Native data sources are called with the FFM API, which requires Java 22 or later, so the + # other jobs, on Java 17, skip the tests that load a library. The tests build their library + # with the C++ compiler of the runner. + - name: Test native data sources + run: | + export MAVEN_OPTS="-Xss64m -Xmx4g -Xms4g -XX:ReservedCodeCacheSize=128m -Dorg.slf4j.simpleLogger.defaultLogLevel=WARN" + export MAVEN_CLI_OPTS="--no-transfer-progress" + ./build/mvn $MAVEN_CLI_OPTS -pl sql/core -Dtest=none -DwildcardSuites=org.apache.spark.sql.execution.datasources.v2.ffi.NativeDataSourceSuite,org.apache.spark.sql.execution.datasources.v2.ffi.NativeDataSourceClusterSuite test + ./build/mvn $MAVEN_CLI_OPTS -pl sql/connect/server -Dtest=none -DwildcardSuites=org.apache.spark.sql.connect.NativeDataSourceConnectSuite test # Documentation build docs: diff --git a/LICENSE-binary b/LICENSE-binary index 7051e09a6383..53dec42594ef 100644 --- a/LICENSE-binary +++ b/LICENSE-binary @@ -301,6 +301,7 @@ joda-time:joda-time net.sf.opencsv:opencsv net.sf.supercsv:super-csv net.sf.jpam:jpam +org.apache.arrow:arrow-c-data org.apache.arrow:arrow-compression org.apache.arrow:arrow-format org.apache.arrow:arrow-memory-core diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 3a796a031afa..e7ecef803e7b 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -2306,6 +2306,12 @@ ], "sqlState" : "0A000" }, + "DATA_SOURCE_CONTINUOUS_WRITE_NOT_SUPPORTED" : { + "message" : [ + ": Continuous write is not supported." + ], + "sqlState" : "0A000" + }, "DATA_SOURCE_DELTA_BATCH_WRITE_NOT_SUPPORTED" : { "message" : [ ": Delta batch write is not supported." @@ -4792,6 +4798,57 @@ ], "sqlState" : "42000" }, + "INVALID_NATIVE_DATA_SOURCE_LIBRARY" : { + "message" : [ + "Invalid native data source library :" + ], + "subClass" : { + "CANNOT_LOAD" : { + "message" : [ + "" + ] + }, + "REPLACED" : { + "message" : [ + "The library was replaced after Spark loaded it: Spark loaded the build with the SHA-256 checksum , but the file now has the checksum . A JVM cannot load another build of a library from the same path, so restart the Spark application to use the new build." + ] + }, + "UNSUPPORTED_ABI_VERSION" : { + "message" : [ + "It requires version of the native data source interface, but Spark supports version ." + ] + } + }, + "sqlState" : "22023" + }, + "INVALID_NATIVE_DATA_SOURCE_PACKAGE" : { + "message" : [ + "Invalid native data source package :" + ], + "subClass" : { + "INVALID_LIBRARY" : { + "message" : [ + "Cannot load the library : " + ] + }, + "INVALID_MANIFEST" : { + "message" : [ + "" + ] + }, + "UNSUPPORTED_ABI_VERSION" : { + "message" : [ + "It requires version of the native data source interface, but Spark supports version ." + ] + }, + "UNSUPPORTED_PLATFORM" : { + "message" : [ + "It has no library for the platform . It has libraries for: ." + ] + } + }, + "sqlState" : "22023" + }, "INVALID_NON_DETERMINISTIC_EXPRESSIONS" : { "message" : [ "The operator expects a deterministic expression, but the actual expression is ." @@ -6262,6 +6319,42 @@ ], "sqlState" : "0A000" }, + "NATIVE_DATA_SOURCE_ERROR" : { + "message" : [ + "Failed to native data source : " + ], + "sqlState" : "38000" + }, + "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND" : { + "message" : [ + "Cannot find the native data source library with the SHA-256 checksum , which is installed at on the driver. Install the same build of the library on every node, at the same path or in the native library path." + ], + "sqlState" : "42K02" + }, + "NATIVE_DATA_SOURCE_PACKAGE_CONFLICT" : { + "message" : [ + "More than one native data source package under `spark.sql.dataSource.native.paths` provides the data source : . Remove all but one of them." + ], + "sqlState" : "42710" + }, + "NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND" : { + "message" : [ + "Cannot find the native data source package with the SHA-256 checksum . Add the package with `spark.addArtifact`, or install the same package under `spark.sql.dataSource.native.paths` on every node." + ], + "sqlState" : "42K02" + }, + "NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION" : { + "message" : [ + "Native data sources require Java or later, but Spark runs on Java ." + ], + "sqlState" : "0A000" + }, + "NATIVE_DATA_SOURCE_UNSUPPORTED_PLATFORM" : { + "message" : [ + "Native data sources are not supported on the platform : the JNI library of the C data interface of Arrow Java, with which Spark exchanges their data, failed to load: " + ], + "sqlState" : "0A000" + }, "NEAREST_BY_JOIN" : { "message" : [ "Invalid nearest-by join." diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index 50f978819c3d..b088d623ea7c 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -21,6 +21,7 @@ antlr4-runtime/4.13.1//antlr4-runtime-4.13.1.jar aopalliance-repackaged/3.0.6//aopalliance-repackaged-3.0.6.jar arpack/3.2.0//arpack-3.2.0.jar arpack_combined_all/0.1//arpack_combined_all-0.1.jar +arrow-c-data/19.0.0//arrow-c-data-19.0.0.jar arrow-compression/19.0.0//arrow-compression-19.0.0.jar arrow-format/19.0.0//arrow-format-19.0.0.jar arrow-memory-core/19.0.0//arrow-memory-core-19.0.0.jar diff --git a/dev/make-distribution.sh b/dev/make-distribution.sh index 4a2bd0991012..fc73cbc9ac05 100755 --- a/dev/make-distribution.sh +++ b/dev/make-distribution.sh @@ -223,6 +223,11 @@ if [ -f "$DISTDIR"/jars/datasketches-memory-3.0.2.jar ]; then zip -d "$DISTDIR"/jars/datasketches-memory-3.0.2.jar org/apache/datasketches/memory/internal/ResourceImpl.class fi +# Copy the C header that the libraries of native data sources implement +mkdir -p "$DISTDIR/include" +cp "$SPARK_HOME/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h" \ + "$DISTDIR/include/" + # Only create the yarn directory if the yarn artifacts were built. if [ -f "$SPARK_HOME"/common/network-yarn/target/scala*/spark-*-yarn-shuffle.jar ]; then mkdir "$DISTDIR/yarn" diff --git a/docs/_data/menu-sql.yaml b/docs/_data/menu-sql.yaml index a100f0ccb1ce..311261454611 100644 --- a/docs/_data/menu-sql.yaml +++ b/docs/_data/menu-sql.yaml @@ -75,6 +75,19 @@ url: sql-v2-data-sources.html#streaming - text: Further Reading url: sql-v2-data-sources.html#further-reading +- text: Native Data Sources + url: sql-native-data-sources.html + subitems: + - text: Overview + url: sql-native-data-sources.html#overview + - text: How It Works + url: sql-native-data-sources.html#how-it-works + - text: Packages + url: sql-native-data-sources.html#packages + - text: Example in Rust + url: sql-native-data-sources.html#example-in-rust + - text: Example in C++ + url: sql-native-data-sources.html#example-in-c - text: Performance Tuning url: sql-performance-tuning.html subitems: diff --git a/docs/sql-native-data-sources.md b/docs/sql-native-data-sources.md new file mode 100644 index 000000000000..4bc9b894f50c --- /dev/null +++ b/docs/sql-native-data-sources.md @@ -0,0 +1,798 @@ +--- +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 C interface declared in the header +[`spark_datasource.h`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h). 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. + +Spark calls native data sources with the +[Foreign Function and Memory API](https://openjdk.org/jeps/454) of Java, which is final from Java +22, so they require Java 22 or later on the driver and the executors. With an earlier version of +Java, a query that uses a native data source fails with +`NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION`. Spark exchanges the data with the C data interface +of Arrow Java, whose JNI library is available for Linux with glibc and macOS on x86_64 and aarch64, +and for Windows on x86_64. On other platforms, such as Linux on ppc64le or s390x, or Linux with +musl, as Alpine Linux, a query that uses a native data source fails with +`NATIVE_DATA_SOURCE_UNSUPPORTED_PLATFORM`. + +### When to Use a Native Data Source + +A native data source fits a source whose client exists only in native code, such as an in-house +storage system written in C++ or the C SDK of a vendor, and that Spark needs to read and write. +Such a source then needs no JVM connector: the same library works with all the releases of Spark +that support its version of the interface, and with every Spark API, including Spark Connect +clients in any language. + +A native data source provides what Data Source V2 offers to data sources, except for the features +that version 1 of the interface does not cover yet, such as row-level operations: see +[Limitations](#limitations). Extensions of Spark outside of Data Source V2, such as table +functions, SQL commands and optimizer rules, are out of the scope of this interface. A source that +needs them is better served by a JVM connector, which can still read and write through a native +library with the Arrow C stream interface. + +## How It Works + +A native data source library exports a single C function, `spark_ds_init`, which returns a table +of function pointers, a `SparkDsLibrary`, that implements the data source: it needs no JVM headers +or bindings. Each library has its own table, so any number of native data sources can be loaded +side by side. When Spark loads a library, it calls `spark_ds_init` with a `SparkDsHost`, which +gives the library the version of Spark and a function to log messages to the log of Spark. + +A Spark distribution contains the header in its `include` directory, as +`$SPARK_HOME/include/spark_datasource.h`, and the `spark-sql` JAR contains it as +`org/apache/spark/sql/datasource/spark_datasource.h`. The header only depends on the C standard +library, so a library can also copy it into its source tree, to build without Spark, for example +from the source of a release: + +```bash +curl -LO {{site.SPARK_GITHUB_URL}}/raw/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h +``` + +A library built against the header works with all the releases of Spark that support its version, +`SPARK_DS_ABI_VERSION`: see [Compatibility](#compatibility). + +The schemas are Arrow schemas, with the types of version 1 that `spark_datasource.h` lists, such +as `int64` for `bigint`, `utf8` for `string` and `timestamp[us]` for `timestamp`: Spark fails an +operation whose schema has another type, such as an interval, before it calls the library. The data +that a library returns must agree with the nullability of the schema to read. + +Spark calls the functions of the table as follows: + +- **Data sources**: for each schema inference, scan and write, Spark creates a data source on the + driver with `create`, which receives the name of the data source and the options of the user in + a `SparkDsCreateInfo`, and returns the capabilities of the data source, such as + `SPARK_DS_BATCH_READ` or `SPARK_DS_STREAMING_WRITE`, and closes it right after with `close`. + Spark reports the operations that a data source does not declare as unsupported. It also creates + a data source with the options of the table while it analyzes a query, to reject a batch write + that the data source does not support before it computes the data to write. When the user does + not specify a schema, Spark asks for it with `schema`. +- **Batch reads**: on the driver, Spark creates a reader with `create_reader`, pushes operations + down with `push_predicates`, `push_limit` and `prune_columns`, plans the partitions with + `partitions`, and serializes the state of the reader with `serialize_reader`. Then it calls + `read` on the executors, in one task for each partition, and `read` exports an + `ArrowArrayStream`. Its callbacks can be called from any thread, but never concurrently. +- **Streaming reads**: on the driver, Spark creates a stream reader with `create_stream_reader`. + For each micro-batch, it asks for the `latest_offset`, plans the partitions with + `stream_partitions`, and calls `read` on the executors. Offsets are defined by the library, as + compact JSON on a single line: Spark stores each offset as a line of the offset log of the + checkpoint, and rejects an offset that is not valid JSON or contains a line break, such as + pretty-printed JSON. +- **Writes**: on the driver, right before the tasks of the write (or of a micro-batch) start, Spark + creates a writer with `create_writer`, which gets a `SparkDsWriteInfo` with the mode of the write, + `SPARK_DS_WRITE_APPEND` or `SPARK_DS_WRITE_TRUNCATE`, and the ID of the query of a streaming + write. With adaptive query execution, it first runs the shuffle stages of the data to write, and + creates no writer if they fail, but the last stage runs in the tasks of the write, so a writer can + still be aborted because computing the data failed: do irreversible work, such as removing the + existing data for an overwrite, in `commit`. On the executors, each task creates a data writer + with `create_data_writer`, writes the batches of its data with `write`, and calls + `commit_data_writer`, or `abort_data_writer` if the task or the commit fails. Then the driver + calls `commit` with the messages of the tasks, or `abort` if the write failed. A streaming write + commits each micro-batch with its epoch ID, and the ID of the query stays the same when the query + restarts from its checkpoint: together, they identify the output of a micro-batch, so that a + library can skip an epoch that it already committed. A micro-batch streaming query creates a new + stream writer for each micro-batch, and Spark closes the previous one when it creates the next + one. Because Spark is not told when a streaming query stops, the last stream writer of a query is + closed once it is no longer referenced, or when the query is restarted. + +The library defines the format of the partitions, of the state of its readers and writers, and of +the commit messages, which Spark sends to the executors and back as byte arrays. The library +returns them, and the offsets, by passing them to a `SparkDsSink` callback, which copies them. A +function reports an error by returning a non-zero value, with the message returned by +`last_error`, and Spark reports it as a `NATIVE_DATA_SOURCE_ERROR`. A C++ exception or a Rust panic +must not cross the interface. Each handle that a `create` function returns is released exactly +once, by the `close` function of its kind. The close functions cannot fail, and they can run after +the query, so do the work that can fail, such as flushing data, in `commit`. Strings are UTF-8, and +the JSON of the predicates is ASCII, with the other characters escaped. + +A library only sets the functions of the operations it supports, and leaves the others `NULL`. A +data source must not declare a capability whose functions are `NULL`: + +| Capability | Required functions | Optional functions | +|-----------|--------------------|--------------------| +| All | `create`, `close` | `last_error`, `schema` | +| `SPARK_DS_BATCH_READ` | `create_reader`, `partitions`, `read`, `close_reader` | `push_predicates`, `push_limit`, `prune_columns`, `serialize_reader` | +| `SPARK_DS_MICRO_BATCH_READ` | `create_stream_reader`, `initial_offset`, `latest_offset`, `stream_partitions`, `read`, `close_stream_reader` | `serialize_stream_reader`, `commit_offset` | +| `SPARK_DS_BATCH_WRITE`, `SPARK_DS_STREAMING_WRITE` | `create_writer`, `commit`, `abort`, `close_writer`, `create_data_writer`, `write`, `commit_data_writer`, `abort_data_writer`, `close_data_writer` | `serialize_writer` | + +`SPARK_DS_TRUNCATE` declares that a writer supports `SPARK_DS_WRITE_TRUNCATE`: the save mode +`overwrite` of a batch write, and the output mode `complete` of a streaming write. + +The comments of [`spark_datasource.h`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h) +specify each function, including the ownership of the Arrow structs and the JSON format of the +pushed predicates. `push_predicates` tells Spark how the reader uses each predicate: not at all, +only to skip data (`SPARK_DS_PREDICATE_PARTIAL`), or to only return the rows for which it is true +(`SPARK_DS_PREDICATE_COMPLETE`). Spark shows the predicates that the reader uses as pushed in the +plan of the query, and evaluates all the pushed predicates again after reading either way, so a +library can use them to skip data, for example files or row groups whose statistics show that no +row matches, but it must never skip a row for which the predicates are true under the semantics of +Spark SQL: NaN values are equal to each other and greater than any other value, `-0.0` is equal to +`0.0`, comparisons with null are null, strings compare by their UTF-8 bytes, and binary values by +their unsigned bytes. The comparisons of IEEE 754 and the total orders of Arrow libraries, such as +the comparison kernels of arrow-rs, differ for some of these values. Spark does not pass +comparisons and `IN` on arrays, maps and structs, which it compares with other rules, such as the +nulls in them equal to each other. The jar of Spark SQL contains test cases for a library, +`org/apache/spark/sql/datasource/predicate-cases.json`: predicates as Spark pushes them, with the +rows of a table for which they are true: + +```bash +unzip -p "$SPARK_HOME"/jars/spark-sql_*.jar org/apache/spark/sql/datasource/predicate-cases.json +``` + +The arrays that a library exports must have an offset of 0: copy a sliced array, such as the +result of `RecordBatch::slice` in arrow-rs, before exporting it. A batch must have no null rows, +since Arrow Java ignores the validity of the struct of a batch: mark the values of its columns as +null instead. 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. + +`read` and `create_data_writer` get a `SparkDsTaskContext`, with the IDs of the task, and its +partition and attempt, and an `is_interrupted` function that returns whether the task was killed, +for example because its query was canceled. Spark cannot interrupt native code: when a task is +killed, Spark stops reading only once the `get_next` function of the stream returns. So a read or +a write that runs for long, or blocks, for example on network I/O, should check `is_interrupted` +regularly and use timeouts. The task reaper, `spark.task.reaper.enabled` and +`spark.task.reaper.killTimeout`, kills an executor whose killed task does not end. + +### Compatibility + +The function table, `SparkDsLibrary`, and the structs that Spark passes to the library, +`SparkDsHost`, `SparkDsCreateInfo`, `SparkDsTaskContext` and `SparkDsWriteInfo`, start with their +size, so that they can grow. The version of the interface, 1, is `SPARK_DS_ABI_VERSION`, which the +table and the manifest of a package declare. Within a version: + +- fields are only appended to these structs, and the existing fields and the semantics of the + existing functions never change; +- Spark only reads the functions that fit in the size of the table of a library, so a library built + against an earlier header keeps working with later releases of Spark, with the functions that it + does not know as `NULL`. A library must only read the fields that fit in the size of a struct + that Spark passes, so that it also runs on earlier releases of Spark that support its version; +- new values may appear where the header says so: for example, later releases may push down + predicates with other names, which a library must ignore if it does not know them. Spark only + uses a new capability or write mode with a data source that declares it. + +An incompatible change comes with a new version: `spark_ds_init` gets the newest version that Spark +supports, in `SparkDsHost`, and returns a table of a version that the library implements. 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 functions or fields appended to the structs of the interface: + +- statistics: Spark estimates the size of a native scan with `spark.sql.defaultSizeInBytes`, so it + does not plan a native scan as the broadcast side of a join, unless the query has a broadcast + hint. Adaptive query execution can still broadcast it at run time, from the size of a shuffle of + the scan; +- custom metrics, shown in the Spark UI; +- admission control for streaming reads, such as a maximum number of rows for each micro-batch; +- continuous processing (`Trigger.Continuous`): a continuous query that reads from or writes to a + native data source fails when it starts; +- the pushdown of aggregates and of top-N queries (`ORDER BY` with `LIMIT`), and reporting the + partitioning and the sort order of the data, with which Spark avoids shuffles and sorts, as in + storage-partitioned joins; +- row-level operations, `UPDATE`, `DELETE` and `MERGE INTO`, and metadata columns; +- memory accounting: see below. + +Extensions of Spark outside of Data Source V2, such as table functions, SQL commands and optimizer +rules, are out of the scope of this interface: see +[When to Use a Native Data Source](#when-to-use-a-native-data-source). + +### Memory + +Spark does not manage the memory of native data sources. The memory that a library allocates, and +the Arrow buffers that it exports, which Spark reads without copying them, are not accounted by the +memory manager of Spark, nor limited by `spark.memory.offHeap.size`. Neither are the Arrow batches +that Spark writes, which end at `spark.sql.execution.arrow.maxRecordsPerBatch` rows or once they +hold `spark.sql.execution.arrow.maxBytesPerBatch` bytes, while the size of the batches that a +library reads is up to the library. This memory comes out of the memory overhead of the driver and +the executors, `spark.driver.memoryOverhead` and `spark.executor.memoryOverhead`, which are 10% of +their memory, and at least 384 MiB, by default. Increase it for libraries that use much memory, and +bound the memory that a library uses, for example with the size of its batches: on YARN or +Kubernetes, an executor that uses more memory than it requested is killed. + +## Finding Native Data Sources + +When no Java or Python data source has the name that a query uses, Spark looks for a native data +source with that name: in the packages of the session, and then in the installed libraries. A Java +data source with the same name takes precedence over a native one, and so does a Python data +source. + +### Installed Libraries + +Like a Python data source installed in the Python path, a native data source library installed +like other native libraries is found automatically, without any configuration. The library is +named after the data source, in lower case, with the prefix `spark_datasource_`: the library of +`rust_range` is `libspark_datasource_rust_range.so` on Linux, +`libspark_datasource_rust_range.dylib` on macOS and `spark_datasource_rust_range.dll` on Windows. +Spark passes the name of the data source to the `create` function of the library. A library that +implements several data sources is installed under each of their names, for example with symbolic +links. + +Spark looks for the library in the native library path, which consists of, in order: + +1. the directories of `java.library.path`, where the JVM finds native libraries. They include + `LD_LIBRARY_PATH` on Linux, `DYLD_LIBRARY_PATH` on macOS and `PATH` on Windows, which also + contain the directories set with `spark.driver.extraLibraryPath` and + `spark.executor.extraLibraryPath`. +2. the `lib` directory of each installation prefix in `PATH`: `/lib` for each directory + `/bin` in `PATH`. + +So Spark finds the libraries where native libraries are usually installed: + +- `/usr/local/lib`, the default directory of `make install`, CMake (`cmake --install`) and + [cargo-c](https://github.com/lu-zero/cargo-c) (`cargo cinstall`), and of Homebrew on Intel macs; +- `/opt/homebrew/lib`, for Homebrew on Apple silicon; +- `$CONDA_PREFIX/lib` and `$VIRTUAL_ENV/lib`, in an activated conda environment or Python virtual + environment; +- `/usr/lib`, for the packages of a Linux distribution. + +Cargo does not install libraries: copy the library that `cargo build --release` builds into one of +these directories. A library in any other directory is found with `LD_LIBRARY_PATH`, or with +`spark.driver.extraLibraryPath` and `spark.executor.extraLibraryPath`. + +The executors load the library from the same path as the driver, or else find it in their own +native library path, which they get from their own environment, such as the `PATH` of the +executor processes. So install the library on every node, for example in the container image of +the cluster. An executor only loads the same build of the library as the driver, which it checks +with the SHA-256 checksum of the file, because the partitions and the state that the driver +serializes are only meant for that build: a query fails with +`NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND` on a node where a different build is installed. + +A JVM loads a library from a path only once. So after a new build is installed at the same path, +the queries of an application that loaded the previous build fail with +`INVALID_NATIVE_DATA_SOURCE_LIBRARY.REPLACED`, on the driver or on the executors that loaded it, +rather than running mixed builds: restart the application to use the new build. + +### Packages + +A native data source package distributes a library with a session, for one or more platforms. It +is a zip file with the extension `.sparkpkg`, which contains a manifest named +`spark-native-datasource.json` and the libraries: + +```json +{ + "abiVersion": 1, + "dataSources": ["rust_range"], + "libraries": { + "linux-x86_64": "linux-x86_64/libspark_datasource_rust_range.so", + "linux-aarch64": "linux-aarch64/libspark_datasource_rust_range.so", + "osx-aarch64": "osx-aarch64/libspark_datasource_rust_range.dylib" + } +} +``` + +- `abiVersion` is the oldest version of `spark_datasource.h` that the library implements, + currently 1. Spark uses a package whose `abiVersion` it supports, or is older than the newest + version that it supports, and `spark_ds_init` returns a table of a version that both support. + Spark ignores the fields of the manifest that it does not know, so that later releases can add + fields. +- `dataSources` lists the names of the data sources that the library implements. The names are + case-insensitive, and Spark passes the name to `create`. +- `libraries` maps each platform, `-` with an `os` of `linux`, `osx` or `windows` and an + `arch` of `x86_64` or `aarch64`, to the path of the library in the package. `linux` is Linux with + glibc: Linux with musl, as Alpine Linux, is `linux--musl`, such as `linux-x86_64-musl`, on + which native data sources do not run yet. + +A session finds the packages in two places, and the executors get them automatically: + +- **Artifacts of a session**: `spark.addArtifact` adds a package to a session, in Scala, Java and + Python, and with Spark Connect: + + ```python + spark.addArtifact("/path/to/rust_range.sparkpkg") + spark.read.format("rust_range").option("end", 100).load().show() + ``` + +- **Configured paths**: the packages listed in `spark.sql.dataSource.native.paths`, or in the + directories it lists, on the driver. A package can be replaced in place: the next queries use + the new version. + +The packages of a session take precedence over the installed libraries, and the package that a +session added last over its other packages, also the configured ones: a session cannot remove an +artifact, or add a file with the same name again, so it upgrades a package by adding the new +version under another file name, such as `rust_range-1.1.sparkpkg`. Distinct file names are also +needed because the sessions that are not isolated share the files they add. A query fails if more +than one configured package provides its data source. Spark ignores, and logs, a package that it +cannot read or whose version it does not support, unless its manifest lists the data source that a +query uses and no other package or installed library provides it. + +## Example in Rust + +This data source reads the numbers `[0, end)` as a column `id`, in two partitions. It only depends +on the `arrow-array` and `arrow-schema` crates of [arrow-rs](https://github.com/apache/arrow-rs), +which export schemas and streams with `FFI_ArrowSchema` and `FFI_ArrowArrayStream`. It declares +`SparkDsLibrary` as a `#[repr(C)]` struct up to `read`, the last function that it implements, since +Spark only reads the functions that fit in the size of the table, and leaves the functions that it +does not implement `None`. The fields must follow the order of `spark_datasource.h`, which +[bindgen](https://github.com/rust-lang/rust-bindgen) can also generate from the header. +`Cargo.toml`, where the name of the library is the name that Spark finds when it is installed: + +```toml +[package] +name = "rust_range" +version = "0.1.0" +edition = "2021" + +[lib] +name = "spark_datasource_rust_range" +crate-type = ["cdylib"] + +[dependencies] +arrow-array = { version = "58", features = ["ffi"] } +arrow-schema = { version = "58", features = ["ffi"] } +``` + +`src/lib.rs`: + +{% raw %} +```rust +//! A native data source for Apache Spark, written in Rust. It reads the numbers [0, end) as a +//! column `id`, in two partitions. The option `end` defaults to 10. + +use std::cell::RefCell; +use std::ffi::{c_char, c_int, c_void, CStr, CString}; +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::sync::Arc; + +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator}; +use arrow_schema::ffi::FFI_ArrowSchema; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; + +type Result = std::result::Result>; + +/// `SparkDsSink`: receives a byte array returned to Spark, and copies it. +type Sink = unsafe extern "C" fn(sink: *mut c_void, data: *const u8, length: i64); + +/// A function of `SparkDsLibrary` that this data source does not implement, left NULL. +type Unused = Option; + +/// `SparkDsCreateInfo` of `spark_datasource.h`: the name and the options of a data source. +#[repr(C)] +pub struct SparkDsCreateInfo { + size: i64, + name: *const c_char, + option_keys: *const *const c_char, + option_values: *const *const c_char, + num_options: i64, +} + +/// `SparkDsLibrary` of `spark_datasource.h`, up to `read`, the last function that this data +/// source implements: Spark only reads the functions that fit in `size`. +#[repr(C)] +pub struct SparkDsLibrary { + size: i64, + abi_version: i32, + last_error: unsafe extern "C" fn() -> *const c_char, + create: unsafe extern "C" fn( + info: *const SparkDsCreateInfo, + capabilities: *mut u64, + source: *mut *mut c_void, + ) -> c_int, + schema: unsafe extern "C" fn(source: *mut c_void, out: *mut FFI_ArrowSchema) -> c_int, + create_reader: unsafe extern "C" fn( + source: *mut c_void, + schema: *mut FFI_ArrowSchema, + reader: *mut *mut c_void, + ) -> c_int, + create_stream_reader: Unused, + create_writer: Unused, + close: unsafe extern "C" fn(source: *mut c_void), + push_predicates: Unused, + push_limit: Unused, + prune_columns: Unused, + partitions: unsafe extern "C" fn(reader: *mut c_void, add: Sink, sink: *mut c_void) -> c_int, + serialize_reader: Unused, + close_reader: unsafe extern "C" fn(reader: *mut c_void), + initial_offset: Unused, + latest_offset: Unused, + stream_partitions: Unused, + serialize_stream_reader: Unused, + commit_offset: Unused, + close_stream_reader: Unused, + read: unsafe extern "C" fn( + task: *const c_void, + reader_state: *const u8, + reader_state_length: i64, + partition: *const u8, + partition_length: i64, + out: *mut FFI_ArrowArrayStream, + ) -> c_int, +} + +const SPARK_DS_BATCH_READ: u64 = 1; + +static LIBRARY: SparkDsLibrary = SparkDsLibrary { + size: std::mem::size_of::() as i64, + abi_version: 1, + last_error, + create, + schema, + create_reader, + create_stream_reader: None, + create_writer: None, + close, + push_predicates: None, + push_limit: None, + prune_columns: None, + partitions, + serialize_reader: None, + close_reader, + initial_offset: None, + latest_offset: None, + stream_partitions: None, + serialize_stream_reader: None, + commit_offset: None, + close_stream_reader: None, + read, +}; + +/// The function that Spark calls when it loads the library. +#[no_mangle] +pub unsafe extern "C" fn spark_ds_init( + _host: *const c_void, + library: *mut *const SparkDsLibrary, +) -> c_int { + *library = &LIBRARY; + 0 +} + +struct RangeSource { + end: i64, +} + +struct RangeReader { + end: i64, +} + +thread_local! { + static LAST_ERROR: RefCell = RefCell::new(CString::default()); +} + +fn arrow_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])) +} + +/// Runs `f`, and reports its error or panic to Spark: neither may cross into the JVM. +fn call(f: impl FnOnce() -> Result<()>) -> c_int { + let message = match catch_unwind(AssertUnwindSafe(f)) { + Ok(Ok(())) => return 0, + Ok(Err(e)) => e.to_string(), + Err(_) => "the data source panicked".to_string(), + }; + LAST_ERROR.with(|error| *error.borrow_mut() = CString::new(message).unwrap_or_default()); + 1 +} + +unsafe extern "C" fn last_error() -> *const c_char { + LAST_ERROR.with(|error| error.borrow().as_ptr()) +} + +unsafe extern "C" fn create( + info: *const SparkDsCreateInfo, + capabilities: *mut u64, + source: *mut *mut c_void, +) -> c_int { + call(|| { + let info = &*info; + let mut range = RangeSource { end: 10 }; + for i in 0..info.num_options as usize { + if CStr::from_ptr(*info.option_keys.add(i)).to_str()? == "end" { + range.end = CStr::from_ptr(*info.option_values.add(i)).to_str()?.parse()?; + } + } + *capabilities = SPARK_DS_BATCH_READ; + *source = Box::into_raw(Box::new(range)) as *mut c_void; + Ok(()) + }) +} + +unsafe extern "C" fn schema(_source: *mut c_void, out: *mut FFI_ArrowSchema) -> c_int { + call(|| { + std::ptr::write(out, FFI_ArrowSchema::try_from(arrow_schema().as_ref())?); + Ok(()) + }) +} + +unsafe extern "C" fn create_reader( + source: *mut c_void, + _schema: *mut FFI_ArrowSchema, + reader: *mut *mut c_void, +) -> c_int { + call(|| { + let end = (*(source as *const RangeSource)).end; + *reader = Box::into_raw(Box::new(RangeReader { end })) as *mut c_void; + Ok(()) + }) +} + +unsafe extern "C" fn close(source: *mut c_void) { + drop(Box::from_raw(source as *mut RangeSource)); +} + +unsafe extern "C" fn partitions(reader: *mut c_void, add: Sink, sink: *mut c_void) -> c_int { + call(|| { + let end = (*(reader as *const RangeReader)).end; + for (start, end) in [(0, end / 2), (end / 2, end)] { + let bytes = [start.to_le_bytes(), end.to_le_bytes()].concat(); + add(sink, bytes.as_ptr(), bytes.len() as i64); + } + Ok(()) + }) +} + +unsafe extern "C" fn close_reader(reader: *mut c_void) { + drop(Box::from_raw(reader as *mut RangeReader)); +} + +unsafe extern "C" fn read( + _task: *const c_void, + _reader_state: *const u8, + _reader_state_length: i64, + partition: *const u8, + partition_length: i64, + out: *mut FFI_ArrowArrayStream, +) -> c_int { + call(|| { + let bytes = std::slice::from_raw_parts(partition, partition_length as usize); + let start = i64::from_le_bytes(bytes[0..8].try_into()?); + let end = i64::from_le_bytes(bytes[8..16].try_into()?); + let ids = Int64Array::from_iter_values(start..end); + let batch = RecordBatch::try_new(arrow_schema(), vec![Arc::new(ids)])?; + let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema()); + std::ptr::write(out, FFI_ArrowArrayStream::new(Box::new(batches))); + Ok(()) + }) +} +``` +{% endraw %} + +Build the library, here on macOS, and install it: + +```bash +cargo build --release +cp target/release/libspark_datasource_rust_range.dylib /usr/local/lib/ +``` + +Spark then finds the data source `rust_range` automatically: + +```python +spark.read.format("rust_range").option("end", 100).load().show() +``` + +Or package it, to add it to a session with `spark.addArtifact`: + +```bash +mkdir -p package/osx-aarch64 +cp target/release/libspark_datasource_rust_range.dylib package/osx-aarch64/ +cat > package/spark-native-datasource.json < +#include +#include +#include +#include + +#include +#include + +#include "spark_datasource.h" + +namespace { + +struct RangeSource { + int64_t end = 10; +}; + +struct RangeReader { + int64_t end; +}; + +thread_local std::string last_error; + +std::shared_ptr 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 +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()); +} + +const char* LastError() { return last_error.c_str(); } + +int Create(const SparkDsCreateInfo* info, uint64_t* capabilities, void** source) { + return Call([&] { + auto range = std::make_unique(); + for (int64_t i = 0; i < info->num_options; i++) { + if (std::string(info->option_keys[i]) == "end") { + range->end = std::stoll(info->option_values[i]); + } + } + *capabilities = SPARK_DS_BATCH_READ; + *source = range.release(); + }); +} + +int GetSchema(void*, struct ArrowSchema* out) { + return Call([&] { Check(arrow::ExportSchema(*Schema(), out)); }); +} + +int CreateReader(void* source, struct ArrowSchema*, void** reader) { + return Call([&] { *reader = new RangeReader{static_cast(source)->end}; }); +} + +void Close(void* source) { delete static_cast(source); } + +int Partitions(void* reader, SparkDsSink add_partition, void* sink) { + return Call([&] { + int64_t end = static_cast(reader)->end; + int64_t bounds[2][2] = {{0, end / 2}, {end / 2, end}}; + for (const auto& partition : bounds) { + add_partition(sink, reinterpret_cast(partition), sizeof(partition)); + } + }); +} + +void CloseReader(void* reader) { delete static_cast(reader); } + +int Read(const SparkDsTaskContext*, const uint8_t*, int64_t, const uint8_t* partition, int64_t, + struct ArrowArrayStream* out) { + return Call([&] { + int64_t bounds[2]; + std::memcpy(bounds, partition, sizeof(bounds)); + arrow::Int64Builder ids; + for (int64_t id = bounds[0]; id < bounds[1]; id++) Check(ids.Append(id)); + std::shared_ptr 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)); + }); +} + +// The functions of the data source. The ones that it does not implement are NULL. +const SparkDsLibrary kLibrary = { + .size = sizeof(SparkDsLibrary), + .abi_version = SPARK_DS_ABI_VERSION, + .last_error = LastError, + .create = Create, + .schema = GetSchema, + .create_reader = CreateReader, + .close = Close, + .partitions = Partitions, + .close_reader = CloseReader, + .read = Read, +}; + +} // namespace + +extern "C" SPARK_DS_EXPORT int spark_ds_init(const SparkDsHost*, const SparkDsLibrary** library) { + *library = &kLibrary; + return 0; +} +``` +{% endraw %} + +Build the library, here on macOS with Arrow C++ 24, which requires C++20, and install it, or +package it as above: + +```bash +c++ -std=c++20 -O2 -shared -fPIC -I"$SPARK_HOME/include" -I"$ARROW_HOME/include" \ + -L"$ARROW_HOME/lib" -larrow -o libspark_datasource_cpp_range.dylib cpp_range.cc +cp libspark_datasource_cpp_range.dylib /usr/local/lib/ +``` + +A library loaded by Spark must find its dependencies, such as `libarrow`, on every node: link them +statically, or install them on the nodes. + +## Configuration + +| Property Name | Default | Meaning | +|---------------|---------|---------| +| `spark.sql.dataSource.native.enabled` | true | Whether Spark loads native data sources. It is a static configuration, so it can only be set when the Spark application starts. | +| `spark.sql.dataSource.native.paths` | (none) | Comma-separated list of native data source packages, and of directories that contain them, on the local file system of the driver. They take precedence over the installed libraries. | + +## Security + +A native data source runs native code in the driver and executor processes, with their +privileges, and is not sandboxed. Spark finds installed libraries by their file name, in the +directories of the native library path, and loads one when a query uses its name and no Java or +Python data source has that name, much like it imports the Python data sources installed as +`pyspark_*` packages. So only install libraries and add packages that you trust, and do not let +untrusted users write to the directories of the native library path, or to the directories in +`spark.sql.dataSource.native.paths`. A user who can add artifacts to a session, for example a Spark +Connect client, can also add a package, as they can add JAR files. Set +`spark.sql.dataSource.native.enabled` to false to prevent loading native data sources, for example +on a server shared by several users. diff --git a/docs/sql-v2-data-sources.md b/docs/sql-v2-data-sources.md index 4b4d0f35cb4e..df1674e98ea1 100644 --- a/docs/sql-v2-data-sources.md +++ b/docs/sql-v2-data-sources.md @@ -40,7 +40,7 @@ Compared to the earlier Data Source V1 API, DSV2 offers: | Feature | Description | |---------|-------------| -| **Java API** | The connector interfaces are pure Java (`org.apache.spark.sql.connector`), removing the Scala dependency that DSV1 required. A [Python Data Source API](api/python/reference/pyspark.sql/api/pyspark.sql.datasource.DataSource.html) (`pyspark.sql.datasource`) is also available as a wrapper for lightweight connectors written entirely in Python. | +| **Java API** | The connector interfaces are pure Java (`org.apache.spark.sql.connector`), removing the Scala dependency that DSV1 required. A [Python Data Source API](api/python/reference/pyspark.sql/api/pyspark.sql.datasource.DataSource.html) (`pyspark.sql.datasource`) is also available as a wrapper for lightweight connectors written entirely in Python. [Native data sources](sql-native-data-sources.html) can be written entirely in native code, such as Rust or C++. | | **Catalog integration** | Connectors can expose namespaces, tables, views, and functions natively through Spark SQL. | | **Operator pushdown** | Connectors can accept pushed-down filters, required columns, aggregates, limits, offsets, and more. | | **Report partitioning and ordering** | Connectors can report the physical layout of data so that Spark can avoid unnecessary shuffles and sorts. | @@ -473,5 +473,7 @@ A `Table` that declares `STREAMING_WRITE` supports streaming writes through - [Data Sources](sql-data-sources.html) for the user-facing guide to built-in data sources (DSV1). - [Python Data Source API](api/python/reference/pyspark.sql/api/pyspark.sql.datasource.DataSource.html) for writing lightweight connectors entirely in Python. +- [Native Data Sources](sql-native-data-sources.html) for writing connectors entirely in native + code, such as Rust or C++. - [Storage Partition Join](sql-performance-tuning.html#storage-partition-join) for how DSV2 partitioning reporting enables join optimizations. diff --git a/pom.xml b/pom.xml index da0530888357..c2eb3de50b8c 100644 --- a/pom.xml +++ b/pom.xml @@ -2570,6 +2570,25 @@ + + org.apache.arrow + arrow-c-data + ${arrow.version} + + + com.fasterxml.jackson.core + jackson-annotations + + + com.fasterxml.jackson.core + jackson-core + + + io.netty + netty-common + + + org.apache.arrow arrow-compression diff --git a/python/pyspark/sql/connect/client/artifact.py b/python/pyspark/sql/connect/client/artifact.py index 3ab3df4ff593..f39362a0e153 100644 --- a/python/pyspark/sql/connect/client/artifact.py +++ b/python/pyspark/sql/connect/client/artifact.py @@ -237,7 +237,8 @@ def _parse_artifacts( name = f"{name}#{parsed.fragment}" artifact = new_archive_artifact(name, LocalFile(local_path)) - elif file: + elif file or name.endswith(".sparkpkg"): + # Native data source packages are added as files. artifact = new_file_artifact(name, LocalFile(local_path)) elif name.endswith(".jar"): artifact = new_jar_artifact(name, LocalFile(local_path)) diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index 984b02ef6e56..916d4a1640ea 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -41,6 +41,7 @@ no_type_check, overload, ) +from urllib.parse import urlparse from pyspark.conf import SparkConf from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError @@ -2267,10 +2268,14 @@ def addArtifacts( .. versionchanged:: 4.0.0 Supports Spark Classic. + .. versionchanged:: 4.4.0 + Supports native data source packages (``.sparkpkg`` files). + Parameters ---------- *path : tuple of str - Artifact's URIs to add. + Artifact's URIs to add. A native data source package (a ``.sparkpkg`` file) makes + the native data sources it contains available to the session. pyfile : bool Whether to add them as Python dependencies such as .py, .egg, .zip or .jar files. The pyfiles are directly inserted into the path when executing Python functions @@ -2292,6 +2297,18 @@ def addArtifacts( "condition": "True together", }, ) + # Native data source packages are artifacts of the session in the JVM, whatever the + # flags, as with Spark Connect. The JVM resolves file URIs, such as + # file:///tmp/a.sparkpkg, but resolves a relative path against its own working directory, + # so it gets an absolute path instead. Anything else is a path, such as C:/a.sparkpkg on + # Windows, or a file name with a colon, such as a:1.sparkpkg. + for p in path: + if p.endswith(".sparkpkg"): + if urlparse(p).scheme == "file": + self._jsparkSession.addArtifact(p) + else: + self._jsparkSession.addArtifact(os.path.abspath(p)) + path = tuple(p for p in path if not p.endswith(".sparkpkg")) for p in path: normalized_path = os.path.abspath(p) target_dir = os.path.join( diff --git a/python/pyspark/sql/tests/connect/client/test_artifact.py b/python/pyspark/sql/tests/connect/client/test_artifact.py index 6e5ed7174595..286f35d539a6 100644 --- a/python/pyspark/sql/tests/connect/client/test_artifact.py +++ b/python/pyspark/sql/tests/connect/client/test_artifact.py @@ -447,6 +447,21 @@ def test_cache_artifact(self): self.assertEqual(actualHash, expected_hash) self.assertEqual(self.artifact_manager.is_cached_artifact(expected_hash), True) + def test_native_data_source_package(self): + with tempfile.TemporaryDirectory(prefix="test_native_data_source_package") as d: + package_path = os.path.join(d, "my_source.sparkpkg") + with open(package_path, "wb") as f: + f.write(b"package") + requests = list( + self.artifact_manager._create_requests( + package_path, pyfile=False, archive=False, file=False + ) + ) + self.assertEqual(len(requests), 1) + artifacts = requests[0].batch.artifacts + self.assertEqual(len(artifacts), 1) + self.assertEqual(artifacts[0].name, os.path.join("files", "my_source.sparkpkg")) + def test_add_not_existing_artifact(self): with tempfile.TemporaryDirectory(prefix="test_add_not_existing_artifact") as d: with self.assertRaises(FileNotFoundError): diff --git a/python/pyspark/sql/tests/test_artifact.py b/python/pyspark/sql/tests/test_artifact.py index 408b33f1ed74..73f18abc9676 100644 --- a/python/pyspark/sql/tests/test_artifact.py +++ b/python/pyspark/sql/tests/test_artifact.py @@ -16,6 +16,7 @@ # import os import tempfile +from pathlib import Path from pyspark.errors import PySparkRuntimeError from pyspark.sql.functions import assert_true, lit, udf @@ -105,6 +106,47 @@ def test_artifacts_cannot_be_overwritten(self): messageParameters={"normalized_path": pyfile_path}, ) + def test_add_native_data_source_package(self): + # A new session, so that the packages do not affect the other tests. + spark = self.spark.newSession() + with tempfile.TemporaryDirectory(prefix="test_add_native_data_source_package") as d: + + def new_file(name): + file_path = os.path.join(d, name) + with open(file_path, "wb") as f: + f.write(name.encode("utf-8")) + return file_path + + spark.addArtifact(new_file("my_source.sparkpkg")) + # Packages are added to the session whatever the flags, as with Spark Connect, also + # together with other files. + spark.addArtifact(new_file("my_file_source.sparkpkg"), file=True) + spark.addArtifacts(new_file("my_other_source.sparkpkg"), new_file("other.txt")) + # A URI, and paths relative to the working directory of Python, which is not the + # one of the JVM, also with a colon, which is not the scheme of a URI. + spark.addArtifact(Path(new_file("my_uri_source.sparkpkg")).as_uri()) + new_file("my_relative_source.sparkpkg") + new_file("my_source:1.sparkpkg") + cwd = os.getcwd() + os.chdir(d) + try: + spark.addArtifact("my_relative_source.sparkpkg") + spark.addArtifact("my_source:1.sparkpkg") + finally: + os.chdir(cwd) + packages = spark._jsparkSession.artifactManager().getNativeDataSourcePackages()._1() + self.assertEqual( + sorted(packages.apply(i).getName() for i in range(packages.size())), + [ + "my_file_source.sparkpkg", + "my_other_source.sparkpkg", + "my_relative_source.sparkpkg", + "my_source.sparkpkg", + "my_source:1.sparkpkg", + "my_uri_source.sparkpkg", + ], + ) + def test_add_zipped_package(self): self.check_add_zipped_package(self.spark) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/Artifact.scala b/sql/api/src/main/scala/org/apache/spark/sql/Artifact.scala index 40c7e532c27d..e4fb52e92839 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/Artifact.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/Artifact.scala @@ -38,6 +38,10 @@ private[sql] object Artifact { val CLASS_PREFIX: Path = Paths.get("classes") val JAR_PREFIX: Path = Paths.get("jars") val CACHE_PREFIX: Path = Paths.get("cache") + val FILE_PREFIX: Path = Paths.get("files") + + /** The file extension of native data source packages, which are added as file artifacts. */ + val NATIVE_DATA_SOURCE_PACKAGE_EXTENSION = ".sparkpkg" def newArtifactFromExtension( fileName: String, @@ -48,13 +52,15 @@ private[sql] object Artifact { newJarArtifact(targetFilePath, storage) case cf if cf.endsWith(".class") => newClassArtifact(targetFilePath, storage) + case pkg if pkg.endsWith(NATIVE_DATA_SOURCE_PACKAGE_EXTENSION) => + newFileArtifact(targetFilePath, storage) case other => throw new UnsupportedOperationException(s"Unsupported file format: $other") } } def parseArtifacts(uri: URI): Seq[Artifact] = { - // Currently only local files with extensions .jar and .class are supported. + // Currently only local files with extensions .jar, .class and .sparkpkg are supported. uri.getScheme match { case "file" => val path = Paths.get(uri) @@ -80,6 +86,10 @@ private[sql] object Artifact { newArtifact(CLASS_PREFIX, ".class", targetFilePath, storage) } + def newFileArtifact(targetFilePath: Path, storage: LocalData): Artifact = { + newArtifact(FILE_PREFIX, "", targetFilePath, storage) + } + def newCacheArtifact(id: String, storage: LocalData): Artifact = { newArtifact(CACHE_PREFIX, "", Paths.get(id), storage) } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala index 8aa554a879a1..cce193fe50db 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -568,6 +568,9 @@ abstract class SparkSession extends Serializable with Closeable { * Ivy settings. Otherwise, Ivy URIs are resolved by the server, which does not search Maven or * Ivy repositories local to the client. * + * Since 4.4.0, it also supports native data source packages: local files with the extension + * .sparkpkg, which make the native data sources that they contain available to the session. + * * @since 4.0.0 */ @Experimental @@ -583,6 +586,9 @@ abstract class SparkSession extends Serializable with Closeable { * Ivy settings. Otherwise, Ivy URIs are resolved by the server, which does not search Maven or * Ivy repositories local to the client. * + * Since 4.4.0, it also supports native data source packages: local files with the extension + * .sparkpkg, which make the native data sources that they contain available to the session. + * * @since 4.0.0 */ @Experimental @@ -593,7 +599,9 @@ abstract class SparkSession extends Serializable with Closeable { * specified by `target` under the session's working directory of that particular file * extension. * - * Supported target file extensions are .jar and .class. + * Supported target file extensions are .jar and .class, and since 4.4.0, .sparkpkg for native + * data source packages, which make the native data sources that they contain available to the + * session. * * ==Example== * {{{ @@ -613,7 +621,9 @@ abstract class SparkSession extends Serializable with Closeable { * Add a single artifact to the session while preserving the directory structure specified by * `target` under the session's working directory of that particular file extension. * - * Supported target file extensions are .jar and .class. + * Supported target file extensions are .jar and .class, and since 4.4.0, .sparkpkg for native + * data source packages, which make the native data sources that they contain available to the + * session. * * ==Example== * {{{ @@ -639,6 +649,9 @@ abstract class SparkSession extends Serializable with Closeable { * Ivy settings. Otherwise, Ivy URIs are resolved by the server, which does not search Maven or * Ivy repositories local to the client. * + * Since 4.4.0, it also supports native data source packages: local files with the extension + * .sparkpkg, which make the native data sources that they contain available to the session. + * * @since 4.0.0 */ @Experimental diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala index c16a7d20dfac..7c82867a2c3f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala @@ -2020,6 +2020,12 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat cause = Some(cause)) } + def nativeDataSourcePackageConflictError(name: String, packages: Seq[String]): Throwable = { + new AnalysisException( + errorClass = "NATIVE_DATA_SOURCE_PACKAGE_CONFLICT", + messageParameters = Map("name" -> name, "packages" -> packages.mkString(", "))) + } + def findMultipleDataSourceError(provider: String, sourceNames: Seq[String]): Throwable = { new AnalysisException( errorClass = "MULTIPLE_DATA_SOURCES", diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala index fd44b835a0ac..c8c3e8cb0418 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala @@ -802,6 +802,66 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE cause = error) } + def nativeDataSourceError( + action: String, name: String, msg: String, cause: Throwable): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "NATIVE_DATA_SOURCE_ERROR", + messageParameters = Map("action" -> action, "name" -> name, "msg" -> msg), + cause = cause) + } + + def invalidNativeDataSourcePackageError( + path: String, + subClass: String, + messageParameters: Map[String, String], + cause: Throwable = null): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = s"INVALID_NATIVE_DATA_SOURCE_PACKAGE.$subClass", + messageParameters = Map("path" -> path) ++ messageParameters, + cause = cause) + } + + def nativeDataSourcePackageNotFoundError( + fileName: String, checksum: String): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND", + messageParameters = Map("fileName" -> fileName, "checksum" -> checksum)) + } + + def invalidNativeDataSourceLibraryError( + path: String, + subClass: String, + messageParameters: Map[String, String], + cause: Throwable = null): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = s"INVALID_NATIVE_DATA_SOURCE_LIBRARY.$subClass", + messageParameters = Map("path" -> path) ++ messageParameters, + cause = cause) + } + + def nativeDataSourceLibraryNotFoundError( + fileName: String, checksum: String, path: String): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND", + messageParameters = Map("fileName" -> fileName, "checksum" -> checksum, "path" -> path)) + } + + def nativeDataSourceUnsupportedJavaVersionError( + required: Int, current: Int): SparkUnsupportedOperationException = { + new SparkUnsupportedOperationException( + errorClass = "NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION", + messageParameters = Map("required" -> required.toString, "current" -> current.toString)) + } + + def nativeDataSourceUnsupportedPlatformError( + platform: String, cause: Throwable): SparkUnsupportedOperationException = { + val e = new SparkUnsupportedOperationException( + errorClass = "NATIVE_DATA_SOURCE_UNSUPPORTED_PLATFORM", + messageParameters = Map("platform" -> platform, "reason" -> cause.toString)) + e.initCause(cause) + e + } + def removedClassInSpark2Error(className: String, e: Throwable): SparkClassNotFoundException = { new SparkClassNotFoundException( errorClass = "_LEGACY_ERROR_TEMP_2052", diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index d3211aba4986..787b58857116 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -7207,6 +7207,19 @@ object SQLConf { .booleanConf .createWithDefault(false) + val NATIVE_DATA_SOURCE_PATHS = buildConf("spark.sql.dataSource.native.paths") + .doc("Comma-separated list of native data source packages (.sparkpkg files), and of " + + "directories that contain them, on the local file system of the driver. Spark adds a " + + "package to the artifacts of the session when one of its data sources is used, so that " + + "the executors get it too. Packages can also be added with spark.addArtifact. These " + + "packages take precedence over the native data source libraries installed in the native " + + "library path. See spark.sql.dataSource.native.enabled.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .stringConf + .toSequence + .createWithDefault(Nil) + val CSV_FILTER_PUSHDOWN_ENABLED = buildConf("spark.sql.csv.filterPushdown.enabled") .doc("When true, enable filter pushdown to CSV datasource.") .version("3.0.0") diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala index e02c7296f937..6da375360c4e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala @@ -464,6 +464,22 @@ object StaticSQLConf { .checkValue(_ >= 0, "The maximum total size must not be negative.") .createWithDefault(128 * 1024) // 128 KiB + val NATIVE_DATA_SOURCE_ENABLED = + buildStaticConf("spark.sql.dataSource.native.enabled") + .doc("When true, Spark finds and loads native data sources: data sources implemented in " + + "native code such as Rust or C++. The library of a data source, such as my_source, is " + + "found when it is installed in the native library path as spark_datasource_my_source, " + + "for example libspark_datasource_my_source.so in /usr/local/lib on Linux. The native " + + "library path consists of the directories of java.library.path, and the lib directory " + + "of each installation prefix in PATH. Libraries can also be distributed in native data " + + "source packages (.sparkpkg files), added with spark.addArtifact or under " + + "spark.sql.dataSource.native.paths. Loading a native data source runs its native code " + + "in the driver and executor processes.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(true) + val KAFKA_DISALLOWED_OPTIONS = buildStaticConf("spark.sql.kafka.disallowedOptions") .internal() diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala new file mode 100644 index 000000000000..97c4b3ae6104 --- /dev/null +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/NativeDataSourceConnectSuite.scala @@ -0,0 +1,105 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.connect + +import java.nio.file.Files +import java.util.UUID + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.SparkException +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.execution.datasources.v2.ffi.NativeDataSourceTestUtils._ +import org.apache.spark.util.Utils + +/** + * Tests native data sources with Spark Connect: a client adds packages to its session with + * `spark.addArtifact`, and the server finds the data sources in the packages of the session or in + * the installed libraries. The tests are skipped without a C++ compiler, or before Java 22. + */ +class NativeDataSourceConnectSuite extends SparkConnectServerTest { + + private lazy val library = compile("test_native_datasource") + private lazy val installedLibrary = + compile("test_native_datasource_installed", "-DNAME_PREFIX=\"installed_\"") + + private def nativeTest(name: String)(f: => Unit): Unit = test(name) { + assume(toolchain.isDefined, "A C++ compiler and Java 22 or later are required.") + f + } + + private def withNewSession(f: SparkSession => Unit): Unit = { + withSession(sessionId = UUID.randomUUID().toString)(f) + } + + nativeTest("add packages with spark.addArtifact") { + val pkg = createPackage(library, dataSources("")) + withNewSession { session => + session.addArtifact(pkg.getPath) + val df = session.read.format("native_range").option("end", "3").load() + assert(df.select("id").collect().map(_.getLong(0)).sorted.toSeq == Seq(0L, 1L, 2L)) + + val dir = Utils.createTempDir() + session + .range(2) + .selectExpr("id", "concat('v', id) AS name") + .write + .format("native_sink") + .option("path", dir.getPath) + .mode("append") + .save() + assert( + dir + .listFiles() + .filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala) + .toSet == Set("0,v0", "1,v1")) + + try { + session.sql("CREATE TABLE native_connect_table USING native_range OPTIONS (end 2)") + assert( + session + .sql("SELECT id FROM native_connect_table") + .collect() + .map(_.getLong(0)) + .sorted + .toSeq == Seq(0L, 1L)) + } finally { + session.sql("DROP TABLE IF EXISTS native_connect_table") + } + } + + // The packages of a session are not visible to the other sessions. + withNewSession { other => + val e = intercept[SparkException] { + other.read.format("native_range").load().collect() + } + assert(e.getCondition == "DATA_SOURCE_NOT_FOUND") + } + } + + nativeTest("find installed libraries automatically") { + val dir = Utils.createTempDir() + installLibrary(installedLibrary, "installed_native_range", dir) + withLibraryPath(dir) { + withNewSession { session => + val df = session.read.format("installed_native_range").option("end", "3").load() + assert(df.select("id").collect().map(_.getLong(0)).sorted.toSeq == Seq(0L, 1L, 2L)) + } + } + } +} diff --git a/sql/core/pom.xml b/sql/core/pom.xml index 281546b901a7..7f3db15f70fd 100644 --- a/sql/core/pom.xml +++ b/sql/core/pom.xml @@ -298,6 +298,10 @@ org.apache.arrow arrow-compression + + org.apache.arrow + arrow-c-data + org.apache.spark spark-udf-worker-proto_${scala.binary.version} diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java new file mode 100644 index 000000000000..66c1a46b7bee --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -0,0 +1,982 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.datasources.v2.ffi; + +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.MethodType; +import java.lang.reflect.Array; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.BooleanSupplier; + +import org.apache.spark.internal.LogKeys; +import org.apache.spark.internal.MDC; +import org.apache.spark.internal.SparkLogger; +import org.apache.spark.internal.SparkLoggerFactory; +import org.apache.spark.unsafe.Platform; + +/** + * A native data source library, which implements the C interface in + * {@code org/apache/spark/sql/datasource/spark_datasource.h}, called through the Foreign Function + * and Memory (FFM) API of Java 22 and later. + *

+ * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached + * through method handles, looked up once when the first library is loaded. The library exports + * {@code spark_ds_init}, which returns its table of functions, {@code SparkDsLibrary}. Each + * function of the table is bound to a method handle that takes and returns primitive values only, + * with pointers as {@code long} addresses, so the calls themselves use no reflection. The memory + * that Spark passes to the library is allocated with {@link Platform}, at the offsets of the + * structs of the header, which {@link #layout()} lists for the tests. + *

+ * Calling a function that the library leaves NULL in its table throws + * {@link MissingFunctionException}. A function that fails throws a {@link RuntimeException} with + * the message of {@code last_error}. + */ +final class NativeLibrary { + + private static final SparkLogger LOGGER = SparkLoggerFactory.getLogger(NativeLibrary.class); + + /** The version of the interface that Spark supports, {@code SPARK_DS_ABI_VERSION}. */ + static final int ABI_VERSION = 1; + + /** The first version of Java with the final FFM API. */ + static final int MINIMUM_JAVA_VERSION = 22; + + /** Whether this JVM can load native data source libraries. */ + static boolean isSupported() { + return Runtime.version().feature() >= MINIMUM_JAVA_VERSION; + } + + // The capabilities of a data source, SPARK_DS_BATCH_READ and the others. + static final long BATCH_READ = 1; + static final long MICRO_BATCH_READ = 2; + static final long BATCH_WRITE = 4; + static final long STREAMING_WRITE = 8; + static final long TRUNCATE = 16; + + // The modes of a write, SPARK_DS_WRITE_APPEND and SPARK_DS_WRITE_TRUNCATE. + static final int WRITE_APPEND = 0; + static final int WRITE_TRUNCATE = 1; + + /** A predicate that the reader does not use, SPARK_DS_PREDICATE_UNUSED. */ + static final int PREDICATE_UNUSED = 0; + + // The C types of the parameters and the results of the functions. + private enum C { VOID, BOOL, I32, I64, PTR } + + /** The functions of the table {@code SparkDsLibrary}, in their order in the table. */ + enum Function { + LAST_ERROR("last_error", C.PTR), + CREATE("create", C.I32, C.PTR, C.PTR, C.PTR), + SCHEMA("schema", C.I32, C.PTR, C.PTR), + CREATE_READER("create_reader", C.I32, C.PTR, C.PTR, C.PTR), + CREATE_STREAM_READER("create_stream_reader", C.I32, C.PTR, C.PTR, C.PTR), + CREATE_WRITER("create_writer", C.I32, C.PTR, C.PTR, C.PTR, C.PTR), + CLOSE("close", C.VOID, C.PTR), + PUSH_PREDICATES("push_predicates", C.I32, C.PTR, C.PTR, C.I64, C.PTR), + PUSH_LIMIT("push_limit", C.I32, C.PTR, C.I32, C.PTR), + PRUNE_COLUMNS("prune_columns", C.I32, C.PTR, C.PTR, C.I64, C.PTR), + PARTITIONS("partitions", C.I32, C.PTR, C.PTR, C.PTR), + SERIALIZE_READER("serialize_reader", C.I32, C.PTR, C.PTR, C.PTR), + CLOSE_READER("close_reader", C.VOID, C.PTR), + INITIAL_OFFSET("initial_offset", C.I32, C.PTR, C.PTR, C.PTR), + LATEST_OFFSET("latest_offset", C.I32, C.PTR, C.PTR, C.PTR), + STREAM_PARTITIONS("stream_partitions", C.I32, C.PTR, C.PTR, C.PTR, C.PTR, C.PTR), + SERIALIZE_STREAM_READER("serialize_stream_reader", C.I32, C.PTR, C.PTR, C.PTR), + COMMIT_OFFSET("commit_offset", C.I32, C.PTR, C.PTR), + CLOSE_STREAM_READER("close_stream_reader", C.VOID, C.PTR), + READ("read", C.I32, C.PTR, C.PTR, C.I64, C.PTR, C.I64, C.PTR), + SERIALIZE_WRITER("serialize_writer", C.I32, C.PTR, C.PTR, C.PTR), + COMMIT("commit", C.I32, C.PTR, C.I64, C.PTR, C.PTR, C.I64), + ABORT("abort", C.I32, C.PTR, C.I64, C.PTR, C.PTR, C.I64), + CLOSE_WRITER("close_writer", C.VOID, C.PTR), + CREATE_DATA_WRITER("create_data_writer", C.I32, C.PTR, C.PTR, C.I64, C.I64, C.PTR), + WRITE("write", C.I32, C.PTR, C.PTR, C.PTR), + COMMIT_DATA_WRITER("commit_data_writer", C.I32, C.PTR, C.PTR, C.PTR), + ABORT_DATA_WRITER("abort_data_writer", C.I32, C.PTR), + CLOSE_DATA_WRITER("close_data_writer", C.VOID, C.PTR); + + /** The name of the field of the function in the table. */ + final String fieldName; + private final C result; + private final C[] parameters; + + Function(String fieldName, C result, C... parameters) { + this.fieldName = fieldName; + this.result = result; + this.parameters = parameters; + } + + /** The offset of the function in the table, after the size and the version of the table. */ + long offset() { + return TABLE_FUNCTIONS + 8L * ordinal(); + } + } + + // The layouts of the structs of the header, on the 64-bit platforms that Spark supports. Each + // struct starts with its size. + private static final long TABLE_ABI_VERSION = 8; + private static final long TABLE_FUNCTIONS = 16; + private static final long HOST_SIZE = 32; + private static final long HOST_ABI_VERSION = 8; + private static final long HOST_SPARK_VERSION = 16; + private static final long HOST_LOG = 24; + private static final long CREATE_INFO_SIZE = 40; + private static final long CREATE_INFO_NAME = 8; + private static final long CREATE_INFO_OPTION_KEYS = 16; + private static final long CREATE_INFO_OPTION_VALUES = 24; + private static final long CREATE_INFO_NUM_OPTIONS = 32; + private static final long TASK_SIZE = 40; + private static final long TASK_STAGE_ID = 8; + private static final long TASK_STAGE_ATTEMPT_NUMBER = 12; + private static final long TASK_PARTITION_ID = 16; + private static final long TASK_ATTEMPT_NUMBER = 20; + private static final long TASK_ATTEMPT_ID = 24; + private static final long TASK_IS_INTERRUPTED = 32; + private static final long WRITE_INFO_SIZE = 24; + private static final long WRITE_INFO_MODE = 8; + private static final long WRITE_INFO_QUERY_ID = 16; + + /** + * The offsets of the fields of the structs that Spark reads or writes, by {@code struct.field}, + * and the sizes of the structs, by {@code struct}, for the tests to compare with the header. + */ + static Map layout() { + Map layout = new LinkedHashMap<>(); + layout.put("SparkDsLibrary", TABLE_FUNCTIONS + 8L * Function.values().length); + layout.put("SparkDsLibrary.size", 0L); + layout.put("SparkDsLibrary.abi_version", TABLE_ABI_VERSION); + for (Function function : Function.values()) { + layout.put("SparkDsLibrary." + function.fieldName, function.offset()); + } + layout.put("SparkDsHost", HOST_SIZE); + layout.put("SparkDsHost.size", 0L); + layout.put("SparkDsHost.abi_version", HOST_ABI_VERSION); + layout.put("SparkDsHost.spark_version", HOST_SPARK_VERSION); + layout.put("SparkDsHost.log", HOST_LOG); + layout.put("SparkDsCreateInfo", CREATE_INFO_SIZE); + layout.put("SparkDsCreateInfo.size", 0L); + layout.put("SparkDsCreateInfo.name", CREATE_INFO_NAME); + layout.put("SparkDsCreateInfo.option_keys", CREATE_INFO_OPTION_KEYS); + layout.put("SparkDsCreateInfo.option_values", CREATE_INFO_OPTION_VALUES); + layout.put("SparkDsCreateInfo.num_options", CREATE_INFO_NUM_OPTIONS); + layout.put("SparkDsTaskContext", TASK_SIZE); + layout.put("SparkDsTaskContext.size", 0L); + layout.put("SparkDsTaskContext.stage_id", TASK_STAGE_ID); + layout.put("SparkDsTaskContext.stage_attempt_number", TASK_STAGE_ATTEMPT_NUMBER); + layout.put("SparkDsTaskContext.partition_id", TASK_PARTITION_ID); + layout.put("SparkDsTaskContext.attempt_number", TASK_ATTEMPT_NUMBER); + layout.put("SparkDsTaskContext.task_attempt_id", TASK_ATTEMPT_ID); + layout.put("SparkDsTaskContext.is_interrupted", TASK_IS_INTERRUPTED); + layout.put("SparkDsWriteInfo", WRITE_INFO_SIZE); + layout.put("SparkDsWriteInfo.size", 0L); + layout.put("SparkDsWriteInfo.mode", WRITE_INFO_MODE); + layout.put("SparkDsWriteInfo.query_id", WRITE_INFO_QUERY_ID); + return layout; + } + + private final int abiVersion; + // The functions of the table, by ordinal, null for the functions that the library leaves NULL. + private final MethodHandle[] functions = new MethodHandle[Function.values().length]; + + private NativeLibrary(long table) throws Throwable { + long size = Platform.getLong(null, table); + if (size < TABLE_FUNCTIONS) { + throw new LoadException( + "spark_ds_init returned a function table with a size of " + size + " bytes.", null, true); + } + abiVersion = Platform.getInt(null, table + TABLE_ABI_VERSION); + // The table of another version of the interface has another layout. Only the functions that + // fit in the size of the table are read: the table of a library built against an earlier + // header ends before the functions that were added later. + if (abiVersion == ABI_VERSION) { + for (Function function : Function.values()) { + if (function.offset() + 8 <= size) { + long address = Platform.getLong(null, table + function.offset()); + if (address != 0) { + functions[function.ordinal()] = Ffm.downcall( + Ffm.OF_ADDRESS.invoke(address), function.result, function.parameters); + } + } + } + } + } + + /** + * Loads the library at the given path, and gets its table with {@code spark_ds_init}. The + * library stays loaded until the JVM exits. + * + * @param sparkVersion the version of Spark, which the library gets in {@code SparkDsHost} + * @throws LoadException if the library cannot be loaded + */ + static NativeLibrary load(String path, String sparkVersion) throws Throwable { + if (!isSupported()) { + throw new UnsupportedOperationException( + "Native data sources require Java " + MINIMUM_JAVA_VERSION + " or later."); + } + Object lookup; + try { + lookup = Ffm.LIBRARY_LOOKUP.invoke(Path.of(path), Ffm.GLOBAL_ARENA); + } catch (IllegalArgumentException | IllegalCallerException e) { + // The library is not a valid library, or native access is not enabled for Spark. + throw new LoadException(e.getMessage(), e); + } catch (LinkageError e) { + // The FFM API could not be initialized, see Ffm: report why, such as an + // IllegalCallerException when native access is not enabled for Spark. + throw new LoadException( + "Cannot use the Foreign Function and Memory API: " + rootCause(e), e); + } + Optional init = (Optional) Ffm.FIND.invoke(lookup, "spark_ds_init"); + if (init.isEmpty()) { + throw new LoadException( + "The library does not export the function spark_ds_init.", null, true); + } + MethodHandle initFunction = Ffm.downcall(init.get(), C.I32, C.PTR, C.PTR); + // The library may keep the host until the JVM exits, even if spark_ds_init fails. + long host = Ffm.host(path, sparkVersion); + long out = Platform.allocateMemory(8); + try { + Platform.putLong(null, out, 0); + int status = (int) initFunction.invokeExact(host, out); + if (status != 0) { + throw new LoadException( + "spark_ds_init failed with the error code " + status + ".", null, true); + } + long table = Platform.getLong(null, out); + if (table == 0) { + throw new LoadException("spark_ds_init returned no function table.", null, true); + } + return new NativeLibrary(table); + } finally { + Platform.freeMemory(out); + } + } + + /** The method handles of the FFM API, initialized when the first library is loaded. */ + private static final class Ffm { + static final Class LAYOUT_CLASS; + static final Object LINKER; + static final Object GLOBAL_ARENA; + static final Object NO_OPTIONS; + static final MethodHandle LIBRARY_LOOKUP; + static final MethodHandle FIND; + static final MethodHandle DOWNCALL; + static final MethodHandle UPCALL; + static final MethodHandle FUNCTION_OF; + static final MethodHandle FUNCTION_OF_VOID; + // (long) -> MemorySegment and (MemorySegment) -> long. + static final MethodHandle OF_ADDRESS; + static final MethodHandle ADDRESS_OF; + static final Object BOOL; + static final Object INT32; + static final Object INT64; + static final Object POINTER; + // NativeLibrary.log, which SparkDsHost.log calls with the path of its library. + static final MethodHandle LOG; + // The addresses of the C functions that call Java, which live as long as the JVM: + // SparkDsSink, which calls NativeLibrary.sink, and SparkDsTaskContext.is_interrupted, which + // calls NativeLibrary.isInterrupted. + static final long SINK; + static final long IS_INTERRUPTED; + + static { + try { + // A full-privilege lookup, so that the restricted methods of the FFM API see this class + // as their caller, in the unnamed module that Spark enables native access for. + MethodHandles.Lookup lookup = MethodHandles.lookup(); + Class linker = Class.forName("java.lang.foreign.Linker"); + Class option = Class.forName("java.lang.foreign.Linker$Option"); + Class symbolLookup = Class.forName("java.lang.foreign.SymbolLookup"); + Class arena = Class.forName("java.lang.foreign.Arena"); + Class segment = Class.forName("java.lang.foreign.MemorySegment"); + Class layout = Class.forName("java.lang.foreign.MemoryLayout"); + Class layoutArray = Array.newInstance(layout, 0).getClass(); + Class valueLayout = Class.forName("java.lang.foreign.ValueLayout"); + Class descriptor = Class.forName("java.lang.foreign.FunctionDescriptor"); + Class options = Array.newInstance(option, 0).getClass(); + + LAYOUT_CLASS = layout; + LINKER = lookup.findStatic(linker, "nativeLinker", MethodType.methodType(linker)) + .invoke(); + GLOBAL_ARENA = lookup.findStatic(arena, "global", MethodType.methodType(arena)).invoke(); + NO_OPTIONS = Array.newInstance(option, 0); + LIBRARY_LOOKUP = lookup.findStatic(symbolLookup, "libraryLookup", + MethodType.methodType(symbolLookup, Path.class, arena)); + FIND = lookup.findVirtual(symbolLookup, "find", + MethodType.methodType(Optional.class, String.class)); + DOWNCALL = lookup.findVirtual(linker, "downcallHandle", + MethodType.methodType(MethodHandle.class, segment, descriptor, options)) + .asFixedArity(); + UPCALL = lookup.findVirtual(linker, "upcallStub", + MethodType.methodType(segment, MethodHandle.class, descriptor, arena, options)) + .asFixedArity(); + FUNCTION_OF = lookup.findStatic(descriptor, "of", + MethodType.methodType(descriptor, layout, layoutArray)).asFixedArity(); + FUNCTION_OF_VOID = lookup.findStatic(descriptor, "ofVoid", + MethodType.methodType(descriptor, layoutArray)).asFixedArity(); + OF_ADDRESS = lookup.findStatic(segment, "ofAddress", + MethodType.methodType(segment, long.class)); + ADDRESS_OF = lookup.findVirtual(segment, "address", MethodType.methodType(long.class)); + BOOL = valueLayout.getField("JAVA_BOOLEAN").get(null); + INT32 = valueLayout.getField("JAVA_INT").get(null); + INT64 = valueLayout.getField("JAVA_LONG").get(null); + POINTER = valueLayout.getField("ADDRESS").get(null); + + LOG = lookup.findStatic(NativeLibrary.class, "log", + MethodType.methodType(void.class, String.class, int.class, long.class)); + SINK = upcall( + lookup.findStatic(NativeLibrary.class, "sink", + MethodType.methodType(void.class, long.class, long.class, long.class)), + C.VOID, C.PTR, C.PTR, C.I64); + IS_INTERRUPTED = upcall( + lookup.findStatic(NativeLibrary.class, "isInterrupted", + MethodType.methodType(boolean.class, long.class)), + C.BOOL, C.PTR); + } catch (Throwable e) { + // The message names the cause, because the JVM only keeps this message for the later + // uses of the class. + throw new IllegalStateException("Cannot initialize the FFM API: " + rootCause(e), e); + } + } + + /** Returns the FunctionDescriptor of a C function with the given types. */ + static Object descriptor(C result, C... parameters) throws Throwable { + Object layouts = Array.newInstance(LAYOUT_CLASS, parameters.length); + for (int i = 0; i < parameters.length; i++) { + Array.set(layouts, i, layout(parameters[i])); + } + return result == C.VOID + ? FUNCTION_OF_VOID.invoke(layouts) + : FUNCTION_OF.invoke(layout(result), layouts); + } + + static Object layout(C type) { + return switch (type) { + case BOOL -> BOOL; + case I32 -> INT32; + case I64 -> INT64; + case PTR -> POINTER; + case VOID -> throw new IllegalArgumentException("void has no layout"); + }; + } + + /** + * Binds the C function at the given address, a MemorySegment, to a method handle whose + * parameters and result have the Java types of the given C types, with {@code long} for + * pointers. + */ + static MethodHandle downcall(Object address, C result, C... parameters) throws Throwable { + MethodHandle handle = (MethodHandle) DOWNCALL.invoke( + LINKER, address, descriptor(result, parameters), NO_OPTIONS); + // Passes and returns the pointers, which are MemorySegments in the FFM API, as addresses. + for (int i = 0; i < parameters.length; i++) { + if (parameters[i] == C.PTR) { + handle = MethodHandles.filterArguments(handle, i, OF_ADDRESS); + } + } + if (result == C.PTR) { + handle = MethodHandles.filterReturnValue(handle, ADDRESS_OF); + } + return handle; + } + + /** + * Returns the address of a C function with the given types that calls `target`, whose + * parameters have the Java types of the C types, with {@code long} for pointers. The + * function lives as long as the JVM. + */ + static long upcall(MethodHandle target, C result, C... parameters) throws Throwable { + for (int i = 0; i < parameters.length; i++) { + if (parameters[i] == C.PTR) { + target = MethodHandles.filterArguments(target, i, ADDRESS_OF); + } + } + Object stub = UPCALL.invoke( + LINKER, target, descriptor(result, parameters), GLOBAL_ARENA, NO_OPTIONS); + return (long) ADDRESS_OF.invoke(stub); + } + + /** + * Allocates the SparkDsHost of a library. The library may keep it until the JVM exits, so it + * is never freed. + */ + static long host(String path, String sparkVersion) throws Throwable { + long log = upcall(MethodHandles.insertArguments(LOG, 0, path), C.VOID, C.I32, C.PTR); + byte[] version = sparkVersion.getBytes(StandardCharsets.UTF_8); + long versionAddress = Platform.allocateMemory(version.length + 1L); + Platform.copyMemory( + version, Platform.BYTE_ARRAY_OFFSET, null, versionAddress, version.length); + Platform.putByte(null, versionAddress + version.length, (byte) 0); + long host = Platform.allocateMemory(HOST_SIZE); + Platform.setMemory(host, (byte) 0, HOST_SIZE); + Platform.putLong(null, host, HOST_SIZE); + Platform.putInt(null, host + HOST_ABI_VERSION, ABI_VERSION); + Platform.putLong(null, host + HOST_SPARK_VERSION, versionAddress); + Platform.putLong(null, host + HOST_LOG, log); + return host; + } + } + + private static Throwable rootCause(Throwable e) { + Throwable cause = e; + while (cause.getCause() != null && cause.getCause() != cause) { + cause = cause.getCause(); + } + return cause; + } + + /** Thrown when a library cannot be loaded, or cannot be used. */ + static final class LoadException extends RuntimeException { + /** + * Whether the library was loaded, in which case it stays loaded: Spark found that it cannot + * use it afterwards, for example because {@code spark_ds_init} failed. + */ + final boolean loaded; + + LoadException(String message, Throwable cause) { + this(message, cause, false); + } + + LoadException(String message, Throwable cause, boolean loaded) { + super(message, cause); + this.loaded = loaded; + } + } + + /** Thrown when a function is called that the library leaves NULL in its table, with its name. */ + static final class MissingFunctionException extends RuntimeException { + MissingFunctionException(String name) { + super(name); + } + } + + /** + * Called by the library, through {@code SparkDsHost.log}, with the path of the library. It must + * not throw: an exception thrown by an upcall crashes the JVM. + */ + private static void log(String path, int level, long message) { + try { + String text = message == 0 ? "" : new String(readCString(message), StandardCharsets.UTF_8); + switch (level) { + case 1 -> LOGGER.error("Native data source library {}: {}", + MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text)); + case 2 -> LOGGER.warn("Native data source library {}: {}", + MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text)); + case 4 -> LOGGER.debug("Native data source library {}: {}", path, text); + case 5 -> LOGGER.trace("Native data source library {}: {}", path, text); + default -> LOGGER.info("Native data source library {}: {}", + MDC.of(LogKeys.PATH, path), MDC.of(LogKeys.MESSAGE, text)); + } + } catch (Throwable e) { + // The message is lost. + } + } + + // --------------------------------------------------------------------------------------------- + // Calls + // --------------------------------------------------------------------------------------------- + + private MethodHandle function(Function function) { + MethodHandle handle = functions[function.ordinal()]; + if (handle == null) { + throw new MissingFunctionException(function.fieldName); + } + return handle; + } + + /** Throws the error that the library reported if a function failed. */ + private void check(int status, Function function) throws Throwable { + if (status == 0) { + return; + } + MethodHandle lastError = functions[Function.LAST_ERROR.ordinal()]; + long address = lastError == null ? 0 : (long) lastError.invokeExact(); + String message = address == 0 ? "" : new String(readCString(address), StandardCharsets.UTF_8); + throw new RuntimeException(message.isEmpty() + ? function.fieldName + " failed with the error code " + status + "." + : message); + } + + private static long readHandle(long out, Function function) { + long handle = Platform.getLong(null, out); + if (handle == 0) { + throw new IllegalStateException(function.fieldName + " returned a NULL handle."); + } + return handle; + } + + /** The version of the interface that the table of the library implements. */ + int abiVersion() { + return abiVersion; + } + + /** Whether the library implements the function, which it does not leave NULL in its table. */ + boolean hasFunction(Function function) { + return functions[function.ordinal()] != null; + } + + /** A data source handle, with the capabilities of the data source. */ + static final class DataSourceHandle { + final long handle; + final long capabilities; + + DataSourceHandle(long handle, long capabilities) { + this.handle = handle; + this.capabilities = capabilities; + } + } + + DataSourceHandle createDataSource(String name, String[] optionKeys, String[] optionValues) + throws Throwable { + MethodHandle f = function(Function.CREATE); + try (Memory memory = new Memory()) { + long info = memory.zeroed(CREATE_INFO_SIZE); + Platform.putLong(null, info, CREATE_INFO_SIZE); + Platform.putLong(null, info + CREATE_INFO_NAME, memory.string(name)); + Platform.putLong(null, info + CREATE_INFO_OPTION_KEYS, memory.strings(optionKeys)); + Platform.putLong(null, info + CREATE_INFO_OPTION_VALUES, memory.strings(optionValues)); + Platform.putLong(null, info + CREATE_INFO_NUM_OPTIONS, optionKeys.length); + long capabilities = memory.zeroed(8); + long out = memory.zeroed(8); + check((int) f.invokeExact(info, capabilities, out), Function.CREATE); + return new DataSourceHandle( + readHandle(out, Function.CREATE), Platform.getLong(null, capabilities)); + } + } + + void schema(long dataSource, long schemaAddress) throws Throwable { + check((int) function(Function.SCHEMA).invokeExact(dataSource, schemaAddress), + Function.SCHEMA); + } + + long createReader(long dataSource, long schemaAddress) throws Throwable { + return createWithSchema(Function.CREATE_READER, dataSource, schemaAddress); + } + + long createStreamReader(long dataSource, long schemaAddress) throws Throwable { + return createWithSchema(Function.CREATE_STREAM_READER, dataSource, schemaAddress); + } + + private long createWithSchema(Function function, long dataSource, long schemaAddress) + throws Throwable { + MethodHandle f = function(function); + try (Memory memory = new Memory()) { + long out = memory.zeroed(8); + check((int) f.invokeExact(dataSource, schemaAddress, out), function); + return readHandle(out, function); + } + } + + /** + * Creates a writer. + * + * @param mode {@link #WRITE_APPEND} or {@link #WRITE_TRUNCATE} + * @param queryId the ID of the streaming query of a streaming write, or null for a batch write + */ + long createWriter(long dataSource, long schemaAddress, int mode, String queryId) + throws Throwable { + MethodHandle f = function(Function.CREATE_WRITER); + try (Memory memory = new Memory()) { + long info = memory.zeroed(WRITE_INFO_SIZE); + Platform.putLong(null, info, WRITE_INFO_SIZE); + Platform.putInt(null, info + WRITE_INFO_MODE, mode); + Platform.putLong( + null, info + WRITE_INFO_QUERY_ID, queryId == null ? 0 : memory.string(queryId)); + long out = memory.zeroed(8); + check((int) f.invokeExact(dataSource, schemaAddress, info, out), Function.CREATE_WRITER); + return readHandle(out, Function.CREATE_WRITER); + } + } + + void closeDataSource(long dataSource) throws Throwable { + function(Function.CLOSE).invokeExact(dataSource); + } + + /** Returns how the reader uses each predicate, as SPARK_DS_PREDICATE_* values. */ + int[] pushPredicates(long reader, String[] predicates) throws Throwable { + MethodHandle f = function(Function.PUSH_PREDICATES); + try (Memory memory = new Memory()) { + long uses = memory.zeroed(4L * Math.max(1, predicates.length)); + check((int) f.invokeExact(reader, memory.strings(predicates), (long) predicates.length, + uses), Function.PUSH_PREDICATES); + int[] result = new int[predicates.length]; + for (int i = 0; i < result.length; i++) { + result[i] = Platform.getInt(null, uses + 4L * i); + } + return result; + } + } + + boolean pushLimit(long reader, int limit) throws Throwable { + MethodHandle f = function(Function.PUSH_LIMIT); + try (Memory memory = new Memory()) { + long applied = memory.zeroed(1); + check((int) f.invokeExact(reader, limit, applied), Function.PUSH_LIMIT); + return Platform.getByte(null, applied) != 0; + } + } + + boolean pruneColumns(long reader, String[] columnNames) throws Throwable { + MethodHandle f = function(Function.PRUNE_COLUMNS); + try (Memory memory = new Memory()) { + long pruned = memory.zeroed(1); + check((int) f.invokeExact(reader, memory.strings(columnNames), (long) columnNames.length, + pruned), Function.PRUNE_COLUMNS); + return Platform.getByte(null, pruned) != 0; + } + } + + byte[][] partitions(long reader) throws Throwable { + MethodHandle f = function(Function.PARTITIONS); + try (Sink sink = new Sink()) { + check((int) f.invokeExact(reader, Ffm.SINK, sink.id), Function.PARTITIONS); + return sink.all(); + } + } + + byte[] serializeReader(long reader) throws Throwable { + return serialize(Function.SERIALIZE_READER, reader); + } + + void closeReader(long reader) throws Throwable { + function(Function.CLOSE_READER).invokeExact(reader); + } + + String initialOffset(long streamReader) throws Throwable { + return offset(Function.INITIAL_OFFSET, streamReader); + } + + String latestOffset(long streamReader) throws Throwable { + return offset(Function.LATEST_OFFSET, streamReader); + } + + private String offset(Function function, long streamReader) throws Throwable { + byte[] offset = serialize(function, streamReader); + if (offset == null) { + throw new IllegalStateException(function.fieldName + " returned no offset."); + } + return new String(offset, StandardCharsets.UTF_8); + } + + byte[][] streamPartitions(long streamReader, String start, String end) throws Throwable { + MethodHandle f = function(Function.STREAM_PARTITIONS); + try (Memory memory = new Memory(); Sink sink = new Sink()) { + check((int) f.invokeExact(streamReader, memory.string(start), memory.string(end), + Ffm.SINK, sink.id), Function.STREAM_PARTITIONS); + return sink.all(); + } + } + + byte[] serializeStreamReader(long streamReader) throws Throwable { + return serialize(Function.SERIALIZE_STREAM_READER, streamReader); + } + + void commitOffset(long streamReader, String end) throws Throwable { + MethodHandle f = function(Function.COMMIT_OFFSET); + try (Memory memory = new Memory()) { + check((int) f.invokeExact(streamReader, memory.string(end)), Function.COMMIT_OFFSET); + } + } + + void closeStreamReader(long streamReader) throws Throwable { + function(Function.CLOSE_STREAM_READER).invokeExact(streamReader); + } + + void read(Task task, byte[] readerState, byte[] partition, long streamAddress) + throws Throwable { + MethodHandle f = function(Function.READ); + try (Memory memory = new Memory()) { + check((int) f.invokeExact(task.address, memory.bytes(readerState), + (long) readerState.length, memory.bytes(partition), (long) partition.length, + streamAddress), Function.READ); + } + } + + byte[] serializeWriter(long writer) throws Throwable { + return serialize(Function.SERIALIZE_WRITER, writer); + } + + void commit(long writer, long epochId, byte[][] messages) throws Throwable { + finishWrite(Function.COMMIT, writer, epochId, messages); + } + + void abort(long writer, long epochId, byte[][] messages) throws Throwable { + finishWrite(Function.ABORT, writer, epochId, messages); + } + + private void finishWrite(Function function, long writer, long epochId, byte[][] messages) + throws Throwable { + MethodHandle f = function(function); + try (Memory memory = new Memory()) { + long pointers = memory.zeroed(8L * Math.max(1, messages.length)); + long lengths = memory.zeroed(8L * Math.max(1, messages.length)); + for (int i = 0; i < messages.length; i++) { + if (messages[i] != null) { + Platform.putLong(null, pointers + 8L * i, memory.bytes(messages[i])); + Platform.putLong(null, lengths + 8L * i, messages[i].length); + } + } + check((int) f.invokeExact(writer, epochId, pointers, lengths, (long) messages.length), + function); + } + } + + void closeWriter(long writer) throws Throwable { + function(Function.CLOSE_WRITER).invokeExact(writer); + } + + long createDataWriter(Task task, byte[] writerState, long epochId) throws Throwable { + MethodHandle f = function(Function.CREATE_DATA_WRITER); + try (Memory memory = new Memory()) { + long out = memory.zeroed(8); + check((int) f.invokeExact(task.address, memory.bytes(writerState), + (long) writerState.length, epochId, out), Function.CREATE_DATA_WRITER); + return readHandle(out, Function.CREATE_DATA_WRITER); + } + } + + void write(long dataWriter, long arrayAddress, long schemaAddress) throws Throwable { + check((int) function(Function.WRITE).invokeExact(dataWriter, arrayAddress, schemaAddress), + Function.WRITE); + } + + /** + * Commits a data writer. Then {@code committed} runs right away, before the message is checked: + * if the message is invalid, the commit fails, but the data writer must not be aborted. + */ + byte[] commitDataWriter(long dataWriter, Runnable committed) throws Throwable { + MethodHandle f = function(Function.COMMIT_DATA_WRITER); + try (Sink sink = new Sink()) { + check((int) f.invokeExact(dataWriter, Ffm.SINK, sink.id), Function.COMMIT_DATA_WRITER); + committed.run(); + byte[] message = single(sink, Function.COMMIT_DATA_WRITER); + return message == null ? new byte[0] : message; + } + } + + void abortDataWriter(long dataWriter) throws Throwable { + check((int) function(Function.ABORT_DATA_WRITER).invokeExact(dataWriter), + Function.ABORT_DATA_WRITER); + } + + void closeDataWriter(long dataWriter) throws Throwable { + function(Function.CLOSE_DATA_WRITER).invokeExact(dataWriter); + } + + /** + * Calls a function that returns a single byte array through a sink. Returns null if the + * function does not call the sink. + */ + private byte[] serialize(Function function, long handle) throws Throwable { + MethodHandle f = function(function); + try (Sink sink = new Sink()) { + check((int) f.invokeExact(handle, Ffm.SINK, sink.id), function); + return single(sink, function); + } + } + + /** Returns the single byte array that a function passed to the sink, or null if none. */ + private static byte[] single(Sink sink, Function function) throws Throwable { + byte[][] values = sink.all(); + if (values.length > 1) { + throw new IllegalStateException( + function.fieldName + " returned " + values.length + " values instead of one."); + } + return values.length == 0 ? null : values[0]; + } + + // --------------------------------------------------------------------------------------------- + // Memory + // --------------------------------------------------------------------------------------------- + + /** Native memory for the arguments of a call, freed after the call. */ + private static final class Memory implements AutoCloseable { + private final List allocations = new ArrayList<>(); + + long zeroed(long size) { + long address = Platform.allocateMemory(size); + allocations.add(address); + Platform.setMemory(address, (byte) 0, size); + return address; + } + + /** Copies the bytes. Empty bytes are passed as a valid pointer too. */ + long bytes(byte[] bytes) { + long address = zeroed(Math.max(1, bytes.length)); + Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address, bytes.length); + return address; + } + + /** Copies the string as a NUL-terminated UTF-8 string. */ + long string(String value) { + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + long address = zeroed(bytes.length + 1L); + Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address, bytes.length); + return address; + } + + /** Copies the strings, and returns the address of an array of pointers to them. */ + long strings(String[] values) { + long array = zeroed(8L * Math.max(1, values.length)); + for (int i = 0; i < values.length; i++) { + Platform.putLong(null, array + 8L * i, string(values[i])); + } + return array; + } + + @Override + public void close() { + allocations.forEach(Platform::freeMemory); + } + } + + private static byte[] readCString(long address) { + int length = 0; + while (Platform.getByte(null, address + length) != 0) { + length++; + } + byte[] bytes = new byte[length]; + Platform.copyMemory(null, address, bytes, Platform.BYTE_ARRAY_OFFSET, length); + return bytes; + } + + // --------------------------------------------------------------------------------------------- + // Task contexts + // --------------------------------------------------------------------------------------------- + + // Whether the task of each SparkDsTaskContext was killed, by the address of the struct. + private static final ConcurrentHashMap tasks = new ConcurrentHashMap<>(); + + /** + * The {@code SparkDsTaskContext} of a task that reads a partition or writes data. The reader or + * the data writer closes it once the library no longer uses it, which frees it. + */ + static final class Task implements AutoCloseable { + final long address; + private boolean closed; + + Task( + int stageId, + int stageAttemptNumber, + int partitionId, + int attemptNumber, + long taskAttemptId, + BooleanSupplier interrupted) { + address = Platform.allocateMemory(TASK_SIZE); + Platform.setMemory(address, (byte) 0, TASK_SIZE); + Platform.putLong(null, address, TASK_SIZE); + Platform.putInt(null, address + TASK_STAGE_ID, stageId); + Platform.putInt(null, address + TASK_STAGE_ATTEMPT_NUMBER, stageAttemptNumber); + Platform.putInt(null, address + TASK_PARTITION_ID, partitionId); + Platform.putInt(null, address + TASK_ATTEMPT_NUMBER, attemptNumber); + Platform.putLong(null, address + TASK_ATTEMPT_ID, taskAttemptId); + Platform.putLong(null, address + TASK_IS_INTERRUPTED, Ffm.IS_INTERRUPTED); + tasks.put(address, interrupted); + } + + @Override + public void close() { + if (!closed) { + closed = true; + // A call to is_interrupted that is still running only looks the address up. + tasks.remove(address); + Platform.freeMemory(address); + } + } + } + + /** + * Called by the library, through {@code SparkDsTaskContext.is_interrupted}, with the address of + * the struct. It must not throw: an exception thrown by an upcall crashes the JVM. + */ + private static boolean isInterrupted(long task) { + try { + BooleanSupplier interrupted = tasks.get(task); + // A freed task context belongs to a read or a write that is over. + return interrupted == null || interrupted.getAsBoolean(); + } catch (Throwable e) { + return false; + } + } + + // --------------------------------------------------------------------------------------------- + // Sink + // --------------------------------------------------------------------------------------------- + + private static final AtomicLong nextSinkId = new AtomicLong(1); + private static final ConcurrentHashMap sinks = new ConcurrentHashMap<>(); + + /** + * Collects the byte arrays that a library returns during a call, through the C function + * {@link Ffm#SINK}. The library passes {@link #id} back to identify it. + */ + private static final class Sink implements AutoCloseable { + final long id = nextSinkId.getAndIncrement(); + final List values = new ArrayList<>(); + Throwable error; + + Sink() { + sinks.put(id, this); + } + + byte[][] all() throws Throwable { + if (error != null) { + throw error; + } + return values.toArray(new byte[0][]); + } + + @Override + public void close() { + sinks.remove(id); + } + } + + /** + * Called by the library, through the upcall stub {@link Ffm#SINK}, with a byte array to copy. + * It must not throw: an exception thrown by an upcall crashes the JVM. + */ + private static void sink(long id, long data, long length) { + Sink sink = null; + try { + sink = sinks.get(id); + if (sink == null || sink.error != null) { + return; + } + if (length < 0 || length > Integer.MAX_VALUE - 8 || (data == 0 && length > 0)) { + throw new IllegalStateException("The library returned an invalid byte array, of " + + length + " bytes at the address " + data + "."); + } + byte[] bytes = new byte[(int) length]; + Platform.copyMemory(null, data, bytes, Platform.BYTE_ARRAY_OFFSET, length); + sink.values.add(bytes); + } catch (Throwable e) { + if (sink != null) { + sink.error = e; + } + } + } +} diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json b/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json new file mode 100644 index 000000000000..a37e9d86bec3 --- /dev/null +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json @@ -0,0 +1,49 @@ +{ + "description": "Predicates as Spark passes them to push_predicates of spark_datasource.h, with the ids of the rows of the table for which they are true. A case lists the predicates that Spark pushes for the filter in its SQL, which are combined with AND. A row has a value for each column of the schema; NaN and infinite doubles are strings, floats are written as the doubles of the same values, and binary values are base64 strings.", + "schema": "id INT,i INT,d DOUBLE,s STRING,f FLOAT,b BINARY", + "rows": [ + [0, 1, 0.0, "a", 0.10000000149011612, "AA=="], + [1, 2, -0.0, "", 1.5, "fw=="], + [2, null, "NaN", "\u00E9", null, "gA=="], + [3, 3, "Infinity", "\uD83D\uDE00", 0.10000000149011612, "/w=="], + [4, 1, "-Infinity", null, 0.5, "fwA="], + [5, null, 1.5, "ab", 2.0, ""], + [6, 2, null, "b\u00E9", 0.25, null], + [7, 4, 2.5, "\uFFFD", 3.0, "AQI="] + ], + "cases": [ + {"sql": "d = 0.0D", "predicates": [{"type":"function","name":"=","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":0.0}]}], "rows": [0, 1]}, + {"sql": "d = -0.0D", "predicates": [{"type":"function","name":"=","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":-0.0}]}], "rows": [0, 1]}, + {"sql": "d < 0.0D", "predicates": [{"type":"function","name":"<","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":0.0}]}], "rows": [4]}, + {"sql": "d > 1.0D", "predicates": [{"type":"function","name":">","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":1.0}]}], "rows": [2, 3, 5, 7]}, + {"sql": "d >= 2.5D", "predicates": [{"type":"function","name":">=","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":2.5}]}], "rows": [2, 3, 7]}, + {"sql": "d <> 0.0D", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":"=","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":0.0}]}]}], "rows": [2, 3, 4, 5, 7]}, + {"sql": "d <=> 0.0D", "predicates": [{"type":"function","name":"<=>","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":0.0}]}], "rows": [0, 1]}, + {"sql": "d IN (0.0D, 1.5D)", "predicates": [{"type":"function","name":"IN","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":0.0},{"type":"literal","dataType":"double","value":1.5}]}], "rows": [0, 1, 5]}, + {"sql": "NOT (d > 1.0D)", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":">","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":1.0}]}]}], "rows": [0, 1, 4]}, + {"sql": "NOT (d < 1.0D)", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":"<","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":1.0}]}]}], "rows": [2, 3, 5, 7]}, + {"sql": "d IS NULL", "predicates": [{"type":"function","name":"IS_NULL","children":[{"type":"column","name":["d"]}]}], "rows": [6]}, + {"sql": "d IS NOT NULL", "predicates": [{"type":"function","name":"IS_NOT_NULL","children":[{"type":"column","name":["d"]}]}], "rows": [0, 1, 2, 3, 4, 5, 7]}, + {"sql": "i = 1", "predicates": [{"type":"function","name":"=","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":1}]}], "rows": [0, 4]}, + {"sql": "i = NULL", "predicates": [{"type":"function","name":"=","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":null}]}], "rows": []}, + {"sql": "i <=> NULL", "predicates": [{"type":"function","name":"<=>","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":null}]}], "rows": [2, 5]}, + {"sql": "NOT (i <=> NULL)", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":"<=>","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":null}]}]}], "rows": [0, 1, 3, 4, 6, 7]}, + {"sql": "i IN (1, NULL)", "predicates": [{"type":"function","name":"IN","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":1},{"type":"literal","dataType":"int","value":null}]}], "rows": [0, 4]}, + {"sql": "NOT (i IN (1, NULL))", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":"IN","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":1},{"type":"literal","dataType":"int","value":null}]}]}], "rows": []}, + {"sql": "NOT (i IN (1, 3))", "predicates": [{"type":"function","name":"NOT","children":[{"type":"function","name":"IN","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":1},{"type":"literal","dataType":"int","value":3}]}]}], "rows": [1, 6, 7]}, + {"sql": "i IS NULL OR d > 1.0D", "predicates": [{"type":"function","name":"OR","children":[{"type":"function","name":"IS_NULL","children":[{"type":"column","name":["i"]}]},{"type":"function","name":">","children":[{"type":"column","name":["d"]},{"type":"literal","dataType":"double","value":1.0}]}]}], "rows": [2, 3, 5, 7]}, + {"sql": "i > 1 AND d IS NOT NULL", "predicates": [{"type":"function","name":">","children":[{"type":"column","name":["i"]},{"type":"literal","dataType":"int","value":1}]}, {"type":"function","name":"IS_NOT_NULL","children":[{"type":"column","name":["d"]}]}], "rows": [1, 3, 7]}, + {"sql": "s = '\uD83D\uDE00'", "predicates": [{"type":"function","name":"=","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"\uD83D\uDE00"}]}], "rows": [3]}, + {"sql": "s > '\uFFFD'", "predicates": [{"type":"function","name":">","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"\uFFFD"}]}], "rows": [3]}, + {"sql": "s < 'b'", "predicates": [{"type":"function","name":"<","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"b"}]}], "rows": [0, 1, 5]}, + {"sql": "startswith(s, 'a')", "predicates": [{"type":"function","name":"STARTS_WITH","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"a"}]}], "rows": [0, 5]}, + {"sql": "endswith(s, '\u00E9')", "predicates": [{"type":"function","name":"ENDS_WITH","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"\u00E9"}]}], "rows": [2, 6]}, + {"sql": "contains(s, '\uD83D\uDE00')", "predicates": [{"type":"function","name":"CONTAINS","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":"\uD83D\uDE00"}]}], "rows": [3]}, + {"sql": "s IN ('', '\u00E9')", "predicates": [{"type":"function","name":"IN","children":[{"type":"column","name":["s"]},{"type":"literal","dataType":"string","value":""},{"type":"literal","dataType":"string","value":"\u00E9"}]}], "rows": [1, 2]}, + {"sql": "f = CAST(0.1 AS FLOAT)", "predicates": [{"type":"function","name":"=","children":[{"type":"column","name":["f"]},{"type":"literal","dataType":"float","value":0.10000000149011612}]}], "rows": [0, 3]}, + {"sql": "f > CAST(0.1 AS FLOAT)", "predicates": [{"type":"function","name":">","children":[{"type":"column","name":["f"]},{"type":"literal","dataType":"float","value":0.10000000149011612}]}], "rows": [1, 4, 5, 6, 7]}, + {"sql": "b > X'7F'", "predicates": [{"type":"function","name":">","children":[{"type":"column","name":["b"]},{"type":"literal","dataType":"binary","value":"fw=="}]}], "rows": [2, 3, 4]}, + {"sql": "b < X'80'", "predicates": [{"type":"function","name":"<","children":[{"type":"column","name":["b"]},{"type":"literal","dataType":"binary","value":"gA=="}]}], "rows": [0, 1, 4, 5, 7]}, + {"sql": "(id > 3) = (id < 3)", "predicates": [{"type":"function","name":"=","children":[{"type":"function","name":">","children":[{"type":"column","name":["id"]},{"type":"literal","dataType":"int","value":3}]},{"type":"function","name":"<","children":[{"type":"column","name":["id"]},{"type":"literal","dataType":"int","value":3}]}]}], "rows": [3]} + ] +} diff --git a/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h new file mode 100644 index 000000000000..954555b6b670 --- /dev/null +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -0,0 +1,713 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * spark_datasource.h: the binary interface between Spark and data sources implemented in native + * code, such as Rust or C++. Since Spark 4.4.0. + * + * A native data source is a shared library that exports a single C function, spark_ds_init, + * which returns the table of the functions that implement the data source, a SparkDsLibrary. + * Spark calls them with the Foreign Function and Memory API of Java, so native data sources + * require Java 22 or later. Each library has its own table, so any number of libraries can be + * loaded side by side. Spark adapts a native data source to Data Source V2, with batch and + * micro-batch streaming reads, batch and streaming writes, and predicate, column and limit + * pushdown. + * + * This header is self-contained: it only depends on the C standard library, and defines the + * structs of the Arrow C data interface unless they are already defined, as by Arrow C++. A library + * can copy it into its source tree, or use the copy in the include directory of a Spark + * distribution. A library built against it works with all the releases of Spark that support its + * SPARK_DS_ABI_VERSION. + * + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of my_source is + * spark_datasource_my_source, such as libspark_datasource_my_source.so in /usr/local/lib on Linux. + * Spark looks for it in the directories of java.library.path, which include LD_LIBRARY_PATH on + * Linux, and in the lib directory of each installation prefix in PATH. A library can also be + * distributed in a native data source package: a zip file with the extension .sparkpkg that + * contains a manifest named spark-native-datasource.json and the library built for one or more + * platforms, added to a session with spark.addArtifact or under the paths in + * spark.sql.dataSource.native.paths. See the "Native Data Sources" page of the Spark SQL guide for + * details and examples. + * + * Lifecycle + * --------- + * Spark loads a library once in each JVM, and calls spark_ds_init before any other function. It + * creates a data source with create on the driver for each schema inference, scan and write, and + * closes it right after. create also returns the capabilities of the data source: Spark reports + * the operations that a data source does not declare as unsupported, without calling the library. + * Spark also creates a data source while it analyzes a query, to reject a batch write that the + * data source does not support before it computes the data to write. That data source has the + * options of the table, those of CREATE TABLE for a table in a catalog, which the options of a + * write can override, so Spark checks the capabilities of each operation again. + * + * - Batch reads. On the driver, Spark creates a reader with create_reader, pushes operations down + * by calling, each at most once and in this order, push_predicates, push_limit and + * prune_columns, plans the partitions with partitions, serializes the state of the reader with + * serialize_reader, and closes it. Then it calls read on the executors, in one task for each + * partition. If Spark never plans the partitions, for example to only explain the plan of the + * query, it closes the reader once the reader is no longer referenced. + * - Streaming reads. On the driver, Spark creates a stream reader with create_stream_reader and + * keeps it until the streaming query stops. For each micro-batch, it asks for the latest_offset, + * plans the partitions with stream_partitions and serializes the state with + * serialize_stream_reader, then calls read on the executors. Offsets are defined by the library, + * as compact JSON on a single line: Spark stores each offset as a line of the offset log of the + * checkpoint, so it rejects an offset that is not valid JSON or contains a line break. When a + * streaming query restarts, from its checkpoint, Spark creates a new stream reader and may call + * stream_partitions again for the last range that it planned before, so the partitions of a + * range must read the same data each time. + * - Writes. On the driver, right before it starts the tasks of the write (or of a micro-batch), + * Spark creates a writer with create_writer, and serializes its state with serialize_writer. + * With adaptive query execution, Spark first runs the stages of the data to write that end with + * a shuffle, and creates no writer if they fail. The rest of the data, such as its last stage, + * is computed in the tasks of the write, so a writer may still be aborted because computing the + * data failed: do irreversible work, such as removing the existing data for an overwrite, in + * commit. On the executors, each task creates a data writer with create_data_writer, calls write + * for each batch of at most spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends + * once it holds spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls commit_data_writer, + * or abort_data_writer if the task or the commit fails, before it closes the data writer. Then + * the driver calls commit with the messages of the tasks if they all succeeded, or abort + * otherwise, including when commit fails, and closes the writer. Several attempts of a task may + * write the same partition, after a failure or with speculative execution, each with a data + * writer of its own. Spark lets the attempts of a partition commit one at a time, and passes to + * commit the message of one committed attempt for each partition, so the write consists of the + * data of these attempts only. It aborts the data writers of the other attempts, except on an + * executor that is lost, where a data writer is neither committed nor aborted, and may even have + * committed before its message reached the driver. A streaming write commits or aborts each + * micro-batch (epoch). Spark may commit the same epoch more than once, for example when a + * streaming query restarts after a failure, so committing an epoch that is already committed + * must have no effect, for exactly-once output. The ID of the query, which create_writer gets for + * a streaming write, stays the same when the query restarts from its checkpoint, so the ID of the + * query and the epoch ID identify the output of a micro-batch, also when other queries write to + * the same destination. A micro-batch streaming query creates a new stream writer for each + * micro-batch, and Spark closes the previous stream writer of the query when the query creates a + * new one. Spark is not told when a streaming query stops, so it closes the last stream writer of + * the query once the writer is no longer referenced, or when the query is restarted. Streaming + * writes are micro-batch only: a query with continuous processing fails before create_writer. + * + * Conventions + * ----------- + * - Errors. The functions that return an int return 0 on success, or any other value on failure. + * They then set the message that last_error returns, and Spark reports it as a + * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail, and they may run after the query, + * from another thread: do the work that can fail a query, such as flushing data, in commit. A + * C++ exception or a Rust panic must not cross the interface: catch it and return an error + * instead. + * - Handles. create, create_reader, create_stream_reader, create_writer and create_data_writer + * return an opaque, non-NULL handle in their last argument, typically a pointer: Spark reports a + * NULL handle as an error. Spark releases each handle exactly once, with the close function of + * its kind: close, close_reader, close_stream_reader, close_writer or close_data_writer, and no + * other function releases a handle. A handle is not always released right after its last use: a + * batch reader whose partitions are never planned and the last stream writer of a query are + * released once they are no longer referenced (see Lifecycle above). A reader or writer handle + * must not depend on the data source handle it was created from, because Spark closes the data + * source handle right after creating it. + * - Memory. What Spark passes to a function, such as strings, byte arrays, SparkDsCreateInfo and + * SparkDsWriteInfo, is only valid during the call, except for SparkDsHost and + * SparkDsTaskContext, which document how long they live. Strings are UTF-8 and NUL-terminated: + * Spark fails an operation whose options or column names contain the character U+0000. + * The JSON of the predicates is ASCII: it escapes the other characters, which any JSON parser + * decodes. The library returns strings and byte arrays by passing them to a SparkDsSink, which + * copies them. + * - State. Partitions, the state of readers and writers, and commit messages are byte arrays in a + * format defined by the library. Spark sends them to the executors and back. + * - Arrow data. Schemas and data are exchanged with the Arrow C data interface + * (https://arrow.apache.org/docs/format/CDataInterface.html), as structs allocated by Spark. + * For an output struct, the library initializes it and Spark takes ownership. For an input + * struct, the library may move it to take ownership; otherwise, Spark releases it after the + * call returns, so the library must not use it afterwards. The arrays that the library exports + * must have an offset of 0, at every level: Arrow Java ignores the offsets of the arrays it + * imports, so Spark rejects an array with another offset, such as a zero-copy slice. Copy a + * sliced array before exporting it. Spark also rejects an array with fewer values than its + * parent refers to, such as a column shorter than its batch, and a batch with null rows, since + * Arrow Java ignores the validity of the struct of a batch: mark the values of its columns as + * null instead. As the C data interface requires, an array without a validity buffer must have + * a null_count of 0, and the buffers of an array must hold all of its values: the interface + * does not carry the sizes of the buffers, so Spark cannot check them, and reads past the end + * of a buffer that is too short. + * - Schemas. These are the types of version 1, with their Arrow types and the formats of the C + * data interface: boolean is bool (b), tinyint, smallint, int and bigint are int8 (c), int16 + * (s), int32 (i) and int64 (l), float and double are float32 (f) and float64 (g), decimal(p, s) + * is decimal128(p, s) (d:p,s), string is utf8 (u), binary is binary (z), date is date32 (tdD), + * timestamp is timestamp[us] with the time zone of the session (tsu:), timestamp_ntz is + * timestamp[us] without a time zone (tsu:), and array, map and struct are list (+l), map (+m) + * and struct (+s). Spark fails an operation whose schema has another type, such as an + * interval, variant or time, before it passes the schema to the library: a later version may + * support other types, for the data sources that declare a capability for them. Data read from + * a library may also use large_utf8 (U) and utf8_view (vu) for strings, large_binary (Z) and + * binary_view (vz) for binary data, and any names for the fields of the entries of a map. + * Spark rejects the other Arrow types, such as unsigned integers, decimal256, large lists, + * fixed-size lists and list views, and dictionary encoding. The data must agree with the + * nullability of the schema to read: Spark fails a read whose data has a null where the schema + * declares a column, a field, an element or a map value not nullable. + * - Options. Option keys are passed in lower case, because Spark matches them + * case-insensitively. + * - Optional functions. A library only sets the functions of the operations it supports in its + * table, and leaves the others NULL. A data source must not declare a capability whose + * functions are NULL: Spark then fails the operation. The pushdown functions and the serialize + * functions are optional: without them, Spark pushes nothing down, and the reader or writer has + * no state. So are schema, without which the user must specify the schema, commit_offset and + * last_error. + * - Threads. Spark calls the library from multiple threads, and calls on different handles may + * run concurrently. Calls on the same handle never run concurrently. A handle may be released + * from a thread other than the threads that used it, such as a JVM Cleaner thread when Spark + * releases a handle once it is no longer referenced, for example the last stream writer of a + * streaming query. A SparkDsSink must be called on the thread of the call that it was passed + * to, during that call. The functions of SparkDsHost and SparkDsTaskContext, and the callbacks + * of the stream that read initializes, can be called from any thread. The callbacks of the + * stream never run concurrently. + * + * Compatibility + * ------------- + * The function table and the structs that Spark passes to the library start with their size, so + * that they can grow. Within a SPARK_DS_ABI_VERSION: + * - fields are only appended to the structs, including SparkDsLibrary, and the existing fields + * never change, nor do the semantics of the existing functions; + * - Spark only reads the functions that fit in the size of the table of a library, so a library + * built against an earlier header keeps working, with the functions that it does not know as + * NULL. A library must only read the fields that fit in the size of a struct that Spark passes; + * - new values may appear where this documentation says so, such as the names of pushed + * predicates. Spark only uses a new capability or write mode with a data source that declares + * it. + * An incompatible change uses a new SPARK_DS_ABI_VERSION: spark_ds_init gets the newest version + * that Spark supports, the library returns a table of a version that it implements, and Spark + * rejects a table of a version that it does not support. Spark supports version 1 at least in all + * the releases of Spark 4 from 4.4.0. + */ +#ifndef SPARK_DATASOURCE_H +#define SPARK_DATASOURCE_H + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* The version of this interface. */ +#define SPARK_DS_ABI_VERSION 1 + +/* Exports spark_ds_init from the library, also when the library hides its other symbols. */ +#if defined(_WIN32) +#define SPARK_DS_EXPORT __declspec(dllexport) +#else +#define SPARK_DS_EXPORT __attribute__((visibility("default"))) +#endif + +/* The Arrow C data interface: https://arrow.apache.org/docs/format/CDataInterface.html */ +#ifndef ARROW_C_DATA_INTERFACE +#define ARROW_C_DATA_INTERFACE + +#define ARROW_FLAG_DICTIONARY_ORDERED 1 +#define ARROW_FLAG_NULLABLE 2 +#define ARROW_FLAG_MAP_KEYS_SORTED 4 + +struct ArrowSchema { + const char* format; + const char* name; + const char* metadata; + int64_t flags; + int64_t n_children; + struct ArrowSchema** children; + struct ArrowSchema* dictionary; + void (*release)(struct ArrowSchema*); + void* private_data; +}; + +struct ArrowArray { + int64_t length; + int64_t null_count; + int64_t offset; + int64_t n_buffers; + int64_t n_children; + const void** buffers; + struct ArrowArray** children; + struct ArrowArray* dictionary; + void (*release)(struct ArrowArray*); + void* private_data; +}; + +#endif /* ARROW_C_DATA_INTERFACE */ + +#ifndef ARROW_C_STREAM_INTERFACE +#define ARROW_C_STREAM_INTERFACE + +struct ArrowArrayStream { + int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out); + int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out); + const char* (*get_last_error)(struct ArrowArrayStream*); + void (*release)(struct ArrowArrayStream*); + void* private_data; +}; + +#endif /* ARROW_C_STREAM_INTERFACE */ + +/* + * Receives a string or a byte array that the library returns, and copies it. `sink` is the + * argument passed with the sink to the function that returns it. + */ +typedef void (*SparkDsSink)(void* sink, const uint8_t* data, int64_t length); + +/* --------------------------------------------------------------------------------------------- + * What Spark passes to the library + * --------------------------------------------------------------------------------------------- */ + +/* The levels of the messages that a library logs with SparkDsHost.log. */ +#define SPARK_DS_LOG_ERROR 1 +#define SPARK_DS_LOG_WARN 2 +#define SPARK_DS_LOG_INFO 3 +#define SPARK_DS_LOG_DEBUG 4 +#define SPARK_DS_LOG_TRACE 5 + +/* Spark, as spark_ds_init gets it. It stays valid until the JVM exits. */ +typedef struct SparkDsHost { + /* + * The size of this struct in the header that Spark was built with. A later version of Spark may + * append fields, so a library must only read the fields that fit in it. + */ + int64_t size; + /* The newest version of this interface that Spark supports. */ + int32_t abi_version; + /* The version of Spark, such as "4.4.0". */ + const char* spark_version; + /* + * Logs a message, a NUL-terminated UTF-8 string, at one of the SPARK_DS_LOG_* levels: in the log + * of the driver or of the executor, with the path of the library. + */ + void (*log)(int32_t level, const char* message); +} SparkDsHost; + +/* A data source, as create gets it. */ +typedef struct SparkDsCreateInfo { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* + * The name of the data source, as listed in the package manifest, or in lower case for an + * installed library. A library can implement several data sources. + */ + const char* name; + /* + * The num_options options given by the user, such as spark.read.option("key", "value"), with + * the keys in lower case. + */ + const char* const* option_keys; + const char* const* option_values; + int64_t num_options; +} SparkDsCreateInfo; + +/* + * The task that calls read or create_data_writer, on an executor. It stays valid until the stream + * that read returns is released, or until the data writer is closed. + */ +typedef struct SparkDsTaskContext { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* The ID of the stage of the task, and the attempt number of the stage. */ + int32_t stage_id; + int32_t stage_attempt_number; + /* The partition that the task reads or writes, and the attempt number of the task. */ + int32_t partition_id; + int32_t attempt_number; + /* The ID of the task attempt, unique in the Spark application. */ + int64_t task_attempt_id; + /* + * Returns whether the task was killed, for example because its query was canceled. Spark cannot + * interrupt native code, so a read or a write that runs for long, or blocks, should check it + * regularly, and fail once it is true. + */ + bool (*is_interrupted)(const struct SparkDsTaskContext* task); +} SparkDsTaskContext; + +/* The modes of a write. */ +#define SPARK_DS_WRITE_APPEND 0 +#define SPARK_DS_WRITE_TRUNCATE 1 + +/* A write, as create_writer gets it. */ +typedef struct SparkDsWriteInfo { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* + * SPARK_DS_WRITE_APPEND to add the data to the existing data, or SPARK_DS_WRITE_TRUNCATE to + * replace the existing data: with the save mode "overwrite" of a batch write, and in each + * micro-batch with the output mode "complete" of a streaming write. Spark only passes + * SPARK_DS_WRITE_TRUNCATE to a data source with the capability SPARK_DS_TRUNCATE. + */ + int32_t mode; + /* + * The ID of the streaming query of a streaming write, or NULL for a batch write. It stays the + * same when the query restarts from its checkpoint, so with the epoch IDs, it identifies the + * commits of the query: see Writes above. + */ + const char* query_id; +} SparkDsWriteInfo; + +/* --------------------------------------------------------------------------------------------- + * What the library returns to Spark + * --------------------------------------------------------------------------------------------- */ + +/* The capabilities of a data source, which create returns: the operations that it supports. */ +#define SPARK_DS_BATCH_READ 1 /* batch reads, with create_reader */ +#define SPARK_DS_MICRO_BATCH_READ 2 /* micro-batch streaming reads, with create_stream_reader */ +#define SPARK_DS_BATCH_WRITE 4 /* batch writes, with create_writer */ +#define SPARK_DS_STREAMING_WRITE 8 /* streaming writes, with create_writer */ +#define SPARK_DS_TRUNCATE 16 /* writes with the mode SPARK_DS_WRITE_TRUNCATE */ + +/* How a reader uses a predicate pushed with push_predicates. */ +#define SPARK_DS_PREDICATE_UNUSED 0 /* not at all */ +/* only to skip data, so it may return rows for which the predicate is not true */ +#define SPARK_DS_PREDICATE_PARTIAL 1 +/* to only return the rows for which the predicate is true */ +#define SPARK_DS_PREDICATE_COMPLETE 2 + +/* + * The functions of a library, which spark_ds_init returns. A library sets the functions of the + * operations it supports, and leaves the others NULL: see "Optional functions" above. + */ +typedef struct SparkDsLibrary { + /* + * The size of this struct in the header that the library was built with. Spark only reads the + * functions that fit in it. + */ + int64_t size; + /* The version of this interface that the table implements, SPARK_DS_ABI_VERSION. */ + int32_t abi_version; + + /* + * Returns the message of the last error on the calling thread, as a NUL-terminated UTF-8 string + * that stays valid until the next call on that thread. Optional: without it, Spark only reports + * the value that the function returned. + */ + const char* (*last_error)(void); + + /* ------------------------------------------------------------------------------------------- + * Data source (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Creates a data source. + * + * info: the name of the data source, and the options given by the user + * capabilities: set to the capabilities of the data source, SPARK_DS_* flags combined with |, + * which may depend on its name and its options. It is 0 when create is called. + * source: set to the data source handle + */ + int (*create)(const SparkDsCreateInfo* info, uint64_t* capabilities, void** source); + + /* + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. + * + * out: initialized with a struct type whose fields are the columns + */ + int (*schema)(void* source, struct ArrowSchema* out); + + /* + * Creates a reader for a batch scan. + * + * schema: the schema to read: the one the user specified, or the one returned by schema + * reader: set to the reader handle + */ + int (*create_reader)(void* source, struct ArrowSchema* schema, void** reader); + + /* + * Creates a reader for a micro-batch streaming scan. + * + * schema: the schema to read: the one the user specified, or the one returned by schema + * reader: set to the stream reader handle + */ + int (*create_stream_reader)(void* source, struct ArrowSchema* schema, void** reader); + + /* + * Creates a writer for a batch write, or for a micro-batch of a streaming write. + * + * schema: the schema of the data to write + * info: the mode of the write, and the ID of the query of a streaming write + * writer: set to the writer handle + */ + int (*create_writer)(void* source, struct ArrowSchema* schema, const SparkDsWriteInfo* info, + void** writer); + + /* Releases a data source handle. */ + void (*close)(void* source); + + /* ------------------------------------------------------------------------------------------- + * Batch reader (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Pushes predicates down to the reader. The predicates are combined with AND. + * + * Each predicate is a JSON expression tree. An expression is one of: + * - {"type": "column", "name": ["a", "b"]}: a column. A nested field has a name part for each + * level. + * - {"type": "literal", "dataType": "bigint", "value": 1}: a literal. The data type is a Spark + * SQL type name such as int, string or decimal(10,2). The value is null, a JSON boolean or + * number for boolean and numeric types, a string for string, a base64 string for binary, a + * decimal string for decimal, the days since the epoch for date, and the microseconds since + * the epoch for timestamp and timestamp_ntz. Read a literal at the precision of its data type: + * a float literal is the exact value of the float, written as the double of the same value, + * and the values of the integer types, dates and timestamps are exact 64-bit integers, which + * a double does not hold above 2^53. Spark does not pass a string that is not valid UTF-8. + * - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in + * org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=> + * (equal, with null equal to null), <, <=, > and >= with two children, IN with the value to + * look for and then the values to look in, IS_NULL and IS_NOT_NULL with one child, + * STARTS_WITH, ENDS_WITH and CONTAINS with the string and the string to look for, AND and OR + * with two predicates, and NOT with one. The children of the other functions are columns, + * literals or predicates. + * Spark does not pass predicates that it cannot express this way, such as predicates that use + * arithmetic or other functions, nor predicates on columns that contain strings with a + * non-binary collation, such as UTF8_LCASE: strings are compared by their UTF-8 bytes. Nor does + * it pass comparisons and IN on arrays, maps or structs, whose values Spark compares with other + * rules, for example with the nulls in them equal to each other, but only on values of the + * other types, such as a field of a struct that is an int. IS_NULL and IS_NOT_NULL apply to + * columns of any type. A later version of Spark may pass predicates with other names without + * changing SPARK_DS_ABI_VERSION, so a reader must ignore a predicate that contains a name that + * it does not know, or a data type that it does not support. + * + * Spark evaluates all the predicates again after reading, so a reader may use them only to skip + * data, for example the files whose statistics show that no row matches. But it must never drop + * or skip a row for which the predicates are true under the semantics of Spark SQL, which Spark + * cannot detect. In particular: + * - For floating-point columns, NaN is equal to NaN and greater than any other value, including + * positive infinity, and -0.0 is equal to 0.0. IEEE 754 comparisons treat NaN differently, and + * total orders, such as f64::total_cmp in Rust or the comparison kernels of arrow-rs, also + * order -0.0 before 0.0. For example, d > 1.0 is true for a NaN value of d, and d = 0.0 for + * -0.0. + * - Comparisons with null are null, and a row only matches if the predicates are true. For + * example, NOT (i IN (1, NULL)) matches no row. + * - Strings are compared by their UTF-8 bytes, which order the characters outside the Basic + * Multilingual Plane after all the others, unlike comparisons of UTF-16 strings. + * - Binary values are compared as unsigned bytes, in lexicographic order: X'80' is greater than + * X'7F', and a value is greater than its prefixes. + * - To use only a part of a predicate, a reader may take a part that it does not support as + * true, but only when every function between that part and the root of the predicate is AND + * or OR. Under any other function, such as NOT, or = and <=> between two predicates, that + * would drop rows: (id > 3) = (id < 3) is true for the row whose id is 3, but taking id < 3 + * as true makes it false. + * The file org/apache/spark/sql/datasource/predicate-cases.json in the jar of Spark SQL lists + * predicates as Spark passes them to this function, with the rows of a table for which they are + * true, to test a reader with. + * + * predicates: the count predicates, as JSON + * uses: an array of count values, all SPARK_DS_PREDICATE_UNUSED, to set to how the reader uses + * each predicate, a SPARK_DS_PREDICATE_* value. Spark shows the predicates that the reader + * uses as pushed in the plan of the query, and evaluates all the predicates again after + * reading anyway. It takes other values as SPARK_DS_PREDICATE_PARTIAL. + */ + int (*push_predicates)(void* reader, const char* const* predicates, int64_t count, + int32_t* uses); + + /* + * Pushes a LIMIT down to the reader. Spark only pushes a limit when the query does not filter + * the rows, since it evaluates the predicates after reading. + * + * applied: set to whether the reader returns at most `limit` rows in total. Spark applies the + * limit again after reading either way. + */ + int (*push_limit)(void* reader, int32_t limit, bool* applied); + + /* + * Prunes the columns to read. + * + * columns: the names of the count top-level columns that the query needs, which include the + * columns of the pushed predicates. There may be none, for example for + * SELECT count(*): then the batches only have a number of rows. + * pruned: set to whether read returns exactly these columns, in this order. If not, it returns + * all the columns of the schema to read. + */ + int (*prune_columns)(void* reader, const char* const* columns, int64_t count, bool* pruned); + + /* + * Plans the partitions of the scan. Spark reads each partition in a separate task. + * + * add_partition: called with each partition, serialized in a format defined by the library + */ + int (*partitions)(void* reader, SparkDsSink add_partition, void* sink); + + /* + * Serializes the state of the reader that read needs on the executors, such as the pushed + * predicates. Spark calls it after partitions. + * + * set_state: called once with the state, or not at all if there is none + */ + int (*serialize_reader)(void* reader, SparkDsSink set_state, void* sink); + + /* Releases a batch reader handle. */ + void (*close_reader)(void* reader); + + /* ------------------------------------------------------------------------------------------- + * Stream reader (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Returns the offset to start from when the streaming query has no checkpoint. + * + * set_offset: called once with the offset, as compact JSON on a single line + */ + int (*initial_offset)(void* reader, SparkDsSink set_offset, void* sink); + + /* + * Returns the most recent offset available. + * + * set_offset: called once with the offset, as compact JSON on a single line + */ + int (*latest_offset)(void* reader, SparkDsSink set_offset, void* sink); + + /* + * Plans the partitions that read the data after `start` up to and including `end`, two offsets + * as JSON. + * + * add_partition: called with each partition, serialized in a format defined by the library + */ + int (*stream_partitions)(void* reader, const char* start, const char* end, + SparkDsSink add_partition, void* sink); + + /* + * Serializes the state of the stream reader that read needs on the executors. Spark calls it + * after each call to stream_partitions. + * + * set_state: called once with the state, or not at all if there is none + */ + int (*serialize_stream_reader)(void* reader, SparkDsSink set_state, void* sink); + + /* + * Tells the library that Spark has processed all the data up to and including `end`, an offset + * as JSON, and will not ask for it again. + */ + int (*commit_offset)(void* reader, const char* end); + + /* Releases a stream reader handle, when the streaming query stops. */ + void (*close_stream_reader)(void* reader); + + /* ------------------------------------------------------------------------------------------- + * Partition reader (executor) + * ------------------------------------------------------------------------------------------- */ + + /* + * Reads a partition of a batch or streaming scan. + * + * task: the task that reads the partition, valid until the stream is released + * reader_state: the state returned by serialize_reader or serialize_stream_reader, empty if + * there is none + * partition: a partition returned by partitions or stream_partitions + * out: the stream 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. The + * callbacks of the stream can be called from any thread, but never concurrently. 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 get_next returns, and then releases the + * stream. So a get_next that blocks forever keeps the task running, unless the task reaper + * (spark.task.reaper.enabled and spark.task.reaper.killTimeout) kills the executor. Use + * timeouts for blocking operations, and check task->is_interrupted. + */ + int (*read)(const SparkDsTaskContext* task, const uint8_t* reader_state, + int64_t reader_state_length, const uint8_t* partition, int64_t partition_length, + struct ArrowArrayStream* out); + + /* ------------------------------------------------------------------------------------------- + * Writer (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Serializes the state of the writer that create_data_writer needs on the executors. Spark + * calls it once, after creating the writer. + * + * set_state: called once with the state, or not at all if there is none + */ + int (*serialize_writer)(void* writer, SparkDsSink set_state, void* sink); + + /* + * Commits a write, when all the tasks succeeded. + * + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * messages, message_lengths: the num_messages messages returned by commit_data_writer of each + * task + */ + int (*commit)(void* writer, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages); + + /* + * Aborts a write, when it failed. + * + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * messages, message_lengths: the num_messages messages returned by commit_data_writer of each + * task. A message is NULL if its task did not commit. + */ + int (*abort)(void* writer, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages); + + /* Releases a writer handle. */ + void (*close_writer)(void* writer); + + /* ------------------------------------------------------------------------------------------- + * Data writer (executor) + * ------------------------------------------------------------------------------------------- */ + + /* + * Creates the writer for one task. + * + * task: the task, with the partition that it writes, valid until the data writer is closed + * writer_state: the state returned by serialize_writer, empty if there is none + * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a batch write + * data_writer: set to the data writer handle + */ + int (*create_data_writer)(const SparkDsTaskContext* task, const uint8_t* writer_state, + int64_t writer_state_length, int64_t epoch_id, void** data_writer); + + /* + * Writes a batch of data. + * + * array: a struct array whose fields are the columns + * schema: the schema of the array + */ + int (*write)(void* data_writer, struct ArrowArray* array, struct ArrowSchema* schema); + + /* + * Commits the data written by the task. If it fails, Spark calls abort_data_writer. If it + * succeeds, Spark does not call abort_data_writer, even if it then fails the task, for example + * because set_message was called more than once. + * + * set_message: called once with the message to pass to commit on the driver, or not at all for + * an empty message + */ + int (*commit_data_writer)(void* data_writer, SparkDsSink set_message, void* sink); + + /* Aborts the data written by the task. */ + int (*abort_data_writer)(void* data_writer); + + /* Releases a data writer handle, after its commit or abort. */ + void (*close_data_writer)(void* data_writer); +} SparkDsLibrary; + +/* + * The function that a library exports. Spark calls it once in each JVM, when it loads the library, + * before any other function. + * + * host: Spark, valid until the JVM exits + * library: set to the function table of the library, for the version of this interface that it + * implements, or the newest one up to host->abi_version if it implements several. Spark + * reads the table after the call, so it is typically the address of a static struct. + * + * Returns 0 on success, or any other value if the library cannot run, after it logs why with + * host->log. + */ +SPARK_DS_EXPORT int spark_ds_init(const SparkDsHost* host, const SparkDsLibrary** library); + +#ifdef __cplusplus +} +#endif + +#endif /* SPARK_DATASOURCE_H */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataSourceRegistration.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataSourceRegistration.scala index 9d763edb079e..d725922d95e8 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/DataSourceRegistration.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/DataSourceRegistration.scala @@ -65,7 +65,9 @@ private[sql] class DataSourceRegistration private[sql] (dataSourceManager: DataS if (dataSourceManager.dataSourceExists(name)) return try { - DataSource.lookupDataSource(name, SQLConf.get) + // A Python data source takes precedence over a native data source with the same name, so + // registering it does not depend on the native libraries that happen to be installed. + DataSource.lookupDataSource(name, SQLConf.get, includeNativeDataSources = false) throw QueryCompilationErrors.dataSourceAlreadyExists(name) } catch { case e: SparkClassNotFoundException if e.getCondition == "DATA_SOURCE_NOT_FOUND" => // OK diff --git a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala index 2a03a58502c6..e0fb397d36cd 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala @@ -36,6 +36,7 @@ import org.apache.spark.internal.{Logging, LogKeys} import org.apache.spark.internal.config.{CONNECT_SCALA_UDF_STUB_PREFIXES, EXECUTOR_USER_CLASS_PATH_FIRST, JAR_IVY_CONNECT_TIMEOUT, JAR_IVY_READ_TIMEOUT, JAR_IVY_REPO_PATH, JAR_IVY_SETTING_PATH} import org.apache.spark.sql.Artifact import org.apache.spark.sql.classic.SparkSession +import org.apache.spark.sql.execution.datasources.v2.ffi.NativeDataSourcePackage import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.util.ArtifactUtils import org.apache.spark.storage.{BlockManager, CacheId, StorageLevel} @@ -153,6 +154,28 @@ class ArtifactManager(session: SparkSession) extends AutoCloseable with Logging */ def getPythonIncludes: Seq[String] = pythonIncludeList.asScala.toSeq + /** + * Returns the native data source packages added to this session, as the local files of the + * driver, and the UUID under which the executors store them, if the session is isolated. + */ + private[sql] def getNativeDataSourcePackages: (Seq[File], Option[String]) = { + val packages = sparkContextRelativePaths.asScala.collect { + case (SparkContextResourceType.FILE, path, _) + if path.toString.endsWith(Artifact.NATIVE_DATA_SOURCE_PACKAGE_EXTENSION) => + artifactPath.resolve(path).toFile + }.toSeq + (packages, Option(state).map(_.uuid)) + } + + /** Returns whether a file with the given name was added to this session as a file artifact. */ + private[sql] def hasFileArtifact(fileName: String): Boolean = { + val path = Paths.get("files", fileName) + sparkContextRelativePaths.asScala.exists { + case (SparkContextResourceType.FILE, relativePath, _) => relativePath == path + case _ => false + } + } + protected[sql] def getCachedBlockId(hash: String): Option[CacheId] = { Option(hashToCachedIdMap.get(hash)).map(_.id) } @@ -618,6 +641,9 @@ object ArtifactManager extends Logging { } } + // Forget the native data source packages read from the artifacts folder + Utils.tryLogNonFatalError(NativeDataSourcePackage.forgetPackagesIn(artifactPath.toFile)) + // Clean up artifacts folder try { Utils.deleteRecursively(artifactPath.toFile) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala index f22008d3b232..3183cc8c8c12 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala @@ -32,8 +32,7 @@ import org.apache.spark.sql.connector.catalog.TableCapability._ import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.command.DDLUtils import org.apache.spark.sql.execution.datasources.{DataSource, LogicalRelation} -import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Utils, FileDataSourceV2} -import org.apache.spark.sql.execution.datasources.v2.python.PythonDataSourceV2 +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Utils, FileDataSourceV2, NamedTableProvider} import org.apache.spark.sql.execution.streaming.runtime.StreamingRelation import org.apache.spark.sql.sources.StreamSourceProvider import org.apache.spark.sql.types.StructType @@ -159,7 +158,7 @@ class ResolveDataSource(sparkSession: SparkSession) extends Rule[LogicalPlan] { optionsWithPath.originalMap val dsOptions = new CaseInsensitiveStringMap(finalOptions.asJava) provider match { - case p: PythonDataSourceV2 => p.setShortName(source) + case p: NamedTableProvider => p.setShortName(source) case _ => } val table = diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala index 3d0036f86522..eab75bc3e5bb 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala @@ -629,8 +629,10 @@ final class DataFrameWriter[T] private[sql](ds: Dataset[T]) extends sql.DataFram * Wrap a DataFrameWriter action to track the QueryExecution and time cost, then report to the * user-registered callback functions. */ - private def runCommand(session: SparkSession)(command: LogicalPlan): Unit = { - val qe = new QueryExecution(session, command, df.queryExecution.tracker, + private def runCommand(session: SparkSession)(command: => LogicalPlan): Unit = { + // The command is built with the session active, like a query is analyzed: it looks up the + // data source, and the lookups of the Python and native data sources use the active session. + val qe = new QueryExecution(session, session.withActive(command), df.queryExecution.tracker, shuffleCleanupModeOpt = Some(QueryExecution.determineShuffleCleanupMode(session.sessionState.conf))) qe.assertCommandExecuted() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala index 1301985e3015..0369bcfacfbf 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataStreamWriter.scala @@ -40,8 +40,7 @@ import org.apache.spark.sql.connector.expressions.{ClusterByTransform, FieldRefe import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.command.DDLUtils import org.apache.spark.sql.execution.datasources.{DataSource, DataSourceUtils} -import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Utils, FileDataSourceV2} -import org.apache.spark.sql.execution.datasources.v2.python.PythonDataSourceV2 +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Utils, FileDataSourceV2, NamedTableProvider} import org.apache.spark.sql.execution.streaming._ import org.apache.spark.sql.execution.streaming.runtime.RealTimeModeAllowlist import org.apache.spark.sql.execution.streaming.sources._ @@ -254,51 +253,55 @@ final class DataStreamWriter[T] private[sql](ds: Dataset[T]) extends streaming.D val sink = new ForeachBatchSink[T](foreachBatchWriter, ds.exprEnc) startQuery(sink, extraOptions, catalogTable = catalogTable) } else { - val cls = DataSource.lookupDataSource(source, ds.sparkSession.sessionState.conf) - val disabledSources = - Utils.stringToSeq(ds.sparkSession.sessionState.conf.disabledV2StreamingWriters) - val useV1Source = disabledSources.contains(cls.getCanonicalName) || - // file source v2 does not support streaming yet. - classOf[FileDataSourceV2].isAssignableFrom(cls) - val optionsWithPath = if (path.isEmpty) { extraOptions } else { extraOptions + ("path" -> path.get) } + // The sink is resolved with the session active, like a query is analyzed: the lookups of + // the Python and native data sources use the active session. + val sink = ds.sparkSession.withActive(resolveSink(optionsWithPath)) + startQuery(sink, optionsWithPath, catalogTable = catalogTable) + } + } - val sink = if (classOf[TableProvider].isAssignableFrom(cls) && !useV1Source) { - val provider = cls.getConstructor().newInstance().asInstanceOf[TableProvider] - val sessionOptions = DataSourceV2Utils.extractSessionConfigs( - source = provider, conf = ds.sparkSession.sessionState.conf) - val finalOptions = sessionOptions.filter { case (k, _) => !optionsWithPath.contains(k) } ++ - optionsWithPath.originalMap - val dsOptions = new CaseInsensitiveStringMap(finalOptions.asJava) - // If the source accepts external table metadata, here we pass the schema of input query - // to `getTable`. This is for avoiding schema inference, which can be very expensive. - // If the query schema is not compatible with the existing data, the behavior is undefined. - val outputSchema = if (provider.supportsExternalMetadata()) { - Some(ds.schema) - } else { - None - } - provider match { - case p: PythonDataSourceV2 => p.setShortName(source) - case _ => - } - val table = DataSourceV2Utils.getTableFromProvider( - provider, dsOptions, userSpecifiedSchema = outputSchema) - import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._ - table match { - case table: SupportsWrite if table.supports(STREAMING_WRITE) => - table - case _ => createV1Sink(optionsWithPath) - } + private def resolveSink(optionsWithPath: CaseInsensitiveMap[String]): Table = { + val cls = DataSource.lookupDataSource(source, ds.sparkSession.sessionState.conf) + val disabledSources = + Utils.stringToSeq(ds.sparkSession.sessionState.conf.disabledV2StreamingWriters) + val useV1Source = disabledSources.contains(cls.getCanonicalName) || + // file source v2 does not support streaming yet. + classOf[FileDataSourceV2].isAssignableFrom(cls) + + if (classOf[TableProvider].isAssignableFrom(cls) && !useV1Source) { + val provider = cls.getConstructor().newInstance().asInstanceOf[TableProvider] + val sessionOptions = DataSourceV2Utils.extractSessionConfigs( + source = provider, conf = ds.sparkSession.sessionState.conf) + val finalOptions = sessionOptions.filter { case (k, _) => !optionsWithPath.contains(k) } ++ + optionsWithPath.originalMap + val dsOptions = new CaseInsensitiveStringMap(finalOptions.asJava) + // If the source accepts external table metadata, here we pass the schema of input query + // to `getTable`. This is for avoiding schema inference, which can be very expensive. + // If the query schema is not compatible with the existing data, the behavior is undefined. + val outputSchema = if (provider.supportsExternalMetadata()) { + Some(ds.schema) } else { - createV1Sink(optionsWithPath) + None } - - startQuery(sink, optionsWithPath, catalogTable = catalogTable) + provider match { + case p: NamedTableProvider => p.setShortName(source) + case _ => + } + val table = DataSourceV2Utils.getTableFromProvider( + provider, dsOptions, userSpecifiedSchema = outputSchema) + import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._ + table match { + case table: SupportsWrite if table.supports(STREAMING_WRITE) => + table + case _ => createV1Sink(optionsWithPath) + } + } else { + createV1Sink(optionsWithPath) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala index 52b9e803a64e..77087a65a99d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala @@ -48,7 +48,8 @@ import org.apache.spark.sql.execution.datasources.jdbc.JdbcRelationProvider import org.apache.spark.sql.execution.datasources.json.JsonFileFormat import org.apache.spark.sql.execution.datasources.orc.OrcFileFormat import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat -import org.apache.spark.sql.execution.datasources.v2.FileDataSourceV2 +import org.apache.spark.sql.execution.datasources.v2.{FileDataSourceV2, NamedTableProvider} +import org.apache.spark.sql.execution.datasources.v2.ffi.{NativeDataSourceRegistry, NativeDataSourceV2} import org.apache.spark.sql.execution.datasources.v2.orc.OrcDataSourceV2 import org.apache.spark.sql.execution.datasources.v2.python.PythonDataSourceV2 import org.apache.spark.sql.execution.datasources.xml.XmlFileFormat @@ -666,6 +667,21 @@ object DataSource extends Logging { /** Given a provider name, look up the data source class definition. */ def lookupDataSource(provider: String, conf: SQLConf): Class[_] = { + lookupDataSource(provider, conf, includeNativeDataSources = true) + } + + /** + * Given a provider name, look up the data source class definition. + * + * @param includeNativeDataSources whether to look the name up among the native data sources, + * which come after the Java and Python data sources + */ + private[sql] def lookupDataSource( + provider: String, + conf: SQLConf, + includeNativeDataSources: Boolean): Class[_] = { + // The provider created for a native data source only takes the result of this lookup. + NativeDataSourceRegistry.clearLookup() val provider1 = backwardCompatibilityMap.getOrElse(provider, provider) match { case name if name.equalsIgnoreCase("orc") && conf.getConf(SQLConf.ORC_IMPLEMENTATION) == "native" => @@ -682,6 +698,9 @@ object DataSource extends Logging { val serviceLoader = ServiceLoader.load(classOf[DataSourceRegister], loader) lazy val isUserDefinedDataSource = SparkSession.getActiveSession.exists( _.sessionState.dataSourceManager.dataSourceExists(provider)) + // The provider created for the native data source takes the result of this lookup. + lazy val isNativeDataSource = includeNativeDataSources && + NativeDataSourceRegistry.lookupForProvider(provider, conf).isDefined try { serviceLoader.asScala.filter(_.shortName().equalsIgnoreCase(provider1)).toList match { @@ -703,6 +722,8 @@ object DataSource extends Logging { throw QueryCompilationErrors.failedToFindKafkaDataSourceError(provider1) } else if (isUserDefinedDataSource) { classOf[PythonDataSourceV2] + } else if (isNativeDataSource) { + classOf[NativeDataSourceV2] } else { throw QueryExecutionErrors.dataSourceNotFoundError(provider1, error) } @@ -718,9 +739,9 @@ object DataSource extends Logging { } } case head :: Nil => - // We do not check whether the provider is a Python data source - // (isUserDefinedDataSource) to avoid the lookup cost. Java data sources - // always take precedence over Python user-defined data sources. + // We do not check whether the provider is a Python or native data source + // (isUserDefinedDataSource, isNativeDataSource) to avoid the lookup cost. Java data + // sources always take precedence over Python and native data sources. head.getClass case sources => // There are multiple registered aliases for the input. If there is single datasource @@ -783,7 +804,7 @@ object DataSource extends Logging { case t: TableProvider if !useV1Sources.contains(cls.getCanonicalName.toLowerCase(Locale.ROOT)) => t match { - case p: PythonDataSourceV2 => p.setShortName(provider) + case p: NamedTableProvider => p.setShortName(provider) case _ => } Some(t) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/NamedTableProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/NamedTableProvider.scala new file mode 100644 index 000000000000..badff781627c --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/NamedTableProvider.scala @@ -0,0 +1,30 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.datasources.v2 + +import org.apache.spark.sql.connector.catalog.TableProvider + +/** + * A [[TableProvider]] that serves the data sources registered under their own short names, such + * as Python and native data sources. `DataSource.lookupDataSource` resolves those names to the + * provider class, so the caller has to tell each new instance which name it was created for with + * [[setShortName]] before using it. + */ +trait NamedTableProvider extends TableProvider { + def setShortName(name: String): Unit +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala new file mode 100644 index 000000000000..7c3204f49094 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala @@ -0,0 +1,152 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.columnar + +import java.util + +import org.apache.spark.sql.connector.catalog.TableCapability +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory} +import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** + * A data source that exchanges its data as columnar batches. [[ColumnarTable]] adapts it to Data + * Source V2. It is implemented by the native data sources, see + * `org.apache.spark.sql.execution.datasources.v2.ffi.NativeDataSource`. + * + * A new instance is created on the driver for each schema inference, scan and write. + */ +trait ColumnarDataSource { + + /** Returns the schema of the data, when the user does not specify it. */ + def schema(): StructType + + /** + * Returns the capabilities that [[ColumnarTable]] reports to Data Source V2, which checks some + * of them while it analyzes a query: by default, all the operations that this trait has. + */ + def capabilities(): util.Set[TableCapability] = + util.EnumSet.of(BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE) + + /** Returns the planner of a batch scan of the given schema. */ + def batchPlanner(schema: StructType): ColumnarBatchPlanner + + /** Returns a reader for a micro-batch streaming scan of the given schema. */ + def streamReader(schema: StructType): ColumnarStreamReader + + /** + * Returns a writer for a batch write. + * + * @param overwrite whether to replace the existing data (save mode "overwrite") instead of + * appending to it + */ + def writer(schema: StructType, overwrite: Boolean): ColumnarWriter + + /** + * Returns a writer for a streaming write. A micro-batch streaming query creates a new one for + * each micro-batch. + * + * @param overwrite whether to replace the existing data in each micro-batch (output mode + * "complete") instead of appending to it + * @param queryId the ID of the streaming query + */ + def streamWriter(schema: StructType, overwrite: Boolean, queryId: String): ColumnarStreamWriter +} + +/** + * Plans a batch scan of a [[ColumnarDataSource]], on the driver. Spark pushes operations down by + * calling, each at most once and in this order, `pushPredicates`, `pushLimit` and `pruneColumns`, + * then plans the scan with `plan`, at most once. Only the [[PartitionReaderFactory]] of the + * planned scan is serialized to the executors. + */ +trait ColumnarBatchPlanner { + + /** + * Returns the predicates that the reader uses, to skip data or to only return the rows for which + * they are true. Spark evaluates all the predicates after reading anyway, and shows the returned + * ones as pushed in the plan. + */ + def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = Array.empty + + /** Returns whether the reader returns at most `limit` rows. Spark applies it again anyway. */ + def pushLimit(limit: Int): Boolean = false + + /** + * Returns whether the partitions return exactly the columns of `requiredSchema`, the top-level + * columns that the query needs, including the columns of the pushed predicates, instead of all + * the columns of the read schema. + */ + def pruneColumns(requiredSchema: StructType): Boolean = false + + /** + * Plans the partitions of the scan, after the pushdown. The reader factory reads them on the + * executors, with columnar readers whose batches only need to be valid until the next one. + */ + def plan(): ColumnarPlannedBatch +} + +/** The partitions of a planned batch scan, and the factory that reads them on the executors. */ +case class ColumnarPlannedBatch( + partitions: Array[InputPartition], + readerFactory: PartitionReaderFactory) + +/** + * Reads the data of a [[ColumnarDataSource]] for a micro-batch streaming scan, with offsets in + * JSON. Spark plans the partitions of each micro-batch on the driver, then serializes the reader + * to the executors, where `read` is called for each partition. + */ +trait ColumnarStreamReader extends Serializable { + def initialOffset(): String + + def latestOffset(): String + + /** Plans the partitions that read the data after `start` up to and including `end`. */ + def partitions(start: String, end: String): Array[InputPartition] + + def read(partition: InputPartition): PartitionReader[ColumnarBatch] + + /** Tells the data source that Spark has processed all the data up to and including `end`. */ + def commit(end: String): Unit = {} + + def stop(): Unit = {} +} + +/** + * Writes data to a [[ColumnarDataSource]] for a batch write. Spark serializes the writer to the + * executors, where `createWriter` is called for each task, and then commits or aborts the write + * on the driver. The batches passed to the data writers are backed by Arrow, and are only valid + * during the call. + */ +trait ColumnarWriter extends Serializable { + def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] + + def commit(messages: Array[WriterCommitMessage]): Unit + + def abort(messages: Array[WriterCommitMessage]): Unit +} + +/** Writes data to a [[ColumnarDataSource]] for a streaming write, one micro-batch at a time. */ +trait ColumnarStreamWriter extends Serializable { + def createWriter(partitionId: Int, taskId: Long, epochId: Long): DataWriter[ColumnarBatch] + + def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit + + def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala new file mode 100644 index 000000000000..76115927950b --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala @@ -0,0 +1,246 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.columnar + +import java.util +import java.util.Locale + +import scala.util.Try + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.connector.catalog.{SupportsRead, SupportsWrite, Table, TableCapability} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read._ +import org.apache.spark.sql.connector.read.streaming.{MicroBatchStream, Offset} +import org.apache.spark.sql.connector.write.{LogicalWriteInfo, WriteBuilder} +import org.apache.spark.sql.internal.connector.SupportsMetadata +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** + * The Data Source V2 [[Table]] of a [[ColumnarDataSource]]. It creates a new data source for each + * scan and each write, with their options on top of the options of the table. Its capabilities + * are those that the data source reports with the options of the table, while a scan or a write + * may have other options: the data source checks them again with those. + * + * @param tableOptions the options of the table, given to `TableProvider.getTable`. They are the + * options of `CREATE TABLE ... USING ... OPTIONS (...)` for a table in a + * catalog, while the scans and the writes of the table only get the options + * of their query. + */ +class ColumnarTable( + shortName: String, + tableSchema: StructType, + tableOptions: util.Map[String, String], + createDataSource: CaseInsensitiveStringMap => ColumnarDataSource) + extends Table with SupportsRead with SupportsWrite { + + override def name(): String = shortName + + override def schema(): StructType = tableSchema + + private lazy val tableCapabilities = + createDataSource(withTableOptions(CaseInsensitiveStringMap.empty())).capabilities() + + override def capabilities(): util.Set[TableCapability] = tableCapabilities + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new ColumnarScanBuilder( + shortName, tableSchema, () => createDataSource(withTableOptions(options))) + } + + override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = { + new ColumnarWriteBuilder( + shortName, info, () => createDataSource(withTableOptions(info.options()))) + } + + private def withTableOptions(options: CaseInsensitiveStringMap): CaseInsensitiveStringMap = { + val merged = new util.HashMap[String, String]() + tableOptions.forEach((key, value) => merged.put(key.toLowerCase(Locale.ROOT), value)) + // The keys of a case-insensitive map are in lower case. + options.forEach((key, value) => merged.put(key, value)) + new CaseInsensitiveStringMap(merged) + } +} + +class ColumnarScanBuilder( + shortName: String, + schema: StructType, + createDataSource: () => ColumnarDataSource) + extends ScanBuilder + with SupportsPushDownV2Filters + with SupportsPushDownLimit + with SupportsPushDownRequiredColumns { + + private lazy val dataSource = createDataSource() + + // Created on first use, because a streaming scan neither pushes anything down nor reads in + // batch, and its data source may not support batch scans at all. + private var planner: ColumnarBatchPlanner = _ + + private var pushed: Array[Predicate] = Array.empty + private var pushedLimit: Option[Int] = None + // The limit passed to the reader, and whether the reader applies it. + private var limitResult: Option[(Int, Boolean)] = None + private var readSchema: StructType = schema + + private def batchPlanner: ColumnarBatchPlanner = { + if (planner == null) { + planner = dataSource.batchPlanner(schema) + } + planner + } + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + pushed = batchPlanner.pushPredicates(predicates) + // Spark evaluates all the predicates again after reading, even the ones that the reader + // evaluates completely: Spark SQL compares values in ways that a data source can easily get + // wrong, such as NaN with other values. The pushed predicates are shown in the plan. + predicates + } + + override def pushedPredicates(): Array[Predicate] = pushed + + override def pushLimit(limit: Int): Boolean = limitResult match { + // Spark can push a limit more than once, for example the limit of LIMIT ... OFFSET, but the + // reader is only asked once. A limit at least as large as the first one is already applied. + case Some((firstLimit, isPushed)) => + isPushed && limit >= firstLimit + case None => + val isPushed = batchPlanner.pushLimit(limit) + limitResult = Some(limit -> isPushed) + if (isPushed) { + pushedLimit = Some(limit) + } + isPushed + } + + // Spark always applies the limit again: a data source is free to return more rows. + override def isPartiallyPushed(): Boolean = true + + override def pruneColumns(requiredSchema: StructType): Unit = { + // Only top-level columns are pruned by the data source. Spark prunes the nested fields of the + // columns it reads. + val required = + StructType(requiredSchema.fieldNames.flatMap(name => schema.find(_.name == name))) + if (required.length < schema.length && batchPlanner.pruneColumns(required)) { + readSchema = required + } + } + + // The scan takes over the planner, so it is built once. + private lazy val scan = new ColumnarScan( + shortName, schema, readSchema, dataSource, Option(planner), pushed, pushedLimit) + + override def build(): Scan = scan +} + +/** + * @param pushedDownPlanner the planner that the operations were pushed down to, if any. Otherwise, + * the scan creates one when its batch is planned. + */ +class ColumnarScan( + shortName: String, + fullSchema: StructType, + prunedSchema: StructType, + dataSource: ColumnarDataSource, + pushedDownPlanner: Option[ColumnarBatchPlanner], + pushedPredicates: Array[Predicate], + pushedLimit: Option[Int]) + extends Scan with SupportsMetadata { + + // Spark may call `toBatch` more than once for the same scan; the partitions are planned once. + private lazy val batch = new ColumnarBatchScan( + pushedDownPlanner.getOrElse(dataSource.batchPlanner(fullSchema))) + + override def readSchema(): StructType = prunedSchema + + override def description(): String = shortName + + override def columnarSupportMode(): Scan.ColumnarSupportMode = + Scan.ColumnarSupportMode.SUPPORTED + + override def toBatch: Batch = batch + + override def toMicroBatchStream(checkpointLocation: String): MicroBatchStream = { + new ColumnarMicroBatchStream(dataSource.streamReader(fullSchema)) + } + + override def getMetaData(): Map[String, String] = { + Map( + "PushedPredicates" -> pushedPredicates.mkString("[", ", ", "]"), + "ReadSchema" -> prunedSchema.simpleString + ) ++ pushedLimit.map(limit => "PushedLimit" -> s"LIMIT $limit") + } +} + +/** + * The batch of a scan, which plans it once, when Spark first asks for its partitions or its reader + * factory. A failure is planned once too: the physical plan can ask again, for example from a lazy + * value, and gets the same error. + */ +class ColumnarBatchScan(planner: ColumnarBatchPlanner) extends Batch { + private lazy val planned: Try[ColumnarPlannedBatch] = Try(planner.plan()) + + override def planInputPartitions(): Array[InputPartition] = planned.get.partitions + + override def createReaderFactory(): PartitionReaderFactory = planned.get.readerFactory +} + +case class ColumnarStreamOffset(json: String) extends Offset + +class ColumnarMicroBatchStream(reader: ColumnarStreamReader) extends MicroBatchStream { + + override def initialOffset(): Offset = ColumnarStreamOffset(reader.initialOffset()) + + override def latestOffset(): Offset = ColumnarStreamOffset(reader.latestOffset()) + + override def planInputPartitions(start: Offset, end: Offset): Array[InputPartition] = { + reader.partitions(start.json(), end.json()) + } + + override def createReaderFactory(): PartitionReaderFactory = { + new ColumnarStreamPartitionReaderFactory(reader) + } + + override def deserializeOffset(json: String): Offset = ColumnarStreamOffset(json) + + override def commit(end: Offset): Unit = reader.commit(end.json()) + + override def stop(): Unit = reader.stop() +} + +/** + * A [[ColumnarDataSource]] only produces columnar batches, so Spark never asks for rows: the scan + * reports [[Scan.ColumnarSupportMode.SUPPORTED]]. + */ +abstract class ColumnarOnlyPartitionReaderFactory extends PartitionReaderFactory { + override def createReader(partition: InputPartition): PartitionReader[InternalRow] = { + throw SparkException.internalError("Columnar data sources only support columnar reads.") + } + + override def supportColumnarReads(partition: InputPartition): Boolean = true +} + +class ColumnarStreamPartitionReaderFactory(reader: ColumnarStreamReader) + extends ColumnarOnlyPartitionReaderFactory { + override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = { + reader.read(partition) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala new file mode 100644 index 000000000000..09ff622e7dcd --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala @@ -0,0 +1,258 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.columnar + +import scala.jdk.CollectionConverters._ + +import org.apache.arrow.vector.VectorSchemaRoot + +import org.apache.spark.{SparkContext, SparkUnsupportedOperationException} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.connector.metric.CustomTaskMetric +import org.apache.spark.sql.connector.write._ +import org.apache.spark.sql.connector.write.streaming.{StreamingDataWriterFactory, StreamingWrite} +import org.apache.spark.sql.execution.arrow.ArrowWriter +import org.apache.spark.sql.execution.streaming.continuous.ContinuousExecution +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +class ColumnarWriteBuilder( + shortName: String, + info: LogicalWriteInfo, + createDataSource: () => ColumnarDataSource) + extends WriteBuilder with SupportsTruncate { + + private var overwrite = false + + override def truncate(): WriteBuilder = { + overwrite = true + this + } + + override def build(): Write = new ColumnarWrite(shortName, info, overwrite, createDataSource) +} + +class ColumnarWrite( + shortName: String, + info: LogicalWriteInfo, + overwrite: Boolean, + createDataSource: () => ColumnarDataSource) + extends Write { + + override def description(): String = shortName + + override def toBatch: BatchWrite = { + new ColumnarBatchWrite( + () => createDataSource().writer(info.schema(), overwrite), info.schema()) + } + + override def toStreaming: StreamingWrite = { + new ColumnarStreamingWrite( + shortName, + () => createDataSource().streamWriter(info.schema(), overwrite, info.queryId()), + info.schema()) + } +} + +/** + * A batch write. The writer of the data source is created with the writer factory, right before + * the tasks of the write start: after the shuffle stages that adaptive query execution runs + * before the write, so no writer is created if they fail. The rest of the data, such as its last + * stage, is computed in the tasks of the write, and a failure there, as any later failure, aborts + * the write. So irreversible work belongs in the commit of the writer. + */ +class ColumnarBatchWrite(createWriter: () => ColumnarWriter, schema: StructType) + extends BatchWrite { + + private var writer: ColumnarWriter = _ + + override def createBatchWriterFactory(info: PhysicalWriteInfo): DataWriterFactory = { + if (writer == null) { + writer = createWriter() + } + new ColumnarDataWriterFactory(writer, ArrowBatchingOptions(schema)) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = writer.commit(messages) + + override def abort(messages: Array[WriterCommitMessage]): Unit = { + if (writer != null) { + writer.abort(messages) + } + } +} + +/** + * A streaming write. As for a [[ColumnarBatchWrite]], the writer of the data source is created + * with the writer factory, right before the tasks of the micro-batch start. A stateful + * micro-batch runs without adaptive query execution, so all its stages, as those of any write + * without it, run after the writer is created. + * + * Continuous processing is not supported: its tasks do not commit or abort the data writers of + * the last epoch when the query stops, so it fails before a writer is created. + */ +class ColumnarStreamingWrite( + description: String, + createWriter: () => ColumnarStreamWriter, + schema: StructType) + extends StreamingWrite { + + private var writer: ColumnarStreamWriter = _ + + override def createStreamingWriterFactory( + info: PhysicalWriteInfo): StreamingDataWriterFactory = { + // Continuous processing sets this property before it runs the write, and only it does. + if (SparkContext.getActive.exists( + _.getLocalProperty(ContinuousExecution.EPOCH_COORDINATOR_ID_KEY) != null)) { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_CONTINUOUS_WRITE_NOT_SUPPORTED", Map("description" -> description)) + } + if (writer == null) { + writer = createWriter() + } + new ColumnarStreamingDataWriterFactory(writer, ArrowBatchingOptions(schema)) + } + + override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + writer.commit(epochId, messages) + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + if (writer != null) { + writer.abort(epochId, messages) + } + } +} + +class ColumnarDataWriterFactory(writer: ColumnarWriter, options: ArrowBatchingOptions) + extends DataWriterFactory { + override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = { + new ArrowBatchingDataWriter(writer.createWriter(partitionId, taskId), options) + } +} + +class ColumnarStreamingDataWriterFactory( + writer: ColumnarStreamWriter, + options: ArrowBatchingOptions) + extends StreamingDataWriterFactory { + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[InternalRow] = { + new ArrowBatchingDataWriter(writer.createWriter(partitionId, taskId, epochId), options) + } +} + +/** + * How [[ArrowBatchingDataWriter]] converts rows to Arrow, captured on the driver from the + * session configuration. Like the other Arrow batches of Spark, a batch ends when it reaches + * `spark.sql.execution.arrow.maxRecordsPerBatch` rows or, after a row, holds at least + * `spark.sql.execution.arrow.maxBytesPerBatch` bytes. + */ +case class ArrowBatchingOptions( + schema: StructType, + timeZoneId: String, + maxRecordsPerBatch: Int, + maxBytesPerBatch: Long) + +object ArrowBatchingOptions { + def apply(schema: StructType): ArrowBatchingOptions = { + val conf = SQLConf.get + val maxRecords = conf.arrowMaxRecordsPerBatch + ArrowBatchingOptions( + schema, + conf.sessionLocalTimeZone, + if (maxRecords > 0) maxRecords else Int.MaxValue, + conf.arrowMaxBytesPerBatch) + } +} + +/** + * Converts the rows of a write task to Arrow-backed [[ColumnarBatch]]es for the [[DataWriter]] + * of a [[ColumnarDataSource]]. A new Arrow vector schema root is created for each batch, so a + * native writer that took ownership of the data of a batch can keep it after the call. + */ +class ArrowBatchingDataWriter(delegate: DataWriter[ColumnarBatch], options: ArrowBatchingOptions) + extends DataWriter[InternalRow] { + + private val arrowSchema = ArrowUtils.toArrowSchema( + options.schema, options.timeZoneId, errorOnDuplicatedFieldNames = true, largeVarTypes = false) + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator("columnar data source writer", 0, Long.MaxValue) + private var root: VectorSchemaRoot = _ + private var arrowWriter: ArrowWriter = _ + private var numRows = 0 + + override def write(record: InternalRow): Unit = { + if (root == null) { + root = VectorSchemaRoot.create(arrowSchema, allocator) + arrowWriter = ArrowWriter.create(root) + } + arrowWriter.write(record) + numRows += 1 + // The size is not computed when the limit is the largest possible one, as in Python UDFs. + if (numRows >= options.maxRecordsPerBatch || (options.maxBytesPerBatch < Int.MaxValue && + arrowWriter.sizeInBytes() >= options.maxBytesPerBatch)) { + flush() + } + } + + private def flush(): Unit = { + if (root != null) { + arrowWriter.finish() + val columns = root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector) + try { + delegate.write(new ColumnarBatch(columns.toArray, root.getRowCount)) + } finally { + releaseBatch() + } + } + } + + private def releaseBatch(): Unit = { + if (root != null) { + root.close() + root = null + arrowWriter = null + numRows = 0 + } + } + + override def commit(): WriterCommitMessage = { + flush() + delegate.commit() + } + + override def abort(): Unit = { + releaseBatch() + delegate.abort() + } + + override def close(): Unit = { + try { + delegate.close() + } finally { + releaseBatch() + Utils.tryLogNonFatalError(allocator.close()) + } + } + + override def currentMetricsValues(): Array[CustomTaskMetric] = delegate.currentMetricsValues() +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala new file mode 100644 index 000000000000..3bf04b67ca98 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -0,0 +1,1145 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.lang.ref.{Cleaner, WeakReference} +import java.util +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import com.fasterxml.jackson.core.JsonProcessingException +import com.fasterxml.jackson.databind.DeserializationFeature +import com.fasterxml.jackson.databind.json.JsonMapper +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, CDataDictionaryProvider, Data} +import org.apache.arrow.vector.{FieldVector, ValueVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} + +import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException, TaskContext} +import org.apache.spark.internal.Logging +import org.apache.spark.sql.connector.catalog.TableCapability +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} +import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ +import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.unsafe.Platform +import org.apache.spark.util.Utils + +/** + * A [[ColumnarDataSource]] implemented by a native library, through the C interface in + * `org/apache/spark/sql/datasource/spark_datasource.h`. + * + * @param location where the library is on the driver. The readers and the writers distribute it + * to the executors of the active session when they are created. + * @param name the name of the data source, as listed in the manifest of its package or in lower + * case for an installed library + */ +class NativeDataSource( + location: NativeLibraryLocation, + name: String, + options: CaseInsensitiveStringMap) + extends ColumnarDataSource { + import NativeCalls._ + + override def schema(): StructType = { + val library = NativeLibraries.get(location) + withDataSource(library, "infer the schema of") { dataSource => + val schema = call(name, "infer the schema of") { + optional(NativeArrow.importSchema(address => library.schema(dataSource.handle, address))) + }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) + NativeArrow.checkSupported(schema, name, "infer the schema of") + schema + } + } + + /** + * Returns the batch writes that the data source declares, and always the reads and the + * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source, + * which then fails with errors that do not say why. A scan or a streaming write that the data + * source does not declare fails when Spark plans it instead, on the driver. + */ + override def capabilities(): util.Set[TableCapability] = { + val library = NativeLibraries.get(location) + val declared = withDataSource(library, "check the capabilities of")(_.capabilities) + val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE) + if ((declared & NativeLibrary.BATCH_WRITE) != 0) { + capabilities.add(BATCH_WRITE) + } + if ((declared & NativeLibrary.TRUNCATE) != 0) { + capabilities.add(TRUNCATE) + } + capabilities + } + + override def batchPlanner(schema: StructType): ColumnarBatchPlanner = { + NativeArrow.checkSupported(schema, name, "plan a scan of") + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ, + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createReader(dataSource.handle, address) + } + } + new NativeBatchPlanner( + library, executorLocation, name, schema, new NativeHandle(handle, library.closeReader)) + } + + override def streamReader(schema: StructType): ColumnarStreamReader = { + NativeArrow.checkSupported(schema, name, "plan a scan of") + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", NativeLibrary.MICRO_BATCH_READ, + "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createStreamReader(dataSource.handle, address) + } + } + new NativeDataSourceStreamReader( + executorLocation, name, schema, new NativeHandle(handle, library.closeStreamReader)) + } + + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = new NativeHandle( + createWriter(library, schema, overwrite, queryId = null), library.closeWriter) + new NativeDataSourceWriter(executorLocation, name, handle, serializeWriter(library, handle)) + } + + override def streamWriter( + schema: StructType, + overwrite: Boolean, + queryId: String): ColumnarStreamWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = + new NativeHandle(createWriter(library, schema, overwrite, queryId), library.closeWriter) + val writer = new NativeDataSourceStreamWriter( + executorLocation, name, handle, serializeWriter(library, handle)) + NativeStreamWriters.register(queryId, handle) + writer + } + + /** Runs `f` with a new native data source, which is closed afterwards. */ + private def withDataSource[T](library: NativeLibrary, action: String)( + f: NativeLibrary.DataSourceHandle => T): T = { + // Checked first, so that a data source is never created without a way to release it. + requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE)) + // The keys of a case-insensitive map are in lower case. + val entries = options.entrySet().asScala.toArray + entries.find(e => e.getKey.contains('\u0000') || e.getValue.contains('\u0000')).foreach { e => + throw error(name, "create", s"The option ${NativeArrow.printable(e.getKey)} contains the " + + "character U+0000, which Spark cannot pass to a library.") + } + val dataSource = call(name, "create") { + library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue)) + } + Utils.tryWithSafeFinally(f(dataSource)) { + call(name, action)(library.closeDataSource(dataSource.handle)) + } + } + + /** + * Creates a reader or a writer with `f`, which fails if the data source does not declare the + * capability, or if the library does not implement the functions that it requires. + */ + private def create( + library: NativeLibrary, + action: String, + capability: Long, + unsupportedError: String, + functions: Seq[NativeLibrary.Function])( + f: NativeLibrary.DataSourceHandle => Long): Long = { + withDataSource(library, action) { dataSource => + if ((dataSource.capabilities & capability) == 0) { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) + } + requireFunctions(library, action, functions) + call(name, action)(f(dataSource)) + } + } + + /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */ + private def createWriter( + library: NativeLibrary, + schema: StructType, + overwrite: Boolean, + queryId: String): Long = { + val action = "plan a write to" + NativeArrow.checkSupported(schema, name, action) + val (capability, unsupportedError, overwriteMode) = if (queryId == null) { + (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", + "the save mode overwrite") + } else { + (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", + "the output mode complete") + } + create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) { + dataSource => + if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) { + throw error(name, action, s"The data source does not support $overwriteMode.") + } + val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND + NativeArrow.exportSchema(schema) { address => + library.createWriter(dataSource.handle, address, mode, queryId) + } + } + } + + private def requireFunctions( + library: NativeLibrary, + action: String, + functions: Seq[NativeLibrary.Function]): Unit = { + functions.find(!library.hasFunction(_)).foreach { function => + throw error(name, action, s"The library does not implement ${function.fieldName}.") + } + } + + private def serializeWriter(library: NativeLibrary, handle: NativeHandle): Array[Byte] = { + try { + bytesOrEmpty(call(name, "plan a write to")(optional(library.serializeWriter(handle.get)))) + } catch { + case e: Throwable => + handle.close() + throw e + } + } +} + +object NativeDataSource { + import NativeLibrary.Function._ + + // The functions that the operations of a data source require, besides create and close. The + // others are optional. They are checked on the driver, before the operation starts. + private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ) + private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET, + STREAM_PARTITIONS, CLOSE_STREAM_READER, READ) + private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER, + WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER) +} + +/** + * Plans a batch scan of a native data source, on the driver, with a reader of the library. The + * reader is released once the scan is planned, or when the planner becomes unreachable if the + * scan is never planned: the executors read the partitions with the state that it serialized. + * + * @param executorLocation where the executors find the library + */ +class NativeBatchPlanner( + library: NativeLibrary, + executorLocation: NativeLibraryLocation, + name: String, + schema: StructType, + handle: NativeHandle) + extends ColumnarBatchPlanner { + import NativeCalls._ + + private var readSchema: StructType = schema + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + // Only the predicates that can be expressed in JSON are passed to the library. + val encoded = predicates.zipWithIndex.flatMap { case (predicate, i) => + NativePredicates.toJson(predicate, schema).map(_ -> i) + } + val used = if (encoded.isEmpty) { + Set.empty[Int] + } else { + // How the reader uses each predicate: whether it uses it completely or only to skip data + // does not matter, since Spark evaluates all the predicates again after reading. + call(name, "plan a scan of") { + optional(library.pushPredicates(handle.get, encoded.map(_._1))) + }.map { uses => + encoded.map(_._2).zip(uses).collect { + case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i + }.toSet + }.getOrElse(Set.empty[Int]) + } + predicates.zipWithIndex.collect { case (predicate, i) if used.contains(i) => predicate } + } + + override def pushLimit(limit: Int): Boolean = { + call(name, "plan a scan of")(optional(library.pushLimit(handle.get, limit))).getOrElse(false) + } + + override def pruneColumns(requiredSchema: StructType): Boolean = { + val pruned = call(name, "plan a scan of") { + optional(library.pruneColumns(handle.get, requiredSchema.fieldNames)) + }.getOrElse(false) + if (pruned) { + readSchema = requiredSchema + } + pruned + } + + override def plan(): ColumnarPlannedBatch = { + try { + val partitions = call(name, "plan a scan of")(library.partitions(handle.get)) + val readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeReader(handle.get)))) + ColumnarPlannedBatch( + NativeInputPartition.from(partitions), + new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema)) + } finally { + // The executors only need the serialized state. + handle.close() + } + } +} + +/** Reads the partitions of a planned batch scan of a native data source, on the executors. */ +class NativePartitionReaderFactory( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + readSchema: StructType) + extends ColumnarOnlyPartitionReaderFactory { + override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, readSchema) + } +} + +/** Reads a native data source for a micro-batch streaming scan. */ +class NativeDataSourceStreamReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarStreamReader { + import NativeCalls._ + + // Set on the driver whenever the partitions of a micro-batch are planned, and serialized to + // the executors with the partitions of that micro-batch. + private var readerState: Array[Byte] = Array.emptyByteArray + + @transient private lazy val library = NativeLibraries.get(location) + + override def initialOffset(): String = { + checkOffset("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get))) + } + + override def latestOffset(): String = { + checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get))) + } + + /** + * Checks an offset returned by the given function of the library: it must be JSON on a single + * line, because the offset log of the checkpoint stores each offset as a line. + */ + private def checkOffset(function: String, offset: String): String = { + if (offset.exists(c => c == '\n' || c == '\r')) { + throw error(name, "plan a scan of", s"$function returned an offset with a line break, " + + "but an offset must be JSON on a single line.") + } + val valid = try { + val json = NativeDataSourceStreamReader.mapper.readTree(offset) + json != null && !json.isMissingNode + } catch { + case _: JsonProcessingException => false + } + if (!valid) { + throw error(name, "plan a scan of", s"$function returned an offset that is not valid " + + s"JSON: $offset") + } + offset + } + + override def partitions(start: String, end: String): Array[InputPartition] = { + val partitions = call(name, "plan a scan of") { + library.streamPartitions(handle.get, start, end) + } + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeStreamReader(handle.get)))) + NativeInputPartition.from(partitions) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, schema) + } + + override def commit(end: String): Unit = { + call(name, "commit an offset of")(optional(library.commitOffset(handle.get, end))) + } + + override def stop(): Unit = handle.close() +} + +object NativeDataSourceStreamReader { + // Parses the offsets of the libraries to check them, and rejects anything after the JSON value. + private val mapper = + JsonMapper.builder().enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS).build() +} + +/** Writes to a native data source for a batch write. */ +class NativeDataSourceWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId = -1L) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + // Not closed if the commit fails: Spark then aborts the write, which closes it. + handle.close() + } + + override def abort(messages: Array[WriterCommitMessage]): Unit = { + Utils.tryWithSafeFinally { + call(name, "abort a write to") { + library.abort(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + }(handle.close()) + } +} + +/** + * Writes to a native data source for a streaming write. Its handle is closed when the streaming + * query creates its next stream writer, see [[NativeStreamWriters]]. + */ +class NativeDataSourceStreamWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarStreamWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId) + } + + override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "abort a write to") { + library.abort(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } +} + +/** A partition planned by a native library. */ +case class NativeInputPartition(index: Int, bytes: Array[Byte]) extends InputPartition + +object NativeInputPartition { + def from(partitions: Array[Array[Byte]]): Array[InputPartition] = { + Option(partitions).getOrElse(Array.empty).zipWithIndex.map { case (bytes, i) => + NativeInputPartition(i, Option(bytes).getOrElse(Array.emptyByteArray)) + } + } +} + +/** The message returned by the data writer of a task of a native library. */ +case class NativeWriterCommitMessage(bytes: Array[Byte]) extends WriterCommitMessage + +object NativeWriterCommitMessage { + def toBytes(messages: Array[WriterCommitMessage]): Array[Array[Byte]] = messages.map { + case message: NativeWriterCommitMessage => message.bytes + case null => null + case other => + throw SparkException.internalError(s"Unexpected commit message: ${other.getClass.getName}") + } +} + +/** + * Reads a partition of a native data source: each array of the Arrow stream exported by the + * library is checked, and imported into the Arrow vectors that back the returned batch. + */ +class NativePartitionReader( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + partition: Array[Byte], + schema: StructType) + extends PartitionReader[ColumnarBatch] { + import NativeCalls._ + + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) + // Read here rather than with `Data.importArrayStream`, to check each array before importing it. + private val stream = ArrowArrayStream.allocateNew(allocator) + private val dictionaries = new CDataDictionaryProvider() + // The task that reads the partition, which the library may use until the stream is released. + private var task: NativeLibrary.Task = _ + private var root: VectorSchemaRoot = _ + private var batch: ColumnarBatch = _ + + try { + openStream() + root = VectorSchemaRoot.create(importSchema(), allocator) + checkSchema() + batch = new ColumnarBatch(root.getFieldVectors.asScala.map(columnVector).toArray) + } catch { + case e: Throwable => + close() + throw e + } + + private def openStream(): Unit = { + val library = NativeLibraries.get(location) + task = NativeTask.create() + call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress())) + } + + private def importSchema(): Schema = { + // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash + // the JVM. + val snapshot = stream.snapshot() + if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 || + snapshot.get_last_error == 0) { + throw error(name, "read", "The library did not initialize the stream, or a callback of " + + "the stream is NULL.") + } + val arrowSchema = ArrowSchema.allocateNew(allocator) + try { + call(name, "read")(stream.getSchema(arrowSchema)) + val imported = call(name, "read")(Data.importSchema(allocator, arrowSchema, dictionaries)) + new Schema( + imported.getFields.asScala.map(withMapEntryNames).asJava, imported.getCustomMetadata) + } finally { + arrowSchema.close() + } + } + + /** + * Names the fields of the entries of the maps `key` and `value`, which ArrowColumnVector looks + * up by name. Arrow does not fix these names: arrow-rs names them `keys` and `values`, for + * example. The arrays are imported by position, so the names do not change how they are read. + */ + private def withMapEntryNames(field: Field): Field = { + val children = field.getChildren.asScala.map(withMapEntryNames).toSeq + val renamed = (field.getType, children) match { + case (_: ArrowType.Map, Seq(entries)) if entries.getChildren.size == 2 => + val Seq(key, value) = entries.getChildren.asScala.toSeq + Seq(new Field(entries.getName, entries.getFieldType, Seq( + new Field(MapVector.KEY_NAME, key.getFieldType, key.getChildren), + new Field(MapVector.VALUE_NAME, value.getFieldType, value.getChildren)).asJava)) + case _ => children + } + new Field(field.getName, field.getFieldType, renamed.asJava) + } + + /** + * Checks that the data agrees with the nullability of the schema to read: a column, a field, + * an element or a map value that it declares not nullable has no null, except where its parent + * is null. Spark relies on the nullability, for example to evaluate `isnull` of such a column to + * false without reading it. + */ + private def checkNulls(): Unit = { + root.getFieldVectors.asScala.zip(schema.fields).foreach { case (vector, field) => + checkNulls(vector, field.dataType, field.nullable, field.name, _ => true) + } + } + + /** `referenced` tells whether a value belongs to a value of the parent that is not null. */ + private def checkNulls( + vector: ValueVector, + dataType: DataType, + nullable: Boolean, + column: String, + referenced: Int => Boolean): Unit = { + if (!nullable && vector.getNullCount > 0 && + (0 until vector.getValueCount).exists(i => vector.isNull(i) && referenced(i))) { + throw error(name, "read", s"The column $column has a null, but the schema to read " + + "declares it not nullable.") + } + val valid = (i: Int) => !vector.isNull(i) && referenced(i) + (dataType, vector) match { + case (struct: StructType, structVector: StructVector) => + struct.fields.zipWithIndex.foreach { case (field, i) => + checkNulls(structVector.getChildByOrdinal(i), field.dataType, field.nullable, + s"$column.${field.name}", valid) + } + case (map: MapType, mapVector: MapVector) => + val entries = mapVector.getDataVector.asInstanceOf[StructVector] + val inMaps = elements(mapVector, valid) + // Spark does not allow null keys. + checkNulls(entries.getChildByOrdinal(0), map.keyType, nullable = false, s"$column.key", + inMaps) + checkNulls(entries.getChildByOrdinal(1), map.valueType, map.valueContainsNull, + s"$column.value", inMaps) + case (array: ArrayType, list: ListVector) => + checkNulls(list.getDataVector, array.elementType, array.containsNull, s"$column.element", + elements(list, valid)) + case _ => + } + } + + /** Whether each element of a list belongs to a list that `valid` accepts. */ + private def elements(list: ListVector, valid: Int => Boolean): Int => Boolean = { + lazy val referenced = { + val bits = new java.util.BitSet(list.getDataVector.getValueCount) + (0 until list.getValueCount).foreach { i => + if (valid(i)) bits.set(list.getElementStartIndex(i), list.getElementEndIndex(i)) + } + bits + } + i => referenced.get(i) + } + + private def checkSchema(): Unit = { + def isDictionaryEncoded(field: Field): Boolean = { + field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) + } + root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + + "is not supported.") + } + // The Arrow types that convert to Spark types, but that Spark does not read: ArrowColumnVector + // cannot read decimals of other widths than 128 bits, and the lengths of the elements of a + // list view cannot be checked without a pass over all its lists (see checkArrays). + def unsupportedType(field: Field): Option[ArrowType] = field.getType match { + case decimal: ArrowType.Decimal if decimal.getBitWidth != 128 => Some(decimal) + case listView: ArrowType.ListView => Some(listView) + case _ => field.getChildren.asScala.view.flatMap(unsupportedType).headOption + } + root.getSchema.getFields.asScala.foreach { field => + unsupportedType(field).foreach { arrowType => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"$arrowType, which is not supported.") + } + } + val actual = StructType(root.getSchema.getFields.asScala.map { field => + val dataType = try { + ArrowUtils.fromArrowField(field) + } catch { + case e: SparkUnsupportedOperationException if e.getCondition == "UNSUPPORTED_ARROWTYPE" => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"${e.getMessageParameters.get("typeName")}, which is not supported.") + } + StructField(field.getName, dataType, field.isNullable) + }.toSeq) + // The schema to read as the library got it: Arrow does not carry everything, such as the + // collations of strings, the fields of intervals and user-defined types, which still apply to + // the data that Spark reads. + val expected = ArrowUtils.fromArrowSchema(ArrowUtils.toArrowSchema( + schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = false)) + val matches = actual.length == expected.length && actual.zip(expected).forall { + case (actualField, field) => + actualField.name == field.name && + DataType.equalsIgnoreNullability(actualField.dataType, field.dataType) + } + if (!matches) { + throw error(name, "read", s"The data has the schema ${actual.toDDL}, but the expected " + + s"schema is ${expected.toDDL}.") + } + } + + private def columnVector(vector: FieldVector): ColumnVector = { + try { + new ReusedArrowColumnVector(vector) + } catch { + case NonFatal(e) => + throw QueryExecutionErrors.nativeDataSourceError("read", name, + s"Spark cannot read the column ${vector.getName} of the Arrow type " + + s"${vector.getField.getType}: ${message(e)}", e) + } + } + + override def next(): Boolean = { + val array = ArrowArray.allocateNew(allocator) + try { + call(name, "read")(stream.getNext(array)) + val snapshot = array.snapshot() + if (snapshot.release == 0) { + // The end of the stream. + false + } else { + try { + checkArrays(snapshot) + } catch { + case e: Throwable => + array.release() + throw e + } + call(name, "read")(Data.importIntoVectorSchemaRoot(allocator, array, root, dictionaries)) + checkNulls() + batch.setNumRows(root.getRowCount) + true + } + } finally { + array.close() + } + } + + /** + * Checks the array of a batch, and the arrays of its columns at every level, before importing + * them. Arrow Java trusts the arrays, and would read them wrongly otherwise: + * - An array must have an offset of 0. Arrow Java ignores the offsets of the arrays it imports, + * so it would read a sliced array from the start of its buffers. + * - An array must have at least as many values as its parent refers to: a column as many as + * the batch has rows, a field as many as its struct has values, and the elements of a list or + * a map as many as its last offset says. Arrow Java would extend a shorter array with nulls. + * Only the lengths that the library exported show this: the import sets the length of each + * array to what its parent refers to. + * - An array without a validity buffer has a null_count of 0, as the C data interface requires: + * Arrow Java fails to import it otherwise. + * - The batch itself has no null rows. Arrow Java ignores the validity of the struct of the + * batch, so it would read a null row as a row of the values of its columns. As in Arrow C++, + * which also rejects them, the nulls are counted when the library did not count them. + * The buffers are not checked, since the C data interface does not carry their sizes: Arrow + * Java takes them from the lengths, so the library must export buffers that hold every value. + */ + private def checkArrays(batch: ArrowArray.Snapshot): Unit = { + // The arrays of the children, one for each field of the type: Arrow Java only checks their + // number when an array has children, and imports the missing columns of a batch as nulls. + def children( + array: ArrowArray.Snapshot, + fields: Int, + owner: String, + expected: String): Seq[ArrowArray.Snapshot] = { + if (array.n_children != fields || (fields > 0 && array.children == 0)) { + throw error(name, "read", s"$owner has ${array.n_children} child arrays, but $expected.") + } + (0 until fields).map { i => + val child = Platform.getLong(null, array.children + 8L * i) + if (child == 0) { + throw error(name, "read", s"$owner has a NULL child array.") + } + ArrowArray.wrap(child).snapshot() + } + } + // The last offset of a list or a map, in its second buffer. Arrow Java rejects the array if + // the buffer is missing. + def lastOffset(array: ArrowArray.Snapshot): Long = { + val offsets = if (array.length > 0 && array.n_buffers >= 2) { + Platform.getLong(null, array.buffers + 8L) + } else { + 0L + } + if (offsets == 0) 0 else Platform.getInt(null, offsets + 4L * array.length) + } + // The validity buffer of an array, its first buffer, or 0 if it has none. + def validity(array: ArrowArray.Snapshot): Long = { + if (array.n_buffers >= 1 && array.buffers != 0) Platform.getLong(null, array.buffers) else 0L + } + // The C data interface only allows an array of values without a validity buffer if its + // null_count is 0. Arrow Java takes another null_count as a mix of nulls and values, and then + // fails to load the missing buffer. + def checkValidity(array: ArrowArray.Snapshot, owner: String): Unit = { + if (array.null_count != 0 && array.length > 0 && validity(array) == 0) { + throw error(name, "read", s"$owner has a null_count of ${array.null_count} but no " + + "validity buffer, which the C data interface only allows with a null_count of 0.") + } + } + // The number of values that an array refers to in each of its children. + def referenced(array: ArrowArray.Snapshot, arrowType: ArrowType): Long = arrowType match { + case _: ArrowType.Struct => array.length + case _: ArrowType.List | _: ArrowType.Map => lastOffset(array) + // The other types that Spark reads have no children. List views are rejected before. + case _ => 0 + } + def checkColumn(array: ArrowArray.Snapshot, field: Field, column: String): Unit = { + if (array.offset != 0) { + throw error(name, "read", s"The column $column has an offset of ${array.offset}, but " + + "Spark only reads arrays with an offset of 0. Copy sliced arrays before exporting them.") + } + checkValidity(array, s"The column $column") + val length = referenced(array, field.getType) + val fields = field.getChildren.size + children(array, fields, s"An array of the column $column", s"its type has $fields") + .zip(field.getChildren.asScala).foreach { case (child, childField) => + if (child.length < length) { + throw error(name, "read", s"The column $column contains an array with a length of " + + s"${child.length}, but its parent refers to $length values.") + } + checkColumn(child, childField, column) + } + } + if (batch.offset != 0) { + throw error(name, "read", s"The batch has an offset of ${batch.offset}, but Spark only " + + "reads arrays with an offset of 0. Copy sliced arrays before exporting them.") + } + checkValidity(batch, "The batch") + val nullRows = batch.null_count match { + // The library did not count the nulls, so they are counted from the validity buffer, which + // holds a bit for each row since the offset is 0. A batch without one has no rows here. + case -1 => + val bitmap = validity(batch) + var nulls = 0L + var i = 0L + while (bitmap != 0 && i < batch.length) { + val bits = math.min(8L, batch.length - i).toInt + val valid = Platform.getByte(null, bitmap + i / 8) & ((1 << bits) - 1) + nulls += bits - Integer.bitCount(valid) + i += 8 + } + nulls + case n if n < -1 => + throw error(name, "read", s"The batch has an invalid null_count of $n.") + case n => n + } + if (nullRows != 0) { + throw error(name, "read", s"The struct of the batch marks $nullRows of its rows as null, " + + "but Spark only reads batches without null rows: mark the values of the columns as null " + + "instead.") + } + val columns = root.getSchema.getFields.size + children(batch, columns, "The batch", s"the schema has $columns columns") + .zip(root.getSchema.getFields.asScala).foreach { case (column, field) => + if (column.length < batch.length) { + throw error(name, "read", s"The column ${field.getName} has a length of " + + s"${column.length}, but the batch has a length of ${batch.length}.") + } + checkColumn(column, field, field.getName) + } + } + + override def get(): ColumnarBatch = batch + + override def close(): Unit = { + Utils.tryWithSafeFinally { + Utils.tryWithSafeFinally { + if (root != null) { + root.close() + } + dictionaries.close() + } { + // Releases the stream of the library, unless it was not exported, even if closing the + // vectors failed: the library may use the task until the stream is released. + stream.release() + stream.close() + } + } { + if (task != null) { + task.close() + } + Utils.tryLogNonFatalError(allocator.close()) + } + } +} + +/** + * The Arrow column vector of a column of a [[NativePartitionReader]], which imports all its batches + * into the same Arrow vectors. It is not freeable between batches: Spark must not close it when it + * moves to the next batch, as `ColumnarToRowExec` does with the column vectors that are freeable, + * because closing a struct column also drops its child column vectors. The reader closes the + * Arrow vectors at the end. + */ +private class ReusedArrowColumnVector(vector: ValueVector) extends ArrowColumnVector(vector) { + override def closeIfFreeable(): Unit = {} +} + +/** Writes the batches of a task to the data writer of a native library. */ +class NativeDataWriter( + location: NativeLibraryLocation, + name: String, + writerState: Array[Byte], + partitionId: Int, + taskId: Long, + epochId: Long) + extends DataWriter[ColumnarBatch] { + import NativeCalls._ + + private val library = NativeLibraries.get(location) + // The task, which the library may use until the data writer is closed. + private val task = NativeTask.create(Some(partitionId), Some(taskId)) + private val handle: Long = try { + call(name, "write to")(library.createDataWriter(task, writerState, epochId)) + } catch { + case e: Throwable => + task.close() + throw e + } + // Whether the data writer was committed or aborted. It stays false if the library fails to + // commit, so that Spark aborts the data writer, but becomes true as soon as the library commits, + // even if Spark then fails the commit, for example because of an invalid message. + private var finished = false + private var closed = false + + override def write(batch: ColumnarBatch): Unit = { + if (finished || closed) { + throw SparkException.internalError("The data writer was already committed or aborted.") + } + val vectors = (0 until batch.numCols()).map { i => + batch.column(i) match { + case vector: ArrowColumnVector => vector.getValueVector.asInstanceOf[FieldVector] + case vector => throw SparkException.internalError( + s"Expected an Arrow column vector, but got ${vector.getClass.getName}.") + } + } + val root = new VectorSchemaRoot( + vectors.map(_.getField).asJava, vectors.asJava, batch.numRows()) + val array = ArrowArray.allocateNew(ArrowUtils.rootAllocator) + val arrowSchema = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) + try { + Data.exportVectorSchemaRoot(ArrowUtils.rootAllocator, root, null, array, arrowSchema) + call(name, "write to") { + library.write(handle, array.memoryAddress(), arrowSchema.memoryAddress()) + } + } finally { + // Releases the batch unless the library moved it to take ownership. + array.release() + array.close() + arrowSchema.release() + arrowSchema.close() + } + } + + override def commit(): WriterCommitMessage = { + if (finished || closed) { + throw SparkException.internalError("The data writer was already committed or aborted.") + } + val message = call(name, "write to")(library.commitDataWriter(handle, () => finished = true)) + NativeWriterCommitMessage(bytesOrEmpty(Some(message))) + } + + override def abort(): Unit = { + if (!finished && !closed) { + finished = true + call(name, "write to")(library.abortDataWriter(handle)) + } + } + + override def close(): Unit = { + if (!closed) { + closed = true + Utils.tryWithSafeFinally { + call(name, "write to")(library.closeDataWriter(handle)) + }(task.close()) + } + } +} + +/** Creates the `SparkDsTaskContext` of the running task. */ +private[ffi] object NativeTask { + /** + * @param partitionId the partition that the task reads or writes, by default the partition of + * the task + * @param taskAttemptId the ID of the task attempt, by default the one of the task + */ + def create( + partitionId: Option[Int] = None, + taskAttemptId: Option[Long] = None): NativeLibrary.Task = { + val context = Option(TaskContext.get()) + new NativeLibrary.Task( + context.map(_.stageId()).getOrElse(-1), + context.map(_.stageAttemptNumber()).getOrElse(-1), + partitionId.orElse(context.map(_.partitionId())).getOrElse(-1), + context.map(_.attemptNumber()).getOrElse(0), + taskAttemptId.orElse(context.map(_.taskAttemptId())).getOrElse(-1L), + () => context.exists(_.isInterrupted())) + } +} + +/** + * A handle to a native object. It is released with `close()`, or once it becomes unreachable if + * it was never closed. + */ +final class NativeHandle(handle: Long, release: Long => Unit) { + private val state = new NativeHandle.State(handle, release) + private val cleanable = NativeHandle.cleaner.register(this, state) + + def get: Long = { + val value = state.handle.get() + if (value == 0) { + throw SparkException.internalError("The native handle was already released.") + } + value + } + + def close(): Unit = cleanable.clean() +} + +object NativeHandle extends Logging { + private val cleaner = Cleaner.create() + + // Must not refer to the handle, which would then never become unreachable. + private class State(value: Long, release: Long => Unit) extends Runnable { + val handle = new AtomicLong(value) + + override def run(): Unit = { + val value = handle.getAndSet(0) + if (value != 0) { + try { + release(value) + } catch { + case NonFatal(e) => + logWarning("Failed to release a native data source handle.", e) + } + } + } + } +} + +/** + * The stream writers of the streaming queries, by query ID. The Data Source V2 streaming write + * has no hook for when the query stops, but a micro-batch streaming query creates a new stream + * writer for each micro-batch, after the previous micro-batch committed or aborted its write, and + * a restarted query creates a new one. So the previous stream writer of a query is closed when + * the query creates a new one, and the last one when it becomes unreachable. + */ +object NativeStreamWriters { + // Weak, so that the last stream writer of a query can still become unreachable. + private val writers = new ConcurrentHashMap[String, WeakReference[NativeHandle]]() + + /** Registers the new stream writer of the query, and closes its previous one. */ + def register(queryId: String, handle: NativeHandle): Unit = { + // Forgets the queries whose last stream writer was released. + writers.values().removeIf(_.get() == null) + Option(writers.put(queryId, new WeakReference(handle))).flatMap(r => Option(r.get())) + .foreach(_.close()) + } +} + +/** Helpers to call native data source libraries. */ +object NativeCalls { + /** + * Calls a native library, and rethrows the errors it reports as NATIVE_DATA_SOURCE_ERROR, as + * well as the link errors, such as when Arrow cannot load its own JNI library. + */ + def call[T](name: String, action: String)(f: => T): T = { + try { + f + } catch { + case e: Throwable with SparkThrowable => throw e + case e: NativeLibrary.MissingFunctionException => + throw QueryExecutionErrors.nativeDataSourceError( + action, name, s"The library does not implement ${e.getMessage}.", e) + case e @ (NonFatal(_) | _: LinkageError) => + throw QueryExecutionErrors.nativeDataSourceError(action, name, message(e), e) + } + } + + // How Arrow reports the error message of an Arrow C stream. + private val StreamError = """(?s)\[errno \d+\] CDataJniException\{errno=\d+, message=(.*)\}""".r + + private[ffi] def message(e: Throwable): String = { + Option(e.getMessage).getOrElse(e.toString) match { + case StreamError(message) => message + case message => message + } + } + + /** Calls a function that the library may not implement. Returns None if it does not. */ + def optional[T](f: => T): Option[T] = { + try { + Some(f) + } catch { + case _: NativeLibrary.MissingFunctionException => None + } + } + + /** Returns a NATIVE_DATA_SOURCE_ERROR for an error detected by Spark. */ + def error(name: String, action: String, msg: String): Throwable = { + QueryExecutionErrors.nativeDataSourceError(action, name, msg, null) + } + + def bytesOrEmpty(bytes: Option[Array[Byte]]): Array[Byte] = { + bytes.flatMap(Option(_)).getOrElse(Array.emptyByteArray) + } +} + +/** Exchanges schemas with native libraries through the Arrow C data interface. */ +object NativeArrow { + /** + * Checks that the schema only has the types of version 1 of the interface, which + * `spark_datasource.h` lists with their Arrow types, and no names with the character U+0000, + * which a NUL-terminated string cannot hold. + */ + def checkSupported(schema: StructType, name: String, action: String): Unit = { + def checkType(dataType: DataType, column: String): Unit = dataType match { + case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | + _: DecimalType | _: StringType | BinaryType | DateType | TimestampType | + TimestampNTZType => + case array: ArrayType => checkType(array.elementType, s"$column.element") + case map: MapType => + checkType(map.keyType, s"$column.key") + checkType(map.valueType, s"$column.value") + case struct: StructType => checkFields(struct, s"$column.") + case other => + throw NativeCalls.error(name, action, s"The column $column has the type ${other.sql}, " + + "which native data sources do not support.") + } + def checkFields(struct: StructType, prefix: String): Unit = struct.fields.foreach { field => + if (field.name.contains('\u0000')) { + throw NativeCalls.error(name, action, s"The name of the column " + + s"${printable(prefix + field.name)} contains the character U+0000, which Spark cannot " + + "pass to a library.") + } + checkType(field.dataType, prefix + field.name) + } + checkFields(schema, "") + } + + /** Shows the character U+0000 in a message. */ + def printable(value: String): String = value.replace("\u0000", "\\u0000") + + /** Exports the schema as an `ArrowSchema`, and calls `f` with its address. */ + def exportSchema[T](schema: StructType)(f: Long => T): T = { + val arrowSchema = ArrowUtils.toArrowSchema( + schema, SQLConf.get.sessionLocalTimeZone, errorOnDuplicatedFieldNames = true, + largeVarTypes = false) + val struct = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) + try { + Data.exportSchema(ArrowUtils.rootAllocator, arrowSchema, null, struct) + f(struct.memoryAddress()) + } finally { + // Releases the schema unless the library moved it to take ownership. + struct.release() + struct.close() + } + } + + /** Calls `f` with the address of an `ArrowSchema` to initialize, and imports it. */ + def importSchema(f: Long => Unit): StructType = { + val struct = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) + var imported = false + try { + f(struct.memoryAddress()) + imported = true + // Releases and closes the struct. + ArrowUtils.fromArrowSchema(Data.importSchema(ArrowUtils.rootAllocator, struct, null)) + } finally { + if (!imported) { + struct.release() + struct.close() + } + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala new file mode 100644 index 000000000000..cc49dbc0cd73 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -0,0 +1,502 @@ +/* + * 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, Paths, StandardCopyOption} +import java.nio.file.attribute.BasicFileAttributes +import java.security.MessageDigest +import java.util.{HexFormat, Locale} +import java.util.concurrent.ConcurrentHashMap +import java.util.zip.ZipFile + +import scala.jdk.CollectionConverters._ +import scala.util.{Failure, Try} + +import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} +import com.google.common.cache.{Cache, CacheBuilder} +import org.apache.arrow.c.jni.JniLoader + +import org.apache.spark.SPARK_VERSION +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.{PATH, VERSION_NUM} +import org.apache.spark.sql.Artifact +import org.apache.spark.sql.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 oldest version of the native data source interface that the libraries + * implement. Spark uses a package whose version it supports, or that is older + * than the newest version that it supports, and `spark_ds_init` chooses the + * version of the table. Spark ignores the fields that it does not know. + * @param dataSources the names of the data sources that the libraries implement + * @param libraries the path of the library in the package for each platform, as returned by + * [[NativePlatform.current]] + */ +case class NativeDataSourceManifest( + abiVersion: Int, + dataSources: Seq[String], + libraries: Map[String, String]) + +/** + * Where the library of a native data source is: in a [[NativeDataSourcePackage]], or installed in + * the native library path, as an [[InstalledNativeLibrary]]. It is sent to the executors, which + * load the library with [[NativeLibraries.get]]. + */ +sealed trait NativeLibraryLocation extends Serializable + +/** + * A native data source package: a zip file that contains a [[NativeDataSourceManifest]] and the + * libraries it lists. The SHA-256 checksum of the file identifies the package, so that an + * executor can check that it loads the same package as the driver. + * + * @param path the path of the package file on the driver + * @param artifactUUID the UUID of the session whose artifacts contain the package, if the session + * is isolated. Executors store the artifacts of such a session separately. + * @param distributedFileName the name of the copy of the package in the artifacts of the session, + * if the package is not itself an artifact of the session + */ +case class NativeDataSourcePackage( + path: String, + sha256: String, + manifest: NativeDataSourceManifest, + artifactUUID: Option[String] = None, + distributedFileName: Option[String] = None) extends NativeLibraryLocation { + def fileName: String = new File(path).getName +} + +/** + * A package file that is not a valid native data source package. + * + * @param error the error that explains why the package is invalid + * @param dataSources the data sources listed in the manifest of the package, if it can be read + */ +case class InvalidNativeDataSourcePackage(error: Throwable, dataSources: Seq[String]) + +/** + * A version of a file, to cache what is read from it until it changes: the file is identified by + * its canonical path, and its version by its length, its modification time and its file key, such + * as the inode on Linux, which changes when another file replaces it, as `make install` does. + */ +private[ffi] case class FileVersion(path: String, length: Long, lastModified: Long, key: AnyRef) + +private[ffi] object FileVersion { + def apply(canonical: File): FileVersion = { + val attributes = Files.readAttributes(canonical.toPath, classOf[BasicFileAttributes]) + FileVersion(canonical.getPath, attributes.size(), attributes.lastModifiedTime().toMillis, + attributes.fileKey()) + } +} + +object NativeDataSourcePackage extends Logging { + val FILE_EXTENSION: String = Artifact.NATIVE_DATA_SOURCE_PACKAGE_EXTENSION + val MANIFEST_NAME = "spark-native-datasource.json" + + private[ffi] type ReadResult = Either[InvalidNativeDataSourcePackage, NativeDataSourcePackage] + + // Bounded, because each session stores its artifacts in a directory of its own, so a + // long-running server where many sessions add packages reads many different files. + private val MAX_CACHED_PACKAGES = 1000 + + private val packages: Cache[FileVersion, ReadResult] = + CacheBuilder.newBuilder().maximumSize(MAX_CACHED_PACKAGES).build[FileVersion, ReadResult]() + private val mapper = new ObjectMapper() + + /** Reads the package in the given file. The result is cached until the file changes. */ + def read(file: File): NativeDataSourcePackage = tryRead(file).fold(p => throw p.error, identity) + + /** + * Reads the package in the given file, or returns why it is invalid. The result is cached until + * the file changes, and an invalid package, or a package of a version of the interface that + * Spark does not support, is logged when it is read. + */ + private[ffi] def tryRead(file: File): ReadResult = { + val version = try { + FileVersion(file.getCanonicalFile) + } catch { + case e: IOException => + return Left(InvalidNativeDataSourcePackage( + invalidManifestError(file, s"The package cannot be read: $e", e), Nil)) + } + Option(packages.getIfPresent(version)).getOrElse { + val canonical = new File(version.path) + val result = parse(canonical) + result match { + case Left(invalid) => + logWarning(log"Invalid native data source package ${MDC(PATH, canonical)}.", + invalid.error) + case Right(pkg) if pkg.manifest.abiVersion > NativeLibrary.ABI_VERSION => + logWarning(log"The native data source package ${MDC(PATH, canonical)} implements " + + log"version ${MDC(VERSION_NUM, pkg.manifest.abiVersion)} of the native data source " + + log"interface, which this version of Spark does not support, so it is not used.") + case _ => + } + packages.put(version, result) + result + } + } + + /** + * Forgets the packages read from the files under the given directory, such as the artifacts of + * a session that is cleaned up. + */ + def forgetPackagesIn(dir: File): Unit = { + val prefix = dir.getCanonicalPath + File.separator + packages.asMap().keySet().removeIf(_.path.startsWith(prefix)) + } + + /** Returns whether a package read from the file at the given canonical path is cached. */ + private[ffi] def isCached(canonicalPath: String): Boolean = { + packages.asMap().keySet().asScala.exists(_.path == canonicalPath) + } + + private def invalidManifestError(file: File, reason: String, cause: Throwable = null) = { + QueryExecutionErrors.invalidNativeDataSourcePackageError( + file.getPath, "INVALID_MANIFEST", Map("reason" -> reason), cause) + } + + private def parse(file: File): ReadResult = { + def invalid(reason: String, dataSources: Seq[String] = Nil, cause: Throwable = null) = { + Left(InvalidNativeDataSourcePackage(invalidManifestError(file, reason, cause), dataSources)) + } + try { + val manifest = Utils.tryWithResource(new ZipFile(file)) { zip => + Option(zip.getEntry(MANIFEST_NAME)) match { + case None => invalid(s"The package does not contain the manifest $MANIFEST_NAME.") + case Some(entry) => + val json = Utils.tryWithResource(zip.getInputStream(entry)) { in => + new String(in.readAllBytes(), StandardCharsets.UTF_8) + } + parseManifest(json) match { + case Left(reason) => invalid(reason) + case Right(manifest) => + manifest.libraries.values.find(zip.getEntry(_) == null) match { + case Some(library) => invalid( + s"The package does not contain the library $library.", manifest.dataSources) + case None => Right(manifest) + } + } + } + } + manifest.map(NativeDataSourcePackage(file.getPath, sha256(file), _)) + } catch { + case e: IOException => invalid(s"The package is not a valid zip file: $e", cause = e) + } + } + + private[ffi] def parseManifest(json: String): Either[String, NativeDataSourceManifest] = { + val root: JsonNode = try { + mapper.readTree(json) + } catch { + case e: IOException => return Left(s"The manifest is not valid JSON: ${e.getMessage}") + } + if (root == null || !root.isObject) { + return Left("The manifest must be a JSON object.") + } + val abiVersion = root.get("abiVersion") + if (abiVersion == null || !abiVersion.isInt || abiVersion.asInt() < 1) { + return Left("The manifest must have a positive 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 && !n.asText().contains('\u0000') + }) { + return Left("The manifest must have a field 'dataSources' with a non-empty array of " + + "data source names.") + } + val libraries = root.get("libraries") + if (libraries == null || !libraries.isObject || libraries.isEmpty || + !libraries.properties().asScala.forall(e => e.getValue.isTextual)) { + return Left("The manifest must have a field 'libraries' with an object that maps " + + "platforms to the paths of the libraries in the package.") + } + Right(NativeDataSourceManifest( + abiVersion.asInt(), + dataSources.elements().asScala.map(_.asText()).toSeq, + libraries.properties().asScala.map(e => e.getKey -> e.getValue.asText()).toMap)) + } + + private[ffi] def sha256(file: File): String = { + val digest = MessageDigest.getInstance("SHA-256") + Utils.tryWithResource(new FileInputStream(file)) { in => + val buffer = new Array[Byte](64 * 1024) + var read = in.read(buffer) + while (read != -1) { + digest.update(buffer, 0, read) + read = in.read(buffer) + } + } + HexFormat.of().formatHex(digest.digest()) + } +} + +/** + * A native data source library installed in the native library path, where + * [[NativeDataSourceRegistry]] finds it by its file name. It must be installed on every node. The + * SHA-256 checksum of the file identifies the build of the library, so that an executor can check + * that it loads the same build as the driver: the state that the driver serializes is only + * meant for the same build. + * + * @param path the path of the library on the driver + * @param sha256 the SHA-256 checksum of the library on the driver + */ +case class InstalledNativeLibrary(path: String, sha256: String) extends NativeLibraryLocation + +object InstalledNativeLibrary { + // Bounded, like the cache of the packages, although few libraries are installed. + private val checksums: Cache[FileVersion, String] = + CacheBuilder.newBuilder().maximumSize(1000).build[FileVersion, String]() + + /** The installed library in the given file. */ + def apply(file: File): InstalledNativeLibrary = InstalledNativeLibrary(file.getPath, sha256(file)) + + /** The SHA-256 checksum of the given file. The result is cached until the file changes. */ + private[ffi] def sha256(file: File): String = { + val version = FileVersion(file.getCanonicalFile) + Option(checksums.getIfPresent(version)).getOrElse { + val checksum = NativeDataSourcePackage.sha256(new File(version.path)) + checksums.put(version, checksum) + checksum + } + } +} + +/** The platform names used in [[NativeDataSourceManifest.libraries]]. */ +object NativePlatform { + /** + * The platform of this JVM, such as `linux-x86_64`, `linux-aarch64` or `osx-aarch64`. `linux` is + * Linux with glibc: with musl, as on Alpine Linux, the platform ends with `-musl`, such as + * `linux-x86_64-musl`, since a library built for one does not load with the other. + */ + lazy val current: String = { + val osName = System.getProperty("os.name") + val linux = osName.toLowerCase(Locale.ROOT).startsWith("linux") + name(osName, System.getProperty("os.arch"), musl = linux && usesMusl(processMaps())) + } + + /** The libraries mapped in this process, as listed in `/proc/self/maps`, if it can be read. */ + private def processMaps(): Seq[String] = { + try { + Files.readAllLines(Paths.get("/proc/self/maps")).asScala.toSeq + } catch { + case _: IOException => Nil + } + } + + /** Whether the C library of a process is musl, which it maps as `ld-musl-.so.1`. */ + private[ffi] def usesMusl(maps: Seq[String]): Boolean = maps.exists(_.contains("/ld-musl-")) + + private[ffi] def name(osName: String, osArch: String, musl: Boolean = false): 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 + } + if (os == "linux" && musl) s"$os-$arch-musl" else s"$os-$arch" + } +} + +/** + * The native data source libraries loaded in this JVM. A library is loaded once, on first use, + * after it is extracted from its package if needed, and stays loaded until the JVM exits. + */ +object NativeLibraries extends Logging { + // The libraries of the packages, by the checksum of the package, and the installed libraries, + // by their canonical path, with the checksum of the build that was loaded, so that each library + // file is loaded once. A library that is loaded but rejected stays loaded, so its error is kept: + // it is not loaded again. + private val packagedLibraries = new ConcurrentHashMap[String, Try[NativeLibrary]]() + private val installedLibraries = new ConcurrentHashMap[String, (String, Try[NativeLibrary])]() + + private lazy val extractionDir: File = Utils.createTempDir(namePrefix = "native-datasources") + + /** Returns the library at the given location, and loads it if needed. */ + private[ffi] def get(location: NativeLibraryLocation): NativeLibrary = { + // The libraries are called with the FFM API, which is final only from Java 22. + if (!NativeLibrary.isSupported) { + throw QueryExecutionErrors.nativeDataSourceUnsupportedJavaVersionError( + NativeLibrary.MINIMUM_JAVA_VERSION, Runtime.version().feature()) + } + checkCDataInterface() + getSupported(location) + } + + /** + * Fails unless the JNI library of the C data interface of Arrow Java, with which Spark exchanges + * the data, loads on this platform. Arrow Java ships it for Linux with glibc and macOS on x86_64 + * and aarch64, and for Windows on x86_64; elsewhere, its own error does not say what it is for. + */ + private[ffi] def checkCDataInterface(load: => Unit = JniLoader.get().ensureLoaded()): Unit = { + try { + load + } catch { + // Arrow Java does not have the library for the architecture, or the library does not link, + // as with musl. + case e @ (_: IllegalStateException | _: LinkageError) => + throw QueryExecutionErrors.nativeDataSourceUnsupportedPlatformError( + NativePlatform.current, e) + } + } + + private def getSupported(location: NativeLibraryLocation): NativeLibrary = location match { + case pkg: NativeDataSourcePackage => + getOrLoad(packagedLibraries, pkg.sha256)(loadPackaged(pkg)).get + case installed: InstalledNativeLibrary => + // The file has the checksum of the build of the driver. + val file = locate(installed).getCanonicalFile + val (loadedChecksum, library) = + getOrLoad(installedLibraries, file.getPath)(installed.sha256 -> loadInstalled(file)) + // The JVM cannot load another build from the same path: the library was replaced since. + if (loadedChecksum != installed.sha256) { + throw QueryExecutionErrors.invalidNativeDataSourceLibraryError( + file.getPath, + "REPLACED", + Map("loadedChecksum" -> loadedChecksum, "checksum" -> installed.sha256)) + } + library.get + } + + private def getOrLoad[T](libraries: ConcurrentHashMap[String, T], key: String)( + load: => T): T = { + val library = libraries.get(key) + if (library != null) library else libraries.computeIfAbsent(key, _ => load) + } + + private def loadPackaged(pkg: NativeDataSourcePackage): Try[NativeLibrary] = { + val platform = NativePlatform.current + val libraryPath = pkg.manifest.libraries.getOrElse(platform, { + throw QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, + "UNSUPPORTED_PLATFORM", + Map( + "platform" -> platform, + "platforms" -> pkg.manifest.libraries.keys.toSeq.sorted.mkString(", "))) + }) + val packageFile = locate(pkg) + val libraryFile = new File(new File(extractionDir, pkg.sha256), new File(libraryPath).getName) + if (!libraryFile.isFile) { + extract(packageFile, libraryPath, libraryFile) + } + load( + libraryFile, + (reason, cause) => QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, "INVALID_LIBRARY", Map("library" -> libraryPath, "reason" -> reason), cause), + parameters => QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, "UNSUPPORTED_ABI_VERSION", parameters)) + } + + private def loadInstalled(file: File): Try[NativeLibrary] = { + load( + file, + (reason, cause) => QueryExecutionErrors.invalidNativeDataSourceLibraryError( + file.getPath, "CANNOT_LOAD", Map("reason" -> reason), cause), + parameters => QueryExecutionErrors.invalidNativeDataSourceLibraryError( + file.getPath, "UNSUPPORTED_ABI_VERSION", parameters)) + } + + /** + * Loads a library, and checks the version of the interface that it implements. Throws the error + * if the library cannot be loaded, and returns it if the library is loaded but rejected, so that + * it is kept: the library stays loaded, and Spark only initializes it once. + * + * @param invalidLibrary returns the error for an invalid library, with the reason + * @param unsupportedAbiVersion returns the error for an unsupported version of the interface, + * with the parameters `version` and `supported` + */ + private def load( + file: File, + invalidLibrary: (String, Throwable) => Throwable, + unsupportedAbiVersion: Map[String, String] => Throwable): Try[NativeLibrary] = { + logInfo(log"Loading the native data source library ${MDC(PATH, file)}.") + val library = try { + NativeLibrary.load(file.getPath, SPARK_VERSION) + } catch { + // A library that was loaded stays loaded, and Spark must not initialize it again. + case e: NativeLibrary.LoadException if e.loaded => + return Failure(invalidLibrary(e.getMessage, e)) + case e: NativeLibrary.LoadException => throw invalidLibrary(e.getMessage, e) + } + Try { + if (library.abiVersion() != NativeLibrary.ABI_VERSION) { + throw unsupportedAbiVersion(Map( + "version" -> library.abiVersion().toString, + "supported" -> NativeLibrary.ABI_VERSION.toString)) + } + library + } + } + + /** Finds a copy of the package on this node: the one on the driver, or a distributed one. */ + private def locate(pkg: NativeDataSourcePackage): File = { + (new File(pkg.path) +: NativeDataSourceRegistry.distributedCopies(pkg)) + .find { file => + file.isFile && NativeDataSourcePackage.tryRead(file).exists(_.sha256 == pkg.sha256) + } + .getOrElse { + throw QueryExecutionErrors.nativeDataSourcePackageNotFoundError(pkg.fileName, pkg.sha256) + } + } + + /** + * Finds an installed library on this node with the same checksum as on the driver: where the + * driver found it, or else in the native library path of this node. + */ + private def locate(installed: InstalledNativeLibrary): File = { + val file = new File(installed.path) + (file +: NativeDataSourceRegistry.findInstalledLibraries(file.getName)) + .find(f => f.isFile && InstalledNativeLibrary.sha256(f) == installed.sha256) + .getOrElse { + throw QueryExecutionErrors.nativeDataSourceLibraryNotFoundError( + file.getName, installed.sha256, installed.path) + } + } + + private def extract(packageFile: File, libraryPath: String, target: File): Unit = { + target.getParentFile.mkdirs() + val temp = File.createTempFile(target.getName, ".tmp", target.getParentFile) + try { + Utils.tryWithResource(new ZipFile(packageFile)) { zip => + Utils.tryWithResource(zip.getInputStream(zip.getEntry(libraryPath))) { in => + Files.copy(in, temp.toPath, StandardCopyOption.REPLACE_EXISTING) + } + } + Files.move(temp.toPath, target.toPath, StandardCopyOption.REPLACE_EXISTING) + } finally { + temp.delete() + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala new file mode 100644 index 000000000000..f58f056fb8dd --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -0,0 +1,315 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.io.File +import java.nio.file.Paths +import java.util.Locale +import java.util.concurrent.ConcurrentHashMap + +import org.apache.spark.{SparkEnv, SparkFiles} +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.{DATA_SOURCE, PATH, PATHS} +import org.apache.spark.sql.Artifact +import org.apache.spark.sql.classic.SparkSession +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.util.Utils + +/** + * Finds the native data sources on the driver. Like the Python data sources installed in the + * Python path, the native data source libraries installed in the native library path are found + * automatically, by their file name: the library `spark_datasource_`, such as + * `libspark_datasource_my_source.so` on Linux, implements the data source ``. The active + * session can also add packages with `spark.addArtifact`, or with + * `spark.sql.dataSource.native.paths`, which take precedence over the installed libraries. The + * package that the session added last takes precedence over the other packages. + */ +object NativeDataSourceRegistry extends Logging { + + /** The prefix of the names of the installed native data source libraries. */ + val LIBRARY_NAME_PREFIX = "spark_datasource_" + + // The result of the last lookup by `DataSource.lookupDataSource` on this thread, if it found a + // native data source, by the name it looked up. The provider that the caller creates next for + // that name takes it, so that it does not look the name up again: see [[takeLookup]]. Each + // lookup clears it first, so a provider never takes the result of an earlier lookup. + private val lastLookup = new ThreadLocal[(String, (String, NativeLibraryLocation))] + + /** + * Finds the package or the installed library that provides the native data source with the + * given name, which is case-insensitive. Returns the name of the data source, as listed in the + * manifest of the package or in lower case for an installed library, and where its library is. + */ + def lookup(name: String, conf: SQLConf): Option[(String, NativeLibraryLocation)] = { + if (!conf.getConf(StaticSQLConf.NATIVE_DATA_SOURCE_ENABLED)) { + return None + } + val artifacts = artifactPackageFiles() + val (invalid, valid) = (configuredPackageFiles(conf) ++ artifacts) + .map(NativeDataSourcePackage.tryRead).partitionMap(identity) + // A package of a version of the interface that Spark does not support is unusable too. + val (unsupported, usable) = valid.partitionMap { pkg => + // The version of a package is the oldest that its library implements. + if (pkg.manifest.abiVersion <= NativeLibrary.ABI_VERSION) Right(pkg) else Left(pkg) + } + // The packages are read from their canonical paths, in the order that the session added them. + val addedPaths = artifacts.map(_.getCanonicalPath).toSet + val (added, configured) = usable.partition(pkg => addedPaths.contains(pkg.path)) + lookupAddedPackage(name, added, configured) + .orElse(lookupPackage(name, configured)) + .orElse(lookupInstalledLibrary(name)) + .orElse { + // An unusable package is only reported when no other package or installed library + // provides the data source: a stale copy does not hide a valid one. An invalid package + // whose manifest cannot be read is only logged, when it is read. + unsupported.find(_.manifest.dataSources.exists(_.equalsIgnoreCase(name))).foreach { pkg => + throw QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, + "UNSUPPORTED_ABI_VERSION", + Map( + "version" -> pkg.manifest.abiVersion.toString, + "supported" -> NativeLibrary.ABI_VERSION.toString)) + } + invalid.find(_.dataSources.exists(_.equalsIgnoreCase(name))).foreach(p => throw p.error) + None + } + } + + /** + * Like [[lookup]], and keeps the result for the provider that the caller creates next for the + * given name, which takes it with [[takeLookup]]. + */ + def lookupForProvider(name: String, conf: SQLConf): Option[(String, NativeLibraryLocation)] = { + val result = lookup(name, conf) + result.foreach(found => lastLookup.set(name -> found)) + result + } + + /** + * Takes the result of the last [[lookupForProvider]] on this thread, if it was for the given + * name, and forgets it. + */ + def takeLookup(name: String): Option[(String, NativeLibraryLocation)] = { + val last = Option(lastLookup.get()) + lastLookup.remove() + last.collect { case (lookedUp, found) if lookedUp.equalsIgnoreCase(name) => found } + } + + /** Forgets the result of the last [[lookupForProvider]] on this thread. */ + def clearLookup(): Unit = lastLookup.remove() + + /** + * Finds the package that the session added last among those that provide the data source. As + * a Python data source that is registered again, it replaces the packages that the session + * added before, and the configured ones: a session cannot remove its artifacts, so it upgrades + * a package by adding the new version under another file name. + */ + private def lookupAddedPackage( + name: String, + added: Seq[NativeDataSourcePackage], + configured: Seq[NativeDataSourcePackage]): Option[(String, NativeDataSourcePackage)] = { + def providers(packages: Seq[NativeDataSourcePackage]): Seq[NativeDataSourcePackage] = { + packages.filter(_.manifest.dataSources.exists(_.equalsIgnoreCase(name))) + } + providers(added).lastOption.map { pkg => + val replaced = (providers(added) ++ providers(configured)) + .filter(_.sha256 != pkg.sha256).map(_.path).distinct + if (replaced.nonEmpty && loggedReplacements.add(pkg.path -> name.toLowerCase(Locale.ROOT))) { + logWarning(log"The native data source package ${MDC(PATH, pkg.path)}, which the " + + log"session added last, replaces ${MDC(PATHS, replaced.mkString(", "))} for the data " + + log"source ${MDC(DATA_SOURCE, name)}.") + } + pkg.manifest.dataSources.find(_.equalsIgnoreCase(name)).get -> pkg + } + } + + // The packages that were logged as replacing others for a data source, to only log it once. + private val loggedReplacements = ConcurrentHashMap.newKeySet[(String, String)]() + + private def lookupPackage( + name: String, + packages: Seq[NativeDataSourcePackage]): Option[(String, NativeDataSourcePackage)] = { + val found = packages + .flatMap(pkg => pkg.manifest.dataSources.find(_.equalsIgnoreCase(name)).map(_ -> pkg)) + // The same package can be found more than once, for example when it is in a configured + // directory and was also added as an artifact. + .groupBy(_._2.sha256).values.map(_.head).toSeq + found match { + case Seq() => None + case Seq(result) => Some(result) + case _ => + throw QueryCompilationErrors.nativeDataSourcePackageConflictError( + name, found.map(_._2.path).sorted) + } + } + + private def lookupInstalledLibrary(name: String): Option[(String, InstalledNativeLibrary)] = { + // Other names, such as class names, are not the names of libraries. + if (!name.matches("[A-Za-z0-9_]+")) { + return None + } + val dataSource = name.toLowerCase(Locale.ROOT) + findInstalledLibraries(libraryFileName(dataSource)).headOption + .map(file => dataSource -> InstalledNativeLibrary(file)) + } + + /** + * Finds the libraries with the given file name in the native library path of this node, in the + * order of the path. + */ + private[ffi] def findInstalledLibraries(fileName: String): Seq[File] = { + librarySearchPath(System.getProperty("java.library.path"), System.getenv("PATH")) + .map(new File(_, fileName)) + .filter(_.isFile) + } + + /** + * The file name of the installed library of the native data source with the given name, such + * as `libspark_datasource_my_source.so` on Linux, `libspark_datasource_my_source.dylib` on + * macOS and `spark_datasource_my_source.dll` on Windows. + */ + def libraryFileName(name: String): String = { + System.mapLibraryName(LIBRARY_NAME_PREFIX + name.toLowerCase(Locale.ROOT)) + } + + /** + * The native library path: the directories where native data source libraries are installed, + * in the order they are searched. They are the directories of `java.library.path`, where the + * JVM finds native libraries, which include `LD_LIBRARY_PATH` on Linux, `DYLD_LIBRARY_PATH` on + * macOS and `PATH` on Windows, and then the `lib` directory of each installation prefix in + * `PATH`, such as `/usr/local/lib` for `/usr/local/bin`. Relative directories are ignored. + */ + private[ffi] def librarySearchPath(javaLibraryPath: String, path: String): Seq[File] = { + def directories(value: String): Seq[File] = { + Option(value).toSeq.flatMap(_.split(File.pathSeparator)).map(new File(_)).filter(_.isAbsolute) + } + val prefixLibraries = directories(path) + .filter(_.getName == "bin") + .flatMap(bin => Option(bin.getParentFile)) + .map(new File(_, "lib")) + (directories(javaLibraryPath) ++ prefixLibraries).distinct + } + + /** + * Makes the library available to the executors of the active session, which runs the jobs of + * the scan or the write: for a streaming query, it is the session of the query, a clone of the + * session that started it. Returns where the executors find it. An installed library is already + * installed on every node. + */ + def distribute(location: NativeLibraryLocation): NativeLibraryLocation = location match { + case pkg: NativeDataSourcePackage => + SparkSession.getActiveSession match { + // In local mode, the executors run in the driver and read the package where it is. + case Some(session) if session.sparkContext.isLocal => pkg + case Some(session) => addToArtifacts(session, pkg) + case None => + logWarning(log"Cannot distribute the native data source package " + + log"${MDC(PATH, pkg.path)} to the executors: there is no active session. The " + + log"executors only find it if it is at the same path on every node.") + pkg + } + case installed: InstalledNativeLibrary => installed + } + + /** + * Makes the package available to the executors of the session, as one of its artifacts. A + * package that the session added with `spark.addArtifact` already is one. Any other package, + * such as a package under `spark.sql.dataSource.native.paths`, is copied to the artifacts under + * a name made of its file name and its checksum: the copy does not have the extension of a + * package, so it is not a package of the session that the lookups find, and two different + * packages with the same file name do not collide. + */ + private[ffi] def addToArtifacts( + session: SparkSession, + pkg: NativeDataSourcePackage): NativeDataSourcePackage = { + val artifactManager = session.artifactManager + val (added, artifactUUID) = artifactManager.getNativeDataSourcePackages + if (added.exists(_.getCanonicalPath == pkg.path)) { + pkg.copy(artifactUUID = artifactUUID) + } else { + val copyName = s"${pkg.fileName}.${pkg.sha256}" + // Queries of the session can use the package at the same time, and the artifact manager + // fails to add a file that another thread is adding. + artifactManager.synchronized { + if (!artifactManager.hasFileArtifact(copyName)) { + artifactManager.addLocalArtifacts(Artifact.newFileArtifact( + Paths.get(copyName), new Artifact.LocalFile(Paths.get(pkg.path))) :: Nil) + } + } + pkg.copy(artifactUUID = artifactUUID, distributedFileName = Some(copyName)) + } + } + + /** + * The local copies, on an executor, of a package distributed with [[distribute]], in the order + * to check them. The files of an isolated session are stored in a directory named after the + * UUID of the session; the copies of other sessions are checked last, because the executors + * check the checksum of the package anyway. + */ + def distributedCopies(pkg: NativeDataSourcePackage): Seq[File] = { + if (SparkEnv.get == null) { + Nil + } else { + val fileName = pkg.distributedFileName.getOrElse(pkg.fileName) + val root = new File(SparkFiles.getRootDirectory()) + val sessionCopies = Option(root.listFiles()).toSeq.flatten + .filter(_.isDirectory) + .sortBy(_.getName) + .map(new File(_, fileName)) + (pkg.artifactUUID.map(uuid => new File(new File(root, uuid), fileName)).toSeq ++ + Seq(new File(root, fileName)) ++ sessionCopies).distinct + } + } + + private def configuredPackageFiles(conf: SQLConf): Seq[File] = { + conf.getConf(SQLConf.NATIVE_DATA_SOURCE_PATHS).flatMap { path => + val uri = Utils.resolveURI(path) + val file = new File(uri.getPath) + if (uri.getScheme != null && uri.getScheme != "file") { + logWarning(log"Ignoring the native data source path ${MDC(PATH, path)}: only local " + + log"paths are supported. Add remote packages with spark.addArtifact instead.") + Nil + } else if (file.isDirectory) { + packageFilesIn(file) + } else if (file.isFile) { + Seq(file) + } else { + logWarning(log"The native data source path ${MDC(PATH, path)} does not exist.") + Nil + } + } + } + + private def packageFilesIn(dir: File): Seq[File] = { + Option(dir.listFiles()) match { + case Some(files) => + files + .filter(f => f.isFile && f.getName.endsWith(NativeDataSourcePackage.FILE_EXTENSION)) + .sortBy(_.getName) + .toSeq + case None => + logWarning(log"Ignoring the native data source path ${MDC(PATH, dir)}: the directory " + + log"cannot be read.") + Nil + } + } + + private def artifactPackageFiles(): Seq[File] = { + SparkSession.getActiveSession.toSeq.flatMap(_.artifactManager.getNativeDataSourcePackages._1) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala new file mode 100644 index 000000000000..e9e6f3815594 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala @@ -0,0 +1,69 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.util + +import org.apache.spark.sql.connector.catalog.Table +import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.execution.datasources.v2.NamedTableProvider +import org.apache.spark.sql.execution.datasources.v2.columnar.ColumnarTable +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * The Data Source V2 provider of the native data sources: data sources implemented by a native + * library through the C interface in `org/apache/spark/sql/datasource/spark_datasource.h`. + */ +class NativeDataSourceV2 extends NamedTableProvider { + private var shortName: String = _ + // The result of the lookup that `DataSource.lookupDataSource` made to find this provider class. + private var lookup: Option[(String, NativeLibraryLocation)] = None + + override def setShortName(name: String): Unit = { + assert(shortName == null) + shortName = name + lookup = NativeDataSourceRegistry.takeLookup(name) + } + + // The name of the data source, and where its library is on the driver. Resolved on first use, + // which is while the query is analyzed, with the configuration of its session. + private lazy val resolved: (String, NativeLibraryLocation) = { + assert(shortName != null) + lookup.orElse(NativeDataSourceRegistry.lookup(shortName, SQLConf.get)).getOrElse { + throw QueryCompilationErrors.dataSourceDoesNotExist(shortName) + } + } + + override def inferSchema(options: CaseInsensitiveStringMap): StructType = { + val (name, location) = resolved + new NativeDataSource(location, name, options).schema() + } + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: util.Map[String, String]): Table = { + val (name, location) = resolved + new ColumnarTable( + shortName, schema, properties, options => new NativeDataSource(location, name, options)) + } + + override def supportsExternalMetadata(): Boolean = true +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala new file mode 100644 index 000000000000..dc4b1cf1f56d --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala @@ -0,0 +1,161 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.util.Base64 + +import com.fasterxml.jackson.core.json.JsonWriteFeature +import com.fasterxml.jackson.databind.json.JsonMapper +import com.fasterxml.jackson.databind.node.ObjectNode + +import org.apache.spark.sql.connector.expressions.{Expression, GeneralScalarExpression, Literal, NamedReference} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.types.UTF8String + +/** + * Encodes Data Source V2 predicates as the JSON expression trees passed to + * `push_predicates` of `spark_datasource.h`. + */ +object NativePredicates { + // Writes ASCII only, escaping the other characters, which any JSON parser decodes. + private val mapper = JsonMapper.builder().enable(JsonWriteFeature.ESCAPE_NON_ASCII).build() + + /** + * The names of the predicates that are passed to the libraries, as listed in + * `push_predicates`. They are part of the interface, so a predicate with any other + * name, such as an arithmetic function, is not passed: its semantics, for example on overflow, + * would be part of the interface too, and the names of Data Source V2 grow with new releases. + */ + val FUNCTION_NAMES: Set[String] = Set( + "=", "<>", "<=>", "<", "<=", ">", ">=", "IN", "IS_NULL", "IS_NOT_NULL", + "STARTS_WITH", "ENDS_WITH", "CONTAINS", "AND", "OR", "NOT") + + /** + * The comparisons, which are only passed on values that are not arrays, maps or structs: Spark + * compares those with other rules than the ones of `push_predicates`, for example with the + * nulls in them equal to each other. + */ + private val COMPARISONS: Set[String] = Set("=", "<>", "<=>", "<", "<=", ">", ">=", "IN") + + /** + * Returns the JSON of the predicate on the columns of the given schema, or None if it cannot + * be expressed in JSON with the functions in [[FUNCTION_NAMES]]. + */ + def toJson(predicate: Predicate, schema: StructType): Option[String] = { + encode(predicate, schema).map(mapper.writeValueAsString) + } + + private def encode(expression: Expression, schema: StructType): Option[ObjectNode] = { + expression match { + case reference: NamedReference => + encodeColumn(reference, schema) + + case literal: Literal[_] => + encodeLiteral(literal.value(), literal.dataType()) + + case function: GeneralScalarExpression if FUNCTION_NAMES.contains(function.name()) && + !(COMPARISONS.contains(function.name()) && + function.children().exists(isNested(_, schema))) => + val children = function.children().map(encode(_, schema)) + if (children.forall(_.isDefined)) { + val node = mapper.createObjectNode() + .put("type", "function") + .put("name", function.name()) + val array = node.putArray("children") + children.foreach(child => array.add(child.get)) + Some(node) + } else { + None + } + + case _ => None + } + } + + private def encodeColumn(reference: NamedReference, schema: StructType): Option[ObjectNode] = { + // A source compares strings by their bytes, so a column that contains strings with a + // non-binary collation, such as UTF8_LCASE, is not pushed down. Neither is a column that is + // not in the schema. + val hasNonBinaryCollation = columnType(reference.fieldNames().toSeq, schema).forall { + _.existsRecursively { + case st: StringType => !st.isUTF8BinaryCollation + case _ => false + } + } + if (hasNonBinaryCollation) { + None + } else { + val node = mapper.createObjectNode().put("type", "column") + val name = node.putArray("name") + reference.fieldNames().foreach(name.add) + Some(node) + } + } + + /** Whether the expression is a column, or a nested field, of an array, map or struct type. */ + private def isNested(expression: Expression, schema: StructType): Boolean = expression match { + case reference: NamedReference => + columnType(reference.fieldNames().toSeq, schema).exists { + case _: ArrayType | _: MapType | _: StructType => true + case _ => false + } + case _ => false + } + + /** The type of the column or the nested field with the given name parts. */ + private def columnType(names: Seq[String], dataType: DataType): Option[DataType] = { + (names, dataType) match { + case (Seq(), _) => Some(dataType) + case (name +: rest, struct: StructType) => + struct.fields.find(_.name == name).flatMap(field => columnType(rest, field.dataType)) + case _ => None + } + } + + private def encodeLiteral(value: Any, dataType: DataType): Option[ObjectNode] = { + val node = mapper.createObjectNode() + .put("type", "literal") + .put("dataType", dataType.catalogString) + if (value == null) { + Some(node.putNull("value")) + } else { + val encoded = (value, dataType) match { + case (b: Boolean, BooleanType) => Some(node.put("value", b)) + case (n: Byte, ByteType) => Some(node.put("value", n.toInt)) + case (n: Short, ShortType) => Some(node.put("value", n)) + case (n: Int, IntegerType | DateType) => Some(node.put("value", n)) + case (n: Long, LongType | TimestampType | TimestampNTZType) => Some(node.put("value", n)) + // As the double of the same value, which a reader that parses the numbers of JSON as + // doubles, as most do, gets exactly. The shortest decimal of the float would not be. + case (n: Float, FloatType) if !n.isNaN && !n.isInfinite => + Some(node.put("value", n.toDouble)) + case (n: Double, DoubleType) if !n.isNaN && !n.isInfinite => Some(node.put("value", n)) + // A source compares strings by their bytes, so only push down binary collations, and + // valid UTF-8: a JSON string would replace the invalid bytes. + case (s: UTF8String, st: StringType) if st.isUTF8BinaryCollation && s.isValid => + Some(node.put("value", s.toString)) + case (bytes: Array[Byte], BinaryType) => + Some(node.put("value", Base64.getEncoder.encodeToString(bytes))) + case (d: Decimal, _: DecimalType) => + Some(node.put("value", d.toJavaBigDecimal.toPlainString)) + case _ => None + } + encoded + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala index 9fcf81746fe3..cd04bba0c288 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.execution.datasources.v2.python import org.apache.spark.sql.SparkSession import org.apache.spark.sql.connector.catalog._ import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.execution.datasources.v2.NamedTableProvider import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap @@ -26,10 +27,10 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap /** * Data Source V2 wrapper for Python Data Source. */ -class PythonDataSourceV2 extends TableProvider { +class PythonDataSourceV2 extends NamedTableProvider { private var name: String = _ - def setShortName(str: String): Unit = { + override def setShortName(str: String): Unit = { assert(name == null) name = str } diff --git a/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister b/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister index c1fc7234d7c1..8fcf427ec8a7 100644 --- a/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister +++ b/sql/core/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister @@ -30,3 +30,4 @@ org.apache.spark.sql.streaming.sources.FakeWriteOnly org.apache.spark.sql.streaming.sources.FakeNoWrite org.apache.spark.sql.streaming.sources.FakeWriteSupportingExternalMetadata org.apache.spark.sql.streaming.sources.FakeWriteSupportProviderV1Fallback +org.apache.spark.sql.execution.datasources.v2.columnar.TestRangeProvider diff --git a/sql/core/src/test/resources/native-datasource/test_native_datasource.cc b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc new file mode 100644 index 000000000000..3b12422f29f7 --- /dev/null +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -0,0 +1,1445 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// A native data source library used by NativeDataSourceSuite. It implements the C interface in +// spark_datasource.h without any other dependency, building the Arrow C data interface structs by +// hand. It provides three data sources: +// +// native_range Reads rows [0, end) with columns of several types. Supports predicate, +// column and limit pushdown. With the option "claim_predicates", it reports +// all the predicates as evaluated completely, but ignores them. With the +// option "predicate_use", it reports the predicates that it evaluates with +// another value: "partial" for SPARK_DS_PREDICATE_PARTIAL, as a reader that +// reads the blocks of 10 rows that may match, or a number. With the option +// "extra", the schema has one more column: "map" adds a map whose entry fields +// are named "keys" and "values", as arrow-rs names them, "struct" adds a +// struct, "decimal256" adds a decimal256, "listview" adds a list view, +// "largelist" adds a large list, and "dictionary" adds a dictionary-encoded +// string column, whose data Spark rejects before reading it. With the option +// "slice", the columns of the batches have an offset of 1. With the option +// "short", an array of the batches has one value fewer than its parent refers +// to, although its buffers hold them all: "id" shortens the column id, +// "entries" the entries of the map and their fields, and "values" the values +// of the entries. "children" exports the batches without their columns, with +// n_children = 0. With the option "null_rows", the struct of a batch has a +// validity buffer that marks its first row as null, and a null_count of 1 for +// "count" or -1 for "unknown"; "unknown_valid" has one without nulls and a +// null_count of -1, and "no_buffer" a null_count of -1 without a validity +// buffer. "column" gives the first column, which has no validity buffer, a +// null_count of -1 instead. With the option "wait_for_interrupt", the stream +// of a partition creates the file .waiting, waits until its +// task is killed, and then creates the file and fails. +// native_sink Writes the rows as CSV files under the directory in the option "path". With +// the option "commit_messages", a data writer passes its commit message that +// many times. Creating and closing a data writer create a file named after +// it in _created_data_writers and _closed_data_writers. +// native_counter Streams the rows [0, max_offset) with a single column "id", making "step" +// more rows available each time Spark asks for the latest offset. With the +// option "bad_offset", the latest offset is "none", "multiline" JSON or +// "invalid" JSON. With the option "closed_path", closing the stream reader +// creates that file. +// +// native_range supports batch reads, native_counter micro-batch streaming reads and native_sink +// batch and streaming writes, which may replace the existing data. With the option +// "capabilities", a data source declares these capabilities instead, a number. With the option +// "task_path", the reads and the data writers of native_range and native_sink record their task +// in a file named after its partition, in that directory. With the option "log", creating a data +// source logs its value as a warning, and with the option "spark_version", it fails unless Spark +// has that version. +// +// With the option "fail", a step fails with an error. With the option "zero", a function that +// creates a handle succeeds without setting it: "create", "create_reader", +// "create_stream_reader", "create_writer" or "create_data_writer". +// +// The suite compiles it with these macros to get several distinct libraries: +// +// NAME_PREFIX A prefix for the names of the data sources, "" by default. +// ID_SIGN Multiplies the ids produced by native_range, 1 by default. +// READ_ONLY Leaves out the functions for writes, streaming reads and pushdown. +// ABI_VERSION The version of the table, 1 by default. +// NO_LAST_ERROR Leaves out last_error. +// TABLE_SIZE The size of the table, sizeof(SparkDsLibrary) by default: Spark must not read +// the functions after it. +// INIT_FAILURE 1 to make spark_ds_init fail, after logging why, 2 to make it succeed without +// a table. +// NO_INIT Leaves out spark_ds_init, which makes the library invalid. + +#include "spark_datasource.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#ifndef NAME_PREFIX +#define NAME_PREFIX "" +#endif +#ifndef ID_SIGN +#define ID_SIGN 1 +#endif +#ifndef ABI_VERSION +#define ABI_VERSION 1 +#endif +#ifndef TABLE_SIZE +#define TABLE_SIZE sizeof(SparkDsLibrary) +#endif +#ifndef INIT_FAILURE +#define INIT_FAILURE 0 +#endif + +namespace { + +// --------------------------------------------------------------------------------------------- +// Columns +// --------------------------------------------------------------------------------------------- + +struct ColumnType { + const char* name; + const char* format; +}; + +// The columns of native_range, with their Arrow formats. native_counter only has "id". +const ColumnType kColumns[] = { + {"id", "l"}, // bigint + {"mod", "i"}, // int + {"name", "u"}, // string, null when id is a multiple of 5 + {"value", "g"}, // double + {"flag", "b"}, // boolean + {"day", "tdD"}, // date + {"ts", "tsu:UTC"} // timestamp +}; + +// The columns that native_range adds with the option "extra", by its value. +const std::map kExtraColumns = { + {"map", {"m", "+m"}}, // map: {"k": id} + {"struct", {"st", "+s"}}, // struct: {id, "s"} + {"decimal256", {"dec", "d:10,2,256"}}, // decimal256(10, 2): id + {"listview", {"lv", "+vl"}}, // list_view: [id] + {"largelist", {"ll", "+L"}}, // large_list, without data + {"dictionary", {"dict", "i"}} // dictionary, without data +}; + +const ColumnType* findColumn(const std::string& name) { + for (const auto& column : kColumns) { + if (name == column.name) return &column; + } + for (const auto& extra : kExtraColumns) { + if (name == extra.second.name) return &extra.second; + } + throw std::runtime_error("unknown column " + name); +} + +std::vector allColumns(bool counter, const std::string& extra = "") { + if (counter) return {"id"}; + std::vector names; + for (const auto& column : kColumns) names.push_back(column.name); + if (!extra.empty()) { + auto it = kExtraColumns.find(extra); + if (it == kExtraColumns.end()) throw std::runtime_error("unknown extra column " + extra); + names.push_back(it->second.name); + } + return names; +} + +// --------------------------------------------------------------------------------------------- +// Exporting schemas +// --------------------------------------------------------------------------------------------- + +struct SchemaData { + std::string format; + std::string name; + std::vector> children; + std::vector child_pointers; + std::unique_ptr dictionary; +}; + +void releaseSchema(ArrowSchema* schema) { + auto* data = static_cast(schema->private_data); + for (auto& child : data->children) { + if (child->release != nullptr) child->release(child.get()); + } + if (data->dictionary != nullptr && data->dictionary->release != nullptr) { + data->dictionary->release(data->dictionary.get()); + } + delete data; + schema->release = nullptr; +} + +void initSchema(ArrowSchema* out, const std::string& format, const std::string& name, + int64_t flags) { + auto* data = new SchemaData{format, name, {}, {}, nullptr}; + out->format = data->format.c_str(); + out->name = data->name.c_str(); + out->metadata = nullptr; + out->flags = flags; + out->n_children = 0; + out->children = nullptr; + out->dictionary = nullptr; + out->release = releaseSchema; + out->private_data = data; +} + +// Adds a child to a schema initialized by initSchema, and returns it. +ArrowSchema* addChild(ArrowSchema* parent, const std::string& format, const std::string& name, + int64_t flags) { + auto* data = static_cast(parent->private_data); + auto child = std::make_unique(); + initSchema(child.get(), format, name, flags); + data->child_pointers.push_back(child.get()); + data->children.push_back(std::move(child)); + parent->n_children = static_cast(data->children.size()); + parent->children = data->child_pointers.data(); + return data->children.back().get(); +} + +// Exports a struct type with the given columns. +void exportSchema(const std::vector& columns, ArrowSchema* out) { + initSchema(out, "+s", "", 0); + for (const auto& name : columns) { + ArrowSchema* column = addChild(out, findColumn(name)->format, name, ARROW_FLAG_NULLABLE); + if (name == "m") { + ArrowSchema* entries = addChild(column, "+s", "entries", 0); + addChild(entries, "u", "keys", 0); + addChild(entries, "l", "values", ARROW_FLAG_NULLABLE); + } else if (name == "st") { + addChild(column, "l", "a", ARROW_FLAG_NULLABLE); + addChild(column, "u", "b", ARROW_FLAG_NULLABLE); + } else if (name == "lv" || name == "ll") { + addChild(column, "l", "item", ARROW_FLAG_NULLABLE); + } else if (name == "dict") { + auto* data = static_cast(column->private_data); + data->dictionary = std::make_unique(); + initSchema(data->dictionary.get(), "u", "", ARROW_FLAG_NULLABLE); + column->dictionary = data->dictionary.get(); + } + } +} + +// Reads the column names of a struct schema exported by Spark, and releases it. +std::vector importColumns(ArrowSchema* schema) { + std::vector columns; + for (int64_t i = 0; i < schema->n_children; i++) { + columns.push_back(schema->children[i]->name); + } + schema->release(schema); + return columns; +} + +// --------------------------------------------------------------------------------------------- +// Exporting arrays +// --------------------------------------------------------------------------------------------- + +struct ArrayData { + std::vector> buffers; + std::vector buffer_pointers; + std::vector> children; + std::vector child_pointers; +}; + +void releaseArray(ArrowArray* array) { + auto* data = static_cast(array->private_data); + for (auto& child : data->children) { + if (child->release != nullptr) child->release(child.get()); + } + delete data; + array->release = nullptr; +} + +// Initializes an array that owns `buffers`. An empty buffer is exported as null. +void initArray(ArrowArray* out, int64_t length, int64_t null_count, + std::vector> buffers) { + auto* data = new ArrayData{std::move(buffers), {}, {}, {}}; + for (auto& buffer : data->buffers) { + data->buffer_pointers.push_back(buffer.empty() ? nullptr : buffer.data()); + } + out->length = length; + out->null_count = null_count; + out->offset = 0; + out->n_buffers = static_cast(data->buffers.size()); + out->n_children = 0; + out->buffers = data->buffer_pointers.data(); + out->children = nullptr; + out->dictionary = nullptr; + out->release = releaseArray; + out->private_data = data; +} + +template +std::vector toBytes(const std::vector& values) { + std::vector bytes(values.size() * sizeof(T)); + if (!bytes.empty()) std::memcpy(bytes.data(), values.data(), bytes.size()); + return bytes; +} + +void setBit(std::vector& bitmap, int64_t i) { + bitmap[i / 8] |= static_cast(1 << (i % 8)); +} + +// Adds a child to an array initialized by initArray, and returns it to be initialized. +ArrowArray* addChild(ArrowArray* parent) { + auto* data = static_cast(parent->private_data); + data->children.push_back(std::make_unique()); + data->child_pointers.push_back(data->children.back().get()); + parent->n_children = static_cast(data->children.size()); + parent->children = data->child_pointers.data(); + return data->children.back().get(); +} + +// Exports the values of a column for the ids [lo, hi). +void exportColumn(const std::string& name, int64_t lo, int64_t hi, ArrowArray* out) { + int64_t n = hi - lo; + std::vector no_validity; + if (name == "id") { + std::vector values; + for (int64_t id = lo; id < hi; id++) values.push_back(id * ID_SIGN); + initArray(out, n, 0, {no_validity, toBytes(values)}); + } else if (name == "mod") { + std::vector values; + for (int64_t id = lo; id < hi; id++) values.push_back(static_cast(id % 3)); + initArray(out, n, 0, {no_validity, toBytes(values)}); + } else if (name == "name") { + std::vector validity((n + 7) / 8, 0); + std::vector offsets{0}; + std::string chars; + int64_t null_count = 0; + for (int64_t id = lo; id < hi; id++) { + if (id % 5 == 0) { + null_count++; + } else { + setBit(validity, id - lo); + chars += "n" + std::to_string(id); + } + offsets.push_back(static_cast(chars.size())); + } + std::vector chars_bytes(chars.begin(), chars.end()); + initArray(out, n, null_count, {validity, toBytes(offsets), chars_bytes}); + } else if (name == "value") { + std::vector values; + for (int64_t id = lo; id < hi; id++) values.push_back(static_cast(id) * 0.5); + initArray(out, n, 0, {no_validity, toBytes(values)}); + } else if (name == "flag") { + std::vector bits((n + 7) / 8, 0); + for (int64_t id = lo; id < hi; id++) { + if (id % 2 == 0) setBit(bits, id - lo); + } + initArray(out, n, 0, {no_validity, bits}); + } else if (name == "day") { + std::vector values; + for (int64_t id = lo; id < hi; id++) values.push_back(static_cast(id)); + initArray(out, n, 0, {no_validity, toBytes(values)}); + } else if (name == "ts") { + std::vector values; + for (int64_t id = lo; id < hi; id++) values.push_back(id * 1000000); + initArray(out, n, 0, {no_validity, toBytes(values)}); + } else if (name == "m") { + // A single entry for each row. + std::vector offsets; + for (int64_t i = 0; i <= n; i++) offsets.push_back(static_cast(i)); + initArray(out, n, 0, {no_validity, toBytes(offsets)}); + ArrowArray* entries = addChild(out); + initArray(entries, n, 0, {no_validity}); + std::vector key_offsets{0}; + std::string keys; + std::vector values; + for (int64_t id = lo; id < hi; id++) { + keys += "k" + std::to_string(id); + key_offsets.push_back(static_cast(keys.size())); + values.push_back(id); + } + initArray(addChild(entries), n, 0, + {no_validity, toBytes(key_offsets), std::vector(keys.begin(), keys.end())}); + initArray(addChild(entries), n, 0, {no_validity, toBytes(values)}); + } else if (name == "st") { + initArray(out, n, 0, {no_validity}); + std::vector ids; + std::vector offsets{0}; + std::string chars; + for (int64_t id = lo; id < hi; id++) { + ids.push_back(id); + chars += "s" + std::to_string(id); + offsets.push_back(static_cast(chars.size())); + } + initArray(addChild(out), n, 0, {no_validity, toBytes(ids)}); + initArray(addChild(out), n, 0, + {no_validity, toBytes(offsets), std::vector(chars.begin(), chars.end())}); + } else if (name == "dec") { + // 256-bit little-endian unscaled values of the scale 2. + std::vector words; + for (int64_t id = lo; id < hi; id++) { + int64_t unscaled = id * 100; + int64_t sign = unscaled < 0 ? -1 : 0; + words.insert(words.end(), {unscaled, sign, sign, sign}); + } + initArray(out, n, 0, {no_validity, toBytes(words)}); + } else if (name == "lv") { + // A list of a single element for each row. + std::vector offsets; + std::vector sizes; + std::vector values; + for (int64_t id = lo; id < hi; id++) { + offsets.push_back(static_cast(id - lo)); + sizes.push_back(1); + values.push_back(id); + } + initArray(out, n, 0, {no_validity, toBytes(offsets), toBytes(sizes)}); + initArray(addChild(out), n, 0, {no_validity, toBytes(values)}); + } else { + // Including "ll" and "dict", which Spark rejects from the schema, before reading data. + throw std::runtime_error("no data for the column " + name); + } +} + +// Exports a struct array with the given columns for the ids [lo, hi). A sliced batch exports +// the columns for [lo - 1, hi), with an offset of 1, like a zero-copy slice. `short_array` names +// the array to export with one value fewer than its parent refers to, as the option "short". +void exportBatch(const std::vector& columns, int64_t lo, int64_t hi, bool slice, + const std::string& short_array, ArrowArray* out) { + initArray(out, hi - lo, 0, {std::vector()}); + for (const auto& name : columns) { + ArrowArray* column = addChild(out); + exportColumn(name, slice ? lo - 1 : lo, hi, column); + if (slice) { + column->offset = 1; + column->length = hi - lo; + column->null_count = -1; // Unknown. + } + if (short_array == "id" && name == "id") { + column->length--; + } else if (short_array == "entries" && name == "m") { + // The last offset of the map refers to one entry more than there are. + ArrowArray* entries = column->children[0]; + entries->length--; + for (int64_t i = 0; i < entries->n_children; i++) entries->children[i]->length--; + } else if (short_array == "values" && name == "m") { + column->children[0]->children[1]->length--; + } + } + // The columns are still released with the batch. + if (short_array == "children") out->n_children = 0; +} + +// Sets the nulls of the struct of a batch, as the option "null_rows". +void setBatchNulls(const std::string& null_rows, ArrowArray* batch) { + if (null_rows.empty()) return; + if (null_rows == "column") { + batch->children[0]->null_count = -1; // Its column has no nulls, so no validity buffer. + return; + } + batch->null_count = null_rows == "count" ? 1 : -1; + if (null_rows == "no_buffer") return; + // Every row is valid, but the first one for "count" and "unknown". + auto* data = static_cast(batch->private_data); + data->buffers[0].assign((batch->length + 7) / 8, 0xFF); + if (null_rows != "unknown_valid" && !data->buffers[0].empty()) data->buffers[0][0] = 0xFE; + data->buffer_pointers[0] = data->buffers[0].data(); +} + +// --------------------------------------------------------------------------------------------- +// A minimal JSON parser, for the predicates and the offsets +// --------------------------------------------------------------------------------------------- + +struct Json { + enum Kind { kNull, kBool, kNumber, kString, kArray, kObject } kind = kNull; + bool boolean = false; + double number = 0; + std::string text; + std::vector items; + std::map fields; + + const Json& operator[](const std::string& key) const { + static const Json missing; + auto it = fields.find(key); + return it == fields.end() ? missing : it->second; + } +}; + +class JsonParser { + public: + explicit JsonParser(const std::string& text) : text_(text) {} + + Json parse() { + Json value = parseValue(); + skipSpaces(); + if (pos_ != text_.size()) fail(); + return value; + } + + private: + [[noreturn]] void fail() { throw std::runtime_error("invalid JSON: " + text_); } + + void skipSpaces() { + while (pos_ < text_.size() && std::isspace(static_cast(text_[pos_]))) pos_++; + } + + bool consume(char c) { + skipSpaces(); + if (pos_ < text_.size() && text_[pos_] == c) { + pos_++; + return true; + } + return false; + } + + void expect(char c) { + if (!consume(c)) fail(); + } + + bool consumeWord(const char* word) { + size_t length = std::strlen(word); + if (text_.compare(pos_, length, word) == 0) { + pos_ += length; + return true; + } + return false; + } + + std::string parseString() { + expect('"'); + std::string result; + while (pos_ < text_.size() && text_[pos_] != '"') { + char c = text_[pos_++]; + if (c == '\\') { + if (pos_ >= text_.size()) fail(); + char escaped = text_[pos_++]; + switch (escaped) { + case 'n': result += '\n'; break; + case 't': result += '\t'; break; + case 'u': result += '?'; pos_ += 4; break; + default: result += escaped; + } + } else { + result += c; + } + } + expect('"'); + return result; + } + + Json parseValue() { + skipSpaces(); + Json value; + if (pos_ >= text_.size()) fail(); + char c = text_[pos_]; + if (c == '{') { + value.kind = Json::kObject; + pos_++; + if (!consume('}')) { + do { + skipSpaces(); + std::string key = parseString(); + expect(':'); + value.fields[key] = parseValue(); + } while (consume(',')); + expect('}'); + } + } else if (c == '[') { + value.kind = Json::kArray; + pos_++; + if (!consume(']')) { + do { + value.items.push_back(parseValue()); + } while (consume(',')); + expect(']'); + } + } else if (c == '"') { + value.kind = Json::kString; + value.text = parseString(); + } else if (consumeWord("true")) { + value.kind = Json::kBool; + value.boolean = true; + } else if (consumeWord("false")) { + value.kind = Json::kBool; + } else if (consumeWord("null")) { + value.kind = Json::kNull; + } else { + value.kind = Json::kNumber; + char* end = nullptr; + value.number = std::strtod(text_.c_str() + pos_, &end); + if (end == text_.c_str() + pos_) fail(); + pos_ = end - text_.c_str(); + } + return value; + } + + const std::string& text_; + size_t pos_ = 0; +}; + +// --------------------------------------------------------------------------------------------- +// State, serialized as "key=value" lines +// --------------------------------------------------------------------------------------------- + +using Options = std::map; + +std::string serialize(const Options& options) { + std::string result; + for (const auto& entry : options) result += entry.first + "=" + entry.second + "\n"; + return result; +} + +Options deserialize(const std::string& text) { + Options options; + std::istringstream lines(text); + std::string line; + while (std::getline(lines, line)) { + size_t eq = line.find('='); + if (eq != std::string::npos) options[line.substr(0, eq)] = line.substr(eq + 1); + } + return options; +} + +std::string get(const Options& options, const std::string& key, const std::string& fallback) { + auto it = options.find(key); + return it == options.end() ? fallback : it->second; +} + +int64_t getLong(const Options& options, const std::string& key, int64_t fallback) { + auto it = options.find(key); + return it == options.end() ? fallback : std::stoll(it->second); +} + +std::string join(const std::vector& values) { + std::string result; + for (size_t i = 0; i < values.size(); i++) result += (i > 0 ? "," : "") + values[i]; + return result; +} + +std::vector split(const std::string& text) { + std::vector values; + std::string value; + std::istringstream stream(text); + while (std::getline(stream, value, ',')) values.push_back(value); + return values; +} + +// --------------------------------------------------------------------------------------------- +// Helpers for the C interface +// --------------------------------------------------------------------------------------------- + +thread_local std::string last_error; + +// Spark, as spark_ds_init got it. +const SparkDsHost* host = nullptr; + +// Logs with the log of Spark. A library must only use the fields that fit in the size of the +// struct that Spark passes, which may be of another version of the header. +void hostLog(int32_t level, const std::string& message) { + if (host->size >= static_cast(offsetof(SparkDsHost, log) + sizeof(host->log))) { + host->log(level, message.c_str()); + } +} + +// Runs `f`, and turns a C++ exception into an error: none may cross the boundary. +template +int guarded(F&& f) { + try { + f(); + return 0; + } catch (const std::exception& e) { + last_error = e.what(); + return 1; + } +} + +std::vector toStrings(const char* const* values, int64_t count) { + return std::vector(values, values + count); +} + +std::string toBytes(const uint8_t* data, int64_t length) { + return std::string(reinterpret_cast(data), static_cast(length)); +} + +void sendBytes(SparkDsSink sink_function, void* sink, const std::string& bytes) { + sink_function(sink, reinterpret_cast(bytes.data()), + static_cast(bytes.size())); +} + +[[maybe_unused]] +std::vector toByteStrings(const uint8_t* const* messages, const int64_t* lengths, + int64_t count, bool* has_null) { + std::vector result; + for (int64_t i = 0; i < count; i++) { + if (messages[i] == nullptr) { + *has_null = true; + } else { + result.push_back(toBytes(messages[i], lengths[i])); + } + } + return result; +} + +template +T* fromHandle(void* handle) { + return static_cast(handle); +} + +// --------------------------------------------------------------------------------------------- +// Data sources +// --------------------------------------------------------------------------------------------- + +std::string name(const std::string& base) { return std::string(NAME_PREFIX) + base; } + +struct DataSource { + std::string name; + Options options; +}; + +void maybeFail(const Options& options, const std::string& step) { + if (get(options, "fail", "") == step) throw std::runtime_error("injected failure in " + step); +} + +// Whether the given function succeeds without setting its handle, with the option "zero". +bool returnsZero(const Options& options, const std::string& function) { + return get(options, "zero", "") == function; +} + +void writeFile(const std::string& path, const std::string& content) { + FILE* file = std::fopen(path.c_str(), "w"); + if (file == nullptr) throw std::runtime_error("cannot write " + path); + std::fputs(content.c_str(), file); + std::fclose(file); +} + +void makeDirs(const std::string& path) { + for (size_t pos = path.find('/', 1); pos != std::string::npos; pos = path.find('/', pos + 1)) { + mkdir(path.substr(0, pos).c_str(), 0755); + } + mkdir(path.c_str(), 0755); +} + +// With the option "task_path", records the task of a read or a data writer in a file named after +// its partition: the size of the struct and its fields. +void recordTask(const Options& options, const SparkDsTaskContext* task) { + std::string dir = get(options, "task_path", ""); + if (dir.empty()) return; + makeDirs(dir); + writeFile(dir + "/" + std::to_string(task->partition_id), + std::to_string(task->size) + "," + std::to_string(task->stage_id) + "," + + std::to_string(task->stage_attempt_number) + "," + + std::to_string(task->partition_id) + "," + std::to_string(task->attempt_number) + + "," + std::to_string(task->task_attempt_id)); +} + +// The driver-side reader of native_range and native_counter. +struct Reader { + Options state; // Serialized to the executors. +}; + +// Reads the ids [lo, hi) of the columns in the state, in batches. +struct Stream { + std::vector columns; + int64_t next; + int64_t end; + int64_t batch_size; + bool fail; + bool slice; + std::string short_array; + std::string null_rows; + const SparkDsTaskContext* task; + std::string wait_for_interrupt; + std::string last_error; +}; + +int streamGetSchema(ArrowArrayStream* stream, ArrowSchema* out) { + auto* state = static_cast(stream->private_data); + exportSchema(state->columns, out); + return 0; +} + +int streamGetNext(ArrowArrayStream* stream, ArrowArray* out) { + auto* state = static_cast(stream->private_data); + if (state->fail) { + state->last_error = "injected failure in get_next"; + return EIO; + } + if (!state->wait_for_interrupt.empty()) { + // Waits, as for data that does not come, until the task is killed, for at most a minute. + writeFile(state->wait_for_interrupt + ".waiting", ""); + for (int i = 0; i < 6000 && !state->task->is_interrupted(state->task); i++) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + bool interrupted = state->task->is_interrupted(state->task); + writeFile(state->wait_for_interrupt, interrupted ? "interrupted" : "timed out"); + state->last_error = "the task was killed"; + return ECANCELED; + } + if (state->next >= state->end) { + out->release = nullptr; // The end of the stream. + return 0; + } + int64_t hi = std::min(state->end, state->next + state->batch_size); + exportBatch(state->columns, state->next, hi, state->slice, state->short_array, out); + setBatchNulls(state->null_rows, out); + state->next = hi; + return 0; +} + +const char* streamGetLastError(ArrowArrayStream* stream) { + return static_cast(stream->private_data)->last_error.c_str(); +} + +void streamRelease(ArrowArrayStream* stream) { + delete static_cast(stream->private_data); + stream->release = nullptr; +} + +// The driver-side writer of native_sink. +struct Writer { + Options state; // path, overwrite, fail. + // The file that closeWriter creates, so that the tests can check that the writer of a committed + // or aborted write is closed: _CLOSED for a batch write, _CLOSED_ for a streaming one. + std::string closed_marker; +}; + +[[maybe_unused]] std::string closedMarker(int64_t epoch_id) { + return epoch_id < 0 ? "/_CLOSED" : "/_CLOSED_" + std::to_string(epoch_id); +} + +// Writes the rows of a task as CSV lines to a temporary file, renamed on commit. +struct DataWriter { + std::string temp_path; + FILE* file; + Options state; + bool committed; +}; + +// The data writers that Spark closed, with the path of their write, so that using one of them +// afterwards, which the header forbids, is recorded in the file _DATA_WRITER_USED_AFTER_CLOSE +// rather than crashing. +std::mutex closed_data_writers_lock; +[[maybe_unused]] std::map closed_data_writers; + + +[[maybe_unused]] bool isValid(const ArrowArray* array, int64_t i) { + if (array->null_count == 0 || array->buffers[0] == nullptr) return true; + int64_t bit = array->offset + i; + return (static_cast(array->buffers[0])[bit / 8] >> (bit % 8)) & 1; +} + +// Formats the value at row `i` of a column, or "null". The helpers below are only used by +// writes, which READ_ONLY leaves out. +[[maybe_unused]] +std::string formatValue(const ArrowSchema* schema, const ArrowArray* array, int64_t i) { + if (!isValid(array, i)) return "null"; + std::string format = schema->format; + int64_t row = array->offset + i; + if (format == "l" || format.rfind("ts", 0) == 0) { + return std::to_string(static_cast(array->buffers[1])[row]); + } else if (format == "i" || format == "tdD") { + return std::to_string(static_cast(array->buffers[1])[row]); + } else if (format == "g") { + std::ostringstream out; + out << static_cast(array->buffers[1])[row]; + return out.str(); + } else if (format == "b") { + return ((static_cast(array->buffers[1])[row / 8] >> (row % 8)) & 1) + ? "true" : "false"; + } else if (format == "u") { + const int32_t* offsets = static_cast(array->buffers[1]); + const char* chars = static_cast(array->buffers[2]); + return std::string(chars + offsets[row], static_cast(offsets[row + 1] - offsets[row])); + } + throw std::runtime_error("unsupported Arrow format " + format); +} + +[[maybe_unused]] +std::vector listFiles(const std::string& dir, const std::string& prefix) { + std::vector files; + DIR* handle = opendir(dir.c_str()); + if (handle == nullptr) return files; + while (dirent* entry = readdir(handle)) { + std::string file = entry->d_name; + if (file.rfind(prefix, 0) == 0) files.push_back(file); + } + closedir(handle); + return files; +} + +// --------------------------------------------------------------------------------------------- +// The functions of the table +// --------------------------------------------------------------------------------------------- + +#ifndef NO_LAST_ERROR +const char* lastError(void) { return last_error.c_str(); } +#endif + +// The capabilities of a data source, unless the option "capabilities" declares others. +uint64_t capabilitiesOf(const DataSource& source) { + std::string declared = get(source.options, "capabilities", ""); + if (!declared.empty()) return std::stoull(declared); + if (source.name == name("native_range")) return SPARK_DS_BATCH_READ; +#ifndef READ_ONLY + if (source.name == name("native_counter")) return SPARK_DS_MICRO_BATCH_READ; + if (source.name == name("native_sink")) { + return SPARK_DS_BATCH_WRITE | SPARK_DS_STREAMING_WRITE | SPARK_DS_TRUNCATE; + } +#endif + return 0; +} + +int createSource(const SparkDsCreateInfo* info, uint64_t* capabilities, void** handle) { + return guarded([&] { + // Like the host, the struct may be of another version of the header: only the fields that fit + // in its size are read. + if (info->size < static_cast(offsetof(SparkDsCreateInfo, num_options) + + sizeof(info->num_options))) { + throw std::runtime_error("SparkDsCreateInfo has a size of " + std::to_string(info->size)); + } + auto source = std::make_unique(DataSource{info->name, {}}); + auto keys = toStrings(info->option_keys, info->num_options); + auto values = toStrings(info->option_values, info->num_options); + for (size_t i = 0; i < keys.size(); i++) source->options[keys[i]] = values[i]; + maybeFail(source->options, "create"); + if (source->name != name("native_range") && source->name != name("native_sink") && + source->name != name("native_counter")) { + throw std::runtime_error("unknown data source " + source->name); + } + std::string message = get(source->options, "log", ""); + if (!message.empty()) hostLog(SPARK_DS_LOG_WARN, message); + std::string version = get(source->options, "spark_version", ""); + if (!version.empty() && version != host->spark_version) { + throw std::runtime_error("the version of Spark is " + std::string(host->spark_version) + + ", which supports version " + std::to_string(host->abi_version) + + " of the interface"); + } + *capabilities = capabilitiesOf(*source); + if (returnsZero(source->options, "create")) return; + *handle = source.release(); + }); +} + +int sourceSchema(void* handle, struct ArrowSchema* out) { + return guarded([&] { + auto* source = fromHandle(handle); + maybeFail(source->options, "schema"); + if (source->name == name("native_sink")) { + throw std::runtime_error("native_sink has no schema; it can only be written"); + } + bool counter = source->name == name("native_counter"); + exportSchema(allColumns(counter, get(source->options, "extra", "")), out); + }); +} + +int createReader(void* handle, struct ArrowSchema* schema, void** out) { + return guarded([&] { + auto* source = fromHandle(handle); + if (source->name != name("native_range")) { + throw std::runtime_error(source->name + " does not support batch reads"); + } + auto columns = importColumns(schema); + for (const auto& column : columns) findColumn(column); + if (returnsZero(source->options, "create_reader")) return; + auto* reader = new Reader(); + reader->state["columns"] = join(columns); + reader->state["lo"] = "0"; + reader->state["hi"] = get(source->options, "end", "10"); + reader->state["partitions"] = get(source->options, "partitions", "2"); + reader->state["batch_size"] = get(source->options, "batch_size", "1000"); + for (const char* option : {"fail", "slice", "short", "null_rows", "claim_predicates", + "predicate_use", "wait_for_interrupt", "task_path"}) { + reader->state[option] = get(source->options, option, ""); + } + *out = reader; + }); +} + +void closeSource(void* handle) { delete fromHandle(handle); } + +#ifndef READ_ONLY + +int createStreamReader(void* handle, struct ArrowSchema* schema, void** out) { + return guarded([&] { + auto* source = fromHandle(handle); + if (source->name != name("native_counter")) { + throw std::runtime_error(source->name + " does not support streaming reads"); + } + auto columns = join(importColumns(schema)); + if (returnsZero(source->options, "create_stream_reader")) return; + auto* reader = new Reader(); + reader->state["columns"] = columns; + reader->state["max_offset"] = get(source->options, "max_offset", "10"); + reader->state["step"] = get(source->options, "step", reader->state["max_offset"]); + reader->state["current"] = "0"; + reader->state["batch_size"] = get(source->options, "batch_size", "1000"); + reader->state["commit_path"] = get(source->options, "commit_path", ""); + reader->state["bad_offset"] = get(source->options, "bad_offset", ""); + reader->state["closed_path"] = get(source->options, "closed_path", ""); + *out = reader; + }); +} + +int pushPredicates(void* handle, const char* const* predicates, int64_t count, int32_t* uses) { + return guarded([&] { + auto* reader = fromHandle(handle); + auto json = toStrings(predicates, count); + int64_t end = getLong(reader->state, "hi", 0); + int64_t lo = getLong(reader->state, "lo", 0); + int64_t hi = end; + // With the option "claim_predicates", reports every predicate as evaluated completely, but + // evaluates none of them. + if (reader->state["claim_predicates"] == "true") { + for (size_t i = 0; i < json.size(); i++) uses[i] = SPARK_DS_PREDICATE_COMPLETE; + return; + } + std::string predicate_use = reader->state["predicate_use"]; + bool partial = predicate_use == "partial"; + int32_t use = predicate_use.empty() ? SPARK_DS_PREDICATE_COMPLETE + : partial ? SPARK_DS_PREDICATE_PARTIAL + : std::stoi(predicate_use); + for (size_t i = 0; i < json.size(); i++) { + // Accepts "id " for the comparison operators. + Json predicate = JsonParser(json[i]).parse(); + const Json& children = predicate["children"]; + std::string op = predicate["name"].text; + bool comparison = op == ">" || op == ">=" || op == "<" || op == "<=" || op == "="; + bool ok = ID_SIGN == 1 && predicate["type"].text == "function" && comparison && + children.items.size() == 2 && children.items[0]["type"].text == "column" && + children.items[0]["name"].items.size() == 1 && + children.items[0]["name"].items[0].text == "id" && + children.items[1]["type"].text == "literal" && + children.items[1]["value"].kind == Json::kNumber; + if (ok) { + auto value = static_cast(children.items[1]["value"].number); + if (op == ">") lo = std::max(lo, value + 1); + if (op == ">=") lo = std::max(lo, value); + if (op == "<") hi = std::min(hi, value); + if (op == "<=") hi = std::min(hi, value + 1); + if (op == "=") { + lo = std::max(lo, value); + hi = std::min(hi, value + 1); + } + } + uses[i] = ok ? use : SPARK_DS_PREDICATE_UNUSED; + } + if (partial) { + // Reads the blocks of 10 rows that contain matching rows, with all their rows. + lo = lo / 10 * 10; + hi = std::min(end, (hi + 9) / 10 * 10); + } + reader->state["lo"] = std::to_string(lo); + reader->state["hi"] = std::to_string(std::max(lo, hi)); + reader->state["predicates"] = std::to_string(json.size()); + }); +} + +int pushLimit(void* handle, int32_t limit, bool* applied) { + return guarded([&] { + auto* reader = fromHandle(handle); + // Spark calls push_limit at most once. + if (reader->state.count("limit") > 0) throw std::runtime_error("pushLimit was called twice"); + reader->state["limit"] = std::to_string(limit); + *applied = true; + }); +} + +int pruneColumns(void* handle, const char* const* columns, int64_t count, bool* pruned) { + return guarded([&] { + auto names = toStrings(columns, count); + for (const auto& column : names) findColumn(column); + fromHandle(handle)->state["columns"] = join(names); + *pruned = true; + }); +} + +#endif // READ_ONLY + +int planPartitions(void* handle, SparkDsSink add_partition, void* sink) { + return guarded([&] { + auto* reader = fromHandle(handle); + maybeFail(reader->state, "plan"); + int64_t lo = getLong(reader->state, "lo", 0); + int64_t hi = getLong(reader->state, "hi", 0); + int64_t count = getLong(reader->state, "partitions", 1); + if (reader->state.count("limit") > 0) { + // A pushed limit is applied by reading a single partition with at most `limit` rows. + hi = std::min(hi, lo + getLong(reader->state, "limit", 0)); + count = 1; + } + int64_t step = std::max(1, (hi - lo + count - 1) / count); + for (int64_t start = lo; start < hi; start += step) { + int64_t stop = std::min(hi, start + step); + sendBytes(add_partition, sink, std::to_string(start) + "," + std::to_string(stop)); + } + }); +} + +int serializeReader(void* handle, SparkDsSink set_state, void* sink) { + return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); +} + +void closeReader(void* handle) { delete fromHandle(handle); } + +#ifndef READ_ONLY + +int initialOffset(void*, SparkDsSink set_offset, void* sink) { + return guarded([&] { sendBytes(set_offset, sink, "{\"offset\":0}"); }); +} + +// Each call makes `step` more rows available, up to `max_offset`. +int latestOffset(void* handle, SparkDsSink set_offset, void* sink) { + return guarded([&] { + Options& state = fromHandle(handle)->state; + int64_t latest = std::min(getLong(state, "max_offset", 10), + getLong(state, "current", 0) + getLong(state, "step", 10)); + state["current"] = std::to_string(latest); + std::string bad_offset = get(state, "bad_offset", ""); + if (bad_offset == "none") return; + if (bad_offset == "multiline") { + // Pretty-printed JSON, as serde_json::to_string_pretty writes it. + sendBytes(set_offset, sink, "{\n \"offset\": " + std::to_string(latest) + "\n}"); + } else if (bad_offset == "invalid") { + sendBytes(set_offset, sink, "offset " + std::to_string(latest)); + } else { + sendBytes(set_offset, sink, "{\"offset\":" + std::to_string(latest) + "}"); + } + }); +} + +int streamPartitions(void*, const char* start, const char* end, SparkDsSink add_partition, + void* sink) { + return guarded([&] { + auto lo = static_cast(JsonParser(start).parse()["offset"].number); + auto hi = static_cast(JsonParser(end).parse()["offset"].number); + if (hi > lo) sendBytes(add_partition, sink, std::to_string(lo) + "," + std::to_string(hi)); + }); +} + +int serializeStreamReader(void* handle, SparkDsSink set_state, void* sink) { + return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); +} + +int commitOffset(void* handle, const char* end) { + return guarded([&] { + std::string path = get(fromHandle(handle)->state, "commit_path", ""); + if (!path.empty()) writeFile(path, end); + }); +} + +void closeStreamReader(void* handle) { + std::unique_ptr reader(fromHandle(handle)); + guarded([&] { + std::string path = get(reader->state, "closed_path", ""); + if (!path.empty()) writeFile(path, ""); + }); +} + +#endif // READ_ONLY + +int readPartition(const SparkDsTaskContext* task, const uint8_t* reader_state, + int64_t reader_state_length, const uint8_t* partition, + int64_t partition_length, struct ArrowArrayStream* stream) { + return guarded([&] { + Options state = deserialize(toBytes(reader_state, reader_state_length)); + maybeFail(state, "open"); + // Returns success without initializing the stream, which Spark must reject. + if (get(state, "fail", "") == "uninitialized") return; + auto bounds = split(toBytes(partition, partition_length)); + if (bounds.size() != 2) throw std::runtime_error("invalid partition"); + recordTask(state, task); + stream->get_schema = streamGetSchema; + stream->get_next = streamGetNext; + stream->get_last_error = streamGetLastError; + stream->release = streamRelease; + stream->private_data = new Stream{split(get(state, "columns", "")), std::stoll(bounds[0]), + std::stoll(bounds[1]), getLong(state, "batch_size", 1000), + get(state, "fail", "") == "read", + get(state, "slice", "") == "true", get(state, "short", ""), + get(state, "null_rows", ""), task, + get(state, "wait_for_interrupt", ""), ""}; + }); +} + +#ifndef READ_ONLY + +int createWriter(void* handle, struct ArrowSchema* schema, const SparkDsWriteInfo* info, + void** out) { + return guarded([&] { + auto* source = fromHandle(handle); + if (source->name != name("native_sink")) { + throw std::runtime_error(source->name + " does not support writes"); + } + if (info->mode != SPARK_DS_WRITE_APPEND && info->mode != SPARK_DS_WRITE_TRUNCATE) { + throw std::runtime_error("unknown mode " + std::to_string(info->mode)); + } + // Spark releases the schema after the call, since it is not moved here. + if (schema->n_children == 0) throw std::runtime_error("no columns to write"); + // The location of a table in a catalog is passed as a URI. + std::string path = get(source->options, "path", ""); + if (path.rfind("file://", 0) == 0) { + path = path.substr(7); + } else if (path.rfind("file:", 0) == 0) { + path = path.substr(5); + } + if (path.empty()) throw std::runtime_error("the option 'path' is required"); + if (returnsZero(source->options, "create_writer")) return; + auto* writer = new Writer(); + writer->state["path"] = path; + writer->state["overwrite"] = info->mode == SPARK_DS_WRITE_TRUNCATE ? "true" : "false"; + writer->state["fail"] = get(source->options, "fail", ""); + writer->state["zero"] = get(source->options, "zero", ""); + writer->state["commit_messages"] = get(source->options, "commit_messages", "1"); + writer->state["task_path"] = get(source->options, "task_path", ""); + // Records that a writer was created, which Spark then commits or aborts, and the ID of the + // streaming query. + makeDirs(path); + writeFile(path + "/_CREATED", ""); + if (info->query_id != nullptr) writeFile(path + "/_QUERY_ID", info->query_id); + *out = writer; + }); +} + +int serializeWriter(void* handle, SparkDsSink set_state, void* sink) { + return guarded([&] { sendBytes(set_state, sink, serialize(fromHandle(handle)->state)); }); +} + +int commitWrite(void* handle, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages) { + return guarded([&] { + auto* writer = fromHandle(handle); + std::string path = writer->state["path"]; + maybeFail(writer->state, "commit"); + bool has_null = false; + auto temp_files = toByteStrings(messages, message_lengths, num_messages, &has_null); + if (has_null) throw std::runtime_error("commit got a null message"); + if (writer->state["overwrite"] == "true") { + for (const auto& file : listFiles(path, "part-")) unlink((path + "/" + file).c_str()); + } + for (const auto& temp_file : temp_files) { + std::string file_name = temp_file.substr(temp_file.rfind('/') + 1); + std::rename(temp_file.c_str(), (path + "/part-" + file_name).c_str()); + } + writeFile(path + (epoch_id < 0 ? "/_SUCCESS" : "/_EPOCH_" + std::to_string(epoch_id)), + std::to_string(temp_files.size())); + writer->closed_marker = closedMarker(epoch_id); + }); +} + +// Records the aborted write in _ABORTED for a batch write, or _ABORTED_ for a streaming +// one, with the number of tasks that did not commit, whose messages are NULL. +int abortWrite(void* handle, int64_t epoch_id, const uint8_t* const* messages, + const int64_t* message_lengths, int64_t num_messages) { + return guarded([&] { + auto* writer = fromHandle(handle); + writer->closed_marker = closedMarker(epoch_id); + bool has_null = false; + for (const auto& temp_file : + toByteStrings(messages, message_lengths, num_messages, &has_null)) { + unlink(temp_file.c_str()); + } + int64_t null_messages = 0; + for (int64_t i = 0; i < num_messages; i++) { + if (messages[i] == nullptr) null_messages++; + } + writeFile(writer->state["path"] + + (epoch_id < 0 ? "/_ABORTED" : "/_ABORTED_" + std::to_string(epoch_id)), + std::to_string(null_messages)); + }); +} + +void closeWriter(void* handle) { + std::unique_ptr writer(fromHandle(handle)); + guarded([&] { + if (!writer->closed_marker.empty()) { + writeFile(writer->state["path"] + writer->closed_marker, ""); + } + }); +} + +// Whether Spark uses a data writer that it closed, which the header forbids. +bool usedAfterClose(void* handle) { + std::lock_guard guard(closed_data_writers_lock); + auto closed = closed_data_writers.find(handle); + if (closed == closed_data_writers.end()) return false; + writeFile(closed->second + "/_DATA_WRITER_USED_AFTER_CLOSE", ""); + return true; +} + +int createDataWriter(const SparkDsTaskContext* task, const uint8_t* writer_state, + int64_t writer_state_length, int64_t epoch_id, void** out) { + return guarded([&] { + Options state = deserialize(toBytes(writer_state, writer_state_length)); + if (returnsZero(state, "create_data_writer")) return; + recordTask(state, task); + std::string dir = state["path"] + "/_temporary"; + makeDirs(dir); + std::string temp_path = dir + "/" + std::to_string(epoch_id) + "-" + + std::to_string(task->partition_id) + "-" + + std::to_string(task->task_attempt_id) + ".csv"; + FILE* file = std::fopen(temp_path.c_str(), "w"); + if (file == nullptr) throw std::runtime_error("cannot write " + temp_path); + makeDirs(state["path"] + "/_created_data_writers"); + writeFile(state["path"] + "/_created_data_writers/" + temp_path.substr(dir.size() + 1), ""); + auto* writer = new DataWriter{temp_path, file, state, false}; + { + // The allocator may reuse the address of a closed data writer. + std::lock_guard guard(closed_data_writers_lock); + closed_data_writers.erase(writer); + } + *out = writer; + }); +} + +int writeBatch(void* handle, struct ArrowArray* array, struct ArrowSchema* schema) { + return guarded([&] { + auto* writer = fromHandle(handle); + // The batch is released here, so Spark does not release it after the call. + std::unique_ptr array_owner( + array, [](ArrowArray* a) { a->release(a); }); + std::unique_ptr schema_owner( + schema, [](ArrowSchema* s) { s->release(s); }); + maybeFail(writer->state, "write"); + for (int64_t i = 0; i < array->length; i++) { + std::string line; + for (int64_t c = 0; c < array->n_children; c++) { + if (c > 0) line += ","; + line += formatValue(schema->children[c], array->children[c], array->offset + i); + } + line += "\n"; + std::fputs(line.c_str(), writer->file); + } + }); +} + +int commitDataWriter(void* handle, SparkDsSink set_message, void* sink) { + if (usedAfterClose(handle)) return 0; + return guarded([&] { + auto* writer = fromHandle(handle); + // If the commit fails, Spark aborts the data writer. + maybeFail(writer->state, "commitDataWriter"); + // With the option "commit_messages", the message is passed that many times. + int64_t messages = getLong(writer->state, "commit_messages", 1); + for (int64_t i = 0; i < messages; i++) sendBytes(set_message, sink, writer->temp_path); + std::fclose(writer->file); + writer->file = nullptr; + writer->committed = true; + }); +} + +int abortDataWriter(void* handle) { + if (usedAfterClose(handle)) return 0; + return guarded([&] { + auto* writer = fromHandle(handle); + // Spark does not abort a data writer that committed: records it if it does. + if (writer->committed) { + writeFile(writer->state["path"] + "/_ABORTED_COMMITTED_DATA_WRITER", ""); + return; + } + if (writer->file != nullptr) { + std::fclose(writer->file); + writer->file = nullptr; + } + unlink(writer->temp_path.c_str()); + }); +} + +void closeDataWriter(void* handle) { + if (usedAfterClose(handle)) return; + std::unique_ptr writer(fromHandle(handle)); + guarded([&] { + if (writer->file != nullptr) std::fclose(writer->file); + std::string path = writer->state["path"]; + std::string dir = path + "/_closed_data_writers"; + makeDirs(dir); + writeFile(dir + "/" + writer->temp_path.substr(writer->temp_path.rfind('/') + 1), ""); + std::lock_guard guard(closed_data_writers_lock); + closed_data_writers[handle] = path; + }); +} + +#endif // READ_ONLY + +// The table, which NO_INIT leaves unused. +[[maybe_unused]] SparkDsLibrary makeTable() { + SparkDsLibrary table{}; + table.size = TABLE_SIZE; + table.abi_version = ABI_VERSION; +#ifndef NO_LAST_ERROR + table.last_error = lastError; +#endif + table.create = createSource; + table.schema = sourceSchema; + table.create_reader = createReader; + table.close = closeSource; + table.partitions = planPartitions; + table.serialize_reader = serializeReader; + table.close_reader = closeReader; + table.read = readPartition; +#ifndef READ_ONLY + table.create_stream_reader = createStreamReader; + table.create_writer = createWriter; + table.push_predicates = pushPredicates; + table.push_limit = pushLimit; + table.prune_columns = pruneColumns; + table.initial_offset = initialOffset; + table.latest_offset = latestOffset; + table.stream_partitions = streamPartitions; + table.serialize_stream_reader = serializeStreamReader; + table.commit_offset = commitOffset; + table.close_stream_reader = closeStreamReader; + table.serialize_writer = serializeWriter; + table.commit = commitWrite; + table.abort = abortWrite; + table.close_writer = closeWriter; + table.create_data_writer = createDataWriter; + table.write = writeBatch; + table.commit_data_writer = commitDataWriter; + table.abort_data_writer = abortDataWriter; + table.close_data_writer = closeDataWriter; +#endif + return table; +} + +} // namespace + +// --------------------------------------------------------------------------------------------- +// spark_datasource.h +// --------------------------------------------------------------------------------------------- + +#ifndef NO_INIT + +int spark_ds_init(const SparkDsHost* spark, const SparkDsLibrary** library) { + host = spark; + if (INIT_FAILURE == 1) { + hostLog(SPARK_DS_LOG_ERROR, "injected failure in spark_ds_init"); + return 1; + } + // Succeeds without returning a table, which Spark must reject. + if (INIT_FAILURE == 2) return 0; + static const SparkDsLibrary table = makeTable(); + *library = &table; + return 0; +} + +#endif // NO_INIT diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala new file mode 100644 index 000000000000..9c51b47cc991 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala @@ -0,0 +1,476 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.columnar + +import java.nio.charset.StandardCharsets.UTF_8 +import java.util +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicInteger + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.spark.SparkException +import org.apache.spark.sql.{AnalysisException, QueryTest, Row} +import org.apache.spark.sql.connector.catalog.{Table, TableCapability, TableProvider} +import org.apache.spark.sql.connector.expressions.{Literal, NamedReference, Transform} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} +import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.sources.DataSourceRegister +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +/** + * Tests the Data Source V2 adapter of [[ColumnarDataSource]] with a data source implemented on the + * JVM. `NativeDataSourceSuite` tests the native data sources, which implement it. + */ +class ColumnarDataSourceSuite extends QueryTest with SharedSparkSession { + import testImplicits._ + + private val format = "test_columnar_range" + + private def scan(df: org.apache.spark.sql.classic.DataFrame): ColumnarScan = { + df.queryExecution.optimizedPlan.collectFirst { + case relation: DataSourceV2ScanRelation => relation.scan.asInstanceOf[ColumnarScan] + }.get + } + + private def options(pairs: (String, String)*): CaseInsensitiveStringMap = { + new CaseInsensitiveStringMap(pairs.toMap.asJava) + } + + test("batch read") { + val df = spark.read.format(format).option("end", "10").load() + assert(df.schema == StructType.fromDDL("id BIGINT, name STRING")) + checkAnswer(df, (0L until 10L).map(id => Row(id, s"name$id"))) + // A data source can also be used by the class name of its provider. + checkAnswer( + spark.read.format(classOf[TestRangeProvider].getName).load(), + (0L until 10L).map(id => Row(id, s"name$id"))) + checkAnswer( + spark.read.schema("id BIGINT").format(format).option("end", "3").load(), + Seq(Row(0L), Row(1L), Row(2L))) + } + + test("push down predicates, a limit and the required columns") { + val filtered = spark.read.format(format).option("end", "10").load().filter("id > 6") + checkAnswer(filtered, (7L until 10L).map(id => Row(id, s"name$id"))) + assert(scan(filtered).getMetaData()("PushedPredicates").contains("id > 6")) + + val pruned = spark.read.format(format).option("end", "4").load().select("name") + checkAnswer(pruned, (0L until 4L).map(id => Row(s"name$id"))) + assert(scan(pruned).readSchema() == StructType.fromDDL("name STRING")) + + val limited = spark.read.format(format).option("end", "10").load().limit(2) + assert(limited.collect().length == 2) + assert(scan(limited).getMetaData()("PushedLimit") == "LIMIT 2") + } + + test("tables in a catalog") { + withTable("columnar_range") { + sql(s"CREATE TABLE columnar_range USING $format OPTIONS (end 3, table 'catalog_write')") + // The options of the table are used by the scans of the table, unless a scan sets them. + checkAnswer(spark.table("columnar_range"), (0L until 3L).map(id => Row(id, s"name$id"))) + checkAnswer( + spark.read.option("END", "2").table("columnar_range"), + (0L until 2L).map(id => Row(id, s"name$id"))) + + sql("INSERT INTO columnar_range VALUES (1, 'a')") + assert(TestColumnarSink.rows("catalog_write") == Seq((1L, "a"))) + } + } + + test("batch write") { + withSQLConf(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "3") { + val table = "batch_write" + (0L until 7L).map(id => (id, s"v$id")).toDF("id", "name").coalesce(1) + .write.format(format).option("table", table).mode("append").save() + assert(TestColumnarSink.rows(table) == (0L until 7L).map(id => (id, s"v$id"))) + // The rows are written in Arrow-backed batches of at most 3 rows. + assert(TestColumnarSink.batchSizes(table) == Seq(3, 3, 1)) + + Seq((10L, "a")).toDF("id", "name") + .write.format(format).option("table", table).mode("append").save() + assert(TestColumnarSink.rows(table).length == 8) + Seq((20L, "b")).toDF("id", "name") + .write.format(format).option("table", table).mode("overwrite").save() + assert(TestColumnarSink.rows(table) == Seq((20L, "b"))) + } + } + + test("batch write in batches of a limited size") { + withSQLConf(SQLConf.ARROW_EXECUTION_MAX_BYTES_PER_BATCH.key -> "4096") { + val table = "batch_write_bytes" + val rows = (0L until 10L).map(id => (id, "x" * 1000 + id)) + rows.toDF("id", "name").coalesce(1) + .write.format(format).option("table", table).mode("append").save() + assert(TestColumnarSink.rows(table) == rows) + // A batch ends once it holds at least 4096 bytes, long before it has 10,000 rows: each row + // takes about 1 KB, so a batch has 5 rows. + assert(TestColumnarSink.batchSizes(table) == Seq(5, 5)) + } + } + + test("create the writer once the data to write is computed") { + val table = "failed_upstream" + // The error happens in a shuffle stage, which adaptive query execution runs before the write, + // so no writer is created. Spark would not abort it. + val e = intercept[Exception] { + spark.range(10).selectExpr("id", "if(id >= 0, raise_error('boom'), 'x') AS name") + .repartition(2) + .write.format(format).option("table", table).mode("append").save() + } + assert(e.getMessage.contains("boom")) + assert(!TestColumnarSink.writers.contains(table)) + assert(!TestColumnarSink.aborted.contains(table)) + } + + test("the table reports the capabilities of the data source") { + val table = "unsupported_write" + // Computing the data would fail, in a shuffle stage that adaptive query execution runs before + // the write: Data Source V2 rejects the write while it analyzes the query, before that. + val df = spark.range(10).selectExpr("id", "if(id >= 0, raise_error('boom'), 'x') AS name") + .repartition(2) + checkError( + exception = intercept[AnalysisException] { + df.write.format(format).option("table", table).option("capabilities", "BATCH_READ") + .mode("append").save() + }, + condition = "UNSUPPORTED_FEATURE.TABLE_OPERATION", + parameters = Map("tableName" -> s"`$format`", "operation" -> "batch write")) + checkError( + exception = intercept[AnalysisException] { + df.write.format(format).option("table", table) + .option("capabilities", "BATCH_READ, BATCH_WRITE").mode("overwrite").save() + }, + condition = "UNSUPPORTED_FEATURE.TABLE_OPERATION", + parameters = Map("tableName" -> s"`$format`", "operation" -> "truncate in batch mode")) + assert(!TestColumnarSink.writers.contains(table)) + } + + test("build and plan a batch scan once") { + val dataSource = new TestRangeDataSource(options("end" -> "4")) + val builder = new ColumnarScanBuilder(format, dataSource.schema(), () => dataSource) + val scan = builder.build() + assert(builder.build() eq scan) + val batch = scan.toBatch + assert(scan.toBatch eq batch) + val partitions = batch.planInputPartitions() + assert(partitions.length == 2) + assert(batch.planInputPartitions() eq partitions) + assert(batch.createReaderFactory() eq batch.createReaderFactory()) + assert(dataSource.plans.get() == 1) + + // A failure is planned once too, and Spark gets the same error each time it asks. + val failing = new TestRangeDataSource(options("failPlanning" -> "true")) + val failingBatch = + new ColumnarScanBuilder(format, failing.schema(), () => failing).build().toBatch + val e = intercept[IllegalStateException](failingBatch.planInputPartitions()) + assert(e.getMessage == "injected planning failure") + assert(intercept[IllegalStateException](failingBatch.createReaderFactory()) eq e) + assert(failing.plans.get() == 1) + } + + test("abort a failed write") { + val table = "failed_write" + val e = intercept[SparkException] { + Seq((1L, "a")).toDF("id", "name").write.format(format).option("table", table) + .option("fail", "true").mode("append").save() + } + assert(Utils.exceptionString(e).contains("injected failure")) + assert(TestColumnarSink.aborted.contains(table)) + assert(!TestColumnarSink.committed.containsKey(table)) + } + + test("streaming read and write") { + val table = "streaming_write" + val query = spark.readStream.format(format).option("end", "5").load() + .writeStream + .format(format) + .option("table", table) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try { + query.processAllAvailable() + } finally { + query.stop() + } + assert(TestColumnarSink.rows(table).sorted == (0L until 5L).map(id => (id, s"name$id"))) + assert(TestColumnarSink.epochs.get(table) == Seq(0L)) + } + + test("create a stream writer for each micro-batch") { + val table = "micro_batch_write" + val input = MemoryStream[(Long, String)] + val query = input.toDF().toDF("id", "name") + .writeStream + .format(format) + .option("table", table) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try { + (0L until 3L).foreach { id => + input.addData((id, s"name$id")) + query.processAllAvailable() + } + } finally { + query.stop() + } + assert(TestColumnarSink.rows(table).sorted == (0L until 3L).map(id => (id, s"name$id"))) + assert(TestColumnarSink.epochs.get(table) == Seq(0L, 1L, 2L)) + // Each micro-batch creates its stream writer with the ID of the query, so that the native + // data sources close the stream writer of the previous micro-batch of the query. + assert(TestColumnarSink.streamWriterQueryIds.get(table) == Seq.fill(3)(query.id.toString)) + } +} + +/** Reads the ids [0, end) and their names, and writes to [[TestColumnarSink]]. */ +class TestRangeProvider extends TableProvider with DataSourceRegister { + override def shortName(): String = "test_columnar_range" + + override def inferSchema(options: CaseInsensitiveStringMap): StructType = { + new TestRangeDataSource(options).schema() + } + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: util.Map[String, String]): Table = { + new ColumnarTable(shortName(), schema, properties, new TestRangeDataSource(_)) + } + + override def supportsExternalMetadata(): Boolean = true +} + +class TestRangeDataSource(options: CaseInsensitiveStringMap) extends ColumnarDataSource { + private val end = options.getLong("end", 10) + private val table = options.getOrDefault("table", "default") + // The number of batch scans that the planners of this data source planned. + val plans = new AtomicInteger() + + override def schema(): StructType = StructType.fromDDL("id BIGINT, name STRING") + + // The names of the capabilities, such as "BATCH_READ, BATCH_WRITE", or all of them by default. + override def capabilities(): util.Set[TableCapability] = { + Option(options.get("capabilities")).map { names => + val capabilities = util.EnumSet.noneOf(classOf[TableCapability]) + names.split(",").map(_.trim).foreach(name => capabilities.add(TableCapability.valueOf(name))) + capabilities + }.getOrElse(super.capabilities()) + } + + override def batchPlanner(schema: StructType): ColumnarBatchPlanner = { + new TestRangePlanner(end, schema, options.getBoolean("failPlanning", false), plans) + } + + override def streamReader(schema: StructType): ColumnarStreamReader = { + new TestRangeStreamReader(end, schema) + } + + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + TestColumnarSink.writers.add(table) + new TestSinkWriter(table, overwrite, options.getBoolean("fail", false)) + } + + override def streamWriter( + schema: StructType, + overwrite: Boolean, + queryId: String): ColumnarStreamWriter = { + TestColumnarSink.streamWriterQueryIds.compute( + table, (_, old) => Option(old).getOrElse(Nil) :+ queryId) + new TestSinkStreamWriter(table) + } +} + +case class TestRangePartition(lo: Long, hi: Long) extends InputPartition + +/** Accepts the predicates `id > `. */ +class TestRangePlanner(end: Long, schema: StructType, fail: Boolean, plans: AtomicInteger) + extends ColumnarBatchPlanner { + private var lo = 0L + private var limit = Long.MaxValue + private var readSchema = schema + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + predicates.filter { predicate => + (predicate.name(), predicate.children()) match { + case (">", Array(column: NamedReference, literal: Literal[_])) + if column.fieldNames().sameElements(Array("id")) => + lo = math.max(lo, literal.value().asInstanceOf[Long] + 1) + true + case _ => false + } + } + } + + override def pushLimit(limit: Int): Boolean = { + this.limit = limit + true + } + + override def pruneColumns(requiredSchema: StructType): Boolean = { + readSchema = requiredSchema + true + } + + override def plan(): ColumnarPlannedBatch = { + plans.incrementAndGet() + if (fail) { + throw new IllegalStateException("injected planning failure") + } + val hi = if (limit == Long.MaxValue) end else math.min(end, lo + limit) + val mid = (lo + hi) / 2 + ColumnarPlannedBatch( + Array(TestRangePartition(lo, mid), TestRangePartition(mid, hi)), + new TestRangeReaderFactory(readSchema)) + } +} + +class TestRangeReaderFactory(readSchema: StructType) extends ColumnarOnlyPartitionReaderFactory { + override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = { + val range = partition.asInstanceOf[TestRangePartition] + new TestBatchReader(range.lo, range.hi, readSchema) + } +} + +class TestRangeStreamReader(end: Long, schema: StructType) extends ColumnarStreamReader { + private def offset(json: String): Long = json.stripPrefix("{\"offset\":").stripSuffix("}").toLong + + override def initialOffset(): String = "{\"offset\":0}" + + override def latestOffset(): String = s"""{"offset":$end}""" + + override def partitions(start: String, end: String): Array[InputPartition] = { + Array(TestRangePartition(offset(start), offset(end))) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + val range = partition.asInstanceOf[TestRangePartition] + new TestBatchReader(range.lo, range.hi, schema) + } +} + +/** Returns the ids [lo, hi) and their names in a single batch. */ +class TestBatchReader(lo: Long, hi: Long, schema: StructType) + extends PartitionReader[ColumnarBatch] { + private var consumed = false + private lazy val batch = { + val numRows = (hi - lo).toInt + val vectors = OnHeapColumnVector.allocateColumns(math.max(numRows, 1), schema) + for (i <- 0 until numRows; (name, column) <- schema.fieldNames.zipWithIndex) { + name match { + case "id" => vectors(column).putLong(i, lo + i) + case "name" => vectors(column).putByteArray(i, s"name${lo + i}".getBytes(UTF_8)) + } + } + new ColumnarBatch(vectors.map(v => v: ColumnVector), numRows) + } + + override def next(): Boolean = { + val hasNext = !consumed + consumed = true + hasNext + } + + override def get(): ColumnarBatch = batch + + override def close(): Unit = if (consumed) batch.close() +} + +/** The data written by the test writers, by table. */ +object TestColumnarSink { + val committed = new ConcurrentHashMap[String, Seq[TestCommitMessage]]() + val aborted = ConcurrentHashMap.newKeySet[String]() + // The tables for which a writer was created. + val writers = ConcurrentHashMap.newKeySet[String]() + val epochs = new ConcurrentHashMap[String, Seq[Long]]() + val streamWriterQueryIds = new ConcurrentHashMap[String, Seq[String]]() + + def rows(table: String): Seq[(Long, String)] = committed.get(table).flatMap(_.rows) + + def batchSizes(table: String): Seq[Int] = committed.get(table).flatMap(_.batchSizes) + + def commit(table: String, messages: Array[WriterCommitMessage], overwrite: Boolean): Unit = { + val newMessages = messages.map(_.asInstanceOf[TestCommitMessage]).toSeq + committed.compute(table, (_, old) => { + if (old == null || overwrite) newMessages else old ++ newMessages + }) + } +} + +case class TestCommitMessage(rows: Seq[(Long, String)], batchSizes: Seq[Int]) + extends WriterCommitMessage + +class TestDataWriter(fail: Boolean) extends DataWriter[ColumnarBatch] { + private val rows = ArrayBuffer.empty[(Long, String)] + private val batchSizes = ArrayBuffer.empty[Int] + + override def write(batch: ColumnarBatch): Unit = { + if (fail) { + throw new IllegalStateException("injected failure") + } + assert((0 until batch.numCols()).forall(batch.column(_).isInstanceOf[ArrowColumnVector])) + batchSizes += batch.numRows() + (0 until batch.numRows()).foreach { i => + rows += ((batch.column(0).getLong(i), batch.column(1).getUTF8String(i).toString)) + } + } + + override def commit(): WriterCommitMessage = TestCommitMessage(rows.toSeq, batchSizes.toSeq) + + override def abort(): Unit = {} + + override def close(): Unit = {} +} + +class TestSinkWriter(table: String, overwrite: Boolean, fail: Boolean) extends ColumnarWriter { + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new TestDataWriter(fail) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = { + TestColumnarSink.commit(table, messages, overwrite) + } + + override def abort(messages: Array[WriterCommitMessage]): Unit = { + TestColumnarSink.aborted.add(table) + } +} + +class TestSinkStreamWriter(table: String) extends ColumnarStreamWriter { + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[ColumnarBatch] = new TestDataWriter(fail = false) + + override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + TestColumnarSink.commit(table, messages, overwrite = false) + TestColumnarSink.epochs.compute(table, (_, old) => Option(old).getOrElse(Nil) :+ epochId) + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + TestColumnarSink.aborted.add(table) + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala new file mode 100644 index 000000000000..bc7f369ba930 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala @@ -0,0 +1,148 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.io.File +import java.nio.file.Files + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.classic.{DataFrame, SparkSession} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.tags.ExtendedSQLTest +import org.apache.spark.util.Utils + +/** + * Tests native data sources in local-cluster mode, where the executors run in JVMs of their own: + * Spark distributes the packages to them, and they find, check and load the libraries + * themselves. [[NativeDataSourceSuite]] tests the rest in local mode, where the executors use the + * packages of the driver. + */ +@ExtendedSQLTest +class NativeDataSourceClusterSuite extends QueryTest { + import NativeDataSourceTestUtils._ + + protected var spark: SparkSession = _ + + // Jobs outside of SQL executions get the files of all the sessions, so the sessions are kept + // until the end of the suite, as in NativeDataSourceSuite. + private val sessions = ArrayBuffer.empty[SparkSession] + + private lazy val library = compile("test_native_datasource") + // Another build, so that its package has another checksum, which the executors did not load. + private lazy val otherLibrary = + compile("test_native_datasource_cluster", "-DNAME_PREFIX=\"cluster_\"") + + override def beforeAll(): Unit = { + super.beforeAll() + spark = SparkSession.builder() + .master("local-cluster[2,1,1024]") + .appName("NativeDataSourceClusterSuite") + .getOrCreate() + } + + override def afterAll(): Unit = { + try { + sessions.clear() + if (spark != null) { + spark.stop() + spark = null + } + } finally { + super.afterAll() + } + } + + private def nativeTest(name: String)(f: => Unit): Unit = test(name) { + assume(toolchain.isDefined, "A C++ compiler and Java 22 or later are required.") + try f finally SparkSession.setActiveSession(spark) + } + + // Like Python data sources, native data sources are found in the active session. + private def newSession(): SparkSession = { + val session = spark.newSession() + SparkSession.setActiveSession(session) + sessions += session + session + } + + private def checkRows(df: => DataFrame, expected: Seq[Row]): Unit = { + QueryTest.checkAnswer(df, expected, checkToRDD = false) + } + + private def written(dir: File): Seq[String] = { + dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala).toSeq.sorted + } + + nativeTest("read and write with a package added as an artifact") { + val session = newSession() + session.addArtifact(createPackage(library, dataSources("")).getPath) + // 4 partitions, so that both executors load the library. + val df = session.read.format("native_range") + .option("end", "20").option("partitions", "4").load().select("id") + checkRows(df, (0L until 20L).map(Row(_))) + + val dir = Utils.createTempDir() + session.range(10).repartition(4).selectExpr("id", "concat('v', id) AS name") + .write.format("native_sink").option("path", dir.getPath).mode("append").save() + assert(written(dir) == (0 until 10).map(i => s"$i,v$i").sorted) + + val streamDir = Utils.createTempDir() + val query = session.readStream.format("native_counter") + .option("max_offset", "12").option("step", "5").load() + .writeStream.format("native_sink") + .option("path", streamDir.getPath) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try { + query.processAllAvailable() + } finally { + query.stop() + } + assert(written(streamDir) == (0 until 12).map(_.toString).sorted) + } + + nativeTest("executors load the distributed copy of a configured package") { + val session = newSession() + val pkg = createPackage(otherLibrary, dataSources("cluster_")) + session.conf.set(SQLConf.NATIVE_DATA_SOURCE_PATHS.key, pkg.getPath) + // Prepares the package on the driver as a query would: Spark reads its manifest, loads its + // library and copies it to the artifacts of the session under a name with its checksum. + val found = NativeDataSourceRegistry.lookup("cluster_native_range", session.sessionState.conf) + .map(_._2).get.asInstanceOf[NativeDataSourcePackage] + NativeLibraries.get(found) + val distributed = + NativeDataSourceRegistry.distribute(found).asInstanceOf[NativeDataSourcePackage] + assert(distributed.artifactUUID.contains(session.sessionUUID)) + assert(distributed.distributedFileName.contains(s"${pkg.getName}.${found.sha256}")) + + // Then the file at the path of the driver is no longer a valid package, so the executors + // cannot use it, while the driver keeps the package it read, since the file is written in + // place and keeps its length and modification time. So the query only succeeds if the + // executors find the distributed copy. NativeDataSourceSuite checks that the executors only + // load a valid package with the checksum of the driver. + val lastModified = pkg.lastModified() + Files.write(pkg.toPath, new Array[Byte](pkg.length().toInt)) + assert(pkg.setLastModified(lastModified)) + val df = session.read.format("cluster_native_range") + .option("end", "20").option("partitions", "4").load().select("id") + checkRows(df, (0L until 20L).map(Row(_))) + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala new file mode 100644 index 000000000000..a13a8af7a5b4 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -0,0 +1,2006 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.io.File +import java.nio.charset.StandardCharsets.UTF_8 +import java.nio.file.Files +import java.nio.file.StandardCopyOption.REPLACE_EXISTING +import java.sql.{Date, Timestamp} +import java.time.{Instant, LocalDate} +import java.util.UUID +import java.util.concurrent.{Callable, CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicReference + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import com.fasterxml.jackson.core.json.JsonWriteFeature +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.json.JsonMapper +import org.scalatest.time.SpanSugar._ + +import org.apache.spark.{SPARK_VERSION, SparkClassNotFoundException, SparkFiles, SparkRuntimeException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.api.python.{PythonBroadcast, SimplePythonFunction} +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.sql.{AnalysisException, IntegratedUDFTestUtils, QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.{Literal => CatalystLiteral} +import org.apache.spark.sql.catalyst.plans.logical.Filter +import org.apache.spark.sql.classic.{DataFrame, SparkSession} +import org.apache.spark.sql.connector.expressions.{Expression, FieldReference, GeneralScalarExpression, LiteralValue} +import org.apache.spark.sql.connector.expressions.filter.{And, Predicate} +import org.apache.spark.sql.execution.WholeStageCodegenExec +import org.apache.spark.sql.execution.datasources.DataSource +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2ScanRelation, DataSourceV2Strategy} +import org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarBatchScan, ColumnarScan} +import org.apache.spark.sql.execution.datasources.v2.python.{PythonDataSourceV2, UserDefinedPythonDataSource} +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.sql.streaming.{StreamingQuery, Trigger} +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types._ +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.unsafe.types.UTF8String +import org.apache.spark.util.{ThreadUtils, Utils} + +/** + * Tests native data sources with the library in `native-datasource/test_native_datasource.cc`, + * compiled with the C++ compiler of the machine by [[NativeDataSourceTestUtils]]. The tests that + * load a library are skipped when there is no C++ compiler, or before Java 22. + */ +class NativeDataSourceSuite extends QueryTest with SharedSparkSession { + import NativeDataSourceTestUtils._ + + // Jobs outside of SQL executions get the files of all the sessions, so the sessions are kept + // until the end of the suite: none of their artifacts is removed while another test runs. + private val sessions = ArrayBuffer.empty[SparkSession] + + // The test library and its variants, each implementing differently named data sources. + private lazy val library = compile("test_native_datasource") + private lazy val readOnlyLibrary = + compile("test_native_datasource_ro", "-DREAD_ONLY", "-DNAME_PREFIX=\"ro_\"") + private lazy val negatingLibrary = + compile("test_native_datasource_b", "-DNAME_PREFIX=\"b_\"", "-DID_SIGN=-1") + // Implements the data sources of the test library, as another version of it would. + private lazy val renamedNegatingLibrary = + compile("test_native_datasource_negating", "-DID_SIGN=-1") + private lazy val abiV2Library = + compile("test_native_datasource_v2", "-DABI_VERSION=2", "-DNAME_PREFIX=\"v2_\"") + private lazy val noLastErrorLibrary = + compile("test_native_datasource_ne", "-DNO_LAST_ERROR", "-DNAME_PREFIX=\"ne_\"") + private lazy val smallTableLibrary = compile("test_native_datasource_small", + "-DTABLE_SIZE=offsetof(SparkDsLibrary,read)", "-DNAME_PREFIX=\"small_\"") + private lazy val failingInitLibrary = + compile("test_native_datasource_if", "-DINIT_FAILURE=1", "-DNAME_PREFIX=\"if_\"") + private lazy val noTableLibrary = + compile("test_native_datasource_nt", "-DINIT_FAILURE=2", "-DNAME_PREFIX=\"nt_\"") + private lazy val noInitLibrary = + compile("test_native_datasource_ni", "-DNO_INIT", "-DNAME_PREFIX=\"ni_\"") + private lazy val installedLibrary = + compile("test_native_datasource_installed", "-DNAME_PREFIX=\"installed_\"") + private lazy val installedNegatingLibrary = compile( + "test_native_datasource_installed_b", "-DNAME_PREFIX=\"installed_\"", "-DID_SIGN=-1") + + private lazy val defaultPackage = createPackage(library, dataSources("")) + + /** A new session that finds the native data source packages at the given paths. */ + // Like Python data sources, native data sources are found in the active session. + private def newSession(paths: File*): SparkSession = { + val session = spark.newSession() + SparkSession.setActiveSession(session) + sessions += session + session.conf.set(SQLConf.NATIVE_DATA_SOURCE_PATHS.key, paths.map(_.getPath).mkString(",")) + session + } + + private def nativeTest(name: String)(f: => Unit): Unit = test(name) { + assume(toolchain.isDefined, "A C++ compiler and Java 22 or later are required.") + try f finally SparkSession.setActiveSession(spark) + } + + // Checks the answer without converting the DataFrame to an RDD: a job outside of a SQL + // execution gets the files of all the sessions, and the streaming queries of the tests create + // sessions whose files are removed when they are garbage collected. + private def checkRows(df: => DataFrame, expected: Seq[Row]): Unit = { + QueryTest.checkAnswer(df, expected, checkToRDD = false) + } + + private def range(session: SparkSession, options: (String, String)*): DataFrame = { + session.read.format("native_range").options(options.toMap).load() + } + + private def scan(df: DataFrame): ColumnarScan = { + df.queryExecution.optimizedPlan.collectFirst { + case relation: DataSourceV2ScanRelation => relation.scan.asInstanceOf[ColumnarScan] + }.get + } + + private def expectedRow(id: Long): Row = Row( + id, + (id % 3).toInt, + if (id % 5 == 0) null else s"n$id", + id * 0.5, + id % 2 == 0, + Date.valueOf(LocalDate.ofEpochDay(id)), + Timestamp.from(Instant.ofEpochSecond(id))) + + /** Finds the error with the given condition among the causes of a failure. */ + private def findError(e: Throwable, condition: String): SparkThrowable = { + Iterator.iterate(e)(_.getCause).takeWhile(_ != null) + .collectFirst { case t: SparkThrowable if t.getCondition == condition => t } + .getOrElse(throw e) + } + + /** + * Checks that Spark closed each data writer that native_sink created under `dir`, and did not + * use it afterwards. When a write fails, the tasks that Spark kills may still close their data + * writers after the write returns. + */ + private def checkDataWritersClosed(dir: File): Unit = eventually(timeout(1.minute)) { + def names(subdir: String): Set[String] = + Option(new File(dir, subdir).list()).map(_.toSet).getOrElse(Set.empty) + assert(names("_created_data_writers").nonEmpty) + assert(names("_closed_data_writers") == names("_created_data_writers")) + assert(!new File(dir, "_DATA_WRITER_USED_AFTER_CLOSE").exists()) + } + + // Formats a double as the test library does. + private def format(d: Double): String = if (d == d.toLong) d.toLong.toString else d.toString + + test("parse manifests") { + assert(NativeDataSourcePackage.parseManifest( + """{"abiVersion": 1, "dataSources": ["a", "b"], "libraries": {"linux-x86_64": "x.so"}}""") == + Right(NativeDataSourceManifest(1, Seq("a", "b"), Map("linux-x86_64" -> "x.so")))) + Seq( + "not json" -> "The manifest is not valid JSON", + "[1]" -> "The manifest must be a JSON object.", + """{"dataSources": ["a"], "libraries": {"p": "x"}}""" -> "'abiVersion'", + """{"abiVersion": 0, "dataSources": ["a"], "libraries": {"p": "x"}}""" -> + "positive integer field 'abiVersion'", + // The JSON escape of the character U+0000. + ("""{"abiVersion": 1, "dataSources": ["a""" + "\\u0000" + + """"], "libraries": {"p": "x"}}""") -> "'dataSources'", + """{"abiVersion": 1, "dataSources": [], "libraries": {"p": "x"}}""" -> "'dataSources'", + """{"abiVersion": 1, "dataSources": [1], "libraries": {"p": "x"}}""" -> "'dataSources'", + """{"abiVersion": 1, "dataSources": ["a"]}""" -> "'libraries'", + """{"abiVersion": 1, "dataSources": ["a"], "libraries": {"p": 1}}""" -> "'libraries'" + ).foreach { case (json, error) => + val result = NativeDataSourcePackage.parseManifest(json) + assert(result.isLeft && result.swap.toOption.get.contains(error), s"$json: $result") + } + } + + test("platform names") { + assert(NativePlatform.name("Linux", "amd64") == "linux-x86_64") + assert(NativePlatform.name("Linux", "aarch64") == "linux-aarch64") + assert(NativePlatform.name("Mac OS X", "aarch64") == "osx-aarch64") + assert(NativePlatform.name("Mac OS X", "x86_64") == "osx-x86_64") + assert(NativePlatform.name("Windows 11", "amd64") == "windows-x86_64") + // Linux with musl, as on Alpine Linux, which maps its C library as the loader of musl. + assert(NativePlatform.name("Linux", "amd64", musl = true) == "linux-x86_64-musl") + assert(NativePlatform.name("Linux", "aarch64", musl = true) == "linux-aarch64-musl") + assert(NativePlatform.usesMusl(Seq( + "7f00a000-7f00b000 r-xp 00000000 08:01 123 /usr/lib/jvm/lib/server/libjvm.so", + "7f00c000-7f00d000 r-xp 00000000 08:01 456 /lib/ld-musl-x86_64.so.1"))) + assert(!NativePlatform.usesMusl(Seq( + "7f00c000-7f00d000 r-xp 00000000 08:01 456 /usr/lib/x86_64-linux-gnu/libc.so.6"))) + // The tests run on glibc, macOS or Windows. + assert(NativePlatform.current == NativePlatform.name( + System.getProperty("os.name"), System.getProperty("os.arch"))) + } + + test("fail clearly where the C data interface of Arrow Java does not load") { + // Arrow Java fails like this without a library for the architecture, such as ppc64le, or + // with one that does not link, such as on musl. + Seq( + new IllegalStateException("error loading native libraries: java.io.FileNotFoundException: " + + "arrow_cdata_jni/ppc64le/libarrow_cdata_jni.so"), + new UnsatisfiedLinkError("libarrow_cdata_jni.so: Error relocating: symbol not found") + ).foreach { cause => + val e = intercept[SparkUnsupportedOperationException] { + NativeLibraries.checkCDataInterface(throw cause) + } + checkError( + exception = e, + condition = "NATIVE_DATA_SOURCE_UNSUPPORTED_PLATFORM", + parameters = Map("platform" -> NativePlatform.current, "reason" -> cause.toString)) + assert(e.getCause eq cause) + } + // It loads where the tests run. + NativeLibraries.checkCDataInterface() + } + + test("native library path") { + def path(directories: String*): String = directories.mkString(File.pathSeparator) + // The directories of java.library.path, and then the lib directory of each installation + // prefix in PATH. Relative and duplicate directories are ignored. + assert(NativeDataSourceRegistry.librarySearchPath( + path("/opt/native", "", "relative", "/usr/lib"), + path("/usr/local/bin", "/home/user/.venv/bin", "/usr/sbin", "relative/bin", "/usr/local/bin") + ) == Seq("/opt/native", "/usr/lib", "/usr/local/lib", "/home/user/.venv/lib").map(new File(_))) + assert(NativeDataSourceRegistry.librarySearchPath(null, null).isEmpty) + assert(NativeDataSourceRegistry.libraryFileName("My_Source") == + System.mapLibraryName("spark_datasource_my_source")) + } + + test("encode predicates as JSON") { + def predicate(name: String, children: Expression*): Predicate = { + new Predicate(name, children.toArray) + } + val id = FieldReference(Seq("id")) + val schema = StructType.fromDDL("id BIGINT, a STRUCT, s STRING, " + + "c STRING COLLATE UTF8_LCASE, d STRUCT>") + def toJson(predicate: Predicate): Option[String] = NativePredicates.toJson(predicate, schema) + assert(toJson(predicate(">", id, LiteralValue(5L, LongType))).contains( + """{"type":"function","name":">","children":[{"type":"column","name":["id"]},""" + + """{"type":"literal","dataType":"bigint","value":5}]}""")) + assert(toJson(new And( + predicate("IS_NULL", FieldReference(Seq("a", "b"))), + predicate("IN", id, LiteralValue(1L, LongType), LiteralValue(2L, LongType)) + )).contains( + """{"type":"function","name":"AND","children":[{"type":"function","name":"IS_NULL",""" + + """"children":[{"type":"column","name":["a","b"]}]},{"type":"function","name":"IN",""" + + """"children":[{"type":"column","name":["id"]},""" + + """{"type":"literal","dataType":"bigint","value":1},""" + + """{"type":"literal","dataType":"bigint","value":2}]}]}""")) + // Only the functions listed in push_predicates are passed, so not arithmetic. + NativePredicates.FUNCTION_NAMES.foreach { name => + assert(toJson(predicate(name, id, id)).nonEmpty, name) + } + val plusOne = new GeneralScalarExpression("+", Array(id, LiteralValue(1, IntegerType))) + assert(toJson(predicate("=", id, plusOne)).isEmpty) + assert(toJson(new And(predicate("IS_NULL", id), predicate("=", id, plusOne))).isEmpty) + assert(toJson(predicate("UPPER", FieldReference(Seq("s")))).isEmpty) + + def literal(value: Any, dataType: DataType): Option[String] = { + toJson(predicate("=", id, LiteralValue(value, dataType))) + .map(json => json.substring(json.indexOf("\"dataType\""), json.length - 3)) + } + assert(literal(null, LongType).contains(""""dataType":"bigint","value":null""")) + assert(literal(true, BooleanType).contains(""""dataType":"boolean","value":true""")) + assert(literal(1.toByte, ByteType).contains(""""dataType":"tinyint","value":1""")) + assert(literal(2.toShort, ShortType).contains(""""dataType":"smallint","value":2""")) + assert(literal(1.5f, FloatType).contains(""""dataType":"float","value":1.5""")) + // As the double of the same value, which a reader that parses JSON numbers as doubles gets + // exactly: parsing 0.1 as a double gives another value than the float 0.1. + assert(literal(0.1f, FloatType) + .contains(""""dataType":"float","value":0.10000000149011612""")) + assert(literal(2.5, DoubleType).contains(""""dataType":"double","value":2.5""")) + assert(literal(UTF8String.fromString("a\"b"), StringType) + .contains(""""dataType":"string","value":"a\"b"""")) + // The JSON is ASCII, with the other characters escaped, such as the characters outside the + // Basic Multilingual Plane. + val text = "x" + new String(Character.toChars(0x1F600)) + new String(Character.toChars(0xE9)) + val encoded = literal(UTF8String.fromString(text), StringType).get + assert(encoded.forall(_ < 128), encoded) + assert(new ObjectMapper().readTree(s"{$encoded}").get("value").asText() == text) + assert(literal(Array[Byte](1, 2), BinaryType) + .contains(""""dataType":"binary","value":"AQI="""")) + assert(literal(Decimal("12.30"), DecimalType(10, 2)) + .contains(""""dataType":"decimal(10,2)","value":"12.30"""")) + assert(literal(10, DateType).contains(""""dataType":"date","value":10""")) + assert(literal(10L, TimestampType).contains(""""dataType":"timestamp","value":10""")) + assert(literal(10L, TimestampNTZType).contains(""""dataType":"timestamp_ntz","value":10""")) + // Values that JSON cannot represent, and types without an encoding, are not pushed down. + assert(literal(Double.NaN, DoubleType).isEmpty) + assert(literal(Float.PositiveInfinity, FloatType).isEmpty) + assert(literal(UTF8String.fromString("a"), StringType("UTF8_LCASE")).isEmpty) + // A JSON string would replace the invalid bytes, which Spark compares. + assert(literal(UTF8String.fromBytes(Array(0xFF.toByte)), StringType).isEmpty) + assert(literal(null, ArrayType(IntegerType)).nonEmpty) + assert(literal(Array(1), ArrayType(IntegerType)).isEmpty) + + // A source compares strings by their bytes, so the columns with strings of a non-binary + // collation are not pushed down, and neither are the columns that are not in the schema. + def compare(left: String, right: String): Option[String] = { + toJson(predicate("=", FieldReference(left.split('.').toSeq), + FieldReference(right.split('.').toSeq))) + } + assert(compare("s", "s").nonEmpty) + assert(compare("a.b", "id").nonEmpty) + assert(compare("c", "c").isEmpty) + assert(compare("s", "c").isEmpty) + assert(compare("d", "d").isEmpty) + assert(compare("d.e", "d.e").isEmpty) + assert(compare("missing", "id").isEmpty) + assert(compare("a.missing", "id").isEmpty) + assert(compare("id.b", "id").isEmpty) + } + + /** + * Generates `org/apache/spark/sql/datasource/predicate-cases.json`, which lists predicates as + * Spark passes them to `push_predicates`, with the rows of a table for which Spark + * finds them true, for libraries to test their pushdown with. Regenerate it with + * `SPARK_GENERATE_GOLDEN_FILES=1`. + */ + test("predicate cases") { + val schema = StructType.fromDDL("id INT, i INT, d DOUBLE, s STRING, f FLOAT, b BINARY") + def char(codePoint: Int): String = new String(Character.toChars(codePoint)) + val (emoji, eAcute, replacement) = (char(0x1F600), char(0xE9), char(0xFFFD)) + def bytes(values: Int*): Array[Byte] = values.map(_.toByte).toArray + val rows = Seq( + Row(0, 1, 0.0, "a", 0.1f, bytes(0x00)), + Row(1, 2, -0.0, "", 1.5f, bytes(0x7F)), + Row(2, null, Double.NaN, eAcute, null, bytes(0x80)), + Row(3, 3, Double.PositiveInfinity, emoji, 0.1f, bytes(0xFF)), + Row(4, 1, Double.NegativeInfinity, null, 0.5f, bytes(0x7F, 0x00)), + Row(5, null, 1.5, "ab", 2.0f, bytes()), + Row(6, 2, null, "b" + eAcute, 0.25f, null), + Row(7, 4, 2.5, replacement, 3.0f, bytes(0x01, 0x02))) + val filters = Seq( + // NaN is equal to NaN and greater than any other value, and -0.0 is equal to 0.0. + "d = 0.0D", "d = -0.0D", "d < 0.0D", "d > 1.0D", "d >= 2.5D", "d <> 0.0D", "d <=> 0.0D", + "d IN (0.0D, 1.5D)", "NOT (d > 1.0D)", "NOT (d < 1.0D)", "d IS NULL", "d IS NOT NULL", + // Comparisons with null are null, also under NOT. + "i = 1", "i = NULL", "i <=> NULL", "NOT (i <=> NULL)", "i IN (1, NULL)", + "NOT (i IN (1, NULL))", "NOT (i IN (1, 3))", "i IS NULL OR d > 1.0D", + "i > 1 AND d IS NOT NULL", + // Strings are compared by their UTF-8 bytes. + s"s = '$emoji'", s"s > '$replacement'", "s < 'b'", "startswith(s, 'a')", + s"endswith(s, '$eAcute')", s"contains(s, '$emoji')", s"s IN ('', '$eAcute')", + // A float literal is the exact value of the float, which a double holds. + "f = CAST(0.1 AS FLOAT)", "f > CAST(0.1 AS FLOAT)", + // Binary values are compared as unsigned bytes: X'80', X'FF' and X'7F00', which starts with + // X'7F', are greater than X'7F'. + "b > X'7F'", "b < X'80'", + // Taking a part of a predicate under = as true would drop the row whose id is 3. + "(id > 3) = (id < 3)") + + val mapper = JsonMapper.builder().enable(JsonWriteFeature.ESCAPE_NON_ASCII).build() + def json(value: Any): String = mapper.writeValueAsString(value) + def jsonValue(value: Any): String = value match { + case d: Double if d.isNaN || d.isInfinite => json(d.toString) + // The exact value of the float, as the literals of the predicates. + case f: Float => json(f.toDouble) + case other => json(other) + } + val df = spark.createDataFrame(rows.asJava, schema) + val cases = filters.map { filter => + val filtered = df.where(filter) + // The predicates that Spark pushes for the filter: its conjuncts, with constants folded. + val condition = filtered.queryExecution.analyzed.collectFirst { + case Filter(condition, _) => condition + }.get + val predicates = splitConjunctivePredicates(condition.transformUp { + case e if e.foldable && !e.isInstanceOf[CatalystLiteral] => + CatalystLiteral.create(e.eval(), e.dataType) + }).map { predicate => + DataSourceV2Strategy.translateFilterV2(predicate) + .flatMap(NativePredicates.toJson(_, schema)) + .getOrElse(fail(s"$filter is not pushed down")) + } + val matches = filtered.select("id").collect().map(_.getInt(0)).sorted + s""" {"sql": ${json(filter)}, "predicates": [${predicates.mkString(", ")}], """ + + s""""rows": [${matches.mkString(", ")}]}""" + } + val generated = + s"""{ + | "description": ${json("Predicates as Spark passes them to push_predicates of " + + "spark_datasource.h, with the ids of the rows of the table for which they are true. A " + + "case lists the predicates that Spark pushes for the filter in its SQL, which are " + + "combined with AND. A row has a value for each column of the schema; NaN and infinite " + + "doubles are strings, floats are written as the doubles of the same values, and binary " + + "values are base64 strings.")}, + | "schema": ${json(schema.toDDL)}, + | "rows": [ + |${rows.map(row => row.toSeq.map(jsonValue).mkString(" [", ", ", "]")).mkString(",\n")} + | ], + | "cases": [ + |${cases.mkString(",\n")} + | ] + |} + |""".stripMargin + assert(generated.forall(_ < 128)) + val file = getWorkspaceFilePath("sql", "core", "src", "main", "resources", "org", "apache", + "spark", "sql", "datasource", "predicate-cases.json").toFile + if (regenerateGoldenFiles) { + Files.createDirectories(file.toPath.getParent) + Files.writeString(file.toPath, generated) + } else { + assert(Files.readString(file.toPath) == generated, + s"$file is out of date: regenerate it with SPARK_GENERATE_GOLDEN_FILES=1.") + } + // The file is in the jar of Spark SQL. + assert(getClass.getClassLoader.getResource( + "org/apache/spark/sql/datasource/predicate-cases.json") != null) + } + + test("comparisons and IN on arrays and structs are not pushed") { + val schema = StructType.fromDDL( + "a ARRAY, a2 ARRAY, s STRUCT, s2 STRUCT") + val df = spark.createDataFrame(Seq.empty[Row].asJava, schema) + def toJson(filter: String): Option[String] = { + val condition = df.where(filter).queryExecution.analyzed.collectFirst { + case Filter(condition, _) => condition + }.get + val predicate = DataSourceV2Strategy.translateFilterV2(condition) + .getOrElse(fail(s"$filter is not a Data Source V2 predicate")) + NativePredicates.toJson(predicate, schema) + } + // Spark compares these values with other rules than push_predicates describes, for example + // with the nulls in them equal to each other: array(1, null) = array(1, null) is true. + Seq("a = a2", "a < a2", "a IN (a2)", "s = s2", "s <=> s2").foreach { filter => + assert(toJson(filter).isEmpty, s"$filter is pushed down") + } + // Null checks apply to any type, and a field of a struct has its own type. + Seq("a IS NULL", "s IS NOT NULL", "s.x = 1", "s.x IN (1, 2)").foreach { filter => + assert(toJson(filter).isDefined, s"$filter is not pushed down") + } + } + + test("tell missing functions from link errors") { + // A function that the library leaves NULL in its table is an unsupported operation. + val missing = new NativeLibrary.MissingFunctionException("push_limit") + assert(NativeCalls.optional(throw missing).isEmpty) + checkError( + exception = intercept[SparkRuntimeException](NativeCalls.call("x", "read")(throw missing)), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "read", "name" -> "x", "msg" -> "The library does not implement push_limit.")) + // Link errors are reported as they are, such as when Arrow, which still calls the Arrow C + // structs through JNI, cannot load its JNI library, which it extracts to java.io.tmpdir, for + // example if the directory does not allow it. + val linkError = new UnsatisfiedLinkError("/tmp/libarrow_cdata_jni.so: failed to map segment") + assert(intercept[UnsatisfiedLinkError](NativeCalls.optional(throw linkError)) eq linkError) + checkError( + exception = intercept[SparkRuntimeException](NativeCalls.call("x", "read")(throw linkError)), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map("action" -> "read", "name" -> "x", "msg" -> linkError.getMessage)) + } + + test("close the previous stream writer of a streaming query") { + val released = ArrayBuffer.empty[Long] + def handle(value: Long): NativeHandle = new NativeHandle(value, released += _) + val (first, other, second) = (handle(1), handle(2), handle(3)) + NativeStreamWriters.register("query", first) + NativeStreamWriters.register("other query", other) + assert(released.isEmpty) + NativeStreamWriters.register("query", second) + assert(released == Seq(1L)) + second.close() + other.close() + assert(released == Seq(1L, 3L, 2L)) + } + + test("native data sources require Java 22 or later") { + assume(!NativeLibrary.isSupported, "The Java version supports native data sources.") + val dir = Utils.createTempDir() + // Not loaded: Spark fails before it loads the library. + Files.writeString( + new File(dir, NativeDataSourceRegistry.libraryFileName("old_java")).toPath, "unused") + withLibraryPath(dir) { + checkError( + exception = intercept[SparkUnsupportedOperationException] { + spark.read.format("old_java").load() + }, + condition = "NATIVE_DATA_SOURCE_UNSUPPORTED_JAVA_VERSION", + parameters = Map( + "required" -> "22", "current" -> Runtime.version().feature().toString)) + } + } + + test("the bindings match the layout of spark_datasource.h") { + assume(compiler.isDefined, "A C++ compiler is required.") + val layout = NativeLibrary.layout().asScala.map { case (key, value) => key -> value.longValue } + val prints = layout.keys.map { key => + val value = key.split('.') match { + case Array(struct) => s"sizeof($struct)" + case Array(struct, field) => s"offsetof($struct, $field)" + } + s""" std::printf("%s %zu\\n", "$key", $value);""" + } + val output = runProgram("layout", + s"""#include + |#include + |#include "spark_datasource.h" + | + |int main() { + |${prints.mkString("\n")} + | return 0; + |} + |""".stripMargin) + val actual = output.linesIterator.map { line => + val Array(key, value) = line.split(" ") + key -> value.toLong + }.toMap + assert(actual == layout.toMap) + } + + nativeTest("disable native data sources") { + val conf = new SQLConf + conf.setConf(SQLConf.NATIVE_DATA_SOURCE_PATHS, Seq(defaultPackage.getPath)) + conf.setConf(StaticSQLConf.NATIVE_DATA_SOURCE_ENABLED, false) + assert(NativeDataSourceRegistry.lookup("native_range", conf).isEmpty) + } + + nativeTest("read with an inferred schema") { + val session = newSession(defaultPackage) + val df = range(session, "end" -> "20") + assert(df.schema == StructType.fromDDL( + "id BIGINT, mod INT, name STRING, value DOUBLE, flag BOOLEAN, day DATE, ts TIMESTAMP")) + checkRows(df, (0L until 20L).map(expectedRow)) + // The data source names are case-insensitive, and so are the option keys. + checkRows( + session.read.format("NATIVE_RANGE").option("END", "3").load(), + (0L until 3L).map(expectedRow)) + } + + nativeTest("read in partitions and batches") { + val df = range(newSession(defaultPackage), "end" -> "50", "partitions" -> "4", + "batch_size" -> "3") + assert(df.rdd.getNumPartitions == 4) + checkRows(df, (0L until 50L).map(expectedRow)) + } + + nativeTest("read with a user-specified schema") { + val session = newSession(defaultPackage) + checkRows( + session.read.schema("id BIGINT, name STRING").format("native_range").option("end", "6") + .load(), + (0L until 6L).map(id => Row(id, if (id % 5 == 0) null else s"n$id"))) + + // The Arrow schema does not have the collation, which applies to the data that Spark reads. + val collated = session.read.schema("id BIGINT, name STRING COLLATE UTF8_LCASE") + .format("native_range").option("end", "6").load() + checkRows(collated.filter("name = 'N1'"), Seq(Row(1L, "n1"))) + assert(scan(collated.filter("name = 'N1'")).getMetaData()("PushedPredicates") == "[]") + + val e = intercept[Exception] { + session.read.schema("id INT").format("native_range").load().collect() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "read", + "name" -> "native_range", + "msg" -> "The data has the schema id BIGINT, but the expected schema is id INT.")) + } + + nativeTest("push down predicates") { + val df = range(newSession(defaultPackage), "end" -> "20").filter("id > 5 AND id <= 12") + checkRows(df, (6L to 12L).map(expectedRow)) + val pushed = scan(df).getMetaData()("PushedPredicates") + assert(pushed.contains("id > 5") && pushed.contains("id <= 12"), pushed) + // The scan is described by the name of the data source. + assert(scan(df).description() == "native_range") + assert(!df.queryExecution.executedPlan.toString.contains("NativeDataSource")) + + // A predicate that the library does not accept is evaluated by Spark. + val partly = range(newSession(defaultPackage), "end" -> "20") + .filter("id < 10 AND name LIKE 'n1%'") + checkRows(partly, Seq(1L).map(expectedRow)) + assert(scan(partly).getMetaData()("PushedPredicates").contains("id < 10")) + + // Spark evaluates all the predicates again after reading, even the ones that the library + // reports as evaluated completely: here, a library that ignores them. + val claimed = range(newSession(defaultPackage), "end" -> "20", "claim_predicates" -> "true") + .filter("id > 15") + checkRows(claimed, (16L until 20L).map(expectedRow)) + assert(scan(claimed).getMetaData()("PushedPredicates").contains("id > 15")) + + // A library that only uses the predicates to skip data, here by reading the blocks of 10 + // rows that may match, with all their rows. Spark shows these predicates as pushed too, and + // so it does with a value that it does not know. + Seq("partial", "7").foreach { use => + val skipping = range(newSession(defaultPackage), "end" -> "45", "predicate_use" -> use) + .filter("id > 13 AND id < 27") + checkRows(skipping, (14L until 27L).map(expectedRow)) + val pushed = scan(skipping).getMetaData()("PushedPredicates") + assert(pushed.contains("id > 13") && pushed.contains("id < 27"), pushed) + } + + // The answers of the library are for the predicates that Spark passes to it: here, not the + // first one, which uses another function. + val mapped = range(newSession(defaultPackage), "end" -> "20") + .filter("upper(name) = 'N1' AND id < 5") + checkRows(mapped, Seq(expectedRow(1L))) + assert(scan(mapped).getMetaData()("PushedPredicates") == "[id < 5]") + } + + nativeTest("prune columns") { + val df = range(newSession(defaultPackage), "end" -> "10").select("name", "id") + checkRows(df, (0L until 10L).map(id => Row(if (id % 5 == 0) null else s"n$id", id))) + assert(scan(df).readSchema() == StructType.fromDDL("id BIGINT, name STRING")) + + // The columns of the pushed predicates are read too, since Spark evaluates the predicates + // again after reading, even the ones that the library evaluates completely. + val filtered = range(newSession(defaultPackage), "end" -> "10").filter("id > 7").select("name") + checkRows(filtered, Seq(Row("n8"), Row("n9"))) + assert(scan(filtered).getMetaData()("PushedPredicates").contains("id > 7")) + assert(scan(filtered).readSchema() == StructType.fromDDL("id BIGINT, name STRING")) + + // A count reads no column at all. + assert(range(newSession(defaultPackage), "end" -> "37", "batch_size" -> "10").count() == 37) + } + + nativeTest("push down a limit") { + val df = range(newSession(defaultPackage), "end" -> "100", "partitions" -> "5").limit(3) + assert(df.collect().length == 3) + assert(scan(df).getMetaData()("PushedLimit") == "LIMIT 3") + + // Spark pushes the limit of LIMIT ... OFFSET twice, but the library, which fails if + // pushLimit is called again, is only asked once. + val withOffset = range(newSession(defaultPackage), "end" -> "100", "partitions" -> "5") + .limit(10).offset(5) + checkRows(withOffset, (5L until 10L).map(expectedRow)) + assert(scan(withOffset).getMetaData()("PushedLimit") == "LIMIT 10") + + // Spark evaluates the predicates after reading, so it does not push a limit with them. + val filtered = range(newSession(defaultPackage), "end" -> "100").filter("id > 5").limit(3) + assert(filtered.collect().length == 3) + assert(!scan(filtered).getMetaData().contains("PushedLimit")) + } + + nativeTest("read maps whose entry fields have other names") { + // The library names the fields of the entries `keys` and `values`, as arrow-rs does. + val df = range(newSession(defaultPackage), "end" -> "3", "extra" -> "map").select("id", "m") + assert(df.schema("m").dataType == MapType(StringType, LongType, valueContainsNull = true)) + checkRows(df, (0L until 3L).map(id => Row(id, Map(s"k$id" -> id)))) + } + + nativeTest("read structs in several batches") { + // The reader imports the batches of a partition into the same vectors, which Spark must not + // close when it moves to the next batch: closing a struct drops its child vectors. + val df = range(newSession(defaultPackage), "end" -> "10", "partitions" -> "1", + "batch_size" -> "3", "extra" -> "struct").select("id", "st") + assert(df.queryExecution.executedPlan.exists(_.isInstanceOf[WholeStageCodegenExec])) + checkRows(df, (0L until 10L).map(id => Row(id, Row(id, s"s$id")))) + + // A nested field is read with its top-level column: the library only prunes the top-level + // columns, and Spark the nested fields. + val nested = range(newSession(defaultPackage), "end" -> "4", "extra" -> "struct") + .select("st.b") + checkRows(nested, (0L until 4L).map(id => Row(s"s$id"))) + assert(scan(nested).readSchema() == StructType.fromDDL("st STRUCT")) + } + + nativeTest("reject the data that has nulls where the schema to read has none") { + // The column name is null for the ids that are multiples of 5. + val session = newSession(defaultPackage) + val e = intercept[Exception] { + session.read.schema("id BIGINT, name STRING NOT NULL").format("native_range") + .option("end", "4").load().collect() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "read", + "name" -> "native_range", + "msg" -> "The column name has a null, but the schema to read declares it not nullable.")) + // The fields that are not nullable have no nulls. + checkRows( + session.read.schema("id BIGINT NOT NULL, st STRUCT") + .format("native_range").option("end", "2").option("extra", "struct").load(), + Seq(Row(0L, Row(0L, "s0")), Row(1L, Row(1L, "s1")))) + } + + nativeTest("reject the types and the strings that a library cannot get") { + val session = newSession(defaultPackage) + def check(action: String, name: String, msg: String)(f: => Unit): Unit = { + val e = intercept[Exception](f) + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map("action" -> action, "name" -> name, "msg" -> msg)) + } + // Spark rejects the types that version 1 does not have before it passes a schema to the + // library: here, before create_reader and create_writer. + check("plan a scan of", "native_range", "The column id has the type INTERVAL DAY TO " + + "SECOND, which native data sources do not support.") { + session.read.schema("id INTERVAL DAY TO SECOND").format("native_range").load().collect() + } + val dir = Utils.createTempDir() + check("plan a write to", "native_sink", "The column v.element has the type VARIANT, which " + + "native data sources do not support.") { + session.sql("SELECT array(parse_json('1')) AS v").write.format("native_sink") + .option("path", dir.getPath).mode("append").save() + } + assert(!new File(dir, "_CREATED").exists()) + // A NUL-terminated string cannot hold the character U+0000. + check("create", "native_range", "The option end contains the character U+0000, which Spark " + + "cannot pass to a library.") { + range(session, "end" -> "1\u00002") + } + } + + nativeTest("Arrow UDFs on the columns of a native data source") { + assume(IntegratedUDFTestUtils.shouldTestPandasUDFs) + // The reader imports every batch into the same vectors, and closes them once its partition + // ends, while Spark writes batches to Python before the results of the earlier ones come + // back: the UDF must not read the pass-through columns of a batch after that. + val session = newSession(defaultPackage) + val udf = IntegratedUDFTestUtils.TestScalarPandasUDF( + name = "native_str", returnType = Some(StringType)) + IntegratedUDFTestUtils.registerTestUDF(udf, session) + // One batch, two, and many. The test UDF casts its input to string: on the string column + // name, the cast is removed, so the batches go to Python as they are. + Seq(100, 200, 200000).foreach { end => + val rows = range(session, "end" -> end.toString, "partitions" -> "1", "batch_size" -> "100") + .selectExpr("id", "name", "native_str(name) AS udf_name").collect() + assert(rows.map(_.getLong(0)).toSeq == (0L until end)) + rows.foreach { row => + val id = row.getLong(0) + val name = if (id % 5 == 0) null else s"n$id" + assert(row.getString(1) == name, s"name mismatch for id $id") + assert(row.getString(2) == name, s"UDF mismatch for id $id") + } + } + } + + nativeTest("reject the Arrow data that Spark cannot read") { + def check(msg: String, options: (String, String)*): Unit = { + checkWithSchema(None, msg, options: _*) + } + // The schema to read is given for the types from which Spark cannot infer it either. + def checkWithSchema(schema: Option[String], msg: String, options: (String, String)*): Unit = { + val e = intercept[Exception] { + val reader = newSession(defaultPackage).read.format("native_range") + schema.fold(reader)(s => reader.schema(s)).options(options.toMap + ("end" -> "4")).load() + .collect() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map("action" -> "read", "name" -> "native_range", "msg" -> msg)) + } + // Arrow Java would read a slice from the start of its buffers. + check("The column id has an offset of 1, but Spark only reads arrays with an offset of 0. " + + "Copy sliced arrays before exporting them.", "slice" -> "true") + // Arrow Java would extend an array with fewer values than its parent refers to with nulls. + // Each batch has 2 rows. + check("The column id has a length of 1, but the batch has a length of 2.", "short" -> "id") + // Arrow Java would import the missing columns as nulls. + check("The batch has 0 child arrays, but the schema has 7 columns.", "short" -> "children") + // Arrow Java would read a null row of a batch as a row of the values of its columns. The + // nulls are counted when the library did not count them. + Seq("count", "unknown").foreach { nullRows => + check("The struct of the batch marks 1 of its rows as null, but Spark only reads batches " + + "without null rows: mark the values of the columns as null instead.", + "null_rows" -> nullRows) + } + // The C data interface only allows an array without a validity buffer if its null_count is + // 0, and Arrow Java would fail to import it. + check("The batch has a null_count of -1 but no validity buffer, which the C data interface " + + "only allows with a null_count of 0.", "null_rows" -> "no_buffer") + check("The column id has a null_count of -1 but no validity buffer, which the C data " + + "interface only allows with a null_count of 0.", "null_rows" -> "column") + // A batch without null rows is read, also when the library did not count its nulls. + checkRows(range(newSession(defaultPackage), "end" -> "4", "null_rows" -> "unknown_valid"), + (0L until 4L).map(expectedRow)) + Seq("entries", "values").foreach { array => + check("The column m contains an array with a length of 1, but its parent refers to 2 " + + "values.", "extra" -> "map", "short" -> array) + } + check("The column dec contains the Arrow type Decimal(10, 2, 256), which is not supported.", + "extra" -> "decimal256") + // Arrow Java would also extend the elements of a list view with nulls, and Spark does not + // check their lengths. + check("The column lv contains the Arrow type ListView, which is not supported.", + "extra" -> "listview") + // Types that Spark does not convert, and dictionary encoding. + checkWithSchema(Some("id BIGINT, ll ARRAY"), + "The column ll contains the Arrow type LargeList, which is not supported.", + "extra" -> "largelist") + checkWithSchema(Some("id BIGINT, dict STRING"), + "The column dict is dictionary-encoded, which is not supported.", + "extra" -> "dictionary") + } + + nativeTest("batch write") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + def write(mode: String, ids: Range): Unit = { + import session.implicits._ + ids.map(i => (i.toLong, s"v$i")).toDF("id", "name").repartition(2) + .write.format("native_sink").option("path", dir.getPath).mode(mode).save() + } + def written(): Set[String] = dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala).toSet + + write("append", 0 until 4) + assert(written() == (0 until 4).map(i => s"$i,v$i").toSet) + assert(new File(dir, "_SUCCESS").isFile) + // The writer is closed after the commit, and so are the data writers of the tasks. + assert(new File(dir, "_CLOSED").isFile) + checkDataWritersClosed(dir) + write("append", 4 until 6) + assert(written() == (0 until 6).map(i => s"$i,v$i").toSet) + write("overwrite", 10 until 12) + assert(written() == (10 until 12).map(i => s"$i,v$i").toSet) + + // Nulls and the other types of the test library. + range(session, "end" -> "6").write.format("native_sink").option("path", dir.getPath) + .mode("overwrite").save() + assert(written() == (0 until 6).map { id => + val name = if (id % 5 == 0) "null" else s"n$id" + s"$id,${id % 3},$name,${format(id * 0.5)},${id % 2 == 0},$id,${id * 1000000L}" + }.toSet) + } + + nativeTest("tables in a catalog") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + try { + session.sql("CREATE TABLE native_range_table USING native_range OPTIONS (end 3)") + checkRows(session.table("native_range_table"), (0L until 3L).map(expectedRow)) + checkRows( + session.read.option("end", "2").table("native_range_table"), + (0L until 2L).map(expectedRow)) + + session.sql("CREATE TABLE native_sink_table (id BIGINT, name STRING) USING native_sink " + + s"OPTIONS (path '${dir.getPath}')") + session.sql("INSERT INTO native_sink_table VALUES (1, 'a'), (2, 'b')") + val written = dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala) + assert(written.toSet == Set("1,a", "2,b")) + } finally { + session.sql("DROP TABLE IF EXISTS native_range_table") + session.sql("DROP TABLE IF EXISTS native_sink_table") + } + } + + nativeTest("abort a failed write") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val e = intercept[Exception] { + session.range(10).write.format("native_sink").option("path", dir.getPath) + .option("fail", "write").mode("append").save() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "write to", + "name" -> "native_sink", + "msg" -> "injected failure in write")) + assert(new File(dir, "_ABORTED").isFile) + assert(new File(dir, "_CLOSED").isFile) + assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) + } + + nativeTest("create no writer when computing the data to write fails") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + // The error happens in a shuffle stage, which adaptive query execution runs before the write. + val e = intercept[Exception] { + session.range(10).selectExpr("id", "if(id >= 0, raise_error('boom'), 'x') AS name") + .repartition(2) + .write.format("native_sink").option("path", dir.getPath).mode("append").save() + } + assert(e.getMessage.contains("boom")) + assert(!new File(dir, "_CREATED").exists()) + + session.range(2).write.format("native_sink").option("path", dir.getPath).mode("append").save() + assert(new File(dir, "_CREATED").isFile) + } + + nativeTest("abort a write whose commit failed") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val e = intercept[Exception] { + session.range(10).write.format("native_sink").option("path", dir.getPath) + .option("fail", "commit").mode("append").save() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "commit a write to", + "name" -> "native_sink", + "msg" -> "injected failure in commit")) + assert(e.getSuppressed.isEmpty) + // Spark aborted the write with the writer, which it closed afterwards. + assert(new File(dir, "_ABORTED").isFile) + assert(new File(dir, "_CLOSED").isFile) + } + + nativeTest("abort a data writer whose commit failed") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val e = intercept[Exception] { + session.range(10).repartition(2).write.format("native_sink").option("path", dir.getPath) + .option("fail", "commitDataWriter").mode("append").save() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "write to", + "name" -> "native_sink", + "msg" -> "injected failure in commitDataWriter")) + // abort_data_writer removed the temporary files of the tasks, whose data writers were then + // closed. + assert(new File(dir, "_ABORTED").isFile) + assert(new File(dir, "_temporary").listFiles().isEmpty) + assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) + checkDataWritersClosed(dir) + } + + nativeTest("do not abort a data writer that committed with an invalid message") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + // The library commits the data writer, but passes two commit messages. + val e = intercept[Exception] { + session.range(10).repartition(1).write.format("native_sink").option("path", dir.getPath) + .option("commit_messages", "2").mode("append").save() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "write to", + "name" -> "native_sink", + "msg" -> "commit_data_writer returned 2 values instead of one.")) + // The task failed without aborting the committed data writer, which it closed, and the driver + // aborted the write. + assert(!new File(dir, "_ABORTED_COMMITTED_DATA_WRITER").exists()) + checkDataWritersClosed(dir) + assert(new File(dir, "_ABORTED").isFile) + assert(!dir.listFiles().exists(_.getName.startsWith("part-"))) + } + + nativeTest("abort a micro-batch whose write failed") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val e = intercept[Exception] { + val query = session.readStream.format("native_counter").option("max_offset", "4").load() + .writeStream + .format("native_sink") + .option("path", dir.getPath) + .option("fail", "write") + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try query.processAllAvailable() finally query.stop() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "write to", + "name" -> "native_sink", + "msg" -> "injected failure in write")) + // Spark aborted the first epoch, with a null message for its task, which did not commit. + assert(Files.readString(new File(dir, "_ABORTED_0").toPath) == "1") + assert(!new File(dir, "_EPOCH_0").exists()) + } + + nativeTest("continuous processing fails with a native data source") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val e = intercept[Exception] { + val query = session.readStream.format("rate").load() + .writeStream + .format("native_sink") + .option("path", dir.getPath) + .option("checkpointLocation", Utils.createTempDir().getPath) + .trigger(Trigger.Continuous("1 second")) + .start() + try query.awaitTermination(60000) finally query.stop() + } + checkError( + exception = findError(e, "DATA_SOURCE_CONTINUOUS_WRITE_NOT_SUPPORTED"), + condition = "DATA_SOURCE_CONTINUOUS_WRITE_NOT_SUPPORTED", + parameters = Map("description" -> "native_sink")) + // The query failed before Spark created a writer in the library. + assert(!new File(dir, "_CREATED").exists()) + // A native scan does not report continuous reads, so Spark rejects them when it starts. + val readError = intercept[Exception] { + session.readStream.format("native_counter").load() + .writeStream.format("noop").trigger(Trigger.Continuous("1 second")).start() + } + checkError( + exception = findError(readError, "_LEGACY_ERROR_TEMP_2253"), + condition = "_LEGACY_ERROR_TEMP_2253", + parameters = Map("sourceName" -> "native_counter")) + } + + nativeTest("streaming read and write") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val commitPath = new File(Utils.createTempDir(), "committed") + val closedPath = new File(Utils.createTempDir(), "closed") + val query = session.readStream.format("native_counter") + .option("max_offset", "12").option("step", "5").option("batch_size", "2") + .option("commit_path", commitPath.getPath) + .option("closed_path", closedPath.getPath) + .load() + .writeStream + .format("native_sink") + .option("path", dir.getPath) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try { + query.processAllAvailable() + } finally { + query.stop() + } + // Each call to latestOffset makes 5 more rows available: [0, 5), [5, 10) and [10, 12). + assert((0 to 2).forall(epoch => new File(dir, s"_EPOCH_$epoch").isFile)) + // Each micro-batch creates a new stream writer, which closes the one of the previous + // micro-batch. The last one is closed once it is no longer referenced. + assert((0 to 1).forall(epoch => new File(dir, s"_CLOSED_$epoch").isFile)) + val written = dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala) + assert(written.sorted.toSeq == (0 until 12).map(_.toString).sorted) + // Spark commits the offsets of a micro-batch when it runs the next one. + assert(Files.readString(commitPath.toPath) == "{\"offset\":10}") + // The stream reader is closed when the query stops. + assert(closedPath.isFile) + + // The streaming source can also be read in micro-batches into a memory sink. + val memoryQuery = session.readStream.format("native_counter").option("max_offset", "4") + .load().writeStream.format("memory").queryName("native_counter_memory").start() + try { + memoryQuery.processAllAvailable() + checkRows(session.table("native_counter_memory"), (0L until 4L).map(Row(_))) + } finally { + memoryQuery.stop() + } + } + + nativeTest("reject offsets that are not JSON on a single line") { + // Spark writes each offset as a line of the offset log of the checkpoint. + def check(badOffset: String, msg: String): Unit = { + val e = intercept[Exception] { + val query = newSession(defaultPackage).readStream.format("native_counter") + .option("bad_offset", badOffset).load() + .writeStream + .format("noop") + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try query.processAllAvailable() finally query.stop() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map("action" -> "plan a scan of", "name" -> "native_counter", "msg" -> msg)) + } + check("none", "latest_offset returned no offset.") + check("multiline", "latest_offset returned an offset with a line break, but an " + + "offset must be JSON on a single line.") + check("invalid", "latest_offset returned an offset that is not valid JSON: offset 10") + } + + nativeTest("pass the ID of a streaming query to its stream writers") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val checkpoint = Utils.createTempDir().getPath + // Runs the query until it processed the rows [0, maxOffset). + def run(maxOffset: Int): StreamingQuery = { + val query = session.readStream.format("native_counter") + .option("max_offset", maxOffset.toString).load() + .writeStream + .format("native_sink") + .option("path", dir.getPath) + .option("checkpointLocation", checkpoint) + .start() + try query.processAllAvailable() finally query.stop() + query + } + val queryId = new File(dir, "_QUERY_ID") + val first = run(3) + assert(Files.readString(queryId.toPath) == first.id.toString) + // The ID stays the same when the query restarts from its checkpoint, unlike the ID of the + // run, so that the library can tell the commits of the query from those of other queries. + assert(queryId.delete()) + val restarted = run(6) + assert(restarted.id == first.id && restarted.runId != first.runId) + assert(Files.readString(queryId.toPath) == first.id.toString) + } + + nativeTest("report the errors of native code") { + def check(action: String, msg: String, name: String = "native_range")( + f: SparkSession => Unit): Unit = { + val e = intercept[Exception](f(newSession(defaultPackage))) + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map("action" -> action, "name" -> name, "msg" -> msg)) + } + check("create", "injected failure in create") { session => + range(session, "fail" -> "create") + } + check("infer the schema of", "injected failure in schema") { session => + range(session, "fail" -> "schema") + } + check("plan a scan of", "injected failure in plan") { session => + range(session, "fail" -> "plan").collect() + } + check("read", "injected failure in open") { session => + range(session, "fail" -> "open").collect() + } + check("read", "injected failure in get_next") { session => + range(session, "fail" -> "read").collect() + } + check("read", "The library did not initialize the stream, or a callback of the stream is " + + "NULL.") { session => + range(session, "fail" -> "uninitialized").collect() + } + // An error of the library is not an unsupported operation: here, of a data source that + // declares batch reads that it does not support. + check("plan a scan of", "native_sink does not support batch reads", "native_sink") { + session => session.read.schema("id BIGINT").format("native_sink") + .option("capabilities", NativeLibrary.BATCH_READ.toString).load().collect() + } + // Without last_error, only the value that the function returned is reported. + val e = intercept[Exception] { + newSession(createPackage(noLastErrorLibrary, dataSources("ne_"))) + .read.format("ne_native_range").option("fail", "create").load() + } + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "create", + "name" -> "ne_native_range", + "msg" -> "create failed with the error code 1.")) + } + + nativeTest("a scan whose planning failed fails again with the same error") { + // The reader is released after planning, even if it failed. The query execution of a + // DataFrame keeps the failure of its plan, but a scan can be planned again, for example by + // the lazy values of the physical plan that call it. + val session = newSession(defaultPackage) + val (name, location) = + NativeDataSourceRegistry.lookup("native_range", session.sessionState.conf).get + val options = new CaseInsensitiveStringMap(Map("fail" -> "plan").asJava) + val batch = new ColumnarBatchScan( + new NativeDataSource(location, name, options).batchPlanner(StructType.fromDDL("id BIGINT"))) + (1 to 2).foreach { _ => + checkError( + exception = intercept[SparkRuntimeException](batch.planInputPartitions()), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "plan a scan of", + "name" -> "native_range", + "msg" -> "injected failure in plan")) + } + } + + nativeTest("reject NULL handles") { + def check(function: String, action: String, name: String)(f: SparkSession => Unit): Unit = { + val e = intercept[Exception](f(newSession(defaultPackage))) + checkError( + exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> action, + "name" -> name, + "msg" -> s"$function returned a NULL handle.")) + } + def runStream(start: => StreamingQuery): Unit = { + val query = start + try query.processAllAvailable() finally query.stop() + } + check("create", "create", "native_range") { session => + range(session, "zero" -> "create") + } + check("create_reader", "plan a scan of", "native_range") { session => + range(session, "zero" -> "create_reader").collect() + } + check("create_stream_reader", "plan a scan of", "native_counter") { session => + runStream { + session.readStream.format("native_counter") + .option("zero", "create_stream_reader").load() + .writeStream.format("noop").start() + } + } + Seq("create_writer" -> "plan a write to", "create_data_writer" -> "write to").foreach { + case (function, action) => + check(function, action, "native_sink") { session => + session.range(2).write.format("native_sink") + .option("path", Utils.createTempDir().getPath).option("zero", function) + .mode("append").save() + } + } + check("create_writer", "plan a write to", "native_sink") { session => + runStream { + session.readStream.format("native_counter").load() + .writeStream + .format("native_sink") + .option("path", Utils.createTempDir().getPath) + .option("zero", "create_writer") + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + } + } + } + + nativeTest("operations that a data source does not support") { + val session = newSession(defaultPackage, createPackage(readOnlyLibrary, dataSources("ro_"))) + // The library does not push anything down, so Spark evaluates the filter. + val df = session.read.format("ro_native_range").option("end", "10").load().filter("id > 6") + checkRows(df, (7L until 10L).map(expectedRow)) + assert(scan(df).getMetaData()("PushedPredicates") == "[]") + + def check(condition: String, parameters: Map[String, String])(f: => Unit): Unit = { + val e = intercept[Exception](f) + checkError(exception = findError(e, condition), condition = condition, + parameters = parameters) + } + def runStream(start: => StreamingQuery): Unit = { + val query = start + try query.processAllAvailable() finally query.stop() + } + // A data source declares the operations that it supports with its capabilities. Spark + // reports the others as unsupported, without calling the library, which here does not + // implement them either. The table reports the batch writes to Data Source V2, which rejects + // them while it analyzes the query, before it computes the data to write: here, in a shuffle + // stage that adaptive query execution would run before the write. + check("UNSUPPORTED_FEATURE.TABLE_OPERATION", + Map("tableName" -> "`ro_native_sink`", "operation" -> "batch write")) { + session.range(10).selectExpr("if(id >= 0, raise_error('boom'), id) AS id").repartition(2) + .write.format("ro_native_sink").mode("append").save() + } + check("DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> "ro_native_counter")) { + runStream { + session.readStream.format("ro_native_counter").load().writeStream.format("noop").start() + } + } + check("DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", Map("description" -> "ro_native_sink")) { + runStream { + session.readStream.format("rate").load().writeStream.format("ro_native_sink") + .option("checkpointLocation", Utils.createTempDir().getPath).start() + } + } + // The capabilities are those of each data source: native_counter only supports streaming + // reads, although its library implements batch reads. + check("DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> "native_counter")) { + session.read.format("native_counter").load().collect() + } + // A data source that declares a capability must implement its functions. + check("NATIVE_DATA_SOURCE_ERROR", Map("action" -> "plan a write to", + "name" -> "ro_native_sink", "msg" -> "The library does not implement create_writer.")) { + session.range(1).write.format("ro_native_sink") + .option("capabilities", NativeLibrary.BATCH_WRITE.toString).mode("append").save() + } + // Replacing the existing data requires the capability SPARK_DS_TRUNCATE. + val dir = Utils.createTempDir() + val noTruncate = (NativeLibrary.BATCH_WRITE | NativeLibrary.STREAMING_WRITE).toString + check("UNSUPPORTED_FEATURE.TABLE_OPERATION", + Map("tableName" -> "`native_sink`", "operation" -> "truncate in batch mode")) { + session.range(1).write.format("native_sink").option("path", dir.getPath) + .option("capabilities", noTruncate).mode("overwrite").save() + } + // The table has the capabilities that the data source declares with the options of the table, + // which the options of a write to a table in a catalog can override: the write checks them + // again. + try { + session.sql("CREATE TABLE native_sink_table (id BIGINT) USING native_sink " + + s"OPTIONS (path '${dir.getPath}')") + check("DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", Map("description" -> "native_sink")) { + session.range(1).writeTo("native_sink_table") + .option("capabilities", NativeLibrary.BATCH_READ.toString).append() + } + check("NATIVE_DATA_SOURCE_ERROR", Map("action" -> "plan a write to", + "name" -> "native_sink", + "msg" -> "The data source does not support the save mode overwrite.")) { + session.range(1).writeTo("native_sink_table").option("capabilities", noTruncate) + .overwrite(lit(true)) + } + } finally { + session.sql("DROP TABLE IF EXISTS native_sink_table") + } + check("NATIVE_DATA_SOURCE_ERROR", Map("action" -> "plan a write to", "name" -> "native_sink", + "msg" -> "The data source does not support the output mode complete.")) { + runStream { + session.readStream.format("rate").load().groupBy().count().writeStream + .outputMode("complete") + .format("native_sink") + .option("path", dir.getPath) + .option("capabilities", noTruncate) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + } + } + assert(!new File(dir, "_CREATED").exists()) + } + + nativeTest("read only the functions that fit in the table of a library") { + // The table of the library ends before read, as the table of a library built against an + // earlier header ends before the functions added later: Spark does not read past its end. + val session = newSession(createPackage(smallTableLibrary, dataSources("small_"))) + val df = session.read.format("small_native_range").load() + assert(df.schema.length == 7) + checkError( + exception = findError(intercept[Exception](df.collect()), "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "plan a scan of", + "name" -> "small_native_range", + "msg" -> "The library does not implement read.")) + } + + nativeTest("pass the version of Spark, and log with the log of Spark") { + val session = newSession(defaultPackage) + checkRows(range(session, "end" -> "2", "spark_version" -> SPARK_VERSION).select("id"), + Seq(Row(0L), Row(1L))) + checkError( + exception = findError( + intercept[Exception](range(session, "spark_version" -> "1.0.0")), + "NATIVE_DATA_SOURCE_ERROR"), + condition = "NATIVE_DATA_SOURCE_ERROR", + parameters = Map( + "action" -> "create", + "name" -> "native_range", + "msg" -> (s"the version of Spark is $SPARK_VERSION, which supports version 1 of the " + + "interface"))) + + val logAppender = new LogAppender("native data source logs") + withLogAppender(logAppender, loggerNames = Seq(classOf[NativeLibrary].getName)) { + range(session, "log" -> "a message from native code") + } + assert(logAppender.loggingEvents.exists { event => + val message = event.getMessage.getFormattedMessage + event.getLevel == org.apache.logging.log4j.Level.WARN && + message.contains("a message from native code") && + message.contains(System.mapLibraryName("test_native_datasource")) + }) + } + + nativeTest("pass the task to reads and data writers") { + val session = newSession(defaultPackage) + // Checks the SparkDsTaskContext that the library got in each task, by partition. + def check(dir: File): Unit = { + val tasks = dir.listFiles().map { file => + file.getName.toInt -> Files.readString(file.toPath).split(",").map(_.toLong).toSeq + }.toMap + assert(tasks.keySet == Set(0, 1)) + tasks.foreach { + case (partition, Seq(size, stageId, stageAttempt, partitionId, attempt, taskAttemptId)) => + assert(size == NativeLibrary.layout().get("SparkDsTaskContext").longValue) + assert(partitionId == partition) + assert(stageId >= 0 && stageAttempt == 0 && attempt == 0 && taskAttemptId >= 0) + } + assert(tasks.values.map(_(5)).toSet.size == 2) + } + val readTasks = Utils.createTempDir() + checkRows(range(session, "end" -> "6", "task_path" -> readTasks.getPath).select("id"), + (0L until 6L).map(Row(_))) + check(readTasks) + val writeTasks = Utils.createTempDir() + session.range(6).repartition(2).write.format("native_sink") + .option("path", Utils.createTempDir().getPath) + .option("task_path", writeTasks.getPath) + .mode("append") + .save() + check(writeTasks) + } + + nativeTest("tell a read that its task was killed") { + val session = newSession(defaultPackage) + val marker = new File(Utils.createTempDir(), "interrupted") + val df = range(session, "end" -> "10", "partitions" -> "1", + "wait_for_interrupt" -> marker.getPath) + val jobGroup = s"native-data-source-${UUID.randomUUID()}" + val error = new AtomicReference[Throwable]() + val thread = new Thread(() => { + session.sparkContext.setJobGroup(jobGroup, "a read that waits", interruptOnCancel = true) + try df.collect() catch { + case e: Throwable => error.set(e) + } + }) + thread.start() + // The read waits until its task is killed, which it only learns from is_interrupted: Spark + // cannot interrupt native code. + eventually(timeout(1.minute)) { + assert(new File(marker.getPath + ".waiting").isFile) + } + session.sparkContext.cancelJobGroup(jobGroup) + thread.join(TimeUnit.MINUTES.toMillis(1)) + assert(!thread.isAlive) + assert(error.get() != null) + // The job fails right away, and its task ends once the library sees that it was killed. + eventually(timeout(1.minute)) { + assert(marker.isFile && Files.readString(marker.toPath) == "interrupted") + } + } + + nativeTest("load libraries that export the same functions side by side") { + val session = newSession( + defaultPackage, createPackage(negatingLibrary, dataSources("b_"))) + checkRows(range(session, "end" -> "3").select("id"), Seq(Row(0L), Row(1L), Row(2L))) + checkRows( + session.read.format("b_native_range").option("end", "3").load().select("id"), + Seq(Row(0L), Row(-1L), Row(-2L))) + } + + nativeTest("add packages with spark.addArtifact") { + val session = newSession() + session.addArtifact(defaultPackage.getPath) + val (added, artifactUUID) = session.artifactManager.getNativeDataSourcePackages + assert(added.map(_.getName) == Seq(defaultPackage.getName)) + checkRows(range(session, "end" -> "4"), (0L until 4L).map(expectedRow)) + + // The executors find the package in the files of the session. + val pkg = NativeDataSourcePackage.read(added.head).copy(artifactUUID = artifactUUID) + assert(NativeDataSourceRegistry.distributedCopies(pkg).exists(_.isFile)) + + // A package that the session added is distributed as it is. + assert(NativeDataSourceRegistry.addToArtifacts(session, pkg) == pkg) + + // A package under spark.sql.dataSource.native.paths is copied to the artifacts of the session + // when it is used, except in local mode, where the executors run in the driver. The copy is + // named after the checksum of the package, and it is not a package of the session. + val configuredDir = Utils.createTempDir() + val configured = createPackage(library, dataSources(""), dir = configuredDir, + fileName = defaultPackage.getName) + val other = newSession(configured) + range(other, "end" -> "1").collect() + assert(!other.artifactManager.hasFileArtifact(defaultPackage.getName)) + val configuredPackage = NativeDataSourcePackage.read(configured) + val distributed = NativeDataSourceRegistry.addToArtifacts(other, configuredPackage) + val copyName = s"${defaultPackage.getName}.${configuredPackage.sha256}" + assert(distributed == configuredPackage.copy( + artifactUUID = other.artifactManager.getNativeDataSourcePackages._2, + distributedFileName = Some(copyName))) + assert(other.artifactManager.hasFileArtifact(copyName)) + assert(other.artifactManager.getNativeDataSourcePackages._1.isEmpty) + assert(NativeDataSourceRegistry.distributedCopies(distributed).exists(_.isFile)) + // Adding the same package again does nothing. + assert(NativeDataSourceRegistry.addToArtifacts(other, configuredPackage) == distributed) + + // A package with the same file name as a package of the session does not collide with it. + other.addArtifact(defaultPackage.getPath) + val sameFileName = createPackage(negatingLibrary, dataSources("b_"), + dir = Utils.createTempDir(), fileName = defaultPackage.getName) + val sameFileNameCopy = + NativeDataSourceRegistry.addToArtifacts(other, NativeDataSourcePackage.read(sameFileName)) + assert(sameFileNameCopy.distributedFileName.exists(_.startsWith(defaultPackage.getName + "."))) + + // The executors also check the copies in the files of the other sessions, such as the session + // of a streaming query, which is a clone of the session that started it. + val sessionDir = new File(SparkFiles.getRootDirectory(), s"test-${UUID.randomUUID()}") + try { + val copy = new File(sessionDir, defaultPackage.getName) + sessionDir.mkdirs() + Files.copy(defaultPackage.toPath, copy.toPath) + val copies = NativeDataSourceRegistry.distributedCopies(pkg.copy(artifactUUID = Some("x"))) + assert(copies.take(2) == Seq( + new File(new File(SparkFiles.getRootDirectory(), "x"), defaultPackage.getName), + new File(SparkFiles.getRootDirectory(), defaultPackage.getName))) + assert(copies.contains(copy)) + } finally { + Utils.deleteRecursively(sessionDir) + } + } + + nativeTest("distribute a configured package for concurrent queries") { + // Queries of a session can use a configured package for the first time at the same time. + val configured = createPackage(library, dataSources(""), dir = Utils.createTempDir()) + val session = newSession(configured) + val pkg = NativeDataSourcePackage.read(configured) + val threads = 8 + val pool = ThreadUtils.newDaemonFixedThreadPool(threads, "distribute-package") + try { + val start = new CountDownLatch(1) + val results = (1 to threads).map { _ => + pool.submit(new Callable[NativeDataSourcePackage] { + override def call(): NativeDataSourcePackage = { + start.await() + NativeDataSourceRegistry.addToArtifacts(session, pkg) + } + }) + } + start.countDown() + val distributed = results.map(_.get(1, TimeUnit.MINUTES)).distinct + assert(distributed.length == 1) + assert(session.artifactManager.hasFileArtifact(distributed.head.distributedFileName.get)) + } finally { + pool.shutdownNow() + } + } + + nativeTest("load a package only with the checksum of the driver") { + // The package at the path of the driver is valid, but it is not the package that the driver + // read, for example because it was replaced since, and neither is any distributed copy. + val pkg = NativeDataSourcePackage.read(defaultPackage).copy(sha256 = "0" * 64) + checkError( + exception = intercept[SparkRuntimeException](NativeLibraries.get(pkg)), + condition = "NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND", + parameters = Map("fileName" -> defaultPackage.getName, "checksum" -> pkg.sha256)) + } + + test("distribute a package without an active session") { + val pkg = NativeDataSourcePackage( + "/path/to/x.sparkpkg", "checksum", NativeDataSourceManifest(1, Seq("x"), Map.empty)) + val logAppender = new LogAppender("distribute without an active session") + SparkSession.clearActiveSession() + try { + withLogAppender(logAppender, level = Some(org.apache.logging.log4j.Level.WARN)) { + assert(NativeDataSourceRegistry.distribute(pkg) == pkg) + } + } finally { + SparkSession.setActiveSession(spark) + } + assert(logAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains("there is no active session"))) + } + + nativeTest("write with a session that is not the active one") { + val session = newSession() + session.addArtifact(defaultPackage.getPath) + SparkSession.setActiveSession(spark) + val dir = Utils.createTempDir() + session.range(3).selectExpr("id", "concat('v', id) AS name") + .write.format("native_sink").option("path", dir.getPath).mode("append").save() + val written = dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala) + assert(written.toSet == Set("0,v0", "1,v1", "2,v2")) + + val streamDir = Utils.createTempDir() + val query = session.readStream.format("native_counter").option("max_offset", "3").load() + .writeStream + .format("native_sink") + .option("path", streamDir.getPath) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try { + query.processAllAvailable() + } finally { + query.stop() + } + assert(streamDir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala).sorted.toSeq == Seq("0", "1", "2")) + } + + nativeTest("ignore configured directories that cannot be read") { + val dir = Utils.createTempDir() + createPackage(library, dataSources(""), dir = dir, fileName = "a.sparkpkg") + assert(dir.setReadable(false)) + try { + assume(dir.listFiles() == null, "The directory can still be read, for example by root.") + val session = newSession(dir) + checkError( + exception = intercept[SparkClassNotFoundException] { + session.read.format("no_such_source").load() + }, + condition = "DATA_SOURCE_NOT_FOUND", + parameters = Map("provider" -> "no_such_source")) + val e = intercept[AnalysisException](session.sql("SELECT * FROM missing_db.missing_table")) + assert(e.getCondition == "TABLE_OR_VIEW_NOT_FOUND") + } finally { + dir.setReadable(true) + } + } + + nativeTest("replace a configured package in place") { + val dir = Utils.createTempDir() + val configured = createPackage(library, dataSources(""), dir = dir, fileName = "x.sparkpkg") + val session = newSession(configured) + checkRows(range(session, "end" -> "2"), (0L until 2L).map(expectedRow)) + // On a cluster, the package is copied to the artifacts of the session when it is used. + val first = + NativeDataSourceRegistry.addToArtifacts(session, NativeDataSourcePackage.read(configured)) + + // An operator replaces the package with a new version: here, with one more file. + val path = s"${NativePlatform.current}/${library.getName}" + createZip(dir, "x.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + dataSources(""), Map(NativePlatform.current -> path)).getBytes(UTF_8), + path -> Files.readAllBytes(library.toPath), + "README" -> "version 2".getBytes(UTF_8)) + // The copy of the previous version is not a package of the session, so the session uses the + // new version, which is copied under another name. + checkRows(range(session, "end" -> "3"), (0L until 3L).map(expectedRow)) + val second = + NativeDataSourceRegistry.addToArtifacts(session, NativeDataSourcePackage.read(configured)) + assert(second.sha256 != first.sha256) + assert(second.distributedFileName != first.distributedFileName) + assert(session.artifactManager.getNativeDataSourcePackages._1.isEmpty) + } + + nativeTest("look up a native data source once") { + val session = newSession(defaultPackage) + val provider = DataSource.lookupDataSourceV2("native_range", session.sessionState.conf).get + assert(provider.isInstanceOf[NativeDataSourceV2]) + assert(NativeDataSourceRegistry.takeLookup("native_range").isEmpty) + // The provider uses the result of the lookup that found its class, even if the package is + // no longer configured. + session.conf.unset(SQLConf.NATIVE_DATA_SOURCE_PATHS.key) + assert(provider.inferSchema(CaseInsensitiveStringMap.empty()).fieldNames.head == "id") + + // The result is only kept until the next lookup on the thread, whatever it finds, so a + // provider never takes the result of an earlier lookup. + session.conf.set(SQLConf.NATIVE_DATA_SOURCE_PATHS.key, defaultPackage.getPath) + assert(DataSource.lookupDataSource("native_range", session.sessionState.conf) == + classOf[NativeDataSourceV2]) + assert(DataSource.lookupDataSource("json", session.sessionState.conf) != null) + assert(NativeDataSourceRegistry.takeLookup("native_range").isEmpty) + } + + nativeTest("Python data sources take precedence") { + val dir = Utils.createTempDir() + installLibrary(installedLibrary, "installed_native_range", dir) + withLibraryPath(dir) { + val session = newSession() + val conf = session.sessionState.conf + assert(DataSource.lookupDataSource("installed_native_range", conf) == + classOf[NativeDataSourceV2]) + // A Python data source can be registered with the name of an installed native data + // source, and takes precedence over it. + val pythonDataSource = UserDefinedPythonDataSource(SimplePythonFunction( + command = Seq.empty, + envVars = new java.util.HashMap[String, String](), + pythonIncludes = java.util.Collections.emptyList[String](), + pythonExec = "python3", + pythonVer = "3", + broadcastVars = java.util.Collections.emptyList[Broadcast[PythonBroadcast]](), + accumulator = null)) + session.dataSource.registerPython("installed_native_range", pythonDataSource) + assert(DataSource.lookupDataSource("installed_native_range", conf) == + classOf[PythonDataSourceV2]) + } + } + + nativeTest("find installed libraries automatically") { + val dir = Utils.createTempDir() + val installed = installLibrary(installedLibrary, "installed_native_range", dir) + // A library that implements several data sources is installed under each of their names. + Files.createSymbolicLink( + new File(dir, NativeDataSourceRegistry.libraryFileName("installed_native_sink")).toPath, + installed.toPath) + withLibraryPath(dir) { + // The session neither adds a package nor configures a path. The name is case-insensitive. + val session = newSession() + checkRows( + session.read.format("Installed_Native_Range").option("end", "3").load().select("id"), + Seq(Row(0L), Row(1L), Row(2L))) + val output = Utils.createTempDir() + session.range(2).selectExpr("id", "concat('v', id) AS name") + .write.format("installed_native_sink").option("path", output.getPath).mode("append").save() + assert(output.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala).toSet == Set("0,v0", "1,v1")) + + // A package of the session takes precedence over an installed library. + val other = newSession(createPackage(installedNegatingLibrary, dataSources("installed_"))) + checkRows( + other.read.format("installed_native_range").option("end", "3").load().select("id"), + Seq(Row(0L), Row(-1L), Row(-2L))) + + // An executor finds the library in its native library path, if it is not where the driver + // found it. + val checksum = InstalledNativeLibrary(installed).sha256 + assert(checksum == NativeDataSourcePackage.sha256(installed)) + val elsewhere = new File(Utils.createTempDir(), installed.getName).getPath + assert(NativeLibraries.get(InstalledNativeLibrary(elsewhere, checksum)) eq + NativeLibraries.get(InstalledNativeLibrary(installed))) + + // An executor only loads the same build as the driver, which serialized the state for its + // build: neither the library at the same path nor the one in its native library path is + // loaded if the driver found another build. + val otherBuild = NativeDataSourcePackage.sha256(installedNegatingLibrary) + Seq(installed.getPath, elsewhere).foreach { path => + checkError( + exception = intercept[SparkRuntimeException] { + NativeLibraries.get(InstalledNativeLibrary(path, otherBuild)) + }, + condition = "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND", + parameters = + Map("fileName" -> installed.getName, "checksum" -> otherBuild, "path" -> path)) + } + } + + val missing = new File(Utils.createTempDir(), installed.getName).getPath + checkError( + exception = intercept[SparkRuntimeException] { + NativeLibraries.get(InstalledNativeLibrary(missing, "0")) + }, + condition = "NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND", + parameters = Map("fileName" -> installed.getName, "checksum" -> "0", "path" -> missing)) + checkError( + exception = intercept[SparkClassNotFoundException] { + newSession().read.format("installed_native_range").load() + }, + condition = "DATA_SOURCE_NOT_FOUND", + parameters = Map("provider" -> "installed_native_range")) + } + + nativeTest("reject an installed library replaced in place") { + val dir = Utils.createTempDir() + val installed = installLibrary(installedLibrary, "installed_native_range", dir) + withLibraryPath(dir) { + def read(): DataFrame = { + newSession().read.format("installed_native_range").option("end", "2").load().select("id") + } + checkRows(read(), Seq(Row(0L), Row(1L))) + val loaded = NativeDataSourcePackage.sha256(installed) + + // An operator installs another build at the same path, as `make install` does. The JVM + // cannot load it from the same path, so the driver and the executors that loaded the + // previous build reject it, rather than running a build that does not match the state + // that the driver serializes. + Files.copy(installedNegatingLibrary.toPath, installed.toPath, REPLACE_EXISTING) + val replaced = NativeDataSourcePackage.sha256(installed) + checkError( + exception = intercept[SparkRuntimeException](read()), + condition = "INVALID_NATIVE_DATA_SOURCE_LIBRARY.REPLACED", + parameters = Map( + "path" -> installed.getCanonicalPath, + "loadedChecksum" -> loaded, + "checksum" -> replaced)) + + // The library works again once the build that was loaded is installed again. + Files.copy(installedLibrary.toPath, installed.toPath, REPLACE_EXISTING) + checkRows(read(), Seq(Row(0L), Row(1L))) + } + } + + nativeTest("find packages in directories") { + val dir = Utils.createTempDir() + createPackage(library, dataSources(""), dir = dir, fileName = "a.sparkpkg") + // Files without the extension are ignored. + createZip(dir, "ignored.zip", "x" -> Array.emptyByteArray) + checkRows(range(newSession(dir), "end" -> "2"), (0L until 2L).map(expectedRow)) + } + + nativeTest("conflicting packages") { + val dir = Utils.createTempDir() + val first = createPackage(library, dataSources(""), dir = dir, fileName = "a.sparkpkg") + val second = createPackage(library, Seq("native_range"), dir = dir, fileName = "b.sparkpkg") + checkError( + exception = intercept[AnalysisException](range(newSession(dir))), + condition = "NATIVE_DATA_SOURCE_PACKAGE_CONFLICT", + parameters = Map( + "name" -> "native_range", + "packages" -> Seq(first, second).map(_.getCanonicalPath).mkString(", "))) + // The same package found twice is not a conflict: here, in the configured paths and in the + // artifacts of the session. + val session = newSession(first) + session.addArtifact(first.getPath) + checkRows(range(session, "end" -> "2"), Seq(expectedRow(0L), expectedRow(1L))) + } + + nativeTest("the package that a session added last replaces the other packages") { + // A session cannot remove an artifact, nor add another file with the same name, so it + // upgrades a package by adding the new version under another file name. + val session = newSession(defaultPackage) + def ids(): DataFrame = range(session, "end" -> "3").select("id") + checkRows(ids(), Seq(Row(0L), Row(1L), Row(2L))) + val appender = new LogAppender("replaced packages") + withLogAppender(appender) { + // It replaces the configured package... + val negating = createPackage(renamedNegatingLibrary, dataSources(""), + fileName = "range-1.0.sparkpkg") + session.addArtifact(negating.getPath) + checkRows(ids(), Seq(Row(0L), Row(-1L), Row(-2L))) + checkRows(ids(), Seq(Row(0L), Row(-1L), Row(-2L))) + // ... and the packages that the session added before. + session.addArtifact(createPackage(library, dataSources(""), + fileName = "range-1.1.sparkpkg").getPath) + checkRows(ids(), Seq(Row(0L), Row(1L), Row(2L))) + } + // Each replacement is logged once. + val logged = appender.loggingEvents.map(_.getMessage.getFormattedMessage) + .filter(_.contains("which the session added last")) + assert(logged.size == 2, logged) + assert(logged.forall(_.contains("for the data source native_range")), logged) + } + + nativeTest("Java data sources take precedence") { + val session = newSession(createPackage(library, Seq("json"))) + val dir = Utils.createTempDir().getPath + "/json" + session.range(2).write.json(dir) + checkRows(session.read.format("json").load(dir), Seq(Row(0L), Row(1L))) + } + + nativeTest("invalid packages") { + val dir = Utils.createTempDir() + def check(pkg: File, subClass: String, params: Map[String, String]): Unit = { + checkError( + exception = intercept[SparkRuntimeException](range(newSession(pkg)).collect()), + condition = s"INVALID_NATIVE_DATA_SOURCE_PACKAGE.$subClass", + parameters = params) + } + + val missingLibrary = createZip(dir, "missing_library.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + Seq("native_range"), Map(NativePlatform.current -> "lib.so")).getBytes(UTF_8)) + check(missingLibrary, "INVALID_MANIFEST", Map( + "path" -> missingLibrary.getCanonicalPath, + "reason" -> "The package does not contain the library lib.so.")) + + val newerAbi = createPackage(library, Seq("native_range"), dir = dir, + fileName = "newer_abi.sparkpkg", abiVersion = 2) + check(newerAbi, "UNSUPPORTED_ABI_VERSION", Map( + "path" -> newerAbi.getCanonicalPath, "version" -> "2", "supported" -> "1")) + + val otherPlatform = createPackage(library, Seq("native_range"), dir = dir, + fileName = "other_platform.sparkpkg", platform = "other-platform") + check(otherPlatform, "UNSUPPORTED_PLATFORM", Map( + "path" -> otherPlatform.getCanonicalPath, + "platform" -> NativePlatform.current, + "platforms" -> "other-platform")) + + val path = s"${NativePlatform.current}/${System.mapLibraryName("not_a_library")}" + val notALibrary = createZip(dir, "not_a_library.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + Seq("native_range"), Map(NativePlatform.current -> path)).getBytes(UTF_8), + path -> "not a library".getBytes(UTF_8)) + val e = intercept[SparkRuntimeException](range(newSession(notALibrary))) + assert(e.getCondition == "INVALID_NATIVE_DATA_SOURCE_PACKAGE.INVALID_LIBRARY") + assert(e.getMessageParameters.get("path") == notALibrary.getCanonicalPath) + assert(e.getMessageParameters.get("library") == path) + + // The library implements another version of the interface than its manifest says. + val mismatchedAbi = createPackage(abiV2Library, dataSources("v2_"), dir = dir, + fileName = "mismatched_abi.sparkpkg") + checkError( + exception = intercept[SparkRuntimeException] { + newSession(mismatchedAbi).read.format("v2_native_range").load() + }, + condition = "INVALID_NATIVE_DATA_SOURCE_PACKAGE.UNSUPPORTED_ABI_VERSION", + parameters = Map( + "path" -> mismatchedAbi.getCanonicalPath, "version" -> "2", "supported" -> "1")) + } + + nativeTest("ignore invalid packages that do not provide the data source") { + val dir = Utils.createTempDir() + createPackage(library, dataSources(""), dir = dir, fileName = "a.sparkpkg") + createZip(dir, "no_manifest.sparkpkg", "x" -> Array.emptyByteArray) + Files.writeString(new File(dir, "not_a_zip.sparkpkg").toPath, "not a zip") + val otherSource = createZip(dir, "other_source.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + Seq("other_source"), Map(NativePlatform.current -> "lib.so")).getBytes(UTF_8)) + val session = newSession(dir) + checkRows(range(session, "end" -> "2"), (0L until 2L).map(expectedRow)) + checkError( + exception = intercept[SparkClassNotFoundException] { + session.read.format("com.typo.Missing").load() + }, + condition = "DATA_SOURCE_NOT_FOUND", + parameters = Map("provider" -> "com.typo.Missing")) + val e = intercept[AnalysisException](session.sql("SELECT * FROM missing_db.missing_table")) + assert(e.getCondition == "TABLE_OR_VIEW_NOT_FOUND") + // The invalid package whose manifest lists the data source reports why it is invalid, when + // no other package or installed library provides the data source. + checkError( + exception = intercept[SparkRuntimeException] { + session.read.format("other_source").load() + }, + condition = "INVALID_NATIVE_DATA_SOURCE_PACKAGE.INVALID_MANIFEST", + parameters = Map( + "path" -> otherSource.getCanonicalPath, + "reason" -> "The package does not contain the library lib.so.")) + } + + nativeTest("stale packages do not hide a valid one") { + // A shared directory still has broken and outdated copies of the package of native_range. + val dir = Utils.createTempDir() + createZip(dir, "native_range-1.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + dataSources(""), Map(NativePlatform.current -> "lib.so")).getBytes(UTF_8)) + createPackage(library, dataSources(""), dir = dir, fileName = "native_range-2.sparkpkg", + abiVersion = 2) + val session = newSession(dir) + checkError( + exception = intercept[SparkRuntimeException](range(session)), + condition = "INVALID_NATIVE_DATA_SOURCE_PACKAGE.UNSUPPORTED_ABI_VERSION", + parameters = Map( + "path" -> new File(dir, "native_range-2.sparkpkg").getCanonicalPath, + "version" -> "2", + "supported" -> "1")) + + // A valid package added to the session takes precedence over them, without a conflict. + session.addArtifact(defaultPackage.getPath) + checkRows(range(session, "end" -> "2"), (0L until 2L).map(expectedRow)) + + // So does an installed library. + val libraries = Utils.createTempDir() + installLibrary(installedLibrary, "installed_native_range", libraries) + createZip(dir, "installed-1.sparkpkg", + NativeDataSourcePackage.MANIFEST_NAME -> manifestJson( + Seq("installed_native_range"), Map(NativePlatform.current -> "lib.so")).getBytes(UTF_8)) + withLibraryPath(libraries) { + checkRows( + newSession(dir).read.format("installed_native_range").option("end", "2").load() + .select("id"), + Seq(Row(0L), Row(1L))) + } + } + + nativeTest("log the packages of a version that Spark does not support") { + // An operator left an outdated package in a configured directory. + val dir = Utils.createTempDir() + val outdated = createPackage(library, Seq("other_native_range"), dir = dir, + fileName = "outdated.sparkpkg", abiVersion = 2) + val logAppender = new LogAppender("packages of an unsupported version") + withLogAppender(logAppender, level = Some(org.apache.logging.log4j.Level.WARN)) { + checkRows(range(newSession(dir, defaultPackage), "end" -> "2"), + (0L until 2L).map(expectedRow)) + } + assert(logAppender.loggingEvents.exists { event => + val message = event.getMessage.getFormattedMessage + message.contains(outdated.getCanonicalPath) && + message.contains("implements version 2 of the native data source interface") + }) + } + + nativeTest("forget the packages of a session that is cleaned up") { + val session = spark.newSession() + session.addArtifact(defaultPackage.getPath) + val added = session.artifactManager.getNativeDataSourcePackages._1.head.getCanonicalPath + NativeDataSourcePackage.read(new File(added)) + assert(NativeDataSourcePackage.isCached(added)) + session.artifactManager.cleanUpResourcesForTesting() + assert(!NativeDataSourcePackage.isCached(added)) + + // Only the packages under the directory are forgotten. + val dir = Utils.createTempDir() + val pkg = createPackage(library, dataSources(""), dir = dir).getCanonicalPath + NativeDataSourcePackage.read(new File(pkg)) + NativeDataSourcePackage.forgetPackagesIn(new File(dir.getPath + "x")) + assert(NativeDataSourcePackage.isCached(pkg)) + NativeDataSourcePackage.forgetPackagesIn(dir) + assert(!NativeDataSourcePackage.isCached(pkg)) + } + + nativeTest("invalid installed libraries") { + val dir = Utils.createTempDir() + val notALibrary = new File(dir, NativeDataSourceRegistry.libraryFileName("not_a_library")) + Files.writeString(notALibrary.toPath, "not a library") + val newerAbi = installLibrary(abiV2Library, "v2_native_range", dir) + withLibraryPath(dir) { + val e = intercept[SparkRuntimeException](newSession().read.format("not_a_library").load()) + assert(e.getCondition == "INVALID_NATIVE_DATA_SOURCE_LIBRARY.CANNOT_LOAD") + assert(e.getMessageParameters.get("path") == notALibrary.getCanonicalPath) + + // The library stays loaded, and the error is the same when it is used again. + (1 to 2).foreach { _ => + checkError( + exception = intercept[SparkRuntimeException] { + newSession().read.format("v2_native_range").load() + }, + condition = "INVALID_NATIVE_DATA_SOURCE_LIBRARY.UNSUPPORTED_ABI_VERSION", + parameters = Map( + "path" -> newerAbi.getCanonicalPath, "version" -> "2", "supported" -> "1")) + } + } + } + + nativeTest("reject the libraries that spark_ds_init does not initialize") { + def check(library: File, prefix: String, reason: String): Unit = { + val pkg = createPackage(library, dataSources(prefix)) + // The library stays loaded, and the error is the same when it is used again: Spark + // initializes a library only once. + (1 to 2).foreach { _ => + checkError( + exception = intercept[SparkRuntimeException] { + newSession(pkg).read.format(prefix + "native_range").load() + }, + condition = "INVALID_NATIVE_DATA_SOURCE_PACKAGE.INVALID_LIBRARY", + parameters = Map( + "path" -> pkg.getCanonicalPath, + "library" -> s"${NativePlatform.current}/${library.getName}", + "reason" -> reason)) + } + } + val logAppender = new LogAppender("spark_ds_init") + withLogAppender(logAppender, loggerNames = Seq(classOf[NativeLibrary].getName)) { + check(failingInitLibrary, "if_", "spark_ds_init failed with the error code 1.") + } + // The library logged why it failed, once. + assert(logAppender.loggingEvents.count { event => + event.getLevel == org.apache.logging.log4j.Level.ERROR && + event.getMessage.getFormattedMessage.contains("injected failure in spark_ds_init") + } == 1) + check(noTableLibrary, "nt_", "spark_ds_init returned no function table.") + check(noInitLibrary, "ni_", "The library does not export the function spark_ds_init.") + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala new file mode 100644 index 000000000000..938f838c6f6b --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala @@ -0,0 +1,167 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.io.{File, FileOutputStream} +import java.nio.charset.StandardCharsets.UTF_8 +import java.nio.file.Files +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicInteger +import java.util.zip.{ZipEntry, ZipOutputStream} + +import org.apache.spark.util.Utils + +/** + * Builds the native data source library of the tests, + * `native-datasource/test_native_datasource.cc`, with the C++ compiler of the machine, and + * packages or installs it. The tests that load the library are skipped when [[toolchain]] is not + * defined, which includes the versions of Java that cannot load native data sources. + */ +object NativeDataSourceTestUtils { + + private lazy val buildDir = Utils.createTempDir(namePrefix = "native-datasource-test") + + private val libraries = new ConcurrentHashMap[String, File]() + + // The executors store the files of the sessions that are not isolated in the same directory, so + // each package has its own name. + private val packageCount = new AtomicInteger() + + /** The C++ compiler, if available, and the include directory with `spark_datasource.h`. */ + lazy val compiler: Option[(String, Seq[String])] = { + Seq("c++", "g++", "clang++").find(Utils.checkCommandAvailable).map { compiler => + val header = new File(buildDir, "spark_datasource.h") + val in = getClass.getClassLoader + .getResourceAsStream("org/apache/spark/sql/datasource/spark_datasource.h") + Utils.tryWithResource(in)(Files.copy(_, header.toPath)) + (compiler, Seq(buildDir.getPath)) + } + } + + /** [[compiler]], if available on a version of Java that can load native data sources. */ + lazy val toolchain: Option[(String, Seq[String])] = + compiler.filter(_ => NativeLibrary.isSupported) + + /** + * Compiles the test library with the given preprocessor definitions, once for each name. The + * library implements the data sources [[dataSources]], with the prefix `NAME_PREFIX`. + */ + def compile(name: String, defines: String*): File = { + libraries.computeIfAbsent(name, _ => build(name, defines)) + } + + private def build(name: String, defines: Seq[String]): File = { + val (compiler, includes) = toolchain.get + val source = new File(buildDir, "test_native_datasource.cc") + if (!source.exists()) { + val in = getClass.getClassLoader + .getResourceAsStream("native-datasource/test_native_datasource.cc") + Utils.tryWithResource(in)(Files.copy(_, source.toPath)) + } + val library = new File(buildDir, System.mapLibraryName(name)) + // The symbols are hidden, as libraries are often built, so only spark_ds_init is exported. + run(Seq(compiler, "-std=c++17", "-O1", "-shared", "-fPIC", "-fvisibility=hidden") ++ + includes.map("-I" + _) ++ defines ++ Seq("-o", library.getPath, source.getPath)) + library + } + + /** Compiles a C++ program that includes `spark_datasource.h`, runs it and returns its output. */ + def runProgram(name: String, source: String): String = { + val (compiler, includes) = this.compiler.get + val sourceFile = new File(buildDir, s"$name.cc") + Files.writeString(sourceFile.toPath, source) + val program = new File(buildDir, name) + run(Seq(compiler, "-std=c++17") ++ includes.map("-I" + _) ++ + Seq("-o", program.getPath, sourceFile.getPath)) + run(Seq(program.getPath)) + } + + /** Runs a command, checks that it succeeds and returns its output. */ + private def run(command: Seq[String]): String = { + val process = new ProcessBuilder(command: _*).redirectErrorStream(true).start() + val output = new String(process.getInputStream.readAllBytes(), UTF_8) + assert(process.waitFor() == 0, s"Failed to run ${command.mkString(" ")}: $output") + output + } + + /** The data sources of the test library compiled with the given `NAME_PREFIX`. */ + def dataSources(prefix: String): Seq[String] = + Seq("native_range", "native_sink", "native_counter").map(prefix + _) + + def manifestJson( + dataSources: Seq[String], + libraries: Map[String, String], + abiVersion: Int = 1): String = { + val names = dataSources.map(n => "\"" + n + "\"").mkString(", ") + val paths = libraries.map { case (p, path) => "\"" + p + "\": \"" + path + "\"" } + .mkString(", ") + s"""{"abiVersion": $abiVersion, "dataSources": [$names], "libraries": {$paths}}""" + } + + def createZip(dir: File, fileName: String, entries: (String, Array[Byte])*): File = { + dir.mkdirs() + val file = new File(dir, fileName) + Utils.tryWithResource(new ZipOutputStream(new FileOutputStream(file))) { zip => + entries.foreach { case (path, bytes) => + zip.putNextEntry(new ZipEntry(path)) + zip.write(bytes) + zip.closeEntry() + } + } + file + } + + /** Creates a package with the library built for this platform. */ + def createPackage( + library: File, + dataSources: Seq[String], + dir: File = Utils.createTempDir(), + fileName: String = s"test-${packageCount.incrementAndGet()}.sparkpkg", + abiVersion: Int = 1, + platform: String = NativePlatform.current): File = { + val path = s"$platform/${library.getName}" + createZip( + dir, + fileName, + NativeDataSourcePackage.MANIFEST_NAME -> + manifestJson(dataSources, Map(platform -> path), abiVersion).getBytes(UTF_8), + path -> Files.readAllBytes(library.toPath)) + } + + /** Installs a library as the library of the given native data source, in a directory. */ + def installLibrary(library: File, dataSource: String, dir: File): File = { + val file = new File(dir, NativeDataSourceRegistry.libraryFileName(dataSource)) + Files.copy(library.toPath, file.toPath) + file + } + + /** Runs `f` with a directory added to `java.library.path`, in the native library path. */ + def withLibraryPath[T](dir: File)(f: => T): T = { + val original = System.getProperty("java.library.path") + System.setProperty( + "java.library.path", (dir.getPath +: Option(original).toSeq).mkString(File.pathSeparator)) + try { + f + } finally { + if (original == null) { + System.clearProperty("java.library.path") + } else { + System.setProperty("java.library.path", original) + } + } + } +}