Add Parquet-based trajectory storage with ~25x compression

Migrate trajectory files from JSONL to Parquet format for improved
storage efficiency and query performance. Key changes:

- Add `art migrate` CLI command for migrating existing JSONL files
- Backend now writes trajectories as Parquet with ZSTD compression
- Auto-migrate JSONL files to Parquet on model registration
- Update benchmarking utilities to read Parquet using DuckDB
- Add trajectory_migration.py with legacy JSONL support and migration tools
- Add duckdb and pyarrow dependencies

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: Cursor Bot

pyproject.toml

@@ -32,6 +32,8 @@ backend = [
     "setuptools>=78.1.0",
     "wandb==0.22.1",
     "polars>=1.26.0",
+    "duckdb>=1.0.0",
+    "pyarrow>=15.0.0",
     "transformers==4.53.2",
     "trl==0.20.0",
     "nbclient>=0.10.1",
@@ -108,6 +110,8 @@ dev = [
     "pytest-xdist>=3.8.0",
     "pyright[nodejs]>=1.1.403",
     "pytest-asyncio>=1.1.0",
+    "duckdb>=1.0.0",
+    "pyarrow>=15.0.0",
 ]
 
 [tool.uv.sources]

src/art/cli.py

@@ -1,6 +1,7 @@
 import json
 import socket
-from typing import Any, AsyncIterator
+from pathlib import Path
+from typing import Any, AsyncIterator, Optional
 
 import pydantic
 import typer
@@ -26,6 +27,95 @@
 app = typer.Typer()
 
 
+@app.command()
+def migrate(
+    path: Path = typer.Argument(..., help="Path to model dir, project dir, or trajectories dir"),
+    dry_run: bool = typer.Option(False, "--dry-run", "-n", help="Show what would be migrated without making changes"),
+    keep_jsonl: bool = typer.Option(False, "--keep-jsonl", help="Keep original JSONL files after conversion"),
+    verbose: bool = typer.Option(False, "--verbose", "-v", help="Print progress for each file"),
+) -> None:
+    """
+    Migrate trajectory files from JSONL to Parquet format.
+
+    This command converts old .jsonl trajectory files to the new .parquet format,
+    which provides ~25x compression and ~20x faster queries.
+
+    Examples:
+        art migrate /path/to/.art/project/models/my-model
+        art migrate /path/to/.art/project --dry-run
+        art migrate /path/to/trajectories --keep-jsonl --verbose
+    """
+    from .utils.trajectory_migration import (
+        migrate_model_dir,
+        migrate_trajectories_dir,
+    )
+
+    if not path.exists():
+        typer.echo(f"Error: Path does not exist: {path}", err=True)
+        raise typer.Exit(1)
+
+    # Determine what kind of path this is
+    if (path / "trajectories").exists():
+        # This is a model directory
+        typer.echo(f"Migrating model directory: {path}")
+        result = migrate_model_dir(
+            path,
+            delete_originals=not keep_jsonl,
+            dry_run=dry_run,
+            progress_callback=lambda f: typer.echo(f"  {f}") if verbose else None,
+        )
+    elif path.name == "trajectories" or any(path.glob("*/[0-9]*.jsonl")):
+        # This is a trajectories directory
+        typer.echo(f"Migrating trajectories directory: {path}")
+        result = migrate_trajectories_dir(
+            path,
+            delete_originals=not keep_jsonl,
+            dry_run=dry_run,
+            progress_callback=lambda f: typer.echo(f"  {f}") if verbose else None,
+        )
+    elif (path / "models").exists():
+        # This is a project directory
+        typer.echo(f"Migrating project directory: {path}")
+        from .utils.trajectory_migration import MigrationResult
+
+        result = MigrationResult()
+        models_dir = path / "models"
+        for model_dir in models_dir.iterdir():
+            if model_dir.is_dir():
+                if verbose:
+                    typer.echo(f"Processing model: {model_dir.name}")
+                model_result = migrate_model_dir(
+                    model_dir,
+                    delete_originals=not keep_jsonl,
+                    dry_run=dry_run,
+                    progress_callback=lambda f: typer.echo(f"    {f}") if verbose else None,
+                )
+                result = result + model_result
+    else:
+        typer.echo(f"Error: Could not determine path type. Expected a model, project, or trajectories directory.", err=True)
+        raise typer.Exit(1)
+
+    # Print summary
+    if dry_run:
+        typer.echo(f"\n[DRY RUN] Would migrate {result.files_migrated} files")
+        if result.bytes_before > 0:
+            typer.echo(f"  Estimated space savings: {result.space_saved / 1024 / 1024:.1f} MB")
+    else:
+        typer.echo(f"\nMigrated {result.files_migrated} files")
+        if result.files_skipped > 0:
+            typer.echo(f"Skipped {result.files_skipped} files")
+        if result.bytes_before > 0 and result.bytes_after > 0:
+            typer.echo(f"Space saved: {result.space_saved / 1024 / 1024:.1f} MB ({result.compression_ratio:.1f}x compression)")
+
+    if result.errors:
+        typer.echo(f"\nErrors ({len(result.errors)}):", err=True)
+        for error in result.errors[:10]:
+            typer.echo(f"  {error}", err=True)
+        if len(result.errors) > 10:
+            typer.echo(f"  ... and {len(result.errors) - 10} more errors", err=True)
+        raise typer.Exit(1)
+
+
 @app.command()
 def run(host: str = "0.0.0.0", port: int = 7999) -> None:
     """Run the ART CLI."""

src/art/local/backend.py

@@ -43,7 +43,7 @@
     pull_model_from_s3,
     push_model_to_s3,
 )
-from art.utils.trajectory_logging import serialize_trajectory_groups
+from art.utils.trajectory_logging import write_trajectory_groups_parquet
 from mp_actors import close_proxy, move_to_child_process
 
 from .. import dev
@@ -125,6 +125,11 @@ async def register(
         with open(f"{output_dir}/model.json", "w") as f:
             json.dump(model.model_dump(), f)
 
+        # Auto-migrate any old JSONL trajectory files to Parquet
+        from art.utils.trajectory_migration import auto_migrate_on_register
+
+        auto_migrate_on_register(output_dir)
+
         # Initialize wandb and weave early if this is a trainable model
         if model.trainable and "WANDB_API_KEY" in os.environ:
             _ = self._get_wandb_run(model)
@@ -356,11 +361,10 @@ async def _log(
 
         # Get the file name for the current iteration, or default to 0 for non-trainable models
         iteration = self.__get_step(model)
-        file_name = f"{iteration:04d}.jsonl"
+        file_name = f"{iteration:04d}.parquet"
 
-        # Write the logs to the file
-        with open(f"{parent_dir}/{file_name}", "w") as f:
-            f.write(serialize_trajectory_groups(trajectory_groups))
+        # Write the logs to Parquet file (with ZSTD compression)
+        write_trajectory_groups_parquet(trajectory_groups, f"{parent_dir}/{file_name}")
 
         # Collect all metrics (including reward) across all trajectories
         all_metrics: dict[str, list[float]] = {"reward": [], "exception_rate": []}