play with extension statement

datafusion-examples/examples/sql_parsing.rs

@@ -18,12 +18,16 @@
 use std::fmt::Display;
 
 use datafusion::error::Result;
-use datafusion_sql::{parser::{CopyToSource, CopyToStatement, DFParser, Statement}, sqlparser::{keywords::Keyword, parser::ParserError, tokenizer::Token}};
+use datafusion_sql::{
+    parser::{CopyToSource, CopyToStatement, DFParser, Statement},
+    sqlparser::{keywords::Keyword, parser::ParserError, tokenizer::Token},
+};
 
 /// This example demonstrates how to use the DFParser to parse a statement in a custom way
 #[tokio::main]
 async fn main() -> Result<()> {
-    let mut my_parser = MyParser::new("COPY source_table TO 'file.fasta' STORED AS FASTA")?;
+    let mut my_parser =
+        MyParser::new("COPY source_table TO 'file.fasta' STORED AS FASTA")?;
 
     let my_statement = my_parser.parse_statement()?;
 
@@ -72,7 +76,6 @@ impl MyParser<'_> {
             }
         }
     }
-
 }
 
 enum MyStatement {
@@ -121,6 +124,10 @@ impl From<CopyToStatement> for MyCopyToStatement {
 
 impl Display for MyCopyToStatement {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "COPY {} TO '{}' STORED AS FASTA", self.source, self.target)
+        write!(
+            f,
+            "COPY {} TO '{}' STORED AS FASTA",
+            self.source, self.target
+        )
     }
 }

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

@@ -1761,6 +1761,9 @@ impl SessionState {
                 DFStatement::Explain(explain) => {
                     visit_statement(&explain.statement, visitor)
                 }
+                DFStatement::Extension(ext) => {
+                    ext.visit(visitor);
+                }
             }
         }
 

datafusion/sql/src/parser.rs

@@ -18,10 +18,14 @@
 //! [`DFParser`]: DataFusion SQL Parser based on [`sqlparser`]
 
 use std::collections::{HashMap, VecDeque};
-use std::fmt;
+use std::fmt::{self, Debug, Display};
+use std::ops::ControlFlow;
 use std::str::FromStr;
+use std::sync::Arc;
 
 use datafusion_common::parsers::CompressionTypeVariant;
+use datafusion_common::DataFusionError;
+use datafusion_expr::LogicalPlan;
 use sqlparser::{
     ast::{
         ColumnDef, ColumnOptionDef, ObjectName, OrderByExpr, Query,
@@ -230,6 +234,51 @@ impl fmt::Display for CreateExternalTable {
     }
 }
 
+
+pub trait UserDefinedStatement: Debug + Display {
+    fn to_logical_plan(&self) -> Result<LogicalPlan, DataFusionError>;
+
+    fn visit(&self, visitor: &mut RelationVisitor) -> ControlFlow<sqlparser::ast::Visitor>;
+}
+
+// Implement Debug, Clone, and PartialEq for ExtensionStatement manually if needed.
+#[derive(Debug, Clone)]
+pub struct ExtensionStatement {
+    pub statement: Arc<dyn UserDefinedStatement>,
+}
+
+impl ExtensionStatement {
+    pub fn new(statement: Arc<dyn UserDefinedStatement>) -> Self {
+        Self { statement }
+    }
+
+    pub fn to_logical_plan(&self) -> Result<LogicalPlan, DataFusionError> {
+        self.statement.to_logical_plan()
+    }
+}
+
+impl Display for ExtensionStatement {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.statement)
+    }
+}
+
+// Implement PartialEq manually to handle comparison logic since dyn UserDefinedStatement
+// cannot directly use PartialEq due to object safety.
+impl PartialEq for ExtensionStatement {
+    fn eq(&self, other: &Self) -> bool {
+        // Implement your logic to compare two ExtensionStatement instances.
+        // This might involve downcasting Arc<dyn UserDefinedStatement> if you need
+        // to compare the statements and if such a comparison is meaningful.
+        // Note: Actual comparison logic would depend on your specific use case.
+        Arc::ptr_eq(&self.statement, &other.statement)
+    }
+}
+
+// Eq requires PartialEq but has no additional requirements.
+impl Eq for ExtensionStatement {}
+
+
 /// DataFusion SQL Statement.
 ///
 /// This can either be a [`Statement`] from [`sqlparser`] from a
@@ -247,6 +296,8 @@ pub enum Statement {
     CopyTo(CopyToStatement),
     /// EXPLAIN for extensions
     Explain(ExplainStatement),
+    /// User defined extension statement
+    Extension(ExtensionStatement),
 }
 
 impl fmt::Display for Statement {
@@ -256,6 +307,7 @@ impl fmt::Display for Statement {
             Statement::CreateExternalTable(stmt) => write!(f, "{stmt}"),
             Statement::CopyTo(stmt) => write!(f, "{stmt}"),
             Statement::Explain(stmt) => write!(f, "{stmt}"),
+            Statement::Extension(stmt) => write!(f, "{stmt}"),
         }
     }
 }

datafusion/sql/src/statement.rs

@@ -21,8 +21,7 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use crate::parser::{
-    CopyToSource, CopyToStatement, CreateExternalTable, DFParser, ExplainStatement,
-    LexOrdering, Statement as DFStatement,
+    CopyToSource, CopyToStatement, CreateExternalTable, DFParser, ExplainStatement, ExtensionStatement, LexOrdering, Statement as DFStatement
 };
 use crate::planner::{
     object_name_to_qualifier, ContextProvider, PlannerContext, SqlToRel,
@@ -151,9 +150,15 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 analyze,
                 statement,
             }) => self.explain_to_plan(verbose, analyze, *statement),
+            DFStatement::Extension(ext) => self.extension_to_plan(ext),
         }
     }
 
+    /// Generate a logical plan from an Extension SQL statement
+    pub fn extension_to_plan(&self, ext: ExtensionStatement) -> Result<LogicalPlan> {
+        ext.to_logical_plan()
+    }
+
     /// Generate a logical plan from an SQL statement
     pub fn sql_statement_to_plan(&self, statement: Statement) -> Result<LogicalPlan> {
         self.sql_statement_to_plan_with_context_impl(