feat: integrate new table provider with DataSink

Signed-off-by: Robert Pack <robstar.pack@gmail.com>

Author: roeap

crates/core/src/delta_datafusion/table_provider.rs

@@ -19,7 +19,7 @@ use datafusion::datasource::physical_plan::{
     FileGroup, FileSource, wrap_partition_type_in_dict, wrap_partition_value_in_dict,
 };
 use datafusion::datasource::physical_plan::{FileScanConfigBuilder, ParquetSource};
-use datafusion::datasource::sink::{DataSink, DataSinkExec};
+use datafusion::datasource::sink::DataSinkExec;
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::logical_expr::dml::InsertOp;
@@ -32,7 +32,6 @@ use datafusion::physical_plan::filter_pushdown::{FilterDescription, FilterPushdo
 use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet};
 use datafusion::physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, PlanProperties,
-    stream::RecordBatchStreamAdapter,
 };
 use datafusion::{
     catalog::Session,
@@ -43,9 +42,8 @@ use datafusion::{
     scalar::ScalarValue,
 };
 use delta_kernel::Version;
+use futures::TryStreamExt as _;
 use futures::future::BoxFuture;
-use futures::{StreamExt as _, TryStreamExt as _};
-use itertools::Itertools;
 use object_store::ObjectMeta;
 use serde::{Deserialize, Serialize};
 use url::Url;
@@ -58,184 +56,18 @@ use crate::delta_datafusion::{
     DataFusionMixins as _, LogDataHandler, get_null_of_arrow_type, register_store,
     to_correct_scalar_value,
 };
-use crate::kernel::schema::cast::cast_record_batch;
-use crate::kernel::transaction::{CommitBuilder, PROTOCOL};
-use crate::kernel::{Action, Add, EagerSnapshot, Snapshot};
+use crate::kernel::transaction::PROTOCOL;
+use crate::kernel::{Add, EagerSnapshot, Snapshot};
 use crate::logstore::LogStore;
-use crate::operations::write::writer::{DeltaWriter, WriterConfig};
-use crate::protocol::{DeltaOperation, SaveMode};
-use crate::table::config::TablePropertiesExt;
+use crate::protocol::SaveMode;
 use crate::table::normalize_table_url;
 use crate::{DeltaResult, DeltaTable, DeltaTableError, logstore::LogStoreRef};
 
+mod data_sink;
 pub(crate) mod next;
 
 const PATH_COLUMN: &str = "__delta_rs_path";
 
-/// DataSink implementation for delta lake
-/// This uses DataSinkExec to handle the insert operation
-/// Implements writing streams of RecordBatches to delta.
-#[derive(Debug)]
-pub struct DeltaDataSink {
-    /// The log store
-    log_store: LogStoreRef,
-    /// The snapshot
-    snapshot: EagerSnapshot,
-    /// The save mode
-    save_mode: SaveMode,
-    /// The schema
-    schema: SchemaRef,
-    /// Metrics for monitoring throughput
-    metrics: ExecutionPlanMetricsSet,
-}
-
-/// A [`DataSink`] implementation for writing to Delta Lake.
-///
-/// `DeltaDataSink` is used by [`DataSinkExec`] during query execution to
-/// stream [`RecordBatch`]es into a Delta table. It encapsulates everything
-/// needed to perform an insert/append/overwrite operation, including
-/// transaction log access, snapshot state, and session configuration.
-impl DeltaDataSink {
-    /// Create a new `DeltaDataSink`
-    pub fn new(log_store: LogStoreRef, snapshot: EagerSnapshot, save_mode: SaveMode) -> Self {
-        Self {
-            log_store,
-            schema: snapshot.read_schema(),
-            snapshot,
-            save_mode,
-            metrics: ExecutionPlanMetricsSet::new(),
-        }
-    }
-
-    /// Create a streaming transformed version of the input that converts dictionary columns
-    /// This is used to convert dictionary columns to their native types
-    fn create_converted_stream(
-        &self,
-        input: SendableRecordBatchStream,
-        target_schema: SchemaRef,
-    ) -> SendableRecordBatchStream {
-        use futures::StreamExt;
-
-        let schema_for_closure = Arc::clone(&target_schema);
-        let converted_stream = input.map(move |batch_result| {
-            batch_result.and_then(|batch| {
-                cast_record_batch(&batch, Arc::clone(&schema_for_closure), false, true)
-                    .map_err(|e| DataFusionError::External(Box::new(e)))
-            })
-        });
-
-        Box::pin(RecordBatchStreamAdapter::new(
-            target_schema,
-            converted_stream,
-        ))
-    }
-}
-
-/// Implementation of the `DataSink` trait for `DeltaDataSink`
-/// This is used to write the data to the delta table
-/// It implements the `DataSink` trait and is used by the `DataSinkExec` node
-/// to write the data to the delta table
-#[async_trait::async_trait]
-impl DataSink for DeltaDataSink {
-    fn as_any(&self) -> &dyn Any {
-        self
-    }
-
-    fn metrics(&self) -> Option<MetricsSet> {
-        Some(self.metrics.clone_inner())
-    }
-
-    fn schema(&self) -> &SchemaRef {
-        &self.schema
-    }
-
-    /// Write the data to the delta table
-    /// This is used for insert into operation
-    async fn write_all(
-        &self,
-        data: SendableRecordBatchStream,
-        _context: &Arc<TaskContext>,
-    ) -> datafusion::common::Result<u64> {
-        let target_schema = self.snapshot.input_schema();
-        let table_props = self.snapshot.table_configuration().table_properties();
-
-        let mut stream = self.create_converted_stream(data, target_schema.clone());
-        let partition_columns = self.snapshot.metadata().partition_columns();
-        let object_store = self.log_store.object_store(None);
-        let total_rows_metric = MetricBuilder::new(&self.metrics).counter("total_rows", 0);
-        let config = WriterConfig::new(
-            self.snapshot.read_schema(),
-            partition_columns.clone(),
-            None,
-            Some(table_props.target_file_size().get() as usize),
-            None,
-            table_props.num_indexed_cols(),
-            table_props
-                .data_skipping_stats_columns
-                .as_ref()
-                .map(|c| c.iter().map(|c| c.to_string()).collect_vec()),
-        );
-
-        let mut writer = DeltaWriter::new(object_store, config);
-        let mut total_rows = 0u64;
-
-        while let Some(batch_result) = stream.next().await {
-            let batch = batch_result?;
-            let batch_rows = batch.num_rows() as u64;
-            total_rows += batch_rows;
-            total_rows_metric.add(batch_rows as usize);
-            writer
-                .write(&batch)
-                .await
-                .map_err(|e| DataFusionError::External(Box::new(e)))?;
-        }
-
-        let mut actions = writer
-            .close()
-            .await
-            .map_err(|e| DataFusionError::External(Box::new(e)))?
-            .into_iter()
-            .map(Action::Add)
-            .collect_vec();
-
-        if self.save_mode == SaveMode::Overwrite {
-            actions.extend(
-                self.snapshot
-                    .file_views(&self.log_store, None)
-                    .map_ok(|f| Action::Remove(f.remove_action(true)))
-                    .try_collect::<Vec<_>>()
-                    .await
-                    .map_err(|e| DataFusionError::External(Box::new(e)))?,
-            );
-        };
-
-        let operation = DeltaOperation::Write {
-            mode: self.save_mode,
-            partition_by: if partition_columns.is_empty() {
-                None
-            } else {
-                Some(partition_columns.clone())
-            },
-            predicate: None,
-        };
-
-        CommitBuilder::default()
-            .with_actions(actions)
-            .build(Some(&self.snapshot), self.log_store.clone(), operation)
-            .await
-            .map_err(|e| DataFusionError::External(Box::new(e)))?;
-
-        Ok(total_rows)
-    }
-}
-
-/// Implementation of the `DisplayAs` trait for `DeltaDataSink`
-impl DisplayAs for DeltaDataSink {
-    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> std::fmt::Result {
-        write!(f, "DeltaDataSink")
-    }
-}
-
 #[derive(Debug, Clone)]
 /// Used to specify if additional metadata columns are exposed to the user
 pub struct DeltaScanConfigBuilder {
@@ -1012,7 +844,7 @@ impl TableProvider for DeltaTableProvider {
         };
 
         let data_sink =
-            DeltaDataSink::new(self.log_store.clone(), self.snapshot.clone(), save_mode);
+            data_sink::DeltaDataSink::new(self.log_store.clone(), self.snapshot.clone(), save_mode);
 
         Ok(Arc::new(DataSinkExec::new(
             input,