@elaraai/e3-types 1.0.77 → 1.0.79

package/src/task.spec.ts

@@ -0,0 +1,208 @@
+/**
+ * Copyright (c) 2025 Elara AI Pty Ltd
+ * Dual-licensed under AGPL-3.0 and commercial license. See LICENSE for details.
+ */
+
+/**
+ * Partition and stream task wire helpers (issue #770): the partition
+ * metadata's merge fields and its triple decoder, the stream metadata's
+ * `merge` mode and its dual decoder, the partition plan round trip, and the
+ * `by` projection shapes the orchestrator evaluates by reading key fields
+ * instead of compiling the projection.
+ */
+
+import { describe, it } from 'node:test';
+import assert from 'node:assert/strict';
+import { ArrayType, BlobType, BooleanType, East, IntegerType, OptionType, StringType, StructType, compareFor, encodeBeast2For, encodeEastIR, isTypeValueEqual, none, some, toEastTypeValue, variant } from '@elaraai/east';
+import type { FunctionIR } from '@elaraai/east';
+import {
+  decodePartitionPlan,
+  decodePartitionTaskMetadata,
+  decodeStreamTaskMetadata,
+  encodePartitionPlan,
+  encodePartitionTaskMetadata,
+  encodeStreamTaskMetadata,
+  partitionProjectionShape,
+  projectKey,
+  projectedKeyType,
+  type PartitionPlan,
+  type PartitionTaskMetadata,
+  type ProjectionShape,
+  type StreamTaskMetadata,
+} from './task.js';
+import { mergeCommandIr } from './stream.js';
+
+describe('PartitionTaskMetadataType', () => {
+  const mergeIr = encodeEastIR(East.function([IntegerType, StringType, StringType], StringType, (_$, _key, acc, value) => acc.concat(value)).toIR());
+  const mergeCommand = encodeEastIR(mergeCommandIr(variant('east_c', { platforms: ['east-c-std'] }), 'function'));
+
+  /** The metadata with every blob as a plain Uint8Array — the decoder hands
+   *  back Buffers, whose prototype strict deep equality would reject. */
+  function plain(meta: PartitionTaskMetadata): PartitionTaskMetadata {
+    const blob = (option: { type: string; value: unknown }) =>
+      option.type === 'some' ? some(new Uint8Array(option.value as Uint8Array)) : none;
+    return { ...meta, by: blob(meta.by), combine: blob(meta.combine), merge: blob(meta.merge), mergeCommand: blob(meta.mergeCommand) };
+  }
+
+  it('round-trips the merge fields and the merge command', () => {
+    const meta: PartitionTaskMetadata = {
+      partitions: 1n, by: none, combine: none, targetPartitionBytes: 1024n,
+      merge: some(mergeIr), mergeSets: false, mergeCommand: some(mergeCommand),
+    };
+    assert.deepEqual(plain(decodePartitionTaskMetadata(encodePartitionTaskMetadata(meta))), meta);
+  });
+
+  it('decodes v1.0.77 metadata, exported before merge commands, with mergeCommand none', () => {
+    const shape = StructType({
+      partitions: IntegerType, by: OptionType(BlobType), combine: OptionType(BlobType), targetPartitionBytes: IntegerType,
+      merge: OptionType(BlobType), mergeSets: BooleanType,
+    });
+    const legacy = encodeBeast2For(shape)({ partitions: 2n, by: none, combine: none, targetPartitionBytes: 512n, merge: some(mergeIr), mergeSets: false });
+    assert.deepEqual(plain(decodePartitionTaskMetadata(legacy)), {
+      partitions: 2n, by: none, combine: none, targetPartitionBytes: 512n, merge: some(mergeIr), mergeSets: false, mergeCommand: none,
+    });
+  });
+
+  it('decodes metadata exported before merge assembly with every merge field off', () => {
+    const shape = StructType({ partitions: IntegerType, by: OptionType(BlobType), combine: OptionType(BlobType), targetPartitionBytes: IntegerType });
+    const legacy = encodeBeast2For(shape)({ partitions: 1n, by: none, combine: none, targetPartitionBytes: 256n });
+    assert.deepEqual(decodePartitionTaskMetadata(legacy), {
+      partitions: 1n, by: none, combine: none, targetPartitionBytes: 256n, merge: none, mergeSets: false, mergeCommand: none,
+    });
+  });
+
+  it('throws for bytes of no known metadata shape', () => {
+    assert.throws(() => decodePartitionTaskMetadata(encodeBeast2For(IntegerType)(7n)));
+  });
+});
+
+describe('StreamTaskMetadataType', () => {
+  it('round-trips the merge mode', () => {
+    const meta: StreamTaskMetadata = { stream: true, emit: 'dict', merge: 'function' };
+    assert.deepEqual(decodeStreamTaskMetadata(encodeStreamTaskMetadata(meta)), meta);
+  });
+
+  it('decodes metadata exported before merge modes as merge none', () => {
+    const legacy = encodeBeast2For(StructType({ stream: BooleanType, emit: StringType }))({ stream: false, emit: 'set' });
+    assert.deepEqual(decodeStreamTaskMetadata(legacy), { stream: false, emit: 'set', merge: 'none' });
+  });
+
+  it('throws for bytes of no known metadata shape', () => {
+    assert.throws(() => decodeStreamTaskMetadata(encodeBeast2For(IntegerType)(7n)));
+  });
+});
+
+describe('PartitionPlanType', () => {
+  it('round-trips a plan, before and after its slices are carved and its merge ranges planned', () => {
+    const plan: PartitionPlan = {
+      partitions: ['a'.repeat(64), 'b'.repeat(64)],
+      boundaries: [0n, 4n, 9n],
+      splits: [[{ seg: 0n, offset: 0n }, { seg: 3n, offset: 17n }, { seg: 8n, offset: 0n }, { seg: 12n, offset: 0n }]],
+      slices: [],
+      merges: [],
+    };
+    assert.deepEqual(decodePartitionPlan(encodePartitionPlan(plan)), plan);
+
+    const carved: PartitionPlan = { ...plan, slices: [['c'.repeat(64), 'd'.repeat(64), 'e'.repeat(64)], ['f'.repeat(64), '0'.repeat(64), '1'.repeat(64)]] };
+    assert.deepEqual(decodePartitionPlan(encodePartitionPlan(carved)), carved);
+
+    const merged: PartitionPlan = {
+      ...carved,
+      merges: [{ partials: ['2'.repeat(64), '3'.repeat(64)], ranges: ['4'.repeat(64), '5'.repeat(64)] }],
+    };
+    assert.deepEqual(decodePartitionPlan(encodePartitionPlan(merged)), merged);
+  });
+
+  it('decodes a plan recorded before merge ranges existed, with none', () => {
+    // The pre-`merges` wire shape, encoded directly.
+    const PreMergesPlanType = StructType({
+      partitions: ArrayType(StringType),
+      boundaries: ArrayType(IntegerType),
+      splits: ArrayType(ArrayType(StructType({ seg: IntegerType, offset: IntegerType }))),
+      slices: ArrayType(ArrayType(StringType)),
+    });
+    const older = encodeBeast2For(PreMergesPlanType)({
+      partitions: ['a'.repeat(64)],
+      boundaries: [0n, 2n],
+      splits: [],
+      slices: [['b'.repeat(64), 'c'.repeat(64)]],
+    });
+    assert.deepEqual(decodePartitionPlan(older), {
+      partitions: ['a'.repeat(64)],
+      boundaries: [0n, 2n],
+      splits: [],
+      slices: [['b'.repeat(64), 'c'.repeat(64)]],
+      merges: [],
+    });
+  });
+});
+
+describe('partition projections', () => {
+  const AtType = StructType({ day: IntegerType, hour: IntegerType });
+  const KeyType = StructType({ region: StringType, at: AtType, sku: StringType });
+  const keyTypeValue = toEastTypeValue(KeyType);
+  type Key = { region: string; at: { day: bigint; hour: bigint }; sku: string };
+  const keys: Key[] = [
+    { region: 'eu', at: { day: 2n, hour: 7n }, sku: 'b' },
+    { region: 'eu', at: { day: 2n, hour: 7n }, sku: 'a' },
+    { region: 'eu', at: { day: 1n, hour: 9n }, sku: 'c' },
+    { region: 'ap', at: { day: 3n, hour: 0n }, sku: 'a' },
+  ];
+
+  /** A `by` projection over the key, compiled, and its shape. */
+  function projection(by: ($: any, key: any) => any): { compiled: (key: Key) => unknown; shape: ProjectionShape | null } {
+    const bundle = East.function([KeyType], undefined, by).toIR();
+    return {
+      compiled: bundle.compile([]) as (key: Key) => unknown,
+      shape: partitionProjectionShape(bundle.ir as FunctionIR),
+    };
+  }
+
+  /** Whether two comparisons agree in sign. */
+  const sameSign = (a: number, b: number): boolean => Math.sign(a) === Math.sign(b);
+
+  it('reads the identity, a leading field, a struct literal of leading fields and a nested first-field path', () => {
+    assert.deepEqual(projection(($, key) => key).shape, { kind: 'fields', names: [] });
+    assert.deepEqual(projection(($, key) => key.region).shape, { kind: 'fields', names: ['region'] });
+    assert.deepEqual(projection(($, key) => ({ region: key.region, at: key.at })).shape, { kind: 'fields', names: ['region', 'at'] });
+    assert.deepEqual(projection(($, key) => key.at.day).shape, { kind: 'path', names: ['at', 'day'] });
+  });
+
+  it('returns null for a projection that computes', () => {
+    assert.equal(projection(($, key) => key.region.concat('-')).shape, null);
+  });
+
+  it('projects the identity to the key itself', () => {
+    const shape: ProjectionShape = { kind: 'fields', names: [] };
+    assert.ok(isTypeValueEqual(projectedKeyType(shape, keyTypeValue), keyTypeValue));
+    assert.equal(projectKey(shape, keys[0]), keys[0]);
+  });
+
+  it('projects a struct literal of leading fields exactly as the compiled projection does', () => {
+    const { compiled, shape } = projection(($, key) => ({ region: key.region, at: key.at }));
+    assert.ok(isTypeValueEqual(projectedKeyType(shape!, keyTypeValue), toEastTypeValue(StructType({ region: StringType, at: AtType }))));
+    for (const key of keys) assert.deepEqual(projectKey(shape!, key), compiled(key));
+  });
+
+  it('projects a nested first-field path exactly as the compiled projection does', () => {
+    const { compiled, shape } = projection(($, key) => key.at.day);
+    assert.ok(isTypeValueEqual(projectedKeyType(shape!, keyTypeValue), toEastTypeValue(IntegerType)));
+    for (const key of keys) assert.equal(projectKey(shape!, key), compiled(key));
+  });
+
+  it('orders a single leading field as the compiled projection does', () => {
+    const { compiled, shape } = projection(($, key) => key.region);
+    const cmp = compareFor(projectedKeyType(shape!, keyTypeValue) as any) as (a: unknown, b: unknown) => number;
+    const cmpCompiled = compareFor(StringType);
+    for (const a of keys) {
+      for (const b of keys) {
+        assert.ok(sameSign(cmp(projectKey(shape!, a), projectKey(shape!, b)), cmpCompiled(compiled(a) as string, compiled(b) as string)));
+      }
+    }
+  });
+
+  it('throws when a shape reads a field the key type does not have', () => {
+    assert.throws(() => projectedKeyType({ kind: 'fields', names: ['store'] }, keyTypeValue), /field 'store'/);
+    assert.throws(() => projectedKeyType({ kind: 'path', names: ['at', 'minute'] }, keyTypeValue), /field 'minute'/);
+  });
+});

