Streaming for the CSV loader

src/datasets/packaged_modules/csv/csv.py

@@ -137,15 +137,16 @@ def _generate_tables(self, files):
         # dtype allows reading an int column as str
         dtype = {name: dtype.to_pandas_dtype() for name, dtype in zip(schema.names, schema.types)} if schema else None
         for file_idx, file in enumerate(files):
-            csv_file_reader = pd.read_csv(file, iterator=True, dtype=dtype, **self.config.read_csv_kwargs)
-
-            try:
-                for batch_idx, df in enumerate(csv_file_reader):
-                    pa_table = pa.Table.from_pandas(df, schema=schema)
-                    # Uncomment for debugging (will print the Arrow table size and elements)
-                    # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
-                    # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
-                    yield (file_idx, batch_idx), pa_table
-            except ValueError as e:
-                logger.error(f"Failed to read file '{csv_file_reader.f}' with error {type(e)}: {e}")
-                raise
+            with open(file, "rb") as f:
+                csv_file_reader = pd.read_csv(f, iterator=True, dtype=dtype, **self.config.read_csv_kwargs)
+
+                try:
+                    for batch_idx, df in enumerate(csv_file_reader):
+                        pa_table = pa.Table.from_pandas(df, schema=schema)
+                        # Uncomment for debugging (will print the Arrow table size and elements)
+                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
+                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
+                        yield (file_idx, batch_idx), pa_table
+                except ValueError as e:
+                    logger.error(f"Failed to read file '{csv_file_reader.f}' with error {type(e)}: {e}")
+                    raise

tests/test_packaged_modules.py

@@ -1,3 +1,4 @@
+import os
 import textwrap
 
 import pyarrow as pa
@@ -60,7 +61,9 @@ def test_csv_generate_tables_raises_error_with_malformed_csv(csv_file, malformed
         for _ in generator:
             pass
     assert any(
-        record.levelname == "ERROR" and f"Failed to read file '{malformed_csv_file}'" in record.message
+        record.levelname == "ERROR"
+        and f"Failed to read file" in record.message
+        and os.path.basename(malformed_csv_file) in record.message
         for record in caplog.records
     )