HyukjinKwon commented on code in PR #59313: URL: https://github.com/apache/spark/pull/59313#discussion_r4255249954
########## sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java: ########## @@ -0,0 +1,398 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.datasource; + +import org.apache.spark.annotation.Evolving; + +/** + * The binary interface between Spark and data sources implemented in native code, such as Rust + * or C++. + * <p> + * A native data source is a shared library that implements the {@code native} methods of this + * class as JNI functions, for example + * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each + * library into its own copy of this class, so any number of libraries that export the same + * functions can be loaded side by side. Spark adapts a native data source to Data Source V2, with + * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and + * limit pushdown. + * <p> + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of {@code my_source} is + * {@code spark_datasource_my_source}, such as {@code libspark_datasource_my_source.so} in + * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of + * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux, and in the + * {@code lib} directory of each installation prefix in {@code PATH}. A library can also be + * distributed in a native data source package: a zip file with the extension {@code .sparkpkg} + * that contains a manifest named {@code spark-native-datasource.json} and the library built for + * one or more platforms, added to a session with {@code spark.addArtifact} or under the paths in + * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL + * guide for details and examples. + * + * <h2>Lifecycle</h2> + * Spark creates a data source with {@link #createDataSource} on the driver for each schema + * inference, scan and write, and closes it right after. + * <ul> + * <li><b>Batch reads.</b> On the driver, Spark creates a reader with {@link #createReader}, + * pushes operations down by calling, each at most once and in this order, + * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions + * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, + * and closes it. Then it calls {@link #read} on the executors, in one task for each + * partition.</li> + * <li><b>Streaming reads.</b> On the driver, Spark creates a stream reader with + * {@link #createStreamReader} and keeps it until the streaming query stops. For each + * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with + * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then + * calls {@link #read} on the executors. Offsets are JSON strings defined by the library.</li> + * <li><b>Writes.</b> On the driver, once it computed the data to write, right before it starts + * the tasks of the write (or of a micro-batch), Spark creates a writer with + * {@link #createWriter} or {@link #createStreamWriter}, and serializes its state with + * {@link #serializeWriter}. If computing the data fails, Spark creates no writer. On the + * executors, each task creates a data writer with {@link #createDataWriter}, calls + * {@link #write} for each batch of at most {@code spark.sql.execution.arrow.maxRecordsPerBatch} + * rows, which also ends once it holds {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, + * and calls {@link #commitDataWriter}, or {@link #abortDataWriter} if the task or the + * commit fails. Then the driver calls {@link #commit} with the messages of the tasks if they + * all succeeded, or {@link #abort} otherwise, including when {@link #commit} fails, and closes + * the writer. A streaming write commits + * or aborts each micro-batch (epoch). A micro-batch streaming query creates a new stream writer + * for each micro-batch, and Spark closes the previous stream writer of the query when the query + * creates a new one. Spark is not told when a streaming query stops, so it closes the last + * stream writer of the query once the writer is no longer referenced, or when the query is + * restarted.</li> + * </ul> + * + * <h2>Conventions</h2> + * <ul> + * <li><b>Handles.</b> The {@code create} methods return an opaque, non-zero handle, typically a + * pointer. Spark releases each handle exactly once with the matching {@code close} method, or + * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers, though not + * always right after the last use of a stream writer (see Writes above). A reader or + * writer handle must not depend on the data source handle it was created from, because Spark + * closes the data source handle right after creating it.</li> + * <li><b>State.</b> Partitions, the state of readers and writers, and commit messages are byte + * arrays in a format defined by the library. Spark sends them to the executors and back.</li> + * <li><b>Errors.</b> To report an error, throw a Java exception, for example with the JNI + * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and return any value. Spark + * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.</li> + * <li><b>Arrow data.</b> Schemas and data are exchanged with the + * <a href="https://arrow.apache.org/docs/format/CDataInterface.html";>Arrow C data + * interface</a>. A {@code long} argument named {@code ...Address} is the address of a C struct + * allocated by Spark. For an output struct, the library initializes it and Spark takes + * ownership. For an input struct, the library may move it to take ownership; otherwise, Spark + * releases it after the call returns, so the library must not use it afterwards. The arrays + * that the library exports must have an offset of 0, at every level: Arrow Java ignores the + * offsets of the arrays it imports, so Spark rejects an array with another offset, such as a + * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with + * fewer values than its parent refers to, such as a column shorter than its batch. The buffers + * of an array must hold all of its values, as the C data interface requires: it does not carry + * the sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that + * is too short.</li> + * <li><b>Schemas.</b> Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: + * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, + * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is + * {@code timestamp[us]} without one. Data read from a library may also use + * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, and any names + * for the fields of the entries of a map, but not unsigned integers, {@code decimal256} or + * dictionary encoding.</li> + * <li><b>Options.</b> Option keys are passed in lower case, because Spark matches them + * case-insensitively.</li> + * <li><b>Optional methods.</b> A library only needs to export the functions for the operations + * it supports. If a function is missing, Spark treats the operation as unsupported, except for + * the pushdown methods and the {@code serialize} methods, which then push nothing down and + * return no state.</li> + * <li><b>Threads.</b> Spark calls the library from multiple threads, and calls on different + * handles may run concurrently. Calls on the same handle never run concurrently.</li> + * </ul> + * + * @since 4.4.0 + */ +@Evolving +public final class NativeBridge { + + /** The version of this interface. It changes when the interface changes incompatibly. */ + public static final int ABI_VERSION = 1; Review Comment: Thanks, done. The class Javadoc of `NativeBridge` now has a "Compatibility" section: the JVM binds by name only, so the class name, the method names and their signatures are the interface. Within an `ABI_VERSION`, the signature and semantics of an existing method never change, new methods are added only as optional methods, and new values appear only where the Javadoc says so (for now, only the names of pushed predicates). An incompatible change uses a new `ABI_VERSION`, with new names or a new class. It also says that Spark supports version 1 at least in all the Spark 4 releases from 4.4.0. The docs page has the same text under "How It Works > Compatibility". On limiting version 1 to `long` addresses and Arrow C structs: I'd keep the current types. Apart from the JNI calling convention, version 1 only uses `long`, `int`, `boolean`, `String`, `byte[]` and arrays of them. Each of these maps directly onto a C type (`int64_t`, `int32_t`, `bool`, a NUL-terminated UTF-8 string, a pointer and a length). So a library can already keep its logic behind a plain C function such as `partitions(reader, out_ptrs, out_lens)`, and its JNI glue is just these conversions; that glue is all that changes when it moves to a C ABI. Encoding options, offsets and partitions as Arrow arrays would make every version 1 library build and parse Arrow arrays for small control-plane values. It would also commit to the shape of a C ABI we haven't designed yet. The data plane already uses the Arrow C structs, and that part carries over to a C ABI as is. ########## sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java: ########## @@ -0,0 +1,398 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.datasource; + +import org.apache.spark.annotation.Evolving; + +/** + * The binary interface between Spark and data sources implemented in native code, such as Rust + * or C++. + * <p> + * A native data source is a shared library that implements the {@code native} methods of this + * class as JNI functions, for example + * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each + * library into its own copy of this class, so any number of libraries that export the same + * functions can be loaded side by side. Spark adapts a native data source to Data Source V2, with + * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and + * limit pushdown. + * <p> + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of {@code my_source} is + * {@code spark_datasource_my_source}, such as {@code libspark_datasource_my_source.so} in + * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of + * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux, and in the + * {@code lib} directory of each installation prefix in {@code PATH}. A library can also be + * distributed in a native data source package: a zip file with the extension {@code .sparkpkg} + * that contains a manifest named {@code spark-native-datasource.json} and the library built for + * one or more platforms, added to a session with {@code spark.addArtifact} or under the paths in + * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL + * guide for details and examples. + * + * <h2>Lifecycle</h2> + * Spark creates a data source with {@link #createDataSource} on the driver for each schema + * inference, scan and write, and closes it right after. + * <ul> + * <li><b>Batch reads.</b> On the driver, Spark creates a reader with {@link #createReader}, + * pushes operations down by calling, each at most once and in this order, + * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions + * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, + * and closes it. Then it calls {@link #read} on the executors, in one task for each + * partition.</li> + * <li><b>Streaming reads.</b> On the driver, Spark creates a stream reader with + * {@link #createStreamReader} and keeps it until the streaming query stops. For each + * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with + * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then + * calls {@link #read} on the executors. Offsets are JSON strings defined by the library.</li> + * <li><b>Writes.</b> On the driver, once it computed the data to write, right before it starts + * the tasks of the write (or of a micro-batch), Spark creates a writer with + * {@link #createWriter} or {@link #createStreamWriter}, and serializes its state with + * {@link #serializeWriter}. If computing the data fails, Spark creates no writer. On the + * executors, each task creates a data writer with {@link #createDataWriter}, calls + * {@link #write} for each batch of at most {@code spark.sql.execution.arrow.maxRecordsPerBatch} + * rows, which also ends once it holds {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, + * and calls {@link #commitDataWriter}, or {@link #abortDataWriter} if the task or the + * commit fails. Then the driver calls {@link #commit} with the messages of the tasks if they + * all succeeded, or {@link #abort} otherwise, including when {@link #commit} fails, and closes + * the writer. A streaming write commits + * or aborts each micro-batch (epoch). A micro-batch streaming query creates a new stream writer + * for each micro-batch, and Spark closes the previous stream writer of the query when the query + * creates a new one. Spark is not told when a streaming query stops, so it closes the last + * stream writer of the query once the writer is no longer referenced, or when the query is + * restarted.</li> + * </ul> + * + * <h2>Conventions</h2> + * <ul> + * <li><b>Handles.</b> The {@code create} methods return an opaque, non-zero handle, typically a + * pointer. Spark releases each handle exactly once with the matching {@code close} method, or + * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers, though not + * always right after the last use of a stream writer (see Writes above). A reader or + * writer handle must not depend on the data source handle it was created from, because Spark + * closes the data source handle right after creating it.</li> + * <li><b>State.</b> Partitions, the state of readers and writers, and commit messages are byte + * arrays in a format defined by the library. Spark sends them to the executors and back.</li> + * <li><b>Errors.</b> To report an error, throw a Java exception, for example with the JNI + * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and return any value. Spark + * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.</li> + * <li><b>Arrow data.</b> Schemas and data are exchanged with the + * <a href="https://arrow.apache.org/docs/format/CDataInterface.html";>Arrow C data + * interface</a>. A {@code long} argument named {@code ...Address} is the address of a C struct + * allocated by Spark. For an output struct, the library initializes it and Spark takes + * ownership. For an input struct, the library may move it to take ownership; otherwise, Spark + * releases it after the call returns, so the library must not use it afterwards. The arrays + * that the library exports must have an offset of 0, at every level: Arrow Java ignores the + * offsets of the arrays it imports, so Spark rejects an array with another offset, such as a + * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with + * fewer values than its parent refers to, such as a column shorter than its batch. The buffers + * of an array must hold all of its values, as the C data interface requires: it does not carry + * the sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that + * is too short.</li> + * <li><b>Schemas.</b> Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: + * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, + * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is + * {@code timestamp[us]} without one. Data read from a library may also use + * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, and any names + * for the fields of the entries of a map, but not unsigned integers, {@code decimal256} or + * dictionary encoding.</li> + * <li><b>Options.</b> Option keys are passed in lower case, because Spark matches them + * case-insensitively.</li> + * <li><b>Optional methods.</b> A library only needs to export the functions for the operations + * it supports. If a function is missing, Spark treats the operation as unsupported, except for + * the pushdown methods and the {@code serialize} methods, which then push nothing down and + * return no state.</li> + * <li><b>Threads.</b> Spark calls the library from multiple threads, and calls on different + * handles may run concurrently. Calls on the same handle never run concurrently.</li> + * </ul> + * + * @since 4.4.0 + */ +@Evolving +public final class NativeBridge { + + /** The version of this interface. It changes when the interface changes incompatibly. */ + public static final int ABI_VERSION = 1; + + private NativeBridge() {} + + /** + * Loads the library that implements the native methods of this copy of the class. Spark calls + * it on the copy of the class it creates for the library. + */ + static void load(String path) { + System.load(path); + } + + /** Returns the version of this interface that the library implements, {@link #ABI_VERSION}. */ + public static native int abiVersion(); + + // --------------------------------------------------------------------------------------------- + // Data source (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Creates a data source with the options given by the user, such as + * {@code spark.read.option("key", "value")}. + * + * @param name the name of the data source, as listed in the package manifest. A library can + * implement several data sources. + * @param optionKeys the option keys, in lower case + * @param optionValues the option values, in the same order as the keys + * @return the data source handle + */ + public static native long createDataSource( + String name, String[] optionKeys, String[] optionValues); + + /** + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. + * + * @param schemaAddress the address of the {@code ArrowSchema} to initialize with a struct type + * whose fields are the columns + */ + public static native void schema(long dataSource, long schemaAddress); + + /** + * Creates a reader for a batch scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the reader handle + */ + public static native long createReader(long dataSource, long schemaAddress); + + /** + * Creates a reader for a micro-batch streaming scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the stream reader handle + */ + public static native long createStreamReader(long dataSource, long schemaAddress); + + /** + * Creates a writer for a batch write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data (save mode "overwrite") instead of + * appending to it + * @return the writer handle + */ + public static native long createWriter(long dataSource, long schemaAddress, boolean overwrite); + + /** + * Creates a writer for a streaming write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data in each micro-batch (output mode + * "complete") instead of appending to it + * @return the writer handle + */ + public static native long createStreamWriter( + long dataSource, long schemaAddress, boolean overwrite); + + /** Releases a data source handle. */ + public static native void closeDataSource(long dataSource); + + // --------------------------------------------------------------------------------------------- + // Batch reader (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Pushes predicates down to the reader. The predicates are combined with AND. + * <p> + * Each predicate is a JSON expression tree. An expression is one of: + * <ul> + * <li>{@code {"type": "column", "name": ["a", "b"]}}: a column. A nested field has a name + * part for each level.</li> + * <li>{@code {"type": "literal", "dataType": "bigint", "value": 1}}: a literal. The data type + * is a Spark SQL type name such as {@code int}, {@code string} or {@code decimal(10,2)}. The + * value is {@code null}, a JSON boolean or number for boolean and numeric types, a string + * for {@code string}, a base64 string for {@code binary}, a decimal string for + * {@code decimal}, the days since the epoch for {@code date}, and the microseconds since the + * epoch for {@code timestamp} and {@code timestamp_ntz}.</li> + * <li>{@code {"type": "function", "name": ">", "children": [...]}}: a predicate or function, + * with the names of {@code org.apache.spark.sql.connector.expressions.filter.Predicate} and + * {@code org.apache.spark.sql.connector.expressions.GeneralScalarExpression}, such as + * {@code =}, {@code <}, {@code IN}, {@code IS_NULL}, {@code STARTS_WITH}, {@code AND}, + * {@code OR} and {@code NOT}.</li> + * </ul> + * Spark does not pass predicates that it cannot express this way, nor predicates on columns + * that contain strings with a non-binary collation, such as {@code UTF8_LCASE}: strings are + * compared by their UTF-8 bytes. + * <p> + * A predicate that the reader evaluates completely must follow the semantics of Spark SQL. In Review Comment: Good point, done. Version 1 now passes only a fixed set of predicates (`NativePredicates.FUNCTION_NAMES`): `=`, `<>`, `<=>`, `<`, `<=`, `>`, `>=`, `IN`, `IS_NULL`, `IS_NOT_NULL`, `STARTS_WITH`, `ENDS_WITH`, `CONTAINS`, `AND`, `OR` and `NOT`, with columns and literals as the leaves. Any other `GeneralScalarExpression`, including `+ - * / %` in ANSI mode, `CASE_WHEN` and `UPPER`, makes the predicate unencodable, so Spark keeps evaluating it. The `pushPredicates` Javadoc lists the names and their children, and it also gives the forward-compatibility rule you suggested: later versions may pass other names without a new `ABI_VERSION`, so a reader must never report a predicate as fully evaluated if it contains a name or a data type it doesn't know. The encoding test checks every name in the list, and that arithmetic (alone or under `AND`) and `UPPER` are not passed. ########## sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java: ########## @@ -0,0 +1,398 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.datasource; + +import org.apache.spark.annotation.Evolving; + +/** + * The binary interface between Spark and data sources implemented in native code, such as Rust + * or C++. + * <p> + * A native data source is a shared library that implements the {@code native} methods of this + * class as JNI functions, for example + * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each + * library into its own copy of this class, so any number of libraries that export the same + * functions can be loaded side by side. Spark adapts a native data source to Data Source V2, with + * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and + * limit pushdown. + * <p> + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of {@code my_source} is + * {@code spark_datasource_my_source}, such as {@code libspark_datasource_my_source.so} in + * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of + * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux, and in the + * {@code lib} directory of each installation prefix in {@code PATH}. A library can also be + * distributed in a native data source package: a zip file with the extension {@code .sparkpkg} + * that contains a manifest named {@code spark-native-datasource.json} and the library built for + * one or more platforms, added to a session with {@code spark.addArtifact} or under the paths in + * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL + * guide for details and examples. + * + * <h2>Lifecycle</h2> + * Spark creates a data source with {@link #createDataSource} on the driver for each schema + * inference, scan and write, and closes it right after. + * <ul> + * <li><b>Batch reads.</b> On the driver, Spark creates a reader with {@link #createReader}, + * pushes operations down by calling, each at most once and in this order, + * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions + * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, + * and closes it. Then it calls {@link #read} on the executors, in one task for each + * partition.</li> + * <li><b>Streaming reads.</b> On the driver, Spark creates a stream reader with + * {@link #createStreamReader} and keeps it until the streaming query stops. For each + * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with + * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then + * calls {@link #read} on the executors. Offsets are JSON strings defined by the library.</li> + * <li><b>Writes.</b> On the driver, once it computed the data to write, right before it starts + * the tasks of the write (or of a micro-batch), Spark creates a writer with + * {@link #createWriter} or {@link #createStreamWriter}, and serializes its state with + * {@link #serializeWriter}. If computing the data fails, Spark creates no writer. On the + * executors, each task creates a data writer with {@link #createDataWriter}, calls + * {@link #write} for each batch of at most {@code spark.sql.execution.arrow.maxRecordsPerBatch} + * rows, which also ends once it holds {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, + * and calls {@link #commitDataWriter}, or {@link #abortDataWriter} if the task or the + * commit fails. Then the driver calls {@link #commit} with the messages of the tasks if they + * all succeeded, or {@link #abort} otherwise, including when {@link #commit} fails, and closes + * the writer. A streaming write commits + * or aborts each micro-batch (epoch). A micro-batch streaming query creates a new stream writer Review Comment: Done, in the Lifecycle section of the Javadoc: - Writes: several attempts of a task may write the same partition (after a failure or with speculation), each with its own data writer. Spark lets the attempts of a partition commit one at a time (the commit coordinator, since `useCommitCoordinator` keeps its default `true`) and passes `commit` the message of one committed attempt per partition, so the write consists of those attempts only. Other attempts are aborted, except on a lost executor, where a data writer is neither committed nor aborted and may even have committed before its message reached the driver. - Streaming writes: Spark may commit the same epoch more than once, for example when a query restarts after a failure, so committing an already committed epoch must have no effect, for exactly-once output. - Streaming reads: after a restart from the checkpoint, Spark creates a new stream reader and may call `streamPartitions` again for the last planned range, so the partitions of a range must read the same data each time. ########## sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java: ########## @@ -0,0 +1,398 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.datasource; + +import org.apache.spark.annotation.Evolving; + +/** + * The binary interface between Spark and data sources implemented in native code, such as Rust + * or C++. + * <p> + * A native data source is a shared library that implements the {@code native} methods of this + * class as JNI functions, for example + * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each + * library into its own copy of this class, so any number of libraries that export the same + * functions can be loaded side by side. Spark adapts a native data source to Data Source V2, with + * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and + * limit pushdown. + * <p> + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of {@code my_source} is + * {@code spark_datasource_my_source}, such as {@code libspark_datasource_my_source.so} in + * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of + * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux, and in the + * {@code lib} directory of each installation prefix in {@code PATH}. A library can also be + * distributed in a native data source package: a zip file with the extension {@code .sparkpkg} + * that contains a manifest named {@code spark-native-datasource.json} and the library built for + * one or more platforms, added to a session with {@code spark.addArtifact} or under the paths in + * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL + * guide for details and examples. + * + * <h2>Lifecycle</h2> + * Spark creates a data source with {@link #createDataSource} on the driver for each schema + * inference, scan and write, and closes it right after. + * <ul> + * <li><b>Batch reads.</b> On the driver, Spark creates a reader with {@link #createReader}, + * pushes operations down by calling, each at most once and in this order, + * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions + * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, + * and closes it. Then it calls {@link #read} on the executors, in one task for each + * partition.</li> + * <li><b>Streaming reads.</b> On the driver, Spark creates a stream reader with + * {@link #createStreamReader} and keeps it until the streaming query stops. For each + * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with + * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then + * calls {@link #read} on the executors. Offsets are JSON strings defined by the library.</li> + * <li><b>Writes.</b> On the driver, once it computed the data to write, right before it starts + * the tasks of the write (or of a micro-batch), Spark creates a writer with + * {@link #createWriter} or {@link #createStreamWriter}, and serializes its state with + * {@link #serializeWriter}. If computing the data fails, Spark creates no writer. On the + * executors, each task creates a data writer with {@link #createDataWriter}, calls + * {@link #write} for each batch of at most {@code spark.sql.execution.arrow.maxRecordsPerBatch} + * rows, which also ends once it holds {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, + * and calls {@link #commitDataWriter}, or {@link #abortDataWriter} if the task or the + * commit fails. Then the driver calls {@link #commit} with the messages of the tasks if they + * all succeeded, or {@link #abort} otherwise, including when {@link #commit} fails, and closes + * the writer. A streaming write commits + * or aborts each micro-batch (epoch). A micro-batch streaming query creates a new stream writer + * for each micro-batch, and Spark closes the previous stream writer of the query when the query + * creates a new one. Spark is not told when a streaming query stops, so it closes the last + * stream writer of the query once the writer is no longer referenced, or when the query is + * restarted.</li> + * </ul> + * + * <h2>Conventions</h2> + * <ul> + * <li><b>Handles.</b> The {@code create} methods return an opaque, non-zero handle, typically a + * pointer. Spark releases each handle exactly once with the matching {@code close} method, or + * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers, though not + * always right after the last use of a stream writer (see Writes above). A reader or + * writer handle must not depend on the data source handle it was created from, because Spark + * closes the data source handle right after creating it.</li> + * <li><b>State.</b> Partitions, the state of readers and writers, and commit messages are byte + * arrays in a format defined by the library. Spark sends them to the executors and back.</li> + * <li><b>Errors.</b> To report an error, throw a Java exception, for example with the JNI + * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and return any value. Spark + * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.</li> + * <li><b>Arrow data.</b> Schemas and data are exchanged with the + * <a href="https://arrow.apache.org/docs/format/CDataInterface.html";>Arrow C data + * interface</a>. A {@code long} argument named {@code ...Address} is the address of a C struct + * allocated by Spark. For an output struct, the library initializes it and Spark takes + * ownership. For an input struct, the library may move it to take ownership; otherwise, Spark + * releases it after the call returns, so the library must not use it afterwards. The arrays + * that the library exports must have an offset of 0, at every level: Arrow Java ignores the + * offsets of the arrays it imports, so Spark rejects an array with another offset, such as a + * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with + * fewer values than its parent refers to, such as a column shorter than its batch. The buffers + * of an array must hold all of its values, as the C data interface requires: it does not carry + * the sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that + * is too short.</li> + * <li><b>Schemas.</b> Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: + * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, + * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is + * {@code timestamp[us]} without one. Data read from a library may also use + * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, and any names + * for the fields of the entries of a map, but not unsigned integers, {@code decimal256} or + * dictionary encoding.</li> + * <li><b>Options.</b> Option keys are passed in lower case, because Spark matches them + * case-insensitively.</li> + * <li><b>Optional methods.</b> A library only needs to export the functions for the operations + * it supports. If a function is missing, Spark treats the operation as unsupported, except for + * the pushdown methods and the {@code serialize} methods, which then push nothing down and + * return no state.</li> + * <li><b>Threads.</b> Spark calls the library from multiple threads, and calls on different + * handles may run concurrently. Calls on the same handle never run concurrently.</li> + * </ul> + * + * @since 4.4.0 + */ +@Evolving +public final class NativeBridge { + + /** The version of this interface. It changes when the interface changes incompatibly. */ + public static final int ABI_VERSION = 1; + + private NativeBridge() {} + + /** + * Loads the library that implements the native methods of this copy of the class. Spark calls + * it on the copy of the class it creates for the library. + */ + static void load(String path) { + System.load(path); + } + + /** Returns the version of this interface that the library implements, {@link #ABI_VERSION}. */ + public static native int abiVersion(); + + // --------------------------------------------------------------------------------------------- + // Data source (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Creates a data source with the options given by the user, such as + * {@code spark.read.option("key", "value")}. + * + * @param name the name of the data source, as listed in the package manifest. A library can + * implement several data sources. + * @param optionKeys the option keys, in lower case + * @param optionValues the option values, in the same order as the keys + * @return the data source handle + */ + public static native long createDataSource( + String name, String[] optionKeys, String[] optionValues); + + /** + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. + * + * @param schemaAddress the address of the {@code ArrowSchema} to initialize with a struct type + * whose fields are the columns + */ + public static native void schema(long dataSource, long schemaAddress); + + /** + * Creates a reader for a batch scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the reader handle + */ + public static native long createReader(long dataSource, long schemaAddress); + + /** + * Creates a reader for a micro-batch streaming scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the stream reader handle + */ + public static native long createStreamReader(long dataSource, long schemaAddress); + + /** + * Creates a writer for a batch write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data (save mode "overwrite") instead of + * appending to it + * @return the writer handle + */ + public static native long createWriter(long dataSource, long schemaAddress, boolean overwrite); + + /** + * Creates a writer for a streaming write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data in each micro-batch (output mode + * "complete") instead of appending to it + * @return the writer handle + */ + public static native long createStreamWriter( + long dataSource, long schemaAddress, boolean overwrite); + + /** Releases a data source handle. */ + public static native void closeDataSource(long dataSource); + + // --------------------------------------------------------------------------------------------- + // Batch reader (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Pushes predicates down to the reader. The predicates are combined with AND. + * <p> + * Each predicate is a JSON expression tree. An expression is one of: + * <ul> + * <li>{@code {"type": "column", "name": ["a", "b"]}}: a column. A nested field has a name + * part for each level.</li> + * <li>{@code {"type": "literal", "dataType": "bigint", "value": 1}}: a literal. The data type + * is a Spark SQL type name such as {@code int}, {@code string} or {@code decimal(10,2)}. The + * value is {@code null}, a JSON boolean or number for boolean and numeric types, a string + * for {@code string}, a base64 string for {@code binary}, a decimal string for + * {@code decimal}, the days since the epoch for {@code date}, and the microseconds since the + * epoch for {@code timestamp} and {@code timestamp_ntz}.</li> + * <li>{@code {"type": "function", "name": ">", "children": [...]}}: a predicate or function, + * with the names of {@code org.apache.spark.sql.connector.expressions.filter.Predicate} and + * {@code org.apache.spark.sql.connector.expressions.GeneralScalarExpression}, such as + * {@code =}, {@code <}, {@code IN}, {@code IS_NULL}, {@code STARTS_WITH}, {@code AND}, + * {@code OR} and {@code NOT}.</li> + * </ul> + * Spark does not pass predicates that it cannot express this way, nor predicates on columns + * that contain strings with a non-binary collation, such as {@code UTF8_LCASE}: strings are + * compared by their UTF-8 bytes. + * <p> + * A predicate that the reader evaluates completely must follow the semantics of Spark SQL. In + * particular, for floating-point columns, NaN is equal to NaN and greater than any other value, + * including positive infinity, and {@code -0.0} is equal to {@code 0.0}, unlike IEEE 754 + * comparisons. For example, {@code d > 1.0} is true for a NaN value of {@code d}. Comparisons + * with null are null, so a row only matches if the predicate is true. + * + * @return for each predicate, whether the reader evaluates it completely, so that Spark does + * not have to evaluate it again. The reader can still use the other predicates to skip + * data. + */ + public static native boolean[] pushPredicates(long reader, String[] predicates); + + /** + * Pushes a LIMIT down to the reader. + * + * @return whether the reader returns at most {@code limit} rows in total. Spark applies the + * limit again after reading either way. + */ + public static native boolean pushLimit(long reader, int limit); + + /** + * Prunes the columns to read. + * + * @param columnNames the names of the top-level columns that the query needs + * @return whether {@link #read} returns exactly these columns, in this order. If not, it + * returns all the columns of the schema to read. + */ + public static native boolean pruneColumns(long reader, String[] columnNames); + + /** + * Plans the partitions of the scan. Spark reads each partition in a separate task. + * + * @return the partitions, each serialized in a format defined by the library + */ + public static native byte[][] partitions(long reader); + + /** + * Serializes the state of the reader that {@link #read} needs on the executors, such as the + * pushed predicates. Spark calls it after {@link #partitions}. + */ + public static native byte[] serializeReader(long reader); + + /** Releases a batch reader handle. */ + public static native void closeReader(long reader); + + // --------------------------------------------------------------------------------------------- + // Stream reader (driver) + // --------------------------------------------------------------------------------------------- + + /** Returns the offset to start from when the streaming query has no checkpoint, as JSON. */ + public static native String initialOffset(long streamReader); + + /** Returns the most recent offset available, as JSON. */ + public static native String latestOffset(long streamReader); + + /** + * Plans the partitions that read the data after {@code start} up to and including {@code end}. + * + * @return the partitions, each serialized in a format defined by the library + */ + public static native byte[][] streamPartitions(long streamReader, String start, String end); + + /** + * Serializes the state of the stream reader that {@link #read} needs on the executors. Spark + * calls it after each call to {@link #streamPartitions}. + */ + public static native byte[] serializeStreamReader(long streamReader); + + /** + * Tells the library that Spark has processed all the data up to and including {@code end}, + * and will not ask for it again. + */ + public static native void commitOffset(long streamReader, String end); + + /** Releases a stream reader handle, when the streaming query stops. */ + public static native void closeStreamReader(long streamReader); + + // --------------------------------------------------------------------------------------------- + // Partition reader (executor) + // --------------------------------------------------------------------------------------------- + + /** + * Reads a partition of a batch or streaming scan. + * + * @param readerState the state returned by {@link #serializeReader} or + * {@link #serializeStreamReader}, or an empty array if there is none + * @param partition a partition returned by {@link #partitions} or {@link #streamPartitions} + * @param streamAddress the address of the {@code ArrowArrayStream} to initialize. Each array + * of the stream is a struct array whose fields are the columns to read: + * the columns of the schema to read, or the pruned columns. Spark calls + * the stream from the thread of the task that reads the partition. + */ + public static native void read(byte[] readerState, byte[] partition, long streamAddress); Review Comment: Done. The `read` Javadoc now says that `get_next` may block (for example on network I/O), but Spark cannot interrupt native code. When the task is killed, Spark only stops reading once `get_next` returns, and then releases the stream. So a `get_next` that blocks forever keeps the task running unless the task reaper (`spark.task.reaper.enabled` and `spark.task.reaper.killTimeout`) kills the executor, and libraries should use timeouts. That is the version 1 contract. A cancellation function can come later as an optional function, and the new "Limitations" section of the docs page lists it. The Threads bullet now says that a handle may be released from a thread other than the ones that used it, such as a JVM `Cleaner` thread when Spark releases a handle that is no longer referenced (for example the last stream writer of a query). Readers, stream readers and writers are all `NativeHandle`s registered with the `Cleaner`. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,791 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.lang.ref.{Cleaner, WeakReference} +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, CDataDictionaryProvider, Data} +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.MapVector +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} + +import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.internal.Logging +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} +import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarDataSource, ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.unsafe.Platform +import org.apache.spark.util.Utils + +/** + * A [[ColumnarDataSource]] implemented by a native library, through + * `org.apache.spark.sql.datasource.NativeBridge`. + * + * @param location where the library is on the driver. The readers and the writers distribute it + * to the executors of the active session when they are created. + * @param name the name of the data source, as listed in the manifest of its package or in lower + * case for an installed library + */ +class NativeDataSource( + location: NativeLibraryLocation, + name: String, + options: CaseInsensitiveStringMap) + extends ColumnarDataSource { + import NativeCalls._ + + override def schema(): StructType = { + val library = NativeLibraries.get(location) + withDataSource(library, "infer the schema of") { dataSource => + call(name, "infer the schema of") { + optional(NativeArrow.importSchema(address => library.schema(dataSource, address))) + }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) + } + } + + override def reader(schema: StructType): ColumnarReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema)(address => library.createReader(dataSource, address)) + } + new NativeDataSourceReader( + executorLocation, name, schema, new NativeHandle(handle, library.closeReader)) + } + + override def streamReader(schema: StructType): ColumnarStreamReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema)(address => library.createStreamReader(dataSource, address)) + } + new NativeDataSourceStreamReader( + executorLocation, name, schema, new NativeHandle(handle, library.closeStreamReader)) + } + + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = new NativeHandle( + create(library, "plan a write to", "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource => + NativeArrow.exportSchema(schema) { address => + library.createWriter(dataSource, address, overwrite) + } + }, + library.closeWriter) + new NativeDataSourceWriter(executorLocation, name, handle, serializeWriter(library, handle)) + } + + override def streamWriter( + schema: StructType, + overwrite: Boolean, + queryId: String): ColumnarStreamWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = new NativeHandle( + create(library, "plan a write to", "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createStreamWriter(dataSource, address, overwrite) + } + }, + library.closeWriter) + val writer = new NativeDataSourceStreamWriter( + executorLocation, name, handle, serializeWriter(library, handle)) + NativeStreamWriters.register(queryId, handle) + writer + } + + /** Runs `f` with a new native data source, which is closed afterwards. */ + private def withDataSource[T](library: NativeLibrary, action: String)(f: Long => T): T = { + // The keys of a case-insensitive map are in lower case. + val entries = options.entrySet().asScala.toArray + val dataSource = call(name, "create") { + library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue)) + } + Utils.tryWithSafeFinally(f(dataSource)) { + call(name, action)(library.closeDataSource(dataSource)) + } + } + + /** Creates a reader or a writer with `f`, which fails if the library does not implement it. */ + private def create(library: NativeLibrary, action: String, unsupportedError: String)( + f: Long => Long): Long = { + withDataSource(library, action) { dataSource => + call(name, action)(optional(f(dataSource))).getOrElse { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) + } + } + } + + private def serializeWriter(library: NativeLibrary, handle: NativeHandle): Array[Byte] = { + try { + bytesOrEmpty(call(name, "plan a write to")(optional(library.serializeWriter(handle.get)))) + } catch { + case e: Throwable => + handle.close() + throw e + } + } +} + +/** Reads a native data source for a batch scan. */ +class NativeDataSourceReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarReader { + import NativeCalls._ + + // Set on the driver when the partitions are planned, and serialized to the executors. + private var readerState: Array[Byte] = Array.emptyByteArray + private var readSchema: StructType = schema + @transient private var plannedPartitions: Array[InputPartition] = _ + + @transient private lazy val library = NativeLibraries.get(location) + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + // Only the predicates that can be expressed in JSON are passed to the library. + val encoded = predicates.zipWithIndex.flatMap { case (predicate, i) => + NativePredicates.toJson(predicate, schema).map(_ -> i) + } + val pushed = if (encoded.isEmpty) { + Set.empty[Int] + } else { + call(name, "plan a scan of") { + optional(library.pushPredicates(handle.get, encoded.map(_._1))) + } match { + case Some(flags) if flags != null && flags.length == encoded.length => + encoded.map(_._2).zip(flags).collect { case (i, true) => i }.toSet + case Some(flags) => + throw error(name, "plan a scan of", s"pushPredicates returned " + + s"${Option(flags).map(_.length).getOrElse(0)} values for ${encoded.length} predicates.") + case None => Set.empty[Int] + } + } + predicates.zipWithIndex.collect { case (predicate, i) if !pushed.contains(i) => predicate } + } + + override def pushLimit(limit: Int): Boolean = { + call(name, "plan a scan of")(optional(library.pushLimit(handle.get, limit))).getOrElse(false) + } + + override def pruneColumns(requiredSchema: StructType): Boolean = { + val pruned = call(name, "plan a scan of") { + optional(library.pruneColumns(handle.get, requiredSchema.fieldNames)) + }.getOrElse(false) + if (pruned) { + readSchema = requiredSchema + } + pruned + } + + override def partitions(): Array[InputPartition] = synchronized { + if (plannedPartitions == null) { + try { + val partitions = call(name, "plan a scan of")(library.partitions(handle.get)) + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeReader(handle.get)))) + plannedPartitions = NativeInputPartition.from(partitions) + } finally { + // The executors only need the serialized state. + handle.close() + } + } + plannedPartitions + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, readSchema) + } +} + +/** Reads a native data source for a micro-batch streaming scan. */ +class NativeDataSourceStreamReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarStreamReader { + import NativeCalls._ + + // Set on the driver whenever the partitions of a micro-batch are planned, and serialized to + // the executors with the partitions of that micro-batch. + private var readerState: Array[Byte] = Array.emptyByteArray + + @transient private lazy val library = NativeLibraries.get(location) + + override def initialOffset(): String = { + call(name, "plan a scan of")(library.initialOffset(handle.get)) + } + + override def latestOffset(): String = { + call(name, "plan a scan of")(library.latestOffset(handle.get)) + } + + override def partitions(start: String, end: String): Array[InputPartition] = { + val partitions = call(name, "plan a scan of") { + library.streamPartitions(handle.get, start, end) + } + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeStreamReader(handle.get)))) + NativeInputPartition.from(partitions) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, schema) + } + + override def commit(end: String): Unit = { + call(name, "commit an offset of")(optional(library.commitOffset(handle.get, end))) + } + + override def stop(): Unit = handle.close() +} + +/** Writes to a native data source for a batch write. */ +class NativeDataSourceWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId = -1L) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + // Not closed if the commit fails: Spark then aborts the write, which closes it. + handle.close() + } + + override def abort(messages: Array[WriterCommitMessage]): Unit = { + Utils.tryWithSafeFinally { + call(name, "abort a write to") { + library.abort(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + }(handle.close()) + } +} + +/** + * Writes to a native data source for a streaming write. Its handle is closed when the streaming + * query creates its next stream writer, see [[NativeStreamWriters]]. + */ +class NativeDataSourceStreamWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarStreamWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId) + } + + override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "abort a write to") { + library.abort(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } +} + +/** A partition planned by a native library. */ +case class NativeInputPartition(index: Int, bytes: Array[Byte]) extends InputPartition + +object NativeInputPartition { + def from(partitions: Array[Array[Byte]]): Array[InputPartition] = { + Option(partitions).getOrElse(Array.empty).zipWithIndex.map { case (bytes, i) => + NativeInputPartition(i, Option(bytes).getOrElse(Array.emptyByteArray)) + } + } +} + +/** The message returned by the data writer of a task of a native library. */ +case class NativeWriterCommitMessage(bytes: Array[Byte]) extends WriterCommitMessage + +object NativeWriterCommitMessage { + def toBytes(messages: Array[WriterCommitMessage]): Array[Array[Byte]] = messages.map { + case message: NativeWriterCommitMessage => message.bytes + case null => null + case other => + throw SparkException.internalError(s"Unexpected commit message: ${other.getClass.getName}") + } +} + +/** + * Reads a partition of a native data source: each array of the Arrow stream exported by the + * library is checked, and imported into the Arrow vectors that back the returned batch. + */ +class NativePartitionReader( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + partition: Array[Byte], + schema: StructType) + extends PartitionReader[ColumnarBatch] { + import NativeCalls._ + + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) + // Read here rather than with `Data.importArrayStream`, to check each array before importing it. + private val stream = ArrowArrayStream.allocateNew(allocator) + private val dictionaries = new CDataDictionaryProvider() + private var root: VectorSchemaRoot = _ + private var batch: ColumnarBatch = _ + + try { + openStream() + root = VectorSchemaRoot.create(importSchema(), allocator) + checkSchema() + batch = new ColumnarBatch(root.getFieldVectors.asScala.map(columnVector).toArray) + } catch { + case e: Throwable => + close() + throw e + } + + private def openStream(): Unit = { + val library = NativeLibraries.get(location) + call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) + .getOrElse { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) + } + } + + private def importSchema(): Schema = { + val arrowSchema = ArrowSchema.allocateNew(allocator) + try { + call(name, "read")(stream.getSchema(arrowSchema)) + val imported = call(name, "read")(Data.importSchema(allocator, arrowSchema, dictionaries)) + new Schema( + imported.getFields.asScala.map(withMapEntryNames).asJava, imported.getCustomMetadata) + } finally { + arrowSchema.close() + } + } + + /** + * Names the fields of the entries of the maps `key` and `value`, which ArrowColumnVector looks + * up by name. Arrow does not fix these names: arrow-rs names them `keys` and `values`, for + * example. The arrays are imported by position, so the names do not change how they are read. + */ + private def withMapEntryNames(field: Field): Field = { + val children = field.getChildren.asScala.map(withMapEntryNames).toSeq + val renamed = (field.getType, children) match { + case (_: ArrowType.Map, Seq(entries)) if entries.getChildren.size == 2 => + val Seq(key, value) = entries.getChildren.asScala.toSeq + Seq(new Field(entries.getName, entries.getFieldType, Seq( + new Field(MapVector.KEY_NAME, key.getFieldType, key.getChildren), + new Field(MapVector.VALUE_NAME, value.getFieldType, value.getChildren)).asJava)) + case _ => children + } + new Field(field.getName, field.getFieldType, renamed.asJava) + } + + private def checkSchema(): Unit = { + def isDictionaryEncoded(field: Field): Boolean = { + field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) + } + root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + + "is not supported.") + } + // The Arrow types that convert to Spark types, but that ArrowColumnVector cannot read. + def unsupportedType(field: Field): Option[ArrowType] = field.getType match { + case decimal: ArrowType.Decimal if decimal.getBitWidth != 128 => Some(decimal) + case _ => field.getChildren.asScala.view.flatMap(unsupportedType).headOption + } + root.getSchema.getFields.asScala.foreach { field => + unsupportedType(field).foreach { arrowType => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"$arrowType, which is not supported.") + } + } + val actual = ArrowUtils.fromArrowSchema(root.getSchema) + // The schema to read as the library got it: Arrow does not carry everything, such as the + // collations of strings, the fields of intervals and user-defined types, which still apply to + // the data that Spark reads. + val expected = ArrowUtils.fromArrowSchema(ArrowUtils.toArrowSchema( + schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = false)) + val matches = actual.length == expected.length && actual.zip(expected).forall { + case (actualField, field) => + actualField.name == field.name && + DataType.equalsIgnoreNullability(actualField.dataType, field.dataType) + } + if (!matches) { + throw error(name, "read", s"The data has the schema ${actual.toDDL}, but the expected " + + s"schema is ${expected.toDDL}.") + } + } + + private def columnVector(vector: FieldVector): ColumnVector = { + try { + new ArrowColumnVector(vector) + } 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) + // A list view has no last offset: its lists refer to ranges of its elements in any order. + case _ => 0 Review Comment: Thanks, good catch. I went with rejecting them: Spark never exports list views, and a check for them would need a pass over all the views. A per-entry check would also have to define how the offsets and sizes of null views count. `checkSchema` now rejects a `ListView` at any level with "The column lv contains the Arrow type ListView, which is not supported.", as for `decimal256`, and the Schemas bullet of `NativeBridge` now lists list views among the unsupported types. The other variants are already safe: `ArrowUtils.fromArrowField` doesn't map large lists, large list views or fixed-size lists, so the schema comparison rejects them. The test library has a new `extra=listview` column, and "reject the Arrow data that Spark cannot read" covers it. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala: ########## @@ -0,0 +1,397 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.io.{File, FileInputStream, IOException} +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, StandardCopyOption} +import java.security.MessageDigest +import java.util.{HexFormat, Locale} +import java.util.concurrent.ConcurrentHashMap +import java.util.zip.ZipFile + +import scala.jdk.CollectionConverters._ +import scala.util.Try + +import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} +import com.google.common.cache.{Cache, CacheBuilder} + +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.PATH +import org.apache.spark.sql.Artifact +import org.apache.spark.sql.datasource.NativeBridge +import org.apache.spark.sql.errors.QueryExecutionErrors +import org.apache.spark.util.Utils + +/** + * The manifest of a native data source package, `spark-native-datasource.json`: + * {{{ + * { + * "abiVersion": 1, + * "dataSources": ["my_source"], + * "libraries": { + * "linux-x86_64": "linux-x86_64/libmy_source.so", + * "osx-aarch64": "osx-aarch64/libmy_source.dylib" + * } + * } + * }}} + * + * @param abiVersion the version of [[NativeBridge]] that the libraries implement + * @param dataSources the names of the data sources that the libraries implement + * @param libraries the path of the library in the package for each platform, as returned by + * [[NativePlatform.current]] + */ +case class NativeDataSourceManifest( + abiVersion: Int, + dataSources: Seq[String], + libraries: Map[String, String]) + +/** + * Where the library of a native data source is: in a [[NativeDataSourcePackage]], or installed in + * the native library path, as an [[InstalledNativeLibrary]]. It is sent to the executors, which + * load the library with [[NativeLibraries.get]]. + */ +sealed trait NativeLibraryLocation extends Serializable + +/** + * A native data source package: a zip file that contains a [[NativeDataSourceManifest]] and the + * libraries it lists. The SHA-256 checksum of the file identifies the package, so that an + * executor can check that it loads the same package as the driver. + * + * @param path the path of the package file on the driver + * @param artifactUUID the UUID of the session whose artifacts contain the package, if the session + * is isolated. Executors store the artifacts of such a session separately. + * @param distributedFileName the name of the copy of the package in the artifacts of the session, + * if the package is not itself an artifact of the session + */ +case class NativeDataSourcePackage( + path: String, + sha256: String, + manifest: NativeDataSourceManifest, + artifactUUID: Option[String] = None, + distributedFileName: Option[String] = None) extends NativeLibraryLocation { + def fileName: String = new File(path).getName +} + +/** + * A package file that is not a valid native data source package. + * + * @param error the error that explains why the package is invalid + * @param dataSources the data sources listed in the manifest of the package, if it can be read + */ +case class InvalidNativeDataSourcePackage(error: Throwable, dataSources: Seq[String]) + +object NativeDataSourcePackage extends Logging { + val FILE_EXTENSION: String = Artifact.NATIVE_DATA_SOURCE_PACKAGE_EXTENSION + val MANIFEST_NAME = "spark-native-datasource.json" + + private case class CacheKey(path: String, length: Long, lastModified: Long) + + private[ffi] type ReadResult = Either[InvalidNativeDataSourcePackage, NativeDataSourcePackage] + + // Bounded, because each session stores its artifacts in a directory of its own, so a + // long-running server where many sessions add packages reads many different files. + private val MAX_CACHED_PACKAGES = 1000 + + private val packages: Cache[CacheKey, ReadResult] = + CacheBuilder.newBuilder().maximumSize(MAX_CACHED_PACKAGES).build[CacheKey, ReadResult]() + private val mapper = new ObjectMapper() + + /** Reads the package in the given file. The result is cached until the file changes. */ + def read(file: File): NativeDataSourcePackage = tryRead(file).fold(p => throw p.error, identity) + + /** + * Reads the package in the given file, or returns why it is invalid. The result is cached until + * the file changes, and an invalid package is logged when it is read. + */ + private[ffi] def tryRead(file: File): ReadResult = { + val canonical = try { + file.getCanonicalFile + } catch { + case e: IOException => + return Left(InvalidNativeDataSourcePackage( + invalidManifestError(file, s"The package cannot be read: $e", e), Nil)) + } + val key = CacheKey(canonical.getPath, canonical.length(), canonical.lastModified()) + Option(packages.getIfPresent(key)).getOrElse { + val result = parse(canonical) + result.swap.foreach { invalid => + logWarning(log"Invalid native data source package ${MDC(PATH, canonical)}.", + invalid.error) + } + packages.put(key, result) + result + } + } + + /** + * Forgets the packages read from the files under the given directory, such as the artifacts of + * a session that is cleaned up. + */ + def forgetPackagesIn(dir: File): Unit = { + val prefix = dir.getCanonicalPath + File.separator + packages.asMap().keySet().removeIf(_.path.startsWith(prefix)) + } + + /** Returns whether a package read from the file at the given canonical path is cached. */ + private[ffi] def isCached(canonicalPath: String): Boolean = { + packages.asMap().keySet().asScala.exists(_.path == canonicalPath) + } + + private def invalidManifestError(file: File, reason: String, cause: Throwable = null) = { + QueryExecutionErrors.invalidNativeDataSourcePackageError( + file.getPath, "INVALID_MANIFEST", Map("reason" -> reason), cause) + } + + private def parse(file: File): ReadResult = { + def invalid(reason: String, dataSources: Seq[String] = Nil, cause: Throwable = null) = { + Left(InvalidNativeDataSourcePackage(invalidManifestError(file, reason, cause), dataSources)) + } + try { + val manifest = Utils.tryWithResource(new ZipFile(file)) { zip => + Option(zip.getEntry(MANIFEST_NAME)) match { + case None => invalid(s"The package does not contain the manifest $MANIFEST_NAME.") + case Some(entry) => + val json = Utils.tryWithResource(zip.getInputStream(entry)) { in => + new String(in.readAllBytes(), StandardCharsets.UTF_8) + } + parseManifest(json) match { + case Left(reason) => invalid(reason) + case Right(manifest) => + manifest.libraries.values.find(zip.getEntry(_) == null) match { + case Some(library) => invalid( + s"The package does not contain the library $library.", manifest.dataSources) + case None => Right(manifest) + } + } + } + } + manifest.map(NativeDataSourcePackage(file.getPath, sha256(file), _)) + } catch { + case e: IOException => invalid(s"The package is not a valid zip file: $e", cause = e) + } + } + + private[ffi] def parseManifest(json: String): Either[String, NativeDataSourceManifest] = { + val root: JsonNode = try { + mapper.readTree(json) + } catch { + case e: IOException => return Left(s"The manifest is not valid JSON: ${e.getMessage}") + } + if (root == null || !root.isObject) { + return Left("The manifest must be a JSON object.") + } + val abiVersion = root.get("abiVersion") + if (abiVersion == null || !abiVersion.isInt) { + return Left("The manifest must have an integer field 'abiVersion'.") + } + val dataSources = root.get("dataSources") + if (dataSources == null || !dataSources.isArray || dataSources.isEmpty || + !dataSources.elements().asScala.forall(n => n.isTextual && n.asText().nonEmpty)) { + return Left("The manifest must have a field 'dataSources' with a non-empty array of " + + "data source names.") + } + val libraries = root.get("libraries") + if (libraries == null || !libraries.isObject || libraries.isEmpty || + !libraries.properties().asScala.forall(e => e.getValue.isTextual)) { + return Left("The manifest must have a field 'libraries' with an object that maps " + + "platforms to the paths of the libraries in the package.") + } + Right(NativeDataSourceManifest( + abiVersion.asInt(), + dataSources.elements().asScala.map(_.asText()).toSeq, + libraries.properties().asScala.map(e => e.getKey -> e.getValue.asText()).toMap)) + } + + private 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. + * + * @param path the path of the library on the driver + */ +case class InstalledNativeLibrary(path: String) extends NativeLibraryLocation + +/** The platform names used in [[NativeDataSourceManifest.libraries]]. */ +object NativePlatform { + /** The platform of this JVM, such as `linux-x86_64`, `linux-aarch64` or `osx-aarch64`. */ + lazy val current: String = name(System.getProperty("os.name"), System.getProperty("os.arch")) + + private[ffi] def name(osName: String, osArch: String): String = { + val os = osName.toLowerCase(Locale.ROOT) match { + case n if n.startsWith("linux") => "linux" + case n if n.startsWith("mac") || n.startsWith("darwin") => "osx" + case n if n.startsWith("windows") => "windows" + case n => n.replaceAll("[^a-z0-9]", "") + } + val arch = osArch.toLowerCase(Locale.ROOT) match { + case "amd64" | "x86_64" | "x64" => "x86_64" + case "aarch64" | "arm64" => "aarch64" + case a => a + } + s"$os-$arch" + } +} + +/** + * The native data source libraries loaded in this JVM. A library is loaded once, on first use, + * after it is extracted from its package if needed, and stays loaded until the JVM exits. + */ +object NativeLibraries extends Logging { + // The libraries of the packages, by the checksum of the package, and the installed libraries, + // by their canonical path: the JVM loads a library file in a single class loader. A library + // that is loaded but rejected stays loaded, so its error is kept: it cannot be loaded again. + private val packagedLibraries = new ConcurrentHashMap[String, Try[NativeLibrary]]() + private val installedLibraries = new ConcurrentHashMap[String, Try[NativeLibrary]]() + + private lazy val extractionDir: File = Utils.createTempDir(namePrefix = "native-datasources") + + /** Returns the library at the given location, and loads it if needed. */ + private[ffi] def get(location: NativeLibraryLocation): NativeLibrary = location match { + case pkg: NativeDataSourcePackage => + getOrLoad(packagedLibraries, pkg.sha256)(loadPackaged(pkg)) + case installed: InstalledNativeLibrary => + val file = locate(installed).getCanonicalFile + getOrLoad(installedLibraries, file.getPath)(loadInstalled(file)) + } + + private def getOrLoad(libraries: ConcurrentHashMap[String, Try[NativeLibrary]], key: String)( + load: => Try[NativeLibrary]): NativeLibrary = { + val library = libraries.get(key) + (if (library != null) library else libraries.computeIfAbsent(key, _ => load)).get + } + + 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. + * + * @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 { + case e: UnsatisfiedLinkError => throw invalidLibrary(e.getMessage, e) + } + Try { + val abiVersion = try { + library.abiVersion() + } catch { + case e: UnsatisfiedLinkError => + throw invalidLibrary("It does not implement NativeBridge.abiVersion.", e) + } + if (abiVersion != NativeBridge.ABI_VERSION) { + throw unsupportedAbiVersion( + Map("version" -> abiVersion.toString, "supported" -> NativeBridge.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: 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) + if (file.isFile) { + file + } else { + NativeDataSourceRegistry.findInstalledLibrary(file.getName).getOrElse { Review Comment: Done, the first option. `InstalledNativeLibrary` now carries the SHA-256 checksum of the driver's file, computed at lookup and cached by (canonical path, length, modification time), like packages. An executor checks the library at the driver's path, and then every library with the same file name in its native library path, and loads only one with the same checksum. Otherwise the query fails with `NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND`, which now includes the checksum and asks for the same build on every node. I made it an error rather than a warning, because a different build would get state serialized for another build, which is undefined behavior. The docs page says so in "Installed Libraries". The test "find installed libraries automatically" covers a driver checksum that matches neither the library at the driver's path nor the one in the library path. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala: ########## @@ -0,0 +1,273 @@ +/* + * 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.datasource.NativeBridge +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_<name>`, such as + * `libspark_datasource_my_source.so` on Linux, implements the data source `<name>`. 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 == NativeBridge.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" -> NativeBridge.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) + findInstalledLibrary(libraryFileName(dataSource)) + .map(file => dataSource -> InstalledNativeLibrary(file.getPath)) + } + + /** Finds the library with the given file name in the native library path of this node. */ + private[ffi] def findInstalledLibrary(fileName: String): Option[File] = { + librarySearchPath(System.getProperty("java.library.path"), System.getenv("PATH")) + .map(new File(_, fileName)) + .find(_.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) Review Comment: Thanks for raising it. We'd like to keep this rule. It is what makes `make install`, `cmake --install`, `cargo cinstall`, Homebrew (`/opt/homebrew/lib`) and conda/venv installs work without configuration. Those directories are not in the default `java.library.path` on Linux or macOS, and automatic discovery of installed libraries was a goal of this API, the same way Python data sources are found in the Python path. On the concerns: - **Different environments on the executors**: the change for your other comment (the SHA-256 check of installed libraries) addresses this. An executor loads only a library with the driver's checksum, from the driver's path or its own native library path, so a different PATH or a different build on a node fails with `NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND` instead of silently loading another build. The docs now say that the executors resolve the path from their own environment, including their own `PATH`. - **What gets picked up**: only a file named exactly `libspark_datasource_<name>.so` (`.dylib`, `.dll`), and only when no Java or Python data source has that name and no package of the session provides it. A stray `~/.cargo/lib` in the search path only matters if it contains such a file. - **Hard to remove**: the rule is documented as the lookup order of the native library path, and it is a fallback after `java.library.path`. We can still narrow it before 4.4.0 if this turns out to be a problem in practice. `spark.driver.extraLibraryPath` / `spark.executor.extraLibraryPath` (part of `java.library.path`) remain the explicit setting, and they come first. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
