feat: Implement native operator registry (#2875)

Author: andygrove

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

@@ -23,7 +23,7 @@ macro_rules! arithmetic_expr_builder {
     ($builder_name:ident, $expr_type:ident, $operator:expr) => {
         pub struct $builder_name;
 
-        impl $crate::execution::planner::traits::ExpressionBuilder for $builder_name {
+        impl $crate::execution::planner::expression_registry::ExpressionBuilder for $builder_name {
             fn build(
                 &self,
                 spark_expr: &datafusion_comet_proto::spark_expression::Expr,
@@ -61,7 +61,8 @@ use crate::execution::{
     expressions::extract_expr,
     operators::ExecutionError,
     planner::{
-        from_protobuf_eval_mode, traits::ExpressionBuilder, BinaryExprOptions, PhysicalPlanner,
+        expression_registry::ExpressionBuilder, from_protobuf_eval_mode, BinaryExprOptions,
+        PhysicalPlanner,
     },
 };
 

native/core/src/execution/operators/mod.rs

@@ -31,6 +31,7 @@ pub use expand::ExpandExec;
 mod iceberg_scan;
 mod parquet_writer;
 pub use parquet_writer::ParquetWriterExec;
+pub mod projection;
 mod scan;
 
 /// Error returned during executing operators.

native/core/src/execution/operators/projection.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.
+
+//! Projection operator builder
+
+use std::sync::Arc;
+
+use datafusion::physical_plan::projection::ProjectionExec;
+use datafusion_comet_proto::spark_operator::Operator;
+use jni::objects::GlobalRef;
+
+use crate::{
+    execution::{
+        operators::{ExecutionError, ScanExec},
+        planner::{operator_registry::OperatorBuilder, PhysicalPlanner},
+        spark_plan::SparkPlan,
+    },
+    extract_op,
+};
+
+/// Builder for Projection operators
+pub struct ProjectionBuilder;
+
+impl OperatorBuilder for ProjectionBuilder {
+    fn build(
+        &self,
+        spark_plan: &Operator,
+        inputs: &mut Vec<Arc<GlobalRef>>,
+        partition_count: usize,
+        planner: &PhysicalPlanner,
+    ) -> Result<(Vec<ScanExec>, Arc<SparkPlan>), ExecutionError> {
+        let project = extract_op!(spark_plan, Projection);
+        let children = &spark_plan.children;
+
+        assert_eq!(children.len(), 1);
+        let (scans, child) = planner.create_plan(&children[0], inputs, partition_count)?;
+
+        // Create projection expressions
+        let exprs: Result<Vec<_>, _> = project
+            .project_list
+            .iter()
+            .enumerate()
+            .map(|(idx, expr)| {
+                planner
+                    .create_expr(expr, child.schema())
+                    .map(|r| (r, format!("col_{idx}")))
+            })
+            .collect();
+
+        let projection = Arc::new(ProjectionExec::try_new(
+            exprs?,
+            Arc::clone(&child.native_plan),
+        )?);
+
+        Ok((
+            scans,
+            Arc::new(SparkPlan::new(spark_plan.plan_id, projection, vec![child])),
+        ))
+    }
+}

native/core/src/execution/planner.rs

@@ -18,7 +18,8 @@
 //! Converts Spark physical plan to DataFusion physical plan
 
 pub mod expression_registry;
-pub mod traits;
+pub mod macros;
+pub mod operator_registry;
 
 use crate::execution::operators::IcebergScanExec;
 use crate::{
@@ -27,6 +28,7 @@ use crate::{
         expressions::subquery::Subquery,
         operators::{ExecutionError, ExpandExec, ParquetWriterExec, ScanExec},
         planner::expression_registry::ExpressionRegistry,
+        planner::operator_registry::OperatorRegistry,
         serde::to_arrow_datatype,
         shuffle::ShuffleWriterExec,
     },
@@ -861,29 +863,19 @@ impl PhysicalPlanner {
         inputs: &mut Vec<Arc<GlobalRef>>,
         partition_count: usize,
     ) -> Result<(Vec<ScanExec>, Arc<SparkPlan>), ExecutionError> {
+        // Try to use the modular registry first - this automatically handles any registered operator types
+        if OperatorRegistry::global().can_handle(spark_plan) {
+            return OperatorRegistry::global().create_plan(
+                spark_plan,
+                inputs,
+                partition_count,
+                self,
+            );
+        }
+
+        // Fall back to the original monolithic match for other operators
         let children = &spark_plan.children;
         match spark_plan.op_struct.as_ref().unwrap() {
-            OpStruct::Projection(project) => {
-                assert_eq!(children.len(), 1);
-                let (scans, child) = self.create_plan(&children[0], inputs, partition_count)?;
-                let exprs: PhyExprResult = project
-                    .project_list
-                    .iter()
-                    .enumerate()
-                    .map(|(idx, expr)| {
-                        self.create_expr(expr, child.schema())
-                            .map(|r| (r, format!("col_{idx}")))
-                    })
-                    .collect();
-                let projection = Arc::new(ProjectionExec::try_new(
-                    exprs?,
-                    Arc::clone(&child.native_plan),
-                )?);
-                Ok((
-                    scans,
-                    Arc::new(SparkPlan::new(spark_plan.plan_id, projection, vec![child])),
-                ))
-            }
             OpStruct::Filter(filter) => {
                 assert_eq!(children.len(), 1);
                 let (scans, child) = self.create_plan(&children[0], inputs, partition_count)?;
@@ -1634,6 +1626,10 @@ impl PhysicalPlanner {
                     Arc::new(SparkPlan::new(spark_plan.plan_id, window_agg, vec![child])),
                 ))
             }
+            _ => Err(GeneralError(format!(
+                "Unsupported or unregistered operator type: {:?}",
+                spark_plan.op_struct
+            ))),
         }
     }
 

native/core/src/execution/planner/expression_registry.rs

@@ -25,7 +25,90 @@ use datafusion::physical_expr::PhysicalExpr;
 use datafusion_comet_proto::spark_expression::{expr::ExprStruct, Expr};
 
 use crate::execution::operators::ExecutionError;
-use crate::execution::planner::traits::{ExpressionBuilder, ExpressionType};
+
+/// Trait for building physical expressions from Spark protobuf expressions
+pub trait ExpressionBuilder: Send + Sync {
+    /// Build a DataFusion physical expression from a Spark protobuf expression
+    fn build(
+        &self,
+        spark_expr: &Expr,
+        input_schema: SchemaRef,
+        planner: &super::PhysicalPlanner,
+    ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError>;
+}
+
+/// Enum to identify different expression types for registry dispatch
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum ExpressionType {
+    // Arithmetic expressions
+    Add,
+    Subtract,
+    Multiply,
+    Divide,
+    IntegralDivide,
+    Remainder,
+    UnaryMinus,
+
+    // Comparison expressions
+    Eq,
+    Neq,
+    Lt,
+    LtEq,
+    Gt,
+    GtEq,
+    EqNullSafe,
+    NeqNullSafe,
+
+    // Logical expressions
+    And,
+    Or,
+    Not,
+
+    // Null checks
+    IsNull,
+    IsNotNull,
+
+    // Bitwise operations
+    BitwiseAnd,
+    BitwiseOr,
+    BitwiseXor,
+    BitwiseShiftLeft,
+    BitwiseShiftRight,
+
+    // Other expressions
+    Bound,
+    Unbound,
+    Literal,
+    Cast,
+    CaseWhen,
+    In,
+    If,
+    Substring,
+    Like,
+    Rlike,
+    CheckOverflow,
+    ScalarFunc,
+    NormalizeNanAndZero,
+    Subquery,
+    BloomFilterMightContain,
+    CreateNamedStruct,
+    GetStructField,
+    ToJson,
+    ToPrettyString,
+    ListExtract,
+    GetArrayStructFields,
+    ArrayInsert,
+    Rand,
+    Randn,
+    SparkPartitionId,
+    MonotonicallyIncreasingId,
+
+    // Time functions
+    Hour,
+    Minute,
+    Second,
+    TruncTimestamp,
+}
 
 /// Registry for expression builders
 pub struct ExpressionRegistry {