package/src/task.ts

@@ -16,7 +16,8 @@
  * specified paths - the task just references locations, not types.
  */
 
-import { StructType, StringType, ArrayType, BlobType, BooleanType, IntegerType, OptionType, ValueTypeOf, decodeBeast2For, encodeBeast2For, none } from '@elaraai/east';
+import { StructType, StringType, ArrayType, BlobType, BooleanType, IntegerType, OptionType, ValueTypeOf, decodeBeast2For, encodeBeast2For, none, variant } from '@elaraai/east';
+import type { EastTypeValue, FunctionIR } from '@elaraai/east';
 import { TreePathType } from './structure.js';
 import { RunnerType } from './runner.js';
 
@@ -141,6 +142,14 @@ export const TASK_KIND_PARTITION = 'partition';
  *  beast2-encoded {@link StreamTaskMetadataType}. */
 export const TASK_KIND_STREAM = 'stream';
 
+/** Task kind of a merge unit — the runner's `merge` command over sorted
+ *  partials of a partitioned task's keyed output (issue #770). The
+ *  orchestrator writes one such task per partitioned task from the
+ *  package's {@link PartitionTaskMetadataType} `mergeCommand`, and runs it
+ *  as an ordinary content-addressed execution per group of partials. It
+ *  carries no metadata. */
+export const TASK_KIND_MERGE = 'merge';
+
 /**
  * Metadata of a {@link TASK_KIND_PARTITION} task.
  *
@@ -171,17 +180,29 @@ export const PartitionTaskMetadataType = StructType({
    * `encodeEastIR` bundle of the per-key merge `(Key, Value, Value) -> Value`
    * for a Dict output; `none` otherwise.
    *
-   * Its presence (or {@link mergeSets}) selects the SEGMENT-MERGE assembly:
-   * partials are walked by their segment fences in the orchestrator, disjoint
-   * segments are byte-copied exactly as a splice does, and only overlapping
-   * ones are decoded, merged and re-encoded. Appended LAST (BEAST2 encodes
+   * Its presence (or {@link mergeSets}) selects the MERGE-TREE assembly:
+   * partials whose key ranges overlap are merged by the task's own runner, in
+   * a tree of merge executions ({@link mergeCommand}) folding equal keys with
+   * this function; disjoint partials are spliced; the orchestrator never
+   * decodes a partial. Appended after `targetPartitionBytes` (BEAST2 encodes
    * struct fields positionally) with a dual decoder — see
    * {@link decodePartitionTaskMetadata}.
    */
   merge: OptionType(BlobType),
