tdcmeehan commented on code in PR #59313: URL: https://github.com/apache/spark/pull/59313#discussion_r4260121143
########## sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java: ########## @@ -0,0 +1,442 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.datasource; + +import org.apache.spark.annotation.Evolving; + +/** + * The binary interface between Spark and data sources implemented in native code, such as Rust + * or C++. + * <p> + * A native data source is a shared library that implements the {@code native} methods of this + * class as JNI functions, for example + * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each + * library into its own copy of this class, so any number of libraries that export the same + * functions can be loaded side by side. Spark adapts a native data source to Data Source V2, with + * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and + * limit pushdown. + * <p> + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of {@code my_source} is + * {@code spark_datasource_my_source}, such as {@code libspark_datasource_my_source.so} in + * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of + * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux, and in the + * {@code lib} directory of each installation prefix in {@code PATH}. A library can also be + * distributed in a native data source package: a zip file with the extension {@code .sparkpkg} + * that contains a manifest named {@code spark-native-datasource.json} and the library built for + * one or more platforms, added to a session with {@code spark.addArtifact} or under the paths in + * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL + * guide for details and examples. + * + * <h2>Lifecycle</h2> + * Spark creates a data source with {@link #createDataSource} on the driver for each schema + * inference, scan and write, and closes it right after. + * <ul> + * <li><b>Batch reads.</b> On the driver, Spark creates a reader with {@link #createReader}, + * pushes operations down by calling, each at most once and in this order, + * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions + * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, + * and closes it. Then it calls {@link #read} on the executors, in one task for each + * partition.</li> + * <li><b>Streaming reads.</b> On the driver, Spark creates a stream reader with + * {@link #createStreamReader} and keeps it until the streaming query stops. For each + * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with + * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then + * calls {@link #read} on the executors. Offsets are JSON strings defined by the library. When + * a streaming query restarts, from its checkpoint, Spark creates a new stream reader and may + * call {@link #streamPartitions} again for the last range that it planned before, so the + * partitions of a range must read the same data each time.</li> + * <li><b>Writes.</b> On the driver, once it computed the data to write, right before it starts + * the tasks of the write (or of a micro-batch), Spark creates a writer with {@link #createWriter} + * or {@link #createStreamWriter}, and serializes its state with {@link #serializeWriter}. If + * computing the data fails, Spark creates no writer. On the executors, each task creates a data + * writer with {@link #createDataWriter}, calls {@link #write} for each batch of at most + * {@code spark.sql.execution.arrow.maxRecordsPerBatch} rows, which also ends once it holds + * {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, and calls {@link #commitDataWriter}, + * or {@link #abortDataWriter} if the task or the commit fails. Then the driver calls + * {@link #commit} with the messages of the tasks if they all succeeded, or {@link #abort} + * otherwise, including when {@link #commit} fails, and closes the writer. Several attempts of a + * task may write the same partition, after a failure or with speculative execution, each with a + * data writer of its own. Spark lets the attempts of a partition commit one at a time, and passes + * to {@link #commit} the message of one committed attempt for each partition, so the write + * consists of the data of these attempts only. It aborts the data writers of the other attempts, + * except on an executor that is lost, where a data writer is neither committed nor aborted, and + * may even have committed before its message reached the driver. A streaming write commits or + * aborts each micro-batch (epoch). Spark may commit the same epoch more than once, for example + * when a streaming query restarts after a failure, so committing an epoch that is already + * committed must have no effect, for exactly-once output. A micro-batch streaming query creates a + * new stream writer for each micro-batch, and Spark closes the previous stream writer of the + * query when the query creates a new one. Spark is not told when a streaming query stops, so it + * closes the last stream writer of the query once the writer is no longer referenced, or when the + * query is restarted.</li> + * </ul> + * + * <h2>Conventions</h2> + * <ul> + * <li><b>Handles.</b> The {@code create} methods return an opaque, non-zero handle, typically a + * pointer. Spark releases each handle exactly once with the matching {@code close} method, or + * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers, though not + * always right after the last use of a stream writer (see Writes above). A reader or + * writer handle must not depend on the data source handle it was created from, because Spark + * closes the data source handle right after creating it.</li> + * <li><b>State.</b> Partitions, the state of readers and writers, and commit messages are byte + * arrays in a format defined by the library. Spark sends them to the executors and back.</li> + * <li><b>Errors.</b> To report an error, throw a Java exception, for example with the JNI + * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and return any value. Spark + * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.</li> + * <li><b>Arrow data.</b> Schemas and data are exchanged with the + * <a href="https://arrow.apache.org/docs/format/CDataInterface.html";>Arrow C data + * interface</a>. A {@code long} argument named {@code ...Address} is the address of a C struct + * allocated by Spark. For an output struct, the library initializes it and Spark takes + * ownership. For an input struct, the library may move it to take ownership; otherwise, Spark + * releases it after the call returns, so the library must not use it afterwards. The arrays + * that the library exports must have an offset of 0, at every level: Arrow Java ignores the + * offsets of the arrays it imports, so Spark rejects an array with another offset, such as a + * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with + * fewer values than its parent refers to, such as a column shorter than its batch. The buffers + * of an array must hold all of its values, as the C data interface requires: it does not carry + * the sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that + * is too short.</li> + * <li><b>Schemas.</b> Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: + * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, + * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is + * {@code timestamp[us]} without one. Data read from a library may also use + * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, and any names + * for the fields of the entries of a map, but not unsigned integers, {@code decimal256}, list + * views or dictionary encoding.</li> + * <li><b>Options.</b> Option keys are passed in lower case, because Spark matches them + * case-insensitively.</li> + * <li><b>Optional methods.</b> A library only needs to export the functions for the operations + * it supports. If a function is missing, Spark treats the operation as unsupported, except for + * the pushdown methods and the {@code serialize} methods, which then push nothing down and + * return no state.</li> + * <li><b>Threads.</b> Spark calls the library from multiple threads, and calls on different + * handles may run concurrently. Calls on the same handle never run concurrently. A handle may + * be released from a thread other than the threads that used it, such as a JVM + * {@code Cleaner} thread when Spark releases a handle once it is no longer referenced, for + * example the last stream writer of a streaming query.</li> + * </ul> + * + * <h2>Compatibility</h2> + * The JVM binds the functions of a library by their names only, so the name of this class, the + * names of its methods and their signatures are the interface. Within an {@link #ABI_VERSION}: + * <ul> + * <li>the signature and the semantics of an existing method never change;</li> + * <li>new methods may be added, only as optional methods, so that a library that does not + * export them keeps working as before;</li> + * <li>new values may appear where this documentation says so, such as the names of pushed + * predicates.</li> + * </ul> + * An incompatible change uses a new {@link #ABI_VERSION}, with new names or a new class. A + * library reports the version it implements with {@link #abiVersion}, and Spark rejects a library + * of a version it does not support. Spark supports version 1 at least in all the releases of + * Spark 4 from 4.4.0. + * + * @since 4.4.0 + */ +@Evolving +public final class NativeBridge { + + /** + * The version of this interface. It changes when the interface changes incompatibly: see + * Compatibility above. + */ + public static final int ABI_VERSION = 1; + + private NativeBridge() {} + + /** + * Loads the library that implements the native methods of this copy of the class. Spark calls + * it on the copy of the class it creates for the library. + */ + static void load(String path) { + System.load(path); + } + + /** Returns the version of this interface that the library implements, {@link #ABI_VERSION}. */ + public static native int abiVersion(); + + // --------------------------------------------------------------------------------------------- + // Data source (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Creates a data source with the options given by the user, such as + * {@code spark.read.option("key", "value")}. + * + * @param name the name of the data source, as listed in the package manifest. A library can + * implement several data sources. + * @param optionKeys the option keys, in lower case + * @param optionValues the option values, in the same order as the keys + * @return the data source handle + */ + public static native long createDataSource( + String name, String[] optionKeys, String[] optionValues); + + /** + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. + * + * @param schemaAddress the address of the {@code ArrowSchema} to initialize with a struct type + * whose fields are the columns + */ + public static native void schema(long dataSource, long schemaAddress); + + /** + * Creates a reader for a batch scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the reader handle + */ + public static native long createReader(long dataSource, long schemaAddress); + + /** + * Creates a reader for a micro-batch streaming scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the stream reader handle + */ + public static native long createStreamReader(long dataSource, long schemaAddress); + + /** + * Creates a writer for a batch write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data (save mode "overwrite") instead of + * appending to it + * @return the writer handle + */ + public static native long createWriter(long dataSource, long schemaAddress, boolean overwrite); + + /** + * Creates a writer for a streaming write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data in each micro-batch (output mode + * "complete") instead of appending to it + * @return the writer handle + */ + public static native long createStreamWriter( + long dataSource, long schemaAddress, boolean overwrite); + + /** Releases a data source handle. */ + public static native void closeDataSource(long dataSource); + + // --------------------------------------------------------------------------------------------- + // Batch reader (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Pushes predicates down to the reader. The predicates are combined with AND. + * <p> + * Each predicate is a JSON expression tree. An expression is one of: + * <ul> + * <li>{@code {"type": "column", "name": ["a", "b"]}}: a column. A nested field has a name + * part for each level.</li> + * <li>{@code {"type": "literal", "dataType": "bigint", "value": 1}}: a literal. The data type + * is a Spark SQL type name such as {@code int}, {@code string} or {@code decimal(10,2)}. The + * value is {@code null}, a JSON boolean or number for boolean and numeric types, a string + * for {@code string}, a base64 string for {@code binary}, a decimal string for + * {@code decimal}, the days since the epoch for {@code date}, and the microseconds since the + * epoch for {@code timestamp} and {@code timestamp_ntz}.</li> + * <li>{@code {"type": "function", "name": ">", "children": [...]}}: a predicate, with the + * name it has in {@code org.apache.spark.sql.connector.expressions.filter.Predicate}. The + * names are {@code =}, {@code <>}, {@code <=>} (equal, with null equal to null), {@code <}, + * {@code <=}, {@code >} and {@code >=} with two children, {@code IN} with the value to look + * for and then the values to look in, {@code IS_NULL} and {@code IS_NOT_NULL} with one + * child, {@code STARTS_WITH}, {@code ENDS_WITH} and {@code CONTAINS} with the string and the + * string to look for, {@code AND} and {@code OR} with two predicates, and {@code NOT} with + * one. The children of the other functions are columns, literals or predicates.</li> + * </ul> + * Spark does not pass predicates that it cannot express this way, such as predicates that use + * arithmetic or other functions, nor predicates on columns that contain strings with a + * non-binary collation, such as {@code UTF8_LCASE}: strings are compared by their UTF-8 bytes. + * A later version of Spark may pass predicates with other names without changing + * {@link #ABI_VERSION}, so a reader must not report a predicate as evaluated completely if it + * contains a name that the reader does not know, or a data type that it does not support. + * <p> + * A predicate that the reader evaluates completely must follow the semantics of Spark SQL. In + * particular, for floating-point columns, NaN is equal to NaN and greater than any other value, + * including positive infinity, and {@code -0.0} is equal to {@code 0.0}, unlike IEEE 754 + * comparisons. For example, {@code d > 1.0} is true for a NaN value of {@code d}. Comparisons + * with null are null, so a row only matches if the predicate is true. + * + * @return for each predicate, whether the reader evaluates it completely, so that Spark does + * not have to evaluate it again. The reader can still use the other predicates to skip + * data. Review Comment: **Spark's comparison semantics shouldn't rest on every library** Following up on [4] (https://github.com/apache/spark/pull/59313#discussion_r4235638438). The Javadoc now states Spark's semantics, but each library still has to reproduce them in its own engine, and the obvious engines differ by default. arrow-rs, which the Rust example uses, compares floats by IEEE 754 totalOrder, which "treats positive and negative zeros as different" (https://docs.rs/arrow-ord/latest/arrow_ord/cmp/fn.eq.html). So a library that evaluates `d = 0.0` with it drops the `-0.0` rows that Spark returns. Null literals are pushed too, so three-valued logic under `NOT` and `IN` is also each library's to get right. Strings are a third case, as @cloud-fan shows on L269 (https://github.com/apache/spark/pull/59313#discussion_r4260023398): a C++ library that reads a literal with `GetStringUTFChars` gets modified UTF-8, so a byte comparison misses rows that Spark returns. This is hard even inside Spark: SPARK-58782 was a null literal that the JDBC pushdown compiled to the stri ng `'null'`. This can go wrong in two ways: - A library returns extra rows. That only matters because Spark drops its own filter when the library reports a predicate as evaluated completely. - A library drops rows, by filtering or skipping under other semantics. Spark can't see that. Could version 1: 1. Keep evaluating every pushed predicate in Spark, even when the library reports it as evaluated completely? That removes the first case for the cost of a filter over rows already read. Trusting a library could come later, as an opt-in. 2. State the remaining rule here: a reader must never drop or skip a row for which a predicate is true under Spark SQL. 3. Publish that rule as test cases (predicate JSON, input, expected rows) covering NaN, `-0.0`, nulls under `NOT` and `IN`, and `<=>`, so that each library can test itself against Spark? Velox avoids the problem by handing connectors filter objects (`common::Filter`, https://github.com/facebookincubator/velox/blob/main/velox/type/Filter.h) that answer both the row test (`testDouble`) and the min/max test used for skipping (`testDoubleRange`), with NaN handled once in the filter, so a connector doesn't implement comparisons itself. A small Spark-provided helper for these 16 predicates, as a header-only C file or a Rust crate, could do the same here. ########## sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java: ########## @@ -0,0 +1,442 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.datasource; + +import org.apache.spark.annotation.Evolving; + +/** + * The binary interface between Spark and data sources implemented in native code, such as Rust + * or C++. + * <p> + * A native data source is a shared library that implements the {@code native} methods of this + * class as JNI functions, for example + * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each + * library into its own copy of this class, so any number of libraries that export the same + * functions can be loaded side by side. Spark adapts a native data source to Data Source V2, with + * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and + * limit pushdown. + * <p> + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of {@code my_source} is + * {@code spark_datasource_my_source}, such as {@code libspark_datasource_my_source.so} in + * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of + * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux, and in the + * {@code lib} directory of each installation prefix in {@code PATH}. A library can also be + * distributed in a native data source package: a zip file with the extension {@code .sparkpkg} + * that contains a manifest named {@code spark-native-datasource.json} and the library built for + * one or more platforms, added to a session with {@code spark.addArtifact} or under the paths in + * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL + * guide for details and examples. + * + * <h2>Lifecycle</h2> + * Spark creates a data source with {@link #createDataSource} on the driver for each schema + * inference, scan and write, and closes it right after. + * <ul> + * <li><b>Batch reads.</b> On the driver, Spark creates a reader with {@link #createReader}, + * pushes operations down by calling, each at most once and in this order, + * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions + * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, + * and closes it. Then it calls {@link #read} on the executors, in one task for each + * partition.</li> + * <li><b>Streaming reads.</b> On the driver, Spark creates a stream reader with + * {@link #createStreamReader} and keeps it until the streaming query stops. For each + * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with + * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then + * calls {@link #read} on the executors. Offsets are JSON strings defined by the library. When + * a streaming query restarts, from its checkpoint, Spark creates a new stream reader and may + * call {@link #streamPartitions} again for the last range that it planned before, so the + * partitions of a range must read the same data each time.</li> + * <li><b>Writes.</b> On the driver, once it computed the data to write, right before it starts + * the tasks of the write (or of a micro-batch), Spark creates a writer with {@link #createWriter} + * or {@link #createStreamWriter}, and serializes its state with {@link #serializeWriter}. If + * computing the data fails, Spark creates no writer. On the executors, each task creates a data + * writer with {@link #createDataWriter}, calls {@link #write} for each batch of at most + * {@code spark.sql.execution.arrow.maxRecordsPerBatch} rows, which also ends once it holds + * {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, and calls {@link #commitDataWriter}, + * or {@link #abortDataWriter} if the task or the commit fails. Then the driver calls + * {@link #commit} with the messages of the tasks if they all succeeded, or {@link #abort} + * otherwise, including when {@link #commit} fails, and closes the writer. Several attempts of a + * task may write the same partition, after a failure or with speculative execution, each with a + * data writer of its own. Spark lets the attempts of a partition commit one at a time, and passes + * to {@link #commit} the message of one committed attempt for each partition, so the write + * consists of the data of these attempts only. It aborts the data writers of the other attempts, + * except on an executor that is lost, where a data writer is neither committed nor aborted, and + * may even have committed before its message reached the driver. A streaming write commits or + * aborts each micro-batch (epoch). Spark may commit the same epoch more than once, for example + * when a streaming query restarts after a failure, so committing an epoch that is already + * committed must have no effect, for exactly-once output. A micro-batch streaming query creates a + * new stream writer for each micro-batch, and Spark closes the previous stream writer of the + * query when the query creates a new one. Spark is not told when a streaming query stops, so it + * closes the last stream writer of the query once the writer is no longer referenced, or when the + * query is restarted.</li> + * </ul> + * + * <h2>Conventions</h2> + * <ul> + * <li><b>Handles.</b> The {@code create} methods return an opaque, non-zero handle, typically a + * pointer. Spark releases each handle exactly once with the matching {@code close} method, or + * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers, though not + * always right after the last use of a stream writer (see Writes above). A reader or + * writer handle must not depend on the data source handle it was created from, because Spark + * closes the data source handle right after creating it.</li> + * <li><b>State.</b> Partitions, the state of readers and writers, and commit messages are byte + * arrays in a format defined by the library. Spark sends them to the executors and back.</li> + * <li><b>Errors.</b> To report an error, throw a Java exception, for example with the JNI + * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and return any value. Spark + * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.</li> + * <li><b>Arrow data.</b> Schemas and data are exchanged with the + * <a href="https://arrow.apache.org/docs/format/CDataInterface.html";>Arrow C data + * interface</a>. A {@code long} argument named {@code ...Address} is the address of a C struct + * allocated by Spark. For an output struct, the library initializes it and Spark takes + * ownership. For an input struct, the library may move it to take ownership; otherwise, Spark + * releases it after the call returns, so the library must not use it afterwards. The arrays + * that the library exports must have an offset of 0, at every level: Arrow Java ignores the + * offsets of the arrays it imports, so Spark rejects an array with another offset, such as a + * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with + * fewer values than its parent refers to, such as a column shorter than its batch. The buffers + * of an array must hold all of its values, as the C data interface requires: it does not carry + * the sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that + * is too short.</li> + * <li><b>Schemas.</b> Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: + * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, + * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is + * {@code timestamp[us]} without one. Data read from a library may also use + * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, and any names + * for the fields of the entries of a map, but not unsigned integers, {@code decimal256}, list + * views or dictionary encoding.</li> + * <li><b>Options.</b> Option keys are passed in lower case, because Spark matches them + * case-insensitively.</li> + * <li><b>Optional methods.</b> A library only needs to export the functions for the operations + * it supports. If a function is missing, Spark treats the operation as unsupported, except for + * the pushdown methods and the {@code serialize} methods, which then push nothing down and + * return no state.</li> + * <li><b>Threads.</b> Spark calls the library from multiple threads, and calls on different + * handles may run concurrently. Calls on the same handle never run concurrently. A handle may + * be released from a thread other than the threads that used it, such as a JVM + * {@code Cleaner} thread when Spark releases a handle once it is no longer referenced, for + * example the last stream writer of a streaming query.</li> + * </ul> + * + * <h2>Compatibility</h2> + * The JVM binds the functions of a library by their names only, so the name of this class, the + * names of its methods and their signatures are the interface. Within an {@link #ABI_VERSION}: + * <ul> + * <li>the signature and the semantics of an existing method never change;</li> + * <li>new methods may be added, only as optional methods, so that a library that does not + * export them keeps working as before;</li> + * <li>new values may appear where this documentation says so, such as the names of pushed + * predicates.</li> + * </ul> + * An incompatible change uses a new {@link #ABI_VERSION}, with new names or a new class. A + * library reports the version it implements with {@link #abiVersion}, and Spark rejects a library + * of a version it does not support. Spark supports version 1 at least in all the releases of + * Spark 4 from 4.4.0. + * + * @since 4.4.0 + */ +@Evolving +public final class NativeBridge { + + /** + * The version of this interface. It changes when the interface changes incompatibly: see + * Compatibility above. + */ + public static final int ABI_VERSION = 1; + + private NativeBridge() {} + + /** + * Loads the library that implements the native methods of this copy of the class. Spark calls + * it on the copy of the class it creates for the library. + */ + static void load(String path) { + System.load(path); + } + + /** Returns the version of this interface that the library implements, {@link #ABI_VERSION}. */ + public static native int abiVersion(); + + // --------------------------------------------------------------------------------------------- + // Data source (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Creates a data source with the options given by the user, such as + * {@code spark.read.option("key", "value")}. + * + * @param name the name of the data source, as listed in the package manifest. A library can + * implement several data sources. + * @param optionKeys the option keys, in lower case + * @param optionValues the option values, in the same order as the keys + * @return the data source handle + */ + public static native long createDataSource( + String name, String[] optionKeys, String[] optionValues); + + /** + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. + * + * @param schemaAddress the address of the {@code ArrowSchema} to initialize with a struct type + * whose fields are the columns + */ + public static native void schema(long dataSource, long schemaAddress); + + /** + * Creates a reader for a batch scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the reader handle + */ + public static native long createReader(long dataSource, long schemaAddress); + + /** + * Creates a reader for a micro-batch streaming scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the stream reader handle + */ + public static native long createStreamReader(long dataSource, long schemaAddress); + + /** + * Creates a writer for a batch write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data (save mode "overwrite") instead of + * appending to it + * @return the writer handle + */ + public static native long createWriter(long dataSource, long schemaAddress, boolean overwrite); + + /** + * Creates a writer for a streaming write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data in each micro-batch (output mode + * "complete") instead of appending to it + * @return the writer handle + */ + public static native long createStreamWriter( + long dataSource, long schemaAddress, boolean overwrite); + + /** Releases a data source handle. */ + public static native void closeDataSource(long dataSource); + + // --------------------------------------------------------------------------------------------- + // Batch reader (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Pushes predicates down to the reader. The predicates are combined with AND. + * <p> + * Each predicate is a JSON expression tree. An expression is one of: + * <ul> + * <li>{@code {"type": "column", "name": ["a", "b"]}}: a column. A nested field has a name + * part for each level.</li> + * <li>{@code {"type": "literal", "dataType": "bigint", "value": 1}}: a literal. The data type + * is a Spark SQL type name such as {@code int}, {@code string} or {@code decimal(10,2)}. The + * value is {@code null}, a JSON boolean or number for boolean and numeric types, a string + * for {@code string}, a base64 string for {@code binary}, a decimal string for + * {@code decimal}, the days since the epoch for {@code date}, and the microseconds since the + * epoch for {@code timestamp} and {@code timestamp_ntz}.</li> + * <li>{@code {"type": "function", "name": ">", "children": [...]}}: a predicate, with the + * name it has in {@code org.apache.spark.sql.connector.expressions.filter.Predicate}. The + * names are {@code =}, {@code <>}, {@code <=>} (equal, with null equal to null), {@code <}, + * {@code <=}, {@code >} and {@code >=} with two children, {@code IN} with the value to look + * for and then the values to look in, {@code IS_NULL} and {@code IS_NOT_NULL} with one + * child, {@code STARTS_WITH}, {@code ENDS_WITH} and {@code CONTAINS} with the string and the + * string to look for, {@code AND} and {@code OR} with two predicates, and {@code NOT} with + * one. The children of the other functions are columns, literals or predicates.</li> + * </ul> + * Spark does not pass predicates that it cannot express this way, such as predicates that use + * arithmetic or other functions, nor predicates on columns that contain strings with a + * non-binary collation, such as {@code UTF8_LCASE}: strings are compared by their UTF-8 bytes. + * A later version of Spark may pass predicates with other names without changing + * {@link #ABI_VERSION}, so a reader must not report a predicate as evaluated completely if it + * contains a name that the reader does not know, or a data type that it does not support. + * <p> + * A predicate that the reader evaluates completely must follow the semantics of Spark SQL. In + * particular, for floating-point columns, NaN is equal to NaN and greater than any other value, + * including positive infinity, and {@code -0.0} is equal to {@code 0.0}, unlike IEEE 754 + * comparisons. For example, {@code d > 1.0} is true for a NaN value of {@code d}. Comparisons + * with null are null, so a row only matches if the predicate is true. + * + * @return for each predicate, whether the reader evaluates it completely, so that Spark does + * not have to evaluate it again. The reader can still use the other predicates to skip + * data. + */ + public static native boolean[] pushPredicates(long reader, String[] predicates); + + /** + * Pushes a LIMIT down to the reader. + * + * @return whether the reader returns at most {@code limit} rows in total. Spark applies the + * limit again after reading either way. + */ + public static native boolean pushLimit(long reader, int limit); + + /** + * Prunes the columns to read. + * + * @param columnNames the names of the top-level columns that the query needs + * @return whether {@link #read} returns exactly these columns, in this order. If not, it + * returns all the columns of the schema to read. + */ + public static native boolean pruneColumns(long reader, String[] columnNames); + + /** + * Plans the partitions of the scan. Spark reads each partition in a separate task. + * + * @return the partitions, each serialized in a format defined by the library + */ + public static native byte[][] partitions(long reader); + + /** + * Serializes the state of the reader that {@link #read} needs on the executors, such as the + * pushed predicates. Spark calls it after {@link #partitions}. + */ + public static native byte[] serializeReader(long reader); + + /** Releases a batch reader handle. */ + public static native void closeReader(long reader); + + // --------------------------------------------------------------------------------------------- + // Stream reader (driver) + // --------------------------------------------------------------------------------------------- + + /** Returns the offset to start from when the streaming query has no checkpoint, as JSON. */ + public static native String initialOffset(long streamReader); + + /** Returns the most recent offset available, as JSON. */ + public static native String latestOffset(long streamReader); + + /** + * Plans the partitions that read the data after {@code start} up to and including {@code end}. + * + * @return the partitions, each serialized in a format defined by the library + */ + public static native byte[][] streamPartitions(long streamReader, String start, String end); + + /** + * Serializes the state of the stream reader that {@link #read} needs on the executors. Spark + * calls it after each call to {@link #streamPartitions}. + */ + public static native byte[] serializeStreamReader(long streamReader); + + /** + * Tells the library that Spark has processed all the data up to and including {@code end}, + * and will not ask for it again. + */ + public static native void commitOffset(long streamReader, String end); + + /** Releases a stream reader handle, when the streaming query stops. */ + public static native void closeStreamReader(long streamReader); + + // --------------------------------------------------------------------------------------------- + // Partition reader (executor) + // --------------------------------------------------------------------------------------------- + + /** + * Reads a partition of a batch or streaming scan. + * + * @param readerState the state returned by {@link #serializeReader} or + * {@link #serializeStreamReader}, or an empty array if there is none + * @param partition a partition returned by {@link #partitions} or {@link #streamPartitions} + * @param streamAddress the address of the {@code ArrowArrayStream} to initialize. Each array + * of the stream is a struct array whose fields are the columns to read: + * the columns of the schema to read, or the pruned columns. Spark calls + * the stream from the thread of the task that reads the partition. + * {@code get_next} may block, for example on network I/O, but Spark + * cannot interrupt native code: when the task is killed, Spark only + * stops reading once {@code get_next} returns, and then releases the + * stream. So a {@code get_next} that blocks forever keeps the task + * running, unless the task reaper ({@code spark.task.reaper.enabled} + * and {@code spark.task.reaper.killTimeout}) kills the executor. Use + * timeouts for blocking operations. + */ + public static native void read(byte[] readerState, byte[] partition, long streamAddress); Review Comment: **Memory used by native libraries** Spark's operators take their memory from the unified memory manager and spill when it's refused, but a library's memory isn't in any pool: the PR doesn't use `TaskMemoryManager`, and `ArrowUtils.rootAllocator` has no limit. So the operators can't account for it, and nothing bounds it. It comes out of `spark.executor.memoryOverhead`, which is meant for "VM overheads, interned strings, other native overheads". On YARN or Kubernetes, a library that goes over gets the container killed, along with every task on the executor. Python workers have the same problem, and Spark's answer was a dedicated budget: `spark.executor.pyspark.memory` is requested as its own executor resource and enforced in the worker with `RLIMIT_AS`. A native library runs inside the executor, so a process limit can't work. The library has to reserve against a budget. Could version 1 include: 1. A budget that Spark requests for native libraries: the off-heap pool (`spark.memory.offHeap.size`, which Spark already requests), or a dedicated setting like `spark.executor.pyspark.memory`. 2. A per-task context in `read` and `createDataWriter`, with `reserve(bytes)`, which returns the bytes granted, and `release(bytes)`, charged to the task against that budget. A refused reservation fails the task with a clear error. The context has to be explicit, since a library's own threads have no `TaskContext`. No spilling at first. The same context could carry a batch size for reads, which today the library alone decides. 3. A section on the docs page on what Spark accounts for and what the library must. CometTaskMemoryManager (https://github.com/apache/datafusion-comet/blob/main/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java) does this for Comet: native memory reserved with `acquireExecutionMemory` and an off-heap consumer, and no spilling. Gluten goes one step further with Velox: it reserves from Spark before growing a Velox memory pool (`ListenableArbitrator`, https://github.com/apache/gluten/blob/main/cpp/velox/memory/VeloxMemoryManager.cc), and its Spark `MemoryConsumer` implements `spill` by making Velox spill (`TreeMemoryConsumer`, https://github.com/apache/gluten/blob/main/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/spark/TreeMemoryConsumer.java), so Spark can also take memory back. I'd put the context in version 1's signatures rather than a later optional variant: libraries written against version 1 without it will run unaccounted, and an optional function only reaches the libraries that adopt it. -- 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]
