Actually add_file in the flush for BufferedJSONWriter

rtyler · rtyler · commit 63d45cf93586 · 2021-05-22T21:58:14.000-07:00
This commit also addresses a number of other code review comments
diff --git a/rust/src/bin/delta-inspect.rs b/rust/src/bin/delta-inspect.rs
@@ -54,7 +54,7 @@ async fn main() -> anyhow::Result<()> {
             let files = if files_matches.is_present("full_path") {
                 table.get_file_paths()
             } else {
-                table.get_files().clone()
+                table.get_files()
             };
 
             files.iter().for_each(|f| println!("{}", f));
diff --git a/rust/src/delta.rs b/rust/src/delta.rs
@@ -1085,9 +1085,16 @@ impl<'a> DeltaTransaction<'a> {
     ///     action for the file added. This should typically be used for txn type actions
     pub async fn add_file(
         &mut self,
-        bytes: &Vec<u8>,
+        bytes: &[u8],
         partitions: Option<Vec<(String, String)>>,
     ) -> Result<(), DeltaTransactionError> {
+        let mut partition_values = HashMap::new();
+        if let Some(partitions) = &partitions {
+            for (key, value) in partitions {
+                partition_values.insert(key.clone(), value.clone());
+            }
+        }
+
         let path = self.generate_parquet_filename(partitions);
         let storage_path = self
             .delta_table
@@ -1109,7 +1116,7 @@ impl<'a> DeltaTransaction<'a> {
         self.actions.push(Action::add(action::Add {
             path,
             size: bytes.len() as i64,
-            partitionValues: HashMap::default(),
+            partitionValues: partition_values,
             partitionValues_parsed: None,
             modificationTime: modification_time,
             dataChange: true,
@@ -1122,13 +1129,12 @@ impl<'a> DeltaTransaction<'a> {
     }
 
     fn generate_parquet_filename(&self, partitions: Option<Vec<(String, String)>>) -> String {
-        let mut path_parts = vec![];
         /*
          * The specific file naming for parquet is not well documented including the preceding five
          * zeros and the trailing c000 string
          *
          */
-        path_parts.push(format!("part-00000-{}-c000.snappy.parquet", Uuid::new_v4()));
+        let mut path_parts = vec![format!("part-00000-{}-c000.snappy.parquet", Uuid::new_v4())];
 
         if let Some(partitions) = partitions {
             for partition in partitions {
diff --git a/rust/src/writer.rs b/rust/src/writer.rs
@@ -7,6 +7,7 @@
 use crate::action::Txn;
 use crate::{DeltaTableError, DeltaTransactionError};
 use arrow::record_batch::RecordBatch;
+use log::*;
 use parquet::arrow::ArrowWriter;
 use parquet::basic::Compression;
 use parquet::file::properties::WriterProperties;
@@ -16,23 +17,19 @@ use std::collections::HashMap;
 use std::convert::TryFrom;
 use std::sync::Arc;
 
-/**
- * BufferedJSONWriter allows for buffering serde_json::Value rows before flushing to parquet files
- * and a Delta transaction
- */
-pub struct BufferedJSONWriter {
+/// BufferedJsonWriter allows for buffering serde_json::Value rows before flushing to parquet files
+/// and a Delta transaction
+pub struct BufferedJsonWriter {
     table: crate::DeltaTable,
     buffer: HashMap<WriterPartition, Vec<Value>>,
     schema: arrow::datatypes::SchemaRef,
     partitions: Vec<String>,
     txns: Vec<Txn>,
 }
 
-impl BufferedJSONWriter {
-    /**
-     * Attempt to construct the BufferedJSONWriter, will fail if the table's metadata is not
-     * present
-     */
+impl BufferedJsonWriter {
+    /// Attempt to construct the BufferedJsonWriter, will fail if the table's metadata is not
+    /// present
     pub fn try_new(table: crate::DeltaTable) -> Result<Self, DeltaTableError> {
         let metadata = table.get_metadata()?.clone();
         let schema = metadata.schema;
@@ -49,38 +46,32 @@ impl BufferedJSONWriter {
         })
     }
 
-    /**
-     * Return the total Values pending in the buffer
-     */
+    /// Return the total Values pending in the buffer
     pub fn count(&self, partitions: &WriterPartition) -> Option<usize> {
         self.buffer.get(&partitions).map(|b| b.len())
     }
 
-    /**
-     * Add a txn action to the buffer
-     */
+    /// Add a txn action to the buffer
     pub fn record_txn(&mut self, txn: Txn) {
         self.txns.push(txn);
     }
 
-    /**
-     * Write a new Value into the buffer
-     */
+    /// Write a new Value into the buffer
     pub fn write(
         &mut self,
         value: Value,
         partitions: WriterPartition,
     ) -> Result<(), DeltaTableError> {
         match partitions {
             WriterPartition::NoPartitions => {
-                if self.partitions.len() > 0 {
+                if !self.partitions.is_empty() {
                     return Err(DeltaTableError::SchemaMismatch {
                         msg: "Table has partitions but noone were supplied on write".to_string(),
                     });
                 }
             }
             WriterPartition::KeyValues { .. } => {
-                if self.partitions.len() == 0 {
+                if self.partitions.is_empty() {
                     return Err(DeltaTableError::SchemaMismatch {
                         msg: "Table has no partitions yet they were supplied on write".to_string(),
                     });
@@ -96,12 +87,10 @@ impl BufferedJSONWriter {
         Ok(())
     }
 
-    /**
-     * Flush the buffer, causing a write of parquet files for each set of partitioned information
-     * as well as any buffered txn actions
-     *
-     * This will create a single transaction in the delta transaction log
-     */
+    /// Flush the buffer, causing a write of parquet files for each set of partitioned information
+    /// as well as any buffered txn actions
+    ///
+    /// This will create a single transaction in the delta transaction log
     pub async fn flush(&mut self) -> Result<(), DeltaTransactionError> {
         use arrow::json::reader::Decoder;
 
@@ -116,21 +105,32 @@ impl BufferedJSONWriter {
                 .next_batch(&mut value_iter)
                 .map_err(|source| DeltaTableError::ArrowError { source })?;
 
-            if record_batch.is_none() {
-                return Ok(());
+            if record_batch.is_some() {
+                let mut pb = ParquetBuffer::try_new(self.schema.clone())?;
+                pb.write_batch(&record_batch.unwrap())?;
+                let _metadata = pb.close();
+                parquet_bufs.push((partitions.clone(), pb.data()));
+            } else {
+                warn!("Attempted to flush an empty RecordBatch from the BufferedJsonWriter");
             }
-
-            let mut pb = ParquetBuffer::try_new(self.schema.clone())?;
-            pb.write_batch(&record_batch.unwrap())?;
-            let _metadata = pb.close();
-            parquet_bufs.push(pb.data());
         }
 
         let mut dtx = self.table.create_transaction(None);
+        for (partitions, buf) in parquet_bufs {
+            match partitions {
+                WriterPartition::NoPartitions => {
+                    dtx.add_file(&buf, None).await?;
+                }
+                WriterPartition::KeyValues { partitions } => {
+                    dtx.add_file(&buf, Some(partitions)).await?;
+                }
+            }
+        }
+
         dtx.add_actions(
             self.txns
                 .drain(0..)
-                .map(|t| crate::action::Action::txn(t))
+                .map(crate::action::Action::txn)
                 .collect(),
         );
 
@@ -200,9 +200,7 @@ impl<'a> Iterator for InMemValueIter<'a> {
     }
 }
 
-/**
- * The type of partition for a row being written to a writer
- */
+/// The type of partition for a row being written to a writer
 #[derive(Clone, Debug, Eq, Hash, PartialEq)]
 pub enum WriterPartition {
     /// The row is not partitioned
@@ -221,7 +219,7 @@ mod tests {
     #[tokio::test]
     async fn test_writer_buffer_nopartition() {
         let table = crate::open_table("./tests/data/delta-0.8.0").await.unwrap();
-        let mut writer = BufferedJSONWriter::try_new(table).unwrap();
+        let mut writer = BufferedJsonWriter::try_new(table).unwrap();
         assert_eq!(writer.count(&WriterPartition::NoPartitions), None);
         let res = writer.write(json!({"hello":"world"}), WriterPartition::NoPartitions);
         assert!(res.is_ok());
@@ -233,15 +231,15 @@ mod tests {
         let table = crate::open_table("./tests/data/delta-0.8.0-partitioned")
             .await
             .unwrap();
-        let mut writer = BufferedJSONWriter::try_new(table).unwrap();
+        let mut writer = BufferedJsonWriter::try_new(table).unwrap();
         let res = writer.write(json!({"hello":"world"}), WriterPartition::NoPartitions);
         assert!(res.is_err());
     }
 
     #[tokio::test]
     async fn test_writer_write_partitions_to_nopartition() {
         let table = crate::open_table("./tests/data/delta-0.8.0").await.unwrap();
-        let mut writer = BufferedJSONWriter::try_new(table).unwrap();
+        let mut writer = BufferedJsonWriter::try_new(table).unwrap();
         let partitions = WriterPartition::KeyValues {
             partitions: vec![("year".to_string(), "2021".to_string())],
         };
