2851: feat: Add expression registry to native planner

[martin-augment]

2851: To review by AI

[coderabbitai]

Walkthrough

This pull request introduces a modular, registry-based architecture for expression handling in the planner. A new arithmetic expression module provides builders for addition, subtraction, multiplication, division, integral division, remainder, and unary minus operations. Concurrently, core traits (ExpressionBuilder, OperatorBuilder) and enums (ExpressionType, OperatorType) are defined in a new traits module. An ExpressionRegistry maps expression types to their corresponding builders. The planner's public API is expanded to expose helper functions and the registry itself, shifting from monolithic match-based expression handling to a pluggable registry-driven approach while preserving fallback paths for non-registered cases.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch pr-2851-2025-12-05-11-28-42

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[augmentcode]

Review completed. No suggestions at this time.

Comment augment review to trigger a new review at any time.

[claude]

Code Review: Expression Registry Refactoring

Overview

This PR introduces a modular expression registry pattern to the native planner, refactoring arithmetic expressions out of the monolithic create_expr method. This is a positive architectural improvement that will make the codebase more maintainable.

---

✅ Strengths

1. Good Architecture: The registry pattern with ExpressionBuilder trait is a clean abstraction that will scale well as more expressions are migrated.

2. Clear Separation of Concerns: Moving arithmetic expressions into their own module (arithmetic.rs) improves code organization.

3. Extensibility: The framework is well-designed for future expansion with placeholder comments for comparison, string, and temporal expressions.

4. Consistent with Existing Patterns: The code follows existing conventions (e.g., unwrap() usage patterns match the codebase).

---

🐛 Critical Issues

1. Copy/Clone Issue in with_exec_id Method (planner.rs:184)

pub fn with_exec_id(self, exec_context_id: i64) -> Self {
    Self {
        exec_context_id,
        partition: self.partition,
        session_ctx: Arc::clone(&self.session_ctx),
        expression_registry: self.expression_registry,  // ❌ PROBLEM HERE
    }
}

Problem: ExpressionRegistry contains HashMap<ExpressionType, Box<dyn ExpressionBuilder>>, which is neither Copy nor Clone. This line attempts to move the registry, but since self is consumed, this will likely fail to compile.

Fix: Add Clone implementation to ExpressionRegistry or use Arc<ExpressionRegistry> in PhysicalPlanner.

Recommended Solution:

// In expression_registry.rs, add Clone implementation
impl Clone for ExpressionRegistry {
    fn clone(&self) -> Self {
        Self::new()  // Re-create registry (builders are stateless)
    }
}

// Then in planner.rs:
expression_registry: self.expression_registry.clone(),

---

⚠️ Issues & Concerns

2. Potential Panic from unwrap() on return_type (arithmetic.rs:201)

expr.return_type
    .as_ref()
    .map(crate::execution::serde::to_arrow_datatype)
    .unwrap(),  // ❌ Will panic if return_type is None

Issue: If return_type is None, this will panic. While existing code uses similar patterns, it's risky.

Recommendation: Either:

• Add a comment explaining why return_type is guaranteed to be Some
• Return a proper error: `ok_or_else(|| ExecutionError::GeneralError("Missing return_type".into()))?"

3. Inconsistent Error Handling in Registry

The registry's get_expression_type method returns errors for unregistered types, but the caller in create_expr (line 254-256) checks can_handle first, making the error path unreachable. This creates dead code.

Suggestion: Simplify the logic:

if let Some(builder) = self.expression_registry.get_builder(spark_expr) {
    return builder.build(spark_expr, input_schema, self);
}
// Fall back to existing match

4. Large ExpressionType Enum (traits.rs:56-124)

The enum lists ALL expression types (even unimplemented ones). This creates maintenance burden and potential confusion.

Recommendation: Only include types that are actually registered in the first phase. Add others as they're implemented.

---

📊 Performance Considerations

5. HashMap Lookup Overhead

Every expression now goes through:

1. can_handle check (HashMap lookup)
2. create_expr call (another HashMap lookup)

Impact: Minimal for most queries, but adds ~2 HashMap lookups per expression node.

Optimization (optional): Combine into single lookup:

if let Some(builder) = self.builders.get(&expr_type) {
    builder.build(spark_expr, input_schema, planner)
} else {
    Err(ExecutionError::GeneralError(...))
}

6. Registry Recreation

ExpressionRegistry::new() recreates and re-registers all builders for each PhysicalPlanner. Consider using Arc<ExpressionRegistry> for better performance.

---

🧪 Test Coverage

7. No Unit Tests for New Code ⚠️

None of the new modules (arithmetic.rs, expression_registry.rs, traits.rs) have unit tests.

Required Tests:

• Registry registration and lookup
• Each ExpressionBuilder implementation
• Error handling for missing fields
• Edge cases (eval modes, null handling)

Example:

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_add_builder_with_ansi_mode() {
        // Test Add expression with ANSI eval mode
    }
    
    #[test]
    fn test_registry_handles_arithmetic_expressions() {
        let registry = ExpressionRegistry::new();
        // Verify all arithmetic types are registered
    }
}

---

🔒 Security Concerns

No security issues identified. The refactoring doesn't introduce new attack surfaces.

---

📝 Minor Issues

8. Missing Documentation: Public methods and structs should have doc comments explaining their purpose and usage.

9. Trait Visibility: OperatorBuilder trait is marked #[allow(dead_code)] but is public. Either implement it or make it private.

10. Import Optimization: traits.rs imports GlobalRef which is only used in the dead OperatorBuilder trait.

---

📋 Summary

Category	Status
Architecture	✅ Excellent
Code Quality	⚠️ Good with issues
Bugs	❌ 1 critical (Copy/Clone)
Performance	⚠️ Minor overhead
Security	✅ No concerns
Tests	❌ Missing

Verdict: Needs Changes before merging. Fix the Copy/Clone issue and add basic test coverage.

---

🎯 Action Items

Must Fix (Blocking):

1. ✅ Fix expression_registry copy/clone issue in with_exec_id
2. ✅ Add at least basic integration tests

Should Fix (Recommended):
3. Handle return_type.unwrap() more safely
4. Remove unimplemented expression types from enum
5. Add doc comments to public APIs

Nice to Have:
6. Consider using Arc<ExpressionRegistry> for performance
7. Optimize registry lookup to single HashMap access

Great work on the refactoring! The architecture is solid and this sets up a good foundation for future improvements. 🚀

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

native/core/src/execution/expressions/arithmetic.rs (1)

195-206: Consider defensive handling for missing return_type.

The .unwrap() on line 201 will panic if return_type is None. While this follows existing patterns in the codebase (e.g., planner.rs line 873), consider using .ok_or_else() for a more descriptive error:

-            let result = create_modulo_expr(
-                left,
-                right,
-                expr.return_type
-                    .as_ref()
-                    .map(crate::execution::serde::to_arrow_datatype)
-                    .unwrap(),
+            let return_type = expr.return_type
+                .as_ref()
+                .map(crate::execution::serde::to_arrow_datatype)
+                .ok_or_else(|| ExecutionError::GeneralError(
+                    "Remainder expression missing return_type".to_string()
+                ))?;
+            let result = create_modulo_expr(
+                left,
+                right,
+                return_type,

This is optional since other builders use the same pattern via create_binary_expr.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between e40f20d and c0e4bf6.

📒 Files selected for processing (5)

• native/core/src/execution/expressions/arithmetic.rs (1 hunks)
• native/core/src/execution/expressions/mod.rs (1 hunks)
• native/core/src/execution/planner.rs (10 hunks)
• native/core/src/execution/planner/expression_registry.rs (1 hunks)
• native/core/src/execution/planner/traits.rs (1 hunks)

🧰 Additional context used

🧠 Learnings (1)

📚 Learning: 2025-11-04T14:26:48.750Z

Learnt from: martin-augment
Repo: martin-augment/datafusion-comet PR: 7
File: native/spark-expr/src/math_funcs/abs.rs:201-302
Timestamp: 2025-11-04T14:26:48.750Z
Learning: In the abs function in native/spark-expr/src/math_funcs/abs.rs (Rust), NULL values for signed integers (Int8, Int16, Int32, Int64) and decimals (Decimal128, Decimal256) should return the argument as-is (e.g., ColumnarValue::Scalar(ScalarValue::Int8(None))) rather than panicking on unwrap().

Applied to files:

• native/core/src/execution/expressions/arithmetic.rs

🧬 Code graph analysis (3)

native/core/src/execution/planner/expression_registry.rs (2)

native/core/src/execution/planner.rs (11)

• expr (647-651)
• expr (701-705)
• expr (1797-1801)
• expr (2259-2263)
• expr (2355-2355)
• expr (2390-2394)
• expr (2528-2528)
• new (170-177)
• new (2560-2572)
• create_expr (246-798)
• default (164-166)

native/spark-expr/src/nondetermenistic_funcs/rand.rs (1)

• None (97-97)

native/core/src/execution/planner/traits.rs (1)

native/core/src/execution/expressions/arithmetic.rs (7)

• build (39-60)
• build (67-88)
• build (95-116)
• build (123-144)
• build (151-175)
• build (182-212)
• build (219-234)

native/core/src/execution/planner.rs (1)

native/core/src/execution/planner/expression_registry.rs (1)

• new (37-44)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)

• GitHub Check: Cursor Bugbot
• GitHub Check: claude-review

🔇 Additional comments (17)

native/core/src/execution/expressions/mod.rs (1)

20-20: LGTM!

The new module export follows the existing pattern and correctly exposes the arithmetic expression builders for use by the registry.

native/core/src/execution/expressions/arithmetic.rs (2)

38-61: LGTM! Clean implementation of the builder pattern.

The AddBuilder correctly delegates to planner.create_binary_expr with the appropriate operator. The pattern is consistent across all arithmetic builders.

---

215-235: LGTM!

The UnaryMinusBuilder correctly uses create_negate_expr with the fail_on_error flag from the protobuf expression.

native/core/src/execution/planner/expression_registry.rs (3)

31-53: Well-designed registry pattern enabling incremental migration.

The registry correctly checks can_handle before attempting dispatch, and the fallback path in planner.rs handles unregistered expressions. This enables gradual migration of expressions to the modular system.

---

74-106: Good modular structure with clear extension points.

The TODO comments and separate registration methods (register_arithmetic_expressions, etc.) provide clear guidance for future contributors to add more expression categories.

---

188-192: LGTM!

The Default implementation correctly delegates to new(), ensuring consistent initialization.

native/core/src/execution/planner.rs (7)

20-21: LGTM!

Clean module structure exposing the new registry and traits modules.

---

147-150: Public visibility is appropriate for cross-module usage.

Making BinaryExprOptions and its field public is necessary for the arithmetic builders to use create_binary_expr_with_options.

---

160-161: LGTM!

The expression_registry field is correctly added to PhysicalPlanner.

---

170-186: Correct initialization and preservation of registry.

The registry is properly initialized in new() and preserved in with_exec_id(), ensuring consistent state across planner instances.

---

251-259: Clean registry-first dispatch with fallback.

The pattern correctly checks can_handle before dispatching to the registry, falling back to the original match-based handling for unregistered expressions. This enables incremental migration while maintaining backward compatibility.

---

826-827: Public visibility enables builder access.

Making create_binary_expr public allows arithmetic builders to leverage this shared implementation.

---

2644-2650: Public visibility with complete enum mapping.

Making from_protobuf_eval_mode public enables the arithmetic builders to convert eval modes. The function correctly maps all three Spark eval modes (Legacy, Try, Ansi).

native/core/src/execution/planner/traits.rs (4)

30-39: LGTM! Well-defined trait for expression building.

The ExpressionBuilder trait has appropriate Send + Sync bounds for concurrent use and a clear contract for building physical expressions from Spark protobuf.

---

41-52: Good forward scaffolding for operator registry.

The OperatorBuilder trait is prepared for a future operator registry similar to the expression registry. The #[allow(dead_code)] annotation correctly acknowledges this is scaffolding.

---

54-125: Comprehensive expression type enumeration.

The ExpressionType enum covers all current expression types with clear categorization (arithmetic, comparison, logical, etc.). This provides type-safe keys for registry dispatch.

---

127-145: Forward scaffolding for operator registry.

The OperatorType enum mirrors the operator types in the planner's create_plan match, preparing for a future modular operator system.

[martin-augment]

Copy/Clone Issue in with_exec_id Method (planner.rs:184)

pub fn with_exec_id(self, exec_context_id: i64) -> Self {
    Self {
        exec_context_id,
        partition: self.partition,
        session_ctx: Arc::clone(&self.session_ctx),
        expression_registry: self.expression_registry,  // ❌ PROBLEM HERE
    }
}

Problem: ExpressionRegistry contains HashMap<ExpressionType, Box<dyn ExpressionBuilder>>, which is neither Copy nor Clone. This line attempts to move the registry, but since self is consumed, this will likely fail to compile.

Fix: Add Clone implementation to ExpressionRegistry or use Arc<ExpressionRegistry> in PhysicalPlanner.

value:annoying; category:bug; feedback: The Claude AI reviewer is not correct! The ownership of this struct is moved, i.e. it is neither copied nor cloned. The build passes and it confirms that there is no such issue here.

[martin-augment]

2. Potential Panic from unwrap() on return_type (arithmetic.rs:201)

expr.return_type
    .as_ref()
    .map(crate::execution::serde::to_arrow_datatype)
    .unwrap(),  // ❌ Will panic if return_type is None

Issue: If return_type is None, this will panic. While existing code uses similar patterns, it's risky.

Recommendation: Either:

• Add a comment explaining why return_type is guaranteed to be Some
• Return a proper error: `ok_or_else(|| ExecutionError::GeneralError("Missing return_type".into()))?"

value:useful; category:bug; feedback: The Claude AI reviewer is correct! Panicking in non-test code is bad practice because it leads to application crash. It is better to return an Err and let the caller decide whether to crash or to handle it somehow.

[github-actions]

Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.