feat: Include time column and table partitions in SAP DatasetConfig

Author: peasee

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

@@ -34,39 +34,26 @@ pub struct DatasetTable {
     pub name: String,
     /// The Arrow schema for the table (without the time column).
     pub schema: SchemaRef,
-    /// The time column for the table, if any.
+    /// The time column for the table.
     ///
-    /// When set, this column is *not* included in [`schema`] — it is appended
+    /// This column is *not* included in [`schema`] — it is appended
     /// during rehydration via [`DatasetTable::rehydrate`].
-    pub time_column: Option<String>,
+    pub time_column: String,
 }
 
 impl DatasetTable {
-    /// Returns the full schema including the time column, if one is configured.
-    ///
-    /// If `time_column` is `None`, this returns the same schema as [`schema`].
+    /// Returns the full schema including the time column.
     pub fn rehydrated_schema(&self) -> SchemaRef {
-        let Some(ref time_col) = self.time_column else {
-            return Arc::clone(&self.schema);
-        };
-
         let ts_type = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into()));
         let mut fields: Vec<_> = self.schema.fields().iter().cloned().collect();
-        fields.push(Arc::new(Field::new(time_col, ts_type, true)));
+        fields.push(Arc::new(Field::new(&self.time_column, ts_type, true)));
         Arc::new(arrow::datatypes::Schema::new(fields))
     }
 
     /// Rehydrate a batch by appending the time column with the current timestamp.
     ///
-    /// If this table has no `time_column`, the batch is returned unchanged.
     /// The batch schema must match [`schema`] (i.e. without the time column).
     pub fn rehydrate(&self, batch: &RecordBatch) -> anyhow::Result<RecordBatch> {
-        if self.time_column.is_none() {
-            anyhow::bail!(
-                "Cannot rehydrate table '{}' without a time column",
-                self.name
-            );
-        }
 
         if batch.schema() != self.schema {
             let mut diffs = Vec::new();
@@ -284,8 +271,7 @@ pub trait Dataset: Send + Sync {
     /// Rehydrate a batch for the given table by appending the time column.
     ///
     /// Uses the table metadata from [`tables()`] to look up the time column name
-    /// and delegates to [`DatasetTable::rehydrate`]. If the table has no time column,
-    /// a rehydration error is returned.
+    /// and delegates to [`DatasetTable::rehydrate`].
     fn rehydrate(&self, table: &str, batch: &RecordBatch) -> anyhow::Result<RecordBatch> {
         let tables = self.tables();
         let dataset_table = tables

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

@@ -102,7 +102,7 @@ impl Dataset for SimpleSequenceDataset {
             DatasetTable {
                 name: "integer_sequence".to_string(),
                 schema: Self::schema(),
-                time_column: Some("inserted_at".to_string()),
+                time_column: "inserted_at".to_string(),
             },
         )])
     }

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

@@ -323,7 +323,7 @@ impl Dataset for TpchDataset {
                     DatasetTable {
                         name: (*name).to_string(),
                         schema: tpch_schema(name),
-                        time_column: Some((*time_col).to_string()),
+                        time_column: (*time_col).to_string(),
                     },
                 )
             })

crates/data-generation/src/ingestor.rs

@@ -73,12 +73,10 @@ impl Ingestor {
                     "location".to_string(),
                     serde_json::Value::String(loc_fn(&name)),
                 );
-                if let Some(ref time_col) = table.time_column {
-                    entry.insert(
-                        "time_column".to_string(),
-                        serde_json::Value::String(time_col.clone()),
-                    );
-                }
+                entry.insert(
+                    "time_column".to_string(),
+                    serde_json::Value::String(table.time_column.clone()),
+                );
                 map.insert(name, serde_json::Value::Object(entry));
             }
             println!("{}", serde_json::Value::Object(map));

crates/etl/src/lib.rs

@@ -189,6 +189,8 @@ impl ETLPipeline {
                     etl_type: EtlType::S3,
                     schema: table.rehydrated_schema(),
                     params: self.target.table_params(&name),
+                    time_column: table.time_column.clone(),
+                    partitions: vec![], // TODO: support dynamically specifying partitioning schemes
                 };
                 (name, config)
             })

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

@@ -150,6 +150,10 @@ pub struct DatasetConfig {
     pub schema: SchemaRef,
     /// ETL-specific configuration parameters
     pub params: HashMap<String, serde_json::Value>,
+    /// The time column to use for append/change-stream capture for the dataset
+    pub time_column: String,
+    /// The table paritioning scheme to use for this dataset
+    pub partitions: Vec<String>,
 }
 
 /// Request to setup a benchmark run with ETL configuration