AODBcRewriter: re-sort tables stored sorted by a reordered reference

Second member of the O2-7098 family. Stage 1b reorders O2track_iu/
O2mfttrack/O2fwdtrack; several other tables (O2fwdtrkcl, O2ambiguoustrack,
O2trackqa_003, O2v0_002, O2cascade_001, O2decay3body) are stored SORTED BY a
reference into one of those (or into O2collision). Their index values were
remapped correctly, but their rows stayed put, so a sorted column came out
unsorted -- the same defect as the split "-1" group that once caused

  [FATAL] Table Tracks_IU index fIndexCollisions has a group with index -1
          that is split by 776

Confirmed on real files: O2fwdtrkcl and O2trackqa_003 came out unsorted in
every DF of two of three anchored-MC samples tested.

Fix: findGroupingColumn() + resortByGroupingColumn(), derived from the data
rather than another hardcoded table list -- "if T.B is non-decreasing on
input and B's referent moved, re-sort T by the remapped B". Leaves
O2mfttrackcov alone (not sorted on input). Iterated to a fixed point, since
O2cascade_001 is sorted by fIndexV0s and O2v0_002 may itself have just
moved.

AODBcRewriterCheckLinks() now also asserts sortedness is preserved. Fixture
extended with O2fwdtrkcl and a second cascade to exercise the
O2v0->O2cascade dependency; against the previous commit the suite fails 3
assertions and 6 link checks, and passes with this one.

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

Author: sawenzel

MC/utils/AODBcRewriter.C

@@ -1075,15 +1075,155 @@ static void stage1b_reorderTrackTables(
   }
 }
 
