diff --git a/.github/actions/rust-test/action.yaml b/.github/actions/rust-test/action.yaml index c39c2dcd4f..a8f93be48b 100644 --- a/.github/actions/rust-test/action.yaml +++ b/.github/actions/rust-test/action.yaml @@ -62,6 +62,15 @@ runs: run: | cargo install cargo-nextest --locked + - name: Build test UDF library (pre-requisite for Rust UDF tests) + shell: bash + run: | + cd native + # comet-test-udfs is `crate-type = ["cdylib"]` and has no test targets, so a + # test build compiles it without emitting the shared library the rust_udf + # tests dlopen. Build it explicitly. + cargo build -p comet-test-udfs + - name: Run Cargo test shell: bash run: | diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 0a71607ead..fd42fcff55 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -215,7 +215,11 @@ jobs: uses: actions/upload-artifact@v7 with: name: native-lib-linux - path: native/target/ci/libcomet.so + # libcomet_test_udfs is the cdylib CometRustUdfSuite loads; it is built by the same + # cargo invocation and the suite finds it next to libcomet. + path: | + native/target/ci/libcomet.so + native/target/ci/libcomet_test_udfs.so retention-days: 1 - name: Save Cargo cache @@ -405,6 +409,7 @@ jobs: org.apache.comet.CometStringDecodeSuite org.apache.comet.CometWidthBucketSuite org.apache.comet.CometUuidExpressionSuite + org.apache.comet.CometRustUdfSuite fail-fast: false name: ${{ matrix.profile.name }} [${{ matrix.suite.name }}] runs-on: ubuntu-24.04 diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 2335588b70..1b6fe7f273 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -82,7 +82,11 @@ jobs: uses: actions/upload-artifact@v7 with: name: native-lib-macos - path: native/target/ci/libcomet.dylib + # libcomet_test_udfs is the cdylib CometRustUdfSuite loads; it is built by the same + # cargo invocation and the suite finds it next to libcomet. + path: | + native/target/ci/libcomet.dylib + native/target/ci/libcomet_test_udfs.dylib retention-days: 1 - name: Save Cargo cache @@ -221,6 +225,7 @@ jobs: org.apache.comet.CometStringDecodeSuite org.apache.comet.CometWidthBucketSuite org.apache.comet.CometUuidExpressionSuite + org.apache.comet.CometRustUdfSuite fail-fast: false name: ${{ matrix.os }}/${{ matrix.profile.name }} [${{ matrix.suite.name }}] diff --git a/docs/source/user-guide/latest/index.rst b/docs/source/user-guide/latest/index.rst index 2cd1c7e16f..d7b665d9ca 100644 --- a/docs/source/user-guide/latest/index.rst +++ b/docs/source/user-guide/latest/index.rst @@ -51,6 +51,7 @@ to read more. Supported Operators Supported Expressions ScalaUDF and Java UDF Support + Custom Rust UDFs .. toctree:: :maxdepth: 1 diff --git a/docs/source/user-guide/latest/rust_udfs.md b/docs/source/user-guide/latest/rust_udfs.md new file mode 100644 index 0000000000..2f1e1eb862 --- /dev/null +++ b/docs/source/user-guide/latest/rust_udfs.md @@ -0,0 +1,241 @@ + + +# Custom Rust UDFs + +Comet can load scalar user-defined functions written in Rust from a shared library and run them +natively, inside the Comet pipeline, with no JVM round trip per row. + +This is different from [Scala UDF and Java UDF Support](scala_java_udfs.md), where the user +function stays on the JVM and Comet dispatches into it. Here the function is compiled Rust that +operates directly on Arrow arrays. + +> **Experimental.** This feature and the ABI it depends on are experimental. Neither +> `comet-udf-sdk` nor `CometRustUDF` is part of Comet's supported API: they fall under +> [everything else is internal](../../about/versioning_policy.md#everything-else-is-internal) in the +> [versioning policy](../../about/versioning_policy.md), so they may change or be removed in any +> release, including a patch release, with no deprecation cycle. Expect to rebuild your UDF library +> against the SDK from each Comet release you upgrade to. It is not yet recommended for production +> use. See [Limitations](#limitations) before adopting it. + +## Writing a UDF + +A UDF library is an ordinary Rust `cdylib` that depends on `comet-udf-sdk` and `arrow`: + +```toml +[package] +name = "my-comet-udfs" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["cdylib"] + +[dependencies] +arrow = "58" +comet-udf-sdk = { git = "https://github.com/apache/datafusion-comet" } +``` + +Implement the `CometCScalarUdf` trait and export it: + +```rust +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int64Array}; +use arrow::datatypes::{DataType, Field}; +use comet_udf_sdk::c_abi::CometCScalarUdf; +use comet_udf_sdk::comet_c_udf_export; + +#[derive(Default)] +pub struct AddOne; + +impl CometCScalarUdf for AddOne { + /// The name the function is registered and called under. + fn name(&self) -> &str { + "add_one" + } + + /// Validate the argument types and declare the output type. Called once + /// per execution, before `invoke`. Returning `Err` fails the query with + /// your message. + fn return_field(&self, args: &[Field]) -> Result { + if args.len() != 1 || args[0].data_type() != &DataType::Int64 { + return Err("add_one expects (Int64) -> Int64".into()); + } + Ok(Field::new("add_one", DataType::Int64, true)) + } + + /// Evaluate one batch. `args` holds one Arrow array per argument. + fn invoke(&self, args: &[ArrayRef], _n_rows: usize) -> Result { + let a = args[0] + .as_any() + .downcast_ref::() + .ok_or("expected an Int64Array")?; + Ok(Arc::new( + a.iter().map(|v| v.map(|x| x + 1)).collect::(), + )) + } +} + +comet_c_udf_export!(AddOne); +``` + +Each type passed to `comet_c_udf_export!` must implement `CometCScalarUdf` and `Default`. One +library may export any number of functions. + +Build it: + +```sh +cargo build --release +# target/release/libmy_comet_udfs.so (Linux) +# target/release/libmy_comet_udfs.dylib (macOS) +``` + +Note that the library depends only on `arrow` and the SDK, not on DataFusion. The ABI is built +purely on the [Arrow C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html), +which keeps your library decoupled from the DataFusion version Comet happens to use. It also means +the same ABI is implementable from C or C++. + +## Registering and calling a UDF + +Register the function on the driver, giving its name, the library path, and its signature: + +```scala +import org.apache.comet.udf.CometRustUDF +import org.apache.spark.sql.types.LongType + +CometRustUDF.register( + spark, + name = "add_one", + libraryPath = "/opt/udfs/libmy_comet_udfs.so", + inputTypes = Seq(LongType), + returnType = LongType) +``` + +Registration loads the library on the driver and verifies that a function with that name exists, so +a bad path or a missing function fails immediately with a clear error rather than at execution time. + +The function is then callable from SQL or the DataFrame API like any other: + +```scala +spark.range(0, 5).selectExpr("add_one(id) AS y").show() +``` + +Your function must be a pure function of its arguments. Comet plans every Rust UDF as immutable, +which lets the optimizer fold a call over constants, evaluate it once and reuse the result, or drop +a repeated call as a common subexpression. `register` therefore rejects `deterministic = false` +rather than accept a function whose volatility it would go on to ignore. + +`libraryPath` is passed to the platform's dynamic loader. An absolute path is what you want in +practice, and it is what the rest of this page assumes, but a bare library name resolves the same +way it would for any other shared object, through `LD_LIBRARY_PATH` on Linux and +`DYLD_LIBRARY_PATH` on macOS. That is a convenience, not a sandbox: Comet does not restrict which +paths may be loaded, so it makes no difference to the trust decision described under +[Limitations](#limitations). + +## Return types + +A UDF does not have one fixed return type. Its `return_field` is called with the actual argument +types and computes the output type on demand, so a single kernel can serve many signatures: the +`echo_c` UDF in Comet's own test library returns whatever type it is given, for every type in the +table below. + +What is fixed is the type you declare to `register`, because Spark needs a concrete `DataType` at +analysis time in order to plan the query. That declaration is per-registration, not per-kernel: +re-registering the same function under different types is supported, and the kernel computes the +matching return type each time. + +The two must agree. Comet checks the declared type against what `return_field` reports at planning +time and fails with both types named if they differ, rather than letting it surface as a type +assertion partway through execution. Nested nullability (`containsNull`, struct field nullability) +is not part of that comparison, since Spark and Arrow disagree about it harmlessly, but everything +that changes how bytes are read is: decimal precision and scale, timestamp unit, and struct field +names and order. + +Watch for Spark's own type promotion when declaring: `cast(id as decimal(10,2)) + 0.25` has type +`decimal(11,2)`, not `decimal(10,2)`, so registering the latter is a mismatch. + +## Supported types + +Arguments and return values may be any of: + +| Spark type | Arrow type | +| ------------------- | ------------------------ | +| `BooleanType` | `Boolean` | +| `ByteType` | `Int8` | +| `ShortType` | `Int16` | +| `IntegerType` | `Int32` | +| `LongType` | `Int64` | +| `FloatType` | `Float32` | +| `DoubleType` | `Float64` | +| `DecimalType(p, s)` | `Decimal128(p, s)` | +| `StringType` | `Utf8` | +| `BinaryType` | `Binary` | +| `DateType` | `Date32` | +| `TimestampType` | `Timestamp(Microsecond)` | +| `TimestampNTZType` | `Timestamp(Microsecond)` | + +Complex types are supported and may be nested arbitrarily: + +| Spark type | Arrow type | +| ------------------------- | --------------------- | +| `ArrayType(t)` | `List(t)` | +| `StructType(f1, f2, ...)` | `Struct(f1, f2, ...)` | +| `MapType(k, v)` | `Map(k, v)` | + +Nulls are preserved in both directions; a null input row arrives as a null slot in the Arrow array +and your output nulls come back to Spark as nulls. + +Not yet supported: `CalendarIntervalType`, `NullType`, `UserDefinedType`, and Arrow extension types +such as Variant and Geometry. + +## Error handling + +Returning `Err(String)` from `return_field` or `invoke` fails the query with your message attached. +This is the intended way to reject bad input. + +Panics in your code are caught at the FFI boundary and converted into query errors, so an `unwrap` +on `None` fails that query rather than taking down the executor. Do not rely on this as a control +flow mechanism: prefer returning `Err`, which produces a much better message. + +## Limitations + +This feature is at an early stage. The current limitations are: + +- **Scalar functions only.** Aggregate, window, and table functions are not supported. +- **Immutable functions only.** A UDF must return the same output for the same input. Comet plans + every Rust UDF with DataFusion's `Volatility::Immutable`, so a function that reads a clock, draws + from an RNG, or carries state across batches may be folded at plan time, evaluated once and + reused, or eliminated as a common subexpression. `register` rejects `deterministic = false`. +- **The library must already be present on every executor**, at the same absolute path given to + `register`. Comet does not distribute it for you: stage it with your image, a mounted volume, or + your cluster's own file distribution, and pass a path that is valid cluster-wide. A path that + exists only on the driver will fail at execution time. +- **Up to 4 arguments** per function. +- **No type coercion.** Arguments arrive as the types the query produces; `return_field` should + reject anything it does not handle. Downcast defensively in `invoke` and return a clear `Err` + rather than assuming a particular array layout. +- **Loading a library is loading native code.** It runs with the full privileges of the executor + process and Comet cannot sandbox it: a bug in a UDF can corrupt memory or crash the executor. + Only register libraries you trust and control. +- Once loaded, a library stays loaded for the life of the process. Replacing the file on disk has + no effect until the executors restart. +- The ABI is versioned and checked strictly at load time. A library built against a different + Comet's SDK is refused with an explicit ABI-mismatch error rather than being loaded unsafely. + Rebuild your UDF library when upgrading Comet. diff --git a/native/Cargo.lock b/native/Cargo.lock index eeaa9da5ce..41e0fdc540 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -1470,6 +1470,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "comet-test-udfs" +version = "1.1.0" +dependencies = [ + "arrow", + "comet-udf-sdk", +] + +[[package]] +name = "comet-udf-sdk" +version = "1.1.0" +dependencies = [ + "arrow", +] + [[package]] name = "comfy-table" version = "7.2.2" @@ -1943,6 +1958,8 @@ dependencies = [ "aws-config", "aws-credential-types", "bytes", + "comet-test-udfs", + "comet-udf-sdk", "criterion", "datafusion", "datafusion-comet-common", @@ -1963,6 +1980,7 @@ dependencies = [ "itertools 0.15.0", "jni 0.22.4", "lazy_static", + "libloading", "log", "log4rs", "mimalloc", diff --git a/native/Cargo.toml b/native/Cargo.toml index a4ad01a3f9..c41ead3ccf 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -16,8 +16,8 @@ # under the License. [workspace] -default-members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle"] -members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle"] +default-members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "comet-udf-sdk", "comet-test-udfs"] +members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "comet-udf-sdk", "comet-test-udfs"] resolver = "2" [workspace.package] diff --git a/native/comet-test-udfs/Cargo.toml b/native/comet-test-udfs/Cargo.toml new file mode 100644 index 0000000000..e07a08bb06 --- /dev/null +++ b/native/comet-test-udfs/Cargo.toml @@ -0,0 +1,35 @@ +# 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] +name = "comet-test-udfs" +version = { workspace = true } +edition = { workspace = true } +rust-version = { workspace = true } +license = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +publish = false +description = "Test UDF cdylib used by Comet's Rust UDF host tests (arrow-ffi based)" + +[lib] +crate-type = ["cdylib"] + +[dependencies] +arrow = { workspace = true } +comet-udf-sdk = { path = "../comet-udf-sdk" } diff --git a/native/comet-test-udfs/src/lib.rs b/native/comet-test-udfs/src/lib.rs new file mode 100644 index 0000000000..3e05ac05b1 --- /dev/null +++ b/native/comet-test-udfs/src/lib.rs @@ -0,0 +1,190 @@ +// 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. + +//! Test UDF cdylib for Comet's Rust UDF host tests. +//! +//! Exports, through the Comet UDF C ABI: +//! +//! - `add_one_c` — `(Int64) -> Int64`, the basic compute path +//! - `echo_c` — identity over any type, used to check that each supported +//! Spark type survives the round trip through the ABI with its nulls +//! - `stringify_c` — `(any) -> Utf8`, which forces the UDF to actually +//! decode the values rather than hand the array straight back +//! - `panics_on_invoke` / `panics_on_return_field` — panic containment +//! +//! Note that this crate depends only on `arrow` and `comet-udf-sdk` — no +//! DataFusion dependency — which is the point of the ABI. +//! +//! Built as `libcomet_test_udfs.{so,dylib}`. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, Int64Array, StringArray}; +use arrow::datatypes::{DataType, Field}; +use arrow::util::display::{ArrayFormatter, FormatOptions}; + +use comet_udf_sdk::c_abi::CometCScalarUdf; +use comet_udf_sdk::comet_c_udf_export; + +/// `add_one` exposed via the C ABI. +pub struct AddOneC; + +impl Default for AddOneC { + fn default() -> Self { + AddOneC + } +} + +impl CometCScalarUdf for AddOneC { + fn name(&self) -> &str { + "add_one_c" + } + + fn return_field(&self, args: &[Field]) -> Result { + if args.len() != 1 { + return Err(format!("add_one_c expects 1 arg, got {}", args.len())); + } + if args[0].data_type() != &DataType::Int64 { + return Err(format!( + "add_one_c expects Int64, got {}", + args[0].data_type() + )); + } + Ok(Field::new("add_one_c", DataType::Int64, true)) + } + + fn invoke(&self, args: &[ArrayRef], _n_rows: usize) -> Result { + let arr = args[0] + .as_any() + .downcast_ref::() + .ok_or_else(|| "expected Int64Array".to_string())?; + let out: Int64Array = arr.iter().map(|v| v.map(|x| x + 1)).collect(); + Ok(Arc::new(out)) + } +} + +/// Identity over any input type: declares the argument's own type as the +/// return type and hands the array back. +/// +/// Used to check that every supported Spark type survives the trip out to +/// the UDF and back through the Arrow C Data Interface, nulls included. +#[derive(Default)] +pub struct EchoC; + +impl CometCScalarUdf for EchoC { + fn name(&self) -> &str { + "echo_c" + } + + fn return_field(&self, args: &[Field]) -> Result { + if args.len() != 1 { + return Err(format!("echo_c expects 1 arg, got {}", args.len())); + } + // Echo the argument's own type, so this works for every type + // including the parameterized ones (decimal, timestamp with tz). + Ok(Field::new("echo_c", args[0].data_type().clone(), true)) + } + + fn invoke(&self, args: &[ArrayRef], _n_rows: usize) -> Result { + Ok(Arc::clone(&args[0])) + } +} + +/// Renders any input array as strings, one per row, preserving nulls. +/// +/// Where `echo_c` only proves the array survives the round trip, this +/// forces the UDF to decode each value, so it catches a type that arrives +/// with the right `DataType` but an unreadable layout. +#[derive(Default)] +pub struct StringifyC; + +impl CometCScalarUdf for StringifyC { + fn name(&self) -> &str { + "stringify_c" + } + + fn return_field(&self, args: &[Field]) -> Result { + if args.len() != 1 { + return Err(format!("stringify_c expects 1 arg, got {}", args.len())); + } + Ok(Field::new("stringify_c", DataType::Utf8, true)) + } + + fn invoke(&self, args: &[ArrayRef], _n_rows: usize) -> Result { + let array = &args[0]; + let options = FormatOptions::default().with_null("__NULL__"); + let formatter = ArrayFormatter::try_new(array.as_ref(), &options) + .map_err(|e| format!("stringify_c cannot format {}: {e}", array.data_type()))?; + let values: StringArray = (0..array.len()) + .map(|i| { + if array.is_null(i) { + None + } else { + Some(formatter.value(i).to_string()) + } + }) + .collect(); + Ok(Arc::new(values)) + } +} + +/// Panics unconditionally when invoked, so host tests can verify that a +/// panic inside user code is caught at the FFI boundary and surfaced as a +/// query error rather than unwinding into the host. +#[derive(Default)] +pub struct PanicsOnInvoke; + +impl CometCScalarUdf for PanicsOnInvoke { + fn name(&self) -> &str { + "panics_on_invoke" + } + + fn return_field(&self, _args: &[Field]) -> Result { + Ok(Field::new("panics_on_invoke", DataType::Int64, true)) + } + + fn invoke(&self, _args: &[ArrayRef], _n_rows: usize) -> Result { + panic!("deliberate panic from user UDF code") + } +} + +/// Panics inside `return_field`, i.e. during planning rather than +/// execution, exercising the other side of the FFI panic boundary. +#[derive(Default)] +pub struct PanicsOnReturnField; + +impl CometCScalarUdf for PanicsOnReturnField { + fn name(&self) -> &str { + "panics_on_return_field" + } + + fn return_field(&self, _args: &[Field]) -> Result { + panic!("deliberate panic from user return_field") + } + + fn invoke(&self, _args: &[ArrayRef], _n_rows: usize) -> Result { + unreachable!("return_field panics first") + } +} + +comet_c_udf_export!( + AddOneC, + EchoC, + StringifyC, + PanicsOnInvoke, + PanicsOnReturnField +); diff --git a/native/comet-udf-sdk/Cargo.toml b/native/comet-udf-sdk/Cargo.toml new file mode 100644 index 0000000000..50830709bd --- /dev/null +++ b/native/comet-udf-sdk/Cargo.toml @@ -0,0 +1,32 @@ +# 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] +name = "comet-udf-sdk" +version = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +edition = { workspace = true } +rust-version = { workspace = true } +license = { workspace = true } +description = "SDK for writing custom Rust UDFs that run inside Apache DataFusion Comet (arrow-ffi based)" + +publish = false + +[dependencies] +arrow = { workspace = true } diff --git a/native/comet-udf-sdk/src/c_abi.rs b/native/comet-udf-sdk/src/c_abi.rs new file mode 100644 index 0000000000..f8e0a09036 --- /dev/null +++ b/native/comet-udf-sdk/src/c_abi.rs @@ -0,0 +1,887 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The Comet UDF C ABI, sedona-style. +//! +//! The wire format is two `#[repr(C)]` structs of function pointers, +//! parameterized only by Arrow's C Data Interface +//! (`FFI_ArrowSchema` / `FFI_ArrowArray`). No DataFusion types appear in +//! the FFI surface, so the user's cdylib only needs a matching `arrow` +//! crate, not a matching `datafusion` version. +//! +//! # Stability +//! +//! These structs are **specific to one Comet version and are not yet ABI +//! stable across Comet releases**. The layouts here may change in any +//! release, including a patch release, with no deprecation cycle: they are +//! internal types under Comet's +//! [versioning policy](https://datafusion.apache.org/comet/about/versioning_policy.html), +//! like every other type the native crates ship. `comet_udf_abi_version` +//! is checked strictly at load time, so a stale cdylib is refused with an +//! explicit error rather than loaded unsafely, but the practical +//! consequence is that a UDF library must be rebuilt against the SDK from +//! the Comet release it will run on. +//! +//! What does *not* have to match is the host's own dependency versions. +//! Only `FFI_ArrowArray` and `FFI_ArrowSchema` cross the boundary, and +//! those are `#[repr(C)]` renderings of the Arrow C Data Interface, which +//! is stable across `arrow` versions. So a cdylib and the Comet host that +//! loads it may be built against different `arrow` versions, and the +//! cdylib needs no `datafusion` dependency at all. The binding constraint +//! is what `comet-udf-sdk` itself compiles against: the SDK is built into +//! your cdylib, so Cargo must be able to unify its `arrow` requirement +//! with yours. +//! +//! # Authoring a UDF +//! +//! Implement [`CometCScalarUdf`] for a type that also implements `Default`, +//! then use the [`comet_c_udf_export!`] macro to emit the discovery entry +//! point: +//! +//! ```ignore +//! use comet_udf_sdk::c_abi::*; +//! use arrow::array::{ArrayRef, Int64Array}; +//! use arrow::datatypes::{DataType, Field}; +//! use std::sync::Arc; +//! +//! #[derive(Default)] +//! pub struct AddOne; +//! impl CometCScalarUdf for AddOne { +//! fn name(&self) -> &str { "add_one_c" } +//! fn return_field(&self, args: &[Field]) -> Result { +//! if args.len() != 1 || args[0].data_type() != &DataType::Int64 { +//! return Err("expected (Int64) -> Int64".into()); +//! } +//! Ok(Field::new("add_one_c", DataType::Int64, true)) +//! } +//! fn invoke(&self, args: &[ArrayRef], _n: usize) -> Result { +//! let a = args[0].as_any().downcast_ref::().unwrap(); +//! Ok(Arc::new(a.iter().map(|v| v.map(|x| x + 1)).collect::())) +//! } +//! } +//! +//! comet_udf_sdk::comet_c_udf_export!(AddOne); +//! ``` + +use std::ffi::{c_char, c_int, c_void}; + +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; + +/// Generic non-zero error code returned by `init` / `execute` to signal +/// failure. The host treats any non-zero return as an error and calls +/// `get_last_error` for the message; the specific code is informational. +const C_ABI_ERR: c_int = 1; + +// -- panic containment ----------------------------------------------------- +// +// Every `extern "C"` function in this module is an unwind boundary. A panic +// that escapes one aborts the whole process (Rust's default `extern "C"` +// unwind behavior since 1.81), which for Comet means killing the executor +// JVM and losing every task on it -- not just the query that used the UDF. +// +// User UDF code is arbitrary and panicking is idiomatic Rust (`unwrap`, +// slice indexing, integer overflow in debug), so the SDK treats a panic in +// user code as an ordinary error: catch it at the boundary, convert it to a +// message, and report it through the same `get_last_error` channel as a +// returned `Err`. The query fails; the executor survives. + +/// Render a caught panic payload as an error message. +fn panic_message(panic: Box) -> String { + let detail = panic + .downcast_ref::<&'static str>() + .map(|s| s.to_string()) + .or_else(|| panic.downcast_ref::().cloned()) + .unwrap_or_else(|| "".to_string()); + format!("panic in UDF code: {detail}") +} + +/// Run `f`, converting a panic into `Err(message)`. +fn catch_panic(f: impl FnOnce() -> Result) -> Result { + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) { + Ok(result) => result, + Err(panic) => Err(panic_message(panic)), + } +} + +/// Run an infallible `f` (typically a release/cleanup callback), containing +/// any panic. Used where the ABI gives us no way to report an error and +/// aborting would be a worse outcome than leaking. +/// +/// There is no error channel to return this on and no logging facade in the +/// SDK's dependencies (`arrow` is the only one, deliberately), so the panic +/// is reported on stderr. Silently swallowing it would leave a leaked +/// allocation or a half-run destructor with nothing at all to show for it. +fn catch_panic_infallible(context: &str, f: impl FnOnce()) { + if let Err(panic) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) { + eprintln!("comet-udf-sdk: {context}: {}", panic_message(panic)); + } +} + +// -- factory struct -------------------------------------------------------- + +/// Factory for [`CometCScalarKernelImpl`] instances. +/// +/// Lives in a registry, may be cloned across an FFI boundary. Calls to +/// `function_name` and `new_impl` must be thread-safe (the implementation +/// is responsible for any internal synchronization). +/// +/// `#[repr(C)]` layout, matched by the host loader. Adding new fields +/// requires bumping `COMET_UDF_ABI_VERSION`. +#[repr(C)] +pub struct CometCScalarKernel { + /// Return the function name this kernel implements as a NUL-terminated + /// UTF-8 C string. The pointer must remain valid for the lifetime of + /// the [`CometCScalarKernel`]. + /// + /// May be `None`, in which case the kernel is treated as anonymous and + /// won't be discoverable by name. (Comet always sets this; field is + /// optional for parity with sedona's design.) + pub function_name: Option *const c_char>, + + /// Initialize a new [`CometCScalarKernelImpl`] into `out`. Called once + /// per execution, on the thread that will then drive `init`/`execute`. + pub new_impl: + Option, + + /// Release this kernel. After release, all callbacks must be set to + /// `None`. Called when the host's `LoadedLibrary` is dropped. + pub release: Option, + + /// Implementation-private data, opaque to the host. + pub private_data: *mut c_void, +} + +// SAFETY: `CometCScalarKernel` is a thin wrapper around C function +// pointers with caller-defined synchronization semantics; the trait impls +// are required so loaded kernels can be referenced from multi-threaded +// host code. Implementations of the FFI must respect thread safety as +// described in the doc comments. +unsafe impl Send for CometCScalarKernel {} +unsafe impl Sync for CometCScalarKernel {} + +impl Default for CometCScalarKernel { + fn default() -> Self { + Self { + function_name: None, + new_impl: None, + release: None, + private_data: std::ptr::null_mut(), + } + } +} + +impl Drop for CometCScalarKernel { + fn drop(&mut self) { + if let Some(release) = self.release.take() { + // SAFETY: release is the FFI-defined cleanup callback; + // implementations must reset `release` to None per the contract. + unsafe { release(self) }; + } + } +} + +// -- per-execution instance struct ---------------------------------------- + +/// Per-execution instance produced by [`CometCScalarKernel::new_impl`]. +/// +/// Not thread-safe; the caller must serialize access. Typically used on +/// one thread for one batch then dropped. +#[repr(C)] +pub struct CometCScalarKernelImpl { + /// Compute the return type from arg types and (optionally) bound + /// scalar arguments. + /// + /// On success, `out` is populated with the return type as an + /// `FFI_ArrowSchema` and the function returns 0. On failure, returns + /// a non-zero errno and the host calls `get_last_error` to retrieve + /// the message. + /// + /// `arg_types` points to an array of `n_args` `*const FFI_ArrowSchema`. + /// `scalar_args` may be NULL (no scalars) or point to an array of + /// `n_args` `*mut FFI_ArrowArray`, each of length 1 (or NULL when + /// the corresponding argument is not a scalar). Implementations may + /// take ownership of scalar entries by replacing them with released + /// arrays. + pub init: Option< + unsafe extern "C" fn( + *mut CometCScalarKernelImpl, + arg_types: *const *const FFI_ArrowSchema, + scalar_args: *const *mut FFI_ArrowArray, + n_args: i64, + out: *mut FFI_ArrowSchema, + ) -> c_int, + >, + + /// Execute one batch. + /// + /// `args` points to an array of `n_args` `*mut FFI_ArrowArray`. + /// Each input must have length `n_rows` or length 1 (scalar broadcast). + /// On success writes the result into `out` and returns 0. + pub execute: Option< + unsafe extern "C" fn( + *mut CometCScalarKernelImpl, + args: *const *mut FFI_ArrowArray, + n_args: i64, + n_rows: i64, + out: *mut FFI_ArrowArray, + ) -> c_int, + >, + + /// Return the last error message produced by `init` or `execute`. + /// + /// Returns NULL if there is no error. The pointer is valid until the + /// next call to any method on this instance (or `release`). + pub get_last_error: Option *const c_char>, + + /// Release this instance. After release `release` must be `None`. + pub release: Option, + + /// Implementation-private data, opaque to the host. + pub private_data: *mut c_void, +} + +impl Default for CometCScalarKernelImpl { + fn default() -> Self { + Self { + init: None, + execute: None, + get_last_error: None, + release: None, + private_data: std::ptr::null_mut(), + } + } +} + +impl Drop for CometCScalarKernelImpl { + fn drop(&mut self) { + if let Some(release) = self.release.take() { + // SAFETY: per the FFI contract `release` cleans up + // private_data and resets `release` to None. + unsafe { release(self) }; + } + } +} + +// -- discovery list -------------------------------------------------------- + +/// List of kernels exposed by a cdylib via `comet_c_udf_list_v1`. +/// +/// Ownership of the underlying `CometCScalarKernel` array is transferred +/// to the host: the host invokes each kernel's `release` and then frees +/// the list via `release_list`. +#[repr(C)] +pub struct CometCScalarKernelList { + /// Pointer to the kernel array, or null if `len == 0`. + pub kernels: *mut CometCScalarKernel, + /// Number of kernels in `kernels`. + pub len: i64, + /// Free the array of kernels. Implementations must invoke each + /// kernel's `release` first, then release the array storage. + pub release: Option, +} + +impl Default for CometCScalarKernelList { + fn default() -> Self { + Self { + kernels: std::ptr::null_mut(), + len: 0, + release: None, + } + } +} + +impl Drop for CometCScalarKernelList { + fn drop(&mut self) { + if let Some(release) = self.release.take() { + // SAFETY: `release` is responsible for freeing each kernel and + // the array storage that backs `kernels`. + unsafe { release(self) }; + } + } +} + +// -- high-level Rust trait + adapter -------------------------------------- + +use arrow::array::ArrayRef; +use arrow::datatypes::Field; + +/// High-level Rust trait the user implements to author a UDF. +/// +/// Adapted to the C ABI by [`ExportedScalarKernel`]. +/// +/// # Only immutable functions are supported +/// +/// Comet registers every imported kernel with DataFusion's +/// `Volatility::Immutable`, which asserts that the same inputs always +/// produce the same output. The planner is entitled to act on that: it may +/// evaluate a call once and reuse the result, fold a call over constants at +/// plan time, or eliminate a repeated call as a common subexpression. +/// +/// So `invoke` must be a pure function of its arguments. A kernel that +/// reads a clock, draws from an RNG, or accumulates state across batches +/// will produce results that depend on decisions the optimizer is free to +/// change between releases. There is currently no way to declare such a +/// kernel: `CometRustUDF.register` rejects `deterministic = false` rather +/// than registering a function whose volatility Comet would then ignore. +pub trait CometCScalarUdf: Send + Sync { + /// Stable function name. Returned via `function_name` over the FFI. + fn name(&self) -> &str; + + /// Compute the output `Field` from the input `Field`s. + /// + /// Called once per execution, before `invoke`. May reject input + /// arities or types by returning an error; the host then surfaces + /// the message to the planner. + fn return_field(&self, args: &[Field]) -> Result; + + /// Evaluate one batch of `n_rows` rows. + fn invoke(&self, args: &[ArrayRef], n_rows: usize) -> Result; +} + +/// Wraps a user `CometCScalarUdf` impl as a [`CometCScalarKernel`] +/// suitable for emission via the C ABI discovery list. +pub struct ExportedScalarKernel { + inner: std::sync::Arc, + /// NUL-terminated C string holding the function name. Lifetime is + /// tied to `self` so the pointer returned to the host stays valid. + name_c: std::ffi::CString, +} + +impl ExportedScalarKernel { + /// Wrap `udf` for export. + pub fn new(udf: U) -> Self { + let name_c = std::ffi::CString::new(udf.name().to_string()) + .expect("UDF name must not contain interior NUL bytes"); + Self { + inner: std::sync::Arc::new(udf), + name_c, + } + } +} + +impl From for CometCScalarKernel { + fn from(value: ExportedScalarKernel) -> Self { + let boxed: Box = Box::new(value); + let private = Box::into_raw(boxed) as *mut c_void; + CometCScalarKernel { + function_name: Some(c_factory_function_name), + new_impl: Some(c_factory_new_impl), + release: Some(c_factory_release), + private_data: private, + } + } +} + +unsafe extern "C" fn c_factory_function_name(this: *const CometCScalarKernel) -> *const c_char { + debug_assert!(!this.is_null()); + let this = unsafe { &*this }; + // A released kernel has null private_data and `release: None`; checking + // both catches a call made after release, not just an uninitialized one. + debug_assert!(!this.private_data.is_null() && this.release.is_some()); + let exp = unsafe { &*(this.private_data as *const ExportedScalarKernel) }; + exp.name_c.as_ptr() +} + +unsafe extern "C" fn c_factory_new_impl( + this: *const CometCScalarKernel, + out: *mut CometCScalarKernelImpl, +) { + debug_assert!(!this.is_null()); + debug_assert!(!out.is_null()); + let this = unsafe { &*this }; + debug_assert!(!this.private_data.is_null() && this.release.is_some()); + let exp = unsafe { &*(this.private_data as *const ExportedScalarKernel) }; + // On panic, leave `out` as the default (all callbacks None). The host + // detects the missing `init` and reports it as a load error. + catch_panic_infallible("constructing kernel impl", || { + let impl_state = ExportedScalarKernelImpl { + inner: std::sync::Arc::clone(&exp.inner), + last_arg_fields: None, + last_return_field: None, + last_error: std::ffi::CString::default(), + }; + unsafe { + std::ptr::write(out, CometCScalarKernelImpl::from(impl_state)); + } + }); +} + +unsafe extern "C" fn c_factory_release(this: *mut CometCScalarKernel) { + debug_assert!(!this.is_null()); + let this_ref = unsafe { &mut *this }; + if !this_ref.private_data.is_null() { + // SAFETY: private_data was set via Box::into_raw in + // From; reclaim and drop. Dropping runs the + // user type's Drop, which may panic; contain it rather than abort. + let raw = this_ref.private_data as *mut ExportedScalarKernel; + this_ref.private_data = std::ptr::null_mut(); + catch_panic_infallible("releasing kernel", || drop(unsafe { Box::from_raw(raw) })); + } + this_ref.function_name = None; + this_ref.new_impl = None; + this_ref.release = None; +} + +struct ExportedScalarKernelImpl { + inner: std::sync::Arc, + last_arg_fields: Option>, + last_return_field: Option, + last_error: std::ffi::CString, +} + +impl From for CometCScalarKernelImpl { + fn from(value: ExportedScalarKernelImpl) -> Self { + let boxed = Box::new(value); + let private = Box::into_raw(boxed) as *mut c_void; + CometCScalarKernelImpl { + init: Some(c_kernel_init), + execute: Some(c_kernel_execute), + get_last_error: Some(c_kernel_get_last_error), + release: Some(c_kernel_release), + private_data: private, + } + } +} + +unsafe extern "C" fn c_kernel_init( + this: *mut CometCScalarKernelImpl, + arg_types: *const *const FFI_ArrowSchema, + _scalar_args: *const *mut FFI_ArrowArray, + n_args: i64, + out: *mut FFI_ArrowSchema, +) -> c_int { + debug_assert!(!this.is_null()); + let this_ref = unsafe { &mut *this }; + debug_assert!(this_ref.release.is_some(), "init after release"); + let priv_ptr = this_ref.private_data as *mut ExportedScalarKernelImpl; + debug_assert!(!priv_ptr.is_null()); + let priv_ref = unsafe { &mut *priv_ptr }; + + let n = n_args as usize; + let mut fields = Vec::with_capacity(n); + for i in 0..n { + let schema_ptr = unsafe { *arg_types.add(i) }; + if schema_ptr.is_null() { + priv_ref.last_error = + std::ffi::CString::new(format!("arg #{i} has null FFI_ArrowSchema")) + .unwrap_or_default(); + return C_ABI_ERR; + } + let schema = unsafe { &*schema_ptr }; + match Field::try_from(schema) { + Ok(f) => fields.push(f), + Err(e) => { + priv_ref.last_error = + std::ffi::CString::new(format!("arg #{i}: {e}")).unwrap_or_default(); + return C_ABI_ERR; + } + } + } + + // `return_field` is user code: contain any panic (see "panic containment"). + match catch_panic(|| priv_ref.inner.return_field(&fields)) { + Ok(ret_field) => match FFI_ArrowSchema::try_from(&ret_field) { + Ok(ffi_schema) => { + unsafe { std::ptr::write(out, ffi_schema) }; + priv_ref.last_arg_fields = Some(fields); + priv_ref.last_return_field = Some(ret_field); + 0 + } + Err(e) => { + priv_ref.last_error = std::ffi::CString::new(format!("encoding return type: {e}")) + .unwrap_or_default(); + C_ABI_ERR + } + }, + Err(msg) => { + priv_ref.last_error = std::ffi::CString::new(msg).unwrap_or_default(); + C_ABI_ERR + } + } +} + +unsafe extern "C" fn c_kernel_execute( + this: *mut CometCScalarKernelImpl, + args: *const *mut FFI_ArrowArray, + n_args: i64, + n_rows: i64, + out: *mut FFI_ArrowArray, +) -> c_int { + debug_assert!(!this.is_null()); + let this_ref = unsafe { &mut *this }; + debug_assert!(this_ref.release.is_some(), "execute after release"); + let priv_ptr = this_ref.private_data as *mut ExportedScalarKernelImpl; + debug_assert!(!priv_ptr.is_null()); + let priv_ref = unsafe { &mut *priv_ptr }; + + let arg_fields = match priv_ref.last_arg_fields.as_ref() { + Some(f) => f, + None => { + priv_ref.last_error = + std::ffi::CString::new("execute called before init").unwrap_or_default(); + return C_ABI_ERR; + } + }; + if arg_fields.len() != n_args as usize { + priv_ref.last_error = std::ffi::CString::new(format!( + "execute n_args={} disagrees with init n_args={}", + n_args, + arg_fields.len() + )) + .unwrap_or_default(); + return C_ABI_ERR; + } + + // Take ownership of each input FFI_ArrowArray. + let n = n_args as usize; + let mut arrays: Vec = Vec::with_capacity(n); + for (i, arg_field) in arg_fields.iter().enumerate().take(n) { + let raw = unsafe { *args.add(i) }; + if raw.is_null() { + priv_ref.last_error = + std::ffi::CString::new(format!("arg #{i} FFI_ArrowArray is null")) + .unwrap_or_default(); + return C_ABI_ERR; + } + // SAFETY: raw points at an FFI_ArrowArray owned by the caller; we + // take ownership by reading and zeroing it. + let owned = unsafe { std::ptr::read(raw) }; + unsafe { std::ptr::write(raw, FFI_ArrowArray::empty()) }; + let dt = arg_field.data_type().clone(); + let data = match unsafe { arrow::ffi::from_ffi_and_data_type(owned, dt) } { + Ok(d) => d, + Err(e) => { + priv_ref.last_error = + std::ffi::CString::new(format!("arg #{i} from_ffi: {e}")).unwrap_or_default(); + return C_ABI_ERR; + } + }; + arrays.push(arrow::array::make_array(data)); + } + + // `invoke` is user code: contain any panic (see "panic containment"). + let result = match catch_panic(|| priv_ref.inner.invoke(&arrays, n_rows as usize)) { + Ok(arr) => arr, + Err(msg) => { + priv_ref.last_error = std::ffi::CString::new(msg).unwrap_or_default(); + return C_ABI_ERR; + } + }; + + let ffi_out = FFI_ArrowArray::new(&result.to_data()); + unsafe { std::ptr::write(out, ffi_out) }; + 0 +} + +unsafe extern "C" fn c_kernel_get_last_error(this: *mut CometCScalarKernelImpl) -> *const c_char { + debug_assert!(!this.is_null()); + let this_ref = unsafe { &*this }; + let priv_ptr = this_ref.private_data as *mut ExportedScalarKernelImpl; + if priv_ptr.is_null() { + return std::ptr::null(); + } + let priv_ref = unsafe { &*priv_ptr }; + priv_ref.last_error.as_ptr() +} + +unsafe extern "C" fn c_kernel_release(this: *mut CometCScalarKernelImpl) { + debug_assert!(!this.is_null()); + let this_ref = unsafe { &mut *this }; + if !this_ref.private_data.is_null() { + // Dropping may run user Drop code, which may panic; contain it. + let raw = this_ref.private_data as *mut ExportedScalarKernelImpl; + this_ref.private_data = std::ptr::null_mut(); + catch_panic_infallible("releasing kernel impl", || { + drop(unsafe { Box::from_raw(raw) }) + }); + } + this_ref.init = None; + this_ref.execute = None; + this_ref.get_last_error = None; + this_ref.release = None; +} + +// -- discovery list construction ------------------------------------------ + +/// Build a heap-allocated [`CometCScalarKernelList`] from a vector of +/// [`CometCScalarKernel`]s. Ownership transfers to the caller, who must +/// free via the list's `release` callback. +pub fn build_kernel_list(kernels: Vec) -> CometCScalarKernelList { + if kernels.is_empty() { + return CometCScalarKernelList::default(); + } + let mut boxed = kernels.into_boxed_slice(); + let len = boxed.len() as i64; + let kernels_ptr = boxed.as_mut_ptr(); + std::mem::forget(boxed); + CometCScalarKernelList { + kernels: kernels_ptr, + len, + release: Some(c_list_release), + } +} + +unsafe extern "C" fn c_list_release(list: *mut CometCScalarKernelList) { + debug_assert!(!list.is_null()); + let list_ref = unsafe { &mut *list }; + if list_ref.kernels.is_null() || list_ref.len == 0 { + list_ref.release = None; + return; + } + let len = list_ref.len as usize; + let kernels = list_ref.kernels; + list_ref.kernels = std::ptr::null_mut(); + list_ref.len = 0; + list_ref.release = None; + // SAFETY: kernels was a Box<[CometCScalarKernel]> turned into raw ptr + + // forgotten in build_kernel_list; reconstruct and drop. Each kernel's + // own Drop runs its `release` callback, which may reach user Drop code, + // so contain any panic rather than abort. + catch_panic_infallible("releasing kernel list", || { + drop(unsafe { Box::from_raw(std::ptr::slice_from_raw_parts_mut(kernels, len)) }) + }); +} + +// -- export macro --------------------------------------------------------- + +/// Emit the C-ABI discovery entry points for a list of UDF types. +/// +/// Each type passed must implement [`CometCScalarUdf`] and `Default`. The +/// macro produces: +/// +/// - `extern "C" fn comet_udf_abi_version() -> u32` +/// - `extern "C" fn comet_c_udf_list_v1(out: *mut CometCScalarKernelList) -> i32` +/// +/// Construction runs your `Default` impl and reads your `name()`, so it is +/// wrapped in a panic guard: a panic there is reported to the host as a +/// load failure rather than aborting the process. +/// +/// Your `Cargo.toml` must declare `crate-type = ["cdylib"]`. +#[macro_export] +macro_rules! comet_c_udf_export { + ( $( $ty:ty ),+ $(,)? ) => { + const _: () = { + #[no_mangle] + pub extern "C" fn comet_udf_abi_version() -> u32 { + $crate::COMET_UDF_ABI_VERSION + } + + #[no_mangle] + pub unsafe extern "C" fn comet_c_udf_list_v1( + out: *mut $crate::c_abi::CometCScalarKernelList, + ) -> i32 { + if out.is_null() { return -1; } + let built = std::panic::catch_unwind(|| { + let kernels: Vec<$crate::c_abi::CometCScalarKernel> = vec![ + $( + $crate::c_abi::CometCScalarKernel::from( + $crate::c_abi::ExportedScalarKernel::new( + <$ty as Default>::default() + ) + ), + )+ + ]; + $crate::c_abi::build_kernel_list(kernels) + }); + match built { + Ok(list) => { + std::ptr::write(out, list); + 0 + } + Err(_) => -1, + } + } + }; + }; +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ArrayRef, Int64Array}; + use arrow::datatypes::DataType; + use std::sync::Arc; + + struct AddOne; + impl CometCScalarUdf for AddOne { + fn name(&self) -> &str { + "add_one" + } + fn return_field(&self, args: &[Field]) -> Result { + if args.len() != 1 || args[0].data_type() != &DataType::Int64 { + return Err("expected (Int64) -> Int64".into()); + } + Ok(Field::new("add_one", DataType::Int64, true)) + } + fn invoke(&self, args: &[ArrayRef], _n: usize) -> Result { + let a = args[0] + .as_any() + .downcast_ref::() + .ok_or("not an Int64Array")?; + Ok(Arc::new( + a.iter().map(|v| v.map(|x| x + 1)).collect::(), + )) + } + } + + #[test] + fn adapter_roundtrip() { + let exp = ExportedScalarKernel::new(AddOne); + let kernel: CometCScalarKernel = exp.into(); + + // function_name lookup. + let name_ptr = unsafe { (kernel.function_name.unwrap())(&kernel) }; + let name = unsafe { std::ffi::CStr::from_ptr(name_ptr) }; + assert_eq!(name.to_str().unwrap(), "add_one"); + + // new_impl + init + execute. + let mut impl_state = CometCScalarKernelImpl::default(); + unsafe { + (kernel.new_impl.unwrap())(&kernel, &mut impl_state); + } + + let arg_field = Field::new("x", DataType::Int64, true); + let arg_schema = FFI_ArrowSchema::try_from(&arg_field).unwrap(); + let arg_schema_ptr: *const FFI_ArrowSchema = &arg_schema; + let mut out_schema = FFI_ArrowSchema::empty(); + let rc = unsafe { + (impl_state.init.unwrap())( + &mut impl_state, + &arg_schema_ptr as *const *const FFI_ArrowSchema, + std::ptr::null(), + 1, + &mut out_schema, + ) + }; + assert_eq!(rc, 0); + let out_field = Field::try_from(&out_schema).unwrap(); + assert_eq!(out_field.data_type(), &DataType::Int64); + + // execute. + let input: Arc = Arc::new(Int64Array::from(vec![1, 2, 3])); + let mut input_ffi = FFI_ArrowArray::new(&input.to_data()); + let input_ffi_ptr: *mut FFI_ArrowArray = &mut input_ffi; + let mut out_arr = FFI_ArrowArray::empty(); + let rc = unsafe { + (impl_state.execute.unwrap())( + &mut impl_state, + &input_ffi_ptr as *const *mut FFI_ArrowArray, + 1, + 3, + &mut out_arr, + ) + }; + assert_eq!(rc, 0); + let result_data = + unsafe { arrow::ffi::from_ffi_and_data_type(out_arr, DataType::Int64) }.unwrap(); + let result = arrow::array::make_array(result_data); + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result.values(), &[2, 3, 4]); + + // Releasing impl_state and kernel runs cleanup callbacks; ensure + // their function pointers are cleared. + drop(impl_state); + drop(kernel); + } + + /// A UDF that panics wherever it is told to. + struct Panicky { + in_return_field: bool, + } + + impl CometCScalarUdf for Panicky { + fn name(&self) -> &str { + "panicky" + } + fn return_field(&self, _args: &[Field]) -> Result { + assert!(!self.in_return_field, "boom in return_field"); + Ok(Field::new("panicky", DataType::Int64, true)) + } + fn invoke(&self, _args: &[ArrayRef], _n: usize) -> Result { + panic!("boom in invoke") + } + } + + /// Build a kernel + initialized impl for `udf`, returning the init rc. + fn init_panicky(in_return_field: bool) -> (CometCScalarKernel, CometCScalarKernelImpl, c_int) { + let kernel: CometCScalarKernel = + ExportedScalarKernel::new(Panicky { in_return_field }).into(); + let mut impl_state = CometCScalarKernelImpl::default(); + unsafe { + (kernel.new_impl.unwrap())(&kernel, &mut impl_state); + } + let arg_field = Field::new("x", DataType::Int64, true); + let arg_schema = FFI_ArrowSchema::try_from(&arg_field).unwrap(); + let arg_schema_ptr: *const FFI_ArrowSchema = &arg_schema; + let mut out_schema = FFI_ArrowSchema::empty(); + let rc = unsafe { + (impl_state.init.unwrap())( + &mut impl_state, + &arg_schema_ptr as *const *const FFI_ArrowSchema, + std::ptr::null(), + 1, + &mut out_schema, + ) + }; + (kernel, impl_state, rc) + } + + fn last_error(impl_state: &mut CometCScalarKernelImpl) -> String { + let ptr = unsafe { (impl_state.get_last_error.unwrap())(impl_state) }; + assert!(!ptr.is_null(), "expected an error message"); + unsafe { std::ffi::CStr::from_ptr(ptr) } + .to_str() + .unwrap() + .to_string() + } + + /// A panic in `return_field` must come back as an error code plus a + /// message, not unwind across the `extern "C"` boundary. + #[test] + fn panic_in_return_field_is_contained() { + let (_kernel, mut impl_state, rc) = init_panicky(true); + assert_eq!(rc, C_ABI_ERR); + let msg = last_error(&mut impl_state); + assert!(msg.contains("panic in UDF code"), "msg: {msg}"); + assert!(msg.contains("boom in return_field"), "msg: {msg}"); + } + + /// Same for a panic in `invoke`. + #[test] + fn panic_in_invoke_is_contained() { + let (_kernel, mut impl_state, rc) = init_panicky(false); + assert_eq!(rc, 0, "init should succeed for this case"); + + let input: Arc = Arc::new(Int64Array::from(vec![1, 2, 3])); + let mut input_ffi = FFI_ArrowArray::new(&input.to_data()); + let input_ffi_ptr: *mut FFI_ArrowArray = &mut input_ffi; + let mut out_arr = FFI_ArrowArray::empty(); + let rc = unsafe { + (impl_state.execute.unwrap())( + &mut impl_state, + &input_ffi_ptr as *const *mut FFI_ArrowArray, + 1, + 3, + &mut out_arr, + ) + }; + assert_eq!(rc, C_ABI_ERR); + let msg = last_error(&mut impl_state); + assert!(msg.contains("panic in UDF code"), "msg: {msg}"); + assert!(msg.contains("boom in invoke"), "msg: {msg}"); + } +} diff --git a/native/comet-udf-sdk/src/lib.rs b/native/comet-udf-sdk/src/lib.rs new file mode 100644 index 0000000000..5920141685 --- /dev/null +++ b/native/comet-udf-sdk/src/lib.rs @@ -0,0 +1,73 @@ +// 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. + +//! SDK for writing scalar UDFs in Rust that are loaded and executed by +//! Apache DataFusion Comet, using only Arrow's stable FFI surface. +//! +//! The ABI is a pure C-callable struct of function pointers built only on +//! the Arrow C Data Interface (`FFI_ArrowSchema` / `FFI_ArrowArray`), +//! modeled on Apache Sedona's `SedonaCScalarKernel`. See [`c_abi`] for the +//! authoring guide. +//! +//! # Why not `datafusion-ffi`? +//! +//! Wrapping the user's `ScalarUDFImpl` as `datafusion_ffi::udf::FFI_ScalarUDF` +//! is the obvious alternative, and it hands the author a much larger surface +//! for free: variadic signatures, type coercion, metadata-aware return types. +//! Comet deliberately does not expose it, for one reason: it would couple +//! every user's cdylib to Comet's DataFusion major version. +//! +//! That is not a hypothetical cost. A prototype carried both ABIs side by +//! side; upgrading Comet from DataFusion 53 to 54 removed `as_any` from +//! `ScalarUDFImpl`, which would have forced an edit and a recompile on every +//! user library built against the `datafusion-ffi` flavor. The C ABI here +//! needed no change, because no DataFusion type appears in it. +//! +//! Comet tracks DataFusion closely and upgrades often, so a UDF ABI pinned to +//! the DataFusion version would break users on a cadence they do not control +//! and cannot opt out of. Keeping the FFI surface to Arrow's C Data Interface +//! means a UDF compiled today keeps working across Comet upgrades, and the +//! same ABI is implementable from C, C++, or any language that speaks the +//! Arrow C Data Interface. +//! +//! The tradeoff is a smaller surface: authors implement [`c_abi::CometCScalarUdf`] +//! and get scalar functions over Arrow arrays, not the full `ScalarUDFImpl` +//! feature set. + +#![warn(missing_docs)] + +/// Discovery ABI version. Bumped on any backwards-incompatible change to +/// the discovery entry-point signatures or to the FFI structs they yield. +pub const COMET_UDF_ABI_VERSION: u32 = 1; + +pub mod c_abi; + +/// Symbol name of the discovery entry point exported by every cdylib. +pub const C_ABI_DISCOVERY_SYMBOL: &str = "comet_c_udf_list_v1"; + +/// Symbol name of the ABI version probe exported by every cdylib. +pub const ABI_VERSION_SYMBOL: &str = "comet_udf_abi_version"; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn abi_version_is_one() { + assert_eq!(COMET_UDF_ABI_VERSION, 1); + } +} diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 7c88b260c3..775a1c769f 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -77,6 +77,8 @@ iceberg-storage-opendal = { workspace = true } reqsign-core = { workspace = true } serde_json = "1.0" uuid = "1.23.3" +libloading = "0.8" +comet-udf-sdk = { path = "../comet-udf-sdk" } [target.'cfg(target_os = "linux")'.dependencies] procfs = "0.18.0" @@ -92,6 +94,7 @@ lazy_static = "1.4" assertables = "10" hex = "0.4.3" datafusion-functions-nested = { version = "54.1.0" } +comet-test-udfs = { path = "../comet-test-udfs" } [features] backtrace = ["datafusion/backtrace"] diff --git a/native/core/build.rs b/native/core/build.rs index bfda157b9a..da42b85a5c 100644 --- a/native/core/build.rs +++ b/native/core/build.rs @@ -40,4 +40,29 @@ fn main() { println!("cargo:rustc-link-search=native={server}"); } } + + // Expose the path of the comet-test-udfs cdylib to test code via + // COMET_TEST_UDFS_LIB. Cargo doesn't propagate cdylib outputs as + // DEP_<...>_OUT_DIR for non-rlib crates, so we compute the path + // from OUT_DIR. + if let Ok(out_dir) = std::env::var("OUT_DIR") { + let out_path = std::path::PathBuf::from(out_dir); + // OUT_DIR is .../target//build//out + let target_profile_dir = out_path + .ancestors() + .nth(3) + .map(|p| p.to_path_buf()) + .unwrap_or_default(); + let dylib_ext = if cfg!(target_os = "macos") { + "dylib" + } else if cfg!(target_os = "windows") { + "dll" + } else { + "so" + }; + let lib_path = target_profile_dir.join(format!("libcomet_test_udfs.{dylib_ext}")); + println!("cargo:rustc-env=COMET_TEST_UDFS_LIB={}", lib_path.display()); + } + println!("cargo:rerun-if-changed=../comet-test-udfs/src/lib.rs"); + println!("cargo:rerun-if-changed=../comet-test-udfs/Cargo.toml"); } diff --git a/native/core/src/comet_rust_udf_bridge.rs b/native/core/src/comet_rust_udf_bridge.rs new file mode 100644 index 0000000000..533e076709 --- /dev/null +++ b/native/core/src/comet_rust_udf_bridge.rs @@ -0,0 +1,89 @@ +// 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. + +//! JNI entry points for driver-side validation of Rust UDF cdylibs. +//! Used by `org.apache.comet.udf.CometRustUdfBridge` on the driver. + +use crate::errors::{try_unwrap_or_throw, CometError}; +use crate::execution::rust_udf::cache::get_or_load; +use crate::execution::rust_udf::loader::LoadedUdf; +use jni::objects::{JClass, JString}; +use jni::sys::jobject; +use jni::EnvUnowned; + +/// Best-effort serialization of a single discovered UDF as JSON. +/// +/// Reporting `args`/`return_type` would require calling the kernel's +/// `init` to discover a return type; that is deferred, so this returns +/// only what the Scala registry needs today (`name`). +fn udf_to_json(udf: &LoadedUdf) -> serde_json::Value { + serde_json::json!({ + "name": udf.name, + }) +} + +/// Validate that `library_path` loads, exposes a UDF named +/// `expected_name`, and return a JSON description of that UDF. Throws +/// on any error. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_udf_CometRustUdfBridge_validateLibrary( + e: EnvUnowned, + _class: JClass, + library_path: JString, + expected_name: JString, +) -> jobject { + try_unwrap_or_throw(&e, |env| { + let path: String = library_path + .try_to_string(env) + .map_err(|e| CometError::Internal(e.to_string()))?; + let name: String = expected_name + .try_to_string(env) + .map_err(|e| CometError::Internal(e.to_string()))?; + let lib = get_or_load(&path).map_err(|e| CometError::Internal(e.to_string()))?; + let udf = lib + .udfs + .iter() + .find(|u| u.name == name) + .ok_or_else(|| CometError::Internal(format!("UDF '{name}' not found in {path}")))?; + let json = udf_to_json(udf).to_string(); + let jstr = env + .new_string(json) + .map_err(|e| CometError::Internal(e.to_string()))?; + Ok(jstr.into_raw()) + }) +} + +/// Return a JSON array describing every UDF exposed by `library_path`. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_udf_CometRustUdfBridge_listUdfs( + e: EnvUnowned, + _class: JClass, + library_path: JString, +) -> jobject { + try_unwrap_or_throw(&e, |env| { + let path: String = library_path + .try_to_string(env) + .map_err(|e| CometError::Internal(e.to_string()))?; + let lib = get_or_load(&path).map_err(|e| CometError::Internal(e.to_string()))?; + let entries: Vec = lib.udfs.iter().map(udf_to_json).collect(); + let json = serde_json::Value::Array(entries).to_string(); + let jstr = env + .new_string(json) + .map_err(|e| CometError::Internal(e.to_string()))?; + Ok(jstr.into_raw()) + }) +} diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs index ec247f72b7..9291eafdaa 100644 --- a/native/core/src/execution/mod.rs +++ b/native/core/src/execution/mod.rs @@ -23,6 +23,7 @@ pub(crate) mod merge_as_partial; pub(crate) mod metrics; pub mod operators; pub(crate) mod planner; +pub mod rust_udf; pub mod serde; pub use datafusion_comet_shuffle as shuffle; pub(crate) mod sort; diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index f20dadf7f3..806dea6edc 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -890,6 +890,73 @@ impl PhysicalPlanner { self.task_context.clone(), ))) } + ExprStruct::RustUdfCall(call) => { + let arg_exprs: Vec> = call + .args + .iter() + .map(|e| self.create_expr(e, Arc::clone(&input_schema))) + .collect::, _>>()?; + + let lib = crate::execution::rust_udf::cache::get_or_load(&call.library_path) + .map_err(|e| { + GeneralError(format!("Rust UDF load '{}': {e}", call.library_path)) + })?; + + let loaded = lib + .udfs + .iter() + .find(|u| u.name == call.name) + .ok_or_else(|| { + GeneralError(format!( + "Rust UDF '{}' not found in '{}'", + call.name, call.library_path + )) + })?; + + let udf = Arc::new(ScalarUDF::new_from_shared_impl(Arc::clone( + &loaded.udf_impl, + ))); + + let return_type = to_arrow_datatype( + call.return_type + .as_ref() + .ok_or_else(|| GeneralError("RustUdfCall missing return_type".into()))?, + ); + + // The declared return type comes from the JVM-side `CometRustUDF.register` call + // and is what Spark planned against; the kernel's own `return_field` is what will + // actually be produced. If they disagree, fail here with both types named rather + // than letting it surface later as a bare type assertion mid-execution. + let arg_types = arg_exprs + .iter() + .map(|e| e.data_type(input_schema.as_ref())) + .collect::, _>>()?; + let kernel_return_type = loaded.udf_impl.return_type(&arg_types)?; + if !crate::execution::rust_udf::return_types_compatible( + &return_type, + &kernel_return_type, + ) { + return Err(GeneralError(format!( + "Rust UDF '{}' was registered as returning {return_type} but its \ + return_field reports {kernel_return_type} for argument types {arg_types:?}. \ + Make the type passed to CometRustUDF.register match what the UDF returns.", + call.name + ))); + } + + // Promise DataFusion the kernel's own type rather than the declared one. The two + // agree up to nested nullability, and using the kernel's avoids tripping + // DataFusion's exact-match assertion on the returned batch. + let return_field = Arc::new(Field::new(&call.name, kernel_return_type, true)); + let expr = Arc::new(ScalarFunctionExpr::new( + &call.name, + udf, + arg_exprs, + return_field, + Arc::new(ConfigOptions::default()), + )); + Ok(expr) + } expr => Err(GeneralError(format!("Not implemented: {expr:?}"))), } } diff --git a/native/core/src/execution/rust_udf/cache.rs b/native/core/src/execution/rust_udf/cache.rs new file mode 100644 index 0000000000..c09449411e --- /dev/null +++ b/native/core/src/execution/rust_udf/cache.rs @@ -0,0 +1,87 @@ +// 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. + +//! Process-wide cache of loaded UDF cdylibs. +//! +//! Same-path lookups always return the same `Arc` for +//! the lifetime of the process — libraries are deliberately never +//! unloaded. Calling `dlclose` while a thread is mid-call would be a +//! use-after-free, and there is no safe point to unload without +//! per-invocation refcounting we don't want on the hot path. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, OnceLock, RwLock}; + +use super::loader::{load, LoadedLibrary, LoaderError}; + +static CACHE: OnceLock>>> = OnceLock::new(); + +fn cache() -> &'static RwLock>> { + CACHE.get_or_init(|| RwLock::new(HashMap::new())) +} + +/// Get an already-loaded library, or load and cache it. +pub fn get_or_load(path: impl AsRef) -> Result, LoaderError> { + let raw = path.as_ref().to_path_buf(); + + if let Some(lib) = cache().read().unwrap().get(&raw).cloned() { + return Ok(lib); + } + + let canonical = raw.canonicalize().unwrap_or_else(|_| raw.clone()); + if canonical != raw { + if let Some(lib) = cache().read().unwrap().get(&canonical).cloned() { + cache().write().unwrap().insert(raw, Arc::clone(&lib)); + return Ok(lib); + } + } + + let mut w = cache().write().unwrap(); + if let Some(lib) = w.get(&canonical).cloned() { + if canonical != raw { + w.insert(raw, Arc::clone(&lib)); + } + return Ok(lib); + } + let loaded = Arc::new(load(&canonical)?); + w.insert(canonical.clone(), Arc::clone(&loaded)); + if canonical != raw { + w.insert(raw, Arc::clone(&loaded)); + } + Ok(loaded) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::execution::rust_udf::test_support::{test_udfs_path, BUILD_HINT}; + + #[test] + fn same_path_returns_same_arc() { + let p = test_udfs_path(); + let a = get_or_load(&p).expect(BUILD_HINT); + let b = get_or_load(&p).expect(BUILD_HINT); + assert!(Arc::ptr_eq(&a, &b)); + } + + #[test] + fn missing_path_propagates_error() { + let err = get_or_load("/no/such/file.dylib").unwrap_err(); + assert!(matches!(err, LoaderError::Open { .. })); + } +} diff --git a/native/core/src/execution/rust_udf/imported_c.rs b/native/core/src/execution/rust_udf/imported_c.rs new file mode 100644 index 0000000000..353d117d9c --- /dev/null +++ b/native/core/src/execution/rust_udf/imported_c.rs @@ -0,0 +1,311 @@ +// 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. + +//! Adapter wrapping a C-ABI [`CometCScalarKernel`] as a DataFusion +//! [`ScalarUDFImpl`]. +//! +//! Lifecycle inside `invoke_with_args`: +//! +//! 1. Build a fresh [`CometCScalarKernelImpl`] via the kernel's `new_impl`. +//! 2. Call `init` with the input field types (and any scalar args) to get +//! the return type. +//! 3. Call `execute` once with the batch. +//! 4. Drop the impl (its `release` callback runs). + +use std::ffi::CStr; +use std::sync::Mutex; + +use arrow::array::ArrayRef; +use arrow::datatypes::{DataType, Field}; +use arrow::ffi::{from_ffi_and_data_type, FFI_ArrowArray, FFI_ArrowSchema}; +use comet_udf_sdk::c_abi::{CometCScalarKernel, CometCScalarKernelImpl}; +use datafusion::common::DataFusionError; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, +}; + +/// Adapter wrapping a [`CometCScalarKernel`] as a DataFusion +/// [`ScalarUDFImpl`]. +pub struct ImportedCScalarUdf { + name: String, + /// Boxed so the kernel's address is stable; held inside a Mutex + /// because the FFI Drop is not Sync-safe under concurrent invocation. + /// The kernel itself is logically immutable post-load — the lock only + /// protects the FFI calls' aliasing rules. (DataFusion serializes + /// invocations of a given ScalarUDFImpl per-batch through + /// invoke_with_args anyway; the lock is defensive.) + kernel: Mutex>, + signature: Signature, +} + +impl ImportedCScalarUdf { + /// Construct from an owned C kernel. + /// + /// Reads the kernel's name via its `function_name` callback and + /// stores it for `name()` lookups; the kernel itself is held inside + /// a mutex. + pub fn try_new(kernel: Box) -> Result { + let function_name_cb = kernel + .function_name + .ok_or_else(|| "kernel.function_name is null".to_string())?; + let _ = kernel + .new_impl + .ok_or_else(|| "kernel.new_impl is null".to_string())?; + + // SAFETY: function_name_cb is the FFI-supplied callback; + // implementations promise the returned pointer is a NUL-terminated + // UTF-8 string valid for the lifetime of the kernel. + let name_ptr = unsafe { function_name_cb(kernel.as_ref() as *const _) }; + if name_ptr.is_null() { + return Err("function_name returned null".into()); + } + let name = unsafe { CStr::from_ptr(name_ptr) } + .to_str() + .map_err(|e| format!("function_name not UTF-8: {e}"))? + .to_string(); + + // Use UserDefined signature: per-call init() is what decides + // whether the input types are acceptable. `coerce_types` is not + // implemented; user must pass exact types from the JVM register call. + // + // Volatility is always Immutable. The signature is built once per + // library load, while determinism is declared per registration, so + // the two do not line up: `CometRustUDF.register` rejects + // `deterministic = false` rather than let a volatile function be + // planned as if it were pure. + let signature = Signature::new(TypeSignature::UserDefined, Volatility::Immutable); + + Ok(Self { + name, + kernel: Mutex::new(kernel), + signature, + }) + } +} + +impl std::fmt::Debug for ImportedCScalarUdf { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ImportedCScalarUdf") + .field("name", &self.name) + .finish() + } +} + +impl PartialEq for ImportedCScalarUdf { + fn eq(&self, other: &Self) -> bool { + self.name == other.name + } +} + +impl Eq for ImportedCScalarUdf {} + +impl std::hash::Hash for ImportedCScalarUdf { + fn hash(&self, state: &mut H) { + self.name.hash(state); + } +} + +impl ScalarUDFImpl for ImportedCScalarUdf { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, args: &[DataType]) -> datafusion::common::Result { + // Build a fresh impl, call init, drop. Done at planning time so + // the planner can know the output type before execution. + let kernel = self.kernel.lock().unwrap(); + let mut impl_state = CometCScalarKernelImpl::default(); + let new_impl_cb = kernel + .new_impl + .ok_or_else(|| DataFusionError::Internal("new_impl is null".into()))?; + // SAFETY: new_impl_cb is the FFI-supplied factory; impl_state is a + // caller-allocated default value the cdylib writes into. + unsafe { + new_impl_cb(kernel.as_ref() as *const _, &mut impl_state); + } + + // Build input fields and FFI schemas. + let fields: Vec = args + .iter() + .map(|dt| Field::new("", dt.clone(), true)) + .collect(); + let ffi_schemas = build_ffi_schemas(&fields)?; + let ffi_schema_ptrs: Vec<*const FFI_ArrowSchema> = + ffi_schemas.iter().map(|s| s as *const _).collect(); + + let init_cb = impl_state + .init + .ok_or_else(|| DataFusionError::Internal("kernel impl missing init".into()))?; + let mut out_schema = FFI_ArrowSchema::empty(); + // SAFETY: pointers are valid for the duration of the call. + let rc = unsafe { + init_cb( + &mut impl_state, + ffi_schema_ptrs.as_ptr(), + std::ptr::null(), + fields.len() as i64, + &mut out_schema, + ) + }; + if rc != 0 { + let msg = read_last_error(&mut impl_state); + return Err(DataFusionError::Plan(format!( + "{}: init failed: {msg}", + self.name + ))); + } + let return_field = Field::try_from(&out_schema) + .map_err(|e| DataFusionError::Internal(format!("decoding return type: {e}")))?; + Ok(return_field.data_type().clone()) + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion::common::Result { + let n_rows = args.number_rows; + let kernel = self.kernel.lock().unwrap(); + + // Build a fresh impl_state; init then execute. + let new_impl_cb = kernel + .new_impl + .ok_or_else(|| DataFusionError::Internal("new_impl is null".into()))?; + let mut impl_state = CometCScalarKernelImpl::default(); + // SAFETY: see return_type. + unsafe { + new_impl_cb(kernel.as_ref() as *const _, &mut impl_state); + } + + // Resolve args to Arrays of length n_rows or 1. + let mut arrays: Vec = Vec::with_capacity(args.args.len()); + for a in args.args { + let arr = match a { + ColumnarValue::Array(arr) => arr, + ColumnarValue::Scalar(s) => s.to_array_of_size(n_rows)?, + }; + arrays.push(arr); + } + + // Build input fields + schemas (the kernel needs init to remember + // the arg types for execute). + let fields: Vec = arrays + .iter() + .map(|a| Field::new("", a.data_type().clone(), true)) + .collect(); + let ffi_schemas = build_ffi_schemas(&fields)?; + let ffi_schema_ptrs: Vec<*const FFI_ArrowSchema> = + ffi_schemas.iter().map(|s| s as *const _).collect(); + + let init_cb = impl_state + .init + .ok_or_else(|| DataFusionError::Internal("kernel impl missing init".into()))?; + let mut out_schema = FFI_ArrowSchema::empty(); + // SAFETY: ffi_schema_ptrs lives for the duration of this call. + let rc = unsafe { + init_cb( + &mut impl_state, + ffi_schema_ptrs.as_ptr(), + std::ptr::null(), + fields.len() as i64, + &mut out_schema, + ) + }; + if rc != 0 { + let msg = read_last_error(&mut impl_state); + return Err(DataFusionError::Execution(format!( + "{}: init failed: {msg}", + self.name + ))); + } + let return_field = Field::try_from(&out_schema) + .map_err(|e| DataFusionError::Internal(format!("decoding return type: {e}")))?; + + // Build FFI arrays. + let mut ffi_arrays: Vec = arrays + .iter() + .map(|a| FFI_ArrowArray::new(&a.to_data())) + .collect(); + let ffi_array_ptrs: Vec<*mut FFI_ArrowArray> = + ffi_arrays.iter_mut().map(|x| x as *mut _).collect(); + + let execute_cb = impl_state + .execute + .ok_or_else(|| DataFusionError::Internal("kernel impl missing execute".into()))?; + let mut out_arr = FFI_ArrowArray::empty(); + // SAFETY: ffi_array_ptrs live for the duration of this call. The + // kernel takes ownership of each input by replacing it with an + // empty FFI_ArrowArray (no-op Drop). + let rc = unsafe { + execute_cb( + &mut impl_state, + ffi_array_ptrs.as_ptr(), + arrays.len() as i64, + n_rows as i64, + &mut out_arr, + ) + }; + + if rc != 0 { + let msg = read_last_error(&mut impl_state); + return Err(DataFusionError::Execution(format!( + "{}: execute failed: {msg}", + self.name + ))); + } + + // Import result. + // SAFETY: out_arr was filled by the cdylib. + let data = unsafe { from_ffi_and_data_type(out_arr, return_field.data_type().clone()) } + .map_err(|e| DataFusionError::Execution(format!("from_ffi: {e}")))?; + let array = arrow::array::make_array(data); + if array.len() != n_rows { + return Err(DataFusionError::Execution(format!( + "{}: returned {} rows, expected {n_rows}", + self.name, + array.len() + ))); + } + Ok(ColumnarValue::Array(array)) + } +} + +fn build_ffi_schemas(fields: &[Field]) -> datafusion::common::Result> { + fields + .iter() + .map(FFI_ArrowSchema::try_from) + .collect::, _>>() + .map_err(|e| DataFusionError::Internal(format!("encoding arg type: {e}"))) +} + +fn read_last_error(impl_state: &mut CometCScalarKernelImpl) -> String { + let cb = match impl_state.get_last_error { + Some(cb) => cb, + None => return "(no get_last_error)".to_string(), + }; + // SAFETY: cb is the FFI-supplied callback. + let ptr = unsafe { cb(impl_state) }; + if ptr.is_null() { + return "(empty)".to_string(); + } + unsafe { CStr::from_ptr(ptr) } + .to_string_lossy() + .into_owned() +} diff --git a/native/core/src/execution/rust_udf/loader.rs b/native/core/src/execution/rust_udf/loader.rs new file mode 100644 index 0000000000..fce329b15d --- /dev/null +++ b/native/core/src/execution/rust_udf/loader.rs @@ -0,0 +1,280 @@ +// 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. + +//! Loader: open a UDF cdylib via libloading, validate the ABI version, +//! discover UDFs via the C-ABI entry point, and produce DataFusion +//! `ScalarUDFImpl` impls for each. +//! +//! See `super::mod.rs` for an overview of the ABI. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use comet_udf_sdk::c_abi::CometCScalarKernelList; +use comet_udf_sdk::{ABI_VERSION_SYMBOL, COMET_UDF_ABI_VERSION, C_ABI_DISCOVERY_SYMBOL}; +use datafusion::logical_expr::ScalarUDFImpl; +use libloading::{Library, Symbol}; + +use super::imported_c::ImportedCScalarUdf; + +/// One loaded UDF: name plus a `ScalarUDFImpl` ready to plug into the +/// planner. +pub struct LoadedUdf { + /// UDF name as exposed by the cdylib. + pub name: String, + /// The `ScalarUDFImpl` adapter the planner will wrap. + pub udf_impl: Arc, +} + +impl std::fmt::Debug for LoadedUdf { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LoadedUdf") + .field("name", &self.name) + .finish() + } +} + +/// Result of loading a UDF cdylib: the live `Library` plus per-UDF +/// adapters. +pub struct LoadedLibrary { + /// Canonicalized path the library was loaded from. + pub path: PathBuf, + /// One entry per UDF, with name and ScalarUDFImpl already built. + /// + /// Declared before `library` on purpose. Struct fields drop in + /// declaration order, and each UDF's drop calls a `release` callback + /// that lives in the library's text: unloading first would call + /// through a dangling pointer. + pub udfs: Vec, + /// The loaded `Library`. Held inside an `Arc` so loaded UDFs can + /// outlive lookups. Library is never unloaded for the process lifetime. + pub library: Arc, +} + +impl std::fmt::Debug for LoadedLibrary { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LoadedLibrary") + .field("path", &self.path) + .field("udfs", &self.udfs) + .finish() + } +} + +/// Errors returned by the loader. +#[derive(Debug)] +pub enum LoaderError { + /// `libloading::Library::new` failed. + Open { + /// Path that was passed to `Library::new`. + path: PathBuf, + /// Underlying error. + source: libloading::Error, + }, + /// `comet_udf_abi_version` is missing or returned an unexpected value. + AbiMismatch { + /// Path of the offending library. + path: PathBuf, + /// Version reported by the cdylib (or `None` if the symbol is missing). + found: Option, + /// Version this host expects. + expected: u32, + }, + /// Library does not expose `comet_c_udf_list_v1`. + NoDiscovery { + /// Path of the offending library. + path: PathBuf, + }, + /// The discovery function returned a non-zero rc, or a kernel entry + /// was malformed. + Discovery { + /// Path of the library. + path: PathBuf, + /// Human-readable reason. + reason: String, + }, +} + +impl std::fmt::Display for LoaderError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + use LoaderError::*; + match self { + Open { path, source } => write!(f, "failed to open {}: {source}", path.display()), + AbiMismatch { + path, + found, + expected, + } => match found { + Some(v) => write!( + f, + "{} reports ABI v{v}, host expects v{expected}", + path.display() + ), + None => write!( + f, + "{} missing required symbol {ABI_VERSION_SYMBOL}", + path.display() + ), + }, + NoDiscovery { path } => write!( + f, + "{} does not export {C_ABI_DISCOVERY_SYMBOL}", + path.display() + ), + Discovery { path, reason } => write!(f, "{}: {reason}", path.display()), + } + } +} + +impl std::error::Error for LoaderError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + LoaderError::Open { source, .. } => Some(source), + _ => None, + } + } +} + +/// Open and validate a UDF cdylib. +pub fn load(path: impl AsRef) -> Result { + let path = path.as_ref().to_path_buf(); + // SAFETY: `Library::new` runs the cdylib's static initializers. We + // accept this risk because user UDF cdylibs are explicitly registered + // by an operator via `CometRustUDF.register`. + let library = unsafe { Library::new(&path) }.map_err(|source| LoaderError::Open { + path: path.clone(), + source, + })?; + + // ABI version probe. + let v = read_abi_version(&library, &path)?; + if v != COMET_UDF_ABI_VERSION { + return Err(LoaderError::AbiMismatch { + path, + found: Some(v), + expected: COMET_UDF_ABI_VERSION, + }); + } + + let udfs = match read_c_kernels(&library, &path)? { + Some(udfs) => udfs, + None => return Err(LoaderError::NoDiscovery { path }), + }; + + Ok(LoadedLibrary { + path, + library: Arc::new(library), + udfs, + }) +} + +fn read_abi_version(lib: &Library, path: &Path) -> Result { + let sym: Symbol u32> = unsafe { + lib.get(ABI_VERSION_SYMBOL.as_bytes()) + } + .map_err(|_| LoaderError::AbiMismatch { + path: path.to_path_buf(), + found: None, + expected: COMET_UDF_ABI_VERSION, + })?; + // SAFETY: comet_udf_abi_version takes no arguments, returns u32, no side effects. + Ok(unsafe { sym() }) +} + +fn read_c_kernels(lib: &Library, path: &Path) -> Result>, LoaderError> { + let sym: Symbol i32> = + match unsafe { lib.get(C_ABI_DISCOVERY_SYMBOL.as_bytes()) } { + Ok(s) => s, + Err(_) => return Ok(None), + }; + let mut list = CometCScalarKernelList::default(); + // SAFETY: list is caller-allocated; the cdylib writes into it via `out`. + let rc = unsafe { sym(&mut list) }; + if rc != 0 { + return Err(LoaderError::Discovery { + path: path.to_path_buf(), + reason: format!("{C_ABI_DISCOVERY_SYMBOL} returned rc={rc}"), + }); + } + let mut udfs = Vec::with_capacity(list.len.max(0) as usize); + if !list.kernels.is_null() && list.len > 0 { + // Move each kernel out of the array into a Box so it owns itself. + // We can't simply read each entry because they implement Drop; + // doing it via std::ptr::read transfers ownership cleanly. + let len = list.len as usize; + for i in 0..len { + // SAFETY: the kernel array was produced by the cdylib's + // `comet_c_udf_export!` and contains `len` valid entries. + // We move each entry out into a Box so its Drop runs when + // the host releases the loaded library. + let raw = unsafe { list.kernels.add(i) }; + let kernel = unsafe { std::ptr::read(raw) }; + // Replace the slot with a default kernel (no callbacks) so + // the array's release doesn't double-free. + unsafe { + std::ptr::write(raw, comet_udf_sdk::c_abi::CometCScalarKernel::default()); + } + let imported = ImportedCScalarUdf::try_new(Box::new(kernel)).map_err(|e| { + LoaderError::Discovery { + path: path.to_path_buf(), + reason: format!("import C kernel idx={i}: {e}"), + } + })?; + udfs.push(LoadedUdf { + name: imported.name().to_string(), + udf_impl: Arc::new(imported), + }); + } + } + // list's Drop releases the array storage. + drop(list); + Ok(Some(udfs)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::execution::rust_udf::test_support::{test_udfs_path, BUILD_HINT}; + + #[test] + fn load_test_udfs_succeeds() { + let lib = load(test_udfs_path()).expect(BUILD_HINT); + let names: Vec<_> = lib.udfs.iter().map(|u| u.name.as_str()).collect(); + assert!(names.contains(&"add_one_c"), "names: {names:?}"); + } + + #[test] + fn missing_path_errors_open() { + let err = load("/no/such/path.so").unwrap_err(); + assert!(matches!(err, LoaderError::Open { .. }), "got: {err:?}"); + } + + /// A library exporting several kernels surfaces all of them. + #[test] + fn all_exported_kernels_are_discovered() { + let lib = load(test_udfs_path()).expect(BUILD_HINT); + let names: Vec<_> = lib.udfs.iter().map(|u| u.name.as_str()).collect(); + for expected in [ + "add_one_c", + "echo_c", + "stringify_c", + "panics_on_invoke", + "panics_on_return_field", + ] { + assert!(names.contains(&expected), "missing {expected} in {names:?}"); + } + } +} diff --git a/native/core/src/execution/rust_udf/mod.rs b/native/core/src/execution/rust_udf/mod.rs new file mode 100644 index 0000000000..3c9565ca82 --- /dev/null +++ b/native/core/src/execution/rust_udf/mod.rs @@ -0,0 +1,148 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Loader and adapters for user-supplied Rust UDF cdylibs registered +//! through `CometRustUDF` on the JVM side. +//! +//! The ABI is built only on Arrow's stable FFI (the C Data Interface): +//! `comet_c_udf_list_v1` returns sedona-style `CometCScalarKernel` +//! factory structs. No DataFusion type appears in the FFI surface, so a +//! user cdylib is not coupled to Comet's DataFusion version and the same +//! ABI is implementable from C or C++. +//! +//! See `comet_udf_sdk` for the rationale behind not exposing +//! `datafusion-ffi` here. + +pub mod cache; +pub mod imported_c; +pub mod loader; + +#[cfg(test)] +pub(crate) mod test_support; + +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, FieldRef}; + +/// Rewrite every nested field of `dt` as nullable. +/// +/// Spark carries `containsNull` / field nullability inside the declared type, so +/// `array` with `containsNull = false` converts to an Arrow `List` whose child field is +/// non-nullable. The arrays Comet actually hands to a UDF normalize those child fields to +/// nullable, so the declared type and the delivered type disagree on nested nullability even when +/// the UDF is behaving perfectly. +/// +/// Used only to compare two types; never to build a type that gets handed to Arrow, which matters +/// because some nested types (map keys) require a non-nullable child. +fn erase_nested_nullability(dt: &DataType) -> DataType { + fn erase(f: &FieldRef) -> FieldRef { + Arc::new(Field::new( + f.name(), + erase_nested_nullability(f.data_type()), + true, + )) + } + match dt { + DataType::List(f) => DataType::List(erase(f)), + DataType::LargeList(f) => DataType::LargeList(erase(f)), + DataType::ListView(f) => DataType::ListView(erase(f)), + DataType::LargeListView(f) => DataType::LargeListView(erase(f)), + DataType::FixedSizeList(f, n) => DataType::FixedSizeList(erase(f), *n), + DataType::Struct(fields) => DataType::Struct(fields.iter().map(erase).collect()), + DataType::Map(f, sorted) => DataType::Map(erase(f), *sorted), + other => other.clone(), + } +} + +/// True if two types agree once nested nullability is disregarded. +/// +/// This is the check applied between the return type registered on the JVM side and the type the +/// UDF's own `return_field` reports. It stays strict about everything that changes how bytes are +/// read (decimal precision and scale, timestamp unit, child ordering and names) while tolerating +/// the nested-nullability drift described on [`erase_nested_nullability`]. +pub fn return_types_compatible(declared: &DataType, actual: &DataType) -> bool { + declared == actual || erase_nested_nullability(declared) == erase_nested_nullability(actual) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::{Fields, TimeUnit}; + + fn list(inner: DataType, nullable: bool) -> DataType { + DataType::List(Arc::new(Field::new("item", inner, nullable))) + } + + #[test] + fn identical_types_are_compatible() { + assert!(return_types_compatible(&DataType::Int32, &DataType::Int32)); + } + + #[test] + fn nested_nullability_is_disregarded() { + assert!(return_types_compatible( + &list(DataType::Int32, false), + &list(DataType::Int32, true) + )); + } + + #[test] + fn nested_nullability_is_disregarded_through_several_levels() { + let non_null = list( + DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, false)])), + false, + ); + let nullable = list( + DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])), + true, + ); + assert!(return_types_compatible(&non_null, &nullable)); + } + + #[test] + fn differing_value_types_are_incompatible() { + assert!(!return_types_compatible( + &list(DataType::Int32, true), + &list(DataType::Int64, true) + )); + } + + /// The case that motivated the check: Spark widens decimal arithmetic, so a UDF registered + /// with the un-widened precision must still be rejected. + #[test] + fn differing_decimal_precision_is_incompatible() { + assert!(!return_types_compatible( + &DataType::Decimal128(10, 2), + &DataType::Decimal128(11, 2) + )); + } + + #[test] + fn differing_timestamp_unit_is_incompatible() { + assert!(!return_types_compatible( + &DataType::Timestamp(TimeUnit::Microsecond, None), + &DataType::Timestamp(TimeUnit::Millisecond, None) + )); + } + + #[test] + fn differing_struct_field_names_are_incompatible() { + let a = DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])); + let b = DataType::Struct(Fields::from(vec![Field::new("b", DataType::Int32, true)])); + assert!(!return_types_compatible(&a, &b)); + } +} diff --git a/native/core/src/execution/rust_udf/test_support.rs b/native/core/src/execution/rust_udf/test_support.rs new file mode 100644 index 0000000000..d5637800c2 --- /dev/null +++ b/native/core/src/execution/rust_udf/test_support.rs @@ -0,0 +1,29 @@ +// 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. + +//! Test helpers shared across the rust_udf submodules. + +/// `comet-test-udfs` is `crate-type = ["cdylib"]` and has no test targets, so a +/// test build compiles it without emitting the shared library these tests +/// dlopen. Name the fix in the failure rather than leaving a bare dlopen error. +pub(crate) const BUILD_HINT: &str = "run `cargo build -p comet-test-udfs` first"; + +/// Path to the `comet-test-udfs` cdylib, baked in at build time by +/// `core/build.rs`. +pub(crate) fn test_udfs_path() -> std::path::PathBuf { + std::path::PathBuf::from(env!("COMET_TEST_UDFS_LIB")) +} diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index 6cfe33223f..2ea399bae0 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -66,6 +66,7 @@ pub mod jvm_bridge { use errors::{try_unwrap_or_throw, CometError, CometResult}; pub mod cloud; +pub mod comet_rust_udf_bridge; pub mod execution; pub mod parquet; // this module is for non release only. Intended for debugging/profiling purposes diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index 04b20b5df3..c01a6836af 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -94,6 +94,7 @@ message Expr { Shuffle shuffle = 72; RandStr rand_str = 73; Uuid uuid = 74; + RustUdfCall rust_udf_call = 75; } reserved 20; @@ -618,3 +619,27 @@ message JvmScalarUdf { // Whether the result column may contain nulls. bool return_nullable = 4; } + +// Call to a user-supplied Rust UDF loaded from a cdylib. +// +// The native side resolves (library_path, name) against its loaded-library +// cache, looks up the kernel by name, and invokes it through whichever ABI +// flavor (C ABI / datafusion-ffi) the cdylib registered the kernel under. +message RustUdfCall { + // Function name as registered through CometRustUDF.register on the JVM + // side; matched against names exposed by the cdylib. + string name = 1; + // Filesystem path of the cdylib. + string library_path = 2; + // Argument expressions, evaluated before invocation. + repeated Expr args = 3; + // Expected return type, declared at register time on the JVM side. + DataType return_type = 4; + // Whether the call is deterministic (mirrors Spark's deterministic flag). + // + // Always true today: registration rejects deterministic = false, because the + // native side plans every Rust UDF with Volatility::Immutable and has no way + // to express a volatile one. The field is carried so that honoring it later + // does not need a wire change. + bool deterministic = 5; +} diff --git a/pom.xml b/pom.xml index a50d94543a..b17865bb59 100644 --- a/pom.xml +++ b/pom.xml @@ -942,6 +942,9 @@ under the License. file:src/test/resources/log4j2.properties true ${project.build.directory}/tmp + + ${comet.test.udfs.lib} diff --git a/spark/src/main/java/org/apache/comet/udf/CometRustUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometRustUdfBridge.java new file mode 100644 index 0000000000..3d40f7b13a --- /dev/null +++ b/spark/src/main/java/org/apache/comet/udf/CometRustUdfBridge.java @@ -0,0 +1,41 @@ +/* + * 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.comet.udf; + +import org.apache.comet.NativeBase; + +/** JNI bridge for driver-side Rust UDF library validation. */ +public final class CometRustUdfBridge extends NativeBase { + private CometRustUdfBridge() {} + + /** + * Validate that {@code libraryPath} loads, exposes a UDF named {@code expectedName}, and return a + * JSON description of that UDF. Throws RuntimeException on any error. + * + *

The returned JSON has the form: {@code {"name":"add_one_c","abi":"c-abi"}} + */ + public static native String validateLibrary(String libraryPath, String expectedName); + + /** + * Return a JSON array describing every UDF exposed by {@code libraryPath}. Each element has the + * same shape as the return value of {@link #validateLibrary}. + */ + public static native String listUdfs(String libraryPath); +} diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index 26fe0c591a..349b9374f9 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -29,6 +29,7 @@ import org.apache.comet.CometSparkSessionExtensions.{withCodegenDispatchExpr, wi import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType} +import org.apache.comet.udf.CometRustUdfRegistry import org.apache.comet.udf.codegen.CometScalaUDFCodegen /** @@ -53,8 +54,45 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen */ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { - override def convert(expr: ScalaUDF, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = - emitJvmCodegenDispatch(expr, inputs, binding) + override def convert(expr: ScalaUDF, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { + // First check if this udfName is a registered Rust UDF -- those get emitted as RustUdfCall + // and dispatched to the loaded cdylib rather than the JVM codegen dispatcher. + expr.udfName.flatMap(CometRustUdfRegistry.instance.get) match { + case Some(meta) => + emitRustUdfCall(expr, meta.libraryPath, meta.returnType, inputs, binding) + case None => + emitJvmCodegenDispatch(expr, inputs, binding) + } + } + + private def emitRustUdfCall( + expr: ScalaUDF, + libraryPath: String, + returnType: org.apache.spark.sql.types.DataType, + inputs: Seq[Attribute], + binding: Boolean): Option[Expr] = { + val name = expr.udfName.get + val argProtos = expr.children.map(c => exprToProtoInternal(c, inputs, binding)) + if (argProtos.exists(_.isEmpty)) { + withFallbackReason( + expr, + "one or more Rust UDF arguments are not supported", + expr.children: _*) + return None + } + val returnTypeProto = serializeDataType(returnType).getOrElse { + withFallbackReason(expr, s"return type $returnType not serializable", expr) + return None + } + val callBuilder = ExprOuterClass.RustUdfCall + .newBuilder() + .setName(name) + .setLibraryPath(libraryPath) + .setReturnType(returnTypeProto) + .setDeterministic(expr.deterministic) + argProtos.foreach(a => callBuilder.addArgs(a.get)) + Some(ExprOuterClass.Expr.newBuilder().setRustUdfCall(callBuilder.build()).build()) + } /** * Bind `expr`, closure-serialize it, and emit a `JvmScalarUdf` proto routed through diff --git a/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala new file mode 100644 index 0000000000..19289df6ae --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala @@ -0,0 +1,150 @@ +/* + * 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.comet.udf + +import scala.util.Try + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.expressions.UserDefinedFunction +import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.types.DataType + +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.node.ObjectNode + +/** + * Entry point for registering Rust scalar UDFs with Comet. + * + * The UDF cdylib is built against the `comet-udf-sdk` crate and exposes its functions through an + * ABI built only on the Arrow C Data Interface, so a compiled UDF is not tied to Comet's + * DataFusion version. + * + * This is an experimental API. It is deliberately not annotated + * `org.apache.comet.annotation.Public`, so it sits outside the enumerated public API in Comet's + * [[https://datafusion.apache.org/comet/about/versioning_policy.html versioning policy]] and + * carries no compatibility guarantee: it may change or be removed in any release, including a + * patch release, with no deprecation cycle. + */ +object CometRustUDF { + + private val mapper: ObjectMapper = new ObjectMapper() + + /** + * Register a single Rust UDF with an explicit signature. + * + * Validates the library on the driver (loads it, confirms a UDF named `name` exists). On + * success a stub Spark catalog UDF is installed (so SQL/DataFrame name resolution succeeds) and + * the driver-side registry is updated. + * + * Executors do not consult the driver's registry: the library path travels with the plan in the + * `RustUdfCall` proto, and each executor loads the library itself on first use. The path must + * therefore be valid on every executor, not just the driver. + * + * `deterministic` must be `true`. Comet plans every imported kernel as immutable, so a + * nondeterministic UDF cannot yet be expressed; passing `false` fails here rather than silently + * planning the function as pure. + */ + def register( + spark: SparkSession, + name: String, + libraryPath: String, + inputTypes: Seq[DataType], + returnType: DataType, + deterministic: Boolean = true): Unit = { + if (!deterministic) { + // The native signature is built once per library load with + // Volatility::Immutable, while determinism is declared per registration, so the + // flag cannot be honored without reworking how kernels are cached. Until then a + // `false` here would let DataFusion constant-fold or CSE a call the user told us + // was not safe to reuse. + throw new IllegalArgumentException( + s"Rust UDF '$name': deterministic = false is not supported yet. Comet plans Rust UDFs " + + "as immutable, so a nondeterministic function may be constant-folded or eliminated " + + "as a common subexpression. See https://github.com/apache/datafusion-comet/issues/5249") + } + val described = describeOne(libraryPath, name) + require(described.name == name, s"unexpected name from native: ${described.name}") + installCatalogStub(spark, name, inputTypes, returnType, deterministic) + val meta = RustUdfMetadata(libraryPath, inputTypes, returnType, deterministic) + CometRustUdfRegistry.instance.register(name, meta) + } + + // -------- internals -------- + + private case class Described(name: String) + + private def describeOne(libraryPath: String, name: String): Described = { + val json = + invokeBridge(() => CometRustUdfBridge.validateLibrary(libraryPath, name), libraryPath) + parseDescribed(json) + } + + private def invokeBridge(call: () => String, libraryPath: String): String = { + Try(call()).recover { case t: Throwable => throw classifyNativeError(libraryPath, t) }.get + } + + private def parseDescribed(json: String): Described = { + val node = mapper.readTree(json).asInstanceOf[ObjectNode] + Described(name = node.get("name").asText()) + } + + private def classifyNativeError(libraryPath: String, t: Throwable): RuntimeException = { + val m = Option(t.getMessage).getOrElse("") + if (m.contains("ABI") || m.contains("missing required symbol") || + m.contains("comet_udf_abi_version") || m.contains("does not export")) { + new CometRustUdfAbiException(m) + } else if (m.contains("not found in")) { + new java.util.NoSuchElementException(m) + } else { + new CometRustUdfLoadException(s"failed to load $libraryPath: $m", t) + } + } + + private def installCatalogStub( + spark: SparkSession, + name: String, + inputTypes: Seq[DataType], + returnType: DataType, + deterministic: Boolean): Unit = { + val arity = inputTypes.size + val u: UserDefinedFunction = arity match { + case 0 => + udf(() => throw new CometRustUdfNotEvaluatedException(name), returnType) + case 1 => + udf((_: Any) => throw new CometRustUdfNotEvaluatedException(name), returnType) + case 2 => + udf((_: Any, _: Any) => throw new CometRustUdfNotEvaluatedException(name), returnType) + case 3 => + udf( + (_: Any, _: Any, _: Any) => throw new CometRustUdfNotEvaluatedException(name), + returnType) + case 4 => + udf( + (_: Any, _: Any, _: Any, _: Any) => throw new CometRustUdfNotEvaluatedException(name), + returnType) + case n => + throw new IllegalArgumentException( + s"Rust UDF '$name' arity $n not supported by stub. Reduce arity " + + "or open a feature request to extend stub coverage.") + } + val finalUdf = if (deterministic) u else u.asNondeterministic() + spark.udf.register(name, finalUdf) + } +} diff --git a/spark/src/main/scala/org/apache/comet/udf/CometRustUdfExceptions.scala b/spark/src/main/scala/org/apache/comet/udf/CometRustUdfExceptions.scala new file mode 100644 index 0000000000..2d67d0ad81 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/udf/CometRustUdfExceptions.scala @@ -0,0 +1,51 @@ +/* + * 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.comet.udf + +import org.apache.comet.{CometNativeException, CometRuntimeException} + +/** Thrown when a Rust UDF dynamic library cannot be opened. */ +class CometRustUdfLoadException(msg: String, cause: Throwable = null) + extends CometNativeException(msg) { + if (cause != null) initCause(cause) +} + +/** + * Thrown when a Rust UDF library exposes the wrong ABI version or is missing required discovery + * symbols. + */ +class CometRustUdfAbiException(msg: String) extends CometNativeException(msg) + +/** + * Thrown when the declared signature does not match what the library reports via + * `comet_*_udf_list_v1`. + */ +class CometRustUdfSignatureException(msg: String) extends CometRuntimeException(msg) + +/** + * Thrown by the catalog stub if a registered Rust UDF is invoked on the JVM (which means Comet's + * plan rule did not replace it). + */ +class CometRustUdfNotEvaluatedException(name: String) + extends CometRuntimeException( + s"Rust UDF '$name' must run inside Comet native execution; the JVM " + + "stub was invoked, which means Comet did not replace this expression " + + "with a native call. Check that Comet is enabled for the operator " + + "hosting this expression.") diff --git a/spark/src/main/scala/org/apache/comet/udf/CometRustUdfRegistry.scala b/spark/src/main/scala/org/apache/comet/udf/CometRustUdfRegistry.scala new file mode 100644 index 0000000000..2e277b6c01 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/udf/CometRustUdfRegistry.scala @@ -0,0 +1,53 @@ +/* + * 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.comet.udf + +import java.util.concurrent.ConcurrentHashMap + +import org.apache.spark.sql.types.DataType + +/** Metadata for a registered Rust UDF. */ +case class RustUdfMetadata( + libraryPath: String, + inputTypes: Seq[DataType], + returnType: DataType, + deterministic: Boolean) + +/** + * Driver-side registry of Rust UDFs. Looked up by `QueryPlanSerde` to recognize names that should + * be emitted as `RustUdfCall` instead of attempted as JVM-evaluated `ScalaUDF`s. + */ +class CometRustUdfRegistry { + private val byName = new ConcurrentHashMap[String, RustUdfMetadata]() + + /** Register or replace metadata for a name. */ + def register(name: String, meta: RustUdfMetadata): Unit = + byName.put(name, meta) + + /** Return metadata for a name, if registered. */ + def get(name: String): Option[RustUdfMetadata] = + Option(byName.get(name)) +} + +object CometRustUdfRegistry { + + /** Process-wide singleton. */ + lazy val instance: CometRustUdfRegistry = new CometRustUdfRegistry +} diff --git a/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala new file mode 100644 index 0000000000..90db98a16a --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala @@ -0,0 +1,312 @@ +/* + * 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.comet + +import java.io.File +import java.util.Locale + +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.types._ + +import org.apache.comet.udf.CometRustUDF + +/** + * End-to-end integration suite: register a Rust UDF, run a Spark query, verify the result. + * + * Requires the `comet-test-udfs` cdylib, which is found automatically under `native/target` and + * can be overridden with `-Dcomet.test.udfs.lib=`. + * + * Note that these tests are self-guarding on native execution: `CometRustUDF.register` installs a + * catalog stub that throws if Spark ever evaluates the UDF itself, so a silent fallback to Spark + * fails the test rather than passing. + * + * To run locally: + * {{{ + * cargo build -p comet-test-udfs --manifest-path native/Cargo.toml + * ./mvnw test -Dsuites="org.apache.comet.CometRustUdfSuite" -Dtest=none + * }}} + */ +class CometRustUdfSuite extends CometTestBase { + + private lazy val libPath: String = { + val overridden = Option(System.getProperty("comet.test.udfs.lib")) + // An undefined Maven property reaches the forked JVM as the literal "null", and an + // unsubstituted one as "${comet.test.udfs.lib}". Neither is a path. + .map(_.trim) + .filter(p => p.nonEmpty && p != "null" && !p.startsWith("$")) + + overridden.orElse(CometRustUdfSuite.discoverBuiltLibrary()).getOrElse { + if (sys.env.contains("CI")) { + // In CI the cdylib is staged alongside libcomet, so its absence means the native build or + // the artifact upload changed, not that someone forgot a flag. Fail rather than skip. + fail( + s"${CometRustUdfSuite.libraryFileName} was not found under native/target. CI stages it " + + "next to libcomet, so a missing library means the native build or the artifact " + + "upload has changed.") + } else { + cancel(s"${CometRustUdfSuite.libraryFileName} not built; run " + + "`cargo build -p comet-test-udfs --manifest-path native/Cargo.toml` to run this suite") + } + } + } + + test("add_one_c returns id + 1 for a range") { + CometRustUDF.register(spark, "add_one_c", libPath, Seq(LongType), LongType) + val df = spark.range(0, 5).selectExpr("add_one_c(id) AS y") + val out = df.collect().map(_.getLong(0)).sorted.toSeq + assert(out == Seq(1L, 2L, 3L, 4L, 5L)) + } + + test("panic inside UDF invoke surfaces as a query error, not a crash") { + CometRustUDF.register(spark, "panics_on_invoke", libPath, Seq(LongType), LongType) + val e = intercept[Exception] { + spark.range(0, 5).selectExpr("panics_on_invoke(id) AS y").collect() + } + assert( + stackTraceContains(e, "deliberate panic from user UDF code"), + s"panic message not propagated: $e") + } + + test("panic inside UDF return_field surfaces as a query error, not a crash") { + CometRustUDF.register(spark, "panics_on_return_field", libPath, Seq(LongType), LongType) + val e = intercept[Exception] { + spark.range(0, 5).selectExpr("panics_on_return_field(id) AS y").collect() + } + assert( + stackTraceContains(e, "deliberate panic from user return_field"), + s"panic message not propagated: $e") + } + + /** True if `needle` appears anywhere in the exception's cause chain. */ + private def stackTraceContains(e: Throwable, needle: String): Boolean = { + Iterator + .iterate(e)(_.getCause) + .takeWhile(_ != null) + .exists(t => Option(t.getMessage).exists(_.contains(needle))) + } + + // ---------- type coverage ---------- + + /** + * One case per supported Spark type: the type itself, and a SQL expression producing a value of + * that type from the `id` column of `spark.range`. + * + * The declared type is asserted against the frame's real schema before it is registered, so a + * case whose expression does not produce the type it claims fails loudly rather than testing + * the wrong thing. + */ + private val typeCases: Seq[(DataType, String)] = Seq( + (BooleanType, "id % 2 = 0"), + (ByteType, "cast(id as byte)"), + (ShortType, "cast(id as short)"), + (IntegerType, "cast(id as int)"), + (LongType, "id"), + (FloatType, "cast(id as float) + 0.5f"), + (DoubleType, "cast(id as double) + 0.5"), + (StringType, "concat('s', cast(id as string))"), + (BinaryType, "cast(concat('b', cast(id as string)) as binary)"), + (DateType, "date_add(date'2024-01-01', cast(id as int))"), + (TimestampType, "cast(date_add(date'2024-01-01', cast(id as int)) as timestamp)"), + (TimestampNTZType, "cast(timestamp_ntz'2024-01-01 12:00:00' as timestamp_ntz)"), + // The outer cast pins the result to decimal(10,2): Spark widens the precision of the addition + // itself to decimal(11,2), which would not match the type registered below. + (DecimalType(10, 2), "cast(cast(id as decimal(10,2)) + 0.25 as decimal(10,2))"), + // Complex types. `containsNull` / `valueContainsNull` / field nullability are part of the type + // and must match what the expression actually produces, hence the explicit constructors. + (ArrayType(IntegerType, containsNull = false), "array(cast(id as int), cast(id + 1 as int))"), + (ArrayType(StringType, containsNull = false), "array(concat('s', cast(id as string)))"), + ( + MapType(StringType, IntegerType, valueContainsNull = false), + "map('k', cast(id as int), 'j', cast(id + 1 as int))"), + ( + StructType( + Seq( + StructField("a", IntegerType, nullable = false), + StructField("b", StringType, nullable = false))), + "named_struct('a', cast(id as int), 'b', concat('s', cast(id as string)))"), + // One level of nesting, to check the FFI carries child arrays rather than just top-level ones. + ( + ArrayType( + StructType(Seq(StructField("a", IntegerType, nullable = false))), + containsNull = false), + "array(named_struct('a', cast(id as int)))"), + ( + StructType( + Seq(StructField("xs", ArrayType(IntegerType, containsNull = false), nullable = false))), + "named_struct('xs', array(cast(id as int), cast(id + 1 as int)))")) + + /** + * A 4-row frame with a single column `c` of the given type, where the last row is null so every + * case also covers null handling across the FFI boundary. + */ + private def typedFrame(valueExpr: String) = + spark.range(0, 4).selectExpr(s"case when id = 3 then null else $valueExpr end as c") + + /** Normalize for comparison: byte arrays do not compare by value as `Any`. */ + private def normalize(v: Any): Any = v match { + case b: Array[Byte] => b.toSeq + case other => other + } + + for ((dataType, valueExpr) <- typeCases) { + test(s"echo_c round-trips ${dataType.simpleString} including nulls") { + val df = typedFrame(valueExpr) + assert( + df.schema.head.dataType == dataType, + s"test expression produced ${df.schema.head.dataType}, not $dataType") + CometRustUDF.register(spark, "echo_c", libPath, Seq(dataType), dataType) + val expected = df.collect().map(r => normalize(r.get(0))).toSeq + val actual = df.selectExpr("echo_c(c) AS y").collect().map(r => normalize(r.get(0))).toSeq + assert(actual == expected, s"round trip changed values for ${dataType.simpleString}") + assert(expected.last == null, "expected a null in the last row") + } + + test(s"stringify_c reads ${dataType.simpleString} values") { + CometRustUDF.register(spark, "stringify_c", libPath, Seq(dataType), StringType) + val df = typedFrame(valueExpr) + val inputs = df.collect().map(r => normalize(r.get(0))).toSeq + val rendered = df.selectExpr("stringify_c(c) AS y").collect().map(r => r.get(0)).toSeq + + assert(rendered.length == inputs.length) + // The UDF must decode each value, so a non-null input yields a non-empty rendering and a + // null input stays null. The exact text is arrow's formatting, not Spark's, so it is not + // asserted here. + inputs.zip(rendered).foreach { case (in, out) => + if (in == null) { + assert(out == null, s"null input rendered as $out for ${dataType.simpleString}") + } else { + assert(out != null, s"non-null input $in rendered as null") + assert( + out.asInstanceOf[String].nonEmpty, + s"non-null input $in rendered empty for ${dataType.simpleString}") + } + } + } + } + + test("one kernel computes its return type on demand and serves many types") { + // echo_c has no fixed return type of its own: its return_field derives one from the argument + // types on every call. The type declared to `register` is what Spark plans against, so it is + // per-registration rather than per-kernel, and the same kernel serves a different type after + // re-registering. + CometRustUDF.register(spark, "echo_c", libPath, Seq(LongType), LongType) + assert( + spark.range(0, 3).selectExpr("echo_c(id) AS y").collect().map(_.getLong(0)).toSeq == + Seq(0L, 1L, 2L)) + + CometRustUDF.register(spark, "echo_c", libPath, Seq(StringType), StringType) + val strings = spark + .range(0, 3) + .selectExpr("echo_c(concat('s', cast(id as string))) AS y") + .collect() + .map(_.getString(0)) + .toSeq + assert(strings == Seq("s0", "s1", "s2")) + + val arrayType = ArrayType(IntegerType, containsNull = false) + CometRustUDF.register(spark, "echo_c", libPath, Seq(arrayType), arrayType) + val arrays = spark + .range(0, 2) + .selectExpr("echo_c(array(cast(id as int), cast(id + 1 as int))) AS y") + .collect() + .map(_.getSeq[Int](0)) + .toSeq + assert(arrays == Seq(Seq(0, 1), Seq(1, 2))) + } + + test("a declared return type that disagrees with the UDF names both types") { + // echo_c returns its argument's type, so declaring a different return type is a mismatch. + CometRustUDF.register(spark, "echo_c", libPath, Seq(LongType), StringType) + val e = intercept[Exception] { + spark.range(0, 4).selectExpr("echo_c(id) AS y").collect() + } + assert(stackTraceContains(e, "was registered as returning"), s"unhelpful error: $e") + assert(stackTraceContains(e, "CometRustUDF.register"), s"error lacks guidance: $e") + } + + test("echo_c rejects a call whose argument count it does not accept") { + CometRustUDF.register(spark, "echo_c", libPath, Seq(LongType), LongType) + // The catalog stub is arity-1, so a 2-arg call is rejected during analysis. + intercept[Exception] { + spark.range(0, 2).selectExpr("echo_c(id, id) AS y").collect() + } + } + + test("registering a nondeterministic UDF is refused") { + // Comet plans every Rust UDF as immutable, so accepting this would let the optimizer + // constant-fold or CSE a call the caller told us was not safe to reuse. Refuse at + // registration rather than silently ignore the flag. + val e = intercept[IllegalArgumentException] { + CometRustUDF.register( + spark, + "echo_c", + libPath, + Seq(LongType), + LongType, + deterministic = false) + } + assert(e.getMessage.contains("deterministic = false is not supported"), s"unclear: $e") + + // The check runs before any library work, so it fires on a path that does not exist + // rather than reporting a load failure first. + val early = intercept[IllegalArgumentException] { + CometRustUDF.register( + spark, + "echo_c", + "/no/such/library.so", + Seq(LongType), + LongType, + deterministic = false) + } + assert( + early.getMessage.contains("deterministic = false is not supported"), + s"unclear: $early") + } +} + +object CometRustUdfSuite { + + /** Platform file name of the test cdylib built by the `comet-test-udfs` crate. */ + val libraryFileName: String = + if (System.getProperty("os.name", "").toLowerCase(Locale.ROOT).contains("mac")) { + "libcomet_test_udfs.dylib" + } else { + "libcomet_test_udfs.so" + } + + /** + * Locate the test cdylib under `native/target`. + * + * The working directory differs between a reactor build and a single-module run, so walk up a + * few levels looking for the `native/target` tree. `release` is where CI stages the downloaded + * artifact, `ci` and `debug` cover local builds. + */ + def discoverBuiltLibrary(): Option[String] = { + val roots = Iterator + .iterate(new File(".").getCanonicalFile)(_.getParentFile) + .takeWhile(_ != null) + .take(4) + val candidates = for { + root <- roots + profile <- Seq("release", "ci", "debug") + } yield new File(root, s"native/target/$profile/$libraryFileName") + candidates.find(_.isFile).map(_.getAbsolutePath) + } +}