Logging

Author: krinart

crates/checkpointer/src/main.rs

@@ -86,6 +86,10 @@ struct Cli {
     /// Directory to write checkpoint parquet files into
     #[arg(long, default_value = "./checkpoints")]
     checkpoint_dir: PathBuf,
+
+    /// Local directory to check for pre-extracted data before downloading from S3
+    #[arg(long, default_value = "./spicebench-data")]
+    data_dir: PathBuf,
 }
 
 #[cfg(feature = "duckdb")]
@@ -105,6 +109,32 @@ impl Cli {
     }
 }
 
+/// Log the row count for each table in the DuckDB sink.
+#[cfg(feature = "duckdb")]
+async fn log_table_row_counts(
+    sink: &DuckDBSink,
+    table_names: &[String],
+    checkpoint_idx: usize,
+) -> anyhow::Result<()> {
+    for table in table_names {
+        let sql = format!("SELECT COUNT(*) AS cnt FROM {table}");
+        let batches = sink.query(&sql).await?;
+        let count: i64 = batches
+            .first()
+            .and_then(|b| {
+                b.column(0)
+                    .as_any()
+                    .downcast_ref::<arrow::array::Int64Array>()
+                    .map(|a| a.value(0))
+            })
+            .unwrap_or(0);
+        if table == "lineitem" {
+            eprintln!("[checkpoint] Checkpoint {checkpoint_idx} | {table}: {count} rows");
+        }
+    }
+    Ok(())
+}
+
 /// Run all checkpoint queries against the DuckDB sink and write each result
 /// set to a parquet file at `<checkpoint_dir>/<checkpoint_idx>/<query_idx>.parquet`.
 #[cfg(feature = "duckdb")]
@@ -198,20 +228,34 @@ async fn main() -> anyhow::Result<()> {
 
         let source_config = cli.source_config();
         let version_prefix = source_config.prefix.clone();
-        let archive_storage: Arc<dyn DataStorage> = Arc::new(S3Storage::new(&source_config)?);
 
-        // Download and extract the archive to a temporary local directory.
-        let extract_dir = tempfile::tempdir()?;
-        ETLPipeline::download(archive_storage, extract_dir.path()).await?;
-        let source: Arc<dyn DataStorage> = Arc::new(FileStorage::new(extract_dir.path()));
+        // Check local directory first; download from S3 only if version.json is missing.
+        let (_extract_dir_handle, extract_path);
+        let version_json_path = cli.data_dir.join("version.json");
+        if version_json_path.exists() {
+            tracing::info!(
+                data_dir = %cli.data_dir.display(),
+                "Using pre-existing local data, skipping S3 download"
+            );
+            _extract_dir_handle = None;
+            extract_path = cli.data_dir.clone();
+        } else {
+            tracing::info!("Local data not found at {}, downloading from S3", cli.data_dir.display());
+            let archive_storage: Arc<dyn DataStorage> = Arc::new(S3Storage::new(&source_config)?);
+            let tmp = tempfile::tempdir()?;
+            ETLPipeline::download(archive_storage, tmp.path()).await?;
+            extract_path = tmp.path().to_path_buf();
+            _extract_dir_handle = Some(tmp);
+        }
+        let source: Arc<dyn DataStorage> = Arc::new(FileStorage::new(&extract_path));
 
         // Read version metadata to derive dataset config and mutations.
         let version_metadata = source.read_version_metadata().await?.ok_or_else(|| {
-        anyhow::anyhow!(
-            "No version.json found in extracted data at {}. Was data generation run for this version?",
-            extract_dir.path().display()
-        )
-    })?;
+            anyhow::anyhow!(
+                "No version.json found in extracted data at {}. Was data generation run for this version?",
+                extract_path.display()
+            )
+        })?;
 
         let dataset_source = DatasetSource::from_dataset_type(&version_metadata.dataset_type)?;
         let dataset_config = version_metadata.dataset_config();
