Parallel per-type load writes + omnigraph optimize/cleanup CLI (#46)

* Parallel per-type load writes + omnigraph optimize/cleanup CLI

## MR-677.3 — parallel per-type load writes

The load path already groups records into one RecordBatch per type and
makes one Lance commit per table (loader::mod.rs:249-..), but those
commits ran sequentially. Wrap node and edge write loops in
`futures::stream::buffered(N)` against a new helper
`write_batches_concurrently`. Concurrency tunable via
`OMNIGRAPH_LOAD_CONCURRENCY` (default 8).

## MR-676 — `omnigraph optimize` and `omnigraph cleanup`

New CLI subcommands that walk every node + edge table in the repo:

- `omnigraph optimize <uri>` — runs Lance `compact_files` on each
  table to merge small fragments into fewer larger ones.
- `omnigraph cleanup <uri> --keep N | --older-than 7d --confirm` —
  runs Lance `cleanup_old_versions` to prune historical manifests +
  unique fragments. Requires `--confirm` because it's destructive.
  Supports both count-based and time-based retention (or both AND'd
  together). Time uses chrono `DateTime<Utc>` (added as a workspace
  dep, default-features off).

Both commands run their per-table loops in parallel (8-way bounded,
`OMNIGRAPH_MAINTENANCE_CONCURRENCY` env override). Smoke-tested
against the 114-table prod graph: optimize went 7m15s sequential
→ 1m28s parallel. cleanup --keep 1 removed 137 historical versions
across 114 tables in 1m57s without disrupting `/healthz` or query
responses.

Public API on `Omnigraph`:

  pub async fn optimize(&mut self) -> Result<Vec<TableOptimizeStats>>
  pub async fn cleanup(&mut self, opts: CleanupPolicyOptions)
      -> Result<Vec<TableCleanupStats>>

All 10 existing loader tests still pass.

Closes MR-676.
Partially addresses MR-677 (the .3 — parallel by type — piece;
MR-677.1 is for the `omnigraph embed` path, not load, since load
doesn't call Gemini directly. .2 was already in place).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* chore: regenerate openapi.json

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
This commit is contained in:
Andrew Altshuler 2026-04-25 14:22:14 +03:00 committed by GitHub
parent 628bc2e607
commit 74eb5a5380
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 498 additions and 31 deletions

9
Cargo.lock generated
View file

@ -4596,7 +4596,7 @@ dependencies = [
[[package]]
name = "omnigraph-cli"
version = "0.3.0"
version = "0.3.1"
dependencies = [
"assert_cmd",
"clap",
@ -4616,7 +4616,7 @@ dependencies = [
[[package]]
name = "omnigraph-compiler"
version = "0.3.0"
version = "0.3.1"
dependencies = [
"ahash",
"arrow-array",
@ -4637,7 +4637,7 @@ dependencies = [
[[package]]
name = "omnigraph-engine"
version = "0.3.0"
version = "0.3.1"
dependencies = [
"arrow-array",
"arrow-cast",
@ -4646,6 +4646,7 @@ dependencies = [
"arrow-select",
"async-trait",
"base64",
"chrono",
"fail",
"futures",
"lance",
@ -4674,7 +4675,7 @@ dependencies = [
[[package]]
name = "omnigraph-server"
version = "0.3.0"
version = "0.3.1"
dependencies = [
"async-trait",
"aws-config",