From c057693d9085872a04f2b84e7d2cb0eec93a73c5 Mon Sep 17 00:00:00 2001
From: Lars Bogner <upyxo@student.kit.edu>
Date: Mon, 22 Jun 2026 08:53:04 +0200
Subject: [PATCH] Add tqdm progress bar to predict

---
 giant/cli.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/giant/cli.py b/giant/cli.py
index 2fedc4b..9cc42a5 100644
--- a/giant/cli.py
+++ b/giant/cli.py
@@ -9,6 +9,7 @@ from typing_extensions import Annotated
 
 import pyarrow as pa
 import pyarrow.parquet as pq
+from tqdm import tqdm
 
 from giant import config as gconfig
 from giant.constants import (
@@ -343,8 +344,10 @@ def predict(
 
     writer: pq.ParquetWriter | None = None
     total = 0
+    total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files)
     chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks
 
+    bar = tqdm(total=total_rows, desc="predict", unit="row", dynamic_ncols=True)
     for path in files:
         for chunk in chunk_iter(path):
             N = len(chunk["event_id"])
@@ -455,7 +458,9 @@ def predict(
                 writer = pq.ParquetWriter(out, table.schema)
             writer.write_table(table)
             total += N
+            bar.update(N)
 
+    bar.close()
     if writer is not None:
         writer.close()