sparkforensics-cli 0.1.0

package/vendor-core/parser-worker.js

@@ -0,0 +1,249 @@
+import { Gunzip } from './vendor/fflate.js';
+import { createLz4BlockDecoder } from './lz4-block.js';
+import { Decompress as ZstdDecompress } from './vendor/fzstd.js';
+import { createSnappyBlockDecoder } from './snappy-block.js';
+import { createState, dispatchLine, buildChunkDecoder, emitParseCompletion,                  } from './event-handlers.js';
+import { TASK_FIELD_NAMES } from './stage-quantiles.js';
+import { runParseFromUrl, sniffCodec } from './shs-fetch.js';
+
+export {
+  buildChunkDecoder, createState, normalizeSparkProperties, parseSparkMemoryMB, extractResources,
+  accumulateTask, resolvePlanTree, startApplication, updateEnvironment, startJob, endJob, submitStage,
+  mergeStageRddInfo, recordStageExecutorMetrics, startSqlExecution, endSqlExecution,
+  applyDriverAccumUpdates, addExecutor, removeExecutor, processEvent, dispatchLine,
+  collectStageExecutorMetrics,
+} from './event-handlers.js';
+
+export {
+  FIELDS, TASK_FIELD_NAMES, finalizeStage, computeFieldQuantiles, computeDurationQuantiles, classifySpill,
+} from './stage-quantiles.js';
+
+export {
+  runParseFromUrl, sniffCodec, decodeShsArchive,
+} from './shs-fetch.js';
+
+export { naturalCompare, reassembleRollingEntries } from './rolling-log-reassembly.js';
+
+// 512 KB rather than a larger round number: Spark event logs commonly
+// compress ~15-20x, so a coarser chunk decompresses into a burst of tens of
+// thousands of lines reported at one unmoving pct before jumping: smaller
+// reads give the progress bar more checkpoints to advance through instead.
+const CHUNK_SIZE = 512 * 1024;
+const MIN_PROGRESS_STEPS = 100;
+// Emitting every 2000 lines was the other half of the stair-step: a highly
+// compressed chunk can decode thousands of lines in one go, so 2000 was
+// rarely a mid-chunk boundary: it just meant a handful of progress
+// messages for the whole file, each landing on a different currentPct.
+const PROGRESS_EMIT_LINES = 300;
+
+                                     
+                                                     
+
+// Minimal shape streamFile/runParse/runParseFiles actually read off `file`
+// (name, size, slice(start,end).arrayBuffer()): narrower than the full DOM
+// `File` interface. A real `File` (the browser Worker path, via
+// WorkerIncomingMessage below) satisfies this structurally, but so does the
+// plain object src/cli/collect-run.ts's nodeFileFromPath builds for Node
+// (which has no DOM `File` constructor to build a real one from). Widened
+// from the plan's literal `File` to match that real second caller rather
+// than forcing a cast at its call site.
+                   
+               
+               
+                                                                             
+  
+
+// Incoming message shapes accepted by the worker's `self.onmessage` handler.
+                            
+                                 
+                                              
+                                         
+                                                                     
+
+// fflate's Gunzip and fzstd's Decompress are untyped vendor JS (plain
+// prototype classes, not `class` declarations), so TS can't infer a
+// construct signature for them; this local shape is just enough to type
+// the two call sites below without touching the vendored files (mirrors
+// shs-fetch.ts's identical StreamingDecoder shim).
+                                                                           
+                                                                                           
+
+// Stream one File's (possibly compressed) bytes through the codec dispatch,
+// in `chunkSize` slices, invoking `onChunk` with each decompressed buffer as
+// it becomes available. Shared by runParse (single file) and runParseFiles
+// (rolling multi-file directories): the codec is sniffed per-file since
+// Spark's file-rolling never spans a compressor's own framing.
+async function streamFile(
+  file            ,
+  onChunk                                           ,
+  chunkSize        ,
+)                {
+  const header = new Uint8Array(await file.slice(0, Math.min(8, file.size)).arrayBuffer());
+  const codec = sniffCodec(header);
+
+  // Codec decoders' push() can synchronously invoke their onChunk callback,
+  // which was bound once, before the loop below has a `start` to report:
+  // route pct through this shared mutable slot instead, updated ahead of
+  // each push so compressed codecs report bytes-read progress too, not just
+  // the uncompressed branch.
+  let currentPct = 0;
+  const gunzip = codec === 'gz' ? new (Gunzip                                   )((inflated) => onChunk(inflated, currentPct)) : null;
+  const lz4 = codec === 'lz4' ? createLz4BlockDecoder((inflated) => onChunk(inflated, currentPct)) : null;
+  const zstd = codec === 'zstd' ? new (ZstdDecompress                                   )((inflated) => onChunk(inflated, currentPct)) : null;
+  const snappy = codec === 'snappy' ? createSnappyBlockDecoder((inflated) => onChunk(inflated, currentPct)) : null;
+
+  // A fixed read size gives too few progress checkpoints on smaller files
+  // (e.g. a 3 MB log at 512 KB reads is only ~6 steps, ~15% jumps): cap the
+  // read size so every file gets at least MIN_PROGRESS_STEPS reads, however
+  // small, while chunkSize still bounds it above for large-file I/O.
+  const stepSize = Math.max(1, Math.min(chunkSize, Math.ceil(file.size / MIN_PROGRESS_STEPS)));
+
+  let offset = 0;
+  while (offset < file.size) {
+    const start = offset;
+    const slice = new Uint8Array(await file.slice(start, start + stepSize).arrayBuffer());
+    offset += stepSize;
+    const final = offset >= file.size;
+    currentPct = start / file.size;
+    if (gunzip) gunzip.push(slice, final);
+    else if (lz4) lz4.push(slice);
+    else if (zstd) zstd.push(slice, final);
+    else if (snappy) snappy.push(slice);
+    else onChunk(slice, currentPct);
+  }
+  if (lz4) lz4.end();
+  if (snappy) snappy.end();
+}
+
+// Parse a dropped `File`, always streaming to honor the core invariant: the
+// decompressed task-event stream must never be buffered whole. The file is read
+// in `chunkSize` slices; a compressed log (gzip, Zstandard, Spark LZ4Block, or
+// Spark's Snappy framing) is inflated incrementally so only one decompressed
+// chunk is live at a time, exactly like the uncompressed path.
+// `chunkSize` is injectable purely so tests can force multi-chunk streaming.
+export async function runParse(
+  file            ,
+  state             ,
+  { emit = (msg         ) => self.postMessage(msg), chunkSize = CHUNK_SIZE }          = {},
+)                {
+  if (file.size === 0) {
+    emit({ type: 'error', message: 'File is empty.' });
+    return;
+  }
+
+  const decoder = buildChunkDecoder();
+  let linesProcessed = 0;
+  const feed = (bytes            , pct         ) => {
+    for (const line of decoder.decode(bytes)) {
+      dispatchLine(line, state, emit);
+      linesProcessed++;
+      if (linesProcessed % PROGRESS_EMIT_LINES === 0) {
+        emit({ type: 'progress', pct: pct ?? null, linesProcessed });
+      }
+    }
+  };
+
+  try {
+    await streamFile(file, feed, chunkSize);
+  } catch (e) {
+    const message = e instanceof Error ? e.message : String(e);
+    emit({ type: 'error', message: `Could not decompress "${file.name}": ${message}` });
+    return;
+  }
+
+  for (const line of decoder.flush()) {
+    dispatchLine(line, state, emit);
+  }
+
+  if (!state.app) {
+    emit({ type: 'error', message: 'Not a Spark event log: SparkListenerApplicationStart not found.' });
+    return;
+  }
+
+  emitParseCompletion(state, emit, linesProcessed);
+}
+
+// Parse a rolling `eventlog_v2_*` directory: an ordered array of File objects
+// (already deduped/sorted by reassembleRollingEntries) representing
+// `events_<index>_...` files. Reuses the same chunked streaming codec
+// dispatch as runParse, once per file in sequence, but keeps ONE NDJSON
+// line-decoder alive across all files: nothing in the rolling-log format
+// guarantees a file-roll boundary lands on a line boundary.
+export async function runParseFiles(
+  files              ,
+  state             ,
+  { emit = (msg         ) => self.postMessage(msg), chunkSize = CHUNK_SIZE }          = {},
+)                {
+  if (files.length === 0) {
+    emit({ type: 'error', message: 'Rolling event-log directory contained no event files.' });
+    return;
+  }
+
+  const decoder = buildChunkDecoder();
+  let linesProcessed = 0;
+  const totalSize = files.reduce((sum, f) => sum + f.size, 0);
+  let bytesBeforeCurrentFile = 0;
+  let currentFileSize = 0;
+  const feed = (bytes            , pct         ) => {
+    for (const line of decoder.decode(bytes)) {
+      dispatchLine(line, state, emit);
+      linesProcessed++;
+      if (linesProcessed % PROGRESS_EMIT_LINES === 0) {
+        const overallPct = totalSize > 0 ? (bytesBeforeCurrentFile + (pct ?? 0) * currentFileSize) / totalSize : null;
+        emit({ type: 'progress', pct: overallPct, linesProcessed });
+      }
+    }
+  };
+
+  for (const file of files) {
+    currentFileSize = file.size;
+    try {
+      await streamFile(file, feed, chunkSize);
+    } catch (e) {
+      const message = e instanceof Error ? e.message : String(e);
+      emit({ type: 'error', message: `Could not decompress "${file.name}": ${message}` });
+      return;
+    }
+    bytesBeforeCurrentFile += file.size;
+  }
+
+  for (const line of decoder.flush()) {
+    dispatchLine(line, state, emit);
+  }
+
+  if (!state.app) {
+    emit({ type: 'error', message: 'Not a Spark event log: SparkListenerApplicationStart not found.' });
+    return;
+  }
+
+  emitParseCompletion(state, emit, linesProcessed);
+}
+
+// ─── Worker message bus (only active when running as a Web Worker) ──────────────
+
+const isWorker = typeof WorkerGlobalScope !== 'undefined' && self instanceof WorkerGlobalScope;
+
+if (isWorker) {
+  let workerState                     = null;
+
+  self.onmessage = async ({ data }                                     ) => {
+    if (data.type === 'parse') {
+      workerState = createState();
+      await runParse(data.file, workerState);
+    } else if (data.type === 'parseFromUrl') {
+      workerState = createState();
+      await runParseFromUrl(data.request, workerState);
+    } else if (data.type === 'parseFiles') {
+      workerState = createState();
+      await runParseFiles(data.files, workerState);
+    } else if (data.type === 'getTaskData') {
+      const { stageId, reqId } = data;
+      const stored = workerState?.taskStore.get(stageId) ?? new Float64Array(0);
+      self.postMessage({
+        type: 'taskData', stageId, reqId,
+        metrics: stored.slice(),
+        fieldNames: TASK_FIELD_NAMES,
+      });
+    }
+  };
+}

