make etl hive sink faster (#130)

* make etl hive sink faster

* chore: auto-fix cargo fmt + clippy

* Update s3_hive.rs

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: peasee

crates/etl/src/lib.rs

@@ -487,6 +487,8 @@ impl ETLPipeline {
                             prefix = config.prefix
                         )
                     }),
+                    time_column: Some(CREATED_AT_COLUMN.to_string()),
+                    partition_columns: self.dataset.partition_columns(&name),
                 };
 
                 (name.clone(), config)

crates/etl/src/sink/s3_hive.rs

@@ -27,6 +27,7 @@ use object_store::path::Path as ObjectPath;
 use parquet::arrow::ArrowWriter;
 use parquet::basic::Compression;
 use parquet::file::properties::WriterProperties;
+use tokio::task::JoinSet;
 
 use super::{InsertOp, Sink};
 
@@ -93,70 +94,75 @@ impl S3HiveSink {
             partition_columns,
         })
     }
+}
 
-    /// Writes a single partition's batch to S3 as a Parquet file.
-    async fn write_partition(
-        &self,
-        table_name: &str,
-        batch_id: u64,
-        partition_path: &str,
-        batch: &RecordBatch,
-        effective_partition_columns: &[String],
-        partition_idx: usize,
-    ) -> anyhow::Result<()> {
-        // Strip partition columns from written data — they are encoded in the path.
-        let batch_without_partition = strip_columns(batch, effective_partition_columns)?;
-
-        if batch_without_partition.num_columns() == 0 {
-            anyhow::bail!(
-                "Cannot write table '{table_name}' with partition columns {:?}: no columns would remain in parquet output",
-                effective_partition_columns
-            );
-        }
+/// Writes a single partition batch to S3 as a Parquet file.
+///
+/// Parquet encoding (CPU-bound) is offloaded to `spawn_blocking` so that the
+/// async executor is not blocked during compression. The resulting bytes are
+/// then uploaded with a single `PUT`.
+async fn write_partition_task(
+    store: Arc<dyn ObjectStore>,
+    prefix: String,
+    table_name: String,
+    batch_id: u64,
+    partition_path: String,
+    batch: RecordBatch,
+    effective_partition_columns: Vec<String>,
+    partition_idx: usize,
+) -> anyhow::Result<()> {
+    // Strip partition columns — they are encoded in the path.
+    let batch_without_partition = strip_columns(&batch, &effective_partition_columns)?;
+
+    if batch_without_partition.num_columns() == 0 {
+        anyhow::bail!(
+            "Cannot write table '{table_name}' with partition columns {effective_partition_columns:?}: \
+             no columns would remain in parquet output"
+        );
+    }
 
-        let path = if self.prefix.is_empty() {
-            if partition_path.is_empty() {
-                ObjectPath::from(format!(
-                    "{table_name}/batch-{batch_id:06}-{partition_idx:04}.parquet"
-                ))
-            } else {
-                ObjectPath::from(format!(
-                    "{table_name}/{partition_path}/batch-{batch_id:06}-{partition_idx:04}.parquet"
-                ))
-            }
+    let path = if prefix.is_empty() {
+        if partition_path.is_empty() {
+            ObjectPath::from(format!(
+                "{table_name}/batch-{batch_id:06}-{partition_idx:04}.parquet"
+            ))
         } else {
-            if partition_path.is_empty() {
-                ObjectPath::from(format!(
-                    "{}/{table_name}/batch-{batch_id:06}-{partition_idx:04}.parquet",
-                    self.prefix
-                ))
-            } else {
-                ObjectPath::from(format!(
-                    "{}/{table_name}/{partition_path}/batch-{batch_id:06}-{partition_idx:04}.parquet",
-                    self.prefix
-                ))
-            }
-        };
-
+            ObjectPath::from(format!(
+                "{table_name}/{partition_path}/batch-{batch_id:06}-{partition_idx:04}.parquet"
+            ))
+        }
+    } else if partition_path.is_empty() {
+        ObjectPath::from(format!(
+            "{prefix}/{table_name}/batch-{batch_id:06}-{partition_idx:04}.parquet",
+        ))
+    } else {
+        ObjectPath::from(format!(
+            "{prefix}/{table_name}/{partition_path}/batch-{batch_id:06}-{partition_idx:04}.parquet",
+        ))
+    };
+
+    // Encode to Parquet + Snappy on a blocking thread so the async executor
+    // is not stalled during CPU-intensive compression.
+    let buf = tokio::task::spawn_blocking(move || -> anyhow::Result<Vec<u8>> {
         let props = WriterProperties::builder()
             .set_compression(Compression::SNAPPY)
             .build();
-
         let mut buf = Vec::new();
-        {
-            let mut writer =
-                ArrowWriter::try_new(&mut buf, batch_without_partition.schema(), Some(props))?;
-            writer.write(&batch_without_partition)?;
-            writer.close()?;
-        }
-
-        self.store
-            .put(&path, buf.into())
-            .await
-            .map_err(|e| anyhow::anyhow!("S3 PUT failed for {path}: {e}"))?;
-
-        Ok(())
-    }
+        let mut writer =
+            ArrowWriter::try_new(&mut buf, batch_without_partition.schema(), Some(props))?;
+        writer.write(&batch_without_partition)?;
+        writer.close()?;
+        Ok(buf)
+    })
+    .await
+    .map_err(|e| anyhow::anyhow!("Parquet encoding task panicked: {e}"))??;
+
+    store
+        .put(&path, buf.into())
+        .await
+        .map_err(|e| anyhow::anyhow!("S3 PUT failed for {path}: {e}"))?;
+
+    Ok(())
 }
 
 #[async_trait]
@@ -204,16 +210,27 @@ impl Sink for S3HiveSink {
         // Group rows by distinct partition tuples in configured column order.
         let partitions = partition_batch(&batch, &partition_columns_with_idx)?;
 
-        for (idx, (partition_path, partition_batch)) in partitions.iter().enumerate() {
-            self.write_partition(
-                table_name,
+        // Spawn all partition writes concurrently. Each task owns its data
+        // so there is no contention, and S3 PUTs for different paths are
+        // fully independent.
+        let mut join_set: JoinSet<anyhow::Result<()>> = JoinSet::new();
+
+        for (idx, (partition_path, partition_batch)) in partitions.into_iter().enumerate() {
+            join_set.spawn(write_partition_task(
+                Arc::clone(&self.store),
+                self.prefix.clone(),
+                table_name.to_string(),
                 batch_id,
                 partition_path,
                 partition_batch,
-                &effective_partition_columns,
+                effective_partition_columns.clone(),
                 idx,
-            )
-            .await?;
+            ));
+        }
+
+        // Collect results; propagate the first error encountered.
+        while let Some(result) = join_set.join_next().await {
+            result.map_err(|e| anyhow::anyhow!("Partition write task panicked: {e}"))??;
         }
 
         Ok(())

crates/system-adapter-protocol/src/lib.rs

@@ -135,6 +135,10 @@ pub struct DatasetConfig {
     pub schema: SchemaRef,
     /// Dataset S3 location (e.g. "s3://my-bucket/path/to/data/")
     pub location: Option<String>,
+    /// Optional column name to use as the ingestion time for metrics tracking
+    pub time_column: Option<String>,
+    /// Optional list of columns to use for partitioning the dataset in storage
+    pub partition_columns: Vec<String>,
 }
 
 /// Request to setup a benchmark run.