-  /** Whether a Set output assembles by segment merge (union). The Set twin of
-   *  {@link merge}, which needs no function. Appended LAST. */
+  /** Whether a Set output assembles by the merge tree, keeping one of equal
+   *  elements (`--union`). The Set twin of {@link merge}, which needs no
+   *  function. */
   mergeSets: BooleanType,
+  /**
+   * `encodeEastIR` bundle of the merge command `(input_paths, output_path) ->
+   * argv` — `mergeCommandIr` over the task's runner, in `function` mode with
+   * {@link merge} and `union` mode with {@link mergeSets} — that the
+   * orchestrator's merge units execute; `none` in splice and combine modes.
+   * Built at export, so the merge unit task is a package object like any
+   * other and nothing is synthesized at run time. Appended LAST with a
+   * triple decoder — see {@link decodePartitionTaskMetadata}.
+   */
+  mergeCommand: OptionType(BlobType),
 });
 export type PartitionTaskMetadataType = typeof PartitionTaskMetadataType;
 
@@ -199,29 +220,49 @@ const PreMergePartitionTaskMetadataType = StructType({
   targetPartitionBytes: IntegerType,
 });
 
+/**
+ * The v1.0.77 partition metadata wire shape — `merge` and `mergeSets` but no
+ * `mergeCommand` — kept only so {@link decodePartitionTaskMetadata} can read
+ * tasks exported by that SDK; their merge units cannot run.
+ */
+const PreMergeCommandPartitionTaskMetadataType = StructType({
+  partitions: IntegerType,
+  by: OptionType(BlobType),
+  combine: OptionType(BlobType),
+  targetPartitionBytes: IntegerType,
+  merge: OptionType(BlobType),
+  mergeSets: BooleanType,
+});
+
 /** Encode a {@link PartitionTaskMetadataType} value for `TaskObject.metadata`. */
 export const encodePartitionTaskMetadata: (value: PartitionTaskMetadata) => Uint8Array =
   encodeBeast2For(PartitionTaskMetadataType);
 
 const decodeCurrentPartitionMetadata = decodeBeast2For(PartitionTaskMetadataType);
