feat: add a custom ExprPlanner to support colon operator

examples/cli.rs

@@ -1,12 +1,13 @@
 use anyhow::{Context, Result};
 use arrow::array::{ArrayRef, RecordBatch, StringArray};
 use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::execution::FunctionRegistry;
 use datafusion::logical_expr::ScalarUDF;
 use datafusion::prelude::*;
 use datafusion_variant::{
-    CastToVariantUdf, IsVariantNullUdf, JsonToVariantUdf, VariantGetUdf, VariantListConstruct,
-    VariantListInsert, VariantObjectConstruct, VariantObjectInsert, VariantObjectKeys,
-    VariantPretty, VariantToJsonUdf,
+    CastToVariantUdf, IsVariantNullUdf, JsonToVariantUdf, VariantExprPlanner, VariantGetUdf,
+    VariantListConstruct, VariantListInsert, VariantObjectConstruct, VariantObjectInsert,
+    VariantObjectKeys, VariantPretty, VariantToJsonUdf,
 };
 use flate2::read::GzDecoder;
 use rustyline::error::ReadlineError;
@@ -97,7 +98,7 @@ async fn main() -> Result<()> {
     let ctx = {
         let setup_start = Instant::now();
 
-        let ctx = SessionContext::new();
+        let mut ctx = SessionContext::new();
         let schema = Schema::new(vec![Field::new("json_data", DataType::Utf8, false)]);
         let string_array: ArrayRef = Arc::new(StringArray::from(json_strings));
         let batch = RecordBatch::try_new(Arc::new(schema), vec![string_array])?;
@@ -119,6 +120,7 @@ async fn main() -> Result<()> {
         ctx.register_udf(ScalarUDF::new_from_impl(VariantListInsert::default()));
         ctx.register_udf(ScalarUDF::new_from_impl(VariantObjectInsert::default()));
         ctx.register_udf(ScalarUDF::new_from_impl(VariantObjectKeys::default()));
+        ctx.register_expr_planner(Arc::new(VariantExprPlanner))?;
 
         let setup_duration = setup_start.elapsed();
         println!(

src/expr_planner.rs

@@ -0,0 +1,36 @@
+use std::sync::Arc;
+
+use datafusion::{
+    common::DFSchema,
+    error::DataFusionError,
+    logical_expr::{
+        ScalarUDF,
+        expr::ScalarFunction,
+        planner::{ExprPlanner, PlannerResult, RawBinaryExpr},
+    },
+    prelude::Expr,
+    sql::sqlparser::ast::BinaryOperator,
+};
+
+use crate::VariantGetUdf;
+
+#[derive(Debug)]
+pub struct VariantExprPlanner;
+
+impl ExprPlanner for VariantExprPlanner {
+    fn plan_binary_op(
+        &self,
+        expr: RawBinaryExpr,
+        _schema: &DFSchema,
+    ) -> Result<PlannerResult<RawBinaryExpr>, DataFusionError> {
+        match &expr.op {
+            BinaryOperator::Custom(s) if s == ":" => Ok(PlannerResult::Planned(
+                Expr::ScalarFunction(ScalarFunction::new_udf(
+                    Arc::new(ScalarUDF::new_from_impl(VariantGetUdf::default())),
+                    vec![expr.left, expr.right],
+                )),
+            )),
+            _ => Ok(PlannerResult::Original(expr)),
+        }
+    }
+}

src/lib.rs

@@ -3,6 +3,7 @@
 mod shared;
 
 mod cast_to_variant;
+mod expr_planner;
 mod impl_variant_get;
 mod is_variant_null;
 mod json_to_variant;
@@ -19,6 +20,7 @@ mod variant_pretty;
 mod variant_to_json;
 
 pub use cast_to_variant::*;
+pub use expr_planner::*;
 pub use is_variant_null::*;
 pub use json_to_variant::*;
 pub use variant_get::*;