From 0b40d0786e26c19816dcbf422b4f7f8df3ad5e95 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Sun, 25 Oct 2026 08:39:48 +0100 Subject: [PATCH 01/10] feat(udf): native Rust scalar UDFs end to end Sessions accept native Rust scalar functions through a small SDK (krishiv-udf-sdk): a trait with name, version, signature, return type, properties and a batch invoke over Arrow arrays plus a plain context (partition index, ANSI flag, session time zone). Plans reference UDFs by value (FunctionId::Udf { name, version }), which enters the canonical encoding (encoding version 5) and the structural hash. Resolution goes through a SessionFunctions registry wrapping the built-ins and the session library: bare names prefer built-ins (Spark's search path), while system.session.f and session.f reach UDFs only and system.builtin.f reaches built-ins only; unknown names keep UNRESOLVED_ROUTINE. The analyzer maps signatures onto its arity and input types so implicit casts match same-input built-ins. Execution binds calls in krishiv-functions with the same checks as built-ins (argument checks, all-literal broadcast, result length/type/nullability), unions null masks for Strict, and maps every user error, contract break or panic to FAILED_EXECUTE_UDF; a missing UDF on the executing side is the new KRISHIV_UDF_NOT_AVAILABLE naming name@version. Constant folding covers deterministic calls of literals in the analyzer (failures become Constant::Error) and the optimizer (failures stay in the plan), through the same binding. The plan codec gains the rust_udf semantic feature. Search-path precedence follows the Spark 4.2 run cited in the task (built-in wins, system.builtin qualification, UNRESOLVED_ROUTINE text); not re-verified with a new Spark run here. InvokeContext.partition is always 0: the engine has no stable partition ids. Session::register_udf takes the UDF by value; aggregate UDFs, dynamic library loading and UDF access to sessions remain out of scope. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- Cargo.lock | 10 + .../krishiv-common/krishiv-error-classes.json | 6 + .../src/translate/expression.rs | 2 +- .../src/translate/expression/tests.rs | 10 +- crates/krishiv-exec/Cargo.toml | 5 + crates/krishiv-exec/benches/udf.rs | 126 +++ .../krishiv-exec/src/physical/instantiate.rs | 21 +- crates/krishiv-exec/src/planner/build.rs | 6 +- crates/krishiv-exec/src/planner/eval/func.rs | 57 ++ crates/krishiv-exec/src/planner/eval/mod.rs | 155 ++- crates/krishiv-exec/src/planner/mod.rs | 13 + crates/krishiv-functions/Cargo.toml | 2 + crates/krishiv-functions/src/scalar/mod.rs | 2 + crates/krishiv-functions/src/scalar/udf.rs | 680 +++++++++++++ crates/krishiv-optimizer/Cargo.toml | 1 + crates/krishiv-optimizer/src/fold.rs | 192 +++- crates/krishiv-optimizer/src/lib.rs | 72 +- crates/krishiv-optimizer/tests/it/main.rs | 1 + crates/krishiv-optimizer/tests/it/udf.rs | 168 ++++ crates/krishiv-plan-codec/src/decode.rs | 5 + crates/krishiv-plan-codec/src/encode.rs | 4 + crates/krishiv-plan-codec/src/features.rs | 92 +- .../src/proto/krishiv.plan.v1.pb.rs | 21 +- .../tests/common/inventory.rs | 33 + crates/krishiv-plan-codec/tests/features.rs | 40 + .../tests/golden/rust_udf/rust_udf.bin | Bin 0 -> 825 bytes crates/krishiv-plan-codec/tests/roundtrip.rs | 2 +- crates/krishiv-plan/Cargo.toml | 2 + crates/krishiv-plan/src/analyzer/analytics.rs | 14 +- crates/krishiv-plan/src/analyzer/bind.rs | 23 +- crates/krishiv-plan/src/analyzer/catalog.rs | 16 +- crates/krishiv-plan/src/analyzer/exprs.rs | 66 +- crates/krishiv-plan/src/analyzer/fold.rs | 329 ++++++- crates/krishiv-plan/src/analyzer/ops.rs | 2 +- crates/krishiv-plan/src/analyzer/setops.rs | 19 +- crates/krishiv-plan/src/display.rs | 23 +- crates/krishiv-plan/src/encode.rs | 10 +- crates/krishiv-plan/src/expr.rs | 22 +- crates/krishiv-plan/src/function.rs | 22 +- crates/krishiv-plan/src/raises.rs | 2 + crates/krishiv-plan/src/tests/analyzer.rs | 8 +- crates/krishiv-plan/src/tests/dag.rs | 2 +- crates/krishiv-plan/src/tests/display.rs | 2 +- crates/krishiv-plan/src/tests/display_more.rs | 4 +- crates/krishiv-plan/src/tests/encode.rs | 8 +- crates/krishiv-plan/src/tests/nodes.rs | 2 +- crates/krishiv-plan/src/tests/props.rs | 20 +- crates/krishiv-plan/src/tests/rewrite.rs | 2 +- crates/krishiv-plan/src/tests/tpch.rs | 2 +- crates/krishiv-session/Cargo.toml | 1 + crates/krishiv-session/public-api.txt | 33 + crates/krishiv-session/src/column.rs | 5 +- crates/krishiv-session/src/lib.rs | 8 + crates/krishiv-session/src/session.rs | 81 +- crates/krishiv-session/src/udfs.rs | 474 +++++++++ crates/krishiv-session/tests/udf.rs | 905 ++++++++++++++++++ crates/krishiv-sql/src/lower.rs | 16 +- crates/krishiv-sql/src/tests.rs | 4 +- ...cts_lower_or_are_named_as_unsupported.snap | 3 +- crates/krishiv-udf-sdk/Cargo.toml | 4 + crates/krishiv-udf-sdk/src/lib.rs | 80 +- crates/krishiv-udf-sdk/src/library.rs | 180 ++++ crates/krishiv-udf-sdk/src/udf.rs | 380 ++++++++ crates/krishiv/public-api.txt | 11 + crates/krishiv/src/lib.rs | 7 +- proto/krishiv/plan/v1/plan.proto | 13 +- 66 files changed, 4366 insertions(+), 165 deletions(-) create mode 100644 crates/krishiv-exec/benches/udf.rs create mode 100644 crates/krishiv-functions/src/scalar/udf.rs create mode 100644 crates/krishiv-optimizer/tests/it/udf.rs create mode 100644 crates/krishiv-plan-codec/tests/golden/rust_udf/rust_udf.bin create mode 100644 crates/krishiv-session/src/udfs.rs create mode 100644 crates/krishiv-session/tests/udf.rs create mode 100644 crates/krishiv-udf-sdk/src/library.rs create mode 100644 crates/krishiv-udf-sdk/src/udf.rs diff --git a/Cargo.lock b/Cargo.lock index 949fbf95b..891d73b38 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1521,6 +1521,7 @@ dependencies = [ "krishiv-plan", "krishiv-sql", "krishiv-types", + "krishiv-udf-sdk", "loom", "object_store", "proptest", @@ -1540,12 +1541,14 @@ name = "krishiv-functions" version = "0.1.0-alpha" dependencies = [ "arrow-array", + "arrow-buffer", "arrow-schema", "arrow-select", "fancy-regex", "krishiv-common", "krishiv-kernels", "krishiv-types", + "krishiv-udf-sdk", "regex-automata", "regex-syntax", ] @@ -1633,6 +1636,7 @@ dependencies = [ "krishiv-plan", "krishiv-sql", "krishiv-types", + "krishiv-udf-sdk", "proptest", "serde_json", ] @@ -1649,6 +1653,7 @@ dependencies = [ "krishiv-functions", "krishiv-kernels", "krishiv-types", + "krishiv-udf-sdk", "proptest", "serde_json", "thiserror", @@ -1721,6 +1726,7 @@ dependencies = [ "krishiv-plan", "krishiv-sql", "krishiv-types", + "krishiv-udf-sdk", "object_store", "parquet", "serde_json", @@ -1809,6 +1815,10 @@ dependencies = [ [[package]] name = "krishiv-udf-sdk" version = "0.1.0-alpha" +dependencies = [ + "arrow-array", + "krishiv-types", +] [[package]] name = "krishiv-ui-assets" diff --git a/crates/krishiv-common/krishiv-error-classes.json b/crates/krishiv-common/krishiv-error-classes.json index 73af6c060..752ed3d88 100644 --- a/crates/krishiv-common/krishiv-error-classes.json +++ b/crates/krishiv-common/krishiv-error-classes.json @@ -107,6 +107,12 @@ ], "sqlState" : "42K0I" }, + "KRISHIV_UDF_NOT_AVAILABLE" : { + "message" : [ + "The user-defined function @ is not available on this worker: register it on the executing session before running the plan." + ], + "sqlState" : "KD002" + }, "KRISHIV_UNKNOWN_FORMAT_OPTION" : { "message" : [ "The option