+const decodePreMergeCommandPartitionMetadata = decodeBeast2For(PreMergeCommandPartitionTaskMetadataType);
 const decodePreMergePartitionMetadata = decodeBeast2For(PreMergePartitionTaskMetadataType);
 
 /**
  * Decode a `TaskObject.metadata` blob of a {@link TASK_KIND_PARTITION} task,
- * tolerating the pre-`merge` wire format (dual-decode migration).
+ * tolerating the two older wire formats: the v1.0.77 shape without
+ * `mergeCommand`, and the pre-`merge` shape.
  *
  * @param data - the metadata blob
- * @returns the decoded metadata, with `merge`/`mergeSets` defaulted off for
- *   older bytes
+ * @returns the decoded metadata, with `mergeCommand` (and, for the oldest
+ *   bytes, `merge`/`mergeSets`) defaulted off
  */
 export function decodePartitionTaskMetadata(data: Uint8Array): PartitionTaskMetadata {
   try {
     return decodeCurrentPartitionMetadata(data);
   } catch (err) {
     try {
-      return { ...decodePreMergePartitionMetadata(data), merge: none, mergeSets: false };
+      return { ...decodePreMergeCommandPartitionMetadata(data), mergeCommand: none };
     } catch {
-      throw err; // no known shape — surface the current-format error
+      try {
+        return { ...decodePreMergePartitionMetadata(data), merge: none, mergeSets: false, mergeCommand: none };
+      } catch {
+        throw err; // no known shape — surface the current-format error
+      }
     }
   }
 }
