[write] Transaction append data API

kernel/Cargo.toml

@@ -90,6 +90,8 @@ default-engine = [
   "parquet/object_store",
   "reqwest",
   "tokio",
+  "uuid/v4",
+  "uuid/fast-rng",
 ]
 
 developer-visibility = []

kernel/src/engine/arrow_utils.rs

@@ -665,11 +665,11 @@ fn parse_json_impl(json_strings: &StringArray, schema: ArrowSchemaRef) -> DeltaR
 /// serialize an arrow RecordBatch to a JSON string by appending to a buffer.
 // TODO (zach): this should stream data to the JSON writer and output an iterator.
 pub(crate) fn to_json_bytes(
-    data: impl Iterator<Item = Box<dyn EngineData>> + Send,
+    data: impl Iterator<Item = DeltaResult<Box<dyn EngineData>>> + Send,
 ) -> DeltaResult<Vec<u8>> {
     let mut writer = LineDelimitedWriter::new(Vec::new());
     for chunk in data.into_iter() {
-        let arrow_data = ArrowEngineData::try_from_engine_data(chunk)?;
+        let arrow_data = ArrowEngineData::try_from_engine_data(chunk?)?;
         let record_batch = arrow_data.record_batch();
         writer.write(record_batch)?;
     }
@@ -1436,7 +1436,7 @@ mod tests {
             vec![Arc::new(StringArray::from(vec!["string1", "string2"]))],
         )?;
         let data: Box<dyn EngineData> = Box::new(ArrowEngineData::new(data));
-        let json = to_json_bytes(Box::new(std::iter::once(data)))?;
+        let json = to_json_bytes(Box::new(std::iter::once(Ok(data))))?;
         assert_eq!(
             json,
             "{\"string\":\"string1\"}\n{\"string\":\"string2\"}\n".as_bytes()

kernel/src/engine/default/json.rs

@@ -96,7 +96,7 @@ impl<E: TaskExecutor> JsonHandler for DefaultJsonHandler<E> {
     fn write_json_file(
         &self,
         path: &Url,
-        data: Box<dyn Iterator<Item = Box<dyn EngineData>> + Send>,
+        data: Box<dyn Iterator<Item = DeltaResult<Box<dyn EngineData>>> + Send + '_>,
         _overwrite: bool,
     ) -> DeltaResult<()> {
         let buffer = to_json_bytes(data)?;

kernel/src/engine/default/mod.rs

@@ -6,6 +6,7 @@
 //! a separate thread pool, provided by the [`TaskExecutor`] trait. Read more in
 //! the [executor] module.
 
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use self::storage::parse_url_opts;
@@ -16,9 +17,13 @@ use self::executor::TaskExecutor;
 use self::filesystem::ObjectStoreFileSystemClient;
 use self::json::DefaultJsonHandler;
 use self::parquet::DefaultParquetHandler;
+use super::arrow_data::ArrowEngineData;
 use super::arrow_expression::ArrowExpressionHandler;
+use crate::schema::Schema;
+use crate::transaction::WriteContext;
 use crate::{
-    DeltaResult, Engine, ExpressionHandler, FileSystemClient, JsonHandler, ParquetHandler,
+    DeltaResult, Engine, EngineData, ExpressionHandler, FileSystemClient, JsonHandler,
+    ParquetHandler,
 };
 
 pub mod executor;
@@ -108,6 +113,32 @@ impl<E: TaskExecutor> DefaultEngine<E> {
     pub fn get_object_store_for_url(&self, _url: &Url) -> Option<Arc<DynObjectStore>> {
         Some(self.store.clone())
     }
+
+    pub async fn write_parquet(
+        &self,
+        data: &ArrowEngineData,
+        write_context: &WriteContext,
+        partition_values: HashMap<String, String>,
+        data_change: bool,
+    ) -> DeltaResult<Box<dyn EngineData>> {
+        let transform = write_context.logical_to_physical();
+        let input_schema: Schema = data.record_batch().schema().try_into()?;
+        let output_schema = write_context.schema();
+        let logical_to_physical_expr = self.get_expression_handler().get_evaluator(
+            input_schema.into(),
+            transform.clone(),
+            output_schema.clone().into(),
+        );
+        let physical_data = logical_to_physical_expr.evaluate(data)?;
+        self.parquet
+            .write_parquet_file(
+                write_context.target_dir(),
+                physical_data,
+                partition_values,
+                data_change,
+            )
+            .await
+    }
 }
 
 impl<E: TaskExecutor> Engine for DefaultEngine<E> {

kernel/src/engine/default/parquet.rs

@@ -1,23 +1,30 @@
 //! Default Parquet handler implementation
 
+use std::collections::HashMap;
 use std::ops::Range;
 use std::sync::Arc;
 
+use arrow_array::builder::{MapBuilder, MapFieldNames, StringBuilder};
+use arrow_array::{BooleanArray, Int64Array, RecordBatch, StringArray};
 use futures::StreamExt;
 use object_store::path::Path;
 use object_store::DynObjectStore;
 use parquet::arrow::arrow_reader::{
     ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder,
 };
+use parquet::arrow::arrow_writer::ArrowWriter;
 use parquet::arrow::async_reader::{ParquetObjectReader, ParquetRecordBatchStreamBuilder};
+use uuid::Uuid;
 
 use super::file_stream::{FileOpenFuture, FileOpener, FileStream};
+use crate::engine::arrow_data::ArrowEngineData;
 use crate::engine::arrow_utils::{generate_mask, get_requested_indices, reorder_struct_array};
 use crate::engine::default::executor::TaskExecutor;
 use crate::engine::parquet_row_group_skipping::ParquetRowGroupSkipping;
 use crate::schema::SchemaRef;
 use crate::{
-    DeltaResult, Error, ExpressionRef, FileDataReadResultIterator, FileMeta, ParquetHandler,
+    DeltaResult, EngineData, Error, ExpressionRef, FileDataReadResultIterator, FileMeta,
+    ParquetHandler,
 };
 
 #[derive(Debug)]
@@ -27,6 +34,66 @@ pub struct DefaultParquetHandler<E: TaskExecutor> {
     readahead: usize,
 }
 
+/// Metadata of a parquet file, currently just includes the file metadata but will expand to
+/// include file statistics and other metadata in the future.
+#[derive(Debug)]
+pub struct ParquetWriteMetadata {
+    file_meta: FileMeta,
+}
+
+impl ParquetWriteMetadata {
+    pub fn new(file_meta: FileMeta) -> Self {
+        Self { file_meta }
+    }
+
+    // convert ParquetMetadata into a record batch which matches the 'write_metadata' schema
+    fn as_record_batch(
+        &self,
+        partition_values: HashMap<String, String>,
+        data_change: bool,
+    ) -> DeltaResult<Box<dyn EngineData>> {
+        let ParquetWriteMetadata { file_meta } = self;
+        let FileMeta {
+            location,
+            last_modified,
+            size,
+        } = file_meta;
+        let write_metadata_schema = crate::transaction::get_write_metadata_schema();
+
+        // create the record batch of the write metadata
+        let path = Arc::new(StringArray::from(vec![location.to_string()]));
+        let key_builder = StringBuilder::new();
+        let val_builder = StringBuilder::new();
+        let names = MapFieldNames {
+            entry: "key_value".to_string(),
+            key: "key".to_string(),
+            value: "value".to_string(),
+        };
+        let mut builder = MapBuilder::new(Some(names), key_builder, val_builder);
+        if partition_values.is_empty() {
+            builder.append(true).unwrap();
+        } else {
+            for (k, v) in partition_values {
+                builder.keys().append_value(&k);
+                builder.values().append_value(&v);
+                builder.append(true).unwrap();
+            }
+        }
+        let partitions = Arc::new(builder.finish());
+        // this means max size we can write is i64::MAX (~8EB)
+        let size: i64 = (*size)
+            .try_into()
+            .map_err(|_| Error::generic("Failed to convert parquet metadata 'size' to i64"))?;
+        let size = Arc::new(Int64Array::from(vec![size]));
+        let data_change = Arc::new(BooleanArray::from(vec![data_change]));
+        let modification_time = Arc::new(Int64Array::from(vec![*last_modified]));
+        Ok(Box::new(ArrowEngineData::new(RecordBatch::try_new(
+            Arc::new(write_metadata_schema.as_ref().try_into()?),
+            vec![path, partitions, size, modification_time, data_change],
+        )?)))
+    }
+}
+
 impl<E: TaskExecutor> DefaultParquetHandler<E> {
     pub fn new(store: Arc<DynObjectStore>, task_executor: Arc<E>) -> Self {
         Self {
@@ -43,6 +110,63 @@ impl<E: TaskExecutor> DefaultParquetHandler<E> {
         self.readahead = readahead;
         self
     }
+
+    // Write `data` to `path`/<uuid>.parquet as parquet using ArrowWriter and return the parquet
+    // metadata (where <uuid> is a generated UUIDv4).
+    //
+    // Note: after encoding the data as parquet, this issues a PUT followed by a HEAD to storage in
+    // order to obtain metadata about the object just written.
+    async fn write_parquet(
+        &self,
+        path: &url::Url,
+        data: Box<dyn EngineData>,
+    ) -> DeltaResult<ParquetWriteMetadata> {
+        let batch: Box<_> = ArrowEngineData::try_from_engine_data(data)?;
+        let record_batch = batch.record_batch();
+
+        let mut buffer = vec![];
+        let mut writer = ArrowWriter::try_new(&mut buffer, record_batch.schema(), None)?;
+        writer.write(record_batch)?;
+        writer.close()?; // writer must be closed to write footer
+
+        let size = buffer.len();
+        let name: String = Uuid::new_v4().to_string() + ".parquet";
+        // FIXME test with trailing '/' and omitting?
+        let path = path.join(&name)?;
+
+        self.store
+            .put(&Path::from(path.path()), buffer.into())
+            .await?;
+
+        let metadata = self.store.head(&Path::from(path.path())).await?;
+        let modification_time = metadata.last_modified.timestamp();
+        if size != metadata.size {
+            return Err(Error::generic(format!(
+                "Size mismatch after writing parquet file: expected {}, got {}",
+                size, metadata.size
+            )));
+        }
+
+        let file_meta = FileMeta::new(path, modification_time, size);
+        Ok(ParquetWriteMetadata::new(file_meta))
+    }
+
+    /// Write `data` to `path`/<uuid>.parquet as parquet using ArrowWriter and return the parquet
+    /// metadata as an EngineData batch which matches the [write metadata] schema (where <uuid> is
+    /// a generated UUIDv4).
+    ///
+    /// [write metadata]: crate::transaction::get_write_metadata_schema
+    pub async fn write_parquet_file(
+        &self,
+        path: &url::Url,
+        data: Box<dyn EngineData>,
+        partition_values: HashMap<String, String>,
+        data_change: bool,
+    ) -> DeltaResult<Box<dyn EngineData>> {
+        let parquet_metadata = self.write_parquet(path, data).await?;
+        let write_metadata = parquet_metadata.as_record_batch(partition_values, data_change)?;
+        Ok(write_metadata)
+    }
 }
 
 impl<E: TaskExecutor> ParquetHandler for DefaultParquetHandler<E> {
@@ -297,4 +421,10 @@ mod tests {
         assert_eq!(data.len(), 1);
         assert_eq!(data[0].num_rows(), 10);
     }
+
+    #[test]
+    fn test_into_write_metadata() {}
+
+    #[tokio::test]
+    async fn test_write_parquet() {}
 }

kernel/src/engine/sync/json.rs

@@ -52,7 +52,7 @@ impl JsonHandler for SyncJsonHandler {
     fn write_json_file(
         &self,
         path: &Url,
-        data: Box<dyn Iterator<Item = Box<dyn EngineData>> + Send>,
+        data: Box<dyn Iterator<Item = DeltaResult<Box<dyn EngineData>>> + Send + '_>,
         _overwrite: bool,
     ) -> DeltaResult<()> {
         let path = path
@@ -120,10 +120,10 @@ mod tests {
 
         let url = Url::from_file_path(path.clone()).unwrap();
         handler
-            .write_json_file(&url, Box::new(std::iter::once(data)), false)
+            .write_json_file(&url, Box::new(std::iter::once(Ok(data))), false)
             .expect("write json file");
         assert!(matches!(
-            handler.write_json_file(&url, Box::new(std::iter::once(empty)), false),
+            handler.write_json_file(&url, Box::new(std::iter::once(Ok(empty))), false),
             Err(Error::FileAlreadyExists(_))
         ));
 

kernel/src/lib.rs

@@ -124,6 +124,17 @@ impl PartialOrd for FileMeta {
     }
 }
 
+impl FileMeta {
+    /// Create a new instance of `FileMeta`
+    pub fn new(location: Url, last_modified: i64, size: usize) -> Self {
+        Self {
+            location,
+            last_modified,
+            size,
+        }
+    }
+}
+
 /// Trait for implementing an Expression evaluator.
 ///
 /// It contains one Expression which can be evaluated on multiple ColumnarBatches.
@@ -233,7 +244,7 @@ pub trait JsonHandler: Send + Sync {
     fn write_json_file(
         &self,
         path: &Url,
-        data: Box<dyn Iterator<Item = Box<dyn EngineData>> + Send>,
+        data: Box<dyn Iterator<Item = DeltaResult<Box<dyn EngineData>>> + Send + '_>,
         overwrite: bool,
     ) -> DeltaResult<()>;
 }