chore: refactor array_repeat

native/core/src/execution/planner.rs

@@ -3905,134 +3905,6 @@ mod tests {
         });
     }
 
-    #[test]
-    fn test_array_repeat() {
-        let session_ctx = SessionContext::new();
-        let task_ctx = session_ctx.task_ctx();
-        let planner = PhysicalPlanner::new(Arc::from(session_ctx), 0);
-
-        // Mock scan operator with 3 INT32 columns
-        let op_scan = Operator {
-            plan_id: 0,
-            children: vec![],
-            op_struct: Some(OpStruct::Scan(spark_operator::Scan {
-                fields: vec![
-                    spark_expression::DataType {
-                        type_id: 3, // Int32
-                        type_info: None,
-                    },
-                    spark_expression::DataType {
-                        type_id: 3, // Int32
-                        type_info: None,
-                    },
-                    spark_expression::DataType {
-                        type_id: 3, // Int32
-                        type_info: None,
-                    },
-                ],
-                source: "".to_string(),
-                arrow_ffi_safe: false,
-            })),
-        };
-
-        // Mock expression to read a INT32 column with position 0
-        let array_col = spark_expression::Expr {
-            expr_struct: Some(Bound(spark_expression::BoundReference {
-                index: 0,
-                datatype: Some(spark_expression::DataType {
-                    type_id: 3,
-                    type_info: None,
-                }),
-            })),
-        };
-
-        // Mock expression to read a INT32 column with position 1
-        let array_col_1 = spark_expression::Expr {
-            expr_struct: Some(Bound(spark_expression::BoundReference {
-                index: 1,
-                datatype: Some(spark_expression::DataType {
-                    type_id: 3,
-                    type_info: None,
-                }),
-            })),
-        };
-
-        // Make a projection operator with array_repeat(array_col, array_col_1)
-        let projection = Operator {
-            children: vec![op_scan],
-            plan_id: 0,
-            op_struct: Some(OpStruct::Projection(spark_operator::Projection {
-                project_list: vec![spark_expression::Expr {
-                    expr_struct: Some(ExprStruct::ScalarFunc(spark_expression::ScalarFunc {
-                        func: "array_repeat".to_string(),
-                        args: vec![array_col, array_col_1],
-                        return_type: None,
-                        fail_on_error: false,
-                    })),
-                }],
-            })),
-        };
-
-        // Create a physical plan
-        let (mut scans, datafusion_plan) =
-            planner.create_plan(&projection, &mut vec![], 1).unwrap();
-
-        // Start executing the plan in a separate thread
-        // The plan waits for incoming batches and emitting result as input comes
-        let mut stream = datafusion_plan.native_plan.execute(0, task_ctx).unwrap();
-
-        let runtime = tokio::runtime::Runtime::new().unwrap();
-        // create async channel
-        let (tx, mut rx) = mpsc::channel(1);
-
-        // Send data as input to the plan being executed in a separate thread
-        runtime.spawn(async move {
-            // create data batch
-            // 0, 1, 2
-            // 3, 4, 5
-            // 6, null, null
-            let a = Int32Array::from(vec![Some(0), Some(3), Some(6)]);
-            let b = Int32Array::from(vec![Some(1), Some(4), None]);
-            let c = Int32Array::from(vec![Some(2), Some(5), None]);
-            let input_batch1 = InputBatch::Batch(vec![Arc::new(a), Arc::new(b), Arc::new(c)], 3);
-            let input_batch2 = InputBatch::EOF;
-
-            let batches = vec![input_batch1, input_batch2];
-
-            for batch in batches.into_iter() {
-                tx.send(batch).await.unwrap();
-            }
-        });
-
-        // Wait for the plan to finish executing and assert the result
-        runtime.block_on(async move {
-            loop {
-                let batch = rx.recv().await.unwrap();
-                scans[0].set_input_batch(batch);
-                match poll!(stream.next()) {
-                    Poll::Ready(Some(batch)) => {
-                        assert!(batch.is_ok(), "got error {}", batch.unwrap_err());
-                        let batch = batch.unwrap();
-                        let expected = [
-                            "+--------------+",
-                            "| col_0        |",
-                            "+--------------+",
-                            "| [0]          |",
-                            "| [3, 3, 3, 3] |",
-                            "|              |",
-                            "+--------------+",
-                        ];
-                        assert_batches_eq!(expected, &[batch]);
-                    }
-                    Poll::Ready(None) => {
-                        break;
-                    }
-                    _ => {}
-                }
-            }
-        });
-    }
-
     /// Executes a `test_data_query` SQL query
     /// and saves the result into a temp folder using parquet format
     /// Read the file back to the memory using a custom schema

native/spark-expr/src/array_funcs/mod.rs

@@ -16,13 +16,11 @@
 // under the License.
 
 mod array_insert;
-mod array_repeat;
 mod get_array_struct_fields;
 mod list_extract;
 mod size;
 
 pub use array_insert::ArrayInsert;
-pub use array_repeat::spark_array_repeat;
 pub use get_array_struct_fields::GetArrayStructFields;
 pub use list_extract::ListExtract;
 pub use size::{spark_size, SparkSizeFunc};

native/spark-expr/src/comet_scalar_funcs.rs

@@ -20,9 +20,9 @@ use crate::math_funcs::abs::abs;
 use crate::math_funcs::checked_arithmetic::{checked_add, checked_div, checked_mul, checked_sub};
 use crate::math_funcs::modulo_expr::spark_modulo;
 use crate::{
-    spark_array_repeat, spark_ceil, spark_decimal_div, spark_decimal_integral_div, spark_floor,
-    spark_isnan, spark_lpad, spark_make_decimal, spark_read_side_padding, spark_round, spark_rpad,
-    spark_unhex, spark_unscaled_value, EvalMode, SparkBitwiseCount, SparkContains, SparkDateDiff,
+    spark_ceil, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_isnan,
+    spark_lpad, spark_make_decimal, spark_read_side_padding, spark_round, spark_rpad, spark_unhex,
+    spark_unscaled_value, EvalMode, SparkBitwiseCount, SparkContains, SparkDateDiff,
     SparkDateTrunc, SparkMakeDate, SparkSizeFunc, SparkStringSpace,
 };
 use arrow::datatypes::DataType;
@@ -169,10 +169,6 @@ pub fn create_comet_physical_fun_with_eval_mode(
             let func = Arc::new(spark_isnan);
             make_comet_scalar_udf!("isnan", func, without data_type)
         }
-        "array_repeat" => {
-            let func = Arc::new(spark_array_repeat);
-            make_comet_scalar_udf!("array_repeat", func, without data_type)
-        }
         "spark_modulo" => {
             let func = Arc::new(spark_modulo);
             make_comet_scalar_udf!("spark_modulo", func, without data_type, fail_on_error)