@@ -229,28 +270,302 @@ export function decodePartitionTaskMetadata(data: Uint8Array): PartitionTaskMeta
 /**
  * Metadata of a {@link TASK_KIND_STREAM} task.
  *
- * The task's wire `inputs` are laid out `[function_ir, stream?, ...inputs]`;
- * the compiled body takes one trailing `emit` parameter beyond the wire
- * inputs, and the runner writes the `-o` file from the emit sink instead of
- * the body's (Null) return value.
+ * The task's wire `inputs` are laid out `[function_ir, merge_ir?, stream?,
+ * ...inputs]` — `merge_ir` only in `function` merge mode; the compiled body
+ * takes one trailing `emit` parameter beyond the wire inputs, and the runner
+ * writes the `-o` file from the emit sink instead of the body's (Null) return
+ * value.
  */
 export const StreamTaskMetadataType = StructType({
-  /** Whether wire input index 1 is the streamed input (producer tasks have
-   *  no streamed input). */
+  /** Whether the first input after the IRs is the streamed input (producer
+   *  tasks have no streamed input). */
   stream: BooleanType,
   /** The output collection kind the emit sink writes: `"array"`, `"set"`, or
    *  `"dict"` — element/key/value types come from the body IR's emit
    *  parameter. */
   emit: StringType,
+  /**
+   * How the emit sink treats equal keys, a `StreamMergeMode`: `"none"` —
+   * a duplicate key is an error; `"function"` — wire input 1 is the merge IR
+   * `(K, V, V) -> V`, folding equal Dict keys; `"union"` — equal Set elements
+   * collapse. Appended LAST (BEAST2 encodes struct fields positionally) with a
+   * dual decoder — see {@link decodeStreamTaskMetadata}.
+   */
+  merge: StringType,
 });
 export type StreamTaskMetadataType = typeof StreamTaskMetadataType;
 
 export type StreamTaskMetadata = ValueTypeOf<typeof StreamTaskMetadataType>;
 
