fix: streaming load_from_db with fetchmany(10000) instead of fetchall() (fixes OOM on 4M rows)

analysis/data_loader.py

@@ -5,6 +5,7 @@ point is :func:`load_csv_directory`, which globs files, detects BOM per file,
 validates schema consistency, and returns a merged :class:`pl.LazyFrame`.
 """
 import glob
+import io
 import logging
 import os
 import sqlite3
@@ -269,6 +270,22 @@ def _str_to_polars_dtype(s: str) -> pl.DataType:
     return _STR_TO_DTYPE.get(base, pl.Utf8)
 
 
+def _rows_to_csv_df(rows: list[sqlite3.Row]) -> pl.DataFrame:
+    """Convert a batch of sqlite3.Row objects to a DataFrame via CSV fallback.
+
+    Loads all columns as Utf8 to handle mixed-type SQLite columns.
+    """
+    import csv as _csv_mod
+
+    buf = io.StringIO()
+    writer = _csv_mod.writer(buf)
+    writer.writerow(rows[0].keys())
+    for r in rows:
+        writer.writerow(str(v) if v is not None else '' for v in r)
+    buf.seek(0)
+    return pl.read_csv(buf, infer_schema_length=0)
+
+
 def load_from_db(
     table_name: str,
     schema_overrides: dict[str, pl.DataType | str] | None = None,
@@ -300,22 +317,50 @@ def load_from_db(
             return None
 
         conn.row_factory = sqlite3.Row
-        rows = conn.execute(f'SELECT * FROM "{table_name}"').fetchall()
+        cursor = conn.execute(f'SELECT * FROM "{table_name}"')
-        if not rows:
+
+        # Peek first row to detect if table is empty
+        first_row = cursor.fetchone()
+        if first_row is None:
             return pl.DataFrame().lazy()
-        try:
+
-            df = pl.from_dicts([dict(r) for r in rows])
+        batch_rows = [first_row]
-        except Exception:
+        dfs = []
-            # Fallback: load all columns as Utf8 to handle mixed-type SQLite columns,
+        _use_csv_fallback = False
-            # then apply schema_overrides to restore intended types.
+
-            import io, csv as _csv_mod
+        while True:
-            buf = io.StringIO()
+            chunk = cursor.fetchmany(10000)
-            writer = _csv_mod.writer(buf)
+            if not chunk:
-            writer.writerow(rows[0].keys())
+                # Flush remaining rows
-            for r in rows:
+                if batch_rows:
-                writer.writerow(str(v) if v is not None else '' for v in r)
+                    if _use_csv_fallback:
-            buf.seek(0)
+                        dfs.append(_rows_to_csv_df(batch_rows))
-            df = pl.read_csv(buf, infer_schema_length=0)
+                    else:
+                        try:
+                            dfs.append(pl.from_dicts([dict(r) for r in batch_rows]))
+                        except Exception:
+                            _use_csv_fallback = True
+                            dfs.append(_rows_to_csv_df(batch_rows))
+                break
+
+            batch_rows.extend(chunk)
+            if len(batch_rows) >= 10000:
+                if _use_csv_fallback:
+                    dfs.append(_rows_to_csv_df(batch_rows))
+                else:
+                    try:
+                        dfs.append(pl.from_dicts([dict(r) for r in batch_rows]))
+                    except Exception:
+                        _use_csv_fallback = True
+                        # Re-process current batch with CSV fallback
+                        dfs.append(_rows_to_csv_df(batch_rows))
+                batch_rows = []
+
+        # Concat all batches
+        if len(dfs) == 1:
+            df = dfs[0]
+        else:
+            df = pl.concat(dfs, how='vertical')
 
         # Apply type overrides to restore original column types
         if schema_overrides: