Custom operator support

datafusion/core/src/execution/context/mod.rs

@@ -60,7 +60,7 @@ use datafusion_execution::registry::SerializerRegistry;
 use datafusion_expr::{
     expr_rewriter::FunctionRewrite,
     logical_plan::{DdlStatement, Statement},
-    Expr, UserDefinedLogicalNode, WindowUDF,
+    Expr, ParseCustomOperator, UserDefinedLogicalNode, WindowUDF,
 };
 
 // backwards compatibility
@@ -1390,6 +1390,19 @@ impl FunctionRegistry for SessionContext {
     ) -> Result<()> {
         self.state.write().register_function_rewrite(rewrite)
     }
+
+    fn register_parse_custom_operator(
+        &mut self,
+        parse_custom_operator: Arc<dyn ParseCustomOperator>,
+    ) -> Result<()> {
+        self.state
+            .write()
+            .register_parse_custom_operator(parse_custom_operator)
+    }
+
+    fn parse_custom_operators(&self) -> Vec<Arc<dyn ParseCustomOperator>> {
+        self.state.read().parse_custom_operators()
+    }
 }
 
 /// Create a new task context instance from SessionContext

datafusion/core/src/execution/session_state.rs

@@ -56,8 +56,8 @@ use datafusion_expr::registry::{FunctionRegistry, SerializerRegistry};
 use datafusion_expr::simplify::SimplifyInfo;
 use datafusion_expr::var_provider::{is_system_variables, VarType};
 use datafusion_expr::{
-    AggregateUDF, Explain, Expr, ExprSchemable, LogicalPlan, ScalarUDF, TableSource,
-    WindowUDF,
+    AggregateUDF, Explain, Expr, ExprSchemable, LogicalPlan, ParseCustomOperator,
+    ScalarUDF, TableSource, WindowUDF,
 };
 use datafusion_optimizer::simplify_expressions::ExprSimplifier;
 use datafusion_optimizer::{
@@ -91,6 +91,8 @@ pub struct SessionState {
     session_id: String,
     /// Responsible for analyzing and rewrite a logical plan before optimization
     analyzer: Analyzer,
+    /// Provides support for parsing custom SQL operators, e.g. `->>` or `?`
+    parse_custom_operators: Vec<Arc<dyn ParseCustomOperator>>,
     /// Responsible for optimizing a logical plan
     optimizer: Optimizer,
     /// Responsible for optimizing a physical execution plan
@@ -221,6 +223,7 @@ impl SessionState {
         let mut new_self = SessionState {
             session_id,
             analyzer: Analyzer::new(),
+            parse_custom_operators: vec![],
             optimizer: Optimizer::new(),
             physical_optimizers: PhysicalOptimizer::new(),
             query_planner: Arc::new(DefaultQueryPlanner {}),
@@ -543,7 +546,7 @@ impl SessionState {
     /// Convert an AST Statement into a LogicalPlan
     pub async fn statement_to_plan(
         &self,
-        statement: datafusion_sql::parser::Statement,
+        statement: Statement,
     ) -> datafusion_common::Result<LogicalPlan> {
         let references = self.resolve_table_references(&statement)?;
 
@@ -573,6 +576,7 @@ impl SessionState {
         ParserOptions {
             parse_float_as_decimal: sql_parser_options.parse_float_as_decimal,
             enable_ident_normalization: sql_parser_options.enable_ident_normalization,
+            parse_custom_operator: self.parse_custom_operators(),
         }
     }
 
@@ -1074,6 +1078,18 @@ impl FunctionRegistry for SessionState {
         self.analyzer.add_function_rewrite(rewrite);
         Ok(())
     }
+
+    fn register_parse_custom_operator(
+        &mut self,
+        parse_custom_operator: Arc<dyn ParseCustomOperator>,
+    ) -> datafusion_common::Result<()> {
+        self.parse_custom_operators.push(parse_custom_operator);
+        Ok(())
+    }
+
+    fn parse_custom_operators(&self) -> Vec<Arc<dyn ParseCustomOperator>> {
+        self.parse_custom_operators.clone()
+    }
 }
 
 impl OptimizerConfig for SessionState {

datafusion/core/tests/user_defined/mod.rs

@@ -15,6 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+/// Tests for custom operator parsing and substitution
+mod user_defined_custom_operators;
 /// Tests for user defined Scalar functions
 mod user_defined_scalar_functions;
 

datafusion/core/tests/user_defined/user_defined_custom_operators.rs

@@ -0,0 +1,169 @@
+// 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 arrow_array::RecordBatch;
+use std::sync::Arc;
+
+use datafusion::arrow::datatypes::DataType;
+use datafusion::common::config::ConfigOptions;
+use datafusion::common::tree_node::Transformed;
+use datafusion::common::{assert_batches_eq, DFSchema};
+use datafusion::error::Result;
+use datafusion::execution::FunctionRegistry;
+use datafusion::logical_expr::expr_rewriter::FunctionRewrite;
+use datafusion::logical_expr::{
+    CustomOperator, Operator, ParseCustomOperator, WrapCustomOperator,
+};
+use datafusion::prelude::*;
+use datafusion::sql::sqlparser::ast::BinaryOperator;
+
+#[derive(Debug)]
+enum MyCustomOperator {
+    Arrow,
+    LongArrow,
+}
+
+impl std::fmt::Display for MyCustomOperator {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            MyCustomOperator::Arrow => write!(f, "->"),
+            MyCustomOperator::LongArrow => write!(f, "->>"),
+        }
+    }
+}
+
+impl CustomOperator for MyCustomOperator {
+    fn binary_signature(
+        &self,
+        lhs: &DataType,
+        rhs: &DataType,
+    ) -> Result<(DataType, DataType, DataType)> {
+        Ok((lhs.clone(), rhs.clone(), lhs.clone()))
+    }
+
+    fn op_to_sql(&self) -> Result<BinaryOperator> {
+        match self {
+            MyCustomOperator::Arrow => Ok(BinaryOperator::Arrow),
+            MyCustomOperator::LongArrow => Ok(BinaryOperator::LongArrow),
+        }
+    }
+
+    fn name(&self) -> &'static str {
+        match self {
+            MyCustomOperator::Arrow => "Arrow",
+            MyCustomOperator::LongArrow => "LongArrow",
+        }
+    }
+}
+
+impl TryFrom<&str> for MyCustomOperator {
+    type Error = ();
+
+    fn try_from(value: &str) -> std::result::Result<Self, Self::Error> {
+        match value {
+            "Arrow" => Ok(MyCustomOperator::Arrow),
+            "LongArrow" => Ok(MyCustomOperator::LongArrow),
+            _ => Err(()),
+        }
+    }
+}
+
+#[derive(Debug)]
+struct CustomOperatorParser;
+
+impl ParseCustomOperator for CustomOperatorParser {
+    fn name(&self) -> &str {
+        "CustomOperatorParser"
+    }
+
+    fn op_from_ast(&self, op: &BinaryOperator) -> Result<Option<Operator>> {
+        match op {
+            BinaryOperator::Arrow => Ok(Some(MyCustomOperator::Arrow.into())),
+            BinaryOperator::LongArrow => Ok(Some(MyCustomOperator::LongArrow.into())),
+            _ => Ok(None),
+        }
+    }
+
+    fn op_from_name(&self, raw_op: &str) -> Result<Option<Operator>> {
+        if let Ok(op) = MyCustomOperator::try_from(raw_op) {
+            Ok(Some(op.into()))
+        } else {
+            Ok(None)
+        }
+    }
+}
+
+impl FunctionRewrite for CustomOperatorParser {
+    fn name(&self) -> &str {
+        "CustomOperatorParser"
+    }
+
+    fn rewrite(
+        &self,
+        expr: Expr,
+        _schema: &DFSchema,
+        _config: &ConfigOptions,
+    ) -> Result<Transformed<Expr>> {
+        if let Expr::BinaryExpr(bin_expr) = &expr {
+            if let Operator::Custom(WrapCustomOperator(op)) = &bin_expr.op {
+                if let Ok(pg_op) = MyCustomOperator::try_from(op.name()) {
+                    // return BinaryExpr with a different operator
+                    let mut bin_expr = bin_expr.clone();
+                    bin_expr.op = match pg_op {
+                        MyCustomOperator::Arrow => Operator::StringConcat,
+                        MyCustomOperator::LongArrow => Operator::Plus,
+                    };
+                    return Ok(Transformed::yes(Expr::BinaryExpr(bin_expr)));
+                }
+            }
+        }
+        Ok(Transformed::no(expr))
+    }
+}
+
+async fn plan_and_collect(sql: &str) -> Result<Vec<RecordBatch>> {
+    let mut ctx = SessionContext::new();
+    ctx.register_function_rewrite(Arc::new(CustomOperatorParser))?;
+    ctx.register_parse_custom_operator(Arc::new(CustomOperatorParser))?;
+    ctx.sql(sql).await?.collect().await
+}
+
+#[tokio::test]
+async fn test_custom_operators_arrow() {
+    let actual = plan_and_collect("select 'foo'->'bar';").await.unwrap();
+    let expected = [
+        "+----------------------------+",
+        "| Utf8(\"foo\") -> Utf8(\"bar\") |",
+        "+----------------------------+",
+        "| foobar                     |",
+        "+----------------------------+",
+    ];
+    assert_batches_eq!(&expected, &actual);
+}
+
+#[tokio::test]
+async fn test_custom_operators_long_arrow() {
+    let actual = plan_and_collect("select 1->>2;").await.unwrap();
+    let expected = [
+        "+-----------------------+",
+        "| Int64(1) ->> Int64(2) |",
+        "+-----------------------+",
+        "| 3                     |",
+        "+-----------------------+",
+    ];
+    assert_batches_eq!(&expected, &actual);
+}

datafusion/expr/src/lib.rs

@@ -48,6 +48,7 @@ pub mod function;
 pub mod groups_accumulator;
 pub mod interval_arithmetic;
 pub mod logical_plan;
+pub mod parse_custom_operator;
 pub mod registry;
 pub mod simplify;
 pub mod sort_properties;
@@ -76,7 +77,8 @@ pub use function::{
 pub use groups_accumulator::{EmitTo, GroupsAccumulator};
 pub use literal::{lit, lit_timestamp_nano, Literal, TimestampLiteral};
 pub use logical_plan::*;
-pub use operator::Operator;
+pub use operator::{CustomOperator, Operator, WrapCustomOperator};
+pub use parse_custom_operator::ParseCustomOperator;
 pub use partition_evaluator::PartitionEvaluator;
 pub use signature::{
     ArrayFunctionSignature, Signature, TypeSignature, Volatility, TIMEZONE_WILDCARD,