delta-io · aersam · Mar 7, 2024 · Mar 8, 2024 · Mar 8, 2024 · Mar 8, 2024
@@ -71,10 +71,13 @@ pub async fn convert_tpcds_web_returns(input_path: String, table_path: String) -
         .await
         .unwrap();
 
+    let tbl = table.collect().await.unwrap();
+    let _schema = tbl[0].schema().clone();
+
     DeltaOps::try_from_uri(table_path)
         .await
         .unwrap()
-        .write(table.collect().await.unwrap())
+        .write(tbl.into())
         .with_partition_columns(vec!["wr_returned_date_sk"])
         .await
         .unwrap();
@@ -552,7 +555,7 @@ async fn main() {
             ]));
 
             let batch = RecordBatch::try_new(
-                schema,
+                schema.clone(),
                 vec![
                     Arc::new(StringArray::from(group_ids)),
                     Arc::new(StringArray::from(name)),
@@ -566,7 +569,7 @@ async fn main() {
             DeltaOps::try_from_uri(output)
                 .await
                 .unwrap()
-                .write(vec![batch])
+                .write(batch.into())
                 .with_save_mode(SaveMode::Append)
                 .await
                 .unwrap();

@@ -40,6 +40,7 @@ datafusion-common = { workspace = true, optional = true }
 datafusion-proto = { workspace = true, optional = true }
 datafusion-sql = { workspace = true, optional = true }
 datafusion-physical-expr = { workspace = true, optional = true }
+async-channel = { version = "2.2", optional = true }
 datafusion-functions = { workspace = true, optional = true }
 
 # serde
@@ -67,7 +68,6 @@ tokio = { workspace = true, features = [
     "fs",
     "parking_lot",
 ] }
-
 # other deps (these should be organized and pulled into workspace.dependencies as necessary)
 cfg-if = "1"
 dashmap = "5"
@@ -124,6 +124,7 @@ datafusion = [
     "datafusion-sql",
     "datafusion-functions",
     "sqlparser",
+    "async-channel",
 ]
 datafusion-ext = ["datafusion"]
 json = ["parquet/json"]

@@ -1686,6 +1686,7 @@ impl From<Column> for DeltaColumn {
 
 #[cfg(test)]
 mod tests {
+    use crate::operations::write::WriteData;
     use crate::writer::test_utils::get_delta_schema;
     use arrow::array::StructArray;
     use arrow::datatypes::{DataType, Field, Schema};
@@ -1976,7 +1977,7 @@ mod tests {
         .unwrap();
         // write some data
         let table = crate::DeltaOps(table)
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .with_save_mode(crate::protocol::SaveMode::Append)
             .await
             .unwrap();
@@ -2040,7 +2041,7 @@ mod tests {
         .unwrap();
         // write some data
         let table = crate::DeltaOps::new_in_memory()
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .with_save_mode(crate::protocol::SaveMode::Append)
             .await
             .unwrap();

@@ -212,6 +212,7 @@ mod tests {
     use arrow_schema::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
     use datafusion_expr::{col, lit};
 
+    use crate::operations::write::WriteData;
     use crate::writer::test_utils::{create_bare_table, get_arrow_schema, get_record_batch};
     use crate::{DeltaOps, DeltaResult, DeltaTable};
 
@@ -244,7 +245,7 @@ mod tests {
     async fn add_constraint_with_invalid_data() -> DeltaResult<()> {
         let batch = get_record_batch(None, false);
         let write = DeltaOps(create_bare_table())
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .await?;
         let table = DeltaOps(write);
 
@@ -260,7 +261,7 @@ mod tests {
     async fn add_valid_constraint() -> DeltaResult<()> {
         let batch = get_record_batch(None, false);
         let write = DeltaOps(create_bare_table())
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .await?;
         let table = DeltaOps(write);
 
@@ -285,7 +286,7 @@ mod tests {
         // Add constraint by providing a datafusion expression.
         let batch = get_record_batch(None, false);
         let write = DeltaOps(create_bare_table())
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .await?;
         let table = DeltaOps(write);
 
@@ -328,7 +329,7 @@ mod tests {
         )
         .unwrap();
 
-        let table = DeltaOps::new_in_memory().write(vec![batch]).await.unwrap();
+        let table = DeltaOps::new_in_memory().write(batch.into()).await.unwrap();
 
         let mut table = DeltaOps(table)
             .add_constraint()
@@ -351,7 +352,7 @@ mod tests {
     async fn add_conflicting_named_constraint() -> DeltaResult<()> {
         let batch = get_record_batch(None, false);
         let write = DeltaOps(create_bare_table())
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .await?;
         let table = DeltaOps(write);
 
@@ -373,7 +374,7 @@ mod tests {
     async fn write_data_that_violates_constraint() -> DeltaResult<()> {
         let batch = get_record_batch(None, false);
         let write = DeltaOps(create_bare_table())
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .await?;
 
         let table = DeltaOps(write)
@@ -387,7 +388,7 @@ mod tests {
             Arc::new(StringArray::from(vec!["2021-02-02"])),
         ];
         let batch = RecordBatch::try_new(get_arrow_schema(&None), invalid_values)?;
-        let err = table.write(vec![batch]).await;
+        let err = table.write(batch.into()).await;
         assert!(err.is_err());
         Ok(())
     }
@@ -396,11 +397,11 @@ mod tests {
     async fn write_data_that_does_not_violate_constraint() -> DeltaResult<()> {
         let batch = get_record_batch(None, false);
         let write = DeltaOps(create_bare_table())
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .await?;
         let table = DeltaOps(write);
 
-        let err = table.write(vec![batch]).await;
+        let err = table.write(batch.into()).await;
 
         assert!(err.is_ok());
         Ok(())

@@ -328,6 +328,7 @@ impl std::future::IntoFuture for DeleteBuilder {
 
 #[cfg(test)]
 mod tests {
+    use crate::operations::write::WriteData;
     use crate::operations::DeltaOps;
     use crate::protocol::*;
     use crate::writer::test_utils::datafusion::get_data;
@@ -396,7 +397,7 @@ mod tests {
         .unwrap();
         // write some data
         let table = DeltaOps(table)
-            .write(vec![batch.clone()])
+            .write(batch.into())
             .with_save_mode(SaveMode::Append)
             .await
             .unwrap();
@@ -458,7 +459,7 @@ mod tests {
 
         // write some data
         let table = DeltaOps(table)
-            .write(vec![batch])
+            .write(batch.into())
             .with_save_mode(SaveMode::Append)
             .await
             .unwrap();
@@ -482,7 +483,7 @@ mod tests {
 
         // write some data
         let table = DeltaOps(table)
-            .write(vec![batch])
+            .write(batch.into())
             .with_save_mode(SaveMode::Append)
             .await
             .unwrap();
@@ -549,7 +550,7 @@ mod tests {
             )
             .unwrap();
 
-            DeltaOps::new_in_memory().write(vec![batch]).await.unwrap()
+            DeltaOps::new_in_memory().write(batch.into()).await.unwrap()
         }
 
         // Validate behaviour of greater than
@@ -638,7 +639,7 @@ mod tests {
 
         // write some data
         let table = DeltaOps(table)
-            .write(vec![batch])
+            .write(batch.into())
             .with_save_mode(SaveMode::Append)
             .await
             .unwrap();
@@ -696,7 +697,7 @@ mod tests {
 
         // write some data
         let table = DeltaOps(table)
-            .write(vec![batch])
+            .write(batch.into())
             .with_save_mode(SaveMode::Append)
             .await
             .unwrap();
@@ -765,7 +766,10 @@ mod tests {
         ];
         let batches = vec![RecordBatch::try_new(schema.clone(), data).unwrap()];
 
-        let table = DeltaOps::new_in_memory().write(batches).await.unwrap();
+        let table = DeltaOps::new_in_memory()
+            .write(batches.into())
+            .await
+            .unwrap();
 
         let (table, _metrics) = DeltaOps(table)
             .delete()

@@ -99,6 +99,7 @@ impl std::future::IntoFuture for DropConstraintBuilder {
 #[cfg(feature = "datafusion")]
 #[cfg(test)]
 mod tests {
+    use crate::operations::write::WriteData;
     use crate::writer::test_utils::{create_bare_table, get_record_batch};
     use crate::{DeltaOps, DeltaResult, DeltaTable};
 
@@ -121,7 +122,7 @@ mod tests {
     async fn drop_valid_constraint() -> DeltaResult<()> {
         let batch = get_record_batch(None, false);
         let write = DeltaOps(create_bare_table())
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .await?;
         let table = DeltaOps(write);
 
@@ -145,7 +146,7 @@ mod tests {
     async fn drop_invalid_constraint_not_existing() -> DeltaResult<()> {
         let batch = get_record_batch(None, false);
         let write = DeltaOps(create_bare_table())
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .await?;
 
         let table = DeltaOps(write)
@@ -161,7 +162,7 @@ mod tests {
     async fn drop_invalid_constraint_ignore() -> DeltaResult<()> {
         let batch = get_record_batch(None, false);
         let write = DeltaOps(create_bare_table())
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .await?;
 
         let version = write.version();

@@ -82,6 +82,7 @@ impl std::future::IntoFuture for LoadBuilder {
 
 #[cfg(test)]
 mod tests {
+    use crate::operations::write::WriteData;
     use crate::operations::{collect_sendable_stream, DeltaOps};
     use crate::writer::test_utils::{get_record_batch, TestResult};
     use crate::DeltaTableBuilder;
@@ -115,7 +116,9 @@ mod tests {
     #[tokio::test]
     async fn test_write_load() -> TestResult {
         let batch = get_record_batch(None, false);
-        let table = DeltaOps::new_in_memory().write(vec![batch.clone()]).await?;
+        let table = DeltaOps::new_in_memory()
+            .write(batch.clone().into())
+            .await?;
 
         let (_table, stream) = DeltaOps(table).load().await?;
         let data = collect_sendable_stream(stream).await?;
@@ -146,7 +149,9 @@ mod tests {
     #[tokio::test]
     async fn test_load_with_columns() -> TestResult {
         let batch = get_record_batch(None, false);
-        let table = DeltaOps::new_in_memory().write(vec![batch.clone()]).await?;
+        let table = DeltaOps::new_in_memory()
+            .write(batch.clone().into())
+            .await?;
 
         let (_table, stream) = DeltaOps(table).load().with_columns(["id", "value"]).await?;
         let data = collect_sendable_stream(stream).await?;

@@ -1553,6 +1553,7 @@ mod tests {
     use crate::kernel::StructField;
     use crate::operations::merge::generalize_filter;
     use crate::operations::merge::try_construct_early_filter;
+    use crate::operations::write::WriteData;
     use crate::operations::DeltaOps;
     use crate::protocol::*;
     use crate::writer::test_utils::datafusion::get_data;
@@ -1635,7 +1636,7 @@ mod tests {
         .unwrap();
         // write some data
         DeltaOps(table)
-            .write(vec![batch.clone()])
+            .write(batch.clone().into())
             .with_save_mode(SaveMode::Append)
             .await
             .unwrap()
@@ -2924,7 +2925,7 @@ mod tests {
         .unwrap();
 
         let table = DeltaOps(table)
-            .write(vec![batch.clone()])
+            .write(batch.into())
             .with_save_mode(SaveMode::Append)
             .await
             .unwrap();
@@ -3038,7 +3039,7 @@ mod tests {
         .unwrap();
 
         let table = DeltaOps(table)
-            .write(vec![batch1, batch2])
+            .write((vec![batch1, batch2]).into())
             .with_write_batch_size(2)
             .with_save_mode(SaveMode::Append)
             .await
@@ -3149,7 +3150,7 @@ mod tests {
         .unwrap();
 
         let table = DeltaOps(table)
-            .write(vec![batch.clone()])
+            .write(batch.into())
             .with_save_mode(SaveMode::Append)
             .await
             .unwrap();

@@ -33,8 +33,7 @@ use self::{
 };
 #[cfg(feature = "datafusion")]
 pub use ::datafusion::physical_plan::common::collect as collect_sendable_stream;
-#[cfg(feature = "datafusion")]
-use arrow::record_batch::RecordBatch;
+
 use optimize::OptimizeBuilder;
 use restore::RestoreBuilder;
 
@@ -137,8 +136,8 @@ impl DeltaOps {
     /// Write data to Delta table
     #[cfg(feature = "datafusion")]
     #[must_use]
-    pub fn write(self, batches: impl IntoIterator<Item = RecordBatch>) -> WriteBuilder {
-        WriteBuilder::new(self.0.log_store, self.0.state).with_input_batches(batches)
+    pub fn write(self, data: write::WriteData) -> WriteBuilder {
+        WriteBuilder::new(self.0.log_store, self.0.state).with_data(data)
     }
 
     /// Vacuum stale files from delta table

@@ -1270,6 +1270,8 @@ pub(super) mod zorder {
 
         #[cfg(test)]
         mod tests {
+            use crate::operations::write::WriteData;
+
             use super::*;
             use ::datafusion::assert_batches_eq;
             use arrow_array::{Int32Array, StringArray};
@@ -1393,7 +1395,7 @@ pub(super) mod zorder {
                 .unwrap();
                 // write some data
                 let table = crate::DeltaOps::new_in_memory()
-                    .write(vec![batch.clone()])
+                    .write(batch.clone().into())
                     .with_save_mode(crate::protocol::SaveMode::Append)
                     .await
                     .unwrap();