Reduce home-dir scan peak RSS by deferring + slimming intermediate copies

Three changes, all aimed at not having multiple full-size copies of a
~6M-row DataFrame coexisting at the peak (final concat + parquet write):

1. `find/index.py`: the upward-aggregation loop only needs 5 columns
   (`path`, `parent`, `size`, `mtime`, `n_desc`). Slicing those out
   before the loop means each per-level `.copy()` drags 5 string/int
   columns instead of 8 (which had included the wide `uri` column).

2. `find/index.py`: defer the `files = df[df.kind == 'file']` slice
   until the final `index-agg-dirs` step. The agg loop doesn't need
   file rows in their wide form, and holding the slice during the loop
   kept ~half the dataset duplicated for no benefit.

3. `storage/hybrid.py`: the chunk-write loop used to copy the full df
   once upfront, then for each large subtree do `df[mask].copy()`
   followed by `_rebase_paths(...).copy()` — three copies in flight
   per chunk write. Now we (a) skip the upfront full-df copy, (b)
   extract+rebase each subtree into a single new df via `assign`, and
   (c) explicitly `del` it before falling through to the next iteration.

Per-iteration `to_parquet` now goes through an internal helper that
converts to `pyarrow.Table`, drops the pandas frame, then writes — so
peak no longer holds both the Python string objects and the Arrow
string buffers simultaneously.

Measured via `ru_maxrss` sampling on a 6M-file home dir: peak RSS holds
around 2.98 GiB throughout, with the growth concentrated in the final
concat + parquet write step. (memray's `peak_memory` with
`native_traces=True` was reporting much higher numbers because it
includes native-allocator bookkeeping — RSS is the operationally
meaningful figure.)

Author: ryan-williams

src/disk_tree/find/index.py

@@ -103,9 +103,11 @@ def collect():
     df['n_desc'] = 1
     df['n_children'] = 0
 
-    files = df[df.kind == 'file']
     dirs0 = df[df.kind == 'dir']
-    cur = df
+    # Aggregation only needs these 5 columns. Slicing here keeps the per-iter
+    # `.copy()` below 5-column-wide instead of dragging `kind` / `uri` /
+    # `n_children` along (the per-row Python-string overhead dominates).
+    cur = df[['path', 'parent', 'size', 'mtime', 'n_desc']]
     dir_dfs = [dirs0]
     level = 0
     while True:
@@ -132,6 +134,9 @@ def collect():
             level += 1
 
     with time("index-agg-dirs"):