+// ----------------------------------------------------------------------------
+// Re-sorting tables that are STORED SORTED BY a reference
+//
+// Stage 1b exists because a track table grouped by fIndexCollisions stops being
+// sliceable once the collision table is reordered.  Exactly the same is true of
+// every other table stored sorted by a reference into a table this tool
+// reorders — O2v0_002 and O2cascade_001 by fIndexCollisions, O2fwdtrkcl by
+// fIndexFwdTracks, O2ambiguoustrack and O2trackqa_003 by fIndexTracks.
+// Remapping their values while leaving their rows in place turns a sorted
+// column into an unsorted one, which is the same defect that produced
+//   "Table ... index fIndexCollisions has a group with index -1 that is split".
+//
+// Rather than hardcode which tables are grouped by what — the enumeration habit
+// that caused O2-7098 — the decision is derived from the data: IF a column is
+// non-decreasing in the input, THEN it is an ordering the file carries and the
+// output must preserve it.  That is self-maintaining across schema changes, and
+// it correctly leaves O2mfttrackcov alone: its fIndexMFTTracks is not sorted in
+// the input, so there is no ordering to preserve.
+
+struct ResortCandidate {
+  size_t                   planIdx = 0;
+  std::string              tableName;
+  std::string              keyBranch;
+  std::vector<Long64_t>    keyValues;    // raw input values of keyBranch
+  std::vector<std::string> refPrefixes;  // referent of keyBranch, from kIndexRefs
+};
+
+// Read a scalar Int_t-like index column.  Returns false for VLA / fixed-array
+// columns, which are not row orderings.
+static bool readScalarIndexColumn(TTree *t, const char *branch,
+                                  std::vector<Long64_t> &out) {
+  TBranch *br = t->GetBranch(branch);
+  if (!br) return false;
+  TLeaf *leaf = static_cast<TLeaf *>(br->GetListOfLeaves()->At(0));
+  if (!leaf || leaf->GetLeafCount() || leaf->GetLen() != 1) return false;
+  ScalarTag tag = tagOf(leaf);
+  size_t sz = byteSize(tag);
+  if (sz == 0) return false;
+  std::vector<unsigned char> buf(sz, 0);
+  br->SetAddress(buf.data());
+  out.clear();
+  out.reserve(t->GetEntries());
+  for (Long64_t i = 0; i < t->GetEntries(); ++i) {
+    br->GetEntry(i);
+    out.push_back(readAsInt(buf.data(), tag));
+  }
+  br->ResetAddress();
+  return true;
+}
+
+// The convention this tool writes (and Stage 1b establishes): valid values
+// ascending, the -1 "ambiguous" group as one contiguous run at the end.
+static bool isOrderedWithNullsLast(const std::vector<Long64_t> &v) {
+  Long64_t prev = -1;
+  bool seenNull = false;
+  for (auto x : v) {
+    if (x < 0) { seenNull = true; continue; }
+    if (seenNull) return false;   // a valid value after a null: not this convention
+    if (x < prev) return false;
+    prev = x;
+  }
+  return true;
+}
+
+// Pick the column a table is stored sorted by, if any.  fIndexCollisions wins
+// when several qualify, since that is the grouping O2's slicing cache checks.
+static bool findGroupingColumn(TTree *src, std::string &keyBranch,
+                               std::vector<Long64_t> &keyValues,
+                               std::vector<std::string> &refPrefixes) {
+  if (!src || src->GetEntries() < 2) return false;
+  bool found = false;
+  for (auto &[branchName, prefixes] : kIndexRefs) {
+    std::vector<Long64_t> vals;
+    if (!readScalarIndexColumn(src, branchName.c_str(), vals)) continue;
+    if (!isOrderedWithNullsLast(vals)) continue;
+    bool preferred = (branchName == "fIndexCollisions");
+    if (found && !preferred) continue;
+    keyBranch   = branchName;
+    keyValues   = std::move(vals);
+    refPrefixes = prefixes;
+    found = true;
+    if (preferred) break;
+  }
+  return found;
+}
+
+// Re-sort each candidate by its remapped grouping column.  Iterated to a fixed
+// point because these tables reference each other: O2cascade_001 is sorted by
+// fIndexV0s, and O2v0_002 may itself have just been re-sorted.
+static void resortByGroupingColumn(
+    std::vector<TablePlan> &plans,
+    std::unordered_map<std::string, PermMap> &allPerms,
+    const std::vector<ResortCandidate> &candidates) {
+
+  const int kMaxPasses = 8;
+  int pass = 0;
+  for (; pass < kMaxPasses; ++pass) {
+    bool changed = false;
+    for (auto &cand : candidates) {
+      const PermMap *refPerm = findPermFor(allPerms, cand.refPrefixes);
+      if (!refPerm) continue;   // referent absent from this DF
+
+      Long64_t n = (Long64_t)cand.keyValues.size();
+      struct SortEntry { Long64_t key; Long64_t srcRow; };
+      std::vector<SortEntry> entries;
+      entries.reserve(n);
+      for (Long64_t i = 0; i < n; ++i) {
+        Long64_t old = cand.keyValues[i];
+        Long64_t nw  = (old >= 0 && old < (Long64_t)refPerm->size()) ? (*refPerm)[old] : -1;
+        entries.push_back({nw, i});
+      }
+      // Same ordering convention as Stage 1: -1 sinks to a contiguous tail.
+      std::stable_sort(entries.begin(), entries.end(),
+        [](const SortEntry &a, const SortEntry &b) {
+          if (a.key < 0 && b.key >= 0) return false;
+          if (a.key >= 0 && b.key < 0) return true;
+          return a.key < b.key;
+        });
+      std::vector<Long64_t> rowOrder;
+      rowOrder.reserve(n);
+      for (auto &e : entries) rowOrder.push_back(e.srcRow);
+
+      if (rowOrder == plans[cand.planIdx].rowOrder) continue;
+      plans[cand.planIdx].rowOrder = rowOrder;
+      allPerms[cand.tableName]     = permFromRowOrder(n, rowOrder);
+      changed = true;
+      std::cout << "  Re-sort: " << cand.tableName << " by remapped "
+                << cand.keyBranch << " (was sorted by it on input)\n";
+    }
+    if (!changed) break;
+  }
+  if (pass == kMaxPasses)
+    std::cerr << "  [warn] re-sort did not reach a fixed point after "
+              << kMaxPasses << " passes — check for a reference cycle\n";
+}
+
 // Plan every table not yet claimed by an earlier stage: paste-join children
-// follow their parent's row order, everything else keeps its own.  The index
-// columns they carry are NOT enumerated here any more — buildRemaps() derives
-// them from kIndexRefs when processDF writes.
+// follow their parent's row order, everything else keeps its own — unless it is
+// stored sorted by a reference, in which case resortByGroupingColumn() below
+// re-establishes that ordering.  The index columns these tables carry are NOT
+// enumerated here — buildRemaps() derives them from kIndexRefs when processDF
+// writes.
 static void planRemainingTables(
     TDirectory *dirIn, std::vector<TablePlan> &plans,
     std::unordered_map<std::string, PermMap> &allPerms,
     std::unordered_set<std::string> &planned) {
 
+  std::vector<ResortCandidate> resortCandidates;
+
   TIter it(dirIn->GetListOfKeys());
   while (TKey *key = static_cast<TKey *>(it())) {
     if (TString(key->GetClassName()) != "TTree") continue;
@@ -1139,10 +1279,26 @@ static void planRemainingTables(
       std::iota(rowOrder.begin(), rowOrder.end(), 0LL);
     }
 
+    // A table that arrives here with its own row order may still be STORED
+    // SORTED BY one of its index columns; if that column's referent gets
+    // reordered, the sortedness has to be re-established.  Record what is
+    // needed for that; the decision is made below, once every table in this
+    // stage has a permutation.
+    if (!parentPerm) {
+      ResortCandidate cand;
+      if (findGroupingColumn(src, cand.keyBranch, cand.keyValues, cand.refPrefixes)) {
+        cand.planIdx   = plans.size();
+        cand.tableName = tname;
+        resortCandidates.push_back(std::move(cand));
+      }
+    }
+
     allPerms[tname] = permFromRowOrder(nSrc, rowOrder);
     plans.push_back({tname, std::move(rowOrder)});
     planned.insert(tname);
   }
+
+  resortByGroupingColumn(plans, allPerms, resortCandidates);
 }
 
 // ============================================================================
@@ -1826,6 +1982,23 @@ static bool checkLinksDF(TDirectory *din, TDirectory *dout, const char *dfName)
       linksOut.push_back({branchName, &fpOut[refName], nullptr});
     }
 
+    // Ordering preservation: a column that is sorted on input describes a
+    // grouping the file carries, and O2's slicing cache relies on it.  Remapping
+    // the values while leaving the rows in place silently destroys it — the same
+    // defect as the split "-1" group, just in a different table.
+    for (auto &[branchName, prefixes] : kIndexRefs) {
+      std::vector<Long64_t> vIn, vOut;
+      if (!readScalarIndexColumn(tIn, branchName.c_str(), vIn)) continue;
+      if (!isOrderedWithNullsLast(vIn)) continue;   // no ordering to preserve
+      if (!readScalarIndexColumn(tOut, branchName.c_str(), vOut)) continue;
+      if (!isOrderedWithNullsLast(vOut)) {
+        std::cerr << "  [FAIL] " << dfName << ": " << tn << "." << branchName
+                  << " was sorted on input but is not on output"
+                     " (grouping destroyed — slicing will misbehave)\n";
+        ok = false;
+      }
+    }
+
     auto cIn  = linkTupleCounts(tIn,  fpIn[tn],  linksIn);
     auto cOut = linkTupleCounts(tOut, fpOut[tn], linksOut);
 

MC/utils/CLAUDE.md

@@ -49,6 +49,11 @@ practice **every merged MC AO2D since 4 Jun 2026 is affected**.
    test instead of silently mis-linking data.
 5. **`AODBcRewriterCheckLinks()` (Section 11b)** — the check that can actually
    see this bug class; see "Testing" below.
+6. **Re-sort tables stored sorted by a reference** (`findGroupingColumn` /
+   `resortByGroupingColumn`, Section 8). The same family: `O2fwdtrkcl` and
+   `O2trackqa_003` were coming out unsorted in every DF of two of the three
+   sample files, which breaks O2's slicing the same way the split `-1` group
+   did. Derived from the data, not from a list — see resolved gap 1.
 
 Maurice's #2418 has the same shape (defer `O2fwdtrack`, remap via the MFT/Fwd
 perms) and is correct; this generalises it from one table to all of them.
@@ -331,13 +336,12 @@ row whose `fIndexMcCollisions` pointed to a dropped row is also dropped.
 
 These were identified during the refactor but not yet implemented:
 
-### 1. Tables SORTED BY a reference into a reordered table are not re-sorted
+### ~~1. Tables SORTED BY a reference into a reordered table are not re-sorted~~ (RESOLVED)
 
-**This is the remaining member of the O2-7098 family — open, and the natural
-follow-up PR.** Stage 1b reorders `O2track_iu` / `O2mfttrack` / `O2fwdtrack`.
-Several other tables are *stored sorted by* a reference into one of those (or
-into `O2collision`). Their values are now correctly remapped, but their rows are
-left where they were, so the ordering is destroyed:
+Stage 1b reorders `O2track_iu` / `O2mfttrack` / `O2fwdtrack`. Several other
+tables are *stored sorted by* a reference into one of those (or into
+`O2collision`). Their values were remapped but their rows left where they were,
+so the ordering was destroyed:
 
 | table | key | sorted in input? |
 |---|---|---|
@@ -347,16 +351,22 @@ left where they were, so the ordering is destroyed:
 | `O2v0_002`, `O2cascade_001`, `O2decay3body` | `fIndexCollisions` | yes |
 | `O2mfttrackcov` | `fIndexMFTTracks` | **no** — so not an invariant there |
 
-(measured on `example_AOD/AO2D_pre.root`.) This is the same failure mode as the
-split `-1` group that produced the `ArrowTableSlicingCache::validateOrder`
-FATAL — not yet observed in the wild, but structurally identical.
-
-**Suggested fix** — data-derived rather than another hardcoded list: *if `T.B`
-was sorted in the input and `B`'s referent was reordered, re-sort `T` by the
-remapped `B`.* That rule reproduces Stage 1b's behaviour and correctly leaves
-`O2mfttrackcov` alone. It is deliberately **not** in the O2-7098 fix: it changes
-the row order of tables that are currently untouched, so it deserves its own
-review. `AODBcRewriterCheckLinks` is already in place to police it.
+(measured on `example_AOD/AO2D_pre.root`.) Same failure mode as the split `-1`
+group that produced the `ArrowTableSlicingCache::validateOrder` FATAL. **Not
+theoretical**: on `example_AOD/AO2D_pre.root` and `bigger2/`, `O2fwdtrkcl` and
+`O2trackqa_003` came out unsorted in *every* DF.
+
+**Fix**: `findGroupingColumn` + `resortByGroupingColumn` (Section 8), driven by
+the data rather than by another hardcoded list — *if `T.B` is non-decreasing in
+the input and `B`'s referent gets reordered, re-sort `T` by the remapped `B`*.
+Self-maintaining across schema changes, and it correctly leaves `O2mfttrackcov`
+alone (its `fIndexMFTTracks` is not sorted on input, so there is no ordering to
+preserve). Iterated to a fixed point, because these tables reference each other:
+`O2cascade_001` is sorted by `fIndexV0s` and `O2v0_002` may itself have just
+been re-sorted.
+
+Policed by a new check in `AODBcRewriterCheckLinks`: a column sorted on input
+must be sorted on output.
 
 ### ~~2. `fIndexCollisions` inside `O2mccollision` is not remapped~~ (MOOT)
 

MC/utils/tests/makeTestAOD.C

@@ -110,21 +110,31 @@ const int kFwdMatch[]  = {-1, 0, -1, -1, 3, -1};// fwd -> fwd self-reference
 // exactly the metric O2-7098 was reported with.
 const int kFwdMcPart[] = {0, 3, -1, 6, 9, -1};
 
+// V0s and cascades are stored SORTED BY fIndexCollisions, as reconstruction
+// writes them.  The collision table gets reordered, so these have to be
+// re-sorted or the grouping is destroyed.  The cascades are additionally sorted
+// by fIndexV0s, whose referent is itself one of these re-sorted tables — a
+// two-level dependency that a single pass would not resolve.
 const int kNV0        = 2;
 const int kV0Coll[]   = {0, 2};
 const int kV0Pos[]    = {0, 4};
 const int kV0Neg[]    = {1, 5};
 