package/vendor-core/plan-dot.js

@@ -0,0 +1,25 @@
+// Graphviz DOT export of a resolved planTree (sparkDoctor SqlPlanDotWriter).
+// Plain string building, no dependency, no metric annotation yet (deferred,
+// matching sparkDoctor's own unimplemented roadmap item).
+                                           
+import { walkPlanTree } from './plan-tree-walk.js';
+
+function esc(s        )         { return String(s).replace(/\\/g, '\\\\').replace(/"/g, '\\"'); }
+
+export function planTreeToDot(planTree                 , { title = 'plan' }                     = {})         {
+  if (!planTree) return '';
+  const lines = [`digraph "${esc(title)}" {`, '  rankdir=BT;', '  node [shape=box];'];
+  const idOf = new WeakMap                  ();
+  const edges           = [];
+  let counter = 0;
+  walkPlanTree(planTree, (node, parent) => {
+    const id = `n${counter++}`;
+    idOf.set(node, id);
+    const label = node.detail && node.detail !== node.name ? `${node.name}\\n${node.detail}` : node.name;
+    lines.push(`  ${id} [label="${esc(label)}"];`);
+    // Pre-order guarantee: parent's id is already assigned by the time a
+    // child is visited, so a single walk covers both node labels and edges.
+    if (parent) edges.push(`  ${idOf.get(parent)} -> ${id};`);
+  }, { dedupe: true });
+  return [...lines, ...edges, '}'].join('\n');
+}

package/vendor-core/plan-duration-attribution.js

@@ -0,0 +1,185 @@
+// Approximate stage wall-time back to plan operators, splitting the tree at
+// Exchange boundaries into segments and zipping segments (deepest-first) to
+// submission-ordered stage IDs. No exact ground truth exists in Spark's event
+// model: this is inference (dataflint's GraphDurationAttribution approach).
+import { walkPlanTree } from './plan-tree-walk.js';
+                                           
+
+                                                                                                          
+
+function isExchange(node          )          { return /exchange/i.test(node.name); }
+
+function timingMs(node          )                {
+  for (const m of (node.metrics ?? [])) {
+    if (m.metricType === 'timing') return m.value;
+    if (m.metricType === 'nsTiming') return m.value / 1e6;
+  }
+  for (const m of (node.metrics ?? [])) {
+    if (/time/i.test(m.name) && typeof m.value === 'number') return m.value;
+  }
+  return null;
+}
+
+// Group nodes into connected components after cutting each Exchange -> child
+// boundary. Component ids are stable pre-order identities only; depth,
+// traversal order, and parent relationships live in segmentTopology so no
+// consumer has to infer plan topology from an id's numeric value. A fresh id
+// is required for every cut: sibling producer branches at the same depth are
+// separate components, not one shared "depth" segment.
+export function computeSegments(planTree                 )   
+                         
+                               
+                                   
+  {
+  const segments               = [];
+  const segOf = new Map                  ();
+  const segmentTopology                  = new Map();
+  if (!planTree) return { segments, segOf, segmentTopology };
+  let nextSegment = 1;
+  let nextTraversalOrder = 0;
+  walkPlanTree(planTree, (node, parent) => {
+    let seg        ;
+    if (!parent) {
+      seg = 0;
+      segmentTopology.set(seg, { parentIndex: null, depth: 0, traversalOrder: nextTraversalOrder++ });
+    } else if (isExchange(parent)) {
+      const parentIndex = segOf.get(parent) ;
+      seg = nextSegment++;
+      segmentTopology.set(seg, {
+        parentIndex,
+        depth: segmentTopology.get(parentIndex) .depth + 1,
+        traversalOrder: nextTraversalOrder++,
+      });
+    } else {
+      seg = segOf.get(parent) ;
+    }
+    segOf.set(node, seg);
+    (segments[seg] ??= []).push(node);
+  }, { dedupe: true });
+  return { segments, segOf, segmentTopology };
+}
+
+// Deepest component executes first; equal-depth components retain plan
+// traversal order. Earliest-submitted stage is first. Zips them positionally,
+// truncating to whichever list is shorter: a component or stage past the
+// shorter length's cutoff has no pair and is simply absent from `pairs`.
+export function zipSegmentsToStages(
+  segments              ,
+  stagesById                                                             ,
+  stageIds          ,
+  segmentTopology                 ,
+)   
+                             
+                           
+                                                                                             
+  {
+  const orderedSegments = segments.filter(Boolean).map((nodes, i) => ({
+    i,
+    nodes,
+    ...segmentTopology.get(i) ,
+  })).sort((a, b) => b.depth - a.depth || a.traversalOrder - b.traversalOrder);
+  const orderedStages = [...stageIds]
+    .map(id => ({ id, s: stagesById.get(id) }))
+    .filter((x)                                                                         => Boolean(x.s))
+    .sort((a, b) => a.s.submittedAt  - b.s.submittedAt );
+
+  const n = Math.min(orderedSegments.length, orderedStages.length);
+  const pairs                                                                                      = [];
+  for (let k = 0; k < n; k++) {
+    pairs.push({
+      segmentIndex: orderedSegments[k].i,
+      stageId: orderedStages[k].id,
+      nodes: orderedSegments[k].nodes,
+      stage: orderedStages[k].s,
+    });
+  }
+  return { orderedSegments, orderedStages, pairs };
+}
+
+function componentDistance(fromIndex        , toIndex        , segmentTopology                 )         {
+  const fromAncestors = new Map                ();
+  let current                = fromIndex;
+  let distance = 0;
+  while (current != null) {
+    fromAncestors.set(current, distance++);
+    current = segmentTopology.get(current)?.parentIndex ?? null;
+  }
+
+  current = toIndex;
+  distance = 0;
+  while (current != null) {
+    const fromDistance = fromAncestors.get(current);
+    if (fromDistance != null) return fromDistance + distance;
+    current = segmentTopology.get(current)?.parentIndex ?? null;
+    distance++;
+  }
+  return Number.POSITIVE_INFINITY;
+}
+
+// For *display* purposes only (the "Stage N" label on a segment's group box
+// in the plan graph): a segment past zipSegmentsToStages' Math.min(...)
+// cutoff has no timing data of its own to pair on, but the physical plan
+// still places it structurally near a component that DID win a stage slot
+// (most often a small broadcast-builder component immediately upstream of the
+// join component that consumes it). Inherit the structurally nearest strictly
+// paired component's stage id rather than leaving its box an unexplained
+// "Stage —". Distance is the number of parent/child edges in the component
+// tree, never numeric id distance. Neighbor search is always against the
+// ORIGINAL strict pairs, never against other already-filled components, so
+// labels don't drift through a chain of inherited-from-inherited guesses.
+// zipSegmentsToStages/attributeStageDurationToPlan are untouched by this and
+// keep their own strict pairing: duration attribution must never double-
+// count one stage's wall time across two segments.
+export function mapSegmentsToStagesForDisplay(
+  segments              ,
+  stagesById                      ,
+  stageIds          ,
+  segmentTopology                      ,
+)                      {
+  const typedStagesById = stagesById                                                               ;
+  const typedTopology = segmentTopology                   ;
+  const { orderedSegments, pairs } = zipSegmentsToStages(segments, typedStagesById, stageIds, typedTopology);
+  const map = new Map                (pairs.map((p) => [p.segmentIndex, p.stageId]));
+  const pairedIndices = pairs.map((p) => p.segmentIndex);
+  if (pairedIndices.length === 0) return map;
+
+  for (const { i } of orderedSegments                   ) {
+    if (map.has(i)) continue;
+    let nearest = pairedIndices[0];
+    let nearestDist = componentDistance(i, nearest, typedTopology);
+    for (const pi of pairedIndices) {
+      const dist = componentDistance(i, pi, typedTopology);
+      if (dist < nearestDist) { nearestDist = dist; nearest = pi; }
+    }
+    map.set(i, map.get(nearest) );
+  }
+  return map;
+}
+
+export function attributeStageDurationToPlan(
+  planTree                 ,
+  stagesById                      ,
+  sqlExec                         ,
+)                        {
+  const out = new Map                  ();
+  const stageIds = sqlExec?.stageIds ?? [];
+  if (!planTree || stageIds.length === 0) return out;
+
+  const { segments, segmentTopology } = computeSegments(planTree);
+  const typedStagesById = stagesById                                                               ;
+  const { pairs } = zipSegmentsToStages(segments, typedStagesById, stageIds, segmentTopology);
+
+  for (const { nodes, stage } of pairs) {
+    const { submittedAt, completedAt } = stage                                                ;
+    const wall = Math.max(0, completedAt - submittedAt);
+    const weights = nodes.map(timingMs);
+    const totalTiming = weights.reduce        ((a, w) => a + (w ?? 0), 0);
+    nodes.forEach((node, idx) => {
+      const share = totalTiming > 0
+        ? wall * ((weights[idx] ?? 0) / totalTiming)
+        : wall / nodes.length;
+      out.set(node, share);
+    });
+  }
+  return out;
+}

package/vendor-core/plan-graph-model.js

@@ -0,0 +1,171 @@
+// Flattens a resolved planTree into a { nodes, edges } graph shape for the
+// React Flow + Dagre plan graph view ("Plan graph view" in
+// docs-site/contributor-guide/architecture/drill-down.md). Every Exchange
+// node is split into paired write/read halves: the
+// write half stays with its children's producer segment and the read half
+// stays with its parent's consumer segment. A duration share stays on the
+// half in the original source node's attributed segment. Both halves share a
+// `sourceNodeId` back-reference to the pre-split node's id, so any aggregate
+// logic over nodes never double-counts one Exchange.
+import { walkPlanTree } from './plan-tree-walk.js';
+import {
+  classifyNode,
+  buildDurationMap,
+  parseOperatorDetail,
+  getPrimaryMetric,
+} from './plan-node-detail.js';
+import { computeSegments, mapSegmentsToStagesForDisplay } from './plan-duration-attribution.js';
+import { stageIdsForSqlExec } from './detectors.js';
+                                                                                                       
+
+function isExchangeNode(node          )          {
+  return classifyNode(node.name) === 'exchange';
+}
+
+function makeGraphNode({
+  id,
+  sourceNodeId,
+  planNode,
+  category,
+  segmentIndex,
+  splitRole,
+  durationShare,
+}   
+             
+                       
+                     
+                   
+                       
+                                     
+                                
+ )                    {
+  return {
+    id,
+    sourceNodeId,
+    label: planNode.name,
+    category,
+    operatorDetail: parseOperatorDetail(planNode.name, planNode.detail ?? ''),
+    primaryMetric: getPrimaryMetric(planNode.metrics),
+    segmentIndex,
+    splitRole,
+    durationShare: durationShare ?? null,
+  };
+}
+
+export function buildPlanGraphModel(
+  planTree                 ,
+  opts                                                                    ,
+)                 {
+  const { scope, stageId, appModel } = opts;
+  if (!planTree) return { nodes: [], edges: [], segmentIndex: null, segmentCount: 0, scope: 'full', segmentStageIds: new Map                () };
+
+  const stage = appModel.stages.get(stageId);
+  // Coalesce undefined -> null so `sqlExecutionId` has a stable `number | null`
+  // type: every branch below narrows off the same variable instead of
+  // re-deriving it from `stage`, which TS can't re-narrow after the `stage?.`
+  // optional-chain expression above collapses to a plain boolean.
+  const sqlExecutionId = stage?.sqlExecutionId ?? null;
+  const sqlExec = sqlExecutionId != null ? appModel.sql.get(sqlExecutionId) : null;
+
+  const { segments, segOf, segmentTopology } = computeSegments(planTree);
+  const segmentCount = segments.filter(Boolean).length;
+
+  const durationMap =
+    sqlExec && sqlExecutionId != null ? buildDurationMap(planTree, appModel, sqlExec, sqlExecutionId) : null;
+
+  // Segment -> stage-id association, computed once up front regardless of
+  // `scope`: 'segment' needs it below to resolve which segment `stageId`
+  // belongs to, and 'full' needs the same map to label each segment's
+  // compound-group header (Task 11) with its stage id; it must not be
+  // computed only inside the 'segment' branch, or the full-plan view would
+  // have no stage-id data for its group headers. mapSegmentsToStagesForDisplay
+  // fills in segments that lost the strict zip's Math.min(segments, stages)
+  // pairing slot from their nearest paired neighbor, so a group box only
+  // falls back to an unexplained "Stage —" when there's truly no stage to
+  // borrow from (see its own doc comment); the earlier `targetSegmentIndex`
+  // lookup below stays correct either way, since a strictly-paired segment's
+  // canonical entry is always inserted before any inherited duplicate.
+  const segmentStageIds = new Map                ();
+  const linkedStageIds           =
+    sqlExec && sqlExecutionId != null ? stageIdsForSqlExec(sqlExecutionId, appModel.stages) : [];
+  if (linkedStageIds.length) {
+    const stagesById = new Map                                                        ();
+    for (const id of linkedStageIds) {
+      const s = appModel.stages.get(id);
+      if (s) stagesById.set(id, { submittedAt: s.submittedAt, completedAt: s.completedAt });
+    }
+    for (const [segIdx, sId] of mapSegmentsToStagesForDisplay(segments, stagesById, linkedStageIds, segmentTopology)) {
+      segmentStageIds.set(segIdx, sId);
+    }
+  }
+
+  // A full walk always happens first (cheap: one tree pass), splitting
+  // every Exchange and tagging each half's own segmentIndex correctly.
+  // Scope filtering is a post-pass over this output, since a split node's
+  // two halves can land in two different segments (see module doc comment)
+  // and a pre-split per-original-node filter can't express that.
+  const idOf = new WeakMap                  ();
+  let counter = 0;
+  const allNodes                      = [];
+  const allEdges                  = [];
+
+  walkPlanTree(planTree, (node, parent) => {
+    const id = `n${counter++}`;
+    const category = classifyNode(node.name);
+    // computeSegments (above) already walked this same planTree with the
+    // same `{ dedupe: true }` traversal, so every node this walk visits is
+    // guaranteed to already have a segOf entry; the `!` reflects that
+    // invariant, not an actual nullable field.
+    const seg = segOf.get(node) ;
+    const parentId = parent ? idOf.get(parent) : null;
+
+    if (isExchangeNode(node)) {
+      const readId = `${id}-read`;
+      const writeId = `${id}-write`;
+      const writeSegment = node.children?.length ? segOf.get(node.children[0])  : seg;
+      allNodes.push(makeGraphNode({
+        id: readId, sourceNodeId: id, planNode: node, category, segmentIndex: seg, splitRole: 'read',
+        durationShare: durationMap?.get(node),
+      }));
+      allNodes.push(makeGraphNode({
+        id: writeId, sourceNodeId: id, planNode: node, category, segmentIndex: writeSegment, splitRole: 'write',
+      }));
+      allEdges.push({ id: `${readId}=>${writeId}`, source: readId, target: writeId });
+      if (parentId) allEdges.push({ id: `${parentId}->${readId}`, source: parentId, target: readId });
+      idOf.set(node, writeId);
+    } else {
+      allNodes.push(makeGraphNode({
+        id, sourceNodeId: id, planNode: node, category, segmentIndex: seg, splitRole: null,
+        durationShare: durationMap?.get(node),
+      }));
+      if (parentId) allEdges.push({ id: `${parentId}->${id}`, source: parentId, target: id });
+      idOf.set(node, id);
+    }
+  }, { dedupe: true });
+
+  if (scope === 'full') {
+    return { nodes: allNodes, edges: allEdges, segmentIndex: null, segmentCount, scope: 'full', segmentStageIds };
+  }
+
+  let targetSegmentIndex                = null;
+  for (const [segIdx, sId] of segmentStageIds) {
+    if (sId === stageId) { targetSegmentIndex = segIdx; break; }
+  }
+
+  if (targetSegmentIndex === null) {
+    // Segment lookup failed (no sql/stage linkage, or this stage's segment
+    // fell outside the Math.min(segments, stages) truncation), fall back
+    // to full-plan scope rather than showing nothing.
+    return { nodes: allNodes, edges: allEdges, segmentIndex: null, segmentCount, scope: 'full', segmentStageIds };
+  }
+
+  const nodeIds = new Set        ();
+  const nodes = allNodes.filter((n) => {
+    if (n.segmentIndex !== targetSegmentIndex) return false;
+    nodeIds.add(n.id);
+    return true;
+  });
+  const edges = allEdges.filter((e) => nodeIds.has(e.source) && nodeIds.has(e.target));
+
+  return { nodes, edges, segmentIndex: targetSegmentIndex, segmentCount, scope: 'segment', segmentStageIds };
+}