feat: Implement basic ETL pipeline for creates with rehydrations (#36)

* feat: Add ETLPipeline struct

* feat: Update ETL pipeline for E2E

* fix: Add params output for setup_request_datasets

Author: peasee

Cargo.lock

@@ -0,0 +1,5 @@
+# Data Generator
+
+```bash
+cargo run -p data-generation -- run --scale-factor 1 --bucket peasee-indexes --region us-west-2 --prefix raw --num-steps 10
+```

crates/data-generation/README.md

@@ -52,7 +52,7 @@ pub struct CommonArgs {
     pub scale_factor: f64,
 
     /// Number of data generation steps (partitions for TPC-H dbgen)
-    #[arg(long, default_value_t = 100)]
+    #[arg(long, default_value_t = 25)]
     pub num_steps: u16,
 
     /// S3 bucket name
@@ -72,7 +72,7 @@ pub struct CommonArgs {
     pub endpoint: Option<String>,
 
     /// Maximum number of concurrent S3 writes
-    #[arg(long, default_value_t = 8)]
+    #[arg(long, default_value_t = 16)]
     pub max_concurrency: usize,
 }
 

crates/data-generation/src/config.rs

@@ -17,14 +17,16 @@ limitations under the License.
 pub mod tpch;
 pub mod simple_sequence;
 
-use std::collections::HashMap;
+use std::collections::{HashMap, VecDeque};
 use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use arrow::array::{RecordBatch, TimestampMicrosecondArray};
 use arrow::datatypes::{DataType, Field, SchemaRef, TimeUnit};
 use async_trait::async_trait;
 
+use crate::config::DatasetConfig;
+
 /// Metadata about a table in a dataset.
 #[derive(Debug, Clone)]
 pub struct DatasetTable {
@@ -64,11 +66,41 @@ impl DatasetTable {
         }
 
         if batch.schema() != self.schema {
+            let mut diffs = Vec::new();
+            let expected_fields = self.schema.fields();
+            let actual_schema = batch.schema();
+            let actual_fields = actual_schema.fields();
+
+            for (i, expected) in expected_fields.iter().enumerate() {
+                match actual_fields.get(i) {
+                    Some(actual) if actual != expected => {
+                        if actual.name() != expected.name() {
+                            diffs.push(format!("  column {i}: expected name '{}', got '{}'", expected.name(), actual.name()));
+                        }
+                        if actual.data_type() != expected.data_type() {
+                            diffs.push(format!("  column '{}' (index {i}): expected type {:?}, got {:?}", expected.name(), expected.data_type(), actual.data_type()));
+                        }
+                        if actual.is_nullable() != expected.is_nullable() {
+                            diffs.push(format!("  column '{}' (index {i}): expected nullable={}, got nullable={}", expected.name(), expected.is_nullable(), actual.is_nullable()));
+                        }
+                    }
+                    None => {
+                        diffs.push(format!("  column '{}' (index {i}): missing from batch", expected.name()));
+                    }
+                    _ => {}
+                }
+            }
+            for i in expected_fields.len()..actual_fields.len() {
+                diffs.push(format!("  column '{}' (index {i}): unexpected extra column in batch", actual_fields[i].name()));
+            }
+            if expected_fields.len() != actual_fields.len() {
+                diffs.push(format!("  expected {} columns, got {}", expected_fields.len(), actual_fields.len()));
+            }
+
             anyhow::bail!(
-                "Schema mismatch for table '{}': expected {}, got {}",
+                "Schema mismatch for table '{}':\n{}",
                 self.name,
-                self.schema,
-                batch.schema()
+                diffs.join("\n")
             );
         }
 
@@ -91,12 +123,27 @@ impl DatasetTable {
 
 #[async_trait]
 pub trait Dataset: Send + Sync {
+    /// Creates a new instance of this dataset from the given configuration.
+    ///
+    /// This is a factory method that returns an `Arc<dyn Dataset>` without any
+    /// external side-effects beyond initialising in-memory state.
+    ///
+    /// The default implementation returns an error; concrete dataset types
+    /// should override this.
+    fn create(config: &DatasetConfig) -> anyhow::Result<Arc<dyn Dataset>>
+    where
+        Self: Sized + 'static,
+    {
+        let _ = config;
+        anyhow::bail!("create() is not implemented for this dataset type")
+    }
+
     /// Returns the batch IDs that would be produced for a given table after a
     /// successful generation run.
     ///
     /// The default implementation returns `0..num_batches(table)`, but
     /// implementations may override this to customise the ID scheme.
-    fn batch_ids(&self, table: &str) -> Vec<u64> {
+    fn batch_ids(&self, table: &str) -> VecDeque<u64> {
         (0..self.num_batches(table)).collect()
     }
 
@@ -127,9 +174,40 @@ pub trait Dataset: Send + Sync {
         };
 
         if batch.schema() != expected_schema {
+            let mut diffs = Vec::new();
+            let expected_fields = expected_schema.fields();
+            let actual_schema = batch.schema();
+            let actual_fields = actual_schema.fields();
+
+            for (i, expected) in expected_fields.iter().enumerate() {
+                match actual_fields.get(i) {
+                    Some(actual) if actual != expected => {
+                        if actual.name() != expected.name() {
+                            diffs.push(format!("  column {i}: expected name '{}', got '{}'", expected.name(), actual.name()));
+                        }
+                        if actual.data_type() != expected.data_type() {
+                            diffs.push(format!("  column '{}' (index {i}): expected type {:?}, got {:?}", expected.name(), expected.data_type(), actual.data_type()));
+                        }
+                        if actual.is_nullable() != expected.is_nullable() {
+                            diffs.push(format!("  column '{}' (index {i}): expected nullable={}, got nullable={}", expected.name(), expected.is_nullable(), actual.is_nullable()));
+                        }
+                    }
+                    None => {
+                        diffs.push(format!("  column '{}' (index {i}): missing from batch", expected.name()));
+                    }
+                    _ => {}
+                }
+            }
+            for i in expected_fields.len()..actual_fields.len() {
+                diffs.push(format!("  column '{}' (index {i}): unexpected extra column in batch", actual_fields[i].name()));
+            }
+            if expected_fields.len() != actual_fields.len() {
+                diffs.push(format!("  expected {} columns, got {}", expected_fields.len(), actual_fields.len()));
+            }
+
             anyhow::bail!(
-                "Schema mismatch for table '{table}': expected {expected_schema}, got {}",
-                batch.schema()
+                "Schema mismatch for table '{table}':\n{}",
+                diffs.join("\n")
             );
         }
 
@@ -171,7 +249,7 @@ pub trait Dataset: Send + Sync {
 
 #[async_trait]
 impl Dataset for Arc<dyn Dataset> {
-    fn batch_ids(&self, table: &str) -> Vec<u64> {
+    fn batch_ids(&self, table: &str) -> VecDeque<u64> {
         (**self).batch_ids(table)
     }
 

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

@@ -62,6 +62,13 @@ impl SimpleSequenceDataset {
 
 #[async_trait]
 impl Dataset for SimpleSequenceDataset {
+    fn create(config: &DatasetConfig) -> anyhow::Result<Arc<dyn Dataset>>
+    where
+        Self: Sized + 'static,
+    {
+        Ok(Arc::new(Self::new(config)))
+    }
+
     fn num_batches(&self, _table: &str) -> u64 {
         // One batch per step for the single table.
         u64::from(self.num_steps)

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

@@ -46,80 +46,80 @@ const TPCH_TABLE_TIME_COLUMNS: &[(&str, &str)] = &[
 fn tpch_schema(table: &str) -> SchemaRef {
     let fields: Vec<Field> = match table {
         "region" => vec![
-            Field::new("r_regionkey", DataType::Int32, false),
-            Field::new("r_name", DataType::Utf8, false),
+            Field::new("r_regionkey", DataType::Int32, true),
+            Field::new("r_name", DataType::Utf8, true),
             Field::new("r_comment", DataType::Utf8, true),
         ],
         "nation" => vec![
-            Field::new("n_nationkey", DataType::Int32, false),
-            Field::new("n_name", DataType::Utf8, false),
-            Field::new("n_regionkey", DataType::Int32, false),
+            Field::new("n_nationkey", DataType::Int32, true),
+            Field::new("n_name", DataType::Utf8, true),
+            Field::new("n_regionkey", DataType::Int32, true),
             Field::new("n_comment", DataType::Utf8, true),
         ],
         "supplier" => vec![
-            Field::new("s_suppkey", DataType::Int32, false),
-            Field::new("s_name", DataType::Utf8, false),
-            Field::new("s_address", DataType::Utf8, false),
-            Field::new("s_nationkey", DataType::Int32, false),
-            Field::new("s_phone", DataType::Utf8, false),
-            Field::new("s_acctbal", DataType::Decimal128(15, 2), false),
+            Field::new("s_suppkey", DataType::Int64, true),
+            Field::new("s_name", DataType::Utf8, true),
+            Field::new("s_address", DataType::Utf8, true),
+            Field::new("s_nationkey", DataType::Int32, true),
+            Field::new("s_phone", DataType::Utf8, true),
+            Field::new("s_acctbal", DataType::Decimal128(15, 2), true),
             Field::new("s_comment", DataType::Utf8, true),
         ],
         "customer" => vec![
-            Field::new("c_custkey", DataType::Int32, false),
-            Field::new("c_name", DataType::Utf8, false),
-            Field::new("c_address", DataType::Utf8, false),
-            Field::new("c_nationkey", DataType::Int32, false),
-            Field::new("c_phone", DataType::Utf8, false),
-            Field::new("c_acctbal", DataType::Decimal128(15, 2), false),
-            Field::new("c_mktsegment", DataType::Utf8, false),
+            Field::new("c_custkey", DataType::Int64, true),
+            Field::new("c_name", DataType::Utf8, true),
+            Field::new("c_address", DataType::Utf8, true),
+            Field::new("c_nationkey", DataType::Int32, true),
+            Field::new("c_phone", DataType::Utf8, true),
+            Field::new("c_acctbal", DataType::Decimal128(15, 2), true),
+            Field::new("c_mktsegment", DataType::Utf8, true),
             Field::new("c_comment", DataType::Utf8, true),
         ],
         "part" => vec![
-            Field::new("p_partkey", DataType::Int32, false),
-            Field::new("p_name", DataType::Utf8, false),
-            Field::new("p_mfgr", DataType::Utf8, false),
-            Field::new("p_brand", DataType::Utf8, false),
-            Field::new("p_type", DataType::Utf8, false),
-            Field::new("p_size", DataType::Int32, false),
-            Field::new("p_container", DataType::Utf8, false),
-            Field::new("p_retailprice", DataType::Decimal128(15, 2), false),
+            Field::new("p_partkey", DataType::Int64, true),
+            Field::new("p_name", DataType::Utf8, true),
+            Field::new("p_mfgr", DataType::Utf8, true),
+            Field::new("p_brand", DataType::Utf8, true),
+            Field::new("p_type", DataType::Utf8, true),
+            Field::new("p_size", DataType::Int32, true),
+            Field::new("p_container", DataType::Utf8, true),
+            Field::new("p_retailprice", DataType::Decimal128(15, 2), true),
             Field::new("p_comment", DataType::Utf8, true),
         ],
         "partsupp" => vec![
-            Field::new("ps_partkey", DataType::Int32, false),
-            Field::new("ps_suppkey", DataType::Int32, false),
-            Field::new("ps_availqty", DataType::Int32, false),
-            Field::new("ps_supplycost", DataType::Decimal128(15, 2), false),
+            Field::new("ps_partkey", DataType::Int64, true),
+            Field::new("ps_suppkey", DataType::Int64, true),
+            Field::new("ps_availqty", DataType::Int64, true),
+            Field::new("ps_supplycost", DataType::Decimal128(15, 2), true),
             Field::new("ps_comment", DataType::Utf8, true),
         ],
         "orders" => vec![
-            Field::new("o_orderkey", DataType::Int32, false),
-            Field::new("o_custkey", DataType::Int32, false),
-            Field::new("o_orderstatus", DataType::Utf8, false),
-            Field::new("o_totalprice", DataType::Decimal128(15, 2), false),
-            Field::new("o_orderdate", DataType::Date32, false),
-            Field::new("o_orderpriority", DataType::Utf8, false),
-            Field::new("o_clerk", DataType::Utf8, false),
-            Field::new("o_shippriority", DataType::Int32, false),
+            Field::new("o_orderkey", DataType::Int64, true),
+            Field::new("o_custkey", DataType::Int64, true),
+            Field::new("o_orderstatus", DataType::Utf8, true),
+            Field::new("o_totalprice", DataType::Decimal128(15, 2), true),
+            Field::new("o_orderdate", DataType::Date32, true),
+            Field::new("o_orderpriority", DataType::Utf8, true),
+            Field::new("o_clerk", DataType::Utf8, true),
+            Field::new("o_shippriority", DataType::Int32, true),
             Field::new("o_comment", DataType::Utf8, true),
         ],
         "lineitem" => vec![
-            Field::new("l_orderkey", DataType::Int32, false),
-            Field::new("l_partkey", DataType::Int32, false),
-            Field::new("l_suppkey", DataType::Int32, false),
-            Field::new("l_linenumber", DataType::Int32, false),
-            Field::new("l_quantity", DataType::Decimal128(15, 2), false),
-            Field::new("l_extendedprice", DataType::Decimal128(15, 2), false),
-            Field::new("l_discount", DataType::Decimal128(15, 2), false),
-            Field::new("l_tax", DataType::Decimal128(15, 2), false),
-            Field::new("l_returnflag", DataType::Utf8, false),
-            Field::new("l_linestatus", DataType::Utf8, false),
-            Field::new("l_shipdate", DataType::Date32, false),
-            Field::new("l_commitdate", DataType::Date32, false),
-            Field::new("l_receiptdate", DataType::Date32, false),
-            Field::new("l_shipinstruct", DataType::Utf8, false),
-            Field::new("l_shipmode", DataType::Utf8, false),
+            Field::new("l_orderkey", DataType::Int64, true),
+            Field::new("l_partkey", DataType::Int64, true),
+            Field::new("l_suppkey", DataType::Int64, true),
+            Field::new("l_linenumber", DataType::Int64, true),
+            Field::new("l_quantity", DataType::Decimal128(15, 2), true),
+            Field::new("l_extendedprice", DataType::Decimal128(15, 2), true),
+            Field::new("l_discount", DataType::Decimal128(15, 2), true),
+            Field::new("l_tax", DataType::Decimal128(15, 2), true),
+            Field::new("l_returnflag", DataType::Utf8, true),
+            Field::new("l_linestatus", DataType::Utf8, true),
+            Field::new("l_shipdate", DataType::Date32, true),
+            Field::new("l_commitdate", DataType::Date32, true),
+            Field::new("l_receiptdate", DataType::Date32, true),
+            Field::new("l_shipinstruct", DataType::Utf8, true),
+            Field::new("l_shipmode", DataType::Utf8, true),
             Field::new("l_comment", DataType::Utf8, true),
         ],
         _ => unreachable!("unknown TPC-H table: {table}"),
@@ -214,13 +214,21 @@ impl TpchDataset {
 
 #[async_trait]
 impl Dataset for TpchDataset {
+    fn create(config: &DatasetConfig) -> anyhow::Result<Arc<dyn Dataset>>
+    where
+        Self: Sized + 'static,
+    {
+        Ok(Arc::new(Self::new(config)?))
+    }
+
     fn num_batches(&self, table: &str) -> u64 {
         if !TPCH_TABLE_TIME_COLUMNS
             .iter()
             .any(|(name, _)| *name == table)
         {
             return 0;
         }
+
         // TPC-H produces one batch per table per step.
         u64::from(self.num_steps)
     }