-const int kNCasc      = 1;
-const int kCascColl[] = {0};
-const int kCascV0[]   = {0};
-const int kCascBach[] = {2};
+const int kNCasc      = 2;
+const int kCascColl[] = {0, 2};
+const int kCascV0[]   = {0, 1};
+const int kCascBach[] = {2, 4};
 
+// Sorted by fIndexTracks / fIndexFwdTracks respectively — both referents are
+// reordered by Stage 1b.
 const int kNAmb          = 2;
 const int kAmbTrack[]    = {3, 6};
 const int kAmbBCFirst[]  = {0, 3};
 const int kAmbBCLast[]   = {2, 4};
 
+const int kNFwdCl        = 4;
+const int kFwdClTrack[]  = {0, 1, 3, 4};
+
 } // namespace
 
 void makeTestAOD(const char *outFileName = "AO2D_test.root")
@@ -308,6 +318,16 @@ void makeTestAOD(const char *outFileName = "AO2D_test.root")
     t.Write();
   }
 
+  // ---- O2fwdtrkcl (grouped by fIndexFwdTracks, a Stage-1b table) ---------
+  {
+    TTree t("O2fwdtrkcl", "fwd track clusters");
+    int fwd; float x;
+    t.Branch("fIndexFwdTracks", &fwd, "fIndexFwdTracks/I");
+    t.Branch("fX",              &x,   "fX/F");
+    for (int i = 0; i < kNFwdCl; ++i) { fwd = kFwdClTrack[i]; x = 3000.f + i; t.Fill(); }
+    t.Write();
+  }
+
   // ---- O2ambiguoustrack (slice into the deduplicated BC table) ------------
   {
     TTree t("O2ambiguoustrack", "ambiguous tracks");

MC/utils/tests/testAODBcRewriter.C

@@ -263,6 +263,40 @@ int testAODBcRewriter(const char *inFileName  = "AO2D_test.root",
                            " daughter slices still point at the right particles");
   }
 
+  // ---- Grouping preserved for tables stored sorted by a reference ---------
+  // These are written sorted by an index into a table the rewriter reorders.
+  // Remapping the values without re-sorting the rows leaves the column
+  // unsorted, which breaks O2's slicing cache the same way a split "-1" group
+  // does.  The referenced tables here are all genuinely reordered, so a
+  // no-op implementation cannot pass this by accident.
+  {
+    const char *cases[][2] = {
+      {"O2v0_002",         "fIndexCollisions"},
+      {"O2cascade_001",    "fIndexCollisions"},
+      {"O2cascade_001",    "fIndexV0s"},        // referent is itself re-sorted
+      {"O2fwdtrkcl",       "fIndexFwdTracks"},
+      {"O2ambiguoustrack", "fIndexTracks"},
+    };
+    for (auto &c : cases) {
+      auto vIn  = readInt(get(din,  c[0]), c[1]);
+      auto vOut = readInt(get(dout, c[0]), c[1]);
+      auto ordered = [](const std::vector<Int_t> &v) {
+        Int_t prev = -1; bool seenNull = false;
+        for (auto x : v) {
+          if (x < 0) { seenNull = true; continue; }
+          if (seenNull || x < prev) return false;
+          prev = x;
+        }
+        return true;
+      };
+      std::string what = std::string(c[0]) + "." + c[1];
+      if (vIn.empty())        { fail(what + ": column missing from input"); continue; }
+      if (!ordered(vIn))      { fail(what + ": fixture is not sorted on input — fix makeTestAOD.C"); continue; }
+      if (!ordered(vOut))     fail(what + ": sorted on input, unsorted on output (grouping destroyed)");
+      else                    pass(what + ": still sorted after the rewrite");
+    }
+  }
+
   // ---- Slice into the deduplicated BC table -------------------------------
   {
     auto bcIn  = readULong(get(din,  "O2bc_001"), "fGlobalBC");