fix: repair OOM, drop broken table, better compression on fib traces

Author: dbirman

src/biodata_cache/__init__.py

@@ -25,7 +25,6 @@
 from biodata_cache.cache_table_helpers.platform_smartspim import assets_smartspim  # noqa: F401
 from biodata_cache.cache_table_helpers.qc import qc, qc_columns  # noqa: F401
 from biodata_cache.cache_table_helpers.raw_to_derived import raw_to_derived  # noqa: F401
-from biodata_cache.cache_table_helpers.scientist.scientist_rl_fib import scientist_rl_fib  # noqa: F401
 from biodata_cache.cache_table_helpers.source_data import source_data  # noqa: F401
 from biodata_cache.cache_table_helpers.time_to_qc import time_to_qc  # noqa: F401
 from biodata_cache.cache_table_helpers.unique_genotypes import (  # noqa: F401

src/biodata_cache/backend.py

@@ -8,6 +8,8 @@
 import boto3
 import duckdb
 import pandas as pd
+import pyarrow as pa
+import pyarrow.parquet as pq
 
 from biodata_cache.utils import BDC_VERSION, CacheLogMessage
 
@@ -89,7 +91,16 @@ def write(self, table_name: str, data: pd.DataFrame) -> None:
             json_key = f"{_CACHE_ROOT}/{_VERSION_FOLDER}/{table_name}.json"
 
         parquet_buffer = io.BytesIO()
-        data.to_parquet(parquet_buffer, index=False, compression="zstd")
+        table = pa.Table.from_pandas(data, preserve_index=False)
+        float_cols = [f.name for f in table.schema if pa.types.is_floating(f.type)]
+        dict_cols = [f.name for f in table.schema if f.name not in float_cols]
+        pq.write_table(
+            table,
+            parquet_buffer,
+            compression="zstd",
+            use_dictionary=dict_cols if dict_cols else False,
+            column_encoding={col: "BYTE_STREAM_SPLIT" for col in float_cols} or None,
+        )
         parquet_buffer.seek(0)
 
         self.s3_client.put_object(

src/biodata_cache/cache_table_helpers/__init__.py

@@ -14,7 +14,6 @@
     platform_smartspim,
     qc,
     raw_to_derived,
-    scientist,
     source_data,
     time_to_qc,
     unique_genotypes,

src/biodata_cache/cache_table_helpers/platform_fib_traces.py

@@ -179,6 +179,8 @@ def _fetch_subject_fib_traces(subject_id: str) -> pd.DataFrame:
             frames.append(session_df)
 
     df = pd.concat(frames, ignore_index=True) if frames else pd.DataFrame()
+    if not df.empty:
+        df = df.sort_values(["asset_name", "channel", "timestamp", "fiber"]).reset_index(drop=True)
 
     _log(f"Cached fib traces for subject {subject_id} ({len(frames)} sessions, {len(df)} samples)")
     registry.BACKEND.write(cache_key, df)

src/biodata_cache/cache_table_helpers/scientist/__init__.py

@@ -54,7 +54,6 @@
     "curriculum": "behavior_curriculum",
     "platform_qc": "platform_qc",
     "time_to_qc": "time_to_qc",
-    "scientist_rl_fib": "scientist_rl_fib",
     "mouselight": "platform_mouselight",
 }