From ae3983b0dcff8121f3ea2723f3b0af4eeac99513 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 27 May 2026 15:29:35 -0600 Subject: [PATCH 01/11] feat: custom Rust UDFs via arrow-ffi (alternative to #4283) [skip ci] Adds custom Rust scalar UDF support using arrow-only FFI surfaces, as suggested by paleolimbot and timsaucer in feedback on #4283. This is an alternative implementation for comparison; see #4283 for the bespoke-ABI version. Two ABI flavors are provided side-by-side so reviewers can compare: 1. C ABI (sedona-style): pure C-callable struct of function pointers, parameterized only by Arrow C Data Interface (FFI_ArrowSchema / FFI_ArrowArray). Decoupled from datafusion versions; future-portable to C/C++. Modeled on apache/sedona-db's SedonaCScalarKernel header. 2. datafusion-ffi (FFI_ScalarUDF): wraps user's ScalarUDFImpl as FFI_ScalarUDF. Inherits full ScalarUDFImpl surface (variadic signatures, type coercion, metadata-aware return types) for free, at the cost of a major-version pin against datafusion-ffi. A single library may export either or both; loader walks both discovery functions. `comet-test-udfs` exposes `add_one_c` (C ABI) and `add_one_df` (datafusion-ffi) and the e2e suite drives both through Spark. Scope is intentionally minimal for comparison: scalar-only, happy-path e2e tests only (no panic/error/signature-mismatch coverage). The Scala JVM API mirrors #4283 exactly so the comparison is apples-to-apples. Pieces: - native/comet-udf-sdk: SDK with both ABI flavors + export macros - native/comet-test-udfs: cdylib exposing one UDF per ABI - native/core/src/execution/rust_udf: loader, cache, ImportedCScalarUdf - native/core/src/comet_rust_udf_bridge.rs: JNI for validateLibrary/listUdfs - native/proto: RustUdfCall message - spark/.../udf: CometRustUDF.register, registry, exception classes, JNI bridge stub - spark/.../serde/CometScalaUDF: dispatch ScalaUDF to RustUdfCall when udfName is in the registry --- native/Cargo.lock | 243 +++++- native/Cargo.toml | 5 +- native/comet-test-udfs/Cargo.toml | 36 + native/comet-test-udfs/src/lib.rs | 142 ++++ native/comet-udf-sdk/Cargo.toml | 41 + native/comet-udf-sdk/src/c_abi.rs | 707 ++++++++++++++++++ native/comet-udf-sdk/src/df_abi.rs | 239 ++++++ native/comet-udf-sdk/src/lib.rs | 71 ++ native/core/Cargo.toml | 4 + native/core/build.rs | 43 ++ native/core/src/comet_rust_udf_bridge.rs | 97 +++ native/core/src/execution/mod.rs | 1 + native/core/src/execution/planner.rs | 41 + native/core/src/execution/rust_udf/cache.rs | 87 +++ .../core/src/execution/rust_udf/imported_c.rs | 312 ++++++++ native/core/src/execution/rust_udf/loader.rs | 350 +++++++++ native/core/src/execution/rust_udf/mod.rs | 43 ++ .../src/execution/rust_udf/test_support.rs | 24 + native/core/src/lib.rs | 1 + native/proto/src/proto/expr.proto | 20 + pom.xml | 2 + .../apache/comet/udf/CometRustUdfBridge.java | 41 + .../apache/comet/serde/CometScalaUDF.scala | 39 +- .../org/apache/comet/udf/CometRustUDF.scala | 140 ++++ .../comet/udf/CometRustUdfExceptions.scala | 51 ++ .../comet/udf/CometRustUdfRegistry.scala | 59 ++ .../org/apache/comet/CometRustUdfSuite.scala | 66 ++ 27 files changed, 2899 insertions(+), 6 deletions(-) create mode 100644 native/comet-test-udfs/Cargo.toml create mode 100644 native/comet-test-udfs/src/lib.rs create mode 100644 native/comet-udf-sdk/Cargo.toml create mode 100644 native/comet-udf-sdk/src/c_abi.rs create mode 100644 native/comet-udf-sdk/src/df_abi.rs create mode 100644 native/comet-udf-sdk/src/lib.rs create mode 100644 native/core/build.rs create mode 100644 native/core/src/comet_rust_udf_bridge.rs create mode 100644 native/core/src/execution/rust_udf/cache.rs create mode 100644 native/core/src/execution/rust_udf/imported_c.rs create mode 100644 native/core/src/execution/rust_udf/loader.rs create mode 100644 native/core/src/execution/rust_udf/mod.rs create mode 100644 native/core/src/execution/rust_udf/test_support.rs create mode 100644 spark/src/main/java/org/apache/comet/udf/CometRustUdfBridge.java create mode 100644 spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala create mode 100644 spark/src/main/scala/org/apache/comet/udf/CometRustUdfExceptions.scala create mode 100644 spark/src/main/scala/org/apache/comet/udf/CometRustUdfRegistry.scala create mode 100644 spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala diff --git a/native/Cargo.lock b/native/Cargo.lock index 11e9b1ccff..80ff09ddf7 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2,6 +2,54 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "abi_stable" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69d6512d3eb05ffe5004c59c206de7f99c34951504056ce23fc953842f12c445" +dependencies = [ + "abi_stable_derive", + "abi_stable_shared", + "const_panic", + "core_extensions", + "crossbeam-channel", + "generational-arena", + "libloading 0.7.4", + "lock_api", + "parking_lot", + "paste", + "repr_offset", + "rustc_version", + "serde", + "serde_derive", + "serde_json", +] + +[[package]] +name = "abi_stable_derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7178468b407a4ee10e881bc7a328a65e739f0863615cca4429d43916b05e898" +dependencies = [ + "abi_stable_shared", + "as_derive_utils", + "core_extensions", + "proc-macro2", + "quote", + "rustc_version", + "syn 1.0.109", + "typed-arena", +] + +[[package]] +name = "abi_stable_shared" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2b5df7688c123e63f4d4d649cba63f2967ba7f7861b1664fca3f77d3dad2b63" +dependencies = [ + "core_extensions", +] + [[package]] name = "addr2line" version = "0.25.1" @@ -458,6 +506,18 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b0f477b951e452a0b6b4a10b53ccd569042d1d01729b519e02074a9c0958a063" +[[package]] +name = "as_derive_utils" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff3c96645900a44cf11941c111bd08a6573b0e2f9f69bc9264b179d8fae753c4" +dependencies = [ + "core_extensions", + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "assertables" version = "9.9.0" @@ -517,6 +577,15 @@ dependencies = [ "slab", ] +[[package]] +name = "async-ffi" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4de21c0feef7e5a556e51af767c953f0501f7f300ba785cc99c47bdc8081a50" +dependencies = [ + "abi_stable", +] + [[package]] name = "async-global-executor" version = "2.4.1" @@ -1427,7 +1496,7 @@ checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" dependencies = [ "glob", "libc", - "libloading", + "libloading 0.8.9", ] [[package]] @@ -1501,6 +1570,24 @@ dependencies = [ "memchr", ] +[[package]] +name = "comet-test-udfs" +version = "0.17.0" +dependencies = [ + "arrow", + "comet-udf-sdk", + "datafusion", +] + +[[package]] +name = "comet-udf-sdk" +version = "0.17.0" +dependencies = [ + "arrow", + "datafusion", + "datafusion-ffi", +] + [[package]] name = "comfy-table" version = "7.2.2" @@ -1573,6 +1660,15 @@ dependencies = [ "tiny-keccak", ] +[[package]] +name = "const_panic" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e262cdaac42494e3ae34c43969f9cdeb7da178bdb4b66fa6a1ea2edb4c8ae652" +dependencies = [ + "typewit", +] + [[package]] name = "constant_time_eq" version = "0.4.2" @@ -1595,6 +1691,21 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core_extensions" +version = "1.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42bb5e5d0269fd4f739ea6cedaf29c16d81c27a7ce7582008e90eb50dcd57003" +dependencies = [ + "core_extensions_proc_macros", +] + +[[package]] +name = "core_extensions_proc_macros" +version = "1.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "533d38ecd2709b7608fb8e18e4504deb99e9a72879e6aa66373a76d8dc4259ea" + [[package]] name = "cpp_demangle" version = "0.5.1" @@ -1982,6 +2093,8 @@ dependencies = [ "aws-config", "aws-credential-types", "bytes", + "comet-test-udfs", + "comet-udf-sdk", "criterion", "datafusion", "datafusion-comet-common", @@ -1991,6 +2104,7 @@ dependencies = [ "datafusion-comet-shuffle", "datafusion-comet-spark-expr", "datafusion-datasource", + "datafusion-ffi", "datafusion-functions-nested", "datafusion-physical-expr-adapter", "datafusion-spark", @@ -2003,6 +2117,7 @@ dependencies = [ "itertools 0.14.0", "jni 0.22.4", "lazy_static", + "libloading 0.8.9", "log", "log4rs", "mimalloc", @@ -2382,6 +2497,36 @@ dependencies = [ "paste", ] +[[package]] +name = "datafusion-ffi" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b95173344d04ba62755c949bf44f8d1a6e4414cf6392a635db96c07e711b9a3c" +dependencies = [ + "abi_stable", + "arrow", + "arrow-schema", + "async-ffi", + "async-trait", + "datafusion-catalog", + "datafusion-common", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions-aggregate-common", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-proto", + "datafusion-proto-common", + "datafusion-session", + "futures", + "log", + "prost", + "semver", + "tokio", +] + [[package]] name = "datafusion-functions" version = "53.1.0" @@ -2653,6 +2798,45 @@ dependencies = [ "tokio", ] +[[package]] +name = "datafusion-proto" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a387aaef949dc16bb6abc81bd1af850ec7449183aef011214f9724957495738" +dependencies = [ + "arrow", + "chrono", + "datafusion-catalog", + "datafusion-catalog-listing", + "datafusion-common", + "datafusion-datasource", + "datafusion-datasource-arrow", + "datafusion-datasource-csv", + "datafusion-datasource-json", + "datafusion-datasource-parquet", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions-table", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-proto-common", + "object_store", + "prost", + "rand 0.9.4", +] + +[[package]] +name = "datafusion-proto-common" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16e614c7c53a9c304c6a850b821010bb492e57300311835f1180613f9d2c63d9" +dependencies = [ + "arrow", + "datafusion-common", + "prost", +] + [[package]] name = "datafusion-pruning" version = "53.1.0" @@ -3196,6 +3380,15 @@ dependencies = [ "slab", ] +[[package]] +name = "generational-arena" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877e94aff08e743b651baaea359664321055749b398adff8740a7399af7796e7" +dependencies = [ + "cfg-if", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -3971,7 +4164,7 @@ dependencies = [ "java-locator", "jni-macros", "jni-sys 0.4.1", - "libloading", + "libloading 0.8.9", "log", "simd_cesu8", "thiserror 2.0.18", @@ -4156,6 +4349,16 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libloading" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f" +dependencies = [ + "cfg-if", + "winapi", +] + [[package]] name = "libloading" version = "0.8.9" @@ -5555,6 +5758,15 @@ version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +[[package]] +name = "repr_offset" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb1070755bd29dffc19d0971cab794e607839ba2ef4b69a9e6fbc8733c1b72ea" +dependencies = [ + "tstr", +] + [[package]] name = "reqsign" version = "0.16.5" @@ -6752,6 +6964,21 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "tstr" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f8e0294f14baae476d0dd0a2d780b2e24d66e349a9de876f5126777a37bdba7" +dependencies = [ + "tstr_proc_macros", +] + +[[package]] +name = "tstr_proc_macros" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e78122066b0cb818b8afd08f7ed22f7fdbc3e90815035726f0840d0d26c0747a" + [[package]] name = "twox-hash" version = "2.1.2" @@ -6761,6 +6988,12 @@ dependencies = [ "rand 0.9.4", ] +[[package]] +name = "typed-arena" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" + [[package]] name = "typed-builder" version = "0.20.1" @@ -6826,6 +7059,12 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "typewit" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "214ca0b2191785cbc06209b9ca1861e048e39b5ba33574b3cedd58363d5bb5f6" + [[package]] name = "unicode-ident" version = "1.0.24" diff --git a/native/Cargo.toml b/native/Cargo.toml index 9e252a796c..35942caa80 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", "hdfs", "fs-hdfs"] +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", "hdfs", "fs-hdfs", "comet-udf-sdk", "comet-test-udfs"] resolver = "2" [workspace.package] @@ -40,6 +40,7 @@ bytes = { version = "1.11.1" } parquet = { version = "58.3.0", default-features = false, features = ["experimental"] } datafusion = { version = "53.1.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] } datafusion-datasource = { version = "53.1.0" } +datafusion-ffi = { version = "53.1.0" } datafusion-physical-expr-adapter = { version = "53.1.0" } datafusion-spark = { version = "53.1.0", features = ["core"] } datafusion-comet-spark-expr = { path = "spark-expr" } diff --git a/native/comet-test-udfs/Cargo.toml b/native/comet-test-udfs/Cargo.toml new file mode 100644 index 0000000000..90dc50b8a5 --- /dev/null +++ b/native/comet-test-udfs/Cargo.toml @@ -0,0 +1,36 @@ +# 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 } +datafusion = { 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..a5031e503a --- /dev/null +++ b/native/comet-test-udfs/src/lib.rs @@ -0,0 +1,142 @@ +// 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 the same `add_one` semantics through both ABI flavors so the +//! host can verify them side-by-side: +//! +//! - `add_one_c` — exposed via the C ABI (arrow-only FFI, sedona-style) +//! - `add_one_df` — exposed via the datafusion-ffi ABI (`FFI_ScalarUDF`) +//! +//! Built as `libcomet_test_udfs.{so,dylib}`. + +use std::any::Any; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int64Array}; +use arrow::datatypes::{DataType, Field}; +use datafusion::common::DataFusionError; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, + Volatility, +}; + +use comet_udf_sdk::c_abi::CometCScalarUdf; +use comet_udf_sdk::{comet_c_udf_export, comet_df_udf_export}; + +// ---------- C ABI implementation ------------------------------------ + +/// `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)) + } +} + +// ---------- datafusion-ffi implementation --------------------------- + +/// `add_one` exposed via datafusion-ffi. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct AddOneDf { + signature: Signature, +} + +impl AddOneDf { + fn new() -> Self { + Self { + signature: Signature::new( + TypeSignature::Exact(vec![DataType::Int64]), + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for AddOneDf { + fn as_any(&self) -> &dyn Any { + self + } + fn name(&self) -> &str { + "add_one_df" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _args: &[DataType]) -> datafusion::common::Result { + Ok(DataType::Int64) + } + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion::common::Result { + let arr: ArrayRef = match &args.args[0] { + ColumnarValue::Array(a) => Arc::clone(a), + ColumnarValue::Scalar(s) => s.to_array_of_size(args.number_rows)?, + }; + let a = arr + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Execution("expected Int64".into()))?; + let out: Int64Array = a.iter().map(|v| v.map(|x| x + 1)).collect(); + Ok(ColumnarValue::Array(Arc::new(out))) + } +} + +/// Factory used by the discovery macro. +pub fn make_add_one_df() -> Arc { + Arc::new(ScalarUDF::from(AddOneDf::new())) +} + +// ---------- discovery exports -------------------------------------- + +// `comet_c_udf_export!` emits both `comet_c_udf_list_v1` and +// `comet_udf_abi_version`. `comet_df_udf_export!` would emit +// `comet_udf_abi_version` again, so we use the `no_abi_version` form to +// avoid the duplicate symbol. +comet_c_udf_export!(AddOneC); +comet_df_udf_export!(no_abi_version => make_add_one_df); diff --git a/native/comet-udf-sdk/Cargo.toml b/native/comet-udf-sdk/Cargo.toml new file mode 100644 index 0000000000..766a095358 --- /dev/null +++ b/native/comet-udf-sdk/Cargo.toml @@ -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] +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 + +[features] +# Enable the C ABI flavor (Arrow C Data Interface only, sedona-style). +c-abi = [] +# Enable the datafusion-ffi flavor (FFI_ScalarUDF / ForeignScalarUDF). +df-abi = ["dep:datafusion", "dep:datafusion-ffi"] +default = ["c-abi", "df-abi"] + +[dependencies] +arrow = { workspace = true } +datafusion = { workspace = true, optional = true } +datafusion-ffi = { workspace = true, optional = 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..e8069eb369 --- /dev/null +++ b/native/comet-udf-sdk/src/c_abi.rs @@ -0,0 +1,707 @@ +// 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. + +//! Pure C ABI flavor — 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. +//! +//! # Authoring a UDF +//! +//! Implement [`CometCScalarUdf`] for a unit struct and 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; +//! +//! 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; + +// -- 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< + unsafe extern "C" fn(*const CometCScalarKernel, out: *mut CometCScalarKernelImpl), + >, + + /// 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`]. +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 }; + debug_assert!(!this.private_data.is_null()); + 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 }; + let exp = unsafe { &*(this.private_data as *const ExportedScalarKernel) }; + 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. + let _ = unsafe { + Box::from_raw(this_ref.private_data as *mut ExportedScalarKernel) + }; + this_ref.private_data = std::ptr::null_mut(); + } + 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 }; + 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; + } + } + } + + match 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 }; + 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 in 0..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_fields[i].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)); + } + + let result = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + priv_ref.inner.invoke(&arrays, n_rows as usize) + })) { + Ok(Ok(arr)) => arr, + Ok(Err(msg)) => { + priv_ref.last_error = std::ffi::CString::new(msg).unwrap_or_default(); + return C_ABI_ERR; + } + Err(panic) => { + let msg = match panic.downcast_ref::<&'static str>() { + Some(s) => format!("panic: {s}"), + None => match panic.downcast_ref::() { + Some(s) => format!("panic: {s}"), + None => "panic: ".to_string(), + }, + }; + 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() { + let _ = unsafe { + Box::from_raw(this_ref.private_data as *mut ExportedScalarKernelImpl) + }; + this_ref.private_data = std::ptr::null_mut(); + } + 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; + // 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. + let _ = unsafe { + Box::from_raw(std::slice::from_raw_parts_mut(list_ref.kernels, len)) + }; + list_ref.kernels = std::ptr::null_mut(); + list_ref.len = 0; + list_ref.release = None; +} + +// -- 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` +/// +/// 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 kernels: Vec<$crate::c_abi::CometCScalarKernel> = vec![ + $( + $crate::c_abi::CometCScalarKernel::from( + $crate::c_abi::ExportedScalarKernel::new(<$ty as Default>::default()) + ), + )+ + ]; + let list = $crate::c_abi::build_kernel_list(kernels); + std::ptr::write(out, list); + 0 + } + }; + }; +} + +#[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); + } +} diff --git a/native/comet-udf-sdk/src/df_abi.rs b/native/comet-udf-sdk/src/df_abi.rs new file mode 100644 index 0000000000..735ca84842 --- /dev/null +++ b/native/comet-udf-sdk/src/df_abi.rs @@ -0,0 +1,239 @@ +// 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. + +//! datafusion-ffi flavor. +//! +//! Discovery returns a list of `FFI_ScalarUDF` values produced by +//! `datafusion_ffi`. The host imports each via +//! `ForeignScalarUDF::try_from`, yielding a `ScalarUDFImpl` it can plug +//! straight into its existing planner — no further adaptation needed. +//! +//! The user's library inherits the entire `ScalarUDFImpl` surface +//! (signature with type coercion, metadata-aware return types, aliases, +//! etc.) for free, at the cost of a hard major-version pin against +//! `datafusion-ffi`. +//! +//! # Authoring a UDF +//! +//! Implement `ScalarUDFImpl` as you would for any DataFusion UDF, then: +//! +//! ```ignore +//! use comet_udf_sdk::comet_df_udf_export; +//! use datafusion::logical_expr::ScalarUDF; +//! use std::sync::Arc; +//! +//! fn make_add_one_df() -> Arc { +//! // ... build ScalarUDF from your impl ... +//! # unimplemented!() +//! } +//! +//! comet_df_udf_export!(make_add_one_df); +//! ``` + +use std::sync::Arc; + +use datafusion::logical_expr::ScalarUDF; +use datafusion_ffi::udf::FFI_ScalarUDF; + +/// List of UDFs exposed by a cdylib via `comet_df_udf_list_v1`. +/// +/// Ownership of the array transfers to the host, which frees via the +/// list's `release` callback after consuming each FFI_ScalarUDF. +#[repr(C)] +pub struct CometDfUdfList { + /// Pointer to the FFI_ScalarUDF array, or null if `len == 0`. + pub udfs: *mut FFI_ScalarUDF, + /// Number of entries. + pub len: i64, + /// Release the array storage. Each FFI_ScalarUDF entry is released + /// individually by its own Drop / clone_release machinery; this + /// callback only frees the heap-allocated array itself. + pub release: Option, +} + +impl Default for CometDfUdfList { + fn default() -> Self { + Self { udfs: std::ptr::null_mut(), len: 0, release: None } + } +} + +impl Drop for CometDfUdfList { + fn drop(&mut self) { + if let Some(release) = self.release.take() { + // SAFETY: release frees the array storage. The FFI_ScalarUDF + // entries themselves were already dropped (or moved out) by + // the host before invoking release. + unsafe { release(self) }; + } + } +} + +/// Build a heap-allocated `CometDfUdfList` from a vector of `Arc`. +pub fn build_udf_list(udfs: Vec>) -> CometDfUdfList { + if udfs.is_empty() { + return CometDfUdfList::default(); + } + let ffi_udfs: Vec = + udfs.into_iter().map(FFI_ScalarUDF::from).collect(); + let mut boxed = ffi_udfs.into_boxed_slice(); + let len = boxed.len() as i64; + let udfs_ptr = boxed.as_mut_ptr(); + std::mem::forget(boxed); + CometDfUdfList { + udfs: udfs_ptr, + len, + release: Some(c_list_release), + } +} + +unsafe extern "C" fn c_list_release(list: *mut CometDfUdfList) { + debug_assert!(!list.is_null()); + let list_ref = unsafe { &mut *list }; + if list_ref.udfs.is_null() || list_ref.len == 0 { + list_ref.release = None; + return; + } + let len = list_ref.len as usize; + // SAFETY: udfs was a Box<[FFI_ScalarUDF]> turned into raw ptr + + // forgotten in build_udf_list. Reconstruct and drop. Each + // FFI_ScalarUDF entry's Drop runs its own release callback. + let _ = unsafe { + Box::from_raw(std::slice::from_raw_parts_mut(list_ref.udfs, len)) + }; + list_ref.udfs = std::ptr::null_mut(); + list_ref.len = 0; + list_ref.release = None; +} + +/// Emit the datafusion-ffi discovery entry points. +/// +/// Each item is the path of a function `fn() -> Arc`. The +/// macro produces: +/// +/// - `extern "C" fn comet_udf_abi_version() -> u32` (if not already exported) +/// - `extern "C" fn comet_df_udf_list_v1(out: *mut CometDfUdfList) -> i32` +/// +/// Use the `_no_abi_version` variant to suppress the abi_version export +/// when emitting both `comet_c_udf_export!` and `comet_df_udf_export!` from +/// the same crate (`comet_c_udf_export!` already emits it). +#[macro_export] +macro_rules! comet_df_udf_export { + ( $( $factory:path ),+ $(,)? ) => { + $crate::comet_df_udf_export!(@list $($factory),+); + const _: () = { + #[no_mangle] + pub extern "C" fn comet_udf_abi_version() -> u32 { + $crate::COMET_UDF_ABI_VERSION + } + }; + }; + (@list $( $factory:path ),+ $(,)? ) => { + const _: () = { + #[no_mangle] + pub unsafe extern "C" fn comet_df_udf_list_v1( + out: *mut $crate::df_abi::CometDfUdfList, + ) -> i32 { + if out.is_null() { return -1; } + let udfs: Vec> = vec![ + $( $factory(), )+ + ]; + let list = $crate::df_abi::build_udf_list(udfs); + std::ptr::write(out, list); + 0 + } + }; + }; + (no_abi_version => $( $factory:path ),+ $(,)? ) => { + $crate::comet_df_udf_export!(@list $($factory),+); + }; +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ArrayRef, Int64Array}; + use arrow::datatypes::DataType; + use datafusion::common::DataFusionError; + use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, + Volatility, + }; + use std::any::Any; + + #[derive(Debug, PartialEq, Eq, Hash)] + struct AddOne { + signature: Signature, + } + impl AddOne { + fn new() -> Self { + Self { + signature: Signature::new( + TypeSignature::Exact(vec![DataType::Int64]), + Volatility::Immutable, + ), + } + } + } + impl ScalarUDFImpl for AddOne { + fn as_any(&self) -> &dyn Any { + self + } + fn name(&self) -> &str { + "add_one_df" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _: &[DataType]) -> datafusion::common::Result { + Ok(DataType::Int64) + } + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion::common::Result { + let arr: ArrayRef = match &args.args[0] { + ColumnarValue::Array(a) => Arc::clone(a), + ColumnarValue::Scalar(s) => s.to_array_of_size(args.number_rows)?, + }; + let a = arr + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Execution("not Int64".into()))?; + let out: Int64Array = a.iter().map(|v| v.map(|x| x + 1)).collect(); + Ok(ColumnarValue::Array(Arc::new(out))) + } + } + + #[test] + fn build_list_and_round_trip_through_ffi() { + let udf = Arc::new(ScalarUDF::from(AddOne::new())); + let mut list = build_udf_list(vec![udf]); + assert_eq!(list.len, 1); + + // Read the FFI_ScalarUDF entry and import it via the canonical + // From<&FFI_ScalarUDF> for Arc. + let entry: &FFI_ScalarUDF = unsafe { &*list.udfs }; + let foreign: Arc = entry.into(); + assert_eq!(foreign.name(), "add_one_df"); + + // Drop list to exercise the release path. + drop(foreign); + if let Some(rel) = list.release.take() { + unsafe { rel(&mut list) }; + } + } +} diff --git a/native/comet-udf-sdk/src/lib.rs b/native/comet-udf-sdk/src/lib.rs new file mode 100644 index 0000000000..9f4d8ef118 --- /dev/null +++ b/native/comet-udf-sdk/src/lib.rs @@ -0,0 +1,71 @@ +// 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. +//! +//! Two ABI flavors are available, behind cargo features: +//! +//! - **C ABI** (`c-abi` feature, on by default) — 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`. Works with any Arrow producer that speaks the +//! C Data Interface — including C/C++ implementations. +//! +//! - **DataFusion FFI ABI** (`df-abi` feature, on by default) — wraps the +//! user's `ScalarUDFImpl` as `datafusion_ffi::udf::FFI_ScalarUDF`. Larger +//! surface, full DataFusion features (variadic signatures, type coercion, +//! metadata-aware return types) for free, but the user's library is +//! coupled to the matching `datafusion-ffi` major version. +//! +//! A single library may export either or both ABIs; the host loader tries +//! both discovery functions and prefers the C ABI on duplicate names. +//! +//! See `docs/source/user-guide/latest/custom-rust-udfs.md` for an +//! end-to-end walkthrough. + +#![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; + +#[cfg(feature = "c-abi")] +pub mod c_abi; + +#[cfg(feature = "df-abi")] +pub mod df_abi; + +/// Symbol name of the discovery entry point exported by C-ABI cdylibs. +pub const C_ABI_DISCOVERY_SYMBOL: &str = "comet_c_udf_list_v1"; + +/// Symbol name of the discovery entry point exported by datafusion-ffi +/// cdylibs. +pub const DF_ABI_DISCOVERY_SYMBOL: &str = "comet_df_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 c58d446917..ef3befb4bc 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -56,6 +56,7 @@ paste = "1.0.14" datafusion = { workspace = true, features = ["parquet_encryption", "sql"] } datafusion-physical-expr-adapter = { workspace = true } datafusion-datasource = { workspace = true } +datafusion-ffi = { workspace = true } datafusion-spark = { workspace = true } once_cell = "1.18.0" datafusion-comet-common = { workspace = true } @@ -77,6 +78,8 @@ iceberg = { workspace = true } iceberg-storage-opendal = { workspace = true } serde_json = "1.0" uuid = "1.23.0" +libloading = "0.8" +comet-udf-sdk = { path = "../comet-udf-sdk" } [target.'cfg(target_os = "linux")'.dependencies] procfs = "0.18.0" @@ -92,6 +95,7 @@ lazy_static = "1.4" assertables = "9" hex = "0.4.3" datafusion-functions-nested = { version = "53.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 new file mode 100644 index 0000000000..b896fdf82a --- /dev/null +++ b/native/core/build.rs @@ -0,0 +1,43 @@ +// 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. + +fn main() { + // 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..7e6341a775 --- /dev/null +++ b/native/core/src/comet_rust_udf_bridge.rs @@ -0,0 +1,97 @@ +// 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. The +/// `args`/`return_type` fields require calling the C-ABI kernel's `init` +/// to discover a return type — for the comparison PR we punt on that and +/// return only the bits the Scala registry needs (`name`, `abi`). +fn udf_to_json(udf: &LoadedUdf) -> serde_json::Value { + let abi = match udf.abi { + crate::execution::rust_udf::loader::UdfAbi::C => "c-abi", + crate::execution::rust_udf::loader::UdfAbi::DataFusion => "datafusion-ffi", + }; + serde_json::json!({ + "name": udf.name, + "abi": abi, + }) +} + +/// 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 542c3d9536..c172d422d7 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -764,6 +764,47 @@ 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()), + )?); + let return_field = Arc::new(Field::new(&call.name, 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..ecc5ea6ea4 --- /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, lib.clone()); + return Ok(lib); + } + } + + let mut w = cache().write().unwrap(); + if let Some(lib) = w.get(&canonical).cloned() { + if canonical != raw { + w.insert(raw, lib.clone()); + } + return Ok(lib); + } + let loaded = Arc::new(load(&canonical)?); + w.insert(canonical.clone(), loaded.clone()); + if canonical != raw { + w.insert(raw, loaded.clone()); + } + Ok(loaded) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::execution::rust_udf::test_support::test_udfs_path; + + #[test] + fn same_path_returns_same_arc() { + let p = test_udfs_path(); + let a = get_or_load(&p).unwrap(); + let b = get_or_load(&p).unwrap(); + 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..11823d2ec7 --- /dev/null +++ b/native/core/src/execution/rust_udf/imported_c.rs @@ -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. + +//! 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::any::Any; +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. + 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 as_any(&self) -> &dyn Any { + self + } + + 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..eaa5bff1c4 --- /dev/null +++ b/native/core/src/execution/rust_udf/loader.rs @@ -0,0 +1,350 @@ +// 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 and/or datafusion-ffi entry points, and +//! produce DataFusion `ScalarUDFImpl` impls for each. +//! +//! See `super::mod.rs` for an overview of the two flavors. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use comet_udf_sdk::c_abi::CometCScalarKernelList; +use comet_udf_sdk::df_abi::CometDfUdfList; +use comet_udf_sdk::{ + ABI_VERSION_SYMBOL, COMET_UDF_ABI_VERSION, C_ABI_DISCOVERY_SYMBOL, DF_ABI_DISCOVERY_SYMBOL, +}; +use datafusion::logical_expr::ScalarUDFImpl; +use datafusion_ffi::udf::FFI_ScalarUDF; +use libloading::{Library, Symbol}; + +use super::imported_c::ImportedCScalarUdf; + +/// Which ABI flavor a given UDF was loaded through. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UdfAbi { + /// Loaded via the pure-C arrow-ffi flavor (sedona-style). + C, + /// Loaded via datafusion-ffi (`FFI_ScalarUDF`). + DataFusion, +} + +/// One loaded UDF: name, ABI flavor, and a `ScalarUDFImpl` ready to plug +/// into the planner. +pub struct LoadedUdf { + /// UDF name as exposed by the cdylib. + pub name: String, + /// Which ABI this UDF was discovered through. + pub abi: UdfAbi, + /// 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) + .field("abi", &self.abi) + .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, + /// The loaded `Library`. Held inside an `Arc` so loaded UDFs can + /// outlive lookups. Library is never unloaded for the process lifetime. + pub library: Arc, + /// One entry per UDF, with name and ScalarUDFImpl already built. + pub udfs: Vec, +} + +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 exposes neither `comet_c_udf_list_v1` nor + /// `comet_df_udf_list_v1`. + NoDiscovery { + /// Path of the offending library. + path: PathBuf, + }, + /// One of the discovery functions returned a non-zero rc, or a + /// kernel/UDF 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, + "{} exposes neither {C_ABI_DISCOVERY_SYMBOL} nor \ + {DF_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 mut udfs: Vec = Vec::new(); + let mut have_any = false; + + if let Some(c_kernels) = read_c_kernels(&library, &path)? { + have_any = true; + for udf in c_kernels { + udfs.push(udf); + } + } + + if let Some(df_udfs) = read_df_udfs(&library, &path)? { + have_any = true; + for udf in df_udfs { + // Prefer C-ABI registration on collision. + if udfs.iter().any(|u| u.name == udf.name) { + continue; + } + udfs.push(udf); + } + } + + if !have_any { + 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< + unsafe extern "C" fn(*mut CometCScalarKernelList) -> 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(), + abi: UdfAbi::C, + udf_impl: Arc::new(imported), + }); + } + } + // list's Drop releases the array storage. + drop(list); + Ok(Some(udfs)) +} + +fn read_df_udfs( + lib: &Library, + path: &Path, +) -> Result>, LoaderError> { + let sym: Symbol i32> = + match unsafe { lib.get(DF_ABI_DISCOVERY_SYMBOL.as_bytes()) } { + Ok(s) => s, + Err(_) => return Ok(None), + }; + let mut list = CometDfUdfList::default(); + // SAFETY: list is caller-allocated; the cdylib writes into it. + let rc = unsafe { sym(&mut list) }; + if rc != 0 { + return Err(LoaderError::Discovery { + path: path.to_path_buf(), + reason: format!("{DF_ABI_DISCOVERY_SYMBOL} returned rc={rc}"), + }); + } + let mut udfs = Vec::with_capacity(list.len.max(0) as usize); + if !list.udfs.is_null() && list.len > 0 { + let len = list.len as usize; + for i in 0..len { + // SAFETY: the FFI_ScalarUDF array was produced by the cdylib's + // `comet_df_udf_export!` and contains `len` valid entries. + let entry: &FFI_ScalarUDF = unsafe { &*list.udfs.add(i) }; + // Convert to ScalarUDFImpl via the canonical From impl; + // the resulting Arc clones the FFI wrapper internally. + let imp: Arc = entry.into(); + udfs.push(LoadedUdf { + name: imp.name().to_string(), + abi: UdfAbi::DataFusion, + udf_impl: imp, + }); + } + } + // list's Drop releases the array storage; each FFI_ScalarUDF entry + // remains valid for the lifetime of the loaded Library, since the + // resulting ScalarUDFImpls cloned the wrappers via the From impl. + drop(list); + Ok(Some(udfs)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::execution::rust_udf::test_support::test_udfs_path; + + #[test] + fn load_test_udfs_succeeds() { + let lib = load(test_udfs_path()).expect("load"); + let names: Vec<_> = lib.udfs.iter().map(|u| u.name.as_str()).collect(); + assert!(names.contains(&"add_one_c"), "names: {names:?}"); + assert!(names.contains(&"add_one_df"), "names: {names:?}"); + } + + #[test] + fn missing_path_errors_open() { + let err = load("/no/such/path.so").unwrap_err(); + assert!(matches!(err, LoaderError::Open { .. }), "got: {err:?}"); + } + + #[test] + fn c_and_df_abi_have_expected_flavors() { + let lib = load(test_udfs_path()).expect("load"); + let c_udf = lib.udfs.iter().find(|u| u.name == "add_one_c").unwrap(); + let df_udf = lib.udfs.iter().find(|u| u.name == "add_one_df").unwrap(); + assert_eq!(c_udf.abi, UdfAbi::C); + assert_eq!(df_udf.abi, UdfAbi::DataFusion); + } +} 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..9bb1fb345f --- /dev/null +++ b/native/core/src/execution/rust_udf/mod.rs @@ -0,0 +1,43 @@ +// 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. +//! +//! Two ABI flavors are supported, both built only on Arrow's stable FFI +//! (the C Data Interface): +//! +//! - **C ABI** — `comet_c_udf_list_v1` returns sedona-style +//! `CometCScalarKernel` factory structs. Decoupled from datafusion +//! versions; suitable for cross-language (C/C++) extensions. +//! +//! - **datafusion-ffi ABI** — `comet_df_udf_list_v1` returns +//! `FFI_ScalarUDF` values. The user's library inherits the entire +//! `ScalarUDFImpl` surface for free, at the cost of a major-version +//! pin against `datafusion-ffi`. +//! +//! A single library may export either or both. The loader walks both +//! discovery functions and exposes whichever it finds; in the rare case +//! of a name collision the C ABI flavor wins (favoring the +//! version-decoupled implementation). + +pub mod cache; +pub mod imported_c; +pub mod loader; + +#[cfg(test)] +pub(crate) mod test_support; 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..a1ea8ac643 --- /dev/null +++ b/native/core/src/execution/rust_udf/test_support.rs @@ -0,0 +1,24 @@ +// 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. + +/// 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 f4ae0b8834..aaa08e0fc2 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -67,6 +67,7 @@ use errors::{try_unwrap_or_throw, CometError, CometResult}; pub mod execution; pub mod parquet; +pub mod comet_rust_udf_bridge; // this module is for non release only. Intended for debugging/profiling purposes #[cfg(debug_assertions)] pub mod debug; diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index 90e3d87032..7ff3b0f462 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -91,6 +91,7 @@ message Expr { HoursTransform hours_transform = 68; ArraysZip arrays_zip = 69; JvmScalarUdf jvm_scalar_udf = 70; + RustUdfCall rust_udf_call = 71; } // Optional QueryContext for error reporting (contains SQL text and position) @@ -530,3 +531,22 @@ 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). + bool deterministic = 5; +} diff --git a/pom.xml b/pom.xml index 7419fecc92..9e83cae5f5 100644 --- a/pom.xml +++ b/pom.xml @@ -917,6 +917,8 @@ 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 852e80ae44..2ca0c183cc 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -28,6 +28,7 @@ import org.apache.comet.CometSparkSessionExtensions.withInfo 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 /** @@ -52,8 +53,42 @@ 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)) { + withInfo(expr, "one or more Rust UDF arguments are not supported", expr.children: _*) + return None + } + val returnTypeProto = serializeDataType(returnType).getOrElse { + withInfo(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..44d41c2e3d --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala @@ -0,0 +1,140 @@ +/* + * 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 + +/** + * Public entry point for registering Rust scalar UDFs with Comet. + * + * Two ABI flavors are supported transparently -- the cdylib registers each UDF under one of: + * - `c-abi`: pure C / Arrow C Data Interface (sedona-style) + * - `datafusion-ffi`: `FFI_ScalarUDF` + * + * The user always calls `register` / `registerAll`; the native side picks the right ABI by name. + */ +object CometRustUDF { + + /** Spark conf key under which registered Rust UDF entries are propagated to executors. */ + val RUST_UDFS_CONF_KEY = "spark.comet.rustUdfs" + + 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), + * the driver-side registry is updated, and `spark.comet.rustUdfs` is updated so executors see + * the registration. + */ + def register( + spark: SparkSession, + name: String, + libraryPath: String, + inputTypes: Seq[DataType], + returnType: DataType, + deterministic: Boolean = true): Unit = { + 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) + propagateConf(spark) + } + + // -------- internals -------- + + private case class Described(name: String, abi: 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(), abi = node.get("abi").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("exposes neither")) { + 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 " + + s"or open a feature request to extend stub coverage.") + } + val finalUdf = if (deterministic) u else u.asNondeterministic() + spark.udf.register(name, finalUdf) + } + + private def propagateConf(spark: SparkSession): Unit = { + val snapshot = CometRustUdfRegistry.instance.snapshot + val entries = snapshot.toSeq.map { case (name, meta) => + mapper.writeValueAsString(java.util.Map.of("name", name, "libraryPath", meta.libraryPath)) + } + spark.conf.set(RUST_UDFS_CONF_KEY, entries.mkString(";")) + } +} 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..3b239aa6a5 --- /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 " + + s"stub was invoked, which means Comet did not replace this expression " + + s"with a native call. Check that Comet is enabled for the operator " + + s"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..ecc12ab458 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/udf/CometRustUdfRegistry.scala @@ -0,0 +1,59 @@ +/* + * 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 scala.jdk.CollectionConverters._ + +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)) + + /** Snapshot the registered set as an immutable Map. */ + def snapshot: Map[String, RustUdfMetadata] = + byName.asScala.toMap +} + +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..fd4f8553a3 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala @@ -0,0 +1,66 @@ +/* + * 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 org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.types.LongType + +import org.apache.comet.udf.CometRustUDF + +/** + * End-to-end integration suite: register a Rust UDF, run a Spark query, verify the result. + * + * Exercises both ABI flavors that the test cdylib exports: + * - `add_one_c` via the pure C ABI (Arrow C Data Interface only) + * - `add_one_df` via datafusion-ffi (`FFI_ScalarUDF`) + * + * Requires the test cdylib at the path given by the system property `comet.test.udfs.lib`. + * + * To run locally: + * {{{ + * cargo build -p comet-test-udfs --manifest-path native/Cargo.toml + * ./mvnw test -pl spark -DwildcardSuites="CometRustUdfSuite" \ + * -Dcomet.test.udfs.lib=$PWD/native/target/debug/libcomet_test_udfs.dylib + * }}} + */ +class CometRustUdfSuite extends CometTestBase { + + private lazy val libPath: String = { + val p = System.getProperty("comet.test.udfs.lib") + if (p == null) { + cancel("set -Dcomet.test.udfs.lib=; skipping without cdylib") + } + p + } + + test("C ABI: 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("datafusion-ffi: add_one_df returns id + 1 for a range") { + CometRustUDF.register(spark, "add_one_df", libPath, Seq(LongType), LongType) + val df = spark.range(0, 5).selectExpr("add_one_df(id) AS y") + val out = df.collect().map(_.getLong(0)).sorted.toSeq + assert(out == Seq(1L, 2L, 3L, 4L, 5L)) + } +} From 0f35933e8b724dc78748146d07b4829977c69217 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 3 Aug 2026 19:16:37 -0600 Subject: [PATCH 02/11] feat: drop datafusion-ffi flavor, contain panics at the FFI boundary, add user guide Keep a single UDF ABI built only on the Arrow C Data Interface. Carrying two public ABIs would mean two permanent compatibility promises; the datafusion-ffi flavor pins every user cdylib to Comet's DataFusion major version, which the 53 -> 54 upgrade demonstrated by removing as_any from ScalarUDFImpl. The rationale is recorded in comet-udf-sdk's crate docs. Contain panics at every extern "C" entry point. A panic escaping one aborts the process, taking down the executor JVM and every task on it. User UDF code is arbitrary and panicking is idiomatic Rust, so a panic is now converted to a query error via the existing get_last_error channel. Previously only execute was guarded; init, new_impl, the discovery entry point and the release callbacks were not. Add a user guide page describing the feature as experimental, including the executor-side library placement requirement and the trust implications of loading native code. --- docs/source/user-guide/latest/index.rst | 1 + docs/source/user-guide/latest/rust_udfs.md | 171 +++++++++++++ native/Cargo.lock | 183 +------------- native/Cargo.toml | 1 - native/comet-test-udfs/Cargo.toml | 1 - native/comet-test-udfs/src/lib.rs | 94 +++---- native/comet-udf-sdk/Cargo.toml | 9 - native/comet-udf-sdk/src/c_abi.rs | 233 +++++++++++++---- native/comet-udf-sdk/src/df_abi.rs | 236 ------------------ native/comet-udf-sdk/src/lib.rs | 50 ++-- native/core/Cargo.toml | 1 - native/core/src/comet_rust_udf_bridge.rs | 14 +- native/core/src/execution/rust_udf/loader.rs | 118 ++------- native/core/src/execution/rust_udf/mod.rs | 22 +- .../org/apache/comet/udf/CometRustUDF.scala | 14 +- .../org/apache/comet/CometRustUdfSuite.scala | 43 +++- 16 files changed, 492 insertions(+), 699 deletions(-) create mode 100644 docs/source/user-guide/latest/rust_udfs.md delete mode 100644 native/comet-udf-sdk/src/df_abi.rs 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..91befe8b7a --- /dev/null +++ b/docs/source/user-guide/latest/rust_udfs.md @@ -0,0 +1,171 @@ + + +# 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. The API, the ABI, and +> the registration mechanism may change in any release, without the usual deprecation cycle, and a +> UDF library built against one Comet version may need changes to work with the next. 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() +``` + +The declared `inputTypes` and `returnType` must match what the Rust `return_field` accepts and +returns. They are what Spark uses for analysis and planning; the Rust side is what actually runs. + +Set `deterministic = false` when the function is not a pure function of its arguments. + +## 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. +- **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 7d05def9ae..41e0fdc540 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -520,12 +520,6 @@ dependencies = [ "slab", ] -[[package]] -name = "async-ffi" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39cd9de47399986d5b216c6bef9434dfff1689ab61ba8d1e2720dc5fe5c84083" - [[package]] name = "async-global-executor" version = "2.4.1" @@ -1482,7 +1476,6 @@ version = "1.1.0" dependencies = [ "arrow", "comet-udf-sdk", - "datafusion", ] [[package]] @@ -1490,8 +1483,6 @@ name = "comet-udf-sdk" version = "1.1.0" dependencies = [ "arrow", - "datafusion", - "datafusion-ffi", ] [[package]] @@ -1977,7 +1968,6 @@ dependencies = [ "datafusion-comet-shuffle", "datafusion-comet-spark-expr", "datafusion-datasource", - "datafusion-ffi", "datafusion-functions-nested", "datafusion-physical-expr-adapter", "datafusion-spark", @@ -1990,7 +1980,7 @@ dependencies = [ "itertools 0.15.0", "jni 0.22.4", "lazy_static", - "libloading 0.8.9", + "libloading", "log", "log4rs", "mimalloc", @@ -2344,39 +2334,6 @@ dependencies = [ "itertools 0.14.0", ] -[[package]] -name = "datafusion-ffi" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9a4a09f24f1387408810acca070207a9cf6c7fc2368122982a8e7671f643d1e" -dependencies = [ - "arrow", - "arrow-schema", - "async-ffi", - "async-trait", - "chrono", - "datafusion-catalog", - "datafusion-common", - "datafusion-datasource", - "datafusion-execution", - "datafusion-expr", - "datafusion-functions-aggregate-common", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "datafusion-physical-optimizer", - "datafusion-physical-plan", - "datafusion-proto", - "datafusion-proto-common", - "datafusion-session", - "futures", - "libloading 0.9.0", - "log", - "prost", - "semver", - "stabby", - "tokio", -] - [[package]] name = "datafusion-functions" version = "54.1.0" @@ -2644,44 +2601,6 @@ dependencies = [ "tokio", ] -[[package]] -name = "datafusion-proto" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67791dcfacd142a9d95f8f73c2a161094cc936d231d80c235f5017dacf24e84e" -dependencies = [ - "arrow", - "chrono", - "datafusion-catalog", - "datafusion-catalog-listing", - "datafusion-common", - "datafusion-datasource", - "datafusion-datasource-arrow", - "datafusion-datasource-csv", - "datafusion-datasource-json", - "datafusion-datasource-parquet", - "datafusion-execution", - "datafusion-expr", - "datafusion-functions-table", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "datafusion-physical-plan", - "datafusion-proto-common", - "object_store", - "prost", -] - -[[package]] -name = "datafusion-proto-common" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8cd9e80d637891645d074db0f6c650b591117367247deb313fdfb78dff559cb" -dependencies = [ - "arrow", - "datafusion-common", - "prost", -] - [[package]] name = "datafusion-pruning" version = "54.1.0" @@ -3990,7 +3909,7 @@ dependencies = [ "java-locator", "jni-macros", "jni-sys 0.4.1", - "libloading 0.8.9", + "libloading", "log", "simd_cesu8", "thiserror 2.0.19", @@ -4157,16 +4076,6 @@ dependencies = [ "windows-link", ] -[[package]] -name = "libloading" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" -dependencies = [ - "cfg-if", - "windows-link", -] - [[package]] name = "liblzma" version = "0.4.7" @@ -5271,15 +5180,6 @@ dependencies = [ "syn 2.0.119", ] -[[package]] -name = "proc-macro-crate" -version = "3.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" -dependencies = [ - "toml_edit", -] - [[package]] name = "proc-macro2" version = "1.0.107" @@ -6312,12 +6212,6 @@ dependencies = [ "digest 0.11.3", ] -[[package]] -name = "sha2-const-stable" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f179d4e11094a893b82fff208f74d448a7512f99f5a0acbd5c679b705f83ed9" - [[package]] name = "shlex" version = "2.0.1" @@ -6446,40 +6340,6 @@ dependencies = [ "syn 2.0.119", ] -[[package]] -name = "stabby" -version = "72.1.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d53d2428934c46277fafd2d41e39357595aa1e47954c75db2b14ed90632f3cc" -dependencies = [ - "rustversion", - "stabby-abi", -] - -[[package]] -name = "stabby-abi" -version = "72.1.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f375eae680bb54203ee5e47d4cd2ae7b79c0a79ed90919279f38f500ad53f190" -dependencies = [ - "rustc_version", - "rustversion", - "sha2-const-stable", - "stabby-macros", -] - -[[package]] -name = "stabby-macros" -version = "72.1.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea664671a576c5f7e32fee291ac123d82af5e92b0689beb3555347c00c76eef1" -dependencies = [ - "proc-macro-crate", - "proc-macro2", - "quote", - "syn 2.0.119", -] - [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -6838,36 +6698,6 @@ dependencies = [ "tokio", ] -[[package]] -name = "toml_datetime" -version = "1.1.1+spec-1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" -dependencies = [ - "serde_core", -] - -[[package]] -name = "toml_edit" -version = "0.25.13+spec-1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" -dependencies = [ - "indexmap 2.14.0", - "toml_datetime", - "toml_parser", - "winnow", -] - -[[package]] -name = "toml_parser" -version = "1.1.3+spec-1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" -dependencies = [ - "winnow", -] - [[package]] name = "tower" version = "0.5.3" @@ -7521,15 +7351,6 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" -[[package]] -name = "winnow" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" -dependencies = [ - "memchr", -] - [[package]] name = "wit-bindgen" version = "0.57.1" diff --git a/native/Cargo.toml b/native/Cargo.toml index b7414f2bcc..c41ead3ccf 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -40,7 +40,6 @@ bytes = { version = "1.11.1" } parquet = { version = "58.4.0", default-features = false, features = ["experimental"] } datafusion = { version = "54.1.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] } datafusion-datasource = { version = "54.1.0" } -datafusion-ffi = { version = "54.1.0" } datafusion-physical-expr-adapter = { version = "54.1.0" } datafusion-spark = { version = "54.1.0", features = ["core"] } datafusion-comet-spark-expr = { path = "spark-expr" } diff --git a/native/comet-test-udfs/Cargo.toml b/native/comet-test-udfs/Cargo.toml index 90dc50b8a5..e07a08bb06 100644 --- a/native/comet-test-udfs/Cargo.toml +++ b/native/comet-test-udfs/Cargo.toml @@ -32,5 +32,4 @@ crate-type = ["cdylib"] [dependencies] arrow = { workspace = true } -datafusion = { 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 index 5466cffe55..84c0de55bc 100644 --- a/native/comet-test-udfs/src/lib.rs +++ b/native/comet-test-udfs/src/lib.rs @@ -17,11 +17,11 @@ //! Test UDF cdylib for Comet's Rust UDF host tests. //! -//! Exports the same `add_one` semantics through both ABI flavors so the -//! host can verify them side-by-side: +//! Exports `add_one_c`, an `(Int64) -> Int64` scalar function, through the +//! Comet UDF C ABI. //! -//! - `add_one_c` — exposed via the C ABI (arrow-only FFI, sedona-style) -//! - `add_one_df` — exposed via the datafusion-ffi ABI (`FFI_ScalarUDF`) +//! 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}`. @@ -29,16 +29,9 @@ use std::sync::Arc; use arrow::array::{ArrayRef, Int64Array}; use arrow::datatypes::{DataType, Field}; -use datafusion::common::DataFusionError; -use datafusion::logical_expr::{ - ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, - Volatility, -}; use comet_udf_sdk::c_abi::CometCScalarUdf; -use comet_udf_sdk::{comet_c_udf_export, comet_df_udf_export}; - -// ---------- C ABI implementation ------------------------------------ +use comet_udf_sdk::comet_c_udf_export; /// `add_one` exposed via the C ABI. pub struct AddOneC; @@ -77,62 +70,43 @@ impl CometCScalarUdf for AddOneC { } } -// ---------- datafusion-ffi implementation --------------------------- +/// 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; -/// `add_one` exposed via datafusion-ffi. -#[derive(Debug, PartialEq, Eq, Hash)] -pub struct AddOneDf { - signature: Signature, -} +impl CometCScalarUdf for PanicsOnInvoke { + fn name(&self) -> &str { + "panics_on_invoke" + } -impl AddOneDf { - fn new() -> Self { - Self { - signature: Signature::new( - TypeSignature::Exact(vec![DataType::Int64]), - Volatility::Immutable, - ), - } + 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") } } -impl ScalarUDFImpl for AddOneDf { +/// 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 { - "add_one_df" + "panics_on_return_field" } - fn signature(&self) -> &Signature { - &self.signature - } - fn return_type(&self, _args: &[DataType]) -> datafusion::common::Result { - Ok(DataType::Int64) - } - fn invoke_with_args( - &self, - args: ScalarFunctionArgs, - ) -> datafusion::common::Result { - let arr: ArrayRef = match &args.args[0] { - ColumnarValue::Array(a) => Arc::clone(a), - ColumnarValue::Scalar(s) => s.to_array_of_size(args.number_rows)?, - }; - let a = arr - .as_any() - .downcast_ref::() - .ok_or_else(|| DataFusionError::Execution("expected Int64".into()))?; - let out: Int64Array = a.iter().map(|v| v.map(|x| x + 1)).collect(); - Ok(ColumnarValue::Array(Arc::new(out))) + + fn return_field(&self, _args: &[Field]) -> Result { + panic!("deliberate panic from user return_field") } -} -/// Factory used by the discovery macro. -pub fn make_add_one_df() -> Arc { - Arc::new(ScalarUDF::from(AddOneDf::new())) + fn invoke(&self, _args: &[ArrayRef], _n_rows: usize) -> Result { + unreachable!("return_field panics first") + } } -// ---------- discovery exports -------------------------------------- - -// `comet_c_udf_export!` emits both `comet_c_udf_list_v1` and -// `comet_udf_abi_version`. `comet_df_udf_export!` would emit -// `comet_udf_abi_version` again, so we use the `no_abi_version` form to -// avoid the duplicate symbol. -comet_c_udf_export!(AddOneC); -comet_df_udf_export!(no_abi_version => make_add_one_df); +comet_c_udf_export!(AddOneC, PanicsOnInvoke, PanicsOnReturnField); diff --git a/native/comet-udf-sdk/Cargo.toml b/native/comet-udf-sdk/Cargo.toml index 766a095358..50830709bd 100644 --- a/native/comet-udf-sdk/Cargo.toml +++ b/native/comet-udf-sdk/Cargo.toml @@ -28,14 +28,5 @@ description = "SDK for writing custom Rust UDFs that run inside Apache DataFusio publish = false -[features] -# Enable the C ABI flavor (Arrow C Data Interface only, sedona-style). -c-abi = [] -# Enable the datafusion-ffi flavor (FFI_ScalarUDF / ForeignScalarUDF). -df-abi = ["dep:datafusion", "dep:datafusion-ffi"] -default = ["c-abi", "df-abi"] - [dependencies] arrow = { workspace = true } -datafusion = { workspace = true, optional = true } -datafusion-ffi = { workspace = true, optional = true } diff --git a/native/comet-udf-sdk/src/c_abi.rs b/native/comet-udf-sdk/src/c_abi.rs index 5b0b4577e2..fa4d73d57d 100644 --- a/native/comet-udf-sdk/src/c_abi.rs +++ b/native/comet-udf-sdk/src/c_abi.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Pure C ABI flavor — sedona-style. +//! 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 @@ -25,8 +25,9 @@ //! //! # Authoring a UDF //! -//! Implement [`CometCScalarUdf`] for a unit struct and use the -//! [`comet_c_udf_export!`] macro to emit the discovery entry point: +//! 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::*; @@ -34,6 +35,7 @@ //! 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" } @@ -61,6 +63,44 @@ use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; /// `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), swallowing +/// any panic. Used where the ABI gives us no way to report an error and +/// aborting would be a worse outcome than leaking. +fn catch_panic_infallible(f: impl FnOnce()) { + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)); +} + // -- factory struct -------------------------------------------------------- /// Factory for [`CometCScalarKernelImpl`] instances. @@ -317,15 +357,19 @@ unsafe extern "C" fn c_factory_new_impl( debug_assert!(!out.is_null()); let this = unsafe { &*this }; let exp = unsafe { &*(this.private_data as *const ExportedScalarKernel) }; - 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)); - } + // 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(|| { + 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) { @@ -333,9 +377,11 @@ unsafe extern "C" fn c_factory_release(this: *mut CometCScalarKernel) { 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. - let _ = unsafe { Box::from_raw(this_ref.private_data as *mut ExportedScalarKernel) }; + // 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(|| drop(unsafe { Box::from_raw(raw) })); } this_ref.function_name = None; this_ref.new_impl = None; @@ -397,7 +443,8 @@ unsafe extern "C" fn c_kernel_init( } } - match priv_ref.inner.return_field(&fields) { + // `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) }; @@ -476,22 +523,10 @@ unsafe extern "C" fn c_kernel_execute( arrays.push(arrow::array::make_array(data)); } - let result = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - priv_ref.inner.invoke(&arrays, n_rows as usize) - })) { - Ok(Ok(arr)) => arr, - Ok(Err(msg)) => { - priv_ref.last_error = std::ffi::CString::new(msg).unwrap_or_default(); - return C_ABI_ERR; - } - Err(panic) => { - let msg = match panic.downcast_ref::<&'static str>() { - Some(s) => format!("panic: {s}"), - None => match panic.downcast_ref::() { - Some(s) => format!("panic: {s}"), - None => "panic: ".to_string(), - }, - }; + // `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; } @@ -517,8 +552,10 @@ 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() { - let _ = unsafe { Box::from_raw(this_ref.private_data as *mut ExportedScalarKernelImpl) }; + // 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(|| drop(unsafe { Box::from_raw(raw) })); } this_ref.init = None; this_ref.execute = None; @@ -554,13 +591,17 @@ unsafe extern "C" fn c_list_release(list: *mut CometCScalarKernelList) { return; } let len = list_ref.len as usize; - // 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. - let _ = unsafe { Box::from_raw(std::ptr::slice_from_raw_parts_mut(list_ref.kernels, len)) }; + 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(|| { + drop(unsafe { Box::from_raw(std::ptr::slice_from_raw_parts_mut(kernels, len)) }) + }); } // -- export macro --------------------------------------------------------- @@ -573,6 +614,10 @@ unsafe extern "C" fn c_list_release(list: *mut CometCScalarKernelList) { /// - `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 { @@ -588,16 +633,25 @@ macro_rules! comet_c_udf_export { out: *mut $crate::c_abi::CometCScalarKernelList, ) -> i32 { if out.is_null() { return -1; } - let kernels: Vec<$crate::c_abi::CometCScalarKernel> = vec![ - $( - $crate::c_abi::CometCScalarKernel::from( - $crate::c_abi::ExportedScalarKernel::new(<$ty as Default>::default()) - ), - )+ - ]; - let list = $crate::c_abi::build_kernel_list(kernels); - std::ptr::write(out, list); - 0 + 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, + } } }; }; @@ -691,4 +745,91 @@ mod tests { 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/df_abi.rs b/native/comet-udf-sdk/src/df_abi.rs deleted file mode 100644 index 90243d82ec..0000000000 --- a/native/comet-udf-sdk/src/df_abi.rs +++ /dev/null @@ -1,236 +0,0 @@ -// 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. - -//! datafusion-ffi flavor. -//! -//! Discovery returns a list of `FFI_ScalarUDF` values produced by -//! `datafusion_ffi`. The host imports each via -//! `ForeignScalarUDF::try_from`, yielding a `ScalarUDFImpl` it can plug -//! straight into its existing planner — no further adaptation needed. -//! -//! The user's library inherits the entire `ScalarUDFImpl` surface -//! (signature with type coercion, metadata-aware return types, aliases, -//! etc.) for free, at the cost of a hard major-version pin against -//! `datafusion-ffi`. -//! -//! # Authoring a UDF -//! -//! Implement `ScalarUDFImpl` as you would for any DataFusion UDF, then: -//! -//! ```ignore -//! use comet_udf_sdk::comet_df_udf_export; -//! use datafusion::logical_expr::ScalarUDF; -//! use std::sync::Arc; -//! -//! fn make_add_one_df() -> Arc { -//! // ... build ScalarUDF from your impl ... -//! # unimplemented!() -//! } -//! -//! comet_df_udf_export!(make_add_one_df); -//! ``` - -use std::sync::Arc; - -use datafusion::logical_expr::ScalarUDF; -use datafusion_ffi::udf::FFI_ScalarUDF; - -/// List of UDFs exposed by a cdylib via `comet_df_udf_list_v1`. -/// -/// Ownership of the array transfers to the host, which frees via the -/// list's `release` callback after consuming each FFI_ScalarUDF. -#[repr(C)] -pub struct CometDfUdfList { - /// Pointer to the FFI_ScalarUDF array, or null if `len == 0`. - pub udfs: *mut FFI_ScalarUDF, - /// Number of entries. - pub len: i64, - /// Release the array storage. Each FFI_ScalarUDF entry is released - /// individually by its own Drop / clone_release machinery; this - /// callback only frees the heap-allocated array itself. - pub release: Option, -} - -impl Default for CometDfUdfList { - fn default() -> Self { - Self { - udfs: std::ptr::null_mut(), - len: 0, - release: None, - } - } -} - -impl Drop for CometDfUdfList { - fn drop(&mut self) { - if let Some(release) = self.release.take() { - // SAFETY: release frees the array storage. The FFI_ScalarUDF - // entries themselves were already dropped (or moved out) by - // the host before invoking release. - unsafe { release(self) }; - } - } -} - -/// Build a heap-allocated `CometDfUdfList` from a vector of `Arc`. -pub fn build_udf_list(udfs: Vec>) -> CometDfUdfList { - if udfs.is_empty() { - return CometDfUdfList::default(); - } - let ffi_udfs: Vec = udfs.into_iter().map(FFI_ScalarUDF::from).collect(); - let mut boxed = ffi_udfs.into_boxed_slice(); - let len = boxed.len() as i64; - let udfs_ptr = boxed.as_mut_ptr(); - std::mem::forget(boxed); - CometDfUdfList { - udfs: udfs_ptr, - len, - release: Some(c_list_release), - } -} - -unsafe extern "C" fn c_list_release(list: *mut CometDfUdfList) { - debug_assert!(!list.is_null()); - let list_ref = unsafe { &mut *list }; - if list_ref.udfs.is_null() || list_ref.len == 0 { - list_ref.release = None; - return; - } - let len = list_ref.len as usize; - // SAFETY: udfs was a Box<[FFI_ScalarUDF]> turned into raw ptr + - // forgotten in build_udf_list. Reconstruct and drop. Each - // FFI_ScalarUDF entry's Drop runs its own release callback. - let _ = unsafe { Box::from_raw(std::ptr::slice_from_raw_parts_mut(list_ref.udfs, len)) }; - list_ref.udfs = std::ptr::null_mut(); - list_ref.len = 0; - list_ref.release = None; -} - -/// Emit the datafusion-ffi discovery entry points. -/// -/// Each item is the path of a function `fn() -> Arc`. The -/// macro produces: -/// -/// - `extern "C" fn comet_udf_abi_version() -> u32` (if not already exported) -/// - `extern "C" fn comet_df_udf_list_v1(out: *mut CometDfUdfList) -> i32` -/// -/// Use the `_no_abi_version` variant to suppress the abi_version export -/// when emitting both `comet_c_udf_export!` and `comet_df_udf_export!` from -/// the same crate (`comet_c_udf_export!` already emits it). -#[macro_export] -macro_rules! comet_df_udf_export { - ( $( $factory:path ),+ $(,)? ) => { - $crate::comet_df_udf_export!(@list $($factory),+); - const _: () = { - #[no_mangle] - pub extern "C" fn comet_udf_abi_version() -> u32 { - $crate::COMET_UDF_ABI_VERSION - } - }; - }; - (@list $( $factory:path ),+ $(,)? ) => { - const _: () = { - #[no_mangle] - pub unsafe extern "C" fn comet_df_udf_list_v1( - out: *mut $crate::df_abi::CometDfUdfList, - ) -> i32 { - if out.is_null() { return -1; } - let udfs: Vec> = vec![ - $( $factory(), )+ - ]; - let list = $crate::df_abi::build_udf_list(udfs); - std::ptr::write(out, list); - 0 - } - }; - }; - (no_abi_version => $( $factory:path ),+ $(,)? ) => { - $crate::comet_df_udf_export!(@list $($factory),+); - }; -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::{ArrayRef, Int64Array}; - use arrow::datatypes::DataType; - use datafusion::common::DataFusionError; - use datafusion::logical_expr::{ - ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, - Volatility, - }; - - #[derive(Debug, PartialEq, Eq, Hash)] - struct AddOne { - signature: Signature, - } - impl AddOne { - fn new() -> Self { - Self { - signature: Signature::new( - TypeSignature::Exact(vec![DataType::Int64]), - Volatility::Immutable, - ), - } - } - } - impl ScalarUDFImpl for AddOne { - fn name(&self) -> &str { - "add_one_df" - } - fn signature(&self) -> &Signature { - &self.signature - } - fn return_type(&self, _: &[DataType]) -> datafusion::common::Result { - Ok(DataType::Int64) - } - fn invoke_with_args( - &self, - args: ScalarFunctionArgs, - ) -> datafusion::common::Result { - let arr: ArrayRef = match &args.args[0] { - ColumnarValue::Array(a) => Arc::clone(a), - ColumnarValue::Scalar(s) => s.to_array_of_size(args.number_rows)?, - }; - let a = arr - .as_any() - .downcast_ref::() - .ok_or_else(|| DataFusionError::Execution("not Int64".into()))?; - let out: Int64Array = a.iter().map(|v| v.map(|x| x + 1)).collect(); - Ok(ColumnarValue::Array(Arc::new(out))) - } - } - - #[test] - fn build_list_and_round_trip_through_ffi() { - let udf = Arc::new(ScalarUDF::from(AddOne::new())); - let mut list = build_udf_list(vec![udf]); - assert_eq!(list.len, 1); - - // Read the FFI_ScalarUDF entry and import it via the canonical - // From<&FFI_ScalarUDF> for Arc. - let entry: &FFI_ScalarUDF = unsafe { &*list.udfs }; - let foreign: Arc = entry.into(); - assert_eq!(foreign.name(), "add_one_df"); - - // Drop list to exercise the release path. - drop(foreign); - if let Some(rel) = list.release.take() { - unsafe { rel(&mut list) }; - } - } -} diff --git a/native/comet-udf-sdk/src/lib.rs b/native/comet-udf-sdk/src/lib.rs index 9f4d8ef118..5920141685 100644 --- a/native/comet-udf-sdk/src/lib.rs +++ b/native/comet-udf-sdk/src/lib.rs @@ -18,25 +18,35 @@ //! SDK for writing scalar UDFs in Rust that are loaded and executed by //! Apache DataFusion Comet, using only Arrow's stable FFI surface. //! -//! Two ABI flavors are available, behind cargo features: +//! 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. //! -//! - **C ABI** (`c-abi` feature, on by default) — 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`. Works with any Arrow producer that speaks the -//! C Data Interface — including C/C++ implementations. +//! # Why not `datafusion-ffi`? //! -//! - **DataFusion FFI ABI** (`df-abi` feature, on by default) — wraps the -//! user's `ScalarUDFImpl` as `datafusion_ffi::udf::FFI_ScalarUDF`. Larger -//! surface, full DataFusion features (variadic signatures, type coercion, -//! metadata-aware return types) for free, but the user's library is -//! coupled to the matching `datafusion-ffi` major version. +//! 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. //! -//! A single library may export either or both ABIs; the host loader tries -//! both discovery functions and prefers the C ABI on duplicate names. +//! 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. //! -//! See `docs/source/user-guide/latest/custom-rust-udfs.md` for an -//! end-to-end walkthrough. +//! 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)] @@ -44,19 +54,11 @@ /// the discovery entry-point signatures or to the FFI structs they yield. pub const COMET_UDF_ABI_VERSION: u32 = 1; -#[cfg(feature = "c-abi")] pub mod c_abi; -#[cfg(feature = "df-abi")] -pub mod df_abi; - -/// Symbol name of the discovery entry point exported by C-ABI cdylibs. +/// 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 discovery entry point exported by datafusion-ffi -/// cdylibs. -pub const DF_ABI_DISCOVERY_SYMBOL: &str = "comet_df_udf_list_v1"; - /// Symbol name of the ABI version probe exported by every cdylib. pub const ABI_VERSION_SYMBOL: &str = "comet_udf_abi_version"; diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index f2ac7cf13b..775a1c769f 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -56,7 +56,6 @@ paste = "1.0.14" datafusion = { workspace = true, features = ["parquet_encryption", "sql"] } datafusion-physical-expr-adapter = { workspace = true } datafusion-datasource = { workspace = true } -datafusion-ffi = { workspace = true } datafusion-spark = { workspace = true } once_cell = "1.18.0" datafusion-comet-common = { workspace = true } diff --git a/native/core/src/comet_rust_udf_bridge.rs b/native/core/src/comet_rust_udf_bridge.rs index cb53e90f47..533e076709 100644 --- a/native/core/src/comet_rust_udf_bridge.rs +++ b/native/core/src/comet_rust_udf_bridge.rs @@ -25,18 +25,14 @@ use jni::objects::{JClass, JString}; use jni::sys::jobject; use jni::EnvUnowned; -/// Best-effort serialization of a single discovered UDF as JSON. The -/// `args`/`return_type` fields require calling the C-ABI kernel's `init` -/// to discover a return type — for the comparison PR we punt on that and -/// return only the bits the Scala registry needs (`name`, `abi`). +/// 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 { - let abi = match udf.abi { - crate::execution::rust_udf::loader::UdfAbi::C => "c-abi", - crate::execution::rust_udf::loader::UdfAbi::DataFusion => "datafusion-ffi", - }; serde_json::json!({ "name": udf.name, - "abi": abi, }) } diff --git a/native/core/src/execution/rust_udf/loader.rs b/native/core/src/execution/rust_udf/loader.rs index 7be8a1f257..1b93a7f98b 100644 --- a/native/core/src/execution/rust_udf/loader.rs +++ b/native/core/src/execution/rust_udf/loader.rs @@ -16,41 +16,26 @@ // under the License. //! Loader: open a UDF cdylib via libloading, validate the ABI version, -//! discover UDFs via the C-ABI and/or datafusion-ffi entry points, and -//! produce DataFusion `ScalarUDFImpl` impls for each. +//! discover UDFs via the C-ABI entry point, and produce DataFusion +//! `ScalarUDFImpl` impls for each. //! -//! See `super::mod.rs` for an overview of the two flavors. +//! 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::df_abi::CometDfUdfList; -use comet_udf_sdk::{ - ABI_VERSION_SYMBOL, COMET_UDF_ABI_VERSION, C_ABI_DISCOVERY_SYMBOL, DF_ABI_DISCOVERY_SYMBOL, -}; +use comet_udf_sdk::{ABI_VERSION_SYMBOL, COMET_UDF_ABI_VERSION, C_ABI_DISCOVERY_SYMBOL}; use datafusion::logical_expr::ScalarUDFImpl; -use datafusion_ffi::udf::FFI_ScalarUDF; use libloading::{Library, Symbol}; use super::imported_c::ImportedCScalarUdf; -/// Which ABI flavor a given UDF was loaded through. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum UdfAbi { - /// Loaded via the pure-C arrow-ffi flavor (sedona-style). - C, - /// Loaded via datafusion-ffi (`FFI_ScalarUDF`). - DataFusion, -} - -/// One loaded UDF: name, ABI flavor, and a `ScalarUDFImpl` ready to plug -/// into the planner. +/// 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, - /// Which ABI this UDF was discovered through. - pub abi: UdfAbi, /// The `ScalarUDFImpl` adapter the planner will wrap. pub udf_impl: Arc, } @@ -59,7 +44,6 @@ 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) - .field("abi", &self.abi) .finish() } } @@ -104,14 +88,13 @@ pub enum LoaderError { /// Version this host expects. expected: u32, }, - /// Library exposes neither `comet_c_udf_list_v1` nor - /// `comet_df_udf_list_v1`. + /// Library does not expose `comet_c_udf_list_v1`. NoDiscovery { /// Path of the offending library. path: PathBuf, }, - /// One of the discovery functions returned a non-zero rc, or a - /// kernel/UDF entry was malformed. + /// The discovery function returned a non-zero rc, or a kernel entry + /// was malformed. Discovery { /// Path of the library. path: PathBuf, @@ -143,8 +126,7 @@ impl std::fmt::Display for LoaderError { }, NoDiscovery { path } => write!( f, - "{} exposes neither {C_ABI_DISCOVERY_SYMBOL} nor \ - {DF_ABI_DISCOVERY_SYMBOL}", + "{} does not export {C_ABI_DISCOVERY_SYMBOL}", path.display() ), Discovery { path, reason } => write!(f, "{}: {reason}", path.display()), @@ -182,30 +164,10 @@ pub fn load(path: impl AsRef) -> Result { }); } - let mut udfs: Vec = Vec::new(); - let mut have_any = false; - - if let Some(c_kernels) = read_c_kernels(&library, &path)? { - have_any = true; - for udf in c_kernels { - udfs.push(udf); - } - } - - if let Some(df_udfs) = read_df_udfs(&library, &path)? { - have_any = true; - for udf in df_udfs { - // Prefer C-ABI registration on collision. - if udfs.iter().any(|u| u.name == udf.name) { - continue; - } - udfs.push(udf); - } - } - - if !have_any { - return Err(LoaderError::NoDiscovery { path }); - } + let udfs = match read_c_kernels(&library, &path)? { + Some(udfs) => udfs, + None => return Err(LoaderError::NoDiscovery { path }), + }; Ok(LoadedLibrary { path, @@ -268,7 +230,6 @@ fn read_c_kernels(lib: &Library, path: &Path) -> Result>, })?; udfs.push(LoadedUdf { name: imported.name().to_string(), - abi: UdfAbi::C, udf_impl: Arc::new(imported), }); } @@ -278,45 +239,6 @@ fn read_c_kernels(lib: &Library, path: &Path) -> Result>, Ok(Some(udfs)) } -fn read_df_udfs(lib: &Library, path: &Path) -> Result>, LoaderError> { - let sym: Symbol i32> = - match unsafe { lib.get(DF_ABI_DISCOVERY_SYMBOL.as_bytes()) } { - Ok(s) => s, - Err(_) => return Ok(None), - }; - let mut list = CometDfUdfList::default(); - // SAFETY: list is caller-allocated; the cdylib writes into it. - let rc = unsafe { sym(&mut list) }; - if rc != 0 { - return Err(LoaderError::Discovery { - path: path.to_path_buf(), - reason: format!("{DF_ABI_DISCOVERY_SYMBOL} returned rc={rc}"), - }); - } - let mut udfs = Vec::with_capacity(list.len.max(0) as usize); - if !list.udfs.is_null() && list.len > 0 { - let len = list.len as usize; - for i in 0..len { - // SAFETY: the FFI_ScalarUDF array was produced by the cdylib's - // `comet_df_udf_export!` and contains `len` valid entries. - let entry: &FFI_ScalarUDF = unsafe { &*list.udfs.add(i) }; - // Convert to ScalarUDFImpl via the canonical From impl; - // the resulting Arc clones the FFI wrapper internally. - let imp: Arc = entry.into(); - udfs.push(LoadedUdf { - name: imp.name().to_string(), - abi: UdfAbi::DataFusion, - udf_impl: imp, - }); - } - } - // list's Drop releases the array storage; each FFI_ScalarUDF entry - // remains valid for the lifetime of the loaded Library, since the - // resulting ScalarUDFImpls cloned the wrappers via the From impl. - drop(list); - Ok(Some(udfs)) -} - #[cfg(test)] mod tests { use super::*; @@ -327,7 +249,6 @@ mod tests { let lib = load(test_udfs_path()).expect("load"); let names: Vec<_> = lib.udfs.iter().map(|u| u.name.as_str()).collect(); assert!(names.contains(&"add_one_c"), "names: {names:?}"); - assert!(names.contains(&"add_one_df"), "names: {names:?}"); } #[test] @@ -336,12 +257,13 @@ mod tests { assert!(matches!(err, LoaderError::Open { .. }), "got: {err:?}"); } + /// A library exporting several kernels surfaces all of them. #[test] - fn c_and_df_abi_have_expected_flavors() { + fn all_exported_kernels_are_discovered() { let lib = load(test_udfs_path()).expect("load"); - let c_udf = lib.udfs.iter().find(|u| u.name == "add_one_c").unwrap(); - let df_udf = lib.udfs.iter().find(|u| u.name == "add_one_df").unwrap(); - assert_eq!(c_udf.abi, UdfAbi::C); - assert_eq!(df_udf.abi, UdfAbi::DataFusion); + let names: Vec<_> = lib.udfs.iter().map(|u| u.name.as_str()).collect(); + for expected in ["add_one_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 index 9bb1fb345f..c76ccd0f88 100644 --- a/native/core/src/execution/rust_udf/mod.rs +++ b/native/core/src/execution/rust_udf/mod.rs @@ -18,22 +18,14 @@ //! Loader and adapters for user-supplied Rust UDF cdylibs registered //! through `CometRustUDF` on the JVM side. //! -//! Two ABI flavors are supported, both built only on Arrow's stable FFI -//! (the C Data Interface): +//! 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++. //! -//! - **C ABI** — `comet_c_udf_list_v1` returns sedona-style -//! `CometCScalarKernel` factory structs. Decoupled from datafusion -//! versions; suitable for cross-language (C/C++) extensions. -//! -//! - **datafusion-ffi ABI** — `comet_df_udf_list_v1` returns -//! `FFI_ScalarUDF` values. The user's library inherits the entire -//! `ScalarUDFImpl` surface for free, at the cost of a major-version -//! pin against `datafusion-ffi`. -//! -//! A single library may export either or both. The loader walks both -//! discovery functions and exposes whichever it finds; in the rare case -//! of a name collision the C ABI flavor wins (favoring the -//! version-decoupled implementation). +//! See `comet_udf_sdk` for the rationale behind not exposing +//! `datafusion-ffi` here. pub mod cache; pub mod imported_c; diff --git a/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala index 44d41c2e3d..7b57241e78 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala @@ -32,11 +32,11 @@ import com.fasterxml.jackson.databind.node.ObjectNode /** * Public entry point for registering Rust scalar UDFs with Comet. * - * Two ABI flavors are supported transparently -- the cdylib registers each UDF under one of: - * - `c-abi`: pure C / Arrow C Data Interface (sedona-style) - * - `datafusion-ffi`: `FFI_ScalarUDF` + * 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. * - * The user always calls `register` / `registerAll`; the native side picks the right ABI by name. + * This is an experimental API and may change without the usual deprecation cycle. */ object CometRustUDF { @@ -70,7 +70,7 @@ object CometRustUDF { // -------- internals -------- - private case class Described(name: String, abi: String) + private case class Described(name: String) private def describeOne(libraryPath: String, name: String): Described = { val json = @@ -84,13 +84,13 @@ object CometRustUDF { private def parseDescribed(json: String): Described = { val node = mapper.readTree(json).asInstanceOf[ObjectNode] - Described(name = node.get("name").asText(), abi = node.get("abi").asText()) + 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("exposes neither")) { + 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) diff --git a/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala index fd4f8553a3..6ec0f4690f 100644 --- a/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala @@ -27,16 +27,16 @@ import org.apache.comet.udf.CometRustUDF /** * End-to-end integration suite: register a Rust UDF, run a Spark query, verify the result. * - * Exercises both ABI flavors that the test cdylib exports: - * - `add_one_c` via the pure C ABI (Arrow C Data Interface only) - * - `add_one_df` via datafusion-ffi (`FFI_ScalarUDF`) - * * Requires the test cdylib at the path given by the system property `comet.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 -pl spark -DwildcardSuites="CometRustUdfSuite" \ + * ./mvnw test -Dsuites="org.apache.comet.CometRustUdfSuite" -Dtest=none \ * -Dcomet.test.udfs.lib=$PWD/native/target/debug/libcomet_test_udfs.dylib * }}} */ @@ -50,17 +50,38 @@ class CometRustUdfSuite extends CometTestBase { p } - test("C ABI: add_one_c returns id + 1 for a range") { + 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("datafusion-ffi: add_one_df returns id + 1 for a range") { - CometRustUDF.register(spark, "add_one_df", libPath, Seq(LongType), LongType) - val df = spark.range(0, 5).selectExpr("add_one_df(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))) } } From 80d2cba7a00b7dc6f63d83701c307b4d8da0016f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 3 Aug 2026 19:25:11 -0600 Subject: [PATCH 03/11] refactor: remove the write-only spark.comet.rustUdfs conf propagateConf wrote the registered UDF set into spark.comet.rustUdfs on every register call, and nothing ever read it. Executors resolve the library from the library_path carried in the RustUdfCall proto, so the conf implied a propagation mechanism that does not exist. Document what actually happens on the register scaladoc instead. CometRustUdfRegistry.snapshot had no other caller and goes with it. --- .../org/apache/comet/udf/CometRustUDF.scala | 21 ++++++------------- .../comet/udf/CometRustUdfRegistry.scala | 6 ------ 2 files changed, 6 insertions(+), 21 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala index 7b57241e78..985552cb82 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala @@ -40,18 +40,18 @@ import com.fasterxml.jackson.databind.node.ObjectNode */ object CometRustUDF { - /** Spark conf key under which registered Rust UDF entries are propagated to executors. */ - val RUST_UDFS_CONF_KEY = "spark.comet.rustUdfs" - 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), - * the driver-side registry is updated, and `spark.comet.rustUdfs` is updated so executors see - * the registration. + * 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. */ def register( spark: SparkSession, @@ -65,7 +65,6 @@ object CometRustUDF { installCatalogStub(spark, name, inputTypes, returnType, deterministic) val meta = RustUdfMetadata(libraryPath, inputTypes, returnType, deterministic) CometRustUdfRegistry.instance.register(name, meta) - propagateConf(spark) } // -------- internals -------- @@ -129,12 +128,4 @@ object CometRustUDF { val finalUdf = if (deterministic) u else u.asNondeterministic() spark.udf.register(name, finalUdf) } - - private def propagateConf(spark: SparkSession): Unit = { - val snapshot = CometRustUdfRegistry.instance.snapshot - val entries = snapshot.toSeq.map { case (name, meta) => - mapper.writeValueAsString(java.util.Map.of("name", name, "libraryPath", meta.libraryPath)) - } - spark.conf.set(RUST_UDFS_CONF_KEY, entries.mkString(";")) - } } diff --git a/spark/src/main/scala/org/apache/comet/udf/CometRustUdfRegistry.scala b/spark/src/main/scala/org/apache/comet/udf/CometRustUdfRegistry.scala index ecc12ab458..2e277b6c01 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometRustUdfRegistry.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometRustUdfRegistry.scala @@ -21,8 +21,6 @@ package org.apache.comet.udf import java.util.concurrent.ConcurrentHashMap -import scala.jdk.CollectionConverters._ - import org.apache.spark.sql.types.DataType /** Metadata for a registered Rust UDF. */ @@ -46,10 +44,6 @@ class CometRustUdfRegistry { /** Return metadata for a name, if registered. */ def get(name: String): Option[RustUdfMetadata] = Option(byName.get(name)) - - /** Snapshot the registered set as an immutable Map. */ - def snapshot: Map[String, RustUdfMetadata] = - byName.asScala.toMap } object CometRustUdfRegistry { From ce1b0e41853011ee677a769e27fcd07936f9f39b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 3 Aug 2026 19:31:16 -0600 Subject: [PATCH 04/11] feat: cover all non-nested Spark types, validate the declared return type Add echo_c (identity over any type) and stringify_c ((any) -> Utf8) to the test cdylib, and drive both over every non-nested Spark type with a null row: bool, the four int widths, float, double, decimal, string, binary, date, timestamp and timestamp_ntz. echo_c proves the array survives the round trip with its type and nulls; stringify_c forces the UDF to decode the values rather than hand the array straight back. No ABI change was needed, since the FFI surface is type-agnostic, but nothing had demonstrated that beyond Int64. Also validate at plan time that the return type declared through CometRustUDF.register matches what the kernel's return_field reports, naming both types and pointing at register. Registering decimal(10,2) for a column Spark had widened to decimal(11,2) previously surfaced as a bare DataFusion type assertion partway through execution. Complex types (array, struct, map) remain future work and are documented as unsupported. --- docs/source/user-guide/latest/rust_udfs.md | 31 +++++++ native/comet-test-udfs/src/lib.rs | 86 +++++++++++++++++- native/core/src/execution/planner.rs | 19 ++++ native/core/src/execution/rust_udf/loader.rs | 8 +- .../org/apache/comet/CometRustUdfSuite.scala | 91 ++++++++++++++++++- 5 files changed, 229 insertions(+), 6 deletions(-) diff --git a/docs/source/user-guide/latest/rust_udfs.md b/docs/source/user-guide/latest/rust_udfs.md index 91befe8b7a..5498002b77 100644 --- a/docs/source/user-guide/latest/rust_udfs.md +++ b/docs/source/user-guide/latest/rust_udfs.md @@ -136,9 +136,38 @@ spark.range(0, 5).selectExpr("add_one(id) AS y").show() The declared `inputTypes` and `returnType` must match what the Rust `return_field` accepts and returns. They are what Spark uses for analysis and planning; the Rust side is what actually runs. +If the two disagree, the query fails with an error naming both types. + +Watch for Spark's own type promotion here: `cast(id as decimal(10,2)) + 0.25` has type +`decimal(11,2)`, not `decimal(10,2)`, so registering the latter is a mismatch. Set `deterministic = false` when the function is not a pure function of its arguments. +## 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)` | + +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. + +Complex types (`ArrayType`, `StructType`, `MapType`) are not supported yet. + ## Error handling Returning `Err(String)` from `return_field` or `invoke` fails the query with your message attached. @@ -153,6 +182,8 @@ flow mechanism: prefer returning `Err`, which produces a much better message. This feature is at an early stage. The current limitations are: - **Scalar functions only.** Aggregate, window, and table functions are not supported. +- **Non-nested types only.** See [Supported types](#supported-types); complex types are future + work. - **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 diff --git a/native/comet-test-udfs/src/lib.rs b/native/comet-test-udfs/src/lib.rs index 84c0de55bc..3e05ac05b1 100644 --- a/native/comet-test-udfs/src/lib.rs +++ b/native/comet-test-udfs/src/lib.rs @@ -17,8 +17,14 @@ //! Test UDF cdylib for Comet's Rust UDF host tests. //! -//! Exports `add_one_c`, an `(Int64) -> Int64` scalar function, through the -//! Comet UDF C ABI. +//! 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. @@ -27,8 +33,9 @@ use std::sync::Arc; -use arrow::array::{ArrayRef, Int64Array}; +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; @@ -70,6 +77,71 @@ impl CometCScalarUdf for AddOneC { } } +/// 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. @@ -109,4 +181,10 @@ impl CometCScalarUdf for PanicsOnReturnField { } } -comet_c_udf_export!(AddOneC, PanicsOnInvoke, PanicsOnReturnField); +comet_c_udf_export!( + AddOneC, + EchoC, + StringifyC, + PanicsOnInvoke, + PanicsOnReturnField +); diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 375fca0d80..a47d09f51e 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -922,6 +922,25 @@ impl PhysicalPlanner { .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 kernel_return_type != 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 + ))); + } + let return_field = Arc::new(Field::new(&call.name, return_type, true)); let expr = Arc::new(ScalarFunctionExpr::new( &call.name, diff --git a/native/core/src/execution/rust_udf/loader.rs b/native/core/src/execution/rust_udf/loader.rs index 1b93a7f98b..3861b8ce8b 100644 --- a/native/core/src/execution/rust_udf/loader.rs +++ b/native/core/src/execution/rust_udf/loader.rs @@ -262,7 +262,13 @@ mod tests { fn all_exported_kernels_are_discovered() { let lib = load(test_udfs_path()).expect("load"); let names: Vec<_> = lib.udfs.iter().map(|u| u.name.as_str()).collect(); - for expected in ["add_one_c", "panics_on_invoke", "panics_on_return_field"] { + 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/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala index 6ec0f4690f..c12231dd6f 100644 --- a/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala @@ -20,7 +20,7 @@ package org.apache.comet import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.types.LongType +import org.apache.spark.sql.types._ import org.apache.comet.udf.CometRustUDF @@ -84,4 +84,93 @@ class CometRustUdfSuite extends CometTestBase { .takeWhile(_ != null) .exists(t => Option(t.getMessage).exists(_.contains(needle))) } + + // ---------- type coverage ---------- + + /** + * One case per supported non-nested Spark type: the type itself, and a SQL expression producing + * a value of that type from the `id` column of `spark.range`. + * + * Complex types (array, struct, map) are not supported yet and are deliberately absent. + */ + private val primitiveTypeCases: 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))")) + + /** + * 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) <- primitiveTypeCases) { + test(s"echo_c round-trips ${dataType.simpleString} including nulls") { + CometRustUDF.register(spark, "echo_c", libPath, Seq(dataType), dataType) + val df = typedFrame(valueExpr) + 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("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() + } + } } From 299c16f0a119afea498010770460a14b6f8a5a7e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 3 Aug 2026 19:38:16 -0600 Subject: [PATCH 05/11] feat: support complex types, compare return types ignoring nested nullability Arrays, maps and structs work over the ABI, including nesting (array of struct, struct of array), so extend the type matrix to cover them alongside the non-nested types. No ABI change was needed; the FFI surface carries child arrays already. The new plan-time return type check rejected them, and was right to look: Spark carries containsNull and field nullability inside the declared type, so array with containsNull=false converts to a List with a non-nullable child, while the array actually delivered to the UDF has that child normalized to nullable. Compare with nested nullability erased, and promise DataFusion the kernel's own type rather than the declared one so its exact-match assertion compares like with like. The comparison stays strict about everything that changes how bytes are read: decimal precision and scale, timestamp unit, struct field names and order. Also document the return type model, which came up in review. A kernel has no fixed return type: return_field derives one from the argument types on every call, and echo_c serves all 19 types in the matrix. What is fixed is the type declared to register, because Spark needs a concrete DataType to plan against, and that is per-registration rather than per-kernel. --- docs/source/user-guide/latest/rust_udfs.md | 40 +++++-- native/core/src/execution/planner.rs | 10 +- native/core/src/execution/rust_udf/mod.rs | 113 ++++++++++++++++++ .../org/apache/comet/CometRustUdfSuite.scala | 72 +++++++++-- 4 files changed, 217 insertions(+), 18 deletions(-) diff --git a/docs/source/user-guide/latest/rust_udfs.md b/docs/source/user-guide/latest/rust_udfs.md index 5498002b77..0117651a65 100644 --- a/docs/source/user-guide/latest/rust_udfs.md +++ b/docs/source/user-guide/latest/rust_udfs.md @@ -134,14 +134,29 @@ The function is then callable from SQL or the DataFrame API like any other: spark.range(0, 5).selectExpr("add_one(id) AS y").show() ``` -The declared `inputTypes` and `returnType` must match what the Rust `return_field` accepts and -returns. They are what Spark uses for analysis and planning; the Rust side is what actually runs. -If the two disagree, the query fails with an error naming both types. +Set `deterministic = false` when the function is not a pure function of its arguments. -Watch for Spark's own type promotion here: `cast(id as decimal(10,2)) + 0.25` has type -`decimal(11,2)`, not `decimal(10,2)`, so registering the latter is a mismatch. +## Return types -Set `deterministic = false` when the function is not a pure function of its arguments. +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 @@ -163,10 +178,19 @@ Arguments and return values may be any of: | `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. -Complex types (`ArrayType`, `StructType`, `MapType`) are not supported yet. +Not yet supported: `CalendarIntervalType`, `NullType`, `UserDefinedType`, and Arrow extension types +such as Variant and Geometry. ## Error handling @@ -182,8 +206,6 @@ flow mechanism: prefer returning `Err`, which produces a much better message. This feature is at an early stage. The current limitations are: - **Scalar functions only.** Aggregate, window, and table functions are not supported. -- **Non-nested types only.** See [Supported types](#supported-types); complex types are future - work. - **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 diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index a47d09f51e..806dea6edc 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -932,7 +932,10 @@ impl PhysicalPlanner { .map(|e| e.data_type(input_schema.as_ref())) .collect::, _>>()?; let kernel_return_type = loaded.udf_impl.return_type(&arg_types)?; - if kernel_return_type != return_type { + 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:?}. \ @@ -941,7 +944,10 @@ impl PhysicalPlanner { ))); } - let return_field = Arc::new(Field::new(&call.name, return_type, true)); + // 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, diff --git a/native/core/src/execution/rust_udf/mod.rs b/native/core/src/execution/rust_udf/mod.rs index c76ccd0f88..3c9565ca82 100644 --- a/native/core/src/execution/rust_udf/mod.rs +++ b/native/core/src/execution/rust_udf/mod.rs @@ -33,3 +33,116 @@ 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/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala index c12231dd6f..8a87c3f140 100644 --- a/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala @@ -88,12 +88,14 @@ class CometRustUdfSuite extends CometTestBase { // ---------- type coverage ---------- /** - * One case per supported non-nested Spark type: the type itself, and a SQL expression producing - * a value of that type from the `id` column of `spark.range`. + * 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`. * - * Complex types (array, struct, map) are not supported yet and are deliberately absent. + * 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 primitiveTypeCases: Seq[(DataType, String)] = Seq( + private val typeCases: Seq[(DataType, String)] = Seq( (BooleanType, "id % 2 = 0"), (ByteType, "cast(id as byte)"), (ShortType, "cast(id as short)"), @@ -108,7 +110,30 @@ class CometRustUdfSuite extends CometTestBase { (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))")) + (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 @@ -123,10 +148,13 @@ class CometRustUdfSuite extends CometTestBase { case other => other } - for ((dataType, valueExpr) <- primitiveTypeCases) { + for ((dataType, valueExpr) <- typeCases) { test(s"echo_c round-trips ${dataType.simpleString} including nulls") { - CometRustUDF.register(spark, "echo_c", libPath, Seq(dataType), dataType) 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}") @@ -156,6 +184,36 @@ class CometRustUdfSuite extends CometTestBase { } } + 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) From 1d2e60e09fe8e1f9faa944e46393ee511a7dabb2 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 3 Aug 2026 19:57:14 -0600 Subject: [PATCH 06/11] ci: run CometRustUdfSuite in the PR builds The suite cancelled itself whenever -Dcomet.test.udfs.lib was unset, and nothing set it, so it had never run in CI and the ABI had never been exercised against a Linux .so. Have the suite locate the cdylib under native/target instead of requiring the flag, keeping the property as an override. When the library is missing the suite still skips locally, but fails in CI, where its absence means the native build or the artifact upload changed rather than that someone forgot a flag. An undefined Maven property arrives in the forked JVM as the literal string "null", which is why the override needs filtering rather than a plain null check. Upload libcomet_test_udfs alongside libcomet from both native build jobs: the same cargo invocation already produces it, since the crate is a workspace default member, and the test jobs download the artifact rather than building. Register the suite in the expressions bucket of both workflows, as dev/ci/check-suites.py requires. --- .github/workflows/pr_build_linux.yml | 7 ++- .github/workflows/pr_build_macos.yml | 7 ++- pom.xml | 3 +- .../org/apache/comet/CometRustUdfSuite.scala | 61 ++++++++++++++++--- 4 files changed, 68 insertions(+), 10 deletions(-) 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/pom.xml b/pom.xml index 38f5dd9cc5..b17865bb59 100644 --- a/pom.xml +++ b/pom.xml @@ -942,7 +942,8 @@ under the License. file:src/test/resources/log4j2.properties true ${project.build.directory}/tmp - + ${comet.test.udfs.lib} diff --git a/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala index 8a87c3f140..daf898e133 100644 --- a/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala @@ -19,6 +19,9 @@ package org.apache.comet +import java.io.File +import java.util.Locale + import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.types._ @@ -27,7 +30,8 @@ import org.apache.comet.udf.CometRustUDF /** * End-to-end integration suite: register a Rust UDF, run a Spark query, verify the result. * - * Requires the test cdylib at the path given by the system property `comet.test.udfs.lib`. + * 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 @@ -36,18 +40,31 @@ import org.apache.comet.udf.CometRustUDF * To run locally: * {{{ * cargo build -p comet-test-udfs --manifest-path native/Cargo.toml - * ./mvnw test -Dsuites="org.apache.comet.CometRustUdfSuite" -Dtest=none \ - * -Dcomet.test.udfs.lib=$PWD/native/target/debug/libcomet_test_udfs.dylib + * ./mvnw test -Dsuites="org.apache.comet.CometRustUdfSuite" -Dtest=none * }}} */ class CometRustUdfSuite extends CometTestBase { private lazy val libPath: String = { - val p = System.getProperty("comet.test.udfs.lib") - if (p == null) { - cancel("set -Dcomet.test.udfs.lib=; skipping without cdylib") + 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") + } } - p } test("add_one_c returns id + 1 for a range") { @@ -232,3 +249,33 @@ class CometRustUdfSuite extends CometTestBase { } } } + +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) + } +} From d8b8f0988471e837a30331bb1999771d8fbc27ff Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 3 Aug 2026 20:13:07 -0600 Subject: [PATCH 07/11] style: drop redundant string interpolators in the Rust UDF messages scalafix RedundantSyntax flags s-prefixed literals with nothing to interpolate. These were latent: the previous head commit on this branch carried [skip ci], so the lint jobs had never run against these files. --- .../src/main/scala/org/apache/comet/udf/CometRustUDF.scala | 2 +- .../scala/org/apache/comet/udf/CometRustUdfExceptions.scala | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala index 985552cb82..11c347c755 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala @@ -123,7 +123,7 @@ object CometRustUDF { case n => throw new IllegalArgumentException( s"Rust UDF '$name' arity $n not supported by stub. Reduce arity " + - s"or open a feature request to extend stub coverage.") + "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 index 3b239aa6a5..2d67d0ad81 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometRustUdfExceptions.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometRustUdfExceptions.scala @@ -46,6 +46,6 @@ class CometRustUdfSignatureException(msg: String) extends CometRuntimeException( class CometRustUdfNotEvaluatedException(name: String) extends CometRuntimeException( s"Rust UDF '$name' must run inside Comet native execution; the JVM " + - s"stub was invoked, which means Comet did not replace this expression " + - s"with a native call. Check that Comet is enabled for the operator " + - s"hosting this expression.") + "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.") From 5156b77e0c1d7d13a37150261dafe120f51bd205 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 4 Aug 2026 07:47:47 -0600 Subject: [PATCH 08/11] fix: drop a loaded UDF library after its kernels, not before LoadedLibrary declared `library: Arc` ahead of `udfs`. Struct fields drop in declaration order, so dropping a LoadedLibrary ran dlclose first and then invoked each kernel's `release` function pointer, which lives in the text of the library just unloaded. Nothing else holds a clone of that Arc. Unreachable in production, since the process-wide cache never drops a LoadedLibrary, but the loader tests build one via `load()` and drop it. Reorder the fields and say why in a comment, so a later edit does not quietly reintroduce it. --- native/core/src/execution/rust_udf/loader.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/native/core/src/execution/rust_udf/loader.rs b/native/core/src/execution/rust_udf/loader.rs index 3861b8ce8b..db72762e09 100644 --- a/native/core/src/execution/rust_udf/loader.rs +++ b/native/core/src/execution/rust_udf/loader.rs @@ -53,11 +53,16 @@ impl std::fmt::Debug for LoadedUdf { 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, - /// One entry per UDF, with name and ScalarUDFImpl already built. - pub udfs: Vec, } impl std::fmt::Debug for LoadedLibrary { From f20bc313f6fa0cfc499f03487f635f55abb1fbaa Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 4 Aug 2026 07:48:01 -0600 Subject: [PATCH 09/11] fix(ci): build the test UDF cdylib before the Rust test run The rust-test job runs `cargo nextest run` and nothing else. comet-test-udfs is `crate-type = ["cdylib"]` with no test targets, so a test build compiles the crate without ever emitting libcomet_test_udfs, and the three rust_udf tests that dlopen it failed on a clean checkout: execution::rust_udf::cache::tests::same_path_returns_same_arc execution::rust_udf::loader::tests::all_exported_kernels_are_discovered execution::rust_udf::loader::tests::load_test_udfs_succeeds It passed locally only because a previous `cargo build` had left the artifact in target/debug. Build the crate explicitly before the test run, and replace the bare `expect("load")` with a message naming that command, so the next person to hit this reads the fix instead of a dlopen error. --- .github/actions/rust-test/action.yaml | 9 +++++++++ native/core/src/execution/rust_udf/cache.rs | 6 +++--- native/core/src/execution/rust_udf/loader.rs | 6 +++--- native/core/src/execution/rust_udf/test_support.rs | 5 +++++ 4 files changed, 20 insertions(+), 6 deletions(-) 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/native/core/src/execution/rust_udf/cache.rs b/native/core/src/execution/rust_udf/cache.rs index 9f8ef2410c..c09449411e 100644 --- a/native/core/src/execution/rust_udf/cache.rs +++ b/native/core/src/execution/rust_udf/cache.rs @@ -69,13 +69,13 @@ pub fn get_or_load(path: impl AsRef) -> Result, LoaderE #[cfg(test)] mod tests { use super::*; - use crate::execution::rust_udf::test_support::test_udfs_path; + 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).unwrap(); - let b = get_or_load(&p).unwrap(); + let a = get_or_load(&p).expect(BUILD_HINT); + let b = get_or_load(&p).expect(BUILD_HINT); assert!(Arc::ptr_eq(&a, &b)); } diff --git a/native/core/src/execution/rust_udf/loader.rs b/native/core/src/execution/rust_udf/loader.rs index db72762e09..fce329b15d 100644 --- a/native/core/src/execution/rust_udf/loader.rs +++ b/native/core/src/execution/rust_udf/loader.rs @@ -247,11 +247,11 @@ fn read_c_kernels(lib: &Library, path: &Path) -> Result>, #[cfg(test)] mod tests { use super::*; - use crate::execution::rust_udf::test_support::test_udfs_path; + 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("load"); + 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:?}"); } @@ -265,7 +265,7 @@ mod tests { /// A library exporting several kernels surfaces all of them. #[test] fn all_exported_kernels_are_discovered() { - let lib = load(test_udfs_path()).expect("load"); + 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", diff --git a/native/core/src/execution/rust_udf/test_support.rs b/native/core/src/execution/rust_udf/test_support.rs index a1ea8ac643..d5637800c2 100644 --- a/native/core/src/execution/rust_udf/test_support.rs +++ b/native/core/src/execution/rust_udf/test_support.rs @@ -17,6 +17,11 @@ //! 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 { From 3e999c45d17000d5f4182979d848bf848de6ad24 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 4 Aug 2026 07:48:09 -0600 Subject: [PATCH 10/11] fix: reject deterministic = false when registering a Rust UDF `CometRustUDF.register` accepted a `deterministic` flag, installed a nondeterministic Spark stub for it, and carried it in the RustUdfCall proto, but nothing on the native side ever read it: ImportedCScalarUdf hardcodes Volatility::Immutable. A UDF the caller declared nondeterministic was therefore planned as pure and could be constant-folded, evaluated once and reused, or eliminated as a common subexpression. Honoring the flag needs the cached per-library signature to become per registration, which is #5249. Until then, fail the registration rather than let the parameter lie about what Comet does with it. Raised in review: does the proto field's value mean RustUdfs are always immutable? --- .../core/src/execution/rust_udf/imported_c.rs | 6 ++++ native/proto/src/proto/expr.proto | 5 +++ .../org/apache/comet/udf/CometRustUDF.scala | 23 ++++++++++++-- .../org/apache/comet/CometRustUdfSuite.scala | 31 +++++++++++++++++++ 4 files changed, 63 insertions(+), 2 deletions(-) diff --git a/native/core/src/execution/rust_udf/imported_c.rs b/native/core/src/execution/rust_udf/imported_c.rs index f48a4b36f1..353d117d9c 100644 --- a/native/core/src/execution/rust_udf/imported_c.rs +++ b/native/core/src/execution/rust_udf/imported_c.rs @@ -81,6 +81,12 @@ impl ImportedCScalarUdf { // 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 { diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index f6656c7193..c01a6836af 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -636,5 +636,10 @@ message RustUdfCall { // 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/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala index 11c347c755..19289df6ae 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala @@ -30,13 +30,17 @@ import com.fasterxml.jackson.databind.ObjectMapper import com.fasterxml.jackson.databind.node.ObjectNode /** - * Public entry point for registering Rust scalar UDFs with Comet. + * 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 and may change without the usual deprecation cycle. + * 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 { @@ -52,6 +56,10 @@ object CometRustUDF { * 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, @@ -60,6 +68,17 @@ object CometRustUDF { 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) diff --git a/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala index daf898e133..90db98a16a 100644 --- a/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala @@ -248,6 +248,37 @@ class CometRustUdfSuite extends CometTestBase { 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 { From ceb93655db719c307686d0f25d98935fa4ef02f1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 4 Aug 2026 07:48:17 -0600 Subject: [PATCH 11/11] docs: clarify the Rust UDF ABI's stability, purity, and path resolution Review follow-ups on the SDK and the user guide, plus two assertion tweaks: - State that the C ABI structs are specific to one Comet version and are internal under the versioning policy, while the host's own arrow and datafusion versions need not match the cdylib's. - Say on CometCScalarUdf that only immutable functions are supported, and add the same to the user guide's limitations. The guide previously told readers to pass `deterministic = false` for impure functions, which is now refused. - Describe what libraryPath actually does: an absolute path is the sensible choice, but a bare name resolves through the platform loader search path, and neither is a security boundary. - Note that CometRustUDF is deliberately not @Public, so it carries no compatibility guarantee. - Report rather than swallow a panic caught in a release callback: there is no error channel there, so it goes to stderr with a context string. - Check `release.is_some()` alongside private_data in the FFI debug asserts, which catches a call made after release rather than only an uninitialized one. --- docs/source/user-guide/latest/rust_udfs.md | 27 +++++++-- native/comet-udf-sdk/src/c_abi.rs | 70 +++++++++++++++++++--- 2 files changed, 83 insertions(+), 14 deletions(-) diff --git a/docs/source/user-guide/latest/rust_udfs.md b/docs/source/user-guide/latest/rust_udfs.md index 0117651a65..2f1e1eb862 100644 --- a/docs/source/user-guide/latest/rust_udfs.md +++ b/docs/source/user-guide/latest/rust_udfs.md @@ -26,10 +26,13 @@ This is different from [Scala UDF and Java UDF Support](scala_java_udfs.md), whe 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. The API, the ABI, and -> the registration mechanism may change in any release, without the usual deprecation cycle, and a -> UDF library built against one Comet version may need changes to work with the next. It is not yet -> recommended for production use. See [Limitations](#limitations) before adopting it. +> **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 @@ -134,7 +137,17 @@ The function is then callable from SQL or the DataFrame API like any other: spark.range(0, 5).selectExpr("add_one(id) AS y").show() ``` -Set `deterministic = false` when the function is not a pure function of its arguments. +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 @@ -206,6 +219,10 @@ flow mechanism: prefer returning `Err`, which produces a much better message. 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 diff --git a/native/comet-udf-sdk/src/c_abi.rs b/native/comet-udf-sdk/src/c_abi.rs index fa4d73d57d..f8e0a09036 100644 --- a/native/comet-udf-sdk/src/c_abi.rs +++ b/native/comet-udf-sdk/src/c_abi.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! The Comet UDF C ABI — sedona-style. +//! 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 @@ -23,6 +23,29 @@ //! 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`, @@ -94,11 +117,18 @@ fn catch_panic(f: impl FnOnce() -> Result) -> Result { } } -/// Run an infallible `f` (typically a release/cleanup callback), swallowing +/// 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. -fn catch_panic_infallible(f: impl FnOnce()) { - let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)); +/// +/// 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 -------------------------------------------------------- @@ -292,6 +322,21 @@ 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; @@ -344,7 +389,9 @@ impl From for CometCScalarKernel { unsafe extern "C" fn c_factory_function_name(this: *const CometCScalarKernel) -> *const c_char { debug_assert!(!this.is_null()); let this = unsafe { &*this }; - debug_assert!(!this.private_data.is_null()); + // 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() } @@ -356,10 +403,11 @@ unsafe extern "C" fn c_factory_new_impl( 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(|| { + catch_panic_infallible("constructing kernel impl", || { let impl_state = ExportedScalarKernelImpl { inner: std::sync::Arc::clone(&exp.inner), last_arg_fields: None, @@ -381,7 +429,7 @@ unsafe extern "C" fn c_factory_release(this: *mut CometCScalarKernel) { // 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(|| drop(unsafe { Box::from_raw(raw) })); + catch_panic_infallible("releasing kernel", || drop(unsafe { Box::from_raw(raw) })); } this_ref.function_name = None; this_ref.new_impl = None; @@ -418,6 +466,7 @@ unsafe extern "C" fn c_kernel_init( ) -> 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 }; @@ -474,6 +523,7 @@ unsafe extern "C" fn c_kernel_execute( ) -> 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 }; @@ -555,7 +605,9 @@ unsafe extern "C" fn c_kernel_release(this: *mut CometCScalarKernelImpl) { // 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(|| drop(unsafe { Box::from_raw(raw) })); + catch_panic_infallible("releasing kernel impl", || { + drop(unsafe { Box::from_raw(raw) }) + }); } this_ref.init = None; this_ref.execute = None; @@ -599,7 +651,7 @@ unsafe extern "C" fn c_list_release(list: *mut CometCScalarKernelList) { // 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(|| { + catch_panic_infallible("releasing kernel list", || { drop(unsafe { Box::from_raw(std::ptr::slice_from_raw_parts_mut(kernels, len)) }) }); }