Batch results file based on file size (#283)

* add ChunkMap dict wrapper class

* add batch size and chunk size args

* implement new chunk management

* estimate size from existing number of rows

* [copilot] estimate and write across multiple files

* replace get with defaultdfict

* make bytes per row a method of chunkmap

* adjust logic of chunking

* add tests

* fix char

* formatting

* fix typing

* change chunk size to output file size

* replace comment

* rename concat_batch_results

* formatting

* rearrange metrics in from_existing_results

* revert that last commit

* simplify file searching

* change most references to 'chunks' to 'output file' or 'chunked output file'

* store bytes per row in outputfilemap

* format

* fix bug

* fix test bug

* formatting

* change to megabytes

* switch to mb

* remove pre-emptive rotation and adjust limit

* adjust multiple chunk test

* format

* address single row bug

* fix naming

* Update CHANGELOG for version 2.2.0

Author: patricktnast

CHANGELOG.rst

@@ -1,3 +1,7 @@
+**2.2.0 - 01/02/26**
+
+  - Add output data file splitting
+
 **2.1.26 - 12/24/25**
 
   - Fix build failure due to unsafe import of private psutil types

setup.py

@@ -59,6 +59,7 @@
         "vivarium_dependencies[lint]",
         "types-setuptools",
         "types-psutil",
+        "pyarrow-stubs",
     ]
 
     test_requirements = [

src/vivarium_cluster_tools/psimulate/cli.py

@@ -50,6 +50,8 @@ def psimulate() -> None:
     cluster.with_hardware,
     redis_dbs.with_max_workers,
     redis_dbs.with_redis,
+    results.with_batch_size,
+    results.with_output_file_size,
     results.with_no_batch,
     results.backup_freq,
     cli_tools.with_verbose_and_pdb,
@@ -133,6 +135,8 @@ def run(
         ),
         max_workers=options["max_workers"],
         redis_processes=options["redis"],
+        batch_size=options["batch_size"],
+        output_file_size=options["output_file_size"],
         no_batch=options["no_batch"],
         backup_freq=options["backup_freq"],
         extra_args={
@@ -178,6 +182,8 @@ def restart(
         ),
         max_workers=options["max_workers"],
         redis_processes=options["redis"],
+        batch_size=options["batch_size"],
+        output_file_size=options["output_file_size"],
         no_batch=options["no_batch"],
         backup_freq=options["backup_freq"],
         extra_args={
@@ -238,6 +244,8 @@ def expand(
         ),
         max_workers=options["max_workers"],
         redis_processes=options["redis"],
+        batch_size=options["batch_size"],
+        output_file_size=options["output_file_size"],
         no_batch=options["no_batch"],
         backup_freq=options["backup_freq"],
         extra_args={
@@ -316,6 +324,8 @@ def test(
         ),
         max_workers=options["max_workers"],
         redis_processes=options["redis"],
+        batch_size=options["batch_size"],
+        output_file_size=options["output_file_size"],
         no_batch=options["no_batch"],
         backup_freq=options["backup_freq"],
         extra_args={

src/vivarium_cluster_tools/psimulate/results/__init__.py

@@ -4,9 +4,15 @@
 ==================
 
 """
+
 from vivarium_cluster_tools.psimulate.results.cli_options import (
     backup_freq,
+    with_batch_size,
     with_no_batch,
     with_no_cleanup,
+    with_output_file_size,
+)
+from vivarium_cluster_tools.psimulate.results.processing import (
+    OutputFileMap,
+    write_results_batch,
 )
-from vivarium_cluster_tools.psimulate.results.processing import write_results_batch

src/vivarium_cluster_tools/psimulate/results/cli_options.py

@@ -6,12 +6,34 @@
 Command line options for configuring results handling in psimulate runs.
 
 """
+
 import click
 
 from vivarium_cluster_tools.cli_tools import MINUTES_OR_NONE
 
+DEFAULT_BATCH_SIZE = 200
+DEFAULT_OUTPUT_FILE_SIZE_MB = 500.0
+
+with_batch_size = click.option(
+    "--batch-size",
+    type=int,
+    default=DEFAULT_BATCH_SIZE,
+    show_default=True,
+    help="Number of simulation jobs to accumulate before writing results to disk.",
+)
+with_output_file_size = click.option(
+    "--output-file-size",
+    "output_file_size",
+    type=float,
+    default=DEFAULT_OUTPUT_FILE_SIZE_MB,
+    show_default=True,
+    help="Maximum file size in MB per result file (within each type of result). When exceeded, a new file is started.",
+    callback=lambda ctx, param, value: int(value * 1024 * 1024),  # Convert MB to bytes
+)
 with_no_batch = click.option(
-    "--no-batch", is_flag=True, help="Don't batch results, write them as they come in."
+    "--no-batch",
+    is_flag=True,
+    help="Write results immediately as they come in (batch_size=0). Files are still split by size per --output_file_size.",
 )
 with_no_cleanup = click.option(
     "--no-cleanup",