fix(orm): let union sources anchor their own index range

.agents/skills/kitcn/references/features/orm.md

@@ -543,45 +543,47 @@ return await ctx.orm.query.articles.findMany({
 
 ```ts
 return await ctx.orm.query.messages
-  .withIndex("by_from_to")
-  .select()
-  .union([
-    { where: { from: input.me, to: input.them } },
-    { where: { from: input.them, to: input.me } },
-  ])
-  .interleaveBy(["createdAt", "id"])
-  .filter(async (m) => !m.deletedAt)
-  .map(async (m) => ({ ...m, body: m.body.slice(0, 240) }))
-  .paginate({
-    cursor: input.cursor,
-    limit: input.limit,
-    maxScan: 500,
-  });
-```
-
-### Union with index ranges
-
-```ts
-const page = await ctx.orm.query.messages
   .select()
   .union([
     {
       index: {
         name: "by_from_to",
-        range: (q) => q.eq("from", me).eq("to", them),
+        range: (q) => q.eq("from", input.me).eq("to", input.them),
       },
     },
     {
       index: {
         name: "by_from_to",
-        range: (q) => q.eq("from", them).eq("to", me),
+        range: (q) => q.eq("from", input.them).eq("to", input.me),
       },
     },
   ])
   .interleaveBy(["createdAt", "id"])
-  .paginate({ cursor: null, limit: 20 });
+  .filter(async (m) => !m.deletedAt)
+  .map(async (m) => ({ ...m, body: m.body.slice(0, 240) }))
+  .paginate({
+    cursor: input.cursor,
+    limit: input.limit,
+    maxScan: 500,
+  });
 ```
 
+Per source:
+
+- `index: { name, range }` anchors that source on its own range. It overrides
+  the chain-level `.withIndex(...)`; sources that omit it use the chain index.
+- `where` filters that source's rows after the read.
+
+`interleaveBy` fields must be the trailing fields each source is already ordered
+by, so every field before them has to be pinned with `eq` in that source's
+range. Sources may name different indexes when they land on the same trailing
+fields — e.g. `by_author_likes` (authorId, numLikes) with `eq("authorId", ...)`
+merges with `numLikesAndType` (type, numLikes) with `eq("type", ...)` under
+`interleaveBy(["numLikes"])`.
+
+Anchor every source. A shared `.withIndex(...)` plus per-source `where` makes
+each source walk the same range and discard the misses after reading them.
+
 ### Pre-pagination transforms
 
 ```ts

.changeset/cuddly-bats-repeat.md

@@ -0,0 +1,30 @@
+---
+"kitcn": patch
+---
+
+## Features
+
+- Support a per-source `index: { name, range }` on `select().union([...])`, so each source walks its own index range instead of re-walking one shared range and filtering the misses in JS.
+
+```ts
+const page = await db.query.messages
+  .select()
+  .union([
+    {
+      index: {
+        name: 'by_from_to',
+        range: (q) => q.eq('from', me).eq('to', them),
+      },
+    },
+    {
+      index: {
+        name: 'by_from_to',
+        range: (q) => q.eq('from', them).eq('to', me),
+      },
+    },
+  ])
+  .interleaveBy(['createdAt', 'id'])
+  .paginate({ cursor: null, limit: 20 });
+```
+
+- Support union sources anchored on different indexes, as long as each source pins its leading fields with `eq` and ends up ordered by the `interleaveBy` fields.

convex/orm/pipeline.test.ts

@@ -2,6 +2,176 @@ import { expect, test } from 'vitest';
 import schema from '../schema';
 import { convexTest, countDocumentReads, runCtx } from '../setup.testing';
 
+async function seedAuthorPartitionedPosts(baseCtx: {
+  db: {
+    insert: (table: 'posts', doc: Record<string, unknown>) => Promise<any>;
+  };
+}) {
+  // Two author partitions that interleave by numLikes, plus noise that a
+  // correctly anchored source must never read.
+  for (const [authorId, likes] of [
+    ['author-a', [1, 3, 5]],
+    ['author-b', [2, 4, 6]],
+  ] as const) {
+    for (const numLikes of likes) {
+      await baseCtx.db.insert('posts', {
+        text: `${authorId}-${numLikes}`,
+        numLikes,
+        type: 'text',
+        authorId,
+      });
+    }
+  }
+  for (let i = 0; i < 20; i++) {
+    await baseCtx.db.insert('posts', {
+      text: `noise-${i}`,
+      numLikes: 100 + i,
+      type: 'noise',
+      authorId: 'author-z',
+    });
+  }
+}
+
+test('union sources anchor their own index range', async () => {
+  const t = convexTest(schema);
+
+  await t.run(seedAuthorPartitionedPosts);
+
+  await t.run(async (baseCtx) => {
+    const ctx = await runCtx(baseCtx);
+    const reads = countDocumentReads(baseCtx);
+    const result = await ctx.orm.query.posts
+      .select()
+      .union([
+        {
+          index: {
+            name: 'by_author_likes',
+            range: (q) => q.eq('authorId', 'author-a'),
+          },
+        },
+        {
+          index: {
+            name: 'by_author_likes',
+            range: (q) => q.eq('authorId', 'author-b'),
+          },
+        },
+      ])
+      .interleaveBy(['numLikes'])
+      .paginate({ cursor: null, limit: 10 });
+
+    expect(result.page.map((p) => p.numLikes)).toEqual([1, 2, 3, 4, 5, 6]);
+    // 6 anchored rows plus each source's one-row lookahead past its range.
+    // The 20 'author-z' rows sit outside both ranges and are never touched.
+    expect(reads.documents).toBeLessThanOrEqual(8);
+  });
+});
+
+test('union sources can anchor different indexes with a shared order suffix', async () => {
+  const t = convexTest(schema);
+
+  await t.run(seedAuthorPartitionedPosts);
+
+  await t.run(async (baseCtx) => {
+    const ctx = await runCtx(baseCtx);
+    // by_author_likes is (authorId, numLikes) and numLikesAndType is
+    // (type, numLikes). Pinning the leading field on each leaves both ordered
+    // by numLikes, which is all interleaveBy needs.
+    const result = await ctx.orm.query.posts
+      .select()
+      .union([
+        {
+          index: {
+            name: 'by_author_likes',
+            range: (q) => q.eq('authorId', 'author-a'),
+          },
+        },
+        {
+          index: {
+            name: 'numLikesAndType',
+            range: (q) => q.eq('type', 'noise').lt('numLikes', 103),
+          },
+        },
+      ])
+      .interleaveBy(['numLikes'])
+      .paginate({ cursor: null, limit: 10 });
+
+    expect(result.page.map((p) => p.text)).toEqual([
+      'author-a-1',
+      'author-a-3',
+      'author-a-5',
+      'noise-0',
+      'noise-1',
+      'noise-2',
+    ]);
+  });
+});
+
+test('a union source index overrides the chain-level withIndex', async () => {
+  const t = convexTest(schema);
+
+  await t.run(seedAuthorPartitionedPosts);
+
+  await t.run(async (baseCtx) => {
+    const ctx = await runCtx(baseCtx);
+    // by_author alone cannot order by numLikes, so this only resolves if the
+    // per-source anchor replaces the chain-level index rather than merging
+    // with it.
+    const result = await ctx.orm.query.posts
+      .withIndex('by_author')
+      .select()
+      .union([
+        {
+          index: {
+            name: 'by_author_likes',
+            range: (q) => q.eq('authorId', 'author-a'),
+          },
+        },
+        {
+          index: {
+            name: 'by_author_likes',
+            range: (q) => q.eq('authorId', 'author-b'),
+          },
+        },
+      ])
+      .interleaveBy(['numLikes'])
+      .paginate({ cursor: null, limit: 10 });
+
+    expect(result.page.map((p) => p.numLikes)).toEqual([1, 2, 3, 4, 5, 6]);
+  });
+});
+
+test('a union source without an index still falls back to the chain index', async () => {
+  const t = convexTest(schema);
+
+  await t.run(seedAuthorPartitionedPosts);
+
+  await t.run(async (baseCtx) => {
+    const ctx = await runCtx(baseCtx);
+    const result = await ctx.orm.query.posts
+      .withIndex('by_author_likes', (q) => q.eq('authorId', 'author-b'))
+      .select()
+      .union([
+        {
+          index: {
+            name: 'by_author_likes',
+            range: (q) => q.eq('authorId', 'author-a'),
+          },
+        },
+        { where: { numLikes: { gt: 2 } } },
+      ])
+      .interleaveBy(['numLikes'])
+      .paginate({ cursor: null, limit: 10 });
+
+    expect(result.page.map((p) => p.text)).toEqual([
+      'author-a-1',
+      'author-a-3',
+      'author-b-4',
+      'author-a-5',
+      'author-b-6',
+    ]);
+  });
+});
+
 test('select chain union can interleave indexed streams', async () => {
   const t = convexTest(schema);
 

convex/orm/query-builder.typecheck.ts

@@ -260,3 +260,34 @@ void db.query.users.min({
 void db.query.users.max({
   field: 'name',
 });
+
+void db.query.posts
+  .select()
+  .union([
+    {
+      index: {
+        name: 'by_author_likes',
+        range: (q) => q.eq('authorId', 'a').gt('numLikes', 3),
+      },
+    },
+    { index: { name: 'numLikesAndType' }, where: { type: 'text' } },
+  ])
+  .interleaveBy(['numLikes'])
+  .paginate({ cursor: null, limit: 20 });
+
+void db.query.posts.select().union([
+  {
+    // @ts-expect-error union source index must name a declared index
+    index: { name: 'by_nothing' },
+  },
+]);
+
+void db.query.posts.select().union([
+  {
+    index: {
+      name: 'by_author_likes',
+      // @ts-expect-error by_author_likes range cannot use a non-index field
+      range: (q) => q.eq('type', 'text'),
+    },
+  },
+]);