@@ -234,7 +278,7 @@ async fn main() -> anyhow::Result<()> {
             bucket = %cli.bucket,
             prefix = %cli.prefix,
             version_prefix = %version_prefix,
-            extract_dir = %extract_dir.path().display(),
+            extract_dir = %extract_path.display(),
             duckdb_path = %cli.duckdb_path.display(),
             scale_factor = version_metadata.scale_factor,
             num_steps = version_metadata.num_steps,
@@ -243,6 +287,9 @@ async fn main() -> anyhow::Result<()> {
             "Starting Checkpointer"
         );
 
+        let mut table_names: Vec<String> = version_metadata.tables.keys().cloned().collect();
+        table_names.sort();
+
         pipeline.initialize().await?;
         pipeline.run(cli.checkpoint_interval_steps as usize).await?;
 
@@ -258,6 +305,7 @@ async fn main() -> anyhow::Result<()> {
                         checkpoint = checkpoint_idx,
                         "Pipeline paused, running checkpoint queries"
                     );
+                    log_table_row_counts(&target, &table_names, checkpoint_idx).await?;
                     run_checkpoint_queries(
                         &target,
                         &checkpoint_queries,
@@ -276,6 +324,7 @@ async fn main() -> anyhow::Result<()> {
                         checkpoint = checkpoint_idx,
                         "Pipeline completed, running final checkpoint queries"
                     );
+                    log_table_row_counts(&target, &table_names, checkpoint_idx).await?;
                     run_checkpoint_queries(
                         &target,
                         &checkpoint_queries,

crates/data-generation/src/storage/file.rs

@@ -157,6 +157,15 @@ impl DataStorage for FileStorage {
         .await
         .map_err(|e| anyhow::anyhow!("spawn_blocking panicked reading parquet: {e}"))??;
 
+        if table_name == "lineitem"
+            && let Some(result) = result.as_ref()
+        {
+            eprintln!(
+                "[etl-read] table={table_name} batch_id={batch_id} rows={}",
+                result.rows_read,
+            );
+        }
+
         Ok(result)
     }
 

crates/etl/src/lib.rs

@@ -1791,6 +1791,14 @@ async fn run_pipeline(
                         }
                     };
 
+                let source_rows: usize = source_batches.iter().map(|b| b.num_rows()).sum();
+                if table_name == "lineitem" {
+                    eprintln!(
+                        "!!!!!!!!!!!!!!! [etl] table={table_name} batch_id={batch_id} source_batches={} source_rows={source_rows} consumed_work_units={consumed_work_units}",
+                        source_batches.len(),
+                    );
+                }
+
                 if source_batches.is_empty() {
                     debug!(
                         table = %table_name,

crates/etl/src/sink/adbc.rs

@@ -15,6 +15,7 @@ limitations under the License.
 */
 
 use std::collections::HashMap;
+use std::sync::Mutex;
 use std::time::Instant;
 
 use adbc_client::{
@@ -88,6 +89,7 @@ pub struct AdbcSink {
     pool: AdbcConnectionPool,
     target_db_catalog: Option<String>,
     target_db_schema: Option<String>,
+    row_counts: Mutex<HashMap<String, u64>>,
     /// Character used to quote SQL identifiers (e.g. '"' for ANSI, '`' for Databricks).
     identifier_quote_char: char,
     /// Whether Int64/UInt64 literals need an `L` suffix (Databricks).
@@ -130,6 +132,7 @@ impl AdbcSink {
             pool,
             target_db_catalog,
             target_db_schema,
+            row_counts: Mutex::new(HashMap::new()),
             identifier_quote_char,
             bigint_suffix,
         })
@@ -788,6 +791,7 @@ impl AdbcSink {
              WHEN MATCHED THEN UPDATE SET {set_clause}"
         )
     }
+
 }
 
 #[async_trait]
@@ -809,6 +813,18 @@ impl Sink for AdbcSink {
             .get()
             .map_err(|e| anyhow::anyhow!("Failed to get ADBC connection from pool: {e}"))?;
 
+        let rows_current = batch.num_rows() as u64;
+        let op_label = match &op {
+            InsertOp::Insert => "insert",
+            InsertOp::Update { .. } => "update",
+            InsertOp::Delete { .. } => "delete",
+        };
+
+        let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S%.3f UTC");
+        // if table_name == "lineitem" {
+        eprintln!("[adbc-write] {now} | {table_name} | {op_label} | rows: {rows_current}");
+        // }
+
         match op {
             InsertOp::Insert => {
                 self.ingest_insert_batch(&mut conn, table_name, batch)?;
@@ -882,6 +898,24 @@ impl Sink for AdbcSink {
             }
         }
 
+        let rows_total = {
+            let mut counts = self.row_counts.lock().unwrap();
+            let total = counts.entry(table_name.to_string()).or_insert(0);
+            match op_label {
+                "insert" => *total += rows_current,
+                "delete" => *total = total.saturating_sub(rows_current),
+                _ => {} // updates don't change row count
+            }
+            *total
+        };
+
+        let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S%.3f UTC");
+        // if table_name == "lineitem" {
+        eprintln!(
+            "[adbc] WRITTEN {now} | {table_name} | {op_label} | rows: {rows_current} | total: {rows_total}"
+        );
+        // }
+
         Ok(())
     }
 }

crates/etl/src/sink/duckdb.rs

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::path::Path;
 use std::sync::{Arc, Mutex};
 
@@ -35,6 +35,7 @@ use super::{InsertOp, Sink};
 pub struct DuckDBSink {
     conn: Arc<Mutex<duckdb::Connection>>,
     created_tables: TokioMutex<HashSet<String>>,
+    row_counts: Mutex<HashMap<String, u64>>,
 }
 
 impl DuckDBSink {
@@ -45,6 +46,7 @@ impl DuckDBSink {
         Ok(Self {
             conn: Arc::new(Mutex::new(conn)),
             created_tables: TokioMutex::new(HashSet::new()),
+            row_counts: Mutex::new(HashMap::new()),
         })
     }
 
@@ -390,6 +392,25 @@ impl Sink for DuckDBSink {
             created.insert(table_name.to_string());
         }
 
+        let rows_current = num_rows as u64;
+        let op_label = match &op {
+            InsertOp::Insert => "insert",
+            InsertOp::Update { .. } => "update",
+            InsertOp::Delete { .. } => "delete",
+        };
+        let rows_total = {
+            let mut counts = self.row_counts.lock().unwrap();
+            let total = counts.entry(table_name.to_string()).or_insert(0);
+            *total += rows_current;
+            *total
+        };
+        let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S%.3f UTC");
+        if table_name == "lineitem" {
+            eprintln!(
+                "[duckdb] WRITTEN {now} | {table_name} | {op_label} | rows: {rows_current} | total: {rows_total}"
+            );
+        }
+
         Ok(())
     }
 }