Properly add support for __created_at column in etl crate (#66)

* Support for e2e latency metric

* Fix

* __create_at column

* Lint

* Fix

Author: krinart

crates/data-generation/src/dataset/simple_sequence.rs

@@ -78,6 +78,9 @@ impl SimpleSequenceDataset {
     /// Returns the static Arrow schema for the `integer_sequence` table.
     ///
     /// Includes change-tracking columns (`_op`, `_op_index`).
+    ///
+    /// The time column (`__created_at`) is not included; it will be added during
+    /// ETL rehydration.
     pub fn schema() -> SchemaRef {
         Arc::new(Schema::new(vec![
             Field::new("id", DataType::Int64, false),

crates/data-generation/src/dataset/tpch.rs

@@ -29,8 +29,8 @@ use crate::dataset::MutationConfig;
 
 use super::{Dataset, DatasetTable};
 
-/// TPC-H table definitions: (table_name, time_column, schema_fn).
-const TPCH_TABLE_TIME_COLUMNS: &[(&str, &str)] = &[
+/// TPC-H table definitions: `(table_name, time_column)`.
+const TPCH_TABLES: &[(&str, &str)] = &[
     ("region", "r_created_at"),
     ("nation", "n_created_at"),
     ("supplier", "s_created_at"),
@@ -217,7 +217,7 @@ impl TpchDataset {
     /// Drop the `_new` tables created by `advance_step()`.
     fn drop_step_tables(&self) -> anyhow::Result<()> {
         let mut sql = String::new();
-        for (table, _) in TPCH_TABLE_TIME_COLUMNS {
+        for (table, _) in TPCH_TABLES {
             sql.push_str(&format!("DROP TABLE IF EXISTS {table}_new;"));
         }
         let conn = self
@@ -256,10 +256,7 @@ impl Dataset for TpchDataset {
     }
 
     fn num_batches(&self, table: &str) -> u64 {
-        if !TPCH_TABLE_TIME_COLUMNS
-            .iter()
-            .any(|(name, _)| *name == table)
-        {
+        if !TPCH_TABLES.iter().any(|(name, _)| *name == table) {
             return 0;
         }
 
@@ -274,7 +271,7 @@ impl Dataset for TpchDataset {
                 .consumed_tables
                 .read()
                 .map_err(|e| anyhow::anyhow!("lock poisoned: {e}"))?;
-            if consumed.len() >= TPCH_TABLE_TIME_COLUMNS.len() {
+            if consumed.len() >= TPCH_TABLES.len() {
                 drop(consumed);
                 if self.current_step.load(Ordering::SeqCst) > 0 {
                     self.drop_step_tables()?;
@@ -302,10 +299,7 @@ impl Dataset for TpchDataset {
         }
 
         // Validate the table name
-        if !TPCH_TABLE_TIME_COLUMNS
-            .iter()
-            .any(|(name, _)| *name == table)
-        {
+        if !TPCH_TABLES.iter().any(|(name, _)| *name == table) {
             anyhow::bail!("Unknown TPC-H table: {table}");
         }
 
@@ -358,7 +352,7 @@ impl Dataset for TpchDataset {
     }
 
     fn tables(&self) -> HashMap<String, DatasetTable> {
-        TPCH_TABLE_TIME_COLUMNS
+        TPCH_TABLES
             .iter()
             .map(|(name, time_col)| {
                 (

crates/etl/src/lib.rs

@@ -16,7 +16,10 @@ limitations under the License.
 
 use std::collections::HashMap;
 use std::sync::Arc;
+use std::time::{SystemTime, UNIX_EPOCH};
 
+use arrow::array::{RecordBatch, TimestampMicrosecondArray};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
 use data_generation::config::DatasetConfig as GenerationDatasetConfig;
 use data_generation::dataset::simple_sequence::SimpleSequenceDataset;
 use data_generation::dataset::tpch::TpchDataset;
@@ -36,6 +39,38 @@ use crate::sink::{InsertOp, Sink};
 
 pub mod sink;
 
+/// Column name appended by the ETL pipeline to every batch.
+const CREATED_AT_COLUMN: &str = "__created_at";
+
+/// Returns a new schema with the `__created_at` timestamp column appended.
+fn schema_with_created_at(schema: &SchemaRef) -> SchemaRef {
+    let mut fields: Vec<_> = schema.fields().iter().cloned().collect();
+    fields.push(Arc::new(Field::new(
+        CREATED_AT_COLUMN,
+        DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
+        true,
+    )));
+    Arc::new(Schema::new(fields))
+}
+
+/// Appends a `__created_at` column (current wall-clock time, microsecond UTC)
+/// to the given batch.
+fn append_created_at(batch: &RecordBatch) -> anyhow::Result<RecordBatch> {
+    let now_us = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .expect("system time before UNIX epoch")
+        .as_micros() as i64;
+
+    let timestamps =
+        TimestampMicrosecondArray::from(vec![Some(now_us); batch.num_rows()]).with_timezone("UTC");
+
+    let new_schema = schema_with_created_at(&batch.schema());
+    let mut columns: Vec<_> = batch.columns().to_vec();
+    columns.push(Arc::new(timestamps));
+
+    Ok(RecordBatch::try_new(new_schema, columns)?)
+}
+
 /// Specifies which dataset implementation to use for the ETL pipeline.
 #[derive(Debug, Clone)]
 pub enum DatasetSource {
@@ -191,7 +226,7 @@ impl ETLPipeline {
             .into_iter()
             .map(|(name, table)| {
                 let config = ProtocolDatasetConfig {
-                    schema: table.rehydrated_schema(),
+                    schema: schema_with_created_at(&table.schema),
                 };
                 (name, config)
             })
@@ -220,7 +255,6 @@ impl ETLPipeline {
 
         let mut join_set: JoinSet<Result<String, String>> = JoinSet::new();
         for table_name in tables.keys() {
-            let dataset = Arc::clone(&self.dataset);
             let source = Arc::clone(&self.data_storage);
             let target = Arc::clone(&self.data_sink);
             let table_name = table_name.clone();
@@ -237,8 +271,8 @@ impl ETLPipeline {
                 let op = sink_op_from_batch_op(&read_result.operation);
 
                 for batch in read_result.batches {
-                    let rehydrated = dataset.rehydrate(&table_name, &batch).map_err(|e| {
-                        format!("rehydrate {table_name} batch {first_batch_id}: {e}")
+                    let rehydrated = append_created_at(&batch).map_err(|e| {
+                        format!("append __created_at to {table_name} batch {first_batch_id}: {e}")
                     })?;
 
                     target
@@ -289,8 +323,8 @@ impl ETLPipeline {
     /// each batch the task:
     ///
     /// 1. Reads the batch from the [`Source`].
-    /// 2. Rehydrates it through the [`Dataset`] (appending time columns, etc.).
-    /// 3. Writes the rehydrated batch to the [`Sink`].
+    /// 2. Appends the `__created_at` timestamp column.
+    /// 3. Writes the enriched batch to the [`Sink`].
     ///
     /// The task transitions to [`PipelineState::Stopped`] when all batches are
     /// processed, the [`CancellationToken`] is triggered, or an error occurs.
@@ -329,7 +363,7 @@ impl ETLPipeline {
         work.sort_by_key(|(_, id)| *id);
 
         let handle = tokio::spawn(async move {
-            let reason = run_pipeline(dataset, source, target, work, cancel).await;
+            let reason = run_pipeline(source, target, work, cancel).await;
             let _ = state_tx.send(PipelineState::Stopped(reason));
         });
 
@@ -355,7 +389,6 @@ impl ETLPipeline {
 /// Groups work items by batch ID and processes all tables within each step
 /// concurrently, checking for cancellation between steps.
 async fn run_pipeline(
-    dataset: Arc<dyn Dataset>,
     data_storage: Arc<dyn DataStorage>,
     data_sink: Arc<dyn Sink>,
     work: Vec<(String, u64)>,
@@ -433,7 +466,6 @@ async fn run_pipeline(
         // Process all tables for this batch_id concurrently.
         let mut join_set: JoinSet<Result<(String, bool), String>> = JoinSet::new();
         for table_name in active_tables {
-            let dataset = Arc::clone(&dataset);
             let data_storage = Arc::clone(&data_storage);
             let data_sink = Arc::clone(&data_sink);
 
@@ -462,18 +494,20 @@ async fn run_pipeline(
 
                 let op = sink_op_from_batch_op(&read_result.operation);
 
-                // 2. Rehydrate each record batch and write to target
+                // 2. Append __created_at and write to target
                 for batch in read_result.batches {
-                    let rehydrated = match dataset.rehydrate(&table_name, &batch) {
+                    let rehydrated = match append_created_at(&batch) {
                         Ok(b) => b,
                         Err(e) => {
                             error!(
                                 table = %table_name,
                                 batch_id,
                                 error = %e,
-                                "Failed to rehydrate batch"
+                                "Failed to append __created_at column"
                             );
-                            return Err(format!("rehydrate {table_name} batch {batch_id}: {e}"));
+                            return Err(format!(
+                                "append __created_at to {table_name} batch {batch_id}: {e}"
+                            ));
                         }
                     };
 

src/metrics.rs

@@ -219,7 +219,6 @@ pub static EFFICIENCY_QUERIES_PER_CORE: LazyLock<Gauge<f64>> = LazyLock::new(||
 
 // --- E2E Latency ---
 
-#[allow(dead_code)]
 pub static E2E_LATENCY_P99_MS: LazyLock<Histogram<f64>> = LazyLock::new(|| {
     meter()
         .f64_histogram("e2e_latency_p99_ms")