+/**
+ * The pre-`merge` stream metadata wire shape, kept only so
+ * {@link decodeStreamTaskMetadata} can read tasks exported before folding
+ * emit sinks existed.
+ */
+const PreMergeStreamTaskMetadataType = StructType({
+  stream: BooleanType,
+  emit: StringType,
+});
+
 /** Encode a {@link StreamTaskMetadataType} value for `TaskObject.metadata`. */
 export const encodeStreamTaskMetadata: (value: StreamTaskMetadata) => Uint8Array =
   encodeBeast2For(StreamTaskMetadataType);
 
-/** Decode a `TaskObject.metadata` blob of a {@link TASK_KIND_STREAM} task. */
-export const decodeStreamTaskMetadata: (data: Uint8Array) => StreamTaskMetadata =
-  decodeBeast2For(StreamTaskMetadataType);
+const decodeCurrentStreamMetadata = decodeBeast2For(StreamTaskMetadataType);
+const decodePreMergeStreamMetadata = decodeBeast2For(PreMergeStreamTaskMetadataType);
+
+/**
+ * Decode a `TaskObject.metadata` blob of a {@link TASK_KIND_STREAM} task,
+ * tolerating the pre-`merge` wire format (dual-decode migration).
+ *
+ * @param data - the metadata blob
+ * @returns the decoded metadata, with `merge` defaulted to `"none"` for older
+ *   bytes
+ */
+export function decodeStreamTaskMetadata(data: Uint8Array): StreamTaskMetadata {
+  try {
+    return decodeCurrentStreamMetadata(data);
+  } catch (err) {
+    try {
+      return { ...decodePreMergeStreamMetadata(data), merge: 'none' };
+    } catch {
+      throw err; // no known shape — surface the current-format error
+    }
+  }
+}
+
+// =============================================================================
+// Partition plan
+// =============================================================================
+
+/** A carve position: a segment index and an element offset within it. */
+const SplitPointType = StructType({ seg: IntegerType, offset: IntegerType });
+
+/**
+ * The plan of one merged component's fan-in (issue #770): the partials whose
+ * key ranges overlap, and the key ranges its merge units run over, as the
+ * hashes of the range blobs the units take as their input.
+ *
+ * @remarks
+ * The ranges are a pure function of the partials and the task's
+ * `targetPartitionBytes` — the boundary keys are fences of the largest
+ * partial — so a re-run whose partials are the same objects plans the same
+ * ranges; recording them lets it skip the planning probes, and because
+ * every unit's inputs are then the same objects, every merge unit
+ * cache-hits.
+ */
+export const MergeRangePlanType = StructType({
+  /** The component's partial hashes, in partition order. */
+  partials: ArrayType(StringType),
+  /** The hashes of the component's range blobs — `Struct{from: Option<K>,
+   *  to: Option<K>}` over the output's key type — in key order; one open
+   *  range for a component that merges whole. */
+  ranges: ArrayType(StringType),
+});
+export type MergeRangePlanType = typeof MergeRangePlanType;
+
+export type MergeRangePlan = ValueTypeOf<typeof MergeRangePlanType>;
+
+/**
+ * The plan of a partitioned execution (issue #770): the partitioned inputs,
+ * where each partition starts in every one of them, the carved slices, and
+ * the ranged fan-in of each merged component.
+ *
+ * @remarks
+ * The step interpreter writes it to the object store as the run goes —
+ * after the map step with the carved slices, and again after the reduce
+ * step has planned each component's merge ranges — and points the `plan`
+ * sidecar of the execution's `(taskHash, inputsHash)` directory at it, so a
+ * re-plan or a resume that computes the same `partitions`/`boundaries`/
+ * `splits` reuses the slices instead of carving them again, partition by
+ * partition, and a component of the same partials reuses its ranges.
+ */
+export const PartitionPlanType = StructType({
+  /** Partitioned input hashes, wire order. */
+  partitions: ArrayType(StringType),
+  /** First segment index of each partition of the primary; boundaries[0] = 0. */
+  boundaries: ArrayType(IntegerType),
+  /** Per secondary (partitions[1..]): the split point of every partition plus the end; length boundaries.length + 1. */
+  splits: ArrayType(ArrayType(SplitPointType)),
+  /** slices[input][partition] object hashes; empty until carved. */
+  slices: ArrayType(ArrayType(StringType)),
+  /** The ranged fan-in of each merged component, once its ranges are
+   *  planned. Appended LAST (BEAST2 encodes struct fields positionally) with
+   *  a dual decoder — see {@link decodePartitionPlan}. */
+  merges: ArrayType(MergeRangePlanType),
+});
+export type PartitionPlanType = typeof PartitionPlanType;
+
+export type PartitionPlan = ValueTypeOf<typeof PartitionPlanType>;
+
+/**
+ * The pre-`merges` plan wire shape, kept only so {@link decodePartitionPlan}
+ * can read plans a run recorded before the fan-in ran per key range.
+ */
+const PreMergesPartitionPlanType = StructType({
+  partitions: ArrayType(StringType),
+  boundaries: ArrayType(IntegerType),
+  splits: ArrayType(ArrayType(SplitPointType)),
+  slices: ArrayType(ArrayType(StringType)),
+});
+
+/** Encode a {@link PartitionPlanType} value for the object store. */
+export const encodePartitionPlan: (value: PartitionPlan) => Uint8Array =
+  encodeBeast2For(PartitionPlanType);
+
+const decodeCurrentPlan = decodeBeast2For(PartitionPlanType);
+const decodePreMergesPlan = decodeBeast2For(PreMergesPartitionPlanType);
+
+/**
+ * Decode a {@link PartitionPlanType} object, tolerating the pre-`merges`
+ * wire shape (dual-decode migration): an older plan decodes with no recorded
+ * merge ranges, so its partition slices are reused and its ranges planned.
+ *
+ * @param data - the plan object's bytes
+ * @returns the decoded plan
+ */
+export function decodePartitionPlan(data: Uint8Array): PartitionPlan {
+  try {
+    return decodeCurrentPlan(data);
+  } catch (err) {
+    try {
+      return { ...decodePreMergesPlan(data), merges: [] };
+    } catch {
+      throw err; // no known shape — surface the current-format error
+    }
+  }
+}
+
+// =============================================================================
+// Partition `by` projections
+// =============================================================================
+
+/**
+ * The shape a partition task's `by` projection reads, as extracted from its
+ * IR by {@link partitionProjectionShape}.
+ *
+ * - `fields`: the identity (`names: []`), or a leading prefix of the key's
+ *   top-level fields — one field, or a struct literal of fields in declared
+ *   order. Validated against each dataset's top-level key field order.
+ * - `path`: a nested leading-field path (`key.a.b`, two or more steps).
+ *   Validated per step: each must read the FIRST field of its level's
+ *   struct, which is what keeps the projection monotone in canonical key
+ *   order.
+ */
+export interface ProjectionShape {
+  /** Top-level fields (or the identity), or a nested first-field path. */
+  kind: 'fields' | 'path';
+  /** The fields read, in order; empty for the identity. */
+  names: string[];
+}
+
+/**
+ * Extracts the shape a `by` projection reads — see {@link ProjectionShape} —
+ * or `null` for any other (unaccepted) shape.
+ *
+ * @param ir - the projection's function IR, `(Key) -> Projection`
+ * @returns the shape, or `null` when the body is not an accepted projection
+ */
+export function partitionProjectionShape(ir: FunctionIR): ProjectionShape | null {
+  const param = ir.value.parameters[0]?.value.name;
+  if (param === undefined) return null;
+
+  // Chase wrappers to the expression the body evaluates to.
+  const unwrap = (node: any): any => {
+    for (;;) {
+      if (node.type === 'Block') {
+        const statements = node.value.statements;
+        if (statements.length === 0) return node;
+        node = statements[statements.length - 1];
+      } else if (node.type === 'Return') {
+        node = node.value.value;
+      } else if (node.type === 'As' || node.type === 'WrapRecursive' || node.type === 'UnwrapRecursive') {
+        node = node.value.value;
+      } else {
+        return node;
+      }
+    }
+  };
+
+  const isParam = (node: any): boolean => node.type === 'Variable' && node.value.name === param;
+  const fieldOf = (node: any): string | null =>
+    node.type === 'GetField' && isParam(unwrap(node.value.struct)) ? node.value.field as string : null;
+  // Walks a GetField chain down to the parameter: `key.a.b` → ['a', 'b'].
+  const pathOf = (node: any): string[] | null => {
+    const path: string[] = [];
+    let cur = node;
+    while (cur.type === 'GetField') {
+      path.unshift(cur.value.field as string);
+      cur = unwrap(cur.value.struct);
+    }
+    return isParam(cur) && path.length > 0 ? path : null;
+  };
+
+  const body = unwrap(ir.value.body as any);
+  if (isParam(body)) return { kind: 'fields', names: [] };
+  const chain = pathOf(body);
+  if (chain !== null) {
+    // A one-step chain is a top-level field read — the `fields` family.
+    return chain.length === 1 ? { kind: 'fields', names: chain } : { kind: 'path', names: chain };
+  }
+  if (body.type === 'Struct') {
+    const fields: string[] = [];
+    for (const { value } of body.value.fields) {
+      const f = fieldOf(unwrap(value));
+      if (f === null) return null;
+      fields.push(f);
+    }
+    return { kind: 'fields', names: fields };
+  }
+  return null;
+}
+
+/**
+ * The East type of a projected key — the type {@link projectKey} values have,
+ * and the one to build their comparator for: the key type itself for the
+ * identity, a struct of the named leading fields in declared order, or the
+ * type at the end of a nested first-field path.
+ *
+ * @param shape - the projection's shape
+ * @param keyType - the partitioned dataset's key type
+ * @returns the projected key type
+ * @throws {Error} When the shape reads a field the key type does not have.
+ *
+ * @example
+ * ```ts
+ * const cmp = compareFor(projectedKeyType(shape, keyTypeValue));
+ * cmp(projectKey(shape, a), projectKey(shape, b));
+ * ```
+ */
+export function projectedKeyType(shape: ProjectionShape, keyType: EastTypeValue): EastTypeValue {
+  if (shape.names.length === 0) return keyType;
+  const fieldType = (level: EastTypeValue, name: string): EastTypeValue => {
+    const field = level.type === 'Struct' ? level.value.find((f) => f.name === name) : undefined;
+    if (field === undefined) {
+      throw new Error(`partition projection reads field '${name}' of a key level that has no such field`);
+    }
+    return field.type as EastTypeValue;
+  };
+  if (shape.kind === 'path') return shape.names.reduce(fieldType, keyType);
+  return variant('Struct', shape.names.map((name) => ({ name, type: fieldType(keyType, name) })));
+}
+
+/**
+ * Projects a partition key through a `by` shape, without compiling the
+ * projection: the key itself for the identity, a struct of the named leading
+ * fields built field by field in declared order, or the value at the end of a
+ * nested first-field path.
+ *
+ * @param shape - the projection's shape
+ * @param key - a key of the partitioned dataset
+ * @returns the projected key, a value of {@link projectedKeyType}
+ */
+export function projectKey(shape: ProjectionShape, key: unknown): unknown {
+  if (shape.names.length === 0) return key;
+  if (shape.kind === 'path') {
+    return shape.names.reduce((level, name) => (level as Record<string, unknown>)[name], key);
+  }
+  const projected: Record<string, unknown> = {};
+  for (const name of shape.names) projected[name] = (key as Record<string, unknown>)[name];
+  return projected;
+}