You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@sedona.apache.org by pa...@apache.org on 2026/08/05 15:55:47 UTC
(sedona-db) branch main updated: feat(c/sedona-extension): Push Expr filters into the TableProvider across the FFI boundary (#1094)
This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new 786595e81 feat(c/sedona-extension): Push Expr filters into the TableProvider across the FFI boundary (#1094)
786595e81 is described below
commit 786595e81e871619c6c5715a47f7abfde6efee01
Author: Dewey Dunnington <de...@dunnington.ca>
AuthorDate: Wed Aug 5 11:55:41 2026 -0400
feat(c/sedona-extension): Push Expr filters into the TableProvider across the FFI boundary (#1094)
---
Cargo.lock | 2 +
Cargo.toml | 1 +
c/sedona-extension/Cargo.toml | 5 +
c/sedona-extension/src/execution_plan.rs | 80 +---
c/sedona-extension/src/expr.rs | 451 ++++++++++++++++++
c/sedona-extension/src/extension.rs | 56 +--
c/sedona-extension/src/lib.rs | 1 +
c/sedona-extension/src/sedona_extension.h | 44 +-
c/sedona-extension/src/table_provider.rs | 382 ++++++++++++++-
c/sedona-extension/src/utils.rs | 51 +-
python/sedonadb/tests/test_dataframe_ffi.py | 108 +++++
rust/sedona-expr/Cargo.toml | 1 +
rust/sedona-expr/src/lib.rs | 1 +
rust/sedona-expr/src/placeholder_udf.rs | 712 ++++++++++++++++++++++++++++
14 files changed, 1779 insertions(+), 116 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index f834d863c..8ac1cca16 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5871,6 +5871,7 @@ dependencies = [
"arrow-buffer",
"arrow-schema",
"datafusion-common",
+ "datafusion-execution",
"datafusion-expr",
"datafusion-physical-expr",
"geo-traits",
@@ -5897,6 +5898,7 @@ dependencies = [
"datafusion-expr",
"datafusion-physical-expr",
"datafusion-physical-plan",
+ "datafusion-proto",
"futures",
"libc",
"sedona-common",
diff --git a/Cargo.toml b/Cargo.toml
index 903078721..df32742d7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -90,6 +90,7 @@ datafusion = { version = "52.5.0", default-features = false }
datafusion-catalog = { version = "52.5.0" }
datafusion-common = { version = "52.5.0", default-features = false }
datafusion-common-runtime = { version = "52.5.0", default-features = false }
+datafusion-proto = { version = "52.5.0", default-features = false }
datafusion-datasource = { version = "52.5.0", default-features = false }
datafusion-datasource-parquet = { version = "52.5.0" }
datafusion-execution = { version = "52.5.0", default-features = false }
diff --git a/c/sedona-extension/Cargo.toml b/c/sedona-extension/Cargo.toml
index 4800913b3..a1a72e404 100644
--- a/c/sedona-extension/Cargo.toml
+++ b/c/sedona-extension/Cargo.toml
@@ -27,6 +27,10 @@ readme.workspace = true
edition.workspace = true
rust-version.workspace = true
+[features]
+default = ["protobuf"]
+protobuf = ["datafusion-proto"]
+
[dependencies]
arrow-array = { workspace = true, features = ["ffi"]}
arrow-schema = { workspace = true, features = ["ffi"]}
@@ -37,6 +41,7 @@ datafusion-execution = { workspace = true }
datafusion-expr = { workspace = true }
datafusion-physical-expr = { workspace = true }
datafusion-physical-plan = { workspace = true }
+datafusion-proto = { workspace = true, optional = true }
futures = { workspace = true }
libc = "0.2.178"
sedona-common = { workspace = true }
diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs
index d8f69d3db..30c5d7dc9 100644
--- a/c/sedona-extension/src/execution_plan.rs
+++ b/c/sedona-extension/src/execution_plan.rs
@@ -29,6 +29,7 @@ use arrow_schema::{ffi::FFI_ArrowSchema, Schema, SchemaRef};
use datafusion_common::{exec_err, Result, Statistics};
use datafusion_execution::TaskContext;
use datafusion_physical_plan::{
+ displayable,
execution_plan::{Boundedness, CardinalityEffect, EmissionType},
metrics::{CustomMetricValue, Metric, MetricValue, MetricsSet},
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
@@ -41,7 +42,9 @@ use crate::extension::{SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanA
use crate::runtime::RuntimeHandle;
use crate::set_ffi_error;
use crate::streaming::{ffi_stream_to_sendable, CancelChecker, StreamingRecordBatchReader};
-use crate::utils::{cstr_from_ptr_or_empty, get_plan_property, get_plan_string_property, ERRNO_OK};
+use crate::utils::{
+ cstr_from_ptr_or_empty, get_plan_property, get_plan_string_property, PropertyValue, ERRNO_OK,
+};
/// Wrapper around an [ExecutionPlan] that can be exported across FFI.
///
@@ -91,36 +94,9 @@ impl ExportedExecutionPlan {
})
}
"debug_string" => Ok(format!("{:?}", self.plan)),
- "display_default" => {
- use std::fmt::Write;
- let mut s = String::new();
- let _ = write!(
- s,
- "{}",
- DisplayAsWrapper(&self.plan, DisplayFormatType::Default)
- );
- Ok(s)
- }
- "display_verbose" => {
- use std::fmt::Write;
- let mut s = String::new();
- let _ = write!(
- s,
- "{}",
- DisplayAsWrapper(&self.plan, DisplayFormatType::Verbose)
- );
- Ok(s)
- }
- "display_tree_render" => {
- use std::fmt::Write;
- let mut s = String::new();
- let _ = write!(
- s,
- "{}",
- DisplayAsWrapper(&self.plan, DisplayFormatType::TreeRender)
- );
- Ok(s)
- }
+ "display_default" => Ok(displayable(self.plan.as_ref()).indent(false).to_string()),
+ "display_verbose" => Ok(displayable(self.plan.as_ref()).indent(true).to_string()),
+ "display_tree_render" => Ok(displayable(self.plan.as_ref()).tree_render().to_string()),
"name" => Ok(self.plan.name().to_string()),
"cardinality_effect" => {
let effect = match self.plan.cardinality_effect() {
@@ -241,12 +217,7 @@ unsafe extern "C" fn c_exec_plan_get_property(
match plan.get_property(&property_str) {
Ok(value) => {
- // Return the string as a single-element string array
- use arrow_array::{builder::StringBuilder, Array};
- let mut builder = StringBuilder::new();
- builder.append_value(&value);
- let array = builder.finish();
- let ffi_array = arrow_array::ffi::FFI_ArrowArray::new(&array.to_data());
+ let ffi_array = PropertyValue::String(value).into_ffi_array();
std::ptr::write(out, ffi_array);
ERRNO_OK
}
@@ -412,17 +383,17 @@ impl ImportedSedonaCExec {
impl DisplayAs for ImportedSedonaCExec {
fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
- let property = match t {
- DisplayFormatType::Default => "display_default",
- DisplayFormatType::Verbose => "display_verbose",
- DisplayFormatType::TreeRender => "display_tree_render",
+ let (property, sep) = match t {
+ DisplayFormatType::Default => ("display_default", ": "),
+ DisplayFormatType::Verbose => ("display_verbose", ": "),
+ DisplayFormatType::TreeRender => ("display_tree_render", "\n"),
};
// Always show the wrapper name, with the inner plan's display info
if let Ok(display_str) = get_plan_string_property(&self.inner, property) {
- write!(f, "ImportedSedonaCExec: {}", display_str)
+ write!(f, "ImportedSedonaCExec{sep}{display_str}")
} else {
- write!(f, "ImportedSedonaCExec")
+ write!(f, "ImportedSedonaCExec (error computing display)")
}
}
}
@@ -678,15 +649,6 @@ impl PlanPropertiesArgs {
}
}
-/// Helper wrapper to format an ExecutionPlan with a specific DisplayFormatType.
-struct DisplayAsWrapper<'a>(&'a Arc<dyn ExecutionPlan>, DisplayFormatType);
-
-impl std::fmt::Display for DisplayAsWrapper<'_> {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- self.0.fmt_as(self.1, f)
- }
-}
-
#[cfg(test)]
mod tests {
use super::*;
@@ -935,18 +897,26 @@ mod tests {
// ImportedSedonaCExec shows itself with the inner plan's display
assert_eq!(
format!("{}", DisplayAsFormat(&imported, DisplayFormatType::Default)),
- "ImportedSedonaCExec: DummyExec: default format"
+ "ImportedSedonaCExec: DummyExec: default format\n"
);
assert_eq!(
format!("{}", DisplayAsFormat(&imported, DisplayFormatType::Verbose)),
- "ImportedSedonaCExec: DummyExec: verbose format with schema"
+ "ImportedSedonaCExec: DummyExec: verbose format with schema\n"
);
assert_eq!(
format!(
"{}",
DisplayAsFormat(&imported, DisplayFormatType::TreeRender)
),
- "ImportedSedonaCExec: DummyExec: tree render format"
+ "\
+ImportedSedonaCExec
+┌───────────────────────────┐
+│ DummyExec │
+│ -------------------- │
+│ DummyExec: tree render │
+│ format │
+└───────────────────────────┘
+"
);
}
diff --git a/c/sedona-extension/src/expr.rs b/c/sedona-extension/src/expr.rs
new file mode 100644
index 000000000..58c5ad012
--- /dev/null
+++ b/c/sedona-extension/src/expr.rs
@@ -0,0 +1,451 @@
+// 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.
+
+use std::{
+ ffi::{c_int, CString},
+ fmt::{Debug, Display},
+ os::raw::{c_char, c_void},
+ ptr::null_mut,
+};
+
+use arrow_array::ffi::{FFI_ArrowArray, FFI_ArrowSchema};
+use arrow_schema::{DataType, Field};
+use datafusion_common::Result;
+use datafusion_execution::FunctionRegistry;
+use datafusion_expr::Expr;
+use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
+
+pub use sedona_expr::placeholder_udf::{
+ PlaceholderRegistry, PlaceholderUDAF, PlaceholderUDF, PlaceholderUDWF,
+};
+
+use crate::{
+ extension::{SedonaCError, SedonaCExprView},
+ set_ffi_error,
+ utils::{
+ call_get_property_schema_impl, cstr_from_ptr_or_empty, parse_ffi_array_to_bytes,
+ PropertyValue, ERRNO_OK,
+ },
+};
+
+/// Wrapper around a [datafusion_expr::Expr] that can be exported across FFI.
+pub struct ExportedExprView<'a> {
+ expr: &'a Expr,
+}
+
+impl<'a> Debug for ExportedExprView<'a> {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("ExportedExprView")
+ .field("expr", &self.expr)
+ .finish()
+ }
+}
+
+impl<'a> ExportedExprView<'a> {
+ /// Create a new ExportedExprView from a reference to a datafusion Expr.
+ ///
+ /// The returned view is valid only as long as the referenced Expr remains alive.
+ pub fn new(expr: &'a Expr) -> Self {
+ Self { expr }
+ }
+
+ /// Get the inner expression reference.
+ pub fn expr(&self) -> &Expr {
+ self.expr
+ }
+
+ /// Get a FFI-compatible view of this expression.
+ ///
+ /// The returned [SedonaCExprView] is valid only as long as this ExportedExprView
+ /// remains alive. The caller must ensure they do not use the FFI view after
+ /// this object is dropped.
+ pub fn as_ffi_view(&self) -> SedonaCExprView {
+ SedonaCExprView {
+ get_property_schema: Some(c_expr_get_property_schema),
+ get_property: Some(c_expr_get_property),
+ reserved: null_mut(),
+ private_data: self as *const ExportedExprView as *const c_void,
+ }
+ }
+
+ fn get_property(&self, property: &str) -> Result<PropertyValue, String> {
+ match property {
+ "debug_string" => Ok(PropertyValue::String(format!("{:?}", self.expr))),
+ "display_string" => Ok(PropertyValue::String(format!("{}", self.expr))),
+ #[cfg(feature = "protobuf")]
+ "datafusion_expr_protobuf" => {
+ use datafusion_proto::bytes::Serializeable;
+ self.expr
+ .to_bytes()
+ .map(|bytes| PropertyValue::Binary(bytes.to_vec()))
+ .map_err(|e| format!("Failed to serialize expression to protobuf: {}", e))
+ }
+ _ => Err(format!("Unknown property: {}", property)),
+ }
+ }
+
+ /// Returns the data type for a given property name.
+ fn get_property_data_type(property: &str) -> DataType {
+ match property {
+ "debug_string" | "display_string" => DataType::Utf8,
+ "datafusion_expr_protobuf" => DataType::Binary,
+ _ => DataType::Utf8, // Default to Utf8 for unknown properties
+ }
+ }
+}
+
+impl<'a> From<&'a Expr> for ExportedExprView<'a> {
+ fn from(expr: &'a Expr) -> Self {
+ Self::new(expr)
+ }
+}
+
+unsafe extern "C" fn c_expr_get_property_schema(
+ _self_: *const SedonaCExprView,
+ property: *const c_char,
+ out: *mut FFI_ArrowSchema,
+ err: *mut SedonaCError,
+) -> c_int {
+ debug_assert!(!out.is_null(), "out pointer is null");
+ let property_str = cstr_from_ptr_or_empty(property);
+ let data_type = ExportedExprView::get_property_data_type(&property_str);
+ let field = Field::new("value", data_type, false);
+ match FFI_ArrowSchema::try_from(&field) {
+ Ok(ffi_schema) => {
+ std::ptr::write(out, ffi_schema);
+ ERRNO_OK
+ }
+ Err(e) => {
+ set_ffi_error!(err, "Failed to convert field to FFI schema: {}", e);
+ libc::EINVAL
+ }
+ }
+}
+
+unsafe extern "C" fn c_expr_get_property(
+ self_: *const SedonaCExprView,
+ property: *const c_char,
+ _args: *const u8,
+ args_len: usize,
+ out: *mut FFI_ArrowArray,
+ err: *mut SedonaCError,
+) -> c_int {
+ debug_assert!(!self_.is_null(), "self pointer is null");
+ debug_assert!(!out.is_null(), "out pointer is null");
+
+ if args_len > 0 {
+ set_ffi_error!(err, "get_property does not accept arguments");
+ return libc::EINVAL;
+ }
+
+ let self_ref = &*self_;
+ debug_assert!(!self_ref.private_data.is_null(), "private_data is null");
+ let exported = &*(self_ref.private_data as *const ExportedExprView);
+ let property_str = cstr_from_ptr_or_empty(property);
+
+ match exported.get_property(&property_str) {
+ Ok(value) => {
+ let ffi_array = value.into_ffi_array();
+ std::ptr::write(out, ffi_array);
+ ERRNO_OK
+ }
+ Err(e) => {
+ set_ffi_error!(err, "{}", e);
+ libc::EINVAL
+ }
+ }
+}
+
+/// A borrowed expression view that can be used across an FFI boundary.
+///
+/// This wraps a reference to a [SedonaCExprView] and provides `Debug` and `Display`
+/// implementations by querying properties from the FFI interface.
+///
+/// The lifetime `'a` represents the lifetime of the underlying expression that
+/// this view references. The view must not be used after the underlying expression
+/// is dropped.
+pub struct ImportedExprView<'a> {
+ inner: &'a SedonaCExprView,
+}
+
+impl<'a> Debug for ImportedExprView<'a> {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ if let Ok(debug_str) = get_expr_view_string_property(self.inner, "debug_string") {
+ f.debug_struct("ImportedExprView")
+ .field("inner", &debug_str)
+ .finish()
+ } else {
+ f.debug_struct("ImportedExprView").finish()
+ }
+ }
+}
+
+impl<'a> Display for ImportedExprView<'a> {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ if let Ok(display_str) = get_expr_view_string_property(self.inner, "display_string") {
+ write!(f, "{}", display_str)
+ } else {
+ write!(f, "ImportedExprView")
+ }
+ }
+}
+
+impl<'a> ImportedExprView<'a> {
+ /// Create a new ImportedExprView from a SedonaCExprView reference.
+ ///
+ /// Returns an error if the SedonaCExprView does not have a valid private_data pointer.
+ pub fn try_new(inner: &'a SedonaCExprView) -> Result<Self> {
+ if inner.private_data.is_null() {
+ return sedona_internal_err!("SedonaCExprView has null private_data (invalid view)");
+ }
+ Ok(Self { inner })
+ }
+
+ /// Import this expression as a logical [Expr]
+ ///
+ /// This recreates a [Expr] object, using a registry to recreate function implementations.
+ /// If a registry is not provided, functions are replaced with [PlaceholderUDF] that may be
+ /// replaced (e.g., using an optimizer rule or locally implemented replacement), inspected
+ /// by parsing functions (e.g., to calculate pruning), or ignored (e.g., for table providers
+ /// that do not support any filters).
+ pub fn to_expr(&self, _registry: Option<&dyn FunctionRegistry>) -> Result<Expr> {
+ #[cfg(feature = "protobuf")]
+ {
+ use datafusion_proto::bytes::Serializeable;
+
+ let bytes = self.get_bytes_property("datafusion_expr_protobuf")?;
+ if let Some(registry) = _registry {
+ Expr::from_bytes_with_registry(&bytes, registry)
+ } else {
+ Expr::from_bytes_with_registry(&bytes, &PlaceholderRegistry)
+ }
+ }
+
+ #[cfg(not(feature = "protobuf"))]
+ {
+ sedona_internal_err!("sedona-expression not built with protobuf enabled")
+ }
+ }
+
+ /// Get a string property from this expression view.
+ pub fn get_string_property(&self, property: &str) -> Result<String> {
+ get_expr_view_string_property(self.inner, property)
+ }
+
+ /// Get a binary property from this expression view.
+ pub fn get_bytes_property(&self, property: &str) -> Result<Vec<u8>> {
+ get_expr_view_bytes_property(self.inner, property)
+ }
+}
+
+/// Get a string property from a [SedonaCExprView].
+pub fn get_expr_view_string_property(expr: &SedonaCExprView, property: &str) -> Result<String> {
+ let Some(get_property) = expr.get_property else {
+ return sedona_internal_err!("SedonaCExprView does not have get_property");
+ };
+
+ let property_cstr = CString::new(property)
+ .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?;
+
+ let mut ffi_array = FFI_ArrowArray::empty();
+ let mut err = SedonaCError::default();
+
+ let code = unsafe {
+ get_property(
+ expr,
+ property_cstr.as_ptr(),
+ std::ptr::null(),
+ 0,
+ &mut ffi_array,
+ &mut err,
+ )
+ };
+
+ if code != ERRNO_OK {
+ return sedona_internal_err!("SedonaCExprView failed to get '{}': {}", property, err);
+ }
+
+ let data_type = get_expr_view_property_data_type(expr, property)?;
+ let bytes = parse_ffi_array_to_bytes(ffi_array, &data_type)?;
+ String::from_utf8(bytes)
+ .map_err(|e| sedona_internal_datafusion_err!("Invalid UTF-8 in '{}': {}", property, e))
+}
+
+/// Get a binary property from a [SedonaCExprView].
+pub fn get_expr_view_bytes_property(expr: &SedonaCExprView, property: &str) -> Result<Vec<u8>> {
+ let Some(get_property) = expr.get_property else {
+ return sedona_internal_err!("SedonaCExprView does not have get_property");
+ };
+
+ let property_cstr = CString::new(property)
+ .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?;
+
+ let mut ffi_array = FFI_ArrowArray::empty();
+ let mut err = SedonaCError::default();
+
+ let code = unsafe {
+ get_property(
+ expr,
+ property_cstr.as_ptr(),
+ std::ptr::null(),
+ 0,
+ &mut ffi_array,
+ &mut err,
+ )
+ };
+
+ if code != ERRNO_OK {
+ return sedona_internal_err!("SedonaCExprView failed to get '{}': {}", property, err);
+ }
+
+ let data_type = get_expr_view_property_data_type(expr, property)?;
+ parse_ffi_array_to_bytes(ffi_array, &data_type)
+}
+
+/// Get the data type for a property from a [SedonaCExprView].
+fn get_expr_view_property_data_type(expr: &SedonaCExprView, property: &str) -> Result<DataType> {
+ let Some(get_property_schema) = expr.get_property_schema else {
+ return Ok(DataType::Utf8);
+ };
+
+ call_get_property_schema_impl(property, |prop, schema, err| unsafe {
+ get_property_schema(expr, prop, schema, err)
+ })
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use datafusion_expr::col;
+
+ #[test]
+ fn test_expr_view_debug_string() {
+ let expr = col("test_column");
+ let exported = ExportedExprView::new(&expr);
+ let view = exported.as_ffi_view();
+ let imported = ImportedExprView::try_new(&view).unwrap();
+
+ let debug_str = imported.get_string_property("debug_string").unwrap();
+ assert_eq!(debug_str, format!("{:?}", expr));
+ }
+
+ #[test]
+ fn test_expr_view_display_string() {
+ let expr = col("test_column");
+ let exported = ExportedExprView::new(&expr);
+ let view = exported.as_ffi_view();
+ let imported = ImportedExprView::try_new(&view).unwrap();
+
+ let display_str = imported.get_string_property("display_string").unwrap();
+ assert_eq!(display_str, format!("{}", expr));
+ }
+
+ #[test]
+ fn test_expr_view_debug_impl() {
+ let expr = col("my_col");
+ let exported = ExportedExprView::new(&expr);
+ let view = exported.as_ffi_view();
+ let imported = ImportedExprView::try_new(&view).unwrap();
+
+ let debug_output = format!("{:?}", imported);
+ assert!(debug_output.contains("ImportedExprView"));
+ assert!(debug_output.contains("my_col"));
+ }
+
+ #[test]
+ fn test_expr_view_display_impl() {
+ let expr = col("my_col");
+ let exported = ExportedExprView::new(&expr);
+ let view = exported.as_ffi_view();
+ let imported = ImportedExprView::try_new(&view).unwrap();
+
+ let display_output = format!("{}", imported);
+ assert_eq!(display_output, "my_col");
+ }
+
+ #[test]
+ fn test_expr_view_unknown_property_error() {
+ let expr = col("test");
+ let exported = ExportedExprView::new(&expr);
+ let view = exported.as_ffi_view();
+ let imported = ImportedExprView::try_new(&view).unwrap();
+
+ let result = imported.get_string_property("nonexistent_property");
+ assert!(result.is_err());
+ let err_msg = result.unwrap_err().to_string();
+ assert!(err_msg.contains("Unknown property"));
+ }
+
+ #[test]
+ fn test_expr_view_with_null_private_data_fails() {
+ let invalid_view = SedonaCExprView::default();
+ let result = ImportedExprView::try_new(&invalid_view);
+ assert!(result.is_err());
+ assert!(result
+ .unwrap_err()
+ .to_string()
+ .contains("null private_data"));
+ }
+
+ #[cfg(feature = "protobuf")]
+ #[test]
+ fn test_expr_view_protobuf_roundtrip() {
+ use datafusion_expr::lit;
+
+ // Create a simple expression: col("x") > 5
+ let expr = col("x").gt(lit(5i32));
+ let exported = ExportedExprView::new(&expr);
+ let view = exported.as_ffi_view();
+ let imported = ImportedExprView::try_new(&view).unwrap();
+
+ // Verify we can deserialize it back to an equivalent expression
+ let decoded = imported.to_expr(None).unwrap();
+ assert_eq!(format!("{}", expr), format!("{}", decoded));
+ }
+
+ #[cfg(feature = "protobuf")]
+ #[test]
+ fn test_expr_view_function_becomes_placeholder_udf() {
+ use datafusion_expr::{lit, ScalarUDF};
+
+ // Create a function that won't exist in PlaceholderRegistry
+ let test_udf = ScalarUDF::new_from_impl(PlaceholderUDF::new("my_custom_func"));
+ let expr = test_udf.call(vec![col("x"), lit(42i32)]);
+
+ let exported = ExportedExprView::new(&expr);
+ let view = exported.as_ffi_view();
+ let imported = ImportedExprView::try_new(&view).unwrap();
+
+ // Deserialize - the function should come back as a PlaceholderUDF
+ let decoded = imported.to_expr(None).unwrap();
+
+ // Verify it's a scalar function with the right name
+ match &decoded {
+ Expr::ScalarFunction(func) => {
+ assert_eq!(func.name(), "my_custom_func");
+ // Verify it's actually a PlaceholderUDF
+ let udf_impl = func.func.inner();
+ assert!(
+ udf_impl.as_any().downcast_ref::<PlaceholderUDF>().is_some(),
+ "Expected PlaceholderUDF, got {:?}",
+ udf_impl
+ );
+ }
+ other => panic!("Expected ScalarFunction, got {:?}", other),
+ }
+ }
+}
diff --git a/c/sedona-extension/src/extension.rs b/c/sedona-extension/src/extension.rs
index 287023f32..15a68fe95 100644
--- a/c/sedona-extension/src/extension.rs
+++ b/c/sedona-extension/src/extension.rs
@@ -220,12 +220,6 @@ pub unsafe fn write_ffi_error(err: *mut SedonaCError, message: &str) {
///
/// This is a convenience wrapper around `write_ffi_error` that supports
/// format strings like `format!()`.
-///
-/// # Example
-///
-/// ```ignore
-/// set_ffi_error!(err, "Failed to parse: {}", e);
-/// ```
#[macro_export]
macro_rules! set_ffi_error {
($err:expr, $msg:expr) => {
@@ -236,13 +230,28 @@ macro_rules! set_ffi_error {
};
}
-/// Raw FFI representation of the SedonaCExpr
+/// Raw FFI representation of a non-owning expression view
+///
+/// This structure provides a read-only view into an expression without taking
+/// ownership. The lifetime of this view is tied to the underlying expression
+/// it references.
+///
+/// Lifetime semantics:
+/// - Exporter: The exporter must ensure the underlying expression remains valid
+/// for the entire lifetime of this view. The exporter is responsible for
+/// managing the lifetime of the underlying expression.
+/// - Importer: The importer must not use this view after the underlying
+/// expression has been released. The importer should not store this view
+/// beyond the scope in which it was provided.
+///
+/// Before using this structure, private_data MUST be checked.
+/// Instances with a NULL private_data are not valid and must not be used.
#[derive(Default)]
#[repr(C)]
-pub struct SedonaCExpr {
+pub struct SedonaCExprView {
pub get_property_schema: Option<
unsafe extern "C" fn(
- self_: *const SedonaCExpr,
+ self_: *const SedonaCExprView,
property: *const c_char,
out: *mut FFI_ArrowSchema,
err: *mut SedonaCError,
@@ -251,9 +260,10 @@ pub struct SedonaCExpr {
pub get_property: Option<
unsafe extern "C" fn(
- self_: *const SedonaCExpr,
+ self_: *const SedonaCExprView,
property: *const c_char,
- args: *const c_char,
+ args: *const u8,
+ args_len: usize,
out: *mut FFI_ArrowArray,
err: *mut SedonaCError,
) -> c_int,
@@ -261,23 +271,11 @@ pub struct SedonaCExpr {
pub reserved: *mut c_void,
- pub release: Option<unsafe extern "C" fn(self_: *mut SedonaCExpr)>,
-
- pub private_data: *mut c_void,
+ pub private_data: *const c_void,
}
-unsafe impl Send for SedonaCExpr {}
-unsafe impl Sync for SedonaCExpr {}
-
-impl Drop for SedonaCExpr {
- fn drop(&mut self) {
- if let Some(releaser) = self.release {
- unsafe { releaser(self) }
- self.release = None;
- self.private_data = null_mut();
- }
- }
-}
+unsafe impl Send for SedonaCExprView {}
+unsafe impl Sync for SedonaCExprView {}
/// Raw FFI representation of the SedonaCExecutionPlanArgs
///
@@ -292,8 +290,10 @@ pub struct SedonaCExecutionPlanArgs {
/// Optional array of execution plans
pub exec_plans: *const *const SedonaCExecutionPlan,
pub num_exec_plans: usize,
- /// Optional array of expressions
- pub exprs: *const *const SedonaCExpr,
+ /// Optional array of expression views
+ ///
+ /// These views are valid only for the duration of the callback invocation.
+ pub exprs: *const *const SedonaCExprView,
pub num_exprs: usize,
pub reserved: *mut c_void,
}
diff --git a/c/sedona-extension/src/lib.rs b/c/sedona-extension/src/lib.rs
index 15dbbfb04..8adf9f55e 100644
--- a/c/sedona-extension/src/lib.rs
+++ b/c/sedona-extension/src/lib.rs
@@ -16,6 +16,7 @@
// under the License.
pub mod execution_plan;
+pub mod expr;
pub mod extension;
pub mod runtime;
pub mod scalar_kernel;
diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h
index 128b231e5..b2ed9a6b1 100644
--- a/c/sedona-extension/src/sedona_extension.h
+++ b/c/sedona-extension/src/sedona_extension.h
@@ -229,9 +229,30 @@ struct SedonaCError {
void (*release)(struct SedonaCError* self);
};
-struct SedonaCExpr {
+/// \brief Non-owning view of an expression
+///
+/// This structure provides a read-only view into an expression without taking
+/// ownership. The lifetime of this view is tied to the underlying expression
+/// it references. The SedonaCExprView may be logical or physical...it is not
+/// typically executed but inspected for the purposes of pruning. The primary
+/// mechanism of recreating an expression is DataFusion Protobuf via the
+/// "datafusion_expr_protobuf" property; however, other types of queries may be
+/// supported by other properties in the future (e.g., rendering SQL or
+/// Substrait or extracting a query bounding box directly).
+///
+/// Lifetime semantics:
+/// - Exporter: The exporter must ensure the underlying expression remains valid
+/// for the entire lifetime of this view. The exporter is responsible for
+/// managing the lifetime of the underlying expression.
+/// - Importer: The importer must not use this view after the underlying
+/// expression has been released. The importer should not store this view
+/// beyond the scope in which it was provided.
+///
+/// Before using this structure, private_data MUST be checked.
+/// Instances with a NULL private_data are not valid and must not be used.
+struct SedonaCExprView {
/// \brief Get the data type of a property
- int (*get_property_schema)(const struct SedonaCExpr* self, const char* property,
+ int (*get_property_schema)(const struct SedonaCExprView* self, const char* property,
struct ArrowSchema* out, struct SedonaCError* err);
/// \brief Extract a serializable property from this expression
@@ -239,17 +260,16 @@ struct SedonaCExpr {
/// This is used to implement PlanProperties and other values that can be
/// easily retrieved and serialized. The data type associated with the out
/// array may be retrieved with the get_property_schema callback.
- int (*get_property)(const struct SedonaCExpr* self, const char* property,
- const char* args, struct ArrowArray* out, struct SedonaCError* err);
+ int (*get_property)(const struct SedonaCExprView* self, const char* property,
+ const uint8_t* args, size_t args_len, struct ArrowArray* out,
+ struct SedonaCError* err);
+ /// \brief Reserved for future use. Must be NULL.
void* reserved;
- /// \brief Release this instance
- ///
- /// Implementations of this callback must set self->release to NULL.
- void (*release)(struct SedonaCExpr* self);
-
/// \brief Opaque implementation-specific data
+ ///
+ /// A NULL value indicates this view is invalid and must not be used.
void* private_data;
};
@@ -269,8 +289,10 @@ struct SedonaCExecutionPlanArgs {
const struct SedonaCExecutionPlan** exec_plans;
size_t num_exec_plans;
- /// \brief Optional array of expressions
- const struct SedonaCExpr** exprs;
+ /// \brief Optional array of expression views
+ ///
+ /// These views are valid only for the duration of the callback invocation.
+ const struct SedonaCExprView** exprs;
size_t num_exprs;
/// \brief Reserved for future use. Must be NULL.
diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs
index aa1e86476..b1a711f9c 100644
--- a/c/sedona-extension/src/table_provider.rs
+++ b/c/sedona-extension/src/table_provider.rs
@@ -17,6 +17,7 @@
use std::{
any::Any,
+ collections::HashSet,
ffi::{c_int, c_void},
fmt::Debug,
ptr::null_mut,
@@ -25,22 +26,30 @@ use std::{
};
use arrow_array::ffi::FFI_ArrowArray;
-use arrow_schema::{ffi::FFI_ArrowSchema, Schema, SchemaRef};
+use arrow_schema::{ffi::FFI_ArrowSchema, DataType, Schema, SchemaRef};
use async_trait::async_trait;
use datafusion_catalog::{Session, TableProvider};
-use datafusion_common::{exec_err, Result, Statistics};
-use datafusion_expr::{Expr, TableType};
+use datafusion_common::{exec_err, plan_err, Result, Statistics};
+use datafusion_execution::FunctionRegistry;
+use datafusion_expr::{Expr, ScalarUDF, TableProviderFilterPushDown, TableType};
use datafusion_physical_plan::ExecutionPlan;
use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
use serde::{Deserialize, Serialize};
-use crate::execution_plan::{ExportedExecutionPlan, ImportedSedonaCExec};
+use crate::expr::{ExportedExprView, ImportedExprView};
use crate::extension::{
- SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs, SedonaCTableProvider,
+ SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs, SedonaCExprView,
+ SedonaCTableProvider,
};
use crate::runtime::RuntimeHandle;
use crate::set_ffi_error;
-use crate::utils::{cstr_from_ptr_or_empty, get_table_provider_string_property, ERRNO_OK};
+use crate::utils::{
+ cstr_from_ptr_or_empty, get_table_provider_string_property, PropertyValue, ERRNO_OK,
+};
+use crate::{
+ execution_plan::{ExportedExecutionPlan, ImportedSedonaCExec},
+ utils::parse_ffi_array_to_bytes,
+};
/// A TableProvider wrapper that can be exported across FFI.
///
@@ -85,6 +94,7 @@ impl ExportedTableProvider {
fn scan(
&self,
projection: Option<Vec<usize>>,
+ filters: Vec<Expr>,
limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
let inner = self.inner.clone();
@@ -95,7 +105,7 @@ impl ExportedTableProvider {
let projection_ref = projection.as_ref();
runtime
.handle()
- .block_on(inner.scan(session.as_ref(), projection_ref, &[], limit))
+ .block_on(inner.scan(session.as_ref(), projection_ref, &filters, limit))
})
.join()
.map_err(|e| sedona_internal_datafusion_err!("Scan thread panicked {e:?}"))?
@@ -115,6 +125,22 @@ impl ExportedTableProvider {
_ => exec_err!("Unknown property: {}", property),
}
}
+
+ fn supports_filters_pushdown(&self, filters: &[Expr]) -> Result<String> {
+ let filter_refs: Vec<&Expr> = filters.iter().collect();
+ let pushdown_results = self.inner.supports_filters_pushdown(&filter_refs)?;
+ let result_strs: Vec<&str> = pushdown_results
+ .iter()
+ .map(|p| match p {
+ TableProviderFilterPushDown::Unsupported => "Unsupported",
+ TableProviderFilterPushDown::Inexact => "Inexact",
+ TableProviderFilterPushDown::Exact => "Exact",
+ })
+ .collect();
+ serde_json::to_string(&result_strs).map_err(|e| {
+ sedona_internal_datafusion_err!("Failed to serialize pushdown results: {}", e)
+ })
+ }
}
impl From<ExportedTableProvider> for SedonaCTableProvider {
@@ -184,7 +210,7 @@ unsafe extern "C" fn c_table_provider_get_property_schema(
unsafe extern "C" fn c_table_provider_get_property(
self_: *const SedonaCTableProvider,
property: *const std::ffi::c_char,
- _args: *mut SedonaCExecutionPlanArgs,
+ args: *mut SedonaCExecutionPlanArgs,
out: *mut FFI_ArrowArray,
err: *mut SedonaCError,
) -> c_int {
@@ -195,14 +221,84 @@ unsafe extern "C" fn c_table_provider_get_property(
let provider = &*(self_ref.private_data as *const ExportedTableProvider);
let property_str = cstr_from_ptr_or_empty(property);
- match provider.get_property(&property_str) {
+ // Handle property requests that require expression arguments
+ let result = if property_str == "supports_filters_pushdown" {
+ // Extract filters from expression views, tracking which ones can be deserialized
+ if args.is_null() {
+ provider.supports_filters_pushdown(&[])
+ } else {
+ let args_ref = &*args;
+ if args_ref.exprs.is_null() || args_ref.num_exprs == 0 {
+ provider.supports_filters_pushdown(&[])
+ } else {
+ let expr_ptrs = std::slice::from_raw_parts(args_ref.exprs, args_ref.num_exprs);
+ let registry = SessionRefRegistry::new(provider.session.as_ref());
+
+ // Try to deserialize each filter; track successes and failures
+ let mut deserialized: Vec<Option<Expr>> = Vec::with_capacity(args_ref.num_exprs);
+ for &expr_ptr in expr_ptrs {
+ if expr_ptr.is_null() {
+ deserialized.push(None);
+ continue;
+ }
+ let expr_view = match ImportedExprView::try_new(&*expr_ptr) {
+ Ok(v) => v,
+ Err(_) => {
+ // Can't even create the view - mark as unsupported
+ deserialized.push(None);
+ continue;
+ }
+ };
+
+ // If deserialization fails (e.g., UDF not available on this side),
+ // mark as None (will become Unsupported)
+ match expr_view.to_expr(Some(®istry)) {
+ Ok(expr) => deserialized.push(Some(expr)),
+ Err(_) => deserialized.push(None),
+ }
+ }
+
+ // Collect only the successfully deserialized filters
+ let valid_filters: Vec<&Expr> =
+ deserialized.iter().filter_map(|opt| opt.as_ref()).collect();
+
+ // Get pushdown results for valid filters from the inner provider
+ let inner_results = provider
+ .inner
+ .supports_filters_pushdown(&valid_filters)
+ .unwrap_or_else(|_| {
+ vec![TableProviderFilterPushDown::Unsupported; valid_filters.len()]
+ });
+
+ // Merge results: use inner results for valid filters, Unsupported for failed ones
+ let mut inner_iter = inner_results.into_iter();
+ let final_results: Vec<&str> = deserialized
+ .iter()
+ .map(|opt| {
+ if opt.is_some() {
+ match inner_iter.next() {
+ Some(TableProviderFilterPushDown::Exact) => "Exact",
+ Some(TableProviderFilterPushDown::Inexact) => "Inexact",
+ _ => "Unsupported",
+ }
+ } else {
+ "Unsupported"
+ }
+ })
+ .collect();
+
+ serde_json::to_string(&final_results).map_err(|e| {
+ sedona_internal_datafusion_err!("Failed to serialize pushdown results: {}", e)
+ })
+ }
+ }
+ } else {
+ provider.get_property(&property_str)
+ };
+
+ match result {
Ok(value) => {
- // Return the string as a single-element string array
- use arrow_array::{builder::StringBuilder, Array};
- let mut builder = StringBuilder::new();
- builder.append_value(&value);
- let array = builder.finish();
- let ffi_array = FFI_ArrowArray::new(&array.to_data());
+ let ffi_array = PropertyValue::String(value).into_ffi_array();
std::ptr::write(out, ffi_array);
ERRNO_OK
}
@@ -249,7 +345,37 @@ unsafe extern "C" fn c_table_provider_scan(
}
};
- match provider.scan(scan_args.projection, scan_args.limit) {
+ // Extract filters from expression views
+ let registry = SessionRefRegistry::new(provider.session.as_ref());
+ let filters: Vec<Expr> = if args_ref.exprs.is_null() || args_ref.num_exprs == 0 {
+ Vec::new()
+ } else {
+ let expr_ptrs = std::slice::from_raw_parts(args_ref.exprs, args_ref.num_exprs);
+ let mut filters = Vec::with_capacity(args_ref.num_exprs);
+ for &expr_ptr in expr_ptrs {
+ if expr_ptr.is_null() {
+ continue;
+ }
+ let expr_view = match ImportedExprView::try_new(&*expr_ptr) {
+ Ok(v) => v,
+ Err(e) => {
+ set_ffi_error!(err, "Failed to import expression view: {}", e);
+ return libc::EINVAL;
+ }
+ };
+
+ match expr_view.to_expr(Some(®istry)) {
+ Ok(expr) => filters.push(expr),
+ Err(e) => {
+ set_ffi_error!(err, "Failed to deserialize filter expression: {}", e);
+ return libc::EINVAL;
+ }
+ }
+ }
+ filters
+ };
+
+ match provider.scan(scan_args.projection, filters, scan_args.limit) {
Ok(plan) => {
let task_ctx = provider.session.task_ctx();
let exported = ExportedExecutionPlan::new(plan, task_ctx, provider.runtime.clone());
@@ -383,11 +509,99 @@ impl TableProvider for ImportedTableProvider {
None
}
+ fn supports_filters_pushdown(
+ &self,
+ filters: &[&Expr],
+ ) -> Result<Vec<TableProviderFilterPushDown>> {
+ let Some(get_property) = self.inner.get_property else {
+ // Default to Unsupported if get_property is not available
+ return Ok(vec![
+ TableProviderFilterPushDown::Unsupported;
+ filters.len()
+ ]);
+ };
+
+ // Create FFI expression views for filters
+ let filter_views: Vec<ExportedExprView> =
+ filters.iter().map(|&e| ExportedExprView::new(e)).collect();
+ let ffi_filter_views: Vec<SedonaCExprView> =
+ filter_views.iter().map(|v| v.as_ffi_view()).collect();
+ let ffi_filter_ptrs: Vec<*const SedonaCExprView> =
+ ffi_filter_views.iter().map(|v| v as *const _).collect();
+
+ let mut ffi_args = SedonaCExecutionPlanArgs {
+ args: std::ptr::null(),
+ args_len: 0,
+ exec_plans: std::ptr::null(),
+ num_exec_plans: 0,
+ exprs: if ffi_filter_ptrs.is_empty() {
+ std::ptr::null()
+ } else {
+ ffi_filter_ptrs.as_ptr()
+ },
+ num_exprs: ffi_filter_ptrs.len(),
+ reserved: null_mut(),
+ };
+
+ let property_cstr = std::ffi::CString::new("supports_filters_pushdown")
+ .map_err(|e| sedona_internal_datafusion_err!("Invalid property name: {}", e))?;
+
+ let mut ffi_array = FFI_ArrowArray::empty();
+ let mut err = SedonaCError::default();
+
+ let code = unsafe {
+ get_property(
+ &self.inner,
+ property_cstr.as_ptr(),
+ &mut ffi_args,
+ &mut ffi_array,
+ &mut err,
+ )
+ };
+
+ if code != ERRNO_OK {
+ // If get_property fails, default to Unsupported
+ return Ok(vec![
+ TableProviderFilterPushDown::Unsupported;
+ filters.len()
+ ]);
+ }
+
+ // Parse the result as a JSON array of strings
+ let bytes = parse_ffi_array_to_bytes(ffi_array, &DataType::Utf8)?;
+ let value = String::from_utf8(bytes).map_err(|e| {
+ sedona_internal_datafusion_err!("Invalid UTF-8 in property value: {}", e)
+ })?;
+ let pushdown_strs: Vec<String> = serde_json::from_str(&value).map_err(|e| {
+ sedona_internal_datafusion_err!("Failed to parse pushdown results: {}", e)
+ })?;
+
+ // Check the number of results
+ if pushdown_strs.len() != filters.len() {
+ return sedona_internal_err!(
+ "Expected {} results for supports_filters_pushdown property but got {}",
+ filters.len(),
+ pushdown_strs.len()
+ );
+ }
+
+ let results = pushdown_strs
+ .iter()
+ .map(|s| match s.as_str() {
+ "Exact" => TableProviderFilterPushDown::Exact,
+ "Inexact" => TableProviderFilterPushDown::Inexact,
+ _ => TableProviderFilterPushDown::Unsupported,
+ })
+ .collect();
+
+ Ok(results)
+ }
+
async fn scan(
&self,
_state: &dyn Session,
projection: Option<&Vec<usize>>,
- _filters: &[Expr],
+ filters: &[Expr],
limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
let Some(scan) = self.inner.scan else {
@@ -401,13 +615,25 @@ impl TableProvider for ImportedTableProvider {
let args_bytes = serde_json::to_vec(&args)
.map_err(|e| sedona_internal_datafusion_err!("Failed to serialize scan args: {}", e))?;
+ // Create FFI expression views for filters
+ let filter_views: Vec<ExportedExprView> =
+ filters.iter().map(ExportedExprView::new).collect();
+ let ffi_filter_views: Vec<SedonaCExprView> =
+ filter_views.iter().map(|v| v.as_ffi_view()).collect();
+ let ffi_filter_ptrs: Vec<*const SedonaCExprView> =
+ ffi_filter_views.iter().map(|v| v as *const _).collect();
+
let mut ffi_args = SedonaCExecutionPlanArgs {
args: args_bytes.as_ptr(),
args_len: args_bytes.len(),
exec_plans: std::ptr::null(),
num_exec_plans: 0,
- exprs: std::ptr::null(),
- num_exprs: 0,
+ exprs: if ffi_filter_ptrs.is_empty() {
+ std::ptr::null()
+ } else {
+ ffi_filter_ptrs.as_ptr()
+ },
+ num_exprs: ffi_filter_ptrs.len(),
reserved: null_mut(),
};
@@ -444,6 +670,62 @@ struct ScanArgs {
pub limit: Option<usize>,
}
+/// [FunctionRegistry] implemented on a dyn [Session]
+///
+/// Many functions pass a reference to the dyn [Session]. This implementation allows
+/// those functions to deserialize a protobuf expression across the FFI boundary.
+struct SessionRefRegistry<'a> {
+ session: &'a dyn Session,
+}
+
+impl<'a> SessionRefRegistry<'a> {
+ pub fn new(session: &'a dyn Session) -> Self {
+ SessionRefRegistry { session }
+ }
+}
+
+impl<'a> FunctionRegistry for SessionRefRegistry<'a> {
+ fn udfs(&self) -> HashSet<String> {
+ self.session.scalar_functions().keys().cloned().collect()
+ }
+
+ fn udafs(&self) -> HashSet<String> {
+ self.session.aggregate_functions().keys().cloned().collect()
+ }
+
+ fn udwfs(&self) -> HashSet<String> {
+ self.session.window_functions().keys().cloned().collect()
+ }
+
+ fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>> {
+ if let Some(func) = self.session.scalar_functions().get(name) {
+ Ok(Arc::clone(func))
+ } else {
+ plan_err!("Can't find scalar function '{name}' in session")
+ }
+ }
+
+ fn udaf(&self, name: &str) -> Result<Arc<datafusion_expr::AggregateUDF>> {
+ if let Some(func) = self.session.aggregate_functions().get(name) {
+ Ok(Arc::clone(func))
+ } else {
+ plan_err!("Can't find aggregate function '{name}' in session")
+ }
+ }
+
+ fn udwf(&self, name: &str) -> Result<Arc<datafusion_expr::WindowUDF>> {
+ if let Some(func) = self.session.window_functions().get(name) {
+ Ok(Arc::clone(func))
+ } else {
+ plan_err!("Can't find window function '{name}' in session")
+ }
+ }
+
+ fn expr_planners(&self) -> Vec<Arc<dyn datafusion_expr::planner::ExprPlanner>> {
+ vec![]
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -611,6 +893,59 @@ mod tests {
.unwrap();
}
+ #[test]
+ fn test_filter_pushdown_in_explain() {
+ let runtime = test_runtime();
+ runtime
+ .block_on(async {
+ let ctx = create_test_context().await?;
+
+ // Get the table provider from the context. Use into_view()
+ // to use the DataFrame implementation of a table, which does
+ // support filter pushdown
+ let table = ctx.table("test_data").await?.into_view();
+
+
+ // Export the table provider
+ let session = Arc::new(ctx.state());
+ let exported = ExportedTableProvider::new(table, session, runtime.clone());
+ let ffi_provider: SedonaCTableProvider = exported.into();
+
+ // Import the table provider
+ let imported = ImportedTableProvider::try_new(ffi_provider)?;
+
+ // Create a new context and register the imported table
+ let ctx2 = SessionContext::new();
+ ctx2.register_table("imported_data", Arc::new(imported))?;
+
+ // Get the explain plan
+ let explain_result = ctx2
+ .sql("EXPLAIN SELECT id FROM imported_data WHERE id > 20")
+ .await?
+ .collect()
+ .await?;
+
+ assert_batches_eq!(
+ &[
+ "+---------------+---------------------------------------------------------------------------------------+",
+ "| plan_type | plan |",
+ "+---------------+---------------------------------------------------------------------------------------+",
+ "| logical_plan | TableScan: imported_data projection=[id], full_filters=[imported_data.id > Int32(20)] |",
+ "| physical_plan | CooperativeExec |",
+ "| | ImportedSedonaCExec: FilterExec: id@0 > 20 |",
+ "| | DataSourceExec: partitions=1, partition_sizes=[5] |",
+ "| | |",
+ "| | |",
+ "+---------------+---------------------------------------------------------------------------------------+",
+ ],
+ &explain_result
+ );
+
+ Ok::<(), datafusion_common::DataFusionError>(())
+ })
+ .unwrap();
+ }
+
#[test]
fn test_roundtrip_sort() {
test_roundtrip_query(
@@ -631,11 +966,18 @@ mod tests {
}
#[test]
+ #[rustfmt::skip]
fn test_roundtrip_limit() {
test_roundtrip_query(
"SELECT id FROM imported_data ORDER BY id LIMIT 3",
&[
- "+----+", "| id |", "+----+", "| 1 |", "| 2 |", "| 3 |", "+----+",
+ "+----+",
+ "| id |",
+ "+----+",
+ "| 1 |",
+ "| 2 |",
+ "| 3 |",
+ "+----+",
],
)
.unwrap();
diff --git a/c/sedona-extension/src/utils.rs b/c/sedona-extension/src/utils.rs
index ae42851bb..56c9b8d57 100644
--- a/c/sedona-extension/src/utils.rs
+++ b/c/sedona-extension/src/utils.rs
@@ -21,6 +21,9 @@ use std::borrow::Cow;
use std::ffi::{c_char, c_int, CStr, CString};
use std::ptr::null_mut;
+use arrow_array::builder::{BinaryBuilder, StringBuilder};
+use arrow_array::ffi::FFI_ArrowArray;
+use arrow_array::Array;
use arrow_schema::ffi::FFI_ArrowSchema;
use arrow_schema::{DataType, Field};
use datafusion_common::Result;
@@ -35,6 +38,50 @@ use crate::extension::{
/// Success return code for FFI functions.
pub const ERRNO_OK: c_int = 0;
+/// The value returned by a property getter.
+///
+/// This enum represents the two primary data types returned by FFI property accessors:
+/// UTF-8 strings and binary data. It provides a type-safe way to handle property values
+/// and can be converted to Arrow arrays for FFI transport.
+#[derive(Debug, Clone)]
+pub enum PropertyValue {
+ /// A UTF-8 string value.
+ String(String),
+ /// A binary (bytes) value.
+ Binary(Vec<u8>),
+}
+
+impl PropertyValue {
+ /// Returns the Arrow [DataType] for this property value.
+ pub fn data_type(&self) -> DataType {
+ match self {
+ PropertyValue::String(_) => DataType::Utf8,
+ PropertyValue::Binary(_) => DataType::Binary,
+ }
+ }
+
+ /// Converts this property value into an FFI-compatible Arrow array.
+ ///
+ /// This is useful for FFI callbacks that need to return property values
+ /// as Arrow arrays across the FFI boundary.
+ pub fn into_ffi_array(self) -> FFI_ArrowArray {
+ match self {
+ PropertyValue::String(value) => {
+ let mut builder = StringBuilder::new();
+ builder.append_value(&value);
+ let array = builder.finish();
+ FFI_ArrowArray::new(&array.to_data())
+ }
+ PropertyValue::Binary(value) => {
+ let mut builder = BinaryBuilder::new();
+ builder.append_value(&value);
+ let array = builder.finish();
+ FFI_ArrowArray::new(&array.to_data())
+ }
+ }
+ }
+}
+
/// Safely convert a C string pointer to a Rust string, treating null as empty.
///
/// # Safety
@@ -152,7 +199,7 @@ where
///
/// The caller provides a closure that performs the actual FFI call with
/// the correct self pointer type.
-fn call_get_property_schema_impl<F>(property: &str, call_ffi: F) -> Result<DataType>
+pub fn call_get_property_schema_impl<F>(property: &str, call_ffi: F) -> Result<DataType>
where
F: FnOnce(*const c_char, *mut FFI_ArrowSchema, *mut SedonaCError) -> c_int,
{
@@ -251,7 +298,7 @@ fn parse_ffi_array<T: DeserializeOwned>(
}
/// Parse an FFI array and return the raw bytes.
-fn parse_ffi_array_to_bytes(
+pub fn parse_ffi_array_to_bytes(
ffi_array: arrow_array::ffi::FFI_ArrowArray,
data_type: &DataType,
) -> Result<Vec<u8>> {
diff --git a/python/sedonadb/tests/test_dataframe_ffi.py b/python/sedonadb/tests/test_dataframe_ffi.py
index 7b230a83a..c08aebdc3 100644
--- a/python/sedonadb/tests/test_dataframe_ffi.py
+++ b/python/sedonadb/tests/test_dataframe_ffi.py
@@ -16,9 +16,11 @@
# under the License.
import pandas as pd
+import pyarrow as pa
import pytest
import sedonadb
+from sedonadb import udf
from sedonadb.testing import skip_if_not_exists
@@ -121,3 +123,109 @@ def test_ffi_roundtrip(geoarrow_data, producer_sql, consumer_sql):
result_over_ffi = sd_consumer.sql(consumer_sql).to_pandas()
pd.testing.assert_frame_equal(result_no_ffi, result_over_ffi)
+
+
+def test_filter_pushdown_into_ffi_producer(geoarrow_data):
+ path = geoarrow_data / "ns-water" / "files" / "ns-water_water-point_geo.parquet"
+ skip_if_not_exists(path)
+
+ sd_producer = sedonadb.connect()
+ sd_consumer = sedonadb.connect()
+
+ # Producer exposes all columns without filtering
+ sd_producer.read_parquet(path).to_view("water_point")
+ df_producer = sd_producer.sql('SELECT "OBJECTID", "FEAT_CODE" FROM water_point')
+
+ # Consumer applies a filter
+ sd_consumer.create_data_frame(df_producer).to_view("df_producer")
+ df_filtered = sd_consumer.sql('SELECT * FROM df_producer WHERE "OBJECTID" < 50')
+
+ # Get the physical plan (second row)
+ plan_df = df_filtered.explain().to_pandas()
+ plan_text = str(plan_df["plan"].iloc[1])
+
+ # The filter should appear inside ImportedSedonaCExec, indicating it was pushed down
+ # Look for the filter predicate appearing after ImportedSedonaCExec in the plan
+ assert "ImportedSedonaCExec" in plan_text, (
+ f"Expected ImportedSedonaCExec in plan:\n{plan_text}"
+ )
+
+ # The filter "OBJECTID < 50" should be pushed into the producer, not applied as a
+ # separate FilterExec on top of the ImportedSedonaCExec
+ # If pushdown works, the filter should appear within the ImportedSedonaCExec display
+ imported_idx = plan_text.find("ImportedSedonaCExec")
+ filter_idx = plan_text.find("FilterExec")
+
+ # If there's a FilterExec, it should NOT be wrapping the ImportedSedonaCExec
+ # (i.e., filter should come after imported in the plan text, meaning it's inside)
+ if filter_idx != -1:
+ # FilterExec appearing before ImportedSedonaCExec means filter is NOT pushed down
+ assert filter_idx > imported_idx, (
+ f"Filter was not pushed down into FFI producer. "
+ f"FilterExec at {filter_idx}, ImportedSedonaCExec at {imported_idx}.\n"
+ f"Plan:\n{plan_text}"
+ )
+
+
+def test_udf_filter_not_pushed_down_into_ffi_producer(geoarrow_data):
+ """Verify that filters using consumer-side UDFs are NOT pushed down into the FFI producer.
+
+ UDFs are session-specific, so a UDF registered on the consumer cannot be
+ evaluated by the producer. The filter must remain on the consumer side.
+ """
+ path = geoarrow_data / "ns-water" / "files" / "ns-water_water-point_geo.parquet"
+ skip_if_not_exists(path)
+
+ sd_producer = sedonadb.connect()
+ sd_consumer = sedonadb.connect()
+
+ # Define a simple UDF on the consumer side only
+ @udf.arrow_udf(pa.bool_(), [pa.int64()])
+ def is_small_id(ids):
+ import pyarrow.compute as pc
+
+ # Convert from SedonaDB internal type to PyArrow array
+ ids_arr = pa.array(ids.to_array())
+ return pc.less(ids_arr, 50)
+
+ sd_consumer.register(is_small_id)
+
+ # Producer exposes all columns without filtering
+ sd_producer.read_parquet(path).to_view("water_point")
+ df_producer = sd_producer.sql('SELECT "OBJECTID", "FEAT_CODE" FROM water_point')
+
+ # Consumer applies a filter using the UDF (which only exists on consumer)
+ sd_consumer.create_data_frame(df_producer).to_view("df_producer")
+ df_filtered = sd_consumer.sql(
+ 'SELECT * FROM df_producer WHERE is_small_id("OBJECTID")'
+ )
+
+ # Get the physical plan (second row)
+ plan_df = df_filtered.explain().to_pandas()
+ plan_text = str(plan_df["plan"].iloc[1])
+
+ assert "ImportedSedonaCExec" in plan_text, (
+ f"Expected ImportedSedonaCExec in plan:\n{plan_text}"
+ )
+
+ # The UDF filter should NOT be pushed down - it must stay on the consumer side
+ # This means FilterExec should appear BEFORE ImportedSedonaCExec in the plan
+ # (FilterExec wraps ImportedSedonaCExec, not inside it)
+ imported_idx = plan_text.find("ImportedSedonaCExec")
+ filter_idx = plan_text.find("FilterExec")
+
+ assert filter_idx != -1, (
+ f"Expected FilterExec in plan (UDF filter should not be pushed down):\n{plan_text}"
+ )
+
+ # FilterExec appearing before ImportedSedonaCExec means it's wrapping it (not pushed)
+ assert filter_idx < imported_idx, (
+ f"UDF filter was incorrectly pushed down into FFI producer. "
+ f"FilterExec at {filter_idx}, ImportedSedonaCExec at {imported_idx}.\n"
+ f"Plan:\n{plan_text}"
+ )
+
+ # Also verify the query actually works and returns correct results
+ result = df_filtered.to_pandas()
+ assert len(result) > 0, "Expected some results"
+ assert all(result["OBJECTID"] < 50), "UDF filter should have been applied"
diff --git a/rust/sedona-expr/Cargo.toml b/rust/sedona-expr/Cargo.toml
index 39c9cf566..d0887d4b1 100644
--- a/rust/sedona-expr/Cargo.toml
+++ b/rust/sedona-expr/Cargo.toml
@@ -39,6 +39,7 @@ arrow-array = { workspace = true }
arrow-buffer = { workspace = true }
arrow-schema = { workspace = true }
datafusion-common = { workspace = true }
+datafusion-execution = { workspace = true }
datafusion-expr = { workspace = true }
datafusion-physical-expr = { workspace = true }
geo-traits = { workspace = true }
diff --git a/rust/sedona-expr/src/lib.rs b/rust/sedona-expr/src/lib.rs
index 7346cf1bc..c865dbdf5 100644
--- a/rust/sedona-expr/src/lib.rs
+++ b/rust/sedona-expr/src/lib.rs
@@ -18,6 +18,7 @@ pub mod aggregate_udf;
pub mod function_set;
pub mod item_crs;
pub mod metadata_preserving_column;
+pub mod placeholder_udf;
pub mod scalar_udf;
pub mod spatial_filter;
pub mod statistics;
diff --git a/rust/sedona-expr/src/placeholder_udf.rs b/rust/sedona-expr/src/placeholder_udf.rs
new file mode 100644
index 000000000..56b700dd1
--- /dev/null
+++ b/rust/sedona-expr/src/placeholder_udf.rs
@@ -0,0 +1,712 @@
+// 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.
+
+//! Placeholder UDFs for expression deserialization
+//!
+//! This module provides placeholder implementations of scalar, aggregate, and window
+//! user-defined functions (UDFs). These are used when deserializing expressions where
+//! the actual execution of the expression is not necessarily important (e.g., for an
+//! implementation of TableProvider that independently parses expressions for pruning,
+//! or an implementation that does not use the filters expression at all).
+
+use std::{
+ collections::HashSet,
+ sync::{Arc, OnceLock},
+};
+
+use arrow_schema::{DataType, Field};
+use datafusion_common::{
+ tree_node::{Transformed, TreeNode, TreeNodeRecursion},
+ Result,
+};
+use datafusion_execution::FunctionRegistry;
+use datafusion_expr::{
+ expr::ScalarFunction,
+ function::{AccumulatorArgs, PartitionEvaluatorArgs, WindowUDFFieldArgs},
+ Accumulator, AggregateUDF, AggregateUDFImpl, Expr, PartitionEvaluator, ScalarUDF,
+ ScalarUDFImpl, Signature, Volatility, WindowFunctionDefinition, WindowUDF, WindowUDFImpl,
+};
+use sedona_common::sedona_internal_err;
+
+/// A [FunctionRegistry] that creates placeholder UDFs for any requested function.
+///
+/// This registry is useful for deserializing expressions where the actual function
+/// implementations are not available or not needed. All functions are resolved to
+/// placeholder stubs that preserve the function name but will error if actually invoked.
+pub struct PlaceholderRegistry;
+
+impl PlaceholderRegistry {
+ /// Check if an expression contains any [PlaceholderUDF] functions.
+ pub fn expr_contains_placeholder(expr: &Expr) -> bool {
+ let mut found = false;
+ expr.apply(|e| {
+ if let Expr::ScalarFunction(func) = e {
+ if func
+ .func
+ .inner()
+ .as_any()
+ .downcast_ref::<PlaceholderUDF>()
+ .is_some()
+ {
+ found = true;
+ return Ok(TreeNodeRecursion::Stop);
+ }
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })
+ .expect("expr_contains_placeholder traversal should not fail");
+ found
+ }
+
+ /// Check if an expression contains any placeholder functions (scalar, aggregate, or window).
+ pub fn expr_contains_any_placeholder(expr: &Expr) -> Result<bool> {
+ let mut found = false;
+ expr.apply(|e| {
+ match e {
+ Expr::ScalarFunction(func) => {
+ if func
+ .func
+ .inner()
+ .as_any()
+ .downcast_ref::<PlaceholderUDF>()
+ .is_some()
+ {
+ found = true;
+ return Ok(TreeNodeRecursion::Stop);
+ }
+ }
+ Expr::AggregateFunction(func) => {
+ if func
+ .func
+ .inner()
+ .as_any()
+ .downcast_ref::<PlaceholderUDAF>()
+ .is_some()
+ {
+ found = true;
+ return Ok(TreeNodeRecursion::Stop);
+ }
+ }
+ Expr::WindowFunction(func) => {
+ if let WindowFunctionDefinition::WindowUDF(ref udf) = func.fun {
+ if udf
+ .inner()
+ .as_any()
+ .downcast_ref::<PlaceholderUDWF>()
+ .is_some()
+ {
+ found = true;
+ return Ok(TreeNodeRecursion::Stop);
+ }
+ }
+ }
+ _ => {}
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+
+ Ok(found)
+ }
+
+ /// Replace all placeholder functions (scalar, aggregate, window) with implementations from the registry.
+ ///
+ /// Returns an error if any placeholder function is not found in the registry.
+ pub fn expr_replace_placeholders(expr: Expr, registry: &dyn FunctionRegistry) -> Result<Expr> {
+ expr.transform_up(|e| {
+ match &e {
+ Expr::ScalarFunction(func) => {
+ if func
+ .func
+ .inner()
+ .as_any()
+ .downcast_ref::<PlaceholderUDF>()
+ .is_some()
+ {
+ let real_udf = registry.udf(func.name())?;
+ let replaced = Expr::ScalarFunction(ScalarFunction {
+ func: real_udf,
+ args: func.args.clone(),
+ });
+ return Ok(Transformed::yes(replaced));
+ }
+ }
+ Expr::AggregateFunction(func) => {
+ if func
+ .func
+ .inner()
+ .as_any()
+ .downcast_ref::<PlaceholderUDAF>()
+ .is_some()
+ {
+ let real_udaf = registry.udaf(func.func.name())?;
+ let replaced = Expr::AggregateFunction(
+ datafusion_expr::expr::AggregateFunction::new_udf(
+ real_udaf,
+ func.params.args.clone(),
+ func.params.distinct,
+ func.params.filter.clone(),
+ func.params.order_by.clone(),
+ func.params.null_treatment,
+ ),
+ );
+ return Ok(Transformed::yes(replaced));
+ }
+ }
+ Expr::WindowFunction(func) => {
+ if let WindowFunctionDefinition::WindowUDF(ref udf) = func.fun {
+ if udf
+ .inner()
+ .as_any()
+ .downcast_ref::<PlaceholderUDWF>()
+ .is_some()
+ {
+ let real_udwf = registry.udwf(udf.name())?;
+ let replaced = Expr::WindowFunction(Box::new(
+ datafusion_expr::expr::WindowFunction {
+ fun: WindowFunctionDefinition::WindowUDF(real_udwf),
+ params: func.params.clone(),
+ },
+ ));
+ return Ok(Transformed::yes(replaced));
+ }
+ }
+ }
+ _ => {}
+ }
+ Ok(Transformed::no(e))
+ })
+ .map(|t| t.data)
+ }
+}
+
+impl FunctionRegistry for PlaceholderRegistry {
+ fn udfs(&self) -> HashSet<String> {
+ HashSet::new()
+ }
+
+ fn udafs(&self) -> HashSet<String> {
+ HashSet::new()
+ }
+
+ fn udwfs(&self) -> HashSet<String> {
+ HashSet::new()
+ }
+
+ fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>> {
+ Ok(Arc::new(ScalarUDF::new_from_impl(PlaceholderUDF::new(
+ name,
+ ))))
+ }
+
+ fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {
+ Ok(Arc::new(AggregateUDF::new_from_impl(PlaceholderUDAF::new(
+ name,
+ ))))
+ }
+
+ fn udwf(&self, name: &str) -> Result<Arc<WindowUDF>> {
+ Ok(Arc::new(WindowUDF::new_from_impl(PlaceholderUDWF::new(
+ name,
+ ))))
+ }
+
+ fn expr_planners(&self) -> Vec<Arc<dyn datafusion_expr::planner::ExprPlanner>> {
+ vec![]
+ }
+}
+
+/// Placeholder [ScalarUDF] that preserves a function name.
+///
+/// This struct is a stub for deserializing expressions where the actual execution
+/// of the expression is not necessarily important.
+#[derive(Debug, Hash, PartialEq, Eq)]
+pub struct PlaceholderUDF {
+ name: String,
+}
+
+impl PlaceholderUDF {
+ /// Create a new placeholder UDF with the given name.
+ pub fn new(name: &str) -> Self {
+ PlaceholderUDF {
+ name: name.to_string(),
+ }
+ }
+}
+
+impl ScalarUDFImpl for PlaceholderUDF {
+ fn as_any(&self) -> &dyn std::any::Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ // In normal usage this is not passed through an optimizer (or if it is, it
+ // will error as soon as its return field is inspected); however, we declare
+ // this as Volatile to be safe.
+ static SIGNATURE_ANY: OnceLock<Signature> = OnceLock::new();
+ SIGNATURE_ANY.get_or_init(|| Signature::variadic_any(Volatility::Volatile))
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ sedona_internal_err!(
+ "Imported placeholder UDF '{}' must be replaced before planning",
+ self.name
+ )
+ }
+
+ fn invoke_with_args(
+ &self,
+ _args: datafusion_expr::ScalarFunctionArgs,
+ ) -> Result<datafusion_expr::ColumnarValue> {
+ sedona_internal_err!(
+ "Imported placeholder UDF '{}' must be replaced before execution",
+ self.name
+ )
+ }
+}
+
+/// Placeholder [AggregateUDF] that preserves a function name.
+///
+/// This struct is a stub for deserializing expressions where the actual execution
+/// of the expression is not necessarily important.
+#[derive(Debug, Hash, PartialEq, Eq)]
+pub struct PlaceholderUDAF {
+ name: String,
+}
+
+impl PlaceholderUDAF {
+ /// Create a new placeholder aggregate UDF with the given name.
+ pub fn new(name: &str) -> Self {
+ PlaceholderUDAF {
+ name: name.to_string(),
+ }
+ }
+}
+
+impl AggregateUDFImpl for PlaceholderUDAF {
+ fn as_any(&self) -> &dyn std::any::Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ static SIGNATURE_ANY: OnceLock<Signature> = OnceLock::new();
+ SIGNATURE_ANY.get_or_init(|| Signature::variadic_any(Volatility::Volatile))
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ sedona_internal_err!(
+ "Imported placeholder UDAF '{}' must be replaced before planning",
+ self.name
+ )
+ }
+
+ fn accumulator(&self, _args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+ sedona_internal_err!(
+ "Imported placeholder UDAF '{}' must be replaced before execution",
+ self.name
+ )
+ }
+}
+
+/// Placeholder [WindowUDF] that preserves a function name.
+///
+/// This struct is a stub for deserializing expressions where the actual execution
+/// of the expression is not necessarily important.
+#[derive(Debug, Hash, PartialEq, Eq)]
+pub struct PlaceholderUDWF {
+ name: String,
+}
+
+impl PlaceholderUDWF {
+ /// Create a new placeholder window UDF with the given name.
+ pub fn new(name: &str) -> Self {
+ PlaceholderUDWF {
+ name: name.to_string(),
+ }
+ }
+}
+
+impl WindowUDFImpl for PlaceholderUDWF {
+ fn as_any(&self) -> &dyn std::any::Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ static SIGNATURE_ANY: OnceLock<Signature> = OnceLock::new();
+ SIGNATURE_ANY.get_or_init(|| Signature::variadic_any(Volatility::Volatile))
+ }
+
+ fn field(&self, _field_args: WindowUDFFieldArgs) -> Result<Arc<Field>> {
+ sedona_internal_err!(
+ "Imported placeholder UDWF '{}' must be replaced before planning",
+ self.name
+ )
+ }
+
+ fn partition_evaluator(
+ &self,
+ _partition_evaluator_args: PartitionEvaluatorArgs,
+ ) -> Result<Box<dyn PartitionEvaluator>> {
+ sedona_internal_err!(
+ "Imported placeholder UDWF '{}' must be replaced before execution",
+ self.name
+ )
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{col, lit, SimpleScalarUDF};
+
+ #[test]
+ fn test_placeholder_registry_returns_empty_function_sets() {
+ let registry = PlaceholderRegistry;
+ assert!(registry.udfs().is_empty());
+ assert!(registry.udafs().is_empty());
+ assert!(registry.udwfs().is_empty());
+ }
+
+ #[test]
+ fn test_placeholder_registry_creates_placeholder_udf() {
+ let registry = PlaceholderRegistry;
+ let udf = registry.udf("test_func").unwrap();
+ assert_eq!(udf.name(), "test_func");
+
+ // Verify it's a PlaceholderUDF
+ let inner = udf.inner();
+ assert!(inner.as_any().downcast_ref::<PlaceholderUDF>().is_some());
+ }
+
+ #[test]
+ fn test_placeholder_registry_creates_placeholder_udaf() {
+ let registry = PlaceholderRegistry;
+ let udaf = registry.udaf("test_agg").unwrap();
+ assert_eq!(udaf.name(), "test_agg");
+
+ // Verify it's a PlaceholderUDAF
+ let inner = udaf.inner();
+ assert!(inner.as_any().downcast_ref::<PlaceholderUDAF>().is_some());
+ }
+
+ #[test]
+ fn test_placeholder_registry_creates_placeholder_udwf() {
+ let registry = PlaceholderRegistry;
+ let udwf = registry.udwf("test_window").unwrap();
+ assert_eq!(udwf.name(), "test_window");
+
+ // Verify it's a PlaceholderUDWF
+ let inner = udwf.inner();
+ assert!(inner.as_any().downcast_ref::<PlaceholderUDWF>().is_some());
+ }
+
+ #[test]
+ fn test_placeholder_registry_expr_planners_empty() {
+ let registry = PlaceholderRegistry;
+ assert!(registry.expr_planners().is_empty());
+ }
+
+ #[test]
+ fn test_placeholder_udf_name() {
+ let udf = PlaceholderUDF::new("my_func");
+ assert_eq!(udf.name(), "my_func");
+ }
+
+ #[test]
+ fn test_placeholder_udf_signature_is_variadic() {
+ let udf = PlaceholderUDF::new("test");
+ let sig = udf.signature();
+ // The signature should accept any arguments
+ assert!(matches!(sig.volatility, Volatility::Volatile));
+ }
+
+ #[test]
+ fn test_placeholder_udf_return_type_errors() {
+ let udf = PlaceholderUDF::new("my_func");
+ let result = udf.return_type(&[DataType::Int32]);
+ assert!(result.is_err());
+ let err = result.unwrap_err().to_string();
+ assert!(err.contains("my_func"));
+ assert!(err.contains("must be replaced before planning"));
+ }
+
+ #[test]
+ fn test_placeholder_udf_equality() {
+ let udf1 = PlaceholderUDF::new("func_a");
+ let udf2 = PlaceholderUDF::new("func_a");
+ let udf3 = PlaceholderUDF::new("func_b");
+
+ assert_eq!(udf1, udf2);
+ assert_ne!(udf1, udf3);
+ }
+
+ #[test]
+ fn test_placeholder_udf_debug() {
+ let udf = PlaceholderUDF::new("debug_test");
+ let debug_str = format!("{:?}", udf);
+ assert!(debug_str.contains("PlaceholderUDF"));
+ assert!(debug_str.contains("debug_test"));
+ }
+
+ #[test]
+ fn test_placeholder_udaf_name() {
+ let udaf = PlaceholderUDAF::new("my_agg");
+ assert_eq!(udaf.name(), "my_agg");
+ }
+
+ #[test]
+ fn test_placeholder_udaf_signature_is_variadic() {
+ let udaf = PlaceholderUDAF::new("test");
+ let sig = udaf.signature();
+ assert!(matches!(sig.volatility, Volatility::Volatile));
+ }
+
+ #[test]
+ fn test_placeholder_udaf_return_type_errors() {
+ let udaf = PlaceholderUDAF::new("my_agg");
+ let result = udaf.return_type(&[DataType::Int32]);
+ assert!(result.is_err());
+ let err = result.unwrap_err().to_string();
+ assert!(err.contains("my_agg"));
+ assert!(err.contains("must be replaced before planning"));
+ }
+
+ #[test]
+ fn test_placeholder_udaf_equality() {
+ let udaf1 = PlaceholderUDAF::new("agg_a");
+ let udaf2 = PlaceholderUDAF::new("agg_a");
+ let udaf3 = PlaceholderUDAF::new("agg_b");
+
+ assert_eq!(udaf1, udaf2);
+ assert_ne!(udaf1, udaf3);
+ }
+
+ #[test]
+ fn test_placeholder_udaf_debug() {
+ let udaf = PlaceholderUDAF::new("debug_agg");
+ let debug_str = format!("{:?}", udaf);
+ assert!(debug_str.contains("PlaceholderUDAF"));
+ assert!(debug_str.contains("debug_agg"));
+ }
+
+ #[test]
+ fn test_placeholder_udwf_name() {
+ let udwf = PlaceholderUDWF::new("my_window");
+ assert_eq!(udwf.name(), "my_window");
+ }
+
+ #[test]
+ fn test_placeholder_udwf_signature_is_variadic() {
+ let udwf = PlaceholderUDWF::new("test");
+ let sig = udwf.signature();
+ assert!(matches!(sig.volatility, Volatility::Volatile));
+ }
+
+ #[test]
+ fn test_placeholder_udwf_return_type_errors() {
+ // Note: WindowUDFImpl uses `field` instead of `return_type`
+ // The field method requires WindowUDFFieldArgs which is hard to construct in tests,
+ // so we just verify the signature exists and the UDF can be created.
+ let udwf = PlaceholderUDWF::new("my_window");
+ assert_eq!(udwf.name(), "my_window");
+ }
+
+ #[test]
+ fn test_placeholder_udwf_partition_evaluator_errors() {
+ // Note: partition_evaluator requires PartitionEvaluatorArgs which is hard to construct
+ // in tests, so we just verify the UDF can be created and wrapped.
+ let udwf = PlaceholderUDWF::new("my_window");
+ let _wrapped = WindowUDF::new_from_impl(udwf);
+ }
+
+ #[test]
+ fn test_placeholder_udwf_equality() {
+ let udwf1 = PlaceholderUDWF::new("win_a");
+ let udwf2 = PlaceholderUDWF::new("win_a");
+ let udwf3 = PlaceholderUDWF::new("win_b");
+
+ assert_eq!(udwf1, udwf2);
+ assert_ne!(udwf1, udwf3);
+ }
+
+ #[test]
+ fn test_placeholder_udwf_debug() {
+ let udwf = PlaceholderUDWF::new("debug_window");
+ let debug_str = format!("{:?}", udwf);
+ assert!(debug_str.contains("PlaceholderUDWF"));
+ assert!(debug_str.contains("debug_window"));
+ }
+
+ #[test]
+ fn test_expr_contains_placeholder_true() {
+ let udf = ScalarUDF::new_from_impl(PlaceholderUDF::new("test_func"));
+ let expr = udf.call(vec![col("x")]);
+
+ assert!(PlaceholderRegistry::expr_contains_placeholder(&expr));
+ }
+
+ #[test]
+ fn test_expr_contains_placeholder_false_for_column() {
+ let expr = col("x");
+ assert!(!PlaceholderRegistry::expr_contains_placeholder(&expr));
+ }
+
+ #[test]
+ fn test_expr_contains_placeholder_nested() {
+ let udf = ScalarUDF::new_from_impl(PlaceholderUDF::new("inner_func"));
+ let inner = udf.call(vec![col("x")]);
+ // Wrap in a binary expression
+ let expr = inner.gt(lit(5i32));
+
+ assert!(PlaceholderRegistry::expr_contains_placeholder(&expr));
+ }
+
+ #[test]
+ fn test_expr_contains_any_placeholder_with_scalar() {
+ let udf = ScalarUDF::new_from_impl(PlaceholderUDF::new("test_func"));
+ let expr = udf.call(vec![col("x")]);
+
+ assert!(PlaceholderRegistry::expr_contains_any_placeholder(&expr).unwrap());
+ }
+
+ #[test]
+ fn test_expr_contains_any_placeholder_with_aggregate() {
+ let udaf = AggregateUDF::new_from_impl(PlaceholderUDAF::new("test_agg"));
+ let expr = udaf.call(vec![col("x")]);
+
+ assert!(PlaceholderRegistry::expr_contains_any_placeholder(&expr).unwrap());
+ }
+
+ #[test]
+ fn test_expr_contains_any_placeholder_false_for_column() {
+ let expr = col("x");
+ assert!(!PlaceholderRegistry::expr_contains_any_placeholder(&expr).unwrap());
+ }
+
+ #[test]
+ fn test_expr_replace_placeholder_udfs() {
+ // Create a simple registry that returns a real UDF implementation
+ // and attempts aggregate and window replacement (these are verbose
+ // to mock so we don't do that here)
+ struct TestRegistry;
+ impl FunctionRegistry for TestRegistry {
+ fn udfs(&self) -> HashSet<String> {
+ HashSet::new()
+ }
+
+ fn udafs(&self) -> HashSet<String> {
+ HashSet::new()
+ }
+
+ fn udwfs(&self) -> HashSet<String> {
+ HashSet::new()
+ }
+
+ fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>> {
+ // Return a real UDF using SimpleScalarUDF with a closure
+ let udf: ScalarUDF = SimpleScalarUDF::new_with_signature(
+ name,
+ Signature::any(1, Volatility::Immutable),
+ DataType::Int32,
+ Arc::new(|_args| Ok(ScalarValue::Int32(Some(42)).into())),
+ )
+ .into();
+ Ok(Arc::new(udf))
+ }
+
+ fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {
+ sedona_internal_err!("attempt to replace udaf '{name}'")
+ }
+
+ fn udwf(&self, name: &str) -> Result<Arc<WindowUDF>> {
+ sedona_internal_err!("attempt to replace udwf '{name}'")
+ }
+
+ fn expr_planners(&self) -> Vec<Arc<dyn datafusion_expr::planner::ExprPlanner>> {
+ vec![]
+ }
+ }
+
+ let udf = ScalarUDF::new_from_impl(PlaceholderUDF::new("original"));
+ let expr = udf.call(vec![col("x")]);
+
+ // Verify it starts as a placeholder
+ assert!(PlaceholderRegistry::expr_contains_placeholder(&expr));
+
+ let registry = TestRegistry;
+ let replaced = PlaceholderRegistry::expr_replace_placeholders(expr, ®istry).unwrap();
+
+ // Verify the function was replaced and is no longer a placeholder
+ assert!(!PlaceholderRegistry::expr_contains_placeholder(&replaced));
+ match &replaced {
+ Expr::ScalarFunction(ScalarFunction { func, .. }) => {
+ assert_eq!(func.name(), "original");
+ // Verify it's NOT a PlaceholderUDF anymore
+ assert!(func
+ .inner()
+ .as_any()
+ .downcast_ref::<PlaceholderUDF>()
+ .is_none());
+ }
+ other => panic!("Expected ScalarFunction, got {:?}", other),
+ }
+
+ // Test aggregate replacement is attempted (errors prove the registry was called)
+ let udaf = AggregateUDF::new_from_impl(PlaceholderUDAF::new("my_aggregate"));
+ let agg_expr = udaf.call(vec![col("x")]);
+ assert!(PlaceholderRegistry::expr_contains_any_placeholder(&agg_expr).unwrap());
+
+ let agg_result = PlaceholderRegistry::expr_replace_placeholders(agg_expr, ®istry);
+ assert!(agg_result.is_err());
+ let agg_err = agg_result.unwrap_err().to_string();
+ assert!(
+ agg_err.contains("attempt to replace udaf 'my_aggregate'"),
+ "Expected error about udaf replacement attempt, got: {}",
+ agg_err
+ );
+
+ // Test window replacement is attempted (errors prove the registry was called)
+ let udwf = WindowUDF::new_from_impl(PlaceholderUDWF::new("my_window"));
+ let window_func = datafusion_expr::expr::WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(Arc::new(udwf)),
+ vec![col("x")],
+ );
+ let win_expr = Expr::WindowFunction(Box::new(window_func));
+ assert!(PlaceholderRegistry::expr_contains_any_placeholder(&win_expr).unwrap());
+
+ let win_result = PlaceholderRegistry::expr_replace_placeholders(win_expr, ®istry);
+ assert!(win_result.is_err());
+ let win_err = win_result.unwrap_err().to_string();
+ assert!(
+ win_err.contains("attempt to replace udwf 'my_window'"),
+ "Expected error about udwf replacement attempt, got: {}",
+ win_err
+ );
+ }
+}