Add unit tests for physical optimizer rule

timsaucer · timsaucer · commit 039a1b429296 · 2026-03-26T06:40:56.000-04:00
diff --git a/datafusion/ffi/src/physical_optimizer.rs b/datafusion/ffi/src/physical_optimizer.rs
@@ -219,3 +219,153 @@ impl PhysicalOptimizerRule for ForeignPhysicalOptimizerRule {
         unsafe { (self.0.schema_check)(&self.0) }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::datatypes::{DataType, Field, Schema};
+    use datafusion_common::config::ConfigOptions;
+    use datafusion_common::error::Result;
+    use datafusion_physical_optimizer::PhysicalOptimizerRule;
+    use datafusion_physical_plan::ExecutionPlan;
+
+    use crate::execution_plan::tests::EmptyExec;
+
+    use super::*;
+
+    #[derive(Debug)]
+    struct NoOpRule {
+        schema_check: bool,
+    }
+
+    impl PhysicalOptimizerRule for NoOpRule {
+        fn optimize(
+            &self,
+            plan: Arc<dyn ExecutionPlan>,
+            _config: &ConfigOptions,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            Ok(plan)
+        }
+
+        fn name(&self) -> &str {
+            "no_op_rule"
+        }
+
+        fn schema_check(&self) -> bool {
+            self.schema_check
+        }
+    }
+
+    fn create_test_plan() -> Arc<dyn ExecutionPlan> {
+        let schema =
+            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
+        Arc::new(EmptyExec::new(schema))
+    }
+
+    #[test]
+    fn test_round_trip_ffi_physical_optimizer_rule() -> Result<()> {
+        for expected_schema_check in [true, false] {
+            let rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> = Arc::new(NoOpRule {
+                schema_check: expected_schema_check,
+            });
+
+            let mut ffi_rule = FFI_PhysicalOptimizerRule::new(rule, None);
+            ffi_rule.library_marker_id = crate::mock_foreign_marker_id;
+
+            let foreign_rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
+                (&ffi_rule).into();
+
+            assert_eq!(foreign_rule.name(), "no_op_rule");
+            assert_eq!(foreign_rule.schema_check(), expected_schema_check);
+        }
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_round_trip_optimize() -> Result<()> {
+        let rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
+            Arc::new(NoOpRule { schema_check: true });
+
+        let mut ffi_rule = FFI_PhysicalOptimizerRule::new(rule, None);
+        ffi_rule.library_marker_id = crate::mock_foreign_marker_id;
+
+        let foreign_rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
+            (&ffi_rule).into();
+
+        let plan = create_test_plan();
+        let config = ConfigOptions::new();
+
+        let optimized = foreign_rule.optimize(plan, &config)?;
+        assert_eq!(optimized.name(), "empty-exec");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_local_bypass() -> Result<()> {
+        let rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
+            Arc::new(NoOpRule { schema_check: true });
+
+        // Without mock marker, local bypass should return the original rule
+        let ffi_rule = FFI_PhysicalOptimizerRule::new(rule, None);
+        let recovered: Arc<dyn PhysicalOptimizerRule + Send + Sync> = (&ffi_rule).into();
+        let any_ref: &dyn std::any::Any = &*recovered;
+        assert!(any_ref.downcast_ref::<NoOpRule>().is_some());
+
+        // With mock marker, should wrap in ForeignPhysicalOptimizerRule
+        let rule2: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
+            Arc::new(NoOpRule { schema_check: true });
+        let mut ffi_rule2 = FFI_PhysicalOptimizerRule::new(rule2, None);
+        ffi_rule2.library_marker_id = crate::mock_foreign_marker_id;
+        let recovered2: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
+            (&ffi_rule2).into();
+        let any_ref2: &dyn std::any::Any = &*recovered2;
+        assert!(
+            any_ref2
+                .downcast_ref::<ForeignPhysicalOptimizerRule>()
+                .is_some()
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_clone() -> Result<()> {
+        let rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
+            Arc::new(NoOpRule { schema_check: true });
+
+        let ffi_rule = FFI_PhysicalOptimizerRule::new(rule, None);
+        let cloned = ffi_rule.clone();
+
+        assert_eq!(unsafe { (ffi_rule.name)(&ffi_rule).as_str() }, unsafe {
+            (cloned.name)(&cloned).as_str()
+        });
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_foreign_rule_rewrap_bypass() -> Result<()> {
+        // When creating an FFI wrapper from a ForeignPhysicalOptimizerRule,
+        // it should return the inner FFI rule rather than double-wrapping.
+        let rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
+            Arc::new(NoOpRule { schema_check: true });
+
+        let mut ffi_rule = FFI_PhysicalOptimizerRule::new(rule, None);
+        ffi_rule.library_marker_id = crate::mock_foreign_marker_id;
+
+        let foreign_rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
+            (&ffi_rule).into();
+
+        // Now wrap the foreign rule back into FFI - should not double-wrap
+        let re_wrapped = FFI_PhysicalOptimizerRule::new(foreign_rule, None);
+        assert_eq!(
+            unsafe { (re_wrapped.name)(&re_wrapped).as_str() },
+            "no_op_rule"
+        );
+
+        Ok(())
+    }
+}
diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs
@@ -39,6 +39,7 @@ use crate::catalog_provider_list::FFI_CatalogProviderList;
 use crate::config::extension_options::FFI_ExtensionOptions;
 use crate::execution_plan::FFI_ExecutionPlan;
 use crate::execution_plan::tests::EmptyExec;
+use crate::physical_optimizer::FFI_PhysicalOptimizerRule;
 use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
 use crate::table_provider::FFI_TableProvider;
 use crate::table_provider_factory::FFI_TableProviderFactory;
@@ -51,6 +52,7 @@ use crate::udwf::FFI_WindowUDF;
 mod async_provider;
 pub mod catalog;
 pub mod config;
+mod physical_optimizer;
 mod sync_provider;
 mod table_provider_factory;
 mod udf_udaf_udwf;
@@ -104,6 +106,8 @@ pub struct ForeignLibraryModule {
 
     pub create_empty_exec: extern "C" fn() -> FFI_ExecutionPlan,
 
+    pub create_physical_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule,
+
     pub version: extern "C" fn() -> u64,
 }
 
@@ -177,6 +181,8 @@ pub fn get_foreign_library_module() -> ForeignLibraryModuleRef {
         create_rank_udwf: create_ffi_rank_func,
         create_extension_options: config::create_extension_options,
         create_empty_exec,
+        create_physical_optimizer_rule:
+            physical_optimizer::create_physical_optimizer_rule,
         version: super::version,
     }
     .leak_into_prefix()
diff --git a/datafusion/ffi/src/tests/physical_optimizer.rs b/datafusion/ffi/src/tests/physical_optimizer.rs
@@ -0,0 +1,54 @@
+// 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::sync::Arc;
+
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_physical_optimizer::PhysicalOptimizerRule;
+use datafusion_physical_plan::ExecutionPlan;
+use datafusion_physical_plan::limit::GlobalLimitExec;
+
+use crate::physical_optimizer::FFI_PhysicalOptimizerRule;
+
+/// A rule that wraps the input plan in a GlobalLimitExec with skip=0, fetch=10.
+/// This produces an observable change in the plan tree that tests can verify.
+#[derive(Debug)]
+struct AddLimitRule;
+
+impl PhysicalOptimizerRule for AddLimitRule {
+    fn optimize(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        _config: &ConfigOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        Ok(Arc::new(GlobalLimitExec::new(plan, 0, Some(10))))
+    }
+
+    fn name(&self) -> &str {
+        "add_limit_rule"
+    }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
+}
+
+pub(crate) extern "C" fn create_physical_optimizer_rule() -> FFI_PhysicalOptimizerRule {
+    let rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> = Arc::new(AddLimitRule);
+    FFI_PhysicalOptimizerRule::new(rule, None)
+}
diff --git a/datafusion/ffi/tests/ffi_physical_optimizer.rs b/datafusion/ffi/tests/ffi_physical_optimizer.rs
@@ -0,0 +1,74 @@
+// 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.
+
+#[cfg(feature = "integration-tests")]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::datatypes::{DataType, Field, Schema};
+    use datafusion_common::DataFusionError;
+    use datafusion_common::config::ConfigOptions;
+    use datafusion_ffi::execution_plan::tests::EmptyExec;
+    use datafusion_ffi::physical_optimizer::ForeignPhysicalOptimizerRule;
+    use datafusion_ffi::tests::utils::get_module;
+    use datafusion_physical_optimizer::PhysicalOptimizerRule;
+    use datafusion_physical_plan::ExecutionPlan;
+
+    fn create_test_plan() -> Arc<dyn ExecutionPlan> {
+        let schema =
+            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
+        Arc::new(EmptyExec::new(schema))
+    }
+
+    #[test]
+    fn test_ffi_physical_optimizer_rule() -> Result<(), DataFusionError> {
+        let module = get_module()?;
+
+        let ffi_rule = module.create_physical_optimizer_rule().ok_or(
+            DataFusionError::NotImplemented(
+                "External module failed to implement create_physical_optimizer_rule"
+                    .to_string(),
+            ),
+        )?();
+
+        let foreign_rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
+            (&ffi_rule).into();
+
+        // Verify the rule is wrapped as a foreign rule
+        let any_ref: &dyn std::any::Any = &*foreign_rule;
+        assert!(
+            any_ref
+                .downcast_ref::<ForeignPhysicalOptimizerRule>()
+                .is_some()
+        );
+
+        // Verify name and schema_check pass through FFI
+        assert_eq!(foreign_rule.name(), "add_limit_rule");
+        assert!(foreign_rule.schema_check());
+
+        // Verify the rule actually transforms the plan
+        let plan = create_test_plan();
+        let config = ConfigOptions::new();
+        let optimized = foreign_rule.optimize(plan, &config)?;
+
+        assert_eq!(optimized.name(), "GlobalLimitExec");
+        assert_eq!(optimized.children().len(), 1);
+        assert_eq!(optimized.children()[0].name(), "empty-exec");
+
+        Ok(())
+    }
+}
