Fix file size determination for data generation script (#1497)

coiled · Mar 28, 2024 · f9eda14 · f9eda14
1 parent c400d36
commit f9eda14
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/tpch/generate_data.py b/tests/tpch/generate_data.py
@@ -273,8 +273,8 @@ def rows_approx_mb(con, table_name, partition_size: str, compression: Compressio
         tmp = pathlib.Path(tmpdir) / "tmp.parquet"
         stmt = f"select * from {table_name} limit {sample_size}"
         df = con.sql(stmt).arrow()
-        write_table(table_name, df, tmp, compression)
-        mb = tmp.stat().st_size
+        file = write_table(table_name, df, tmp, compression)
+        mb = pathlib.Path(file).stat().st_size
     return int(
         (len(table) * ((len(table) / sample_size) * partition_size)) / mb
     ) or len(table)