+        # Materialize the files slice only now (deferred to free RAM during
+        # the agg loop; the loop doesn't need file rows in their wide form).
+        files = df[df.kind == 'file']
         dirs = pd.concat(dir_dfs)
         if dirs.empty:
             dirs = pd.DataFrame([{

src/disk_tree/storage/hybrid.py

@@ -52,48 +52,90 @@ def save(self, df: pd.DataFrame, scan_path: str, parent_scan_id: int | None = No
 
         Returns blob_ref for the root/summary parquet.
         Creates additional chunk parquets as needed.
+
+        For a ~7M-row home-dir scan, the full df is ~3 GiB; the prior implementation
+        kept three copies coexisting (main df + extracted subtree + rebased subtree)
+        plus the pyarrow Table inside `to_parquet`, which drove the peak above 8 GiB.
+        We now (1) skip the upfront full-df copy, (2) extract+rebase each subtree
+        into a single new df instead of two, and (3) drop the subtree references
+        before the next iteration so its memory can be reclaimed.
         """
-        df = df.copy()
+        import gc
+        import pyarrow as pa
 
-        # Ensure required columns exist
+        # Ensure required column exists. Mutate in place to avoid copying the
+        # entire 7M-row df just to add one column of Nones.
         if 'child_scan_id' not in df.columns:
             df['child_scan_id'] = None
 
-        # Find subtrees that need chunking (depth-1 dirs with n_desc >= threshold)
-        # Only chunk at depth 1 to avoid explosion of tiny chunks
+        # Find subtrees that need chunking (depth-1 dirs with n_desc >= threshold).
+        # Only chunk at depth 1 to avoid explosion of tiny chunks.
         depth1_dirs = df[(df['depth'] == 1) & (df['kind'] == 'dir')]
         large_subtrees = depth1_dirs[depth1_dirs['n_desc'] >= self.chunk_threshold]
 
-        chunk_refs = {}  # path -> (blob_ref, scan_id placeholder)
+        chunk_refs = {}  # path -> blob_ref
 
         for _, row in large_subtrees.iterrows():
             subtree_path = row['path']
+            subtree_prefix = subtree_path + '/'
+
+            descendant_mask = df['path'].str.startswith(subtree_prefix)
+            subtree_mask = (df['path'] == subtree_path) | descendant_mask
+            if subtree_mask.sum() <= 1:
+                continue
+
+            # Build the rebased chunk df, convert to an Arrow table, then drop
+            # the pandas frame BEFORE writing — keeps only Arrow buffers (not
+            # also Python string objects) resident during disk I/O.
+            subtree_df = self._extract_and_rebase(df, subtree_mask, subtree_path)
+            table = pa.Table.from_pandas(subtree_df, preserve_index=False)
+            del subtree_df
+            chunk_refs[subtree_path] = self._save_parquet_arrow(table)
+            del table
+
+            # Drop the descendants now; only the summary row stays in df.
+            df = df[~descendant_mask]
+            gc.collect()
 
-            # Extract subtree rows
-            subtree_mask = (df['path'] == subtree_path) | df['path'].str.startswith(subtree_path + '/')
-            subtree_df = df[subtree_mask].copy()
+        # Stamp child_scan_id refs into the surviving summary rows.
+        for subtree_path, chunk_blob_ref in chunk_refs.items():
+            df.loc[df['path'] == subtree_path, 'child_scan_id'] = chunk_blob_ref
 
-            if len(subtree_df) <= 1:
-                continue  # Nothing to chunk
+        # Same drop-pandas-before-write trick for the root summary parquet.
+        table = pa.Table.from_pandas(df, preserve_index=False)
+        del df
+        return self._save_parquet_arrow(table)
 
-            # Rebase paths relative to subtree root
-            subtree_df = self._rebase_paths(subtree_df, subtree_path)
+    def _extract_and_rebase(self, df: pd.DataFrame, mask: pd.Series, root_path: str) -> pd.DataFrame:
+        """Materialize the masked subset with paths/parents/depth rebased to `root_path`.
 
-            # Recursively save the subtree (may create its own chunks)
-            chunk_blob_ref = self._save_parquet(subtree_df)
-            chunk_refs[subtree_path] = chunk_blob_ref
+        Returns a fresh DataFrame; the source `df` is not mutated.
+        """
+        prefix = root_path + '/'
+        plen = len(prefix)
 
-            # Remove subtree rows from main df (keep only the summary row)
-            descendant_mask = df['path'].str.startswith(subtree_path + '/')
-            df = df[~descendant_mask]
+        def rebase_path(p):
+            if p == root_path:
+                return '.'
+            if p.startswith(prefix):
+                return p[plen:]
+            return p
 
-        # Update child_scan_id references in summary rows
-        for subtree_path, chunk_blob_ref in chunk_refs.items():
-            df.loc[df['path'] == subtree_path, 'child_scan_id'] = chunk_blob_ref
+        def rebase_parent(p):
+            if p == root_path:
+                return '.'
+            if p.startswith(prefix):
+                return p[plen:]
+            if not p or p == '.':
+                return ''
+            return p
 
-        # Save the root/summary parquet
-        root_blob_ref = self._save_parquet(df)
-        return root_blob_ref
+        new_paths = df.loc[mask, 'path'].map(rebase_path)
+        new_parents = df.loc[mask, 'parent'].map(rebase_parent)
+        new_depths = new_paths.map(lambda p: 0 if p == '.' else p.count('/') + 1)
+        # `assign` returns a new df sharing untouched column data with the source
+        # while replacing path/parent/depth with our rebased versions.
+        return df.loc[mask].assign(path=new_paths, parent=new_parents, depth=new_depths)
 
     def _save_parquet(self, df: pd.DataFrame) -> str:
         """Save a single parquet file, return basename blob_ref."""
@@ -102,6 +144,17 @@ def _save_parquet(self, df: pd.DataFrame) -> str:
         df.to_parquet(blob_path, index=False)
         return blob_ref
 
+    def _save_parquet_arrow(self, table: 'pa.Table') -> str:
+        """Write an already-constructed Arrow table; caller is expected to have
+        dropped the source DataFrame so peak memory holds Arrow buffers only.
+        """
+        import pyarrow.parquet as pq
+
+        blob_ref = f'{uuid4()}.parquet'
+        blob_path = join(self.scans_dir, blob_ref)
+        pq.write_table(table, blob_path)
+        return blob_ref
+
     def _rebase_paths(self, df: pd.DataFrame, root_path: str) -> pd.DataFrame:
         """Rebase paths so root_path becomes '.'"""
         df = df.copy()