diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index b118c9d2774f9..8261d4e3632f8 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 7051e09a6383a..53dec42594efa 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 3a796a031afaf..fbc26caba7ee7 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -4792,6 +4792,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 +6313,36 @@ ], "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 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" + }, "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 50f978819c3da..b088d623ea7cb 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 4a2bd09910129..f7c29cc2eae81 100755 --- a/dev/make-distribution.sh +++ b/dev/make-distribution.sh @@ -223,6 +223,13 @@ 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, and its messages +mkdir -p "$DISTDIR/include" +cp "$SPARK_HOME/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h" \ + "$DISTDIR/include/" +cp "$SPARK_HOME/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto" \ + "$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 a100f0ccb1ce3..3112614546113 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 0000000000000..138645e859979 --- /dev/null +++ b/docs/sql-native-data-sources.md @@ -0,0 +1,780 @@ +--- +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), with the +messages of `spark_datasource.proto`. 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`. + +### 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 the six C functions of `spark_datasource.h`, which never +change: it needs no JVM headers or bindings. Everything else is in the messages of +[`spark_datasource.proto`]({{site.SPARK_GITHUB_URL}}/blob/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto), +which Spark and the library exchange as [Protocol Buffers](https://protobuf.dev) bytes: the +library generates their code with a Protocol Buffers compiler, such as `protoc` or `prost`. The +data itself is exchanged without copying it, as Arrow C streams. Spark looks the functions up in +each library separately, so any number of native data sources can be loaded side by side. + +| Function | What it does | +|----------|--------------| +| `spark_ds_open` | Opens an object, with an `OpenRequest` that says which object and carries everything that it needs, and replies with an `OpenReply` | +| `spark_ds_call` | Runs a control operation on an object, with a `CallRequest`, and replies with a `CallReply` | +| `spark_ds_stream` | Exchanges the data of an object, as Arrow C streams, with a `StreamRequest` | +| `spark_ds_close` | Releases an object | +| `spark_ds_abi_version` | Returns the version of the interface, `SPARK_DS_ABI_VERSION` | +| `spark_ds_last_error` | Returns the message of the last error on the calling thread | + +A Spark distribution contains the header and the messages in its `include` directory, as +`$SPARK_HOME/include/spark_datasource.h` and `$SPARK_HOME/include/spark_datasource.proto`, and the +`spark-sql` JAR contains the header as `org/apache/spark/sql/datasource/spark_datasource.h`. They +only depend on the C standard library and on Protocol Buffers, so a library can also copy them +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 +curl -LO {{site.SPARK_GITHUB_URL}}/raw/v{{site.SPARK_VERSION_SHORT}}/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto +``` + +A library built against them works with all the releases of Spark that support its version, +`SPARK_DS_ABI_VERSION`: see [Compatibility](#compatibility). + +Spark opens and uses the objects of a library as follows: + +- **Data sources**: for each schema inference, scan and write, Spark opens a data source on the + driver, with the name of the data source and the options of the user. Its `OpenReply` declares + the capabilities of the data source, such as `CAPABILITY_BATCH_READ` or + `CAPABILITY_STREAMING_WRITE`, and Spark reports the operations that a data source does not + declare as unsupported. When the user does not specify a schema, Spark asks for it with + `InferSchema`. Spark closes the data source right after. +- **Batch reads**: on the driver, Spark opens a batch reader, with the schema to read, pushes + operations down with `PushPredicates`, `PushLimit` and `PruneColumns`, plans the partitions with + `PlanPartitions`, and serializes the state of the reader with `Serialize`. Then, on the + executors, it opens a partition reader for each partition, in a task, and `Read` initializes an + `ArrowArrayStream` with the batches of the partition. +- **Streaming reads**: on the driver, Spark opens a stream reader. For each micro-batch, it asks + for the `LatestOffset`, plans the partitions with `PlanStreamPartitions`, and reads them 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 + opens a writer, with the mode of the write, `WRITE_MODE_APPEND` or `WRITE_MODE_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 opens 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 opens a data writer, and `Write` passes it an `ArrowArrayStream` with the + data of the task, which the library reads during the call. Then the task calls + `CommitDataWriter`, or `AbortDataWriter` if the task or the commit fails, and 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 opens a new + writer for each micro-batch, and Spark closes the previous one when it opens the next one. + Because Spark is not told when a streaming query stops, the last 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 bytes. A library replies to a +function by passing the encoded reply to a `SparkDsSink` callback, which copies it, or not at all +for an empty reply, which is valid for every request: for example, a reader that does not support +pushdown replies to `PushPredicates` with nothing, and Spark pushes nothing down. A function reports +an error by returning a non-zero value, with the message returned by `spark_ds_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 object is released exactly once, with `spark_ds_close`, which cannot fail and can +run after the query, so do the work that can fail, such as flushing data, in a commit. Strings are +UTF-8, and the JSON of the predicates is ASCII, with the other characters escaped. + +`CAPABILITY_TRUNCATE` declares that a writer supports `WRITE_MODE_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) +and `spark_datasource.proto` specify each function and message, including the ownership of the +Arrow structs and the JSON format of the pushed predicates. `PushPredicatesReply` tells Spark how +the reader uses each predicate: not at all, only to skip data (`PREDICATE_USE_PARTIAL`), or to +only return the rows for which it is true (`PREDICATE_USE_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, and strings +compare by their UTF-8 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. 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. 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. + +Partition readers and data writers get the `TaskContext` of their task, with its IDs, and the +`is_interrupted` function of `SparkDsHost` returns whether a 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. `SparkDsHost` also has a `log` function, which writes +to the log of the driver or of the executor. + +### Compatibility + +The six functions of `spark_datasource.h` never change, nor do the existing fields of +`SparkDsHost`, which starts with its size and only gets fields appended: a library must only read +the fields that fit in the size of the struct that Spark passes. Everything else grows in the +messages, as Protocol Buffers messages do: + +- a later release of Spark may add fields, messages and values of enums, but never changes or + reuses the number of a field. Both sides ignore the fields that they do not know, so a library + built against an earlier version of the messages keeps working with later releases of Spark, and + a library built against a later version works with earlier releases, without the features that + they do not know; +- Spark only sends a new kind of request to an object whose `OpenReply` declares that it supports + it, so a new request never reaches an older library; +- new values may appear where the messages say so: for example, later releases may push down + predicates with other names, which a library must ignore if it does not know them. + +The version of the interface, 1, is `SPARK_DS_ABI_VERSION`, which `spark_ds_abi_version` returns +and the manifest of a package declares. 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 fields and messages of `spark_datasource.proto`: + +- 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, + and continuous streaming; +- 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 in the requests 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 version of `spark_datasource.h` that the library implements, currently 1. +- `dataSources` lists the names of the data sources that the library implements. The names are + case-insensitive, and Spark passes the name in the requests of the library. +- `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. + +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. Give packages distinct +file names, for example with their version: sessions that are not isolated share the files they +add. 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 depends on +the `arrow-array` and `arrow-schema` crates of [arrow-rs](https://github.com/apache/arrow-rs), +which export streams with `FFI_ArrowArrayStream`, and on [prost](https://github.com/tokio-rs/prost), +whose `prost-build` generates the code of the messages from `spark_datasource.proto` when the +library is built, with `protoc`. It only opens data sources, batch readers and partition readers, +and replies to the requests that it does not support with empty replies. `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"] } +prost = "0.14" + +[build-dependencies] +prost-build = "0.14" +``` + +`build.rs`, which generates the messages from the `include` directory of a Spark distribution: + +```rust +// Generates the code of the messages from spark_datasource.proto, in the include directory of a +// Spark distribution. prost-build runs protoc, which must be installed. +fn main() -> std::io::Result<()> { + let include = format!("{}/include", std::env::var("SPARK_HOME").expect("SPARK_HOME")); + prost_build::compile_protos(&[format!("{include}/spark_datasource.proto")], &[include]) +} +``` + +`src/lib.rs`: + +{% 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, CString}; +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::sync::Arc; + +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use prost::Message; + +/// The messages of `spark_datasource.proto`. +mod proto { + include!(concat!(env!("OUT_DIR"), "/spark.datasource.rs")); +} + +use proto::{call_reply, call_request, data_type, open_request, stream_request}; + +type Result = std::result::Result>; + +/// `SparkDsSink`: receives the reply of a function, and copies it. +type Sink = unsafe extern "C" fn(sink: *mut c_void, data: *const u8, length: i64); + +/// What a handle refers to. +enum Object { + DataSource, + Reader { end: i64 }, + Partition { start: i64, end: i64 }, +} + +thread_local! { + static LAST_ERROR: RefCell = RefCell::new(CString::default()); +} + +fn arrow_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])) +} + +/// The schema of the data, as the Spark SQL types of the messages. +fn spark_schema() -> proto::StructType { + let long = data_type::Kind::Long(data_type::Long {}); + proto::StructType { + fields: vec![proto::StructField { + name: "id".to_string(), + data_type: Some(proto::DataType { kind: Some(long) }), + nullable: false, + }], + } +} + +/// Runs `f`, and reports its error or panic to Spark: neither may cross into the JVM. +fn call(f: impl FnOnce() -> Result<()>) -> c_int { + let message = match catch_unwind(AssertUnwindSafe(f)) { + Ok(Ok(())) => return 0, + Ok(Err(e)) => e.to_string(), + Err(_) => "the data source panicked".to_string(), + }; + LAST_ERROR.with(|error| *error.borrow_mut() = CString::new(message).unwrap_or_default()); + 1 +} + +unsafe fn decode(request: *const u8, length: i64) -> Result { + Ok(M::decode(std::slice::from_raw_parts(request, length as usize))?) +} + +unsafe fn reply(reply: Sink, sink: *mut c_void, message: &impl Message) { + let bytes = message.encode_to_vec(); + reply(sink, bytes.as_ptr(), bytes.len() as i64); +} + +#[no_mangle] +pub extern "C" fn spark_ds_abi_version() -> i32 { + 1 +} + +#[no_mangle] +pub extern "C" fn spark_ds_last_error() -> *const c_char { + LAST_ERROR.with(|error| error.borrow().as_ptr()) +} + +#[no_mangle] +pub unsafe extern "C" fn spark_ds_open( + _host: *const c_void, + request: *const u8, + length: i64, + sink_function: Sink, + sink: *mut c_void, + handle: *mut *mut c_void, +) -> c_int { + call(|| { + let object = match decode::(request, length)?.object { + Some(open_request::Object::DataSource(_)) => { + let capabilities = vec![proto::Capability::BatchRead as i32]; + reply(sink_function, sink, &proto::OpenReply { capabilities }); + Object::DataSource + } + Some(open_request::Object::BatchReader(reader)) => { + let options = reader.data_source.unwrap_or_default().options; + let end = options.get("end").map(|end| end.parse()).transpose()?; + Object::Reader { end: end.unwrap_or(10) } + } + Some(open_request::Object::PartitionReader(reader)) => { + let bytes = reader.partition; + let start = i64::from_le_bytes(bytes[0..8].try_into()?); + let end = i64::from_le_bytes(bytes[8..16].try_into()?); + Object::Partition { start, end } + } + _ => return Err("this data source only supports batch reads".into()), + }; + *handle = Box::into_raw(Box::new(object)) as *mut c_void; + Ok(()) + }) +} + +#[no_mangle] +pub unsafe extern "C" fn spark_ds_call( + handle: *mut c_void, + request: *const u8, + length: i64, + sink_function: Sink, + sink: *mut c_void, +) -> c_int { + call(|| { + let object = &*(handle as *const Object); + let result = match (decode::(request, length)?.call, object) { + (Some(call_request::Call::InferSchema(_)), _) => { + let schema = Some(spark_schema()); + call_reply::Reply::InferSchema(proto::InferSchemaReply { schema }) + } + (Some(call_request::Call::PlanPartitions(_)), Object::Reader { end }) => { + let partitions = [(0, end / 2), (end / 2, *end)] + .iter() + .map(|(start, end)| [start.to_le_bytes(), end.to_le_bytes()].concat()) + .collect(); + call_reply::Reply::PlanPartitions(proto::PlanPartitionsReply { partitions }) + } + // The other requests, such as the pushdown, have empty replies. + _ => return Ok(()), + }; + reply(sink_function, sink, &proto::CallReply { reply: Some(result) }); + Ok(()) + }) +} + +#[no_mangle] +pub unsafe extern "C" fn spark_ds_stream( + handle: *mut c_void, + request: *const u8, + length: i64, + _input: *mut FFI_ArrowArrayStream, + output: *mut FFI_ArrowArrayStream, +) -> c_int { + call(|| { + let Object::Partition { start, end } = &*(handle as *const Object) else { + return Err("not a partition reader".into()); + }; + let stream = decode::(request, length)?.stream; + let Some(stream_request::Stream::Read(_)) = stream else { + return Err("this data source only supports reads".into()); + }; + let ids = Int64Array::from_iter_values(*start..*end); + let batch = RecordBatch::try_new(arrow_schema(), vec![Arc::new(ids)])?; + let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema()); + std::ptr::write(output, FFI_ArrowArrayStream::new(Box::new(batches))); + Ok(()) + }) +} + +#[no_mangle] +pub unsafe extern "C" fn spark_ds_close(handle: *mut c_void) { + drop(Box::from_raw(handle as *mut Object)); +} +``` +{% 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 + +#include "spark_datasource.h" +#include "spark_datasource.pb.h" + +namespace { + +namespace sds = spark::datasource; + +// What a handle refers to. +struct DataSource {}; +struct Reader { + int64_t end; +}; +struct Partition { + int64_t start; + int64_t end; +}; +using Object = std::variant; + +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()); +} + +template +M Decode(const uint8_t* request, int64_t length) { + M message; + if (!message.ParseFromArray(request, static_cast(length))) { + throw std::runtime_error("invalid request"); + } + return message; +} + +void Reply(SparkDsSink reply, void* sink, const google::protobuf::MessageLite& message) { + std::string bytes = message.SerializeAsString(); + reply(sink, reinterpret_cast(bytes.data()), static_cast(bytes.size())); +} + +} // namespace + +extern "C" { + +int32_t spark_ds_abi_version(void) { return SPARK_DS_ABI_VERSION; } + +const char* spark_ds_last_error(void) { return last_error.c_str(); } + +int spark_ds_open(const SparkDsHost*, const uint8_t* request, int64_t length, SparkDsSink reply, + void* sink, void** handle) { + return Call([&] { + auto open = Decode(request, length); + std::unique_ptr object; + if (open.has_data_source()) { + sds::OpenReply capabilities; + capabilities.add_capabilities(sds::CAPABILITY_BATCH_READ); + Reply(reply, sink, capabilities); + object = std::make_unique(DataSource{}); + } else if (open.has_batch_reader()) { + const auto& options = open.batch_reader().data_source().options(); + auto end = options.find("end"); + object = std::make_unique( + Reader{end == options.end() ? 10 : std::stoll(end->second)}); + } else if (open.has_partition_reader()) { + int64_t bounds[2]; + std::memcpy(bounds, open.partition_reader().partition().data(), sizeof(bounds)); + object = std::make_unique(Partition{bounds[0], bounds[1]}); + } else { + throw std::runtime_error("this data source only supports batch reads"); + } + *handle = object.release(); + }); +} + +int spark_ds_call(void* handle, const uint8_t* request, int64_t length, SparkDsSink reply, + void* sink) { + return Call([&] { + auto call = Decode(request, length); + sds::CallReply result; + if (call.has_infer_schema()) { + auto* field = result.mutable_infer_schema()->mutable_schema()->add_fields(); + field->set_name("id"); + field->mutable_data_type()->mutable_long_(); + field->set_nullable(false); + } else if (call.has_plan_partitions()) { + int64_t end = std::get(*static_cast(handle)).end; + int64_t bounds[2][2] = {{0, end / 2}, {end / 2, end}}; + for (const auto& partition : bounds) { + result.mutable_plan_partitions()->add_partitions(partition, sizeof(partition)); + } + } else { + return; // The other requests, such as the pushdown, have empty replies. + } + Reply(reply, sink, result); + }); +} + +int spark_ds_stream(void* handle, const uint8_t* request, int64_t length, struct ArrowArrayStream*, + struct ArrowArrayStream* out) { + return Call([&] { + if (!Decode(request, length).has_read()) { + throw std::runtime_error("this data source only supports reads"); + } + const auto& partition = std::get(*static_cast(handle)); + arrow::Int64Builder ids; + for (int64_t id = partition.start; id < partition.end; 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)); + }); +} + +void spark_ds_close(void* handle) { delete static_cast(handle); } + +} // extern "C" +``` +{% endraw %} + +Generate the messages and build the library, here on macOS with Arrow C++ 24, which requires +C++20, and install it, or package it as above: + +```bash +protoc --cpp_out=. -I"$SPARK_HOME/include" "$SPARK_HOME/include/spark_datasource.proto" +c++ -std=c++20 -O2 -shared -fPIC -I. -I"$SPARK_HOME/include" -I"$ARROW_HOME/include" \ + cpp_range.cc spark_datasource.pb.cc $(pkg-config --cflags --libs protobuf-lite) \ + -L"$ARROW_HOME/lib" -larrow -o libspark_datasource_cpp_range.dylib +cp libspark_datasource_cpp_range.dylib /usr/local/lib/ +``` + +A library loaded by Spark must find its dependencies, such as `libarrow` and `libprotobuf-lite`, on +every node: link them statically, or install them on the nodes. Several native data source +libraries, or several copies of one, such as an installed library and the library of a package, can +share a process. The messages use the lite runtime of Protocol Buffers, which, unlike the full one, +does not register them in a pool of the process, where a second copy would abort the process. But +libraries built against different versions of Protocol Buffers conflict if they link it +dynamically, so link it statically, with hidden symbols. + +## 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 4b4d0f35cb4e4..df1674e98ea1c 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 da05308883570..c2eb3de50b8cd 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 3ab3df4ff5938..f39362a0e153e 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 984b02ef6e562..916d4a1640ea6 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 6e5ed71745950..286f35d539a63 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 408b33f1ed743..73f18abc9676e 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 40c7e532c27db..e4fb52e928391 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 8aa554a879a15..cce193fe50db8 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 c16a7d20dfacf..7c82867a2c3f1 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 fd44b835a0ace..5f940e5126473 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,57 @@ 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 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 d3211aba4986f..787b58857116c 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 e02c7296f9378..6da375360c4e5 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 0000000000000..97c4b3ae61046 --- /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 281546b901a7a..7f3db15f70fdf 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 0000000000000..c9edcf0b0b6b2 --- /dev/null +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java @@ -0,0 +1,630 @@ +/* + * 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 com.google.protobuf.InvalidProtocolBufferException; + +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.sql.execution.datasources.v2.ffi.proto.CallReply; +import org.apache.spark.sql.execution.datasources.v2.ffi.proto.CallRequest; +import org.apache.spark.sql.execution.datasources.v2.ffi.proto.OpenReply; +import org.apache.spark.sql.execution.datasources.v2.ffi.proto.OpenRequest; +import org.apache.spark.sql.execution.datasources.v2.ffi.proto.StreamRequest; +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, with the messages of + * {@code org/apache/spark/sql/datasource/spark_datasource.proto}. + *

+ * Spark compiles for Java 17, where the FFM API does not exist, so the API is only reached + * through method handles, looked up once when the first library is loaded. Each of the six + * functions that the library exports 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. + *

+ * A function that fails throws a {@link RuntimeException} with the message of + * {@code spark_ds_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 newest version of the messages that Spark supports. */ + static final int PROTOCOL_VERSION = 1; + + /** The first version of Java with the final FFM API. */ + static final int MINIMUM_JAVA_VERSION = 22; + + /** Whether this JVM can load native data source libraries. */ + static boolean isSupported() { + return Runtime.version().feature() >= MINIMUM_JAVA_VERSION; + } + + // The C types of the parameters and the results of the functions. + private enum C { VOID, BOOL, I32, I64, PTR } + + // The layout of SparkDsHost on the 64-bit platforms that Spark supports. + private static final long HOST_SIZE = 24; + private static final long HOST_LOG = 8; + private static final long HOST_IS_INTERRUPTED = 16; + + /** + * The offsets of the fields of the structs that Spark 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("SparkDsHost", HOST_SIZE); + layout.put("SparkDsHost.size", 0L); + layout.put("SparkDsHost.log", HOST_LOG); + layout.put("SparkDsHost.is_interrupted", HOST_IS_INTERRUPTED); + return layout; + } + + private final long host; + private final int abiVersion; + private final MethodHandle lastError; + private final MethodHandle open; + private final MethodHandle call; + private final MethodHandle stream; + private final MethodHandle close; + + private NativeLibrary(Object lookup, String path) throws Throwable { + abiVersion = (int) function(lookup, "spark_ds_abi_version", C.I32).invokeExact(); + // The functions of another version of the interface may differ: the caller rejects the + // library before any of them is called. + boolean supported = abiVersion == ABI_VERSION; + lastError = supported ? optionalFunction(lookup, "spark_ds_last_error", C.PTR) : null; + open = supported + ? function(lookup, "spark_ds_open", C.I32, C.PTR, C.PTR, C.I64, C.PTR, C.PTR, C.PTR) + : null; + call = supported + ? function(lookup, "spark_ds_call", C.I32, C.PTR, C.PTR, C.I64, C.PTR, C.PTR) + : null; + stream = supported + ? function(lookup, "spark_ds_stream", C.I32, C.PTR, C.PTR, C.I64, C.PTR, C.PTR) + : null; + close = supported ? function(lookup, "spark_ds_close", C.VOID, C.PTR) : null; + // The library may keep the host until the JVM exits. + host = supported ? Ffm.host(path) : 0; + } + + private static MethodHandle function(Object lookup, String name, C result, C... parameters) + throws Throwable { + MethodHandle function = optionalFunction(lookup, name, result, parameters); + if (function == null) { + throw new LoadException( + "The library does not export the function " + name + ".", null, true); + } + return function; + } + + private static MethodHandle optionalFunction( + Object lookup, String name, C result, C... parameters) throws Throwable { + Optional address = (Optional) Ffm.FIND.invoke(lookup, name); + return address.isEmpty() ? null : Ffm.downcall(address.get(), result, parameters); + } + + /** + * Loads the library at the given path. The library stays loaded until the JVM exits. + * + * @throws LoadException if the library cannot be loaded + */ + static NativeLibrary load(String path) throws Throwable { + if (!isSupported()) { + throw new UnsupportedOperationException( + "Native data sources require Java " + MINIMUM_JAVA_VERSION + " or later."); + } + Object lookup; + try { + lookup = Ffm.LIBRARY_LOOKUP.invoke(Path.of(path), Ffm.GLOBAL_ARENA); + } catch (IllegalArgumentException | IllegalCallerException e) { + // The library is not a valid library, or native access is not enabled for Spark. + throw new LoadException(e.getMessage(), e); + } catch (LinkageError e) { + // The FFM API could not be initialized, see Ffm: report why, such as an + // IllegalCallerException when native access is not enabled for Spark. + throw new LoadException( + "Cannot use the Foreign Function and Memory API: " + rootCause(e), e); + } + return new NativeLibrary(lookup, path); + } + + /** 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 SparkDsHost.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.I64); + } 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) throws Throwable { + long log = upcall(MethodHandles.insertArguments(LOG, 0, path), C.VOID, C.I32, C.PTR); + long host = Platform.allocateMemory(HOST_SIZE); + Platform.setMemory(host, (byte) 0, HOST_SIZE); + Platform.putLong(null, host, HOST_SIZE); + Platform.putLong(null, host + HOST_LOG, log); + Platform.putLong(null, host + HOST_IS_INTERRUPTED, IS_INTERRUPTED); + 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 it does not export a function. + */ + 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; + } + } + + /** + * 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 + // --------------------------------------------------------------------------------------------- + + /** The version of the interface that the library implements. */ + int abiVersion() { + return abiVersion; + } + + /** Throws the error that the library reported if a function failed. */ + private void check(int status, String function) throws Throwable { + if (status == 0) { + return; + } + long address = lastError == null ? 0 : (long) lastError.invokeExact(); + String message = address == 0 ? "" : new String(readCString(address), StandardCharsets.UTF_8); + throw new RuntimeException(message.isEmpty() + ? function + " failed with the error code " + status + "." + : message); + } + + /** An object that the library opened, with its handle and the reply of the library. */ + static final class Opened { + final long handle; + final OpenReply reply; + + Opened(long handle, OpenReply reply) { + this.handle = handle; + this.reply = reply; + } + } + + /** Opens an object. */ + Opened open(OpenRequest request) throws Throwable { + try (Memory memory = new Memory(); Sink sink = new Sink()) { + byte[] bytes = request.toByteArray(); + long out = memory.zeroed(8); + check((int) open.invokeExact(host, memory.bytes(bytes), (long) bytes.length, Ffm.SINK, + sink.id, out), "spark_ds_open"); + long handle = Platform.getLong(null, out); + if (handle == 0) { + throw new IllegalStateException("spark_ds_open returned a NULL handle."); + } + try { + return new Opened(handle, OpenReply.parseFrom(sink.single("spark_ds_open"))); + } catch (InvalidProtocolBufferException e) { + // The library opened the object, but its reply is invalid. + close(handle); + throw new IllegalStateException(invalidReply("spark_ds_open", e), e); + } catch (Throwable e) { + close(handle); + throw e; + } + } + } + + /** Runs a control operation on an object. */ + CallReply call(long handle, CallRequest request) throws Throwable { + return call(handle, request, null); + } + + /** + * Runs a control operation on an object. Then {@code succeeded} runs right away, before the + * reply is checked: for example, a data writer that committed must not be aborted, even if its + * reply is invalid. + */ + CallReply call(long handle, CallRequest request, Runnable succeeded) throws Throwable { + try (Memory memory = new Memory(); Sink sink = new Sink()) { + byte[] bytes = request.toByteArray(); + check((int) call.invokeExact(handle, memory.bytes(bytes), (long) bytes.length, Ffm.SINK, + sink.id), "spark_ds_call"); + if (succeeded != null) { + succeeded.run(); + } + try { + return CallReply.parseFrom(sink.single("spark_ds_call")); + } catch (InvalidProtocolBufferException e) { + throw new IllegalStateException(invalidReply("spark_ds_call", e), e); + } + } + } + + /** + * Exchanges the data of an object. + * + * @param in the address of the input stream, or 0 + * @param out the address of the output stream, or 0 + */ + void stream(long handle, StreamRequest request, long in, long out) throws Throwable { + try (Memory memory = new Memory()) { + byte[] bytes = request.toByteArray(); + check((int) stream.invokeExact(handle, memory.bytes(bytes), (long) bytes.length, in, out), + "spark_ds_stream"); + } + } + + /** Releases the handle of an object. */ + void close(long handle) throws Throwable { + close.invokeExact(handle); + } + + // --------------------------------------------------------------------------------------------- + // 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; + } + + @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; + } + + // --------------------------------------------------------------------------------------------- + // Tasks + // --------------------------------------------------------------------------------------------- + + /** A task that uses the library, and the number of objects that it opened. */ + private static final class Task { + final BooleanSupplier interrupted; + int objects; + + Task(BooleanSupplier interrupted) { + this.interrupted = interrupted; + } + } + + // The tasks that use a library, by the ID of their task attempt, for SparkDsHost.is_interrupted. + private static final ConcurrentHashMap tasks = new ConcurrentHashMap<>(); + + /** + * Registers a task for {@code SparkDsHost.is_interrupted}, until the returned + * {@link AutoCloseable} is closed, once the library no longer uses the object of the task. + */ + static AutoCloseable registerTask(long taskAttemptId, BooleanSupplier interrupted) { + tasks.compute(taskAttemptId, (id, task) -> { + Task registered = task == null ? new Task(interrupted) : task; + registered.objects++; + return registered; + }); + return () -> tasks.computeIfPresent(taskAttemptId, (id, task) -> + --task.objects == 0 ? null : task); + } + + /** + * Called by the library, through {@code SparkDsHost.is_interrupted}, with the ID of a task + * attempt. It must not throw: an exception thrown by an upcall crashes the JVM. + */ + private static boolean isInterrupted(long taskAttemptId) { + try { + Task task = tasks.get(taskAttemptId); + // A task that is not registered has ended, or is not a task of this executor. + return task == null || task.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 reply 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); + } + + /** Returns the reply, or no bytes, an empty message, if the library did not reply. */ + byte[] single(String function) throws Throwable { + if (error != null) { + throw error; + } + if (values.size() > 1) { + throw new IllegalStateException( + function + " replied " + values.size() + " times instead of once."); + } + return values.isEmpty() ? new byte[0] : values.get(0); + } + + @Override + public void close() { + sinks.remove(id); + } + } + + /** + * Called by the library, through the upcall stub {@link Ffm#SINK}, with a reply 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 reply, 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; + } + } + } + + /** Describes an invalid reply of the library. */ + private static String invalidReply(String function, InvalidProtocolBufferException e) { + return function + " returned an invalid reply: " + e.getMessage(); + } +} diff --git a/sql/core/src/main/protobuf/buf.yaml b/sql/core/src/main/protobuf/buf.yaml new file mode 100644 index 0000000000000..986dbaa5d43a2 --- /dev/null +++ b/sql/core/src/main/protobuf/buf.yaml @@ -0,0 +1,22 @@ +# +# 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. +# +version: v1 +build: + excludes: + # The messages of native data sources, which PySpark does not use: dev/streaming-gen-protos.sh + # only generates the Python code of the messages of the Python streaming APIs. + - org/apache/spark/sql/datasource diff --git a/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto b/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto new file mode 100644 index 0000000000000..a296f7d4f5e28 --- /dev/null +++ b/sql/core/src/main/protobuf/org/apache/spark/sql/datasource/spark_datasource.proto @@ -0,0 +1,487 @@ +/* + * 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. + */ + +// The messages that Spark and a native data source exchange through the functions of +// spark_datasource.h, which describes how they are used. Since Spark 4.4.0. +// +// The messages only grow: a later release of Spark may add fields, messages and values of enums, +// but it never changes or reuses the number of a field, so a library built against an earlier +// version of this file keeps working, and ignores the fields that it does not know. Spark only +// sends a new kind of request to an object whose reply to spark_ds_open declares that it supports +// it, so a new request never reaches an older library. +syntax = "proto3"; + +package spark.datasource; + +option java_multiple_files = true; +option java_package = "org.apache.spark.sql.execution.datasources.v2.ffi.proto"; +// Lite messages have no descriptors, which the full C++ runtime registers in a pool of the +// process: two libraries that link the same libprotobuf, or two copies of one library, such as an +// installed one and the one of a package, would register spark_datasource.proto twice, which +// aborts the process. A library should still link its Protocol Buffers runtime statically, with +// hidden symbols, since libraries built against different versions of it can share a process. +option optimize_for = LITE_RUNTIME; + +// --------------------------------------------------------------------------------------------- +// spark_ds_open +// --------------------------------------------------------------------------------------------- + +// The request of spark_ds_open: the object to open, with what it needs. +message OpenRequest { + // The version of Spark, such as "4.4.0". + string spark_version = 1; + // The newest version of these messages that Spark supports: 1. + int32 protocol_version = 2; + + oneof object { + // A data source, on the driver, for each schema inference, scan and write. Spark closes it + // right after. + DataSource data_source = 10; + // A reader for a batch scan, on the driver. + BatchReader batch_reader = 11; + // A reader for a micro-batch streaming scan, on the driver. Spark keeps it until the + // streaming query stops. + StreamReader stream_reader = 12; + // A writer for a batch write, or for a micro-batch of a streaming write, on the driver. + Writer writer = 13; + // The reader of one partition of a batch or streaming scan, on an executor. + PartitionReader partition_reader = 14; + // The writer of one task of a write, on an executor. + DataWriter data_writer = 15; + } +} + +// The reply of spark_ds_open. +message OpenReply { + // For a data source: the operations that it supports, which may depend on its name and its + // options. Spark reports the operations that a data source does not declare as unsupported, + // without opening a reader or a writer. + repeated Capability capabilities = 1; +} + +// An operation that a data source supports. +enum Capability { + CAPABILITY_UNSPECIFIED = 0; + // Batch reads, with batch readers. + CAPABILITY_BATCH_READ = 1; + // Micro-batch streaming reads, with stream readers. + CAPABILITY_MICRO_BATCH_READ = 2; + // Batch writes, with writers. + CAPABILITY_BATCH_WRITE = 3; + // Streaming writes, with writers. + CAPABILITY_STREAMING_WRITE = 4; + // Writes with the mode WRITE_MODE_TRUNCATE. + CAPABILITY_TRUNCATE = 5; +} + +// A data source, with the options given by the user, such as spark.read.option("key", "value"). +message DataSource { + // 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. + string name = 1; + // The options, with the keys in lower case, because Spark matches them case-insensitively. + map options = 2; +} + +message BatchReader { + DataSource data_source = 1; + // The schema to read: the one that the user specified, or the one that the data source infers. + StructType schema = 2; +} + +message StreamReader { + DataSource data_source = 1; + // The schema to read: the one that the user specified, or the one that the data source infers. + StructType schema = 2; +} + +message Writer { + DataSource data_source = 1; + // The schema of the data to write. + StructType schema = 2; + // Whether the write adds to the existing data, or replaces it. + WriteMode mode = 3; + // The ID of the streaming query of a streaming write, and empty 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. + string query_id = 4; +} + +enum WriteMode { + // Adds the data to the existing data. + WRITE_MODE_APPEND = 0; + // Replaces 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 uses it with a + // data source that declares CAPABILITY_TRUNCATE. + WRITE_MODE_TRUNCATE = 1; +} + +message PartitionReader { + // The task that reads the partition. + TaskContext task = 1; + // The state that the batch reader or the stream reader returned in its SerializeReply. + bytes reader_state = 2; + // A partition that the batch reader or the stream reader planned. + bytes partition = 3; +} + +message DataWriter { + // The task that writes, with the partition that it writes. + TaskContext task = 1; + // The state that the writer returned in its SerializeReply. + bytes writer_state = 2; + // The ID of the micro-batch for a streaming write, or -1 for a batch write. + int64 epoch_id = 3; +} + +// A task of Spark, on an executor. +message TaskContext { + // The ID of the stage of the task, and the attempt number of the stage. + int32 stage_id = 1; + int32 stage_attempt_number = 2; + // The partition that the task reads or writes, and the attempt number of the task. + int32 partition_id = 3; + int32 attempt_number = 4; + // The ID of the task attempt, unique in the Spark application. SparkDsHost.is_interrupted takes + // it. + int64 task_attempt_id = 5; +} + +// --------------------------------------------------------------------------------------------- +// spark_ds_call +// --------------------------------------------------------------------------------------------- + +// The request of spark_ds_call: a control operation on an object. The comment of each operation +// says on which object Spark calls it, and the reply has the field of the same name. +message CallRequest { + oneof call { + // On a data source, when the user does not specify a schema. + InferSchema infer_schema = 1; + // On a batch reader, at most once and in this order: push_predicates, push_limit and + // prune_columns. + PushPredicates push_predicates = 2; + PushLimit push_limit = 3; + PruneColumns prune_columns = 4; + // On a batch reader, after the pushdown. + PlanPartitions plan_partitions = 5; + // On a batch reader after plan_partitions, on a stream reader after each + // plan_stream_partitions, and on a writer once, after opening it. + Serialize serialize = 6; + // On a stream reader. + InitialOffset initial_offset = 7; + LatestOffset latest_offset = 8; + PlanStreamPartitions plan_stream_partitions = 9; + CommitOffset commit_offset = 10; + // On a writer. + Commit commit = 11; + Abort abort = 12; + // On a data writer, after the data was written. + CommitDataWriter commit_data_writer = 13; + AbortDataWriter abort_data_writer = 14; + } +} + +// The reply of spark_ds_call. An empty reply is valid for every request: it means that the +// operation has no result, such as a reader that does not use the pushed predicates. +message CallReply { + oneof reply { + InferSchemaReply infer_schema = 1; + PushPredicatesReply push_predicates = 2; + PushLimitReply push_limit = 3; + PruneColumnsReply prune_columns = 4; + PlanPartitionsReply plan_partitions = 5; + SerializeReply serialize = 6; + OffsetReply initial_offset = 7; + OffsetReply latest_offset = 8; + PlanPartitionsReply plan_stream_partitions = 9; + CommitDataWriterReply commit_data_writer = 13; + } +} + +message InferSchema {} + +message InferSchemaReply { + // The schema of the data. Without it, the user must specify the schema. + StructType schema = 1; +} + +// Pushes predicates down to a batch reader. The predicates are combined with AND. +// +// Each predicate is a JSON expression tree. An expression is one of: +// - {"type": "column", "name": ["a", "b"]}: a column. A nested field has a name part for each +// level. +// - {"type": "literal", "dataType": "bigint", "value": 1}: a literal. The data type is a Spark +// SQL type name such as int, string or decimal(10,2). The value is null, a JSON boolean or +// number for boolean and numeric types, a string for string, a base64 string for binary, a +// decimal string for decimal, the days since the epoch for date, and the microseconds since +// the epoch for timestamp and timestamp_ntz. +// - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in +// org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=> +// (equal, with null equal to null), <, <=, > and >= with two children, IN with the value to +// look for and then the values to look in, IS_NULL and IS_NOT_NULL with one child, +// STARTS_WITH, ENDS_WITH and CONTAINS with the string and the string to look for, AND and OR +// with two predicates, and NOT with one. The children of the other functions are columns, +// literals or predicates. +// Spark does not pass predicates that it cannot express this way, such as predicates that use +// arithmetic or other functions, nor predicates on columns that contain strings with a +// non-binary collation, such as UTF8_LCASE: strings are compared by their UTF-8 bytes. The JSON +// is ASCII: it escapes the other characters, which any JSON parser decodes. +// +// 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: see spark_datasource.h. The file +// org/apache/spark/sql/datasource/predicate-cases.json in the jar of Spark SQL lists predicates as +// Spark passes them, with the rows of a table for which they are true, to test a reader with. +message PushPredicates { + repeated string predicates_json = 1; +} + +message PushPredicatesReply { + // How the reader uses each predicate, in the order of the request. Spark shows the predicates + // that the reader uses as pushed in the plan of the query, and evaluates all of them again after + // reading anyway. Missing answers are PREDICATE_USE_UNUSED. + repeated PredicateUse uses = 1; +} + +enum PredicateUse { + // Not at all. + PREDICATE_USE_UNUSED = 0; + // Only to skip data, so the reader may return rows for which the predicate is not true. + PREDICATE_USE_PARTIAL = 1; + // To only return the rows for which the predicate is true. + PREDICATE_USE_COMPLETE = 2; +} + +// Pushes a LIMIT down to a batch reader. Spark only pushes a limit when the query does not filter +// the rows, since it evaluates the predicates after reading. +message PushLimit { + int32 limit = 1; +} + +message PushLimitReply { + // Whether the reader returns at most `limit` rows in total. Spark applies the limit again + // after reading either way. + bool applied = 1; +} + +// Prunes the columns that a batch reader reads. +message PruneColumns { + // The names of the 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. + repeated string columns = 1; +} + +message PruneColumnsReply { + // Whether the partition readers return exactly these columns, in this order. If not, they return + // all the columns of the schema to read. + bool pruned = 1; +} + +// Plans the partitions of a batch scan. Spark reads each partition in a separate task. +message PlanPartitions {} + +message PlanPartitionsReply { + // The partitions, serialized in a format defined by the library. + repeated bytes partitions = 1; +} + +// Serializes the state of a reader or a writer that its partition readers or data writers need +// on the executors, such as the pushed predicates. +message Serialize {} + +message SerializeReply { + // The state, in a format defined by the library, or empty if there is none. + bytes state = 1; +} + +// Returns the offset to start from when the streaming query has no checkpoint. +message InitialOffset {} + +// Returns the most recent offset available. +message LatestOffset {} + +message OffsetReply { + // The offset, 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. + string offset_json = 1; +} + +// Plans the partitions that read the data after `start` up to and including `end`. When a +// streaming query restarts, from its checkpoint, Spark may plan the last range that it planned +// before again, so the partitions of a range must read the same data each time. +message PlanStreamPartitions { + string start_json = 1; + string end_json = 2; +} + +// Tells the library that Spark has processed all the data up to and including `end`, and will not +// ask for it again. +message CommitOffset { + string end_json = 1; +} + +// Commits a write, when all its tasks succeeded. Spark may commit the same epoch of a streaming +// query more than once, for example when the query restarts after a failure, so committing an +// epoch that is already committed must have no effect, for exactly-once output. +message Commit { + // The ID of the micro-batch for a streaming write, or -1 for a batch write. + int64 epoch_id = 1; + // The messages that the data writers of the tasks returned, one for each partition. + repeated TaskMessage messages = 2; +} + +// Aborts a write, when it failed, including when its commit failed. +message Abort { + // The ID of the micro-batch for a streaming write, or -1 for a batch write. + int64 epoch_id = 1; + // The messages that the data writers of the tasks returned. A task that did not commit has no + // message. + repeated TaskMessage messages = 2; +} + +message TaskMessage { + // Whether the task committed, and so whether `message` is set. + bool committed = 1; + bytes message = 2; +} + +// Commits the data that a data writer wrote. If it fails, Spark aborts the data writer. If it +// succeeds, Spark does not abort it, even if it then fails the task, for example because the +// reply is invalid. +message CommitDataWriter {} + +message CommitDataWriterReply { + // The message to pass to the commit of the writer on the driver. + bytes message = 1; +} + +// Aborts the data that a data writer wrote. +message AbortDataWriter {} + +// --------------------------------------------------------------------------------------------- +// spark_ds_stream +// --------------------------------------------------------------------------------------------- + +// The request of spark_ds_stream: a data operation on an object, through Arrow C streams. +message StreamRequest { + oneof stream { + // On a partition reader: initializes the output stream with the batches of the partition. + Read read = 1; + // On a data writer: reads the batches of the input stream and writes them. + Write write = 2; + } +} + +message Read {} + +message Write {} + +// --------------------------------------------------------------------------------------------- +// Schemas +// --------------------------------------------------------------------------------------------- + +// A schema: the columns of the data, as Spark SQL types. The data itself is Arrow, with the types +// of spark.sql.execution.arrow: for example, bigint is int64, string is utf8, timestamp is +// timestamp[us] with a time zone and timestamp_ntz is timestamp[us] without one. +message StructType { + repeated StructField fields = 1; +} + +message StructField { + string name = 1; + DataType data_type = 2; + bool nullable = 3; +} + +// A Spark SQL type. +message DataType { + oneof kind { + Null null = 1; + Boolean boolean = 2; + Byte byte = 3; + Short short = 4; + Integer integer = 5; + Long long = 6; + Float float = 7; + Double double = 8; + Decimal decimal = 9; + String string = 10; + Binary binary = 11; + Date date = 12; + Timestamp timestamp = 13; + TimestampNtz timestamp_ntz = 14; + YearMonthInterval year_month_interval = 15; + DayTimeInterval day_time_interval = 16; + CalendarInterval calendar_interval = 17; + Array array = 18; + Map map = 19; + StructType struct = 20; + } + + message Null {} + message Boolean {} + message Byte {} + message Short {} + message Integer {} + message Long {} + message Float {} + message Double {} + + message Decimal { + int32 precision = 1; + int32 scale = 2; + } + + message String { + // The collation of the strings, such as "UTF8_LCASE", or empty for UTF8_BINARY. + string collation = 1; + } + + message Binary {} + message Date {} + message Timestamp {} + message TimestampNtz {} + + // The fields are those of org.apache.spark.sql.types.YearMonthIntervalType: 0 for YEAR and 1 for + // MONTH. + message YearMonthInterval { + int32 start_field = 1; + int32 end_field = 2; + } + + // The fields are those of org.apache.spark.sql.types.DayTimeIntervalType: 0 for DAY, 1 for + // HOUR, 2 for MINUTE and 3 for SECOND. + message DayTimeInterval { + int32 start_field = 1; + int32 end_field = 2; + } + + message CalendarInterval {} + + message Array { + DataType element_type = 1; + bool contains_null = 2; + } + + message Map { + DataType key_type = 1; + DataType value_type = 2; + bool value_contains_null = 3; + } +} 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 0000000000000..332e36d25273f --- /dev/null +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/predicate-cases.json @@ -0,0 +1,44 @@ +{ + "description": "Predicates as Spark passes them in PushPredicates of spark_datasource.proto, 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.", + "schema": "id INT,i INT,d DOUBLE,s STRING", + "rows": [ + [0, 1, 0.0, "a"], + [1, 2, -0.0, ""], + [2, null, "NaN", "\u00E9"], + [3, 3, "Infinity", "\uD83D\uDE00"], + [4, 1, "-Infinity", null], + [5, null, 1.5, "ab"], + [6, 2, null, "b\u00E9"], + [7, 4, 2.5, "\uFFFD"] + ], + "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]} + ] +} 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 0000000000000..1a29b7a192c9c --- /dev/null +++ b/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h @@ -0,0 +1,310 @@ +/* + * 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 the six functions declared below. They + * never change. Everything that the interface grows with is in the messages of + * spark_datasource.proto, which Spark and the library exchange as Protocol Buffers bytes: which + * object to open, and every control operation on it, such as the pushdown of predicates, the + * planning of partitions, offsets and commits. The data itself is exchanged without copying, as + * Arrow C streams. Spark calls the functions with the Foreign Function and Memory API of Java, so + * native data sources require Java 22 or later. 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 and spark_datasource.proto into its source tree, or use the copies in the include + * directory of a Spark distribution, and generate the code of the messages with any Protocol + * Buffers compiler, such as protoc or prost. A library built against them 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. It opens objects with spark_ds_open, whose OpenRequest + * says which object and carries everything that it needs, calls control operations on them with + * spark_ds_call, exchanges their data with spark_ds_stream, and releases them with spark_ds_close. + * The messages are those of spark_datasource.proto. + * + * - Data sources. On the driver, for each schema inference, scan and write, Spark opens a data + * source, whose OpenReply declares its capabilities: Spark reports the operations that a data + * source does not declare as unsupported, without opening a reader or a writer. When the user + * does not specify a schema, Spark calls InferSchema. Spark closes the data source right after. + * - Batch reads. On the driver, Spark opens a batch reader, calls, each at most once and in this + * order, PushPredicates, PushLimit and PruneColumns, then PlanPartitions and Serialize, and + * closes it. Then, on the executors, it opens a partition reader for each partition, in a task, + * initializes its stream with a Read, reads the stream, releases it and closes the partition + * reader. If Spark never plans the partitions, for example to only explain the plan of the query, + * it closes the batch reader once the reader is no longer referenced. + * - Streaming reads. On the driver, Spark opens a stream reader and keeps it until the streaming + * query stops. For each micro-batch, it calls LatestOffset, PlanStreamPartitions and Serialize, + * then reads the partitions on the executors as for a batch read, and calls CommitOffset once it + * has processed the micro-batch. Offsets are defined by the library, as compact JSON on a single + * line. When a streaming query restarts, from its checkpoint, Spark opens a new stream reader and + * may plan the last range that it planned before again, 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 opens a writer and calls Serialize. With adaptive query execution, Spark first runs the + * stages of the data to write that end with a shuffle, and opens 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 opens a data writer, + * passes its data in one or more calls of Write, each with a stream of batches of at most + * spark.sql.execution.arrow.maxRecordsPerBatch rows, which also end once they hold + * spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls CommitDataWriter, or + * AbortDataWriter 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), which the query ID of the writer and the epoch ID identify, also when the + * query restarts from its checkpoint and when other queries write to the same destination. A + * micro-batch streaming query opens a new writer for each micro-batch, and Spark closes the + * previous writer of the query when the query opens a new one. Spark is not told when a + * streaming query stops, so it closes the last writer of the query once the writer is no longer + * referenced, or when the query is restarted. + * + * Conventions + * ----------- + * - Errors. The functions that return an int return 0 on success, or any other value on failure. + * They then set the message that spark_ds_last_error returns, and Spark reports it as a + * NATIVE_DATA_SOURCE_ERROR. spark_ds_close cannot fail, and it may run after the query, from + * another thread: do the work that can fail a query, such as flushing data, in a commit. A C++ + * exception or a Rust panic must not cross the interface: catch it and return an error instead. + * - Handles. spark_ds_open returns an opaque, non-NULL handle in its last argument, typically a + * pointer: Spark reports a NULL handle as an error. Spark releases each handle exactly once, + * with spark_ds_close, 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 + * writer of a streaming query are released once they are no longer referenced (see Lifecycle + * above). A handle does not depend on another one: each OpenRequest carries everything that the + * object needs. + * - Messages. A request is valid only during the call. The library replies by calling the + * SparkDsSink that it gets once, with the encoded reply, on the thread of the call and during the + * call. Not calling it is an empty reply, which is valid for every request. Both sides ignore the + * fields that they do not know. + * - Memory. What Spark passes to a function is only valid during the call, except for SparkDsHost, + * which stays valid until the JVM exits. + * - Arrow data. Data is exchanged with the Arrow C stream interface + * (https://arrow.apache.org/docs/format/CStreamInterface.html), as structs allocated by Spark. + * For an output stream, the library initializes it and Spark takes ownership: Spark reads it + * from the thread of the task, and releases it before it closes the object. For an input stream, + * the library reads it during the call, on the thread of the call, and may move it, or the arrays + * that it returns, to take ownership; otherwise, Spark releases it after the call returns. The + * arrays that the library exports must have an offset of 0, at every level: Arrow Java ignores + * the offsets of the arrays it imports, so Spark rejects an array with another offset, such as a + * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with + * fewer values than its parent refers to, such as a column shorter than its batch. The buffers of + * an array must hold all of its values, as the C data interface requires: it does not carry the + * sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that is + * too short. + * - Schemas. The messages describe schemas with Spark SQL types, which map to Arrow types as in + * spark.sql.execution.arrow: for example, bigint is int64, string is utf8, timestamp is + * timestamp[us] with a time zone and timestamp_ntz is timestamp[us] without one. Data read from a + * library may also use large_utf8 and utf8_view for strings, large_binary and binary_view for + * binary data, and any names for the fields of the entries of a map. Spark rejects the other + * Arrow types, such as unsigned integers, decimal256, large lists, fixed-size lists and list + * views, and dictionary encoding. + * - 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 writer of a streaming + * query. The functions of SparkDsHost can be called from any thread. + * + * Compatibility + * ------------- + * The six functions of this header and the fields of SparkDsHost never change. SparkDsHost starts + * with its size, and only gets fields appended: a library must only read the fields that fit in + * the size of the struct that Spark passes. The messages only grow, as spark_datasource.proto + * describes, and Spark only sends a new kind of request to an object whose OpenReply declares that + * it supports it. So a library built against an earlier version of these files keeps working with + * later releases of Spark, and a library built against a later version works with earlier + * releases, without the features that they do not know. SPARK_DS_ABI_VERSION is the version of + * the functions of this header, and Spark rejects a library 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, which spark_ds_abi_version returns. */ +#define SPARK_DS_ABI_VERSION 1 + +/* Exports the functions of this header, 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 the reply of a function, encoded as a message of spark_datasource.proto, and copies it. + * `sink` is the argument passed with the sink to the function. + */ +typedef void (*SparkDsSink)(void* sink, const uint8_t* data, int64_t length); + +/* 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_open 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; + /* + * 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); + /* + * Returns whether the task with the given task_attempt_id of its TaskContext was killed, for + * example because its query was canceled, or has ended. 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. When a task is killed, Spark only stops reading once the get_next function of the stream + * returns: 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. + */ + bool (*is_interrupted)(int64_t task_attempt_id); +} SparkDsHost; + +/* Returns SPARK_DS_ABI_VERSION, the version of this interface that the library implements. */ +SPARK_DS_EXPORT int32_t spark_ds_abi_version(void); + +/* + * 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. + */ +SPARK_DS_EXPORT const char* spark_ds_last_error(void); + +/* + * Opens an object: a data source, a batch reader, a stream reader, a writer, a partition reader or + * a data writer. + * + * host: Spark, the same struct for every call in a JVM + * request, length: an OpenRequest + * reply: called with an OpenReply, or not at all for an empty one + * handle: set to the handle of the object + */ +SPARK_DS_EXPORT int spark_ds_open(const SparkDsHost* host, const uint8_t* request, int64_t length, + SparkDsSink reply, void* sink, void** handle); + +/* + * Runs a control operation on an object. + * + * request, length: a CallRequest + * reply: called with a CallReply, or not at all for an empty one + */ +SPARK_DS_EXPORT int spark_ds_call(void* handle, const uint8_t* request, int64_t length, + SparkDsSink reply, void* sink); + +/* + * Exchanges the data of an object, as Arrow C streams. + * + * request, length: a StreamRequest + * in: for a Write, the stream of the batches to write. It is NULL otherwise. + * out: for a Read, the stream to initialize with the batches of the partition: each array of the + * stream is a struct array whose fields are the columns to read, those of the schema to read + * or the pruned columns. It is NULL otherwise. + */ +SPARK_DS_EXPORT int spark_ds_stream(void* handle, const uint8_t* request, int64_t length, + struct ArrowArrayStream* in, struct ArrowArrayStream* out); + +/* Releases the handle of an object. */ +SPARK_DS_EXPORT void spark_ds_close(void* handle); + +#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 9d763edb079e7..d725922d95e83 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 2a03a58502c60..e0fb397d36cdc 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 f22008d3b2321..3183cc8c8c12a 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 3d0036f865225..eab75bc3e5bbd 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 1301985e30158..0369bcfacfbf8 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 52b9e803a64ed..77087a65a99d2 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 0000000000000..badff781627c9 --- /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 0000000000000..f739fdfe3e3d0 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala @@ -0,0 +1,135 @@ +/* + * 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 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.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 a reader for a batch scan of the given schema. */ + def reader(schema: StructType): ColumnarReader + + /** 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 +} + +/** + * Reads the data of a [[ColumnarDataSource]] for a batch scan. Spark pushes operations down by + * calling, each at most once and in this order, `pushPredicates`, `pushLimit` and `pruneColumns`, + * and plans the scan with `partitions`, on the driver. Then it serializes the reader to the + * executors, where `read` is called for each partition. + */ +trait ColumnarReader extends Serializable { + + /** + * 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 `read` returns 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 + + def partitions(): Array[InputPartition] + + /** Reads a partition on an executor. A batch only needs to be valid until the next one. */ + def read(partition: InputPartition): PartitionReader[ColumnarBatch] +} + +/** + * 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 0000000000000..1393713a78f11 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala @@ -0,0 +1,241 @@ +/* + * 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 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.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._ +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. + * + * @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 + + override def capabilities(): util.Set[TableCapability] = + util.EnumSet.of(BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE) + + 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 reader: ColumnarReader = _ + + 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 batchReader: ColumnarReader = { + if (reader == null) { + reader = dataSource.reader(schema) + } + reader + } + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + pushed = batchReader.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 = batchReader.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 && batchReader.pruneColumns(required)) { + readSchema = required + } + } + + override def build(): Scan = { + new ColumnarScan( + shortName, schema, readSchema, () => batchReader, () => dataSource, pushed, pushedLimit) + } +} + +class ColumnarScan( + shortName: String, + fullSchema: StructType, + prunedSchema: StructType, + batchReader: () => ColumnarReader, + dataSource: () => ColumnarDataSource, + 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(batchReader()) + + 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") + } +} + +class ColumnarBatchScan(reader: ColumnarReader) extends Batch { + private lazy val partitions: Array[InputPartition] = reader.partitions() + + override def planInputPartitions(): Array[InputPartition] = partitions + + override def createReaderFactory(): PartitionReaderFactory = { + // The reader is serialized to the executors, so it has to be planned first. + partitions + new ColumnarPartitionReaderFactory(reader) + } +} + +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 ColumnarPartitionReaderFactory(reader: ColumnarReader) + extends ColumnarOnlyPartitionReaderFactory { + override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = { + reader.read(partition) + } +} + +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 0000000000000..8721301852c73 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala @@ -0,0 +1,329 @@ +/* + * 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.memory.BufferAllocator +import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.arrow.vector.types.pojo.Schema + +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.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( + () => 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. + */ +class ColumnarStreamingWrite(createWriter: () => ColumnarStreamWriter, schema: StructType) + extends StreamingWrite { + + private var writer: ColumnarStreamWriter = _ + + override def createStreamingWriterFactory( + info: PhysicalWriteInfo): StreamingDataWriterFactory = { + 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 writeAll(records: java.util.Iterator[InternalRow]): Unit = delegate match { + case streaming: ArrowStreamDataWriter => + // The rows that `write` got before, if any, are written first. + flush() + streaming.writeStream(new ArrowRowReader(records.asScala, arrowSchema, allocator, options)) + case _ => + super.writeAll(records) + } + + 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() +} + +/** + * A data writer of a [[ColumnarDataSource]] that can also write all the data of a task from an + * Arrow reader, which it reads on the calling thread, during the call. + */ +trait ArrowStreamDataWriter extends DataWriter[ColumnarBatch] { + /** Writes all the batches of `reader`. */ + def writeStream(reader: ArrowReader): Unit +} + +/** + * Reads rows as Arrow batches, which end with the limits of [[ArrowBatchingOptions]]. Each batch + * is written to new buffers, so a native writer that took ownership of the arrays of a batch can + * keep them after the next batch. The first error of the rows is kept, see + * [[ArrowRowReader.failure]]. + */ +class ArrowRowReader( + rows: Iterator[InternalRow], + schema: Schema, + allocator: BufferAllocator, + options: ArrowBatchingOptions) + extends ArrowReader(allocator) { + + @volatile private var error: Throwable = _ + + override def loadNextBatch(): Boolean = { + try { + val root = getVectorSchemaRoot + // Releases the buffers of the previous batch, which its exported arrays may still hold. + root.clear() + if (!rows.hasNext) { + false + } else { + val writer = ArrowWriter.create(root) + var numRows = 0 + var full = false + while (!full && rows.hasNext) { + writer.write(rows.next()) + numRows += 1 + // The size is not computed when the limit is the largest possible one, as in Python + // UDFs. + full = numRows >= options.maxRecordsPerBatch || + (options.maxBytesPerBatch < Int.MaxValue && + writer.sizeInBytes() >= options.maxBytesPerBatch) + } + writer.finish() + true + } + } catch { + case e: Throwable => + if (error == null) { + error = e + } + throw e + } + } + + override def bytesRead(): Long = 0L + + override protected def closeReadSource(): Unit = {} + + override protected def readSchema(): Schema = schema +} + +object ArrowRowReader { + /** + * The error of the rows of `reader`, if it is an [[ArrowRowReader]] that failed. An Arrow C + * stream that exports the reader only reports the message of the error. + */ + def failure(reader: ArrowReader): Option[Throwable] = reader match { + case rows: ArrowRowReader => Option(rows.error) + case _ => None + } +} 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 0000000000000..1d48f409f181c --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala @@ -0,0 +1,1161 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.lang.ref.{Cleaner, WeakReference} +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import com.fasterxml.jackson.core.JsonProcessingException +import com.fasterxml.jackson.databind.DeserializationFeature +import com.fasterxml.jackson.databind.json.JsonMapper +import com.google.protobuf.ByteString +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, CDataDictionaryProvider, Data} +import org.apache.arrow.vector.{FieldVector, ValueVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.MapVector +import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} + +import org.apache.spark.{SPARK_VERSION, SparkException, SparkThrowable, SparkUnsupportedOperationException, TaskContext} +import org.apache.spark.internal.Logging +import org.apache.spark.sql.catalyst.util.CollationFactory +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.{ArrowRowReader, ArrowStreamDataWriter, ColumnarDataSource, ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} +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` and the messages of + * `org/apache/spark/sql/datasource/spark_datasource.proto`. + * + * @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._ + + // The data source, as the requests of spark_ds_open carry it. The keys of a case-insensitive + // map are in lower case. + private def dataSource: proto.DataSource = { + val builder = proto.DataSource.newBuilder().setName(name) + options.entrySet().asScala.foreach(e => builder.putOptions(e.getKey, e.getValue)) + builder.build() + } + + override def schema(): StructType = { + val library = NativeLibraries.get(location) + withDataSource(library, "infer the schema of") { (handle, _) => + val reply = call(name, "infer the schema of") { + library.call(handle, NativeMessages.InferSchema) + } + if (!reply.getInferSchema.hasSchema) { + throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name) + } + call(name, "infer the schema of")(NativeSchemas.fromProto(reply.getInferSchema.getSchema)) + } + } + + override def reader(schema: StructType): ColumnarReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = open(library, "plan a scan of", proto.Capability.CAPABILITY_BATCH_READ, + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { + _.setBatchReader(proto.BatchReader.newBuilder() + .setDataSource(dataSource) + .setSchema(NativeSchemas.toProto(schema))) + } + new NativeDataSourceReader( + executorLocation, name, schema, new NativeHandle(handle, library.close)) + } + + override def streamReader(schema: StructType): ColumnarStreamReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = open(library, "plan a scan of", proto.Capability.CAPABILITY_MICRO_BATCH_READ, + "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { + _.setStreamReader(proto.StreamReader.newBuilder() + .setDataSource(dataSource) + .setSchema(NativeSchemas.toProto(schema))) + } + new NativeDataSourceStreamReader( + executorLocation, name, schema, new NativeHandle(handle, library.close)) + } + + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = + new NativeHandle(openWriter(library, schema, overwrite, queryId = null), library.close) + 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(openWriter(library, schema, overwrite, queryId), library.close) + val writer = new NativeDataSourceStreamWriter( + executorLocation, name, handle, serializeWriter(library, handle)) + NativeStreamWriters.register(queryId, handle) + writer + } + + /** + * Runs `f` with a new native data source and its capabilities, and closes the data source + * afterwards. + */ + private def withDataSource[T](library: NativeLibrary, action: String)( + f: (Long, Set[proto.Capability]) => T): T = { + val opened = call(name, "create") { + library.open(NativeMessages.open().setDataSource(dataSource).build()) + } + Utils.tryWithSafeFinally { + f(opened.handle, opened.reply.getCapabilitiesList.asScala.toSet) + } { + call(name, action)(library.close(opened.handle)) + } + } + + /** Returns the capabilities of the data source, and fails if it does not declare `capability`. */ + private def capabilities( + library: NativeLibrary, + action: String, + capability: proto.Capability, + unsupportedError: String): Set[proto.Capability] = { + val capabilities = withDataSource(library, action)((_, capabilities) => capabilities) + if (!capabilities.contains(capability)) { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) + } + capabilities + } + + /** Opens a reader with the request that `f` builds, if the data source declares `capability`. */ + private def open( + library: NativeLibrary, + action: String, + capability: proto.Capability, + unsupportedError: String)( + f: proto.OpenRequest.Builder => proto.OpenRequest.Builder): Long = { + capabilities(library, action, capability, unsupportedError) + call(name, action)(library.open(f(NativeMessages.open()).build()).handle) + } + + /** Opens a writer for a batch write, or for a streaming write with the ID of its query. */ + private def openWriter( + library: NativeLibrary, + schema: StructType, + overwrite: Boolean, + queryId: String): Long = { + val action = "plan a write to" + val (capability, unsupportedError, overwriteMode) = if (queryId == null) { + (proto.Capability.CAPABILITY_BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", + "the save mode overwrite") + } else { + (proto.Capability.CAPABILITY_STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", + "the output mode complete") + } + val declared = capabilities(library, action, capability, unsupportedError) + if (overwrite && !declared.contains(proto.Capability.CAPABILITY_TRUNCATE)) { + throw error(name, action, s"The data source does not support $overwriteMode.") + } + val writer = proto.Writer.newBuilder() + .setDataSource(dataSource) + .setSchema(NativeSchemas.toProto(schema)) + .setMode( + if (overwrite) proto.WriteMode.WRITE_MODE_TRUNCATE else proto.WriteMode.WRITE_MODE_APPEND) + if (queryId != null) { + writer.setQueryId(queryId) + } + call(name, action)(library.open(NativeMessages.open().setWriter(writer).build()).handle) + } + + private def serializeWriter(library: NativeLibrary, handle: NativeHandle): Array[Byte] = { + try { + call(name, "plan a write to") { + library.call(handle.get, NativeMessages.Serialize).getSerialize.getState.toByteArray + } + } catch { + case e: Throwable => + handle.close() + throw e + } + } +} + +/** Reads a native data source for a batch scan. */ +class NativeDataSourceReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarReader { + import NativeCalls._ + + // Set on the driver when the partitions are planned, and serialized to the executors. + private var readerState: Array[Byte] = Array.emptyByteArray + private var readSchema: StructType = schema + @transient private var plannedPartitions: Array[InputPartition] = _ + // Why planning the partitions failed. The handle is released after the first try, but the + // scan can be planned again, for example by the lazy values of the physical plan that call it. + @transient private var planningError: Throwable = _ + + @transient private lazy val library = NativeLibraries.get(location) + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + // Only the predicates that can be expressed in JSON are passed to the library. + val encoded = predicates.zipWithIndex.flatMap { case (predicate, i) => + NativePredicates.toJson(predicate, schema).map(_ -> i) + } + val used = if (encoded.isEmpty) { + Set.empty[Int] + } else { + val request = proto.CallRequest.newBuilder().setPushPredicates( + proto.PushPredicates.newBuilder().addAllPredicatesJson(encoded.map(_._1).toSeq.asJava)) + val uses = call(name, "plan a scan of") { + library.call(handle.get, request.build()).getPushPredicates.getUsesValueList.asScala + } + if (uses.length > encoded.length) { + throw error(name, "plan a scan of", s"PushPredicates returned ${uses.length} answers " + + s"for ${encoded.length} predicates.") + } + // 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. The + // missing answers are PREDICATE_USE_UNUSED, and Spark takes the values that it does not + // know as PREDICATE_USE_PARTIAL. + encoded.map(_._2).zip(uses).collect { + case (i, use) if use != proto.PredicateUse.PREDICATE_USE_UNUSED_VALUE => i + }.toSet + } + predicates.zipWithIndex.collect { case (predicate, i) if used.contains(i) => predicate } + } + + override def pushLimit(limit: Int): Boolean = { + val request = proto.CallRequest.newBuilder() + .setPushLimit(proto.PushLimit.newBuilder().setLimit(limit)) + call(name, "plan a scan of") { + library.call(handle.get, request.build()).getPushLimit.getApplied + } + } + + override def pruneColumns(requiredSchema: StructType): Boolean = { + val request = proto.CallRequest.newBuilder().setPruneColumns( + proto.PruneColumns.newBuilder().addAllColumns(requiredSchema.fieldNames.toSeq.asJava)) + val pruned = call(name, "plan a scan of") { + library.call(handle.get, request.build()).getPruneColumns.getPruned + } + if (pruned) { + readSchema = requiredSchema + } + pruned + } + + override def partitions(): Array[InputPartition] = synchronized { + if (plannedPartitions == null) { + if (planningError != null) { + throw planningError + } + try { + val partitions = call(name, "plan a scan of") { + library.call(handle.get, NativeMessages.PlanPartitions).getPlanPartitions + .getPartitionsList.asScala.map(_.toByteArray).toArray + } + readerState = call(name, "plan a scan of") { + library.call(handle.get, NativeMessages.Serialize).getSerialize.getState.toByteArray + } + plannedPartitions = NativeInputPartition.from(partitions) + } catch { + case e: Throwable => + planningError = e + throw e + } finally { + // The executors only need the serialized state. + handle.close() + } + } + plannedPartitions + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, readSchema) + } +} + +/** Reads a native data source for a micro-batch streaming scan. */ +class NativeDataSourceStreamReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarStreamReader { + import NativeCalls._ + + // Set on the driver whenever the partitions of a micro-batch are planned, and serialized to + // the executors with the partitions of that micro-batch. + private var readerState: Array[Byte] = Array.emptyByteArray + + @transient private lazy val library = NativeLibraries.get(location) + + override def initialOffset(): String = { + checkOffset("InitialOffset", call(name, "plan a scan of") { + library.call(handle.get, NativeMessages.InitialOffset).getInitialOffset.getOffsetJson + }) + } + + override def latestOffset(): String = { + checkOffset("LatestOffset", call(name, "plan a scan of") { + library.call(handle.get, NativeMessages.LatestOffset).getLatestOffset.getOffsetJson + }) + } + + /** + * Checks an offset returned by the given request: it must be JSON on a single line, because the + * offset log of the checkpoint stores each offset as a line. + */ + private def checkOffset(request: String, offset: String): String = { + if (offset.isEmpty) { + throw error(name, "plan a scan of", s"$request returned no offset.") + } + if (offset.exists(c => c == '\n' || c == '\r')) { + throw error(name, "plan a scan of", s"$request 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"$request returned an offset that is not valid " + + s"JSON: $offset") + } + offset + } + + override def partitions(start: String, end: String): Array[InputPartition] = { + val request = proto.CallRequest.newBuilder().setPlanStreamPartitions( + proto.PlanStreamPartitions.newBuilder().setStartJson(start).setEndJson(end)) + val partitions = call(name, "plan a scan of") { + library.call(handle.get, request.build()).getPlanStreamPartitions.getPartitionsList.asScala + .map(_.toByteArray).toArray + } + readerState = call(name, "plan a scan of") { + library.call(handle.get, NativeMessages.Serialize).getSerialize.getState.toByteArray + } + 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 = { + val request = proto.CallRequest.newBuilder() + .setCommitOffset(proto.CommitOffset.newBuilder().setEndJson(end)) + call(name, "commit an offset of")(library.call(handle.get, request.build())) + } + + 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.call(handle.get, NativeMessages.commit(-1L, 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.call(handle.get, NativeMessages.abort(-1L, 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.call(handle.get, NativeMessages.commit(epochId, messages)) + } + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "abort a write to") { + library.call(handle.get, NativeMessages.abort(epochId, 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 { + /** The messages of the tasks, as the requests of a commit or an abort carry them. */ + def toTaskMessages(messages: Array[WriterCommitMessage]): Seq[proto.TaskMessage] = { + messages.toSeq.map { + case message: NativeWriterCommitMessage => + proto.TaskMessage.newBuilder() + .setCommitted(true) + .setMessage(ByteString.copyFrom(message.bytes)) + .build() + case null => proto.TaskMessage.getDefaultInstance + case other => + throw SparkException.internalError(s"Unexpected commit message: ${other.getClass.getName}") + } + } +} + +/** + * Reads a partition of a native data source: each array of the Arrow stream exported by the + * library is checked, and imported into the Arrow vectors that back the returned batch. + */ +class NativePartitionReader( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + partition: Array[Byte], + schema: StructType) + extends PartitionReader[ColumnarBatch] { + import NativeCalls._ + + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) + // Read here rather than with `Data.importArrayStream`, to check each array before importing it. + private val stream = ArrowArrayStream.allocateNew(allocator) + private val dictionaries = new CDataDictionaryProvider() + private var library: NativeLibrary = _ + // The partition reader of the library, which is closed after its stream is released, and the + // registration of the task for SparkDsHost.is_interrupted, until then. + private var partitionReader: Long = 0L + private var task: AutoCloseable = _ + 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 = { + library = NativeLibraries.get(location) + val (taskContext, registration) = NativeTask.create() + task = registration + val request = proto.PartitionReader.newBuilder() + .setTask(taskContext) + .setReaderState(ByteString.copyFrom(readerState)) + .setPartition(ByteString.copyFrom(partition)) + partitionReader = call(name, "read") { + library.open(NativeMessages.open().setPartitionReader(request).build()).handle + } + call(name, "read") { + library.stream(partitionReader, NativeMessages.Read, 0L, stream.memoryAddress()) + } + } + + private def importSchema(): Schema = { + // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash + // the JVM. + val snapshot = stream.snapshot() + if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 || + snapshot.get_last_error == 0) { + throw error(name, "read", "The library did not initialize the stream, or a callback of " + + "the stream is NULL.") + } + val arrowSchema = ArrowSchema.allocateNew(allocator) + try { + call(name, "read")(stream.getSchema(arrowSchema)) + val imported = call(name, "read")(Data.importSchema(allocator, arrowSchema, dictionaries)) + new Schema( + imported.getFields.asScala.map(withMapEntryNames).asJava, imported.getCustomMetadata) + } finally { + arrowSchema.close() + } + } + + /** + * Names the fields of the entries of the maps `key` and `value`, which ArrowColumnVector looks + * up by name. Arrow does not fix these names: arrow-rs names them `keys` and `values`, for + * example. The arrays are imported by position, so the names do not change how they are read. + */ + private def withMapEntryNames(field: Field): Field = { + val children = field.getChildren.asScala.map(withMapEntryNames).toSeq + val renamed = (field.getType, children) match { + case (_: ArrowType.Map, Seq(entries)) if entries.getChildren.size == 2 => + val Seq(key, value) = entries.getChildren.asScala.toSeq + Seq(new Field(entries.getName, entries.getFieldType, Seq( + new Field(MapVector.KEY_NAME, key.getFieldType, key.getChildren), + new Field(MapVector.VALUE_NAME, value.getFieldType, value.getChildren)).asJava)) + case _ => children + } + new Field(field.getName, field.getFieldType, renamed.asJava) + } + + private def checkSchema(): Unit = { + def isDictionaryEncoded(field: Field): Boolean = { + field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) + } + root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + + "is not supported.") + } + // The Arrow types that convert to Spark types, but that Spark does not read: ArrowColumnVector + // cannot read decimals of other widths than 128 bits, and the lengths of the elements of a + // list view cannot be checked without a pass over all its lists (see checkArrays). + def unsupportedType(field: Field): Option[ArrowType] = field.getType match { + case decimal: ArrowType.Decimal if decimal.getBitWidth != 128 => Some(decimal) + case listView: ArrowType.ListView => Some(listView) + case _ => field.getChildren.asScala.view.flatMap(unsupportedType).headOption + } + root.getSchema.getFields.asScala.foreach { field => + unsupportedType(field).foreach { arrowType => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"$arrowType, which is not supported.") + } + } + val actual = StructType(root.getSchema.getFields.asScala.map { field => + val dataType = try { + ArrowUtils.fromArrowField(field) + } catch { + case e: SparkUnsupportedOperationException if e.getCondition == "UNSUPPORTED_ARROWTYPE" => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"${e.getMessageParameters.get("typeName")}, which is not supported.") + } + StructField(field.getName, dataType, field.isNullable) + }.toSeq) + // The schema to read as the library got it: Arrow does not carry everything, such as the + // collations of strings, the fields of intervals and user-defined types, which still apply to + // the data that Spark reads. + val expected = ArrowUtils.fromArrowSchema(ArrowUtils.toArrowSchema( + schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = false)) + val matches = actual.length == expected.length && actual.zip(expected).forall { + case (actualField, field) => + actualField.name == field.name && + DataType.equalsIgnoreNullability(actualField.dataType, field.dataType) + } + if (!matches) { + throw error(name, "read", s"The data has the schema ${actual.toDDL}, but the expected " + + s"schema is ${expected.toDDL}.") + } + } + + private def columnVector(vector: FieldVector): ColumnVector = { + try { + new ReusedArrowColumnVector(vector) + } catch { + case NonFatal(e) => + throw QueryExecutionErrors.nativeDataSourceError("read", name, + s"Spark cannot read the column ${vector.getName} of the Arrow type " + + s"${vector.getField.getType}: ${message(e)}", e) + } + } + + override def next(): Boolean = { + val array = ArrowArray.allocateNew(allocator) + try { + call(name, "read")(stream.getNext(array)) + val snapshot = array.snapshot() + if (snapshot.release == 0) { + // The end of the stream. + false + } else { + try { + checkArrays(snapshot) + } catch { + case e: Throwable => + array.release() + throw e + } + call(name, "read")(Data.importIntoVectorSchemaRoot(allocator, array, root, dictionaries)) + batch.setNumRows(root.getRowCount) + true + } + } finally { + array.close() + } + } + + /** + * Checks the array of a batch, and the arrays of its columns at every level, before importing + * them. Arrow Java trusts the arrays, and would read them wrongly otherwise: + * - An array must have an offset of 0. Arrow Java ignores the offsets of the arrays it imports, + * so it would read a sliced array from the start of its buffers. + * - An array must have at least as many values as its parent refers to: a column as many as + * the batch has rows, a field as many as its struct has values, and the elements of a list or + * a map as many as its last offset says. Arrow Java would extend a shorter array with nulls. + * Only the lengths that the library exported show this: the import sets the length of each + * array to what its parent refers to. + * The buffers are not checked, since the C data interface does not carry their sizes: Arrow + * Java takes them from the lengths, so the library must export buffers that hold every value. + */ + private def checkArrays(batch: ArrowArray.Snapshot): Unit = { + def children(array: ArrowArray.Snapshot): Seq[ArrowArray.Snapshot] = { + (0 until array.n_children.toInt).map { i => + ArrowArray.wrap(Platform.getLong(null, array.children + 8L * i)).snapshot() + } + } + // The last offset of a list or a map, in its second buffer. Arrow Java rejects the array if + // the buffer is missing. + def lastOffset(array: ArrowArray.Snapshot): Long = { + val offsets = if (array.length > 0 && array.n_buffers >= 2) { + Platform.getLong(null, array.buffers + 8L) + } else { + 0L + } + if (offsets == 0) 0 else Platform.getInt(null, offsets + 4L * array.length) + } + // The number of values that an array refers to in each of its children. + def referenced(array: ArrowArray.Snapshot, arrowType: ArrowType): Long = arrowType match { + case _: ArrowType.Struct => array.length + case _: ArrowType.List | _: ArrowType.Map => lastOffset(array) + // The other types that Spark reads have no children. List views are rejected before. + case _ => 0 + } + def checkColumn(array: ArrowArray.Snapshot, field: Field, column: String): Unit = { + if (array.offset != 0) { + throw error(name, "read", s"The column $column has an offset of ${array.offset}, but " + + "Spark only reads arrays with an offset of 0. Copy sliced arrays before exporting them.") + } + val length = referenced(array, field.getType) + children(array).zip(field.getChildren.asScala).foreach { case (child, childField) => + if (child.length < length) { + throw error(name, "read", s"The column $column contains an array with a length of " + + s"${child.length}, but its parent refers to $length values.") + } + checkColumn(child, childField, column) + } + } + if (batch.offset != 0) { + throw error(name, "read", s"The batch has an offset of ${batch.offset}, but Spark only " + + "reads arrays with an offset of 0. Copy sliced arrays before exporting them.") + } + children(batch).zip(root.getSchema.getFields.asScala).foreach { case (column, field) => + if (column.length < batch.length) { + throw error(name, "read", s"The column ${field.getName} has a length of " + + s"${column.length}, but the batch has a length of ${batch.length}.") + } + checkColumn(column, field, field.getName) + } + } + + override def get(): ColumnarBatch = batch + + override def close(): Unit = { + Utils.tryWithSafeFinally { + if (root != null) { + root.close() + } + dictionaries.close() + // Releases the stream of the library, unless it was not exported. + stream.release() + stream.close() + } { + Utils.tryWithSafeFinally { + if (partitionReader != 0L) { + call(name, "read")(library.close(partitionReader)) + } + } { + 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 data 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 ArrowStreamDataWriter { + import NativeCalls._ + + private val library = NativeLibraries.get(location) + // The task, registered for SparkDsHost.is_interrupted until the data writer is closed. + private val (taskContext, task) = NativeTask.create(Some(partitionId), Some(taskId)) + private val handle: Long = try { + val request = proto.DataWriter.newBuilder() + .setTask(taskContext) + .setWriterState(ByteString.copyFrom(writerState)) + .setEpochId(epochId) + call(name, "write to") { + library.open(NativeMessages.open().setDataWriter(request).build()).handle + } + } 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 reply. + private var finished = false + private var closed = false + + private def checkOpen(): Unit = { + if (finished || closed) { + throw SparkException.internalError("The data writer was already committed or aborted.") + } + } + + override def write(batch: ColumnarBatch): Unit = { + checkOpen() + 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}.") + } + } + writeStream(new SingleBatchArrowReader( + new VectorSchemaRoot(vectors.map(_.getField).asJava, vectors.asJava, batch.numRows()))) + } + + override def writeStream(reader: ArrowReader): Unit = { + checkOpen() + val stream = ArrowArrayStream.allocateNew(ArrowUtils.rootAllocator) + try { + Data.exportArrayStream(ArrowUtils.rootAllocator, reader, stream) + try { + call(name, "write to") { + library.stream(handle, NativeMessages.Write, stream.memoryAddress(), 0L) + } + } catch { + // The library reads the batches through the stream, which only reports the message of + // an error of the rows, such as a failure of an upstream operator: that error is thrown. + case e: Throwable => throw ArrowRowReader.failure(reader).getOrElse(e) + } + } finally { + // Releases the stream, and so the reader, unless the library moved it to take ownership. + stream.release() + stream.close() + } + } + + override def commit(): WriterCommitMessage = { + checkOpen() + val reply = call(name, "write to") { + library.call(handle, NativeMessages.CommitDataWriter, () => finished = true) + } + NativeWriterCommitMessage(reply.getCommitDataWriter.getMessage.toByteArray) + } + + override def abort(): Unit = { + if (!finished && !closed) { + finished = true + call(name, "write to")(library.call(handle, NativeMessages.AbortDataWriter)) + } + } + + override def close(): Unit = { + if (!closed) { + closed = true + Utils.tryWithSafeFinally { + call(name, "write to")(library.close(handle)) + }(task.close()) + } + } +} + +/** + * Exports one batch of a [[VectorSchemaRoot]] as an Arrow reader. The vectors of the batch are + * transferred to the reader, which releases them when it is closed. + */ +private class SingleBatchArrowReader(batch: VectorSchemaRoot) + extends ArrowReader(ArrowUtils.rootAllocator) { + private var loaded = false + + override def loadNextBatch(): Boolean = { + if (loaded) { + false + } else { + loaded = true + val root = getVectorSchemaRoot + batch.getFieldVectors.asScala.zip(root.getFieldVectors.asScala).foreach { + case (from, to) => from.makeTransferPair(to).transfer() + } + root.setRowCount(batch.getRowCount) + true + } + } + + override def bytesRead(): Long = 0L + + override protected def closeReadSource(): Unit = {} + + override protected def readSchema(): Schema = batch.getSchema +} + +/** Creates the `TaskContext` message of the running task. */ +private[ffi] object NativeTask { + /** + * Returns the message, and the registration of the task for `SparkDsHost.is_interrupted`, which + * the caller closes once the library no longer uses the object of the task. + * + * @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): (proto.TaskContext, AutoCloseable) = { + val context = Option(TaskContext.get()) + val attemptId = taskAttemptId.orElse(context.map(_.taskAttemptId())).getOrElse(-1L) + val message = proto.TaskContext.newBuilder() + .setStageId(context.map(_.stageId()).getOrElse(-1)) + .setStageAttemptNumber(context.map(_.stageAttemptNumber()).getOrElse(-1)) + .setPartitionId(partitionId.orElse(context.map(_.partitionId())).getOrElse(-1)) + .setAttemptNumber(context.map(_.attemptNumber()).getOrElse(0)) + .setTaskAttemptId(attemptId) + .build() + (message, NativeLibrary.registerTask(attemptId, () => 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 @ (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 + } + } + + /** 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) + } +} + +/** The messages of `spark_datasource.proto` that Spark sends. */ +private[ffi] object NativeMessages { + /** An `OpenRequest`, with the versions of Spark and of the messages. */ + def open(): proto.OpenRequest.Builder = proto.OpenRequest.newBuilder() + .setSparkVersion(SPARK_VERSION) + .setProtocolVersion(NativeLibrary.PROTOCOL_VERSION) + + val InferSchema: proto.CallRequest = proto.CallRequest.newBuilder() + .setInferSchema(proto.InferSchema.getDefaultInstance).build() + val PlanPartitions: proto.CallRequest = proto.CallRequest.newBuilder() + .setPlanPartitions(proto.PlanPartitions.getDefaultInstance).build() + val Serialize: proto.CallRequest = proto.CallRequest.newBuilder() + .setSerialize(proto.Serialize.getDefaultInstance).build() + val InitialOffset: proto.CallRequest = proto.CallRequest.newBuilder() + .setInitialOffset(proto.InitialOffset.getDefaultInstance).build() + val LatestOffset: proto.CallRequest = proto.CallRequest.newBuilder() + .setLatestOffset(proto.LatestOffset.getDefaultInstance).build() + val CommitDataWriter: proto.CallRequest = proto.CallRequest.newBuilder() + .setCommitDataWriter(proto.CommitDataWriter.getDefaultInstance).build() + val AbortDataWriter: proto.CallRequest = proto.CallRequest.newBuilder() + .setAbortDataWriter(proto.AbortDataWriter.getDefaultInstance).build() + val Read: proto.StreamRequest = + proto.StreamRequest.newBuilder().setRead(proto.Read.getDefaultInstance).build() + val Write: proto.StreamRequest = + proto.StreamRequest.newBuilder().setWrite(proto.Write.getDefaultInstance).build() + + def commit(epochId: Long, messages: Array[WriterCommitMessage]): proto.CallRequest = { + proto.CallRequest.newBuilder().setCommit(proto.Commit.newBuilder() + .setEpochId(epochId) + .addAllMessages(NativeWriterCommitMessage.toTaskMessages(messages).asJava)).build() + } + + def abort(epochId: Long, messages: Array[WriterCommitMessage]): proto.CallRequest = { + proto.CallRequest.newBuilder().setAbort(proto.Abort.newBuilder() + .setEpochId(epochId) + .addAllMessages(NativeWriterCommitMessage.toTaskMessages(messages).asJava)).build() + } +} + +/** Converts schemas to and from the Spark SQL types of `spark_datasource.proto`. */ +private[ffi] object NativeSchemas { + import proto.{DataType => P} + + def toProto(schema: StructType): proto.StructType = { + val builder = proto.StructType.newBuilder() + schema.fields.foreach { field => + builder.addFields(proto.StructField.newBuilder() + .setName(field.name) + .setDataType(toProto(field.dataType)) + .setNullable(field.nullable)) + } + builder.build() + } + + def toProto(dataType: DataType): proto.DataType = { + val builder = proto.DataType.newBuilder() + dataType match { + case NullType => builder.setNull(P.Null.getDefaultInstance) + case BooleanType => builder.setBoolean(P.Boolean.getDefaultInstance) + case ByteType => builder.setByte(P.Byte.getDefaultInstance) + case ShortType => builder.setShort(P.Short.getDefaultInstance) + case IntegerType => builder.setInteger(P.Integer.getDefaultInstance) + case LongType => builder.setLong(P.Long.getDefaultInstance) + case FloatType => builder.setFloat(P.Float.getDefaultInstance) + case DoubleType => builder.setDouble(P.Double.getDefaultInstance) + case d: DecimalType => + builder.setDecimal(P.Decimal.newBuilder().setPrecision(d.precision).setScale(d.scale)) + case s: StringType => + val collation = if (s.isUTF8BinaryCollation) { + "" + } else { + CollationFactory.fetchCollation(s.collationId).collationName + } + builder.setString(P.String.newBuilder().setCollation(collation)) + case BinaryType => builder.setBinary(P.Binary.getDefaultInstance) + case DateType => builder.setDate(P.Date.getDefaultInstance) + case TimestampType => builder.setTimestamp(P.Timestamp.getDefaultInstance) + case TimestampNTZType => builder.setTimestampNtz(P.TimestampNtz.getDefaultInstance) + case y: YearMonthIntervalType => + builder.setYearMonthInterval( + P.YearMonthInterval.newBuilder().setStartField(y.startField).setEndField(y.endField)) + case d: DayTimeIntervalType => + builder.setDayTimeInterval( + P.DayTimeInterval.newBuilder().setStartField(d.startField).setEndField(d.endField)) + case CalendarIntervalType => + builder.setCalendarInterval(P.CalendarInterval.getDefaultInstance) + case a: ArrayType => + builder.setArray(P.Array.newBuilder() + .setElementType(toProto(a.elementType)) + .setContainsNull(a.containsNull)) + case m: MapType => + builder.setMap(P.Map.newBuilder() + .setKeyType(toProto(m.keyType)) + .setValueType(toProto(m.valueType)) + .setValueContainsNull(m.valueContainsNull)) + case s: StructType => builder.setStruct(toProto(s)) + case udt: UserDefinedType[_] => return toProto(udt.sqlType) + case other => + throw new IllegalArgumentException( + s"Native data sources do not support the type ${other.catalogString}.") + } + builder.build() + } + + def fromProto(schema: proto.StructType): StructType = { + StructType(schema.getFieldsList.asScala.map { field => + StructField(field.getName, fromProto(field.getDataType), field.getNullable) + }.toSeq) + } + + def fromProto(dataType: proto.DataType): DataType = dataType.getKindCase match { + case P.KindCase.NULL => NullType + case P.KindCase.BOOLEAN => BooleanType + case P.KindCase.BYTE => ByteType + case P.KindCase.SHORT => ShortType + case P.KindCase.INTEGER => IntegerType + case P.KindCase.LONG => LongType + case P.KindCase.FLOAT => FloatType + case P.KindCase.DOUBLE => DoubleType + case P.KindCase.DECIMAL => + DecimalType(dataType.getDecimal.getPrecision, dataType.getDecimal.getScale) + case P.KindCase.STRING => + val collation = dataType.getString.getCollation + if (collation.isEmpty) StringType else StringType(collation) + case P.KindCase.BINARY => BinaryType + case P.KindCase.DATE => DateType + case P.KindCase.TIMESTAMP => TimestampType + case P.KindCase.TIMESTAMP_NTZ => TimestampNTZType + case P.KindCase.YEAR_MONTH_INTERVAL => + val interval = dataType.getYearMonthInterval + YearMonthIntervalType(interval.getStartField.toByte, interval.getEndField.toByte) + case P.KindCase.DAY_TIME_INTERVAL => + val interval = dataType.getDayTimeInterval + DayTimeIntervalType(interval.getStartField.toByte, interval.getEndField.toByte) + case P.KindCase.CALENDAR_INTERVAL => CalendarIntervalType + case P.KindCase.ARRAY => + ArrayType(fromProto(dataType.getArray.getElementType), dataType.getArray.getContainsNull) + case P.KindCase.MAP => + val map = dataType.getMap + MapType(fromProto(map.getKeyType), fromProto(map.getValueType), map.getValueContainsNull) + case P.KindCase.STRUCT => fromProto(dataType.getStruct) + case P.KindCase.KIND_NOT_SET => + throw new IllegalStateException("The library returned a data type without a kind.") + } +} 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 0000000000000..817ce79ed26df --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala @@ -0,0 +1,457 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.io.{File, FileInputStream, IOException} +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, StandardCopyOption} +import java.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.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 version of the native data source interface that the libraries implement + * @param dataSources the names of the data sources that the libraries implement + * @param libraries the path of the library in the package for each platform, as returned by + * [[NativePlatform.current]] + */ +case class NativeDataSourceManifest( + abiVersion: Int, + dataSources: Seq[String], + libraries: Map[String, String]) + +/** + * Where the library of a native data source is: in a [[NativeDataSourcePackage]], or installed in + * the native library path, as an [[InstalledNativeLibrary]]. It is sent to the executors, which + * load the library with [[NativeLibraries.get]]. + */ +sealed trait NativeLibraryLocation extends Serializable + +/** + * A native data source package: a zip file that contains a [[NativeDataSourceManifest]] and the + * libraries it lists. The SHA-256 checksum of the file identifies the package, so that an + * executor can check that it loads the same package as the driver. + * + * @param path the path of the package file on the driver + * @param artifactUUID the UUID of the session whose artifacts contain the package, if the session + * is isolated. Executors store the artifacts of such a session separately. + * @param distributedFileName the name of the copy of the package in the artifacts of the session, + * if the package is not itself an artifact of the session + */ +case class NativeDataSourcePackage( + path: String, + sha256: String, + manifest: NativeDataSourceManifest, + artifactUUID: Option[String] = None, + distributedFileName: Option[String] = None) extends NativeLibraryLocation { + def fileName: String = new File(path).getName +} + +/** + * A package file that is not a valid native data source package. + * + * @param error the error that explains why the package is invalid + * @param dataSources the data sources listed in the manifest of the package, if it can be read + */ +case class InvalidNativeDataSourcePackage(error: Throwable, dataSources: Seq[String]) + +/** + * 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) { + return Left("The manifest must have an integer field 'abiVersion'.") + } + val dataSources = root.get("dataSources") + if (dataSources == null || !dataSources.isArray || dataSources.isEmpty || + !dataSources.elements().asScala.forall(n => n.isTextual && n.asText().nonEmpty)) { + return Left("The manifest must have a field 'dataSources' with a non-empty array of " + + "data source names.") + } + val libraries = root.get("libraries") + if (libraries == null || !libraries.isObject || libraries.isEmpty || + !libraries.properties().asScala.forall(e => e.getValue.isTextual)) { + return Left("The manifest must have a field 'libraries' with an object that maps " + + "platforms to the paths of the libraries in the package.") + } + Right(NativeDataSourceManifest( + abiVersion.asInt(), + dataSources.elements().asScala.map(_.asText()).toSeq, + libraries.properties().asScala.map(e => e.getKey -> e.getValue.asText()).toMap)) + } + + private[ffi] def sha256(file: File): String = { + val digest = MessageDigest.getInstance("SHA-256") + Utils.tryWithResource(new FileInputStream(file)) { in => + val buffer = new Array[Byte](64 * 1024) + var read = in.read(buffer) + while (read != -1) { + digest.update(buffer, 0, read) + read = in.read(buffer) + } + } + HexFormat.of().formatHex(digest.digest()) + } +} + +/** + * A native data source library installed in the native library path, where + * [[NativeDataSourceRegistry]] finds it by its file name. It must be installed on every node. The + * SHA-256 checksum of the file identifies the build of the library, so that an executor can check + * that it loads the same build as the driver: the state that the driver serializes is only + * meant for the same build. + * + * @param path the path of the library on the driver + * @param sha256 the SHA-256 checksum of the library on the driver + */ +case class InstalledNativeLibrary(path: String, sha256: String) extends NativeLibraryLocation + +object InstalledNativeLibrary { + // 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`. */ + lazy val current: String = name(System.getProperty("os.name"), System.getProperty("os.arch")) + + private[ffi] def name(osName: String, osArch: String): String = { + val os = osName.toLowerCase(Locale.ROOT) match { + case n if n.startsWith("linux") => "linux" + case n if n.startsWith("mac") || n.startsWith("darwin") => "osx" + case n if n.startsWith("windows") => "windows" + case n => n.replaceAll("[^a-z0-9]", "") + } + val arch = osArch.toLowerCase(Locale.ROOT) match { + case "amd64" | "x86_64" | "x64" => "x86_64" + case "aarch64" | "arm64" => "aarch64" + case a => a + } + s"$os-$arch" + } +} + +/** + * The native data source libraries loaded in this JVM. A library is loaded once, on first use, + * after it is extracted from its package if needed, and stays loaded until the JVM exits. + */ +object NativeLibraries extends Logging { + // The libraries of the packages, by the checksum of the package, and the installed libraries, + // by their canonical path, 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()) + } + getSupported(location) + } + + 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 gets the same error when it is used again. + * + * @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) + } catch { + // A library that was loaded stays loaded, and gets the same error when it is used 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 0000000000000..71928a6c07647 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala @@ -0,0 +1,279 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.io.File +import java.nio.file.Paths +import java.util.Locale + +import org.apache.spark.{SparkEnv, SparkFiles} +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.PATH +import org.apache.spark.sql.Artifact +import org.apache.spark.sql.classic.SparkSession +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.util.Utils + +/** + * Finds the native data sources on the driver. Like the Python data sources installed in the + * Python path, the native data source libraries installed in the native library path are found + * automatically, by their file name: the library `spark_datasource_`, 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. + */ +object NativeDataSourceRegistry extends Logging { + + /** The prefix of the names of the installed native data source libraries. */ + val LIBRARY_NAME_PREFIX = "spark_datasource_" + + // The result of the last lookup by `DataSource.lookupDataSource` on this thread, if it found a + // native data source, by the name it looked up. The provider that the caller creates next for + // that name takes it, so that it does not look the name up again: see [[takeLookup]]. Each + // lookup clears it first, so a provider never takes the result of an earlier lookup. + private val lastLookup = new ThreadLocal[(String, (String, NativeLibraryLocation))] + + /** + * Finds the package or the installed library that provides the native data source with the + * given name, which is case-insensitive. Returns the name of the data source, as listed in the + * manifest of the package or in lower case for an installed library, and where its library is. + */ + def lookup(name: String, conf: SQLConf): Option[(String, NativeLibraryLocation)] = { + if (!conf.getConf(StaticSQLConf.NATIVE_DATA_SOURCE_ENABLED)) { + return None + } + val (invalid, valid) = (configuredPackageFiles(conf) ++ artifactPackageFiles()) + .map(NativeDataSourcePackage.tryRead).partitionMap(identity) + // A package of a version of the interface that Spark does not support is unusable too. + val (unsupported, usable) = valid.partitionMap { pkg => + if (pkg.manifest.abiVersion == NativeLibrary.ABI_VERSION) Right(pkg) else Left(pkg) + } + lookupPackage(name, usable) + .orElse(lookupInstalledLibrary(name)) + .orElse { + // An unusable package is only reported when no other package or installed library + // provides the data source: a stale copy does not hide a valid one. An invalid package + // whose manifest cannot be read is only logged, when it is read. + unsupported.find(_.manifest.dataSources.exists(_.equalsIgnoreCase(name))).foreach { pkg => + throw QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, + "UNSUPPORTED_ABI_VERSION", + Map( + "version" -> pkg.manifest.abiVersion.toString, + "supported" -> NativeLibrary.ABI_VERSION.toString)) + } + invalid.find(_.dataSources.exists(_.equalsIgnoreCase(name))).foreach(p => throw p.error) + None + } + } + + /** + * Like [[lookup]], and keeps the result for the provider that the caller creates next for the + * given name, which takes it with [[takeLookup]]. + */ + def lookupForProvider(name: String, conf: SQLConf): Option[(String, NativeLibraryLocation)] = { + val result = lookup(name, conf) + result.foreach(found => lastLookup.set(name -> found)) + result + } + + /** + * Takes the result of the last [[lookupForProvider]] on this thread, if it was for the given + * name, and forgets it. + */ + def takeLookup(name: String): Option[(String, NativeLibraryLocation)] = { + val last = Option(lastLookup.get()) + lastLookup.remove() + last.collect { case (lookedUp, found) if lookedUp.equalsIgnoreCase(name) => found } + } + + /** Forgets the result of the last [[lookupForProvider]] on this thread. */ + def clearLookup(): Unit = lastLookup.remove() + + private def lookupPackage( + name: String, + packages: Seq[NativeDataSourcePackage]): Option[(String, NativeDataSourcePackage)] = { + val found = packages + .flatMap(pkg => pkg.manifest.dataSources.find(_.equalsIgnoreCase(name)).map(_ -> pkg)) + // The same package can be found more than once, for example when it is in a configured + // directory and was also added as an artifact. + .groupBy(_._2.sha256).values.map(_.head).toSeq + found match { + case Seq() => None + case Seq(result) => Some(result) + case _ => + throw QueryCompilationErrors.nativeDataSourcePackageConflictError( + 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 0000000000000..e9e6f3815594f --- /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 0000000000000..a84fce136986e --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala @@ -0,0 +1,136 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.util.Base64 + +import com.fasterxml.jackson.core.json.JsonWriteFeature +import com.fasterxml.jackson.databind.json.JsonMapper +import com.fasterxml.jackson.databind.node.ObjectNode + +import org.apache.spark.sql.connector.expressions.{Expression, GeneralScalarExpression, Literal, NamedReference} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.types._ + +/** + * Encodes Data Source V2 predicates as the JSON expression trees passed to + * `PushPredicates` of `spark_datasource.proto`. + */ +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 + * `PushPredicates`. They are part of the interface, so a predicate with any other + * name, such as an arithmetic function, is not passed: its semantics, for example on overflow, + * would be part of the interface too, and the names of Data Source V2 grow with new releases. + */ + val FUNCTION_NAMES: Set[String] = Set( + "=", "<>", "<=>", "<", "<=", ">", ">=", "IN", "IS_NULL", "IS_NOT_NULL", + "STARTS_WITH", "ENDS_WITH", "CONTAINS", "AND", "OR", "NOT") + + /** + * Returns the JSON of the predicate on the columns of the given schema, or None if it cannot + * be expressed in JSON with the functions in [[FUNCTION_NAMES]]. + */ + def toJson(predicate: Predicate, schema: StructType): Option[String] = { + encode(predicate, schema).map(mapper.writeValueAsString) + } + + private def encode(expression: Expression, schema: StructType): Option[ObjectNode] = { + expression match { + case reference: NamedReference => + encodeColumn(reference, schema) + + case literal: Literal[_] => + encodeLiteral(literal.value(), literal.dataType()) + + case function: GeneralScalarExpression if FUNCTION_NAMES.contains(function.name()) => + val children = function.children().map(encode(_, schema)) + if (children.forall(_.isDefined)) { + val node = mapper.createObjectNode() + .put("type", "function") + .put("name", function.name()) + val array = node.putArray("children") + children.foreach(child => array.add(child.get)) + Some(node) + } else { + None + } + + case _ => None + } + } + + private def encodeColumn(reference: NamedReference, schema: StructType): Option[ObjectNode] = { + // A source compares strings by their bytes, so a column that contains strings with a + // non-binary collation, such as UTF8_LCASE, is not pushed down. Neither is a column that is + // not in the schema. + val hasNonBinaryCollation = columnType(reference.fieldNames().toSeq, schema).forall { + _.existsRecursively { + case st: StringType => !st.isUTF8BinaryCollation + case _ => false + } + } + if (hasNonBinaryCollation) { + None + } else { + val node = mapper.createObjectNode().put("type", "column") + val name = node.putArray("name") + reference.fieldNames().foreach(name.add) + Some(node) + } + } + + /** The type of the column or the nested field with the given name parts. */ + private def columnType(names: Seq[String], dataType: DataType): Option[DataType] = { + (names, dataType) match { + case (Seq(), _) => Some(dataType) + case (name +: rest, struct: StructType) => + struct.fields.find(_.name == name).flatMap(field => columnType(rest, field.dataType)) + case _ => None + } + } + + private def encodeLiteral(value: Any, dataType: DataType): Option[ObjectNode] = { + val node = mapper.createObjectNode() + .put("type", "literal") + .put("dataType", dataType.catalogString) + if (value == null) { + Some(node.putNull("value")) + } else { + val encoded = (value, dataType) match { + case (b: Boolean, BooleanType) => Some(node.put("value", b)) + case (n: Byte, ByteType) => Some(node.put("value", n.toInt)) + case (n: Short, ShortType) => Some(node.put("value", n)) + case (n: Int, IntegerType | DateType) => Some(node.put("value", n)) + case (n: Long, LongType | TimestampType | TimestampNTZType) => Some(node.put("value", n)) + case (n: Float, FloatType) if !n.isNaN && !n.isInfinite => Some(node.put("value", n)) + 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. + case (s, st: StringType) if st.isUTF8BinaryCollation => 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 9fcf81746fe30..cd04bba0c288e 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 c1fc7234d7c19..8fcf427ec8a75 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 0000000000000..a74ef7ae2a7bd --- /dev/null +++ b/sql/core/src/test/resources/native-datasource/test_native_datasource.cc @@ -0,0 +1,1840 @@ +/* + * 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 and +// the Protocol Buffers messages of spark_datasource.proto 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 PREDICATE_USE_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. 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, as the numbers of the enum +// Capability separated by commas. 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", opening 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 +// "unknown_fields", every reply has a field that Spark does not know. +// +// With the option "fail", a step fails with an error. With the option "zero", opening an object +// succeeds without setting its handle: "create" for a data source, "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 writes, streaming reads and pushdown: the replies to the pushdown +// requests are empty. +// ABI_VERSION The value of spark_ds_abi_version, 1 by default. +// NO_LAST_ERROR Leaves out spark_ds_last_error. +// NO_OPEN Leaves out spark_ds_open, 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 + +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(); + } + } +} + +// --------------------------------------------------------------------------------------------- +// 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--; + } + } +} + +// --------------------------------------------------------------------------------------------- +// 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_open 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()); + } +} + +// Whether the task with the given ID was killed. +bool isInterrupted(int64_t task_attempt_id) { + return host->size >= static_cast(offsetof(SparkDsHost, is_interrupted) + + sizeof(host->is_interrupted)) && + host->is_interrupted(task_attempt_id); +} + +// 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::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())); +} + +// --------------------------------------------------------------------------------------------- +// Data sources +// --------------------------------------------------------------------------------------------- + +std::string name(const std::string& base) { return std::string(NAME_PREFIX) + base; } + +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); +} + +// 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; + int64_t task_attempt_id; + 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 && !isInterrupted(state->task_attempt_id); i++) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + bool interrupted = isInterrupted(state->task_attempt_id); + 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); + 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; +} + +[[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; +} + +// --------------------------------------------------------------------------------------------- +// Messages +// --------------------------------------------------------------------------------------------- + +// A minimal Protocol Buffers codec for the messages of spark_datasource.proto, so that the test +// library needs nothing but a C++ compiler. A library would rather use the code that a Protocol +// Buffers compiler generates from spark_datasource.proto. +class ProtoReader { + public: + // Keeps the bytes of the message, which are often a temporary string. + explicit ProtoReader(std::string bytes) + : bytes_(std::move(bytes)), + p_(reinterpret_cast(bytes_.data())), + end_(p_ + bytes_.size()) {} + + // Moves to the next field, and returns false at the end of the message. + bool next() { + if (p_ == end_) return false; + uint64_t key = varint(); + field_ = static_cast(key >> 3); + wire_type_ = static_cast(key & 7); + return true; + } + + int field() const { return field_; } + + uint64_t varint() { + uint64_t value = 0; + for (int shift = 0; shift < 64; shift += 7) { + if (p_ == end_) throw std::runtime_error("truncated message"); + uint8_t byte = *p_++; + value |= static_cast(byte & 0x7f) << shift; + if ((byte & 0x80) == 0) return value; + } + throw std::runtime_error("invalid varint"); + } + + int64_t int64() { return static_cast(varint()); } + + // A string, bytes or a message. + std::string bytes() { + uint64_t length = varint(); + if (length > static_cast(end_ - p_)) throw std::runtime_error("truncated message"); + std::string value(reinterpret_cast(p_), static_cast(length)); + p_ += length; + return value; + } + + // Skips a field that the library does not know, as Protocol Buffers requires. + void skip() { + switch (wire_type_) { + case 0: varint(); break; + case 1: advance(8); break; + case 2: bytes(); break; + case 5: advance(4); break; + default: throw std::runtime_error("unsupported wire type " + std::to_string(wire_type_)); + } + } + + private: + void advance(size_t n) { + if (n > static_cast(end_ - p_)) throw std::runtime_error("truncated message"); + p_ += n; + } + + std::string bytes_; + const uint8_t* p_; + const uint8_t* end_; + int field_ = 0; + int wire_type_ = 0; +}; + +class ProtoWriter { + public: + ProtoWriter& varint(int field, uint64_t value) { + key(field, 0); + raw(value); + return *this; + } + + ProtoWriter& bytes(int field, const std::string& value) { + key(field, 2); + raw(value.size()); + out_ += value; + return *this; + } + + ProtoWriter& message(int field, const ProtoWriter& value) { return bytes(field, value.out_); } + + // A packed repeated field of enums. + ProtoWriter& packed(int field, const std::vector& values) { + ProtoWriter packed; + for (int32_t value : values) packed.raw(static_cast(static_cast(value))); + return bytes(field, packed.out_); + } + + const std::string& str() const { return out_; } + + private: + void key(int field, int wire_type) { + raw((static_cast(field) << 3) | static_cast(wire_type)); + } + + void raw(uint64_t value) { + while (value >= 0x80) { + out_ += static_cast((value & 0x7f) | 0x80); + value >>= 7; + } + out_ += static_cast(value); + } + + std::string out_; +}; + +// The fields of the messages, by their number in spark_datasource.proto. +enum OpenField { kSparkVersion = 1, kDataSource = 10, kBatchReader, kStreamReader, kWriter, + kPartitionReader, kDataWriter }; +enum CallField { kInferSchema = 1, kPushPredicates, kPushLimit, kPruneColumns, kPlanPartitions, + kSerialize, kInitialOffset, kLatestOffset, kPlanStreamPartitions, kCommitOffset, + kCommit, kAbort, kCommitDataWriter, kAbortDataWriter }; +enum StreamField { kRead = 1, kWrite }; +enum Capability { kBatchRead = 1, kMicroBatchRead, kBatchWrite, kStreamingWrite, kTruncate }; +enum PredicateUse { kUnused = 0, kPartial, kComplete }; +enum WriteMode { kAppend = 0, kTruncateMode }; + +// With the option "unknown_fields", every reply has a field that Spark does not know, which it +// must ignore. +std::string withUnknownField(const Options& options, const ProtoWriter& reply) { + ProtoWriter result = reply; + if (get(options, "unknown_fields", "") == "true") result.bytes(1000, "unknown"); + return result.str(); +} + +void sendReply(SparkDsSink reply, void* sink, const Options& options, const ProtoWriter& message) { + sendBytes(reply, sink, withUnknownField(options, message)); +} + +// The fields of a message, by number: the last value of each field. +std::map parseMessage(const std::string& bytes) { + std::map fields; + ProtoReader reader(bytes); + while (reader.next()) { + // The fields that the library reads are strings, bytes and messages, or varints, which are + // returned as decimal numbers. + int field = reader.field(); + try { + fields[field] = reader.bytes(); + } catch (const std::exception&) { + throw std::runtime_error("unexpected field " + std::to_string(field)); + } + } + return fields; +} + +// The varints of a message, by number. +std::map parseVarints(const std::string& bytes) { + std::map fields; + ProtoReader reader(bytes); + while (reader.next()) { + int field = reader.field(); + fields[field] = reader.varint(); + } + return fields; +} + +// A TaskContext message. +struct TaskInfo { + int64_t stage_id = 0; + int64_t stage_attempt_number = 0; + int64_t partition_id = 0; + int64_t attempt_number = 0; + int64_t task_attempt_id = 0; +}; + +TaskInfo parseTask(const std::string& bytes) { + auto fields = parseVarints(bytes); + auto field = [&](int number) { + return static_cast(static_cast(fields.count(number) ? fields[number] : 0)); + }; + TaskInfo task; + task.stage_id = field(1); + task.stage_attempt_number = field(2); + task.partition_id = field(3); + task.attempt_number = field(4); + task.task_attempt_id = fields.count(5) ? static_cast(fields[5]) : 0; + return task; +} + +// The names of the columns of a StructType message. +std::vector parseColumns(const std::string& struct_type) { + std::vector columns; + ProtoReader fields(struct_type); + while (fields.next()) { + if (fields.field() != 1) { + fields.skip(); + continue; + } + ProtoReader field(fields.bytes()); + std::string name; + while (field.next()) { + if (field.field() == 1) { + name = field.bytes(); + } else { + field.skip(); + } + } + columns.push_back(name); + } + return columns; +} + +// A DataType message of a type without parameters, by the number of its kind. +ProtoWriter simpleType(int kind) { return ProtoWriter().message(kind, ProtoWriter()); } + +// The DataType message of a column, as Spark infers it from the Arrow type: Spark cannot express +// some of the Arrow types of the extra columns, whose data it rejects. +ProtoWriter columnType(const std::string& name) { + enum { kInteger = 5, kLong, kDouble = 8, kDecimal, kString, kBoolean = 2, kDate = 12, + kTimestamp, kArray = 18, kMap, kStruct }; + auto field = [](const std::string& field_name, const ProtoWriter& type) { + return ProtoWriter().bytes(1, field_name).message(2, type).varint(3, 1); + }; + std::string format = findColumn(name)->format; + if (name == "m") { + return ProtoWriter().message(kMap, ProtoWriter() + .message(1, ProtoWriter().message(kString, ProtoWriter())) + .message(2, simpleType(kLong)) + .varint(3, 1)); + } else if (name == "st") { + return ProtoWriter().message(kStruct, ProtoWriter() + .message(1, field("a", simpleType(kLong))) + .message(1, field("b", ProtoWriter().message(kString, ProtoWriter())))); + } else if (name == "lv" || name == "ll") { + return ProtoWriter().message(kArray, + ProtoWriter().message(1, simpleType(kLong)).varint(2, 1)); + } else if (name == "dict" || format == "u") { + return ProtoWriter().message(kString, ProtoWriter()); + } else if (format == "l") { + return simpleType(kLong); + } else if (format == "i") { + return simpleType(kInteger); + } else if (format == "g") { + return simpleType(kDouble); + } else if (format == "b") { + return simpleType(kBoolean); + } else if (format == "tdD") { + return simpleType(kDate); + } else if (format.rfind("ts", 0) == 0) { + return simpleType(kTimestamp); + } else if (format.rfind("d:", 0) == 0) { + return ProtoWriter().message(kDecimal, ProtoWriter().varint(1, 10).varint(2, 2)); + } + throw std::runtime_error("no Spark type for the column " + name); +} + +// The StructType message of the given columns. +ProtoWriter structType(const std::vector& columns) { + ProtoWriter schema; + for (const auto& column : columns) { + schema.message(1, ProtoWriter().bytes(1, column).message(2, columnType(column)).varint(3, 1)); + } + return schema; +} + +// --------------------------------------------------------------------------------------------- +// Objects +// --------------------------------------------------------------------------------------------- + +// What a handle refers to. +struct Object { + virtual ~Object() = default; +}; + +struct DataSource : Object { + std::string name; + Options options; +}; + +// The batch or stream reader of native_range or native_counter. +struct Reader : Object { + Options state; // Serialized to the executors. +}; + +// The reader of one partition, on an executor. +struct PartitionReader : Object { + Options state; + std::string partition; + TaskInfo task; +}; + +// The writer of native_sink. +struct Writer : Object { + Options state; // path, overwrite, fail. + // The file that closing the writer 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. The commit records +// in _streams the number of streams and of batches that the data writer got. +struct DataWriter : Object { + std::string temp_path; + FILE* file = nullptr; + Options state; + bool committed = false; + int64_t streams = 0; + int64_t batches = 0; +}; + +// 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; +std::map closed_data_writers; + +// 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; +} + +// Returns a new object as a handle. The allocator may reuse the address of a closed data writer. +void* newHandle(Object* object) { + std::lock_guard guard(closed_data_writers_lock); + closed_data_writers.erase(object); + return object; +} + +template +T* as(void* handle, const char* request) { + auto* object = dynamic_cast(static_cast(handle)); + if (object == nullptr) throw std::runtime_error(std::string(request) + " on another object"); + return object; +} + +// With the option "task_path", records the task of a read or a data writer in a file named after +// its partition: its fields. +void recordTask(const Options& options, const TaskInfo& 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.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)); +} + +DataSource parseDataSource(const std::string& bytes) { + DataSource source; + ProtoReader reader(bytes); + while (reader.next()) { + if (reader.field() == 1) { + source.name = reader.bytes(); + } else if (reader.field() == 2) { + // A map is a repeated message of entries, with the key and the value. + ProtoReader entry(reader.bytes()); + std::string key; + std::string value; + while (entry.next()) { + if (entry.field() == 1) { + key = entry.bytes(); + } else if (entry.field() == 2) { + value = entry.bytes(); + } else { + entry.skip(); + } + } + source.options[key] = value; + } else { + reader.skip(); + } + } + return source; +} + +// The capabilities of a data source, unless the option "capabilities" declares others: their +// numbers, separated by commas. +std::vector capabilitiesOf(const DataSource& source) { + auto declared = source.options.find("capabilities"); + if (declared != source.options.end()) { + std::vector capabilities; + for (const auto& value : split(declared->second)) capabilities.push_back(std::stoi(value)); + return capabilities; + } + if (source.name == name("native_range")) return {kBatchRead}; +#ifndef READ_ONLY + if (source.name == name("native_counter")) return {kMicroBatchRead}; + if (source.name == name("native_sink")) return {kBatchWrite, kStreamingWrite, kTruncate}; +#endif + return {}; +} + +// --------------------------------------------------------------------------------------------- +// spark_ds_open +// --------------------------------------------------------------------------------------------- + +void* openDataSource(const std::string& message, const std::string& spark_version, + SparkDsSink reply, void* sink) { + auto source = std::make_unique(parseDataSource(message)); + 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 log = get(source->options, "log", ""); + if (!log.empty()) hostLog(SPARK_DS_LOG_WARN, log); + std::string version = get(source->options, "spark_version", ""); + if (!version.empty() && version != spark_version) { + throw std::runtime_error("the version of Spark is " + spark_version); + } + sendReply(reply, sink, source->options, ProtoWriter().packed(1, capabilitiesOf(*source))); + if (returnsZero(source->options, "create")) return nullptr; + return newHandle(source.release()); +} + +void* openBatchReader(const std::string& message) { + auto fields = parseMessage(message); + DataSource source = parseDataSource(fields[1]); + if (source.name != name("native_range")) { + throw std::runtime_error(source.name + " does not support batch reads"); + } + auto columns = parseColumns(fields[2]); + for (const auto& column : columns) findColumn(column); + if (returnsZero(source.options, "create_reader")) return nullptr; + auto reader = std::make_unique(); + 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", "claim_predicates", "predicate_use", + "wait_for_interrupt", "task_path", "unknown_fields"}) { + reader->state[option] = get(source.options, option, ""); + } + return newHandle(reader.release()); +} + +#ifndef READ_ONLY + +void* openStreamReader(const std::string& message) { + auto fields = parseMessage(message); + DataSource source = parseDataSource(fields[1]); + if (source.name != name("native_counter")) { + throw std::runtime_error(source.name + " does not support streaming reads"); + } + auto columns = join(parseColumns(fields[2])); + if (returnsZero(source.options, "create_stream_reader")) return nullptr; + auto reader = std::make_unique(); + 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", ""); + reader->state["stream"] = "true"; + return newHandle(reader.release()); +} + +void* openWriter(const std::string& message) { + ProtoReader reader(message); + DataSource source; + std::vector columns; + uint64_t mode = kAppend; + std::string query_id; + while (reader.next()) { + switch (reader.field()) { + case 1: source = parseDataSource(reader.bytes()); break; + case 2: columns = parseColumns(reader.bytes()); break; + case 3: mode = reader.varint(); break; + case 4: query_id = reader.bytes(); break; + default: reader.skip(); + } + } + if (source.name != name("native_sink")) { + throw std::runtime_error(source.name + " does not support writes"); + } + if (mode != kAppend && mode != kTruncateMode) { + throw std::runtime_error("unknown mode " + std::to_string(mode)); + } + if (columns.empty()) 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 nullptr; + auto writer = std::make_unique(); + writer->state["path"] = path; + writer->state["overwrite"] = mode == kTruncateMode ? "true" : "false"; + for (const char* option : {"fail", "zero", "task_path", "unknown_fields"}) { + writer->state[option] = get(source.options, option, ""); + } + writer->state["commit_messages"] = get(source.options, "commit_messages", "1"); + // 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 (!query_id.empty()) writeFile(path + "/_QUERY_ID", query_id); + return newHandle(writer.release()); +} + +void* openDataWriter(const std::string& message) { + ProtoReader reader(message); + TaskInfo task; + Options state; + int64_t epoch_id = 0; + while (reader.next()) { + switch (reader.field()) { + case 1: task = parseTask(reader.bytes()); break; + case 2: state = deserialize(reader.bytes()); break; + case 3: epoch_id = reader.int64(); break; + default: reader.skip(); + } + } + if (returnsZero(state, "create_data_writer")) return nullptr; + 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 = std::make_unique(); + writer->temp_path = temp_path; + writer->file = file; + writer->state = state; + return newHandle(writer.release()); +} + +#endif // READ_ONLY + +void* openPartitionReader(const std::string& message) { + ProtoReader reader(message); + auto partition = std::make_unique(); + while (reader.next()) { + switch (reader.field()) { + case 1: partition->task = parseTask(reader.bytes()); break; + case 2: partition->state = deserialize(reader.bytes()); break; + case 3: partition->partition = reader.bytes(); break; + default: reader.skip(); + } + } + maybeFail(partition->state, "open"); + recordTask(partition->state, partition->task); + return newHandle(partition.release()); +} + +// --------------------------------------------------------------------------------------------- +// spark_ds_call +// --------------------------------------------------------------------------------------------- + +void inferSchema(DataSource* source, SparkDsSink reply, void* sink) { + 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"); + auto columns = allColumns(counter, get(source->options, "extra", "")); + sendReply(reply, sink, source->options, + ProtoWriter().message(kInferSchema, ProtoWriter().message(1, structType(columns)))); +} + +#ifndef READ_ONLY + +void pushPredicates(Reader* reader, const std::string& message, SparkDsSink reply, void* sink) { + std::vector json; + ProtoReader request(message); + while (request.next()) { + if (request.field() == 1) { + json.push_back(request.bytes()); + } else { + request.skip(); + } + } + int64_t end = getLong(reader->state, "hi", 0); + int64_t lo = getLong(reader->state, "lo", 0); + int64_t hi = end; + std::vector uses(json.size(), kUnused); + // With the option "claim_predicates", reports every predicate as evaluated completely, but + // evaluates none of them. + if (reader->state["claim_predicates"] == "true") { + std::fill(uses.begin(), uses.end(), kComplete); + } else { + std::string predicate_use = reader->state["predicate_use"]; + bool partial = predicate_use == "partial"; + int32_t use = predicate_use.empty() ? kComplete + : partial ? kPartial + : 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 : kUnused; + } + 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()); + sendReply(reply, sink, reader->state, + ProtoWriter().message(kPushPredicates, ProtoWriter().packed(1, uses))); +} + +void pushLimit(Reader* reader, const std::string& message, SparkDsSink reply, void* sink) { + // Spark pushes a limit at most once. + if (reader->state.count("limit") > 0) throw std::runtime_error("PushLimit was called twice"); + auto fields = parseVarints(message); + reader->state["limit"] = std::to_string(static_cast(fields[1])); + sendReply(reply, sink, reader->state, + ProtoWriter().message(kPushLimit, ProtoWriter().varint(1, 1))); +} + +void pruneColumns(Reader* reader, const std::string& message, SparkDsSink reply, void* sink) { + std::vector names; + ProtoReader request(message); + while (request.next()) { + if (request.field() == 1) { + names.push_back(request.bytes()); + } else { + request.skip(); + } + } + for (const auto& column : names) findColumn(column); + reader->state["columns"] = join(names); + sendReply(reply, sink, reader->state, + ProtoWriter().message(kPruneColumns, ProtoWriter().varint(1, 1))); +} + +#endif // READ_ONLY + +void planPartitions(Reader* reader, SparkDsSink reply, void* sink) { + 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; + } + ProtoWriter partitions; + 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); + partitions.bytes(1, std::to_string(start) + "," + std::to_string(stop)); + } + sendReply(reply, sink, reader->state, ProtoWriter().message(kPlanPartitions, partitions)); +} + +void serializeState(const Options& state, SparkDsSink reply, void* sink) { + sendReply(reply, sink, state, + ProtoWriter().message(kSerialize, ProtoWriter().bytes(1, serialize(state)))); +} + +#ifndef READ_ONLY + +void initialOffset(Reader* reader, SparkDsSink reply, void* sink) { + sendReply(reply, sink, reader->state, + ProtoWriter().message(kInitialOffset, ProtoWriter().bytes(1, "{\"offset\":0}"))); +} + +// Each call makes `step` more rows available, up to `max_offset`. +void latestOffset(Reader* reader, SparkDsSink reply, void* sink) { + Options& state = reader->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; + std::string offset = "{\"offset\":" + std::to_string(latest) + "}"; + if (bad_offset == "multiline") { + // Pretty-printed JSON, as serde_json::to_string_pretty writes it. + offset = "{\n \"offset\": " + std::to_string(latest) + "\n}"; + } else if (bad_offset == "invalid") { + offset = "offset " + std::to_string(latest); + } + sendReply(reply, sink, state, + ProtoWriter().message(kLatestOffset, ProtoWriter().bytes(1, offset))); +} + +void planStreamPartitions(Reader* reader, const std::string& message, SparkDsSink reply, + void* sink) { + auto fields = parseMessage(message); + auto lo = static_cast(JsonParser(fields[1]).parse()["offset"].number); + auto hi = static_cast(JsonParser(fields[2]).parse()["offset"].number); + ProtoWriter partitions; + if (hi > lo) partitions.bytes(1, std::to_string(lo) + "," + std::to_string(hi)); + sendReply(reply, sink, reader->state, + ProtoWriter().message(kPlanStreamPartitions, partitions)); +} + +void commitOffset(Reader* reader, const std::string& message) { + std::string path = get(reader->state, "commit_path", ""); + if (!path.empty()) writeFile(path, parseMessage(message)[1]); +} + +// The epoch ID and the messages of a Commit or an Abort, with whether a task did not commit. +int64_t parseTaskMessages(const std::string& message, std::vector* messages, + int64_t* uncommitted) { + int64_t epoch_id = 0; + ProtoReader request(message); + while (request.next()) { + if (request.field() == 1) { + epoch_id = request.int64(); + } else if (request.field() == 2) { + ProtoReader task(request.bytes()); + bool committed = false; + std::string bytes; + while (task.next()) { + if (task.field() == 1) { + committed = task.varint() != 0; + } else if (task.field() == 2) { + bytes = task.bytes(); + } else { + task.skip(); + } + } + if (committed) { + messages->push_back(bytes); + } else { + (*uncommitted)++; + } + } else { + request.skip(); + } + } + return epoch_id; +} + +void commitWrite(Writer* writer, const std::string& message) { + std::string path = writer->state["path"]; + maybeFail(writer->state, "commit"); + std::vector temp_files; + int64_t uncommitted = 0; + int64_t epoch_id = parseTaskMessages(message, &temp_files, &uncommitted); + if (uncommitted > 0) throw std::runtime_error("commit got a task that did not commit"); + 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. +void abortWrite(Writer* writer, const std::string& message) { + std::vector temp_files; + int64_t uncommitted = 0; + int64_t epoch_id = parseTaskMessages(message, &temp_files, &uncommitted); + writer->closed_marker = closedMarker(epoch_id); + for (const auto& temp_file : temp_files) unlink(temp_file.c_str()); + writeFile(writer->state["path"] + + (epoch_id < 0 ? "/_ABORTED" : "/_ABORTED_" + std::to_string(epoch_id)), + std::to_string(uncommitted)); +} + +void commitDataWriter(DataWriter* writer, SparkDsSink reply, void* sink) { + // If the commit fails, Spark aborts the data writer. + maybeFail(writer->state, "commitDataWriter"); + // With the option "commit_messages", the reply is passed that many times. + int64_t messages = getLong(writer->state, "commit_messages", 1); + for (int64_t i = 0; i < messages; i++) { + sendReply(reply, sink, writer->state, ProtoWriter().message( + kCommitDataWriter, ProtoWriter().bytes(1, writer->temp_path))); + } + std::fclose(writer->file); + writer->file = nullptr; + writer->committed = true; + std::string dir = writer->state["path"] + "/_streams"; + makeDirs(dir); + writeFile(dir + "/" + writer->temp_path.substr(writer->temp_path.rfind('/') + 1), + std::to_string(writer->streams) + "," + std::to_string(writer->batches)); +} + +void abortDataWriter(DataWriter* writer) { + // 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()); +} + +#endif // READ_ONLY + +void call(void* handle, const std::string& request, SparkDsSink reply, void* sink) { + ProtoReader message(request); + while (message.next()) { + int field = message.field(); + if (field < kInferSchema || field > kAbortDataWriter) { + // A request that the library does not know, which Spark does not send to it. + message.skip(); + continue; + } + std::string body = message.bytes(); + switch (field) { + case kInferSchema: return inferSchema(as(handle, "InferSchema"), reply, sink); + case kPlanPartitions: + return planPartitions(as(handle, "PlanPartitions"), reply, sink); + case kSerialize: { + auto* object = static_cast(handle); + if (auto* reader = dynamic_cast(object)) { + return serializeState(reader->state, reply, sink); + } + return serializeState(as(handle, "Serialize")->state, reply, sink); + } +#ifdef READ_ONLY + // A library without pushdown replies to it with empty replies: Spark pushes nothing. + case kPushPredicates: + case kPushLimit: + case kPruneColumns: + return; + default: + throw std::runtime_error("unsupported request " + std::to_string(field)); +#else + case kPushPredicates: + return pushPredicates(as(handle, "PushPredicates"), body, reply, sink); + case kPushLimit: return pushLimit(as(handle, "PushLimit"), body, reply, sink); + case kPruneColumns: + return pruneColumns(as(handle, "PruneColumns"), body, reply, sink); + case kInitialOffset: + return initialOffset(as(handle, "InitialOffset"), reply, sink); + case kLatestOffset: return latestOffset(as(handle, "LatestOffset"), reply, sink); + case kPlanStreamPartitions: + return planStreamPartitions(as(handle, "PlanStreamPartitions"), body, reply, sink); + case kCommitOffset: return commitOffset(as(handle, "CommitOffset"), body); + case kCommit: return commitWrite(as(handle, "Commit"), body); + case kAbort: return abortWrite(as(handle, "Abort"), body); + case kCommitDataWriter: + return commitDataWriter(as(handle, "CommitDataWriter"), reply, sink); + case kAbortDataWriter: + return abortDataWriter(as(handle, "AbortDataWriter")); +#endif + } + } +} + +// --------------------------------------------------------------------------------------------- +// spark_ds_stream +// --------------------------------------------------------------------------------------------- + +void read(PartitionReader* partition, struct ArrowArrayStream* stream) { + const Options& state = partition->state; + // Returns success without initializing the stream, which Spark must reject. + if (get(state, "fail", "") == "uninitialized") return; + auto bounds = split(partition->partition); + if (bounds.size() != 2) throw std::runtime_error("invalid partition"); + 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", ""), + partition->task.task_attempt_id, + get(state, "wait_for_interrupt", ""), ""}; +} + +#ifndef READ_ONLY + +// Writes the batches of the stream. Spark releases the stream after the call, since it is not +// moved here, but the schema and the arrays that the stream returns are the library's to release. +void write(DataWriter* writer, struct ArrowArrayStream* stream) { + maybeFail(writer->state, "write"); + writer->streams++; + ArrowSchema schema; + if (stream->get_schema(stream, &schema) != 0) { + throw std::runtime_error(std::string("get_schema failed: ") + stream->get_last_error(stream)); + } + std::unique_ptr schema_owner( + &schema, [](ArrowSchema* s) { s->release(s); }); + while (true) { + ArrowArray array; + if (stream->get_next(stream, &array) != 0) { + throw std::runtime_error(std::string("get_next failed: ") + stream->get_last_error(stream)); + } + if (array.release == nullptr) break; // The end of the stream. + std::unique_ptr array_owner( + &array, [](ArrowArray* a) { a->release(a); }); + writer->batches++; + 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); + } + } +} + +#endif // READ_ONLY + +// --------------------------------------------------------------------------------------------- +// spark_ds_close +// --------------------------------------------------------------------------------------------- + +void close(void* handle) { + std::unique_ptr object(static_cast(handle)); + guarded([&] { +#ifndef READ_ONLY + if (auto* reader = dynamic_cast(object.get())) { + std::string path = get(reader->state, "closed_path", ""); + if (!path.empty()) writeFile(path, ""); + } else if (auto* writer = dynamic_cast(object.get())) { + if (!writer->closed_marker.empty()) { + writeFile(writer->state["path"] + writer->closed_marker, ""); + } + } else if (auto* data_writer = dynamic_cast(object.get())) { + if (data_writer->file != nullptr) std::fclose(data_writer->file); + std::string path = data_writer->state["path"]; + std::string dir = path + "/_closed_data_writers"; + makeDirs(dir); + std::string temp_path = data_writer->temp_path; + writeFile(dir + "/" + temp_path.substr(temp_path.rfind('/') + 1), ""); + std::lock_guard guard(closed_data_writers_lock); + closed_data_writers[handle] = path; + } +#endif + }); +} + +// spark_ds_open, which NO_OPEN does not export. +[[maybe_unused]] int open(const SparkDsHost* spark, const uint8_t* request, int64_t length, + SparkDsSink reply, void* sink, void** handle) { + return guarded([&] { + host = spark; + // The fields may come in any order, so they are all read before the object is opened. + std::string spark_version; + int object = 0; + std::string message; + ProtoReader reader(toBytes(request, length)); + while (reader.next()) { + if (reader.field() == kSparkVersion) { + spark_version = reader.bytes(); + } else if (reader.field() >= kDataSource && reader.field() <= kDataWriter) { + object = reader.field(); + message = reader.bytes(); + } else { + reader.skip(); + } + } + switch (object) { + case kDataSource: *handle = openDataSource(message, spark_version, reply, sink); break; + case kBatchReader: *handle = openBatchReader(message); break; + case kPartitionReader: *handle = openPartitionReader(message); break; +#ifndef READ_ONLY + case kStreamReader: *handle = openStreamReader(message); break; + case kWriter: *handle = openWriter(message); break; + case kDataWriter: *handle = openDataWriter(message); break; +#endif + default: { + const char* kinds[] = {"data source", "batch reader", "stream reader", "writer", + "partition reader", "data writer"}; + bool known = object >= kDataSource && object <= kDataWriter; + throw std::runtime_error(std::string("the library cannot open a ") + + (known ? kinds[object - kDataSource] : "new kind of object")); + } + } + }); +} + +} // namespace + +// --------------------------------------------------------------------------------------------- +// spark_datasource.h +// --------------------------------------------------------------------------------------------- + +int32_t spark_ds_abi_version(void) { return ABI_VERSION; } + +#ifndef NO_LAST_ERROR +const char* spark_ds_last_error(void) { return last_error.c_str(); } +#endif + +#ifndef NO_OPEN + +int spark_ds_open(const SparkDsHost* spark, const uint8_t* request, int64_t length, + SparkDsSink reply, void* sink, void** handle) { + return open(spark, request, length, reply, sink, handle); +} + +#endif // NO_OPEN + +int spark_ds_call(void* handle, const uint8_t* request, int64_t length, SparkDsSink reply, + void* sink) { + if (usedAfterClose(handle)) return 0; + return guarded([&] { call(handle, toBytes(request, length), reply, sink); }); +} + +int spark_ds_stream(void* handle, const uint8_t* request, int64_t length, + [[maybe_unused]] struct ArrowArrayStream* in, struct ArrowArrayStream* out) { + if (usedAfterClose(handle)) return 0; + return guarded([&] { + ProtoReader message(toBytes(request, length)); + while (message.next()) { + int field = message.field(); + message.skip(); + if (field == kRead) return read(as(handle, "Read"), out); +#ifndef READ_ONLY + if (field == kWrite) return write(as(handle, "Write"), in); +#endif + } + throw std::runtime_error("unsupported stream request"); + }); +} + +void spark_ds_close(void* handle) { + if (usedAfterClose(handle)) return; + close(handle); +} 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 0000000000000..898e1796abe03 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala @@ -0,0 +1,402 @@ +/* + * 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 scala.collection.mutable.ArrayBuffer + +import org.apache.spark.SparkException +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.connector.catalog.{Table, 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 + } + + 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("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") + + override def schema(): StructType = StructType.fromDDL("id BIGINT, name STRING") + + override def reader(schema: StructType): ColumnarReader = new TestRangeReader(end, schema) + + 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 TestRangeReader(end: Long, schema: StructType) extends ColumnarReader { + 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 partitions(): Array[InputPartition] = { + val hi = if (limit == Long.MaxValue) end else math.min(end, lo + limit) + val mid = (lo + hi) / 2 + Array(TestRangePartition(lo, mid), TestRangePartition(mid, hi)) + } + + override def read(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 0000000000000..bc7f369ba930f --- /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 0000000000000..133d04946b796 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala @@ -0,0 +1,1780 @@ +/* + * 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, 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.ColumnarScan +import org.apache.spark.sql.execution.datasources.v2.python.{PythonDataSourceV2, UserDefinedPythonDataSource} +import org.apache.spark.sql.functions.{col, udf} +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.sql.streaming.StreamingQuery +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") + 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 noOpenLibrary = + compile("test_native_datasource_no", "-DNO_OPEN", "-DNAME_PREFIX=\"no_\"") + 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": 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") + assert(NativePlatform.current == NativePlatform.name( + System.getProperty("os.name"), System.getProperty("os.arch"))) + } + + 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 PushPredicates 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""")) + 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) + 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 in `PushPredicates`, 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") + def char(codePoint: Int): String = new String(Character.toChars(codePoint)) + val (emoji, eAcute, replacement) = (char(0x1F600), char(0xE9), char(0xFFFD)) + val rows = Seq( + Row(0, 1, 0.0, "a"), + Row(1, 2, -0.0, ""), + Row(2, null, Double.NaN, eAcute), + Row(3, 3, Double.PositiveInfinity, emoji), + Row(4, 1, Double.NegativeInfinity, null), + Row(5, null, 1.5, "ab"), + Row(6, 2, null, "b" + eAcute), + Row(7, 4, 2.5, replacement)) + 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')") + + 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) + 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 in PushPredicates of " + + "spark_datasource.proto, 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.")}, + | "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("report link errors") { + // 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") + 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 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") + 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" -> "spark_ds_call replied 2 times instead of once.")) + // 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("write the data of a task in one stream") { + val session = newSession(defaultPackage) + session.conf.set(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key, "3") + val dir = Utils.createTempDir() + session.range(10).repartition(1).write.format("native_sink").option("path", dir.getPath) + .mode("append").save() + // The library reads the 10 rows of the task in one stream, in batches of at most 3 rows. + val streams = new File(dir, "_streams").listFiles() + assert(streams.length == 1) + assert(Files.readString(streams.head.toPath) == "1,4") + } + + nativeTest("report the errors of the data that a library writes") { + val session = newSession(defaultPackage) + val dir = Utils.createTempDir() + val failing = udf { (id: Long) => + if (id == 5) throw new IllegalStateException("injected failure in the data") + id + } + val e = intercept[Exception] { + session.range(10).repartition(1).select(failing(col("id")).as("id")).write + .format("native_sink").option("path", dir.getPath).mode("append").save() + } + // The library reads the data through the stream, which only gets the message of the error: + // the error itself is reported, rather than the error that the library returns. + val causes = Iterator.iterate[Throwable](e)(_.getCause).takeWhile(_ != null).toSeq + assert(causes.exists(c => c.isInstanceOf[IllegalStateException] && + c.getMessage == "injected failure in the data"), e) + assert(!causes.exists { + case t: SparkThrowable => t.getCondition == "NATIVE_DATA_SOURCE_ERROR" + case _ => false + }, e) + checkDataWritersClosed(dir) + assert(new File(dir, "_ABORTED").isFile) + } + + 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("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", "LatestOffset returned no offset.") + check("multiline", "LatestOffset returned an offset with a line break, but an " + + "offset must be JSON on a single line.") + check("invalid", "LatestOffset 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", proto.Capability.CAPABILITY_BATCH_READ.getNumber.toString) + .load().collect() + } + // Without spark_ds_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" -> "spark_ds_open 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 reader = + new NativeDataSource(location, name, options).reader(StructType.fromDDL("id BIGINT")) + (1 to 2).foreach { _ => + checkError( + exception = intercept[SparkRuntimeException](reader.partitions()), + 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" -> "spark_ds_open 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. + check("DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", Map("description" -> "ro_native_sink")) { + session.range(1).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 support its operations: here, the library + // cannot open a writer. + check("NATIVE_DATA_SOURCE_ERROR", Map("action" -> "plan a write to", + "name" -> "ro_native_sink", "msg" -> "the library cannot open a writer")) { + session.range(1).write.format("ro_native_sink") + .option("capabilities", proto.Capability.CAPABILITY_BATCH_WRITE.getNumber.toString) + .mode("append").save() + } + // Replacing the existing data requires the capability CAPABILITY_TRUNCATE. + val dir = Utils.createTempDir() + val noTruncate = + Seq(proto.Capability.CAPABILITY_BATCH_WRITE, proto.Capability.CAPABILITY_STREAMING_WRITE) + .map(_.getNumber).mkString(",") + 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).write.format("native_sink").option("path", dir.getPath) + .option("capabilities", noTruncate).mode("overwrite").save() + } + 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("ignore the fields of messages that they do not know") { + // A library built against a later version of the messages may reply with fields that Spark + // does not know: here, every reply has one. + val session = newSession(defaultPackage) + checkRows(range(session, "end" -> "6", "unknown_fields" -> "true").filter("id > 2"), + (3L until 6L).map(expectedRow)) + // And a library built against an earlier version gets requests with fields that it does not + // know: here, a request of a later version, with a field that the library does not know. + val (name, location) = + NativeDataSourceRegistry.lookup("native_range", session.sessionState.conf).get + val library = NativeLibraries.get(location) + val known = NativeMessages.open() + .setDataSource(proto.DataSource.newBuilder().setName(name)) + .build() + // A message with a field of a later version: field 1000, with a string. Spark keeps the fields + // that it does not know of a message that it parses. + val unknown = Array[Byte](0xc2.toByte, 0x3e, 3) ++ "new".getBytes(UTF_8) + val request = proto.OpenRequest.parseFrom(known.toByteArray ++ unknown) + assert(request.toByteArray.length == known.toByteArray.length + unknown.length) + val opened = library.open(request) + try { + assert(opened.reply.getCapabilitiesList.asScala == + Seq(proto.Capability.CAPABILITY_BATCH_READ)) + } finally { + library.close(opened.handle) + } + } + + 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")) + + 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 TaskContext 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(stageId, stageAttempt, partitionId, attempt, taskAttemptId)) => + assert(partitionId == partition) + assert(stageId >= 0 && stageAttempt == 0 && attempt == 0 && taskAttemptId >= 0) + } + assert(tasks.values.map(_(4)).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("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 do not export the functions of the interface") { + val pkg = createPackage(noOpenLibrary, dataSources("no_")) + // The library stays loaded, and the error is the same when it is used again. + (1 to 2).foreach { _ => + checkError( + exception = intercept[SparkRuntimeException] { + newSession(pkg).read.format("no_native_range").load() + }, + condition = "INVALID_NATIVE_DATA_SOURCE_PACKAGE.INVALID_LIBRARY", + parameters = Map( + "path" -> pkg.getCanonicalPath, + "library" -> s"${NativePlatform.current}/${noOpenLibrary.getName}", + "reason" -> "The library does not export the function spark_ds_open.")) + } + } +} 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 0000000000000..4b7e3cf55a49c --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceTestUtils.scala @@ -0,0 +1,168 @@ +/* + * 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 the functions of the header + // are 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) + } + } + } +}