feat: add resolved target to DmlStatement (to eliminate need for table lookup after deserialization) (apache#14631)

* feat: serialize table source to DML proto

* Update datafusion/core/src/dataframe/mod.rs

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

* remove redundant comment

---------

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: 2 people

datafusion/core/src/dataframe/mod.rs

@@ -30,7 +30,9 @@ use crate::arrow::util::pretty;
 use crate::datasource::file_format::csv::CsvFormatFactory;
 use crate::datasource::file_format::format_as_file_type;
 use crate::datasource::file_format::json::JsonFormatFactory;
-use crate::datasource::{provider_as_source, MemTable, TableProvider};
+use crate::datasource::{
+    provider_as_source, DefaultTableSource, MemTable, TableProvider,
+};
 use crate::error::Result;
 use crate::execution::context::{SessionState, TaskContext};
 use crate::execution::FunctionRegistry;
@@ -62,6 +64,7 @@ use datafusion_functions_aggregate::expr_fn::{
 
 use async_trait::async_trait;
 use datafusion_catalog::Session;
+use datafusion_sql::TableReference;
 
 /// Contains options that control how data is
 /// written out from a DataFrame

datafusion/core/src/datasource/listing/table.rs

@@ -1074,7 +1074,7 @@ mod tests {
     use crate::datasource::file_format::json::JsonFormat;
     #[cfg(feature = "parquet")]
     use crate::datasource::file_format::parquet::ParquetFormat;
-    use crate::datasource::{provider_as_source, MemTable};
+    use crate::datasource::{provider_as_source, DefaultTableSource, MemTable};
     use crate::execution::options::ArrowReadOptions;
     use crate::prelude::*;
     use crate::{
@@ -1904,6 +1904,8 @@ mod tests {
         session_ctx.register_table("source", source_table.clone())?;
         // Convert the source table into a provider so that it can be used in a query
         let source = provider_as_source(source_table);
+        let target = session_ctx.table_provider("t").await?;
+        let target = Arc::new(DefaultTableSource::new(target));
         // Create a table scan logical plan to read from the source table
         let scan_plan = LogicalPlanBuilder::scan("source", source, None)?
             .filter(filter_predicate)?

datafusion/core/src/datasource/memory.rs

@@ -367,7 +367,7 @@ impl DataSink for MemSink {
 mod tests {
 
     use super::*;
-    use crate::datasource::provider_as_source;
+    use crate::datasource::{provider_as_source, DefaultTableSource};
     use crate::physical_plan::collect;
     use crate::prelude::SessionContext;
 
@@ -617,6 +617,7 @@ mod tests {
         // Create and register the initial table with the provided schema and data
         let initial_table = Arc::new(MemTable::try_new(schema.clone(), initial_data)?);
         session_ctx.register_table("t", initial_table.clone())?;
+        let target = Arc::new(DefaultTableSource::new(initial_table.clone()));
         // Create and register the source table with the provided schema and inserted data
         let source_table = Arc::new(MemTable::try_new(schema.clone(), inserted_data)?);
         session_ctx.register_table("source", source_table.clone())?;

datafusion/core/src/physical_planner.rs

@@ -24,7 +24,7 @@ use std::sync::Arc;
 use crate::datasource::file_format::file_type_to_format;
 use crate::datasource::listing::ListingTableUrl;
 use crate::datasource::physical_plan::FileSinkConfig;
-use crate::datasource::source_as_provider;
+use crate::datasource::{source_as_provider, DefaultTableSource};
 use crate::error::{DataFusionError, Result};
 use crate::execution::context::{ExecutionProps, SessionState};
 use crate::logical_expr::utils::generate_sort_key;

datafusion/expr/src/logical_plan/dml.rs

@@ -86,6 +86,39 @@ pub struct DmlStatement {
     /// The schema of the output relation
     pub output_schema: DFSchemaRef,
 }
+impl Eq for DmlStatement {}
+impl Hash for DmlStatement {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.table_name.hash(state);
+        self.target.schema().hash(state);
+        self.op.hash(state);
+        self.input.hash(state);
+        self.output_schema.hash(state);
+    }
+}
+
+impl PartialEq for DmlStatement {
+    fn eq(&self, other: &Self) -> bool {
+        self.table_name == other.table_name
+            && self.target.schema() == other.target.schema()
+            && self.op == other.op
+            && self.input == other.input
+            && self.output_schema == other.output_schema
+    }
+}
+
+impl Debug for DmlStatement {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("DmlStatement")
+            .field("table_name", &self.table_name)
+            .field("target", &"...")
+            .field("target_schema", &self.target.schema())
+            .field("op", &self.op)
+            .field("input", &self.input)
+            .field("output_schema", &self.output_schema)
+            .finish()
+    }
+}
 
 impl Debug for DmlStatement {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {

datafusion/proto/proto/datafusion.proto

@@ -276,7 +276,7 @@ message DmlNode{
   Type dml_type = 1;
   LogicalPlanNode input = 2;
   TableReference table_name = 3;
-  datafusion_common.DfSchema schema = 4;
+  LogicalPlanNode target = 5;
 }
 
 message UnnestNode {

datafusion/proto/src/generated/pbjson.rs

@@ -55,8 +55,8 @@ use datafusion::{
 };
 use datafusion_common::file_options::file_type::FileType;
 use datafusion_common::{
-    context, internal_datafusion_err, internal_err, not_impl_err, DataFusionError,
-    Result, TableReference,
+    context, internal_datafusion_err, internal_err, not_impl_err, plan_err,
+    DataFusionError, Result, TableReference, ToDFSchema,
 };
 use datafusion_expr::{
     dml,
@@ -66,10 +66,10 @@ use datafusion_expr::{
         EmptyRelation, Extension, Join, JoinConstraint, Limit, Prepare, Projection,
         Repartition, Sort, SubqueryAlias, TableScan, Values, Window,
     },
-    DistinctOn, DropView, Expr, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr,
-    WindowUDF,
+    AggregateUDF, ColumnUnnestList, DistinctOn, DmlStatement, DropView, Expr, FetchType,
+    LogicalPlan, LogicalPlanBuilder, RecursiveQuery, ScalarUDF, SkipType, SortExpr,
+    Statement, TableSource, Unnest, WindowUDF,
 };
-use datafusion_expr::{AggregateUDF, DmlStatement, RecursiveQuery, Unnest};
 
 use self::to_proto::{serialize_expr, serialize_exprs};
 use crate::logical_plan::to_proto::serialize_sorts;
@@ -233,6 +233,45 @@ fn from_table_reference(
     Ok(table_ref.clone().try_into()?)
 }
 
+/// Converts [LogicalPlan::TableScan] to [TableSource]
+/// method to be used to deserialize nodes
+/// serialized by [from_table_source]
+fn to_table_source(
+    node: &Option<Box<LogicalPlanNode>>,
+    ctx: &SessionContext,
+    extension_codec: &dyn LogicalExtensionCodec,
+) -> Result<Arc<dyn TableSource>> {
+    if let Some(node) = node {
+        match node.try_into_logical_plan(ctx, extension_codec)? {
+            LogicalPlan::TableScan(TableScan { source, .. }) => Ok(source),
+            _ => plan_err!("expected TableScan node"),
+        }
+    } else {
+        plan_err!("LogicalPlanNode should be provided")
+    }
+}
+
+/// converts [TableSource] to [LogicalPlan::TableScan]
+/// using [LogicalPlan::TableScan] was the best approach to
+/// serialize [TableSource] to [LogicalPlan::TableScan]
+fn from_table_source(
+    table_name: TableReference,
+    target: Arc<dyn TableSource>,
+    extension_codec: &dyn LogicalExtensionCodec,
+) -> Result<LogicalPlanNode> {
+    let projected_schema = target.schema().to_dfschema_ref()?;
+    let r = LogicalPlan::TableScan(TableScan {
+        table_name,
+        source: target,
+        projection: None,
+        projected_schema,
+        filters: vec![],
+        fetch: None,
+    });
+
+    LogicalPlanNode::try_from_logical_plan(&r, extension_codec)
+}
+
 impl AsLogicalPlan for LogicalPlanNode {
     fn try_decode(buf: &[u8]) -> Result<Self>
     where
@@ -922,7 +961,7 @@ impl AsLogicalPlan for LogicalPlanNode {
             LogicalPlanType::Dml(dml_node) => Ok(LogicalPlan::Dml(
                 datafusion::logical_expr::DmlStatement::new(
                     from_table_reference(dml_node.table_name.as_ref(), "DML ")?,
-                    Arc::new(convert_required!(dml_node.schema)?),
+                    to_table_source(&dml_node.target, ctx, extension_codec)?,
                     dml_node.dml_type().into(),
                     Arc::new(into_logical_plan!(dml_node.input, ctx, extension_codec)?),
                 ),
@@ -1656,7 +1695,7 @@ impl AsLogicalPlan for LogicalPlanNode {
             )),
             LogicalPlan::Dml(DmlStatement {
                 table_name,
-                table_schema,
+                target,
                 op,
                 input,
                 ..
@@ -1667,7 +1706,11 @@ impl AsLogicalPlan for LogicalPlanNode {
                 Ok(LogicalPlanNode {
                     logical_plan_type: Some(LogicalPlanType::Dml(Box::new(DmlNode {
                         input: Some(Box::new(input)),
-                        schema: Some(table_schema.try_into()?),
+                        target: Some(Box::new(from_table_source(
+                            table_name.clone(),
+                            Arc::clone(target),
+                            extension_codec,
+                        )?)),
                         table_name: Some(table_name.clone().into()),
                         dml_type: dml_type.into(),
                     }))),

datafusion/proto/src/generated/prost.rs

@@ -1227,8 +1227,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         // Do a table lookup to verify the table exists
         let table_ref = self.object_name_to_table_reference(table_name.clone())?;
         let table_source = self.context_provider.get_table_source(table_ref.clone())?;
-        let schema = (*table_source.schema()).clone();
-        let schema = DFSchema::try_from(schema)?;
+        let schema = table_source.schema().to_dfschema_ref()?;
         let scan = LogicalPlanBuilder::scan(
             object_name_to_string(&table_name),
             Arc::clone(&table_source),
@@ -1242,7 +1241,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             Some(predicate_expr) => {
                 let filter_expr =
                     self.sql_to_expr(predicate_expr, &schema, &mut planner_context)?;
-                let schema = Arc::new(schema.clone());
+                let schema = Arc::new(schema);
                 let mut using_columns = HashSet::new();
                 expr_to_columns(&filter_expr, &mut using_columns)?;
                 let filter_expr = normalize_col_with_schemas_and_ambiguity_check(