[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,64 @@ 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 DataFileMetadata {
+    file_meta: FileMeta,
+}
+
+impl DataFileMetadata {
+    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 DataFileMetadata {
+            file_meta:
+                FileMeta {
+                    location,
+                    last_modified,
+                    size,
+                },
+        } = self;
+        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);
+        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 +108,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<DataFileMetadata> {
+        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_millis();
+        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(DataFileMetadata::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> {
@@ -242,9 +364,12 @@ impl FileOpener for PresignedUrlOpener {
 #[cfg(test)]
 mod tests {
     use std::path::PathBuf;
+    use std::time::{SystemTime, UNIX_EPOCH};
 
+    use arrow_array::array::Array;
     use arrow_array::RecordBatch;
-    use object_store::{local::LocalFileSystem, ObjectStore};
+    use object_store::{local::LocalFileSystem, memory::InMemory, ObjectStore};
+    use url::Url;
 
     use crate::engine::arrow_data::ArrowEngineData;
     use crate::engine::default::executor::tokio::TokioBackgroundExecutor;
@@ -297,4 +422,122 @@ mod tests {
         assert_eq!(data.len(), 1);
         assert_eq!(data[0].num_rows(), 10);
     }
+
+    #[test]
+    fn test_as_record_batch() {
+        let location = Url::parse("file:///test_url").unwrap();
+        let size = 1_000_000;
+        let last_modified = 10000000000;
+        let file_metadata = FileMeta::new(location.clone(), last_modified, size as usize);
+        let data_file_metadata = DataFileMetadata::new(file_metadata);
+        let partition_values = HashMap::from([("partition1".to_string(), "a".to_string())]);
+        let data_change = true;
+        let actual = data_file_metadata
+            .as_record_batch(&partition_values, data_change)
+            .unwrap();
+        let actual = ArrowEngineData::try_from_engine_data(actual).unwrap();
+
+        let schema = Arc::new(
+            crate::transaction::get_write_metadata_schema()
+                .as_ref()
+                .try_into()
+                .unwrap(),
+        );
+        let key_builder = StringBuilder::new();
+        let val_builder = StringBuilder::new();
+        let mut partition_values_builder = MapBuilder::new(
+            Some(MapFieldNames {
+                entry: "key_value".to_string(),
+                key: "key".to_string(),
+                value: "value".to_string(),
+            }),
+            key_builder,
+            val_builder,
+        );
+        partition_values_builder.keys().append_value("partition1");
+        partition_values_builder.values().append_value("a");
+        partition_values_builder.append(true).unwrap();
+        let partition_values = partition_values_builder.finish();
+        let expected = RecordBatch::try_new(
+            schema,
+            vec![
+                Arc::new(StringArray::from(vec![location.to_string()])),
+                Arc::new(partition_values),
+                Arc::new(Int64Array::from(vec![size])),
+                Arc::new(Int64Array::from(vec![last_modified])),
+                Arc::new(BooleanArray::from(vec![data_change])),
+            ],
+        )
+        .unwrap();
+
+        assert_eq!(actual.record_batch(), &expected);
+    }
+
+    #[tokio::test]
+    async fn test_write_parquet() {
+        let store = Arc::new(InMemory::new());
+        let parquet_handler =
+            DefaultParquetHandler::new(store.clone(), Arc::new(TokioBackgroundExecutor::new()));
+
+        let data = Box::new(ArrowEngineData::new(
+            RecordBatch::try_from_iter(vec![(
+                "a",
+                Arc::new(Int64Array::from(vec![1, 2, 3])) as Arc<dyn Array>,
+            )])
+            .unwrap(),
+        ));
+
+        let write_metadata = parquet_handler
+            .write_parquet(&Url::parse("memory:///data/").unwrap(), data)
+            .await
+            .unwrap();
+
+        let DataFileMetadata {
+            file_meta:
+                ref parquet_file @ FileMeta {
+                    ref location,
+                    last_modified,
+                    size,
+                },
+        } = write_metadata;
+        let expected_location = Url::parse("memory:///data/").unwrap();
+        let expected_size = 497;
+
+        // check that last_modified is within 10s of now
+        let now: i64 = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_millis()
+            .try_into()
+            .unwrap();
+
+        let filename = location.path().split('/').last().unwrap();
+        assert_eq!(&expected_location.join(filename).unwrap(), location);
+        assert_eq!(expected_size, size);
+        assert!(now - last_modified < 10_000);
+
+        // check we can read back
+        let path = Path::from(location.path());
+        let meta = store.head(&path).await.unwrap();
+        let reader = ParquetObjectReader::new(store.clone(), meta.clone());
+        let physical_schema = ParquetRecordBatchStreamBuilder::new(reader)
+            .await
+            .unwrap()
+            .schema()
+            .clone();
+
+        let data: Vec<RecordBatch> = parquet_handler
+            .read_parquet_files(
+                &[parquet_file.clone()],
+                Arc::new(physical_schema.try_into().unwrap()),
+                None,
+            )
+            .unwrap()
+            .map(into_record_batch)
+            .try_collect()
+            .unwrap();
+
+        assert_eq!(data.len(), 1);
+        assert_eq!(data[0].num_rows(), 3);
+    }
 }