dml: store destination as table source

This patch adds a field `dst` into struct `DmlStatement`. No reason to
do it in the other way than for a table scan, where we store the table
source. Whereas storing a dst adds the ability to analyze the table
for which DML operation is performed. It is useful for analyzer/optimizer
passes writing.

Author: askalt

datafusion-examples/examples/dataframe_output.rs

@@ -15,9 +15,11 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use datafusion::datasource::provider_as_source;
 use datafusion::{dataframe::DataFrameWriteOptions, prelude::*};
 use datafusion_common::config::CsvOptions;
 use datafusion_common::{parsers::CompressionTypeVariant, DataFusionError};
+use datafusion_sql::TableReference;
 
 /// This example demonstrates the various methods to write out a DataFrame to local storage.
 /// See datafusion-examples/examples/external_dependency/dataframe-to-s3.rs for an example
@@ -32,20 +34,26 @@ async fn main() -> Result<(), DataFusionError> {
     df = df.with_column_renamed("column1", "tablecol1").unwrap();
 
     ctx.sql(
-        "create external table 
+        "create external table
     test(tablecol1 varchar)
-    stored as parquet 
+    stored as parquet
     location './datafusion-examples/test_table/'",
     )
     .await?
     .collect()
     .await?;
 
+    let dst = ctx.table_provider(TableReference::bare("test")).await?;
+
     // This is equivalent to INSERT INTO test VALUES ('a'), ('b'), ('c').
     // The behavior of write_table depends on the TableProvider's implementation
     // of the insert_into method.
     df.clone()
-        .write_table("test", DataFrameWriteOptions::new())
+        .write_table(
+            "test",
+            provider_as_source(dst),
+            DataFrameWriteOptions::new(),
+        )
         .await?;
 
     df.clone()

datafusion/core/src/dataframe/mod.rs

@@ -52,7 +52,7 @@ use datafusion_common::config::{CsvOptions, JsonOptions};
 use datafusion_common::{
     plan_err, Column, DFSchema, DataFusionError, ParamValues, SchemaError, UnnestOptions,
 };
-use datafusion_expr::{case, is_null, lit, SortExpr};
+use datafusion_expr::{case, is_null, lit, SortExpr, TableSource};
 use datafusion_expr::{
     utils::COUNT_STAR_EXPANSION, TableProviderFilterPushDown, UNNAMED_TABLE,
 };
@@ -1296,13 +1296,13 @@ impl DataFrame {
     pub async fn write_table(
         self,
         table_name: &str,
+        dst: Arc<dyn TableSource>,
         write_options: DataFrameWriteOptions,
     ) -> Result<Vec<RecordBatch>, DataFusionError> {
-        let arrow_schema = Schema::from(self.schema());
         let plan = LogicalPlanBuilder::insert_into(
             self.plan,
             table_name.to_owned(),
-            &arrow_schema,
+            dst,
             write_options.overwrite,
         )?
         .build()?;

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

@@ -1091,6 +1091,7 @@ mod tests {
     use datafusion_physical_expr::PhysicalSortExpr;
     use datafusion_physical_plan::ExecutionPlanProperties;
 
+    use datafusion_sql::TableReference;
     use tempfile::TempDir;
 
     #[tokio::test]
@@ -1890,6 +1891,11 @@ mod tests {
             _ => panic!("Unrecognized file extension {file_type_ext}"),
         }
 
+        let dst_table = session_ctx
+            .table_provider(TableReference::bare("t"))
+            .await
+            .unwrap();
+
         // Create and register the source table with the provided schema and inserted data
         let source_table = Arc::new(MemTable::try_new(
             schema.clone(),
@@ -1905,8 +1911,13 @@ mod tests {
         // Since logical plan contains a filter, increasing parallelism is helpful.
         // Therefore, we will have 8 partitions in the final plan.
         // Create an insert plan to insert the source data into the initial table
-        let insert_into_table =
-            LogicalPlanBuilder::insert_into(scan_plan, "t", &schema, false)?.build()?;
+        let insert_into_table = LogicalPlanBuilder::insert_into(
+            scan_plan,
+            "t",
+            provider_as_source(dst_table),
+            false,
+        )?
+        .build()?;
         // Create a physical plan from the insert plan
         let plan = session_ctx
             .state()

datafusion/core/src/datasource/memory.rs

@@ -625,8 +625,13 @@ mod tests {
         // Create a table scan logical plan to read from the source table
         let scan_plan = LogicalPlanBuilder::scan("source", source, None)?.build()?;
         // Create an insert plan to insert the source data into the initial table
-        let insert_into_table =
-            LogicalPlanBuilder::insert_into(scan_plan, "t", &schema, false)?.build()?;
+        let insert_into_table = LogicalPlanBuilder::insert_into(
+            scan_plan,
+            "t",
+            provider_as_source(Arc::clone(&initial_table) as Arc<_>),
+            false,
+        )?
+        .build()?;
         // Create a physical plan from the insert plan
         let plan = session_ctx
             .state()

datafusion/core/src/datasource/physical_plan/parquet/mod.rs

@@ -792,6 +792,7 @@ mod tests {
     use crate::datasource::file_format::test_util::scan_format;
     use crate::datasource::listing::{FileRange, ListingOptions};
     use crate::datasource::object_store::ObjectStoreUrl;
+    use crate::datasource::provider_as_source;
     use crate::execution::context::SessionState;
     use crate::physical_plan::displayable;
     use crate::prelude::{ParquetReadOptions, SessionConfig, SessionContext};
@@ -816,6 +817,7 @@ mod tests {
     use datafusion_physical_plan::ExecutionPlanProperties;
 
     use chrono::{TimeZone, Utc};
+    use datafusion_sql::TableReference;
     use futures::StreamExt;
     use object_store::local::LocalFileSystem;
     use object_store::path::Path;
@@ -2024,8 +2026,16 @@ mod tests {
         )
         .await
         .unwrap();
-        df.write_table("my_table", DataFrameWriteOptions::new())
-            .await?;
+        let dst = ctx
+            .table_provider(TableReference::bare("my_table"))
+            .await
+            .unwrap();
+        df.write_table(
+            "my_table",
+            provider_as_source(dst),
+            DataFrameWriteOptions::new(),
+        )
+        .await?;
 
         // create a new context and verify that the results were saved to a partitioned parquet file
         let ctx = SessionContext::new();

datafusion/core/src/physical_planner.rs

@@ -68,8 +68,7 @@ use arrow_array::builder::StringBuilder;
 use arrow_array::RecordBatch;
 use datafusion_common::display::ToStringifiedPlan;
 use datafusion_common::{
-    exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema,
-    ScalarValue,
+    internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema, ScalarValue,
 };
 use datafusion_expr::dml::CopyTo;
 use datafusion_expr::expr::{
@@ -540,36 +539,26 @@ impl DefaultPhysicalPlanner {
                     .await?
             }
             LogicalPlan::Dml(DmlStatement {
-                table_name,
                 op: WriteOp::InsertInto,
+                dst,
                 ..
             }) => {
-                let name = table_name.table();
-                let schema = session_state.schema_for_ref(table_name.clone())?;
-                if let Some(provider) = schema.table(name).await? {
-                    let input_exec = children.one()?;
-                    provider
-                        .insert_into(session_state, input_exec, false)
-                        .await?
-                } else {
-                    return exec_err!("Table '{table_name}' does not exist");
-                }
+                let input_exec = children.one()?;
+                let provider = source_as_provider(dst)?;
+                provider
+                    .insert_into(session_state, input_exec, false)
+                    .await?
             }
             LogicalPlan::Dml(DmlStatement {
-                table_name,
+                dst,
                 op: WriteOp::InsertOverwrite,
                 ..
             }) => {
-                let name = table_name.table();
-                let schema = session_state.schema_for_ref(table_name.clone())?;
-                if let Some(provider) = schema.table(name).await? {
-                    let input_exec = children.one()?;
-                    provider
-                        .insert_into(session_state, input_exec, true)
-                        .await?
-                } else {
-                    return exec_err!("Table '{table_name}' does not exist");
-                }
+                let input_exec = children.one()?;
+                let provider = source_as_provider(dst)?;
+                provider
+                    .insert_into(session_state, input_exec, true)
+                    .await?
             }
             LogicalPlan::Window(Window {
                 input, window_expr, ..

datafusion/expr/src/logical_plan/builder.rs

@@ -302,11 +302,9 @@ impl LogicalPlanBuilder {
     pub fn insert_into(
         input: LogicalPlan,
         table_name: impl Into<TableReference>,
-        table_schema: &Schema,
+        dst: Arc<dyn TableSource>,
         overwrite: bool,
     ) -> Result<Self> {
-        let table_schema = table_schema.clone().to_dfschema_ref()?;
-
         let op = if overwrite {
             WriteOp::InsertOverwrite
         } else {
@@ -315,7 +313,7 @@ impl LogicalPlanBuilder {
 
         Ok(Self::new(LogicalPlan::Dml(DmlStatement::new(
             table_name.into(),
-            table_schema,
+            dst,
             op,
             Arc::new(input),
         ))))

datafusion/expr/src/logical_plan/dml.rs

@@ -24,7 +24,7 @@ use arrow::datatypes::{DataType, Field, Schema};
 use datafusion_common::file_options::file_type::FileType;
 use datafusion_common::{DFSchemaRef, TableReference};
 
-use crate::LogicalPlan;
+use crate::{LogicalPlan, TableSource};
 
 /// Operator that copies the contents of a database to file(s)
 #[derive(Clone)]
@@ -73,12 +73,12 @@ impl Hash for CopyTo {
 
 /// The operator that modifies the content of a database (adapted from
 /// substrait WriteRel)
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[derive(Clone)]
 pub struct DmlStatement {
     /// The table name
     pub table_name: TableReference,
-    /// The schema of the table (must align with Rel input)
-    pub table_schema: DFSchemaRef,
+    /// The operation destination
+    pub dst: Arc<dyn TableSource>,
     /// The type of operation to perform
     pub op: WriteOp,
     /// The relation that determines the tuples to add/remove/modify the schema must match with table_schema
@@ -87,17 +87,49 @@ pub struct DmlStatement {
     pub output_schema: DFSchemaRef,
 }
 
+impl Debug for DmlStatement {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        f.debug_struct("TableScan")
+            .field("table_name", &self.table_name)
+            .field("source", &"...")
+            .field("op", &self.op)
+            .field("input", &self.input)
+            .field("output_schema", &self.output_schema)
+            .finish_non_exhaustive()
+    }
+}
+
+impl PartialEq for DmlStatement {
+    fn eq(&self, other: &Self) -> bool {
+        self.table_name == other.table_name
+            && self.op == other.op
+            && self.input == other.input
+            && self.output_schema == other.output_schema
+    }
+}
+
+impl Eq for DmlStatement {}
+
+impl Hash for DmlStatement {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.table_name.hash(state);
+        self.op.hash(state);
+        self.input.hash(state);
+        self.output_schema.hash(state);
+    }
+}
+
 impl DmlStatement {
     /// Creates a new DML statement with the output schema set to a single `count` column.
     pub fn new(
         table_name: TableReference,
-        table_schema: DFSchemaRef,
+        dst: Arc<dyn TableSource>,
         op: WriteOp,
         input: Arc<LogicalPlan>,
     ) -> Self {
         Self {
             table_name,
-            table_schema,
+            dst,
             op,
             input,
 

datafusion/expr/src/logical_plan/plan.rs

@@ -810,15 +810,15 @@ impl LogicalPlan {
             }
             LogicalPlan::Dml(DmlStatement {
                 table_name,
-                table_schema,
+                dst,
                 op,
                 ..
             }) => {
                 self.assert_no_expressions(expr)?;
                 let input = self.only_input(inputs)?;
                 Ok(LogicalPlan::Dml(DmlStatement::new(
                     table_name.clone(),
-                    Arc::clone(table_schema),
+                    Arc::clone(dst),
                     op.clone(),
                     Arc::new(input),
                 )))
@@ -1987,7 +1987,7 @@ impl LogicalPlan {
                             .map(|i| &input_columns[*i])
                             .collect::<Vec<&Column>>();
                         // get items from input_columns indexed by list_col_indices
-                        write!(f, "Unnest: lists[{}] structs[{}]", 
+                        write!(f, "Unnest: lists[{}] structs[{}]",
                         expr_vec_fmt!(list_type_columns),
                         expr_vec_fmt!(struct_type_columns))
                     }

datafusion/expr/src/logical_plan/tree_node.rs

@@ -248,14 +248,14 @@ impl TreeNode for LogicalPlan {
             }),
             LogicalPlan::Dml(DmlStatement {
                 table_name,
-                table_schema,
+                dst,
                 op,
                 input,
                 output_schema,
             }) => rewrite_arc(input, f)?.update_data(|input| {
                 LogicalPlan::Dml(DmlStatement {
                     table_name,
-                    table_schema,
+                    dst,
                     op,
                     input,
                     output_schema,