fix parquet to arrow script failed when number of samples is small (#301

) key_format becomes negative Co-authored-by: Romain Beaumont <[email protected]>
rom1504 · Jan 6, 2024 · 9fe1276 · 9fe1276
1 parent c4e6615
commit 9fe1276
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/clip_retrieval/clip_back_prepro/parquet_to_arrow.py b/clip_retrieval/clip_back_prepro/parquet_to_arrow.py
@@ -36,7 +36,7 @@ def parquet_to_arrow(parquet_folder, output_arrow_folder, columns_to_return):
     sink = None
     current_batch_count = 0
     batch_counter = 0
-    key_format = int(math.log10(number_samples / 10**10)) + 1
+    key_format = max(0, int(math.log10(number_samples / 10**10))) + 1
     for parquet_files in tqdm(files):
         if sink is None or current_batch_count > 10**10:
             if sink is not None: