perf(kv): chunk oversized IN lists instead of scanning (#2163)

fluxer_api/src/api/database/PostgresKvExactnessFuzz.test.ts

@@ -62,6 +62,14 @@ function rowKeyMatchesPlan(plan: ReturnType<typeof buildCandidatePlan>, rowKey:
 					Buffer.compare(buf, Buffer.from(lower, 'utf8')) >= 0 &&
 					Buffer.compare(buf, Buffer.from(c.upperBounds[i]!, 'utf8')) < 0,
 			);
+		case 'rangeGroups':
+			return c.groups.some((group) =>
+				group.lowerBounds.some(
+					(lower, i) =>
+						Buffer.compare(buf, Buffer.from(lower, 'utf8')) >= 0 &&
+						Buffer.compare(buf, Buffer.from(group.upperBounds[i]!, 'utf8')) < 0,
+				),
+			);
 		case 'partitionKeys':
 			return c.partitionKeys.includes(rowKey);
 	}

fluxer_api/src/api/database/PostgresKvNonSelectMatrix.test.ts

@@ -13,7 +13,7 @@ import cassandra from 'cassandra-driver';
 import {afterAll, beforeAll, describe, expect, it} from 'vitest';
 import {LegacyPostgresKvQueryExecutor} from './__testref__/LegacyPostgresKvQueryExecutor';
 import type {CassandraParams, KvQueryMeta, KvTableSpec, WhereExpr} from './CassandraTypes';
-import {ensurePostgresKvSchema, PostgresKvQueryExecutor} from './PostgresKvQueryExecutor';
+import {buildCandidatePlan, ensurePostgresKvSchema, PostgresKvQueryExecutor} from './PostgresKvQueryExecutor';
 
 type Row = Record<string, unknown>;
 type AnyMeta = KvQueryMeta<Row>;
@@ -300,6 +300,33 @@ describe.skipIf(!dockerAvailable)('postgres kv non-select matrix', () => {
 		expect(await dump(NEXT_TABLE, Two.name)).toBe(await dump(LEGACY_TABLE, Two.name));
 	}, 300_000);
 
+	it('IN lists above the range cap agree once the plan is chunked into groups', async () => {
+		const rows: Array<Row> = [];
+		for (let k = 0; k < 300; k += 1) {
+			for (const c of [0n, 1n]) rows.push({k: BigInt(k), c, v: k});
+		}
+		await seed(Two, rows);
+		const where = [{kind: 'in', col: 'k', param: 'ks'}] as Array<WhereExpr<Row>>;
+		const params = {ks: Array.from({length: 260}, (_, i) => BigInt(i))} as CassandraParams;
+		const sel = {action: 'select', table: Two, where, columns: Two.columns} as AnyMeta;
+		expect(buildCandidatePlan(sel, params).candidates.kind).toBe('rangeGroups');
+		const ls = await legacy.executeQuery({cql: '__chunk_sel__', params, kvMeta: sel});
+		const ns = await next.executeQuery({cql: '__chunk_sel__', params, kvMeta: sel});
+		expect(ns.length).toBe(520);
+		expect(JSON.stringify(ns, bigintJson)).toBe(JSON.stringify(ls, bigintJson));
+		const cnt = {action: 'count', table: Two, where} as AnyMeta;
+		const lc = await legacy.executeQuery({cql: '__chunk_cnt__', params, kvMeta: cnt});
+		const nc = await next.executeQuery({cql: '__chunk_cnt__', params, kvMeta: cnt});
+		expect(nc).toEqual([{count: 520}]);
+		expect(nc).toEqual(lc);
+		const del = {action: 'delete', table: Two, where} as AnyMeta;
+		await legacy.executeQuery({cql: '__chunk_del__', params, kvMeta: del});
+		await next.executeQuery({cql: '__chunk_del__', params, kvMeta: del});
+		const remaining = await dump(NEXT_TABLE, Two.name);
+		expect(remaining).toBe(await dump(LEGACY_TABLE, Two.name));
+		expect(JSON.parse(remaining).length).toBe(80);
+	}, 300_000);
+
 	it('count with a limit or an order by matches legacy', async () => {
 		await seed(Two, [
 			{k: 1n, c: 1n, v: 1},

fluxer_api/src/api/database/PostgresKvQueryExecutor.ts

@@ -20,11 +20,17 @@ interface PageState {
 	offset: number;
 }
 
+interface RangeGroup {
+	lowerBounds: Array<string>;
+	upperBounds: Array<string>;
+}
+
 export type CandidatePlan =
 	| {kind: 'none'}
 	| {kind: 'rowKeys'; rowKeys: Array<string>}
 	| {kind: 'range'; lowerBound: string; upperBound: string}
 	| {kind: 'ranges'; lowerBounds: Array<string>; upperBounds: Array<string>}
+	| {kind: 'rangeGroups'; groups: Array<RangeGroup>}
 	| {kind: 'partitionKeys'; partitionKeys: Array<string>}
 	| {kind: 'scan'};
 
@@ -418,6 +424,15 @@ function keyRangeUpperBound(prefix: string): string {
 	return `${prefix}${KEY_RANGE_UPPER}`;
 }
 
+function prefixRangeGroups(prefixes: ReadonlyArray<string>): Array<RangeGroup> {
+	const groups: Array<RangeGroup> = [];
+	for (let start = 0; start < prefixes.length; start += MAX_PREFIX_RANGES) {
+		const group = prefixes.slice(start, start + MAX_PREFIX_RANGES);
+		groups.push({lowerBounds: group.map(keyRangeLowerBound), upperBounds: group.map(keyRangeUpperBound)});
+	}
+	return groups;
+}
+
 interface PinnedColumn {
 	values: Array<unknown>;
 	segments: Array<string>;
@@ -430,11 +445,28 @@ function pinnedColumns(clauses: ReadonlyArray<PinnedWhereExpr>, params: Cassandr
 		if (values === null || values.length === 0) break;
 		const segments = keySegments(values);
 		if (segments === null) break;
-		pinned.push({values, segments});
+		pinned.push({values, segments: [...new Set(segments)]});
 	}
 	return pinned;
 }
 
+function boundedLeading(pinned: ReadonlyArray<PinnedColumn>, primaryKeyLength: number): number {
+	for (let leading = pinned.length; leading > 0; leading -= 1) {
+		const cap = leading === primaryKeyLength ? MAX_ROW_KEY_COMBINATIONS : MAX_PREFIX_RANGES;
+		if (combinationCount(pinned.slice(0, leading).map((column) => column.segments)) <= cap) return leading;
+	}
+	return 0;
+}
+
+function chunkableLeading(pinned: ReadonlyArray<PinnedColumn>): number {
+	for (let leading = pinned.length; leading > 0; leading -= 1) {
+		if (combinationCount(pinned.slice(0, leading).map((column) => column.segments)) <= MAX_ROW_KEY_COMBINATIONS) {
+			return leading;
+		}
+	}
+	return 0;
+}
+
 function planIsExact(meta: KvQueryMeta, shape: QueryShape, pinned: ReadonlyArray<PinnedColumn>): boolean {
 	if (pinned.length !== shape.clauseCount) return false;
 	if (meta.limit !== undefined || meta.orderBy !== undefined) return false;
@@ -456,13 +488,8 @@ export function buildCandidatePlan(meta: KvQueryMeta, params: CassandraParams):
 	}
 	const primaryKey = meta.table.primaryKey as ReadonlyArray<string>;
 	const pinned = pinnedColumns(shape.leadingClauses, params);
-	let leading = pinned.length;
-	while (leading > 0) {
-		const segmentLists = pinned.slice(0, leading).map((column) => column.segments);
-		const cap = leading === primaryKey.length ? MAX_ROW_KEY_COMBINATIONS : MAX_PREFIX_RANGES;
-		if (combinationCount(segmentLists) <= cap) break;
-		leading -= 1;
-	}
+	let leading = boundedLeading(pinned, primaryKey.length);
+	if (leading === 0) leading = chunkableLeading(pinned);
 	if (leading > 0) {
 		const columns = pinned.slice(0, leading);
 		const prefixes = keyPrefixes(columns.map((column) => column.segments));
@@ -477,6 +504,9 @@ export function buildCandidatePlan(meta: KvQueryMeta, params: CassandraParams):
 				exact,
 			};
 		}
+		if (prefixes.length > MAX_PREFIX_RANGES) {
+			return {candidates: {kind: 'rangeGroups', groups: prefixRangeGroups(prefixes)}, exact};
+		}
 		return {
 			candidates: {
 				kind: 'ranges',
@@ -501,7 +531,16 @@ export function buildCandidatePlan(meta: KvQueryMeta, params: CassandraParams):
 	return {candidates: {kind: 'scan'}, exact: planIsExact(meta, shape, [])};
 }
 
-export function planFragments(plan: CandidatePlan): PlanFragments {
+function rangeFragments(group: RangeGroup): PlanFragments {
+	const params: Array<unknown> = [];
+	const arms = group.lowerBounds.map((lowerBound, index) => {
+		params.push(lowerBound, group.upperBounds[index]);
+		return `(kv.row_key COLLATE "C" >= $${params.length} AND kv.row_key COLLATE "C" < $${params.length + 1})`;
+	});
+	return {predicate: ` AND (${arms.join(' OR ')})`, params};
+}
+
+function planFragments(plan: Exclude<CandidatePlan, {kind: 'rangeGroups'}>): PlanFragments {
 	switch (plan.kind) {
 		case 'rowKeys':
 			return {predicate: ' AND kv.row_key = ANY($2::text[])', params: [plan.rowKeys]};
@@ -510,14 +549,8 @@ export function planFragments(plan: CandidatePlan): PlanFragments {
 				predicate: ' AND kv.row_key COLLATE "C" >= $2 AND kv.row_key COLLATE "C" < $3',
 				params: [plan.lowerBound, plan.upperBound],
 			};
-		case 'ranges': {
-			const params: Array<unknown> = [];
-			const arms = plan.lowerBounds.map((lowerBound, index) => {
-				params.push(lowerBound, plan.upperBounds[index]);
-				return `(kv.row_key COLLATE "C" >= $${params.length} AND kv.row_key COLLATE "C" < $${params.length + 1})`;
-			});
-			return {predicate: ` AND (${arms.join(' OR ')})`, params};
-		}
+		case 'ranges':
+			return rangeFragments(plan);
 		case 'partitionKeys':
 			return plan.partitionKeys.length === 1
 				? {predicate: ' AND kv.partition_key = $2', params: [plan.partitionKeys[0]]}
@@ -527,6 +560,11 @@ export function planFragments(plan: CandidatePlan): PlanFragments {
 	}
 }
 
+export function planFragmentGroups(plan: CandidatePlan): Array<PlanFragments> {
+	if (plan.kind === 'rangeGroups') return plan.groups.map(rangeFragments);
+	return [planFragments(plan)];
+}
+
 function logWarn(details: Record<string, unknown>, message: string): void {
 	try {
 		Logger.warn(details, message);
@@ -802,10 +840,11 @@ export class PostgresKvQueryExecutor {
 		return meta;
 	}
 
-	private async candidates(meta: KvQueryMeta, plan: QueryPlan, db: PostgresQueryable): Promise<Array<StoredRow>> {
-		if (plan.candidates.kind === 'none') return [];
-		if (plan.candidates.kind === 'scan') logFullScan(meta);
-		const fragments = planFragments(plan.candidates);
+	private async candidateGroup(
+		meta: KvQueryMeta,
+		fragments: PlanFragments,
+		db: PostgresQueryable,
+	): Promise<Array<StoredRow>> {
 		const result = await db.query<StoredRow>(
 			`SELECT kv.row_key, kv.row_data FROM ${this.table} kv WHERE kv.table_name = $1${fragments.predicate} AND (kv.expires_at IS NULL OR kv.expires_at > now())`,
 			[meta.table.name, ...fragments.params],
@@ -813,6 +852,18 @@ export class PostgresKvQueryExecutor {
 		return result.rows;
 	}
 
+	private async candidates(meta: KvQueryMeta, plan: QueryPlan, db: PostgresQueryable): Promise<Array<StoredRow>> {
+		if (plan.candidates.kind === 'none') return [];
+		if (plan.candidates.kind === 'scan') logFullScan(meta);
+		const groups = planFragmentGroups(plan.candidates);
+		if (groups.length === 1) return this.candidateGroup(meta, groups[0]!, db);
+		const byRowKey = new Map<string, StoredRow>();
+		for (const fragments of groups) {
+			for (const stored of await this.candidateGroup(meta, fragments, db)) byRowKey.set(stored.row_key, stored);
+		}
+		return [...byRowKey.values()];
+	}
+
 	private matchingRows(meta: KvQueryMeta, stored: ReadonlyArray<StoredRow>, params: CassandraParams): Array<Row> {
 		const required = queryShape(meta).requiredColumns;
 		return stored
@@ -839,12 +890,15 @@ export class PostgresKvQueryExecutor {
 		}
 		if (plan.candidates.kind === 'none') return [{count: 0}];
 		if (plan.candidates.kind === 'scan') logFullScan(meta);
-		const fragments = planFragments(plan.candidates);
-		const result = await db.query<{count: string}>(
-			`SELECT count(*) AS count FROM ${this.table} kv WHERE kv.table_name = $1${fragments.predicate} AND (kv.expires_at IS NULL OR kv.expires_at > now())`,
-			[meta.table.name, ...fragments.params],
-		);
-		return [{count: Number(result.rows[0]?.count ?? 0)}];
+		let total = 0;
+		for (const fragments of planFragmentGroups(plan.candidates)) {
+			const result = await db.query<{count: string}>(
+				`SELECT count(*) AS count FROM ${this.table} kv WHERE kv.table_name = $1${fragments.predicate} AND (kv.expires_at IS NULL OR kv.expires_at > now())`,
+				[meta.table.name, ...fragments.params],
+			);
+			total += Number(result.rows[0]?.count ?? 0);
+		}
+		return [{count: total}];
 	}
 
 	private async upsert(meta: KvQueryMeta, params: CassandraParams, db: PostgresQueryable): Promise<Array<Row>> {
@@ -907,11 +961,12 @@ DO UPDATE SET partition_key = EXCLUDED.partition_key, row_data = EXCLUDED.row_da
 		if (plan.candidates.kind === 'none') return;
 		if (plan.exact) {
 			if (plan.candidates.kind === 'scan') logFullScan(meta);
-			const fragments = planFragments(plan.candidates);
-			await db.query(
-				`DELETE FROM ${this.table} kv WHERE kv.table_name = $1${fragments.predicate} AND (kv.expires_at IS NULL OR kv.expires_at > now())`,
-				[meta.table.name, ...fragments.params],
-			);
+			for (const fragments of planFragmentGroups(plan.candidates)) {
+				await db.query(
+					`DELETE FROM ${this.table} kv WHERE kv.table_name = $1${fragments.predicate} AND (kv.expires_at IS NULL OR kv.expires_at > now())`,
+					[meta.table.name, ...fragments.params],
+				);
+			}
 			return;
 		}
 		const whereColumns = queryShape(meta).whereColumns;

fluxer_api/src/api/database/PostgresKvQueryPlan.test.ts

@@ -33,6 +33,16 @@ function withinRange(rowKey: string, lowerBound: string, upperBound: string): bo
 	return byteCompare(rowKey, lowerBound) >= 0 && byteCompare(rowKey, upperBound) < 0;
 }
 
+function rangeBounds(plan: CandidatePlan): Array<[string, string]> {
+	if (plan.kind === 'ranges') return plan.lowerBounds.map((lower, index) => [lower, plan.upperBounds[index]!]);
+	if (plan.kind === 'rangeGroups') {
+		return plan.groups.flatMap((group) =>
+			group.lowerBounds.map((lower, index): [string, string] => [lower, group.upperBounds[index]!]),
+		);
+	}
+	return [];
+}
+
 function planAccepts(plan: CandidatePlan, rowKey: string, partitionKey: string): boolean {
 	switch (plan.kind) {
 		case 'none':
@@ -42,7 +52,8 @@ function planAccepts(plan: CandidatePlan, rowKey: string, partitionKey: string):
 		case 'range':
 			return withinRange(rowKey, plan.lowerBound, plan.upperBound);
 		case 'ranges':
-			return plan.lowerBounds.some((lower, index) => withinRange(rowKey, lower, plan.upperBounds[index]!));
+		case 'rangeGroups':
+			return rangeBounds(plan).some(([lower, upper]) => withinRange(rowKey, lower, upper));
 		case 'partitionKeys':
 			return plan.partitionKeys.includes(partitionKey);
 		case 'scan':
@@ -51,8 +62,7 @@ function planAccepts(plan: CandidatePlan, rowKey: string, partitionKey: string):
 }
 
 function matchingRangeCount(plan: CandidatePlan, rowKey: string): number {
-	if (plan.kind !== 'ranges') return 0;
-	return plan.lowerBounds.filter((lower, index) => withinRange(rowKey, lower, plan.upperBounds[index]!)).length;
+	return rangeBounds(plan).filter(([lower, upper]) => withinRange(rowKey, lower, upper)).length;
 }
 
 function expectKeyShape(key: string, separators: number): void {
@@ -362,19 +372,58 @@ describe('PostgresKvQueryExecutor plan size caps', () => {
 		expect(plan.candidates.rowKeys.length).toBe(900);
 	});
 
-	it('caps the number of prefix ranges and degrades to a scan instead', () => {
+	it('chunks the prefix ranges into groups instead of degrading to a scan', () => {
 		const under = TripleKey.select({where: TripleKey.where.in('a', 'ids')}).bind({
 			ids: Array.from({length: 256}, (_, i) => `id${i}`),
 		});
 		expect(planOf(under).plan.candidates.kind).toBe('ranges');
 		const over = TripleKey.select({where: TripleKey.where.in('a', 'ids')}).bind({
-			ids: Array.from({length: 257}, (_, i) => `id${i}`),
+			ids: Array.from({length: 700}, (_, i) => `id${i}`),
 		});
 		const {plan} = planOf(over);
+		expect(plan.candidates.kind).toBe('rangeGroups');
+		expect(plan.exact).toBe(true);
+		if (plan.candidates.kind !== 'rangeGroups') return;
+		expect(plan.candidates.groups.map((group) => group.lowerBounds.length)).toEqual([256, 256, 188]);
+		expect(plan.candidates.groups.map((group) => group.upperBounds.length)).toEqual([256, 256, 188]);
+	});
+
+	it('brackets exactly the matching rows once the prefix ranges are chunked', () => {
+		const ids = [...TRICKY_STRINGS, ...Array.from({length: 260}, (_, i) => `id${i}`)];
+		const query = TripleKey.select({where: TripleKey.where.in('a', 'ids')}).bind({ids});
+		expect(checkQuery(query, tripleRows).kind).toBe('rangeGroups');
+	});
+
+	it('falls back to a scan when even the leading column blows the combination bound', () => {
+		const query = TripleKey.select({where: TripleKey.where.in('a', 'ids')}).bind({
+			ids: Array.from({length: 32_769}, (_, i) => `id${i}`),
+		});
+		const {plan} = planOf(query);
 		expect(plan.candidates.kind).toBe('scan');
 		expect(plan.exact).toBe(false);
 	});
 
+	it('applies the cap to distinct IN values rather than to duplicates', () => {
+		const query = TripleKey.select({where: TripleKey.where.in('a', 'ids')}).bind({
+			ids: Array.from({length: 300}, (_, i) => `id${i % 50}`),
+		});
+		const {plan} = planOf(query);
+		expect(plan.candidates.kind).toBe('ranges');
+		if (plan.candidates.kind !== 'ranges') return;
+		expect(plan.candidates.lowerBounds.length).toBe(50);
+	});
+
+	it('applies the row key combination bound to distinct products too', () => {
+		const query = PairKey.select({where: [PairKey.where.in('a', 'ids'), PairKey.where.in('b', 'seqs')]}).bind({
+			ids: Array.from({length: 200}, (_, i) => `id${i % 50}`),
+			seqs: Array.from({length: 200}, (_, i) => BigInt(i % 50)),
+		});
+		const {plan} = planOf(query);
+		expect(plan.candidates.kind).toBe('rowKeys');
+		if (plan.candidates.kind !== 'rowKeys') return;
+		expect(plan.candidates.rowKeys.length).toBe(2500);
+	});
+
 	it('enforces the cap even when the trailing pinned column is an equality', () => {
 		const ids = Array.from({length: 40}, (_, i) => `id${i}`);
 		const query = WideKey.select({

fluxer_api/src/api/database/PostgresKvUpgradeSafety.test.ts

@@ -412,7 +412,7 @@ FROM generate_series($1::bigint, $1::bigint + 19999) u, generate_series(1, 3) s`
 
 		it('emits a multi-range plan that is not slower than the tier-3 scan it replaces', async () => {
 			const {PushSubscriptions} = await import('../Tables');
-			const {buildCandidatePlan, planFragments} = await import('./PostgresKvQueryExecutor');
+			const {buildCandidatePlan, planFragmentGroups} = await import('./PostgresKvQueryExecutor');
 			const {getKvMeta} = await import('./CassandraMetaRegistry');
 			const db = getDefaultPostgresClient();
 			const explain = async (sql: string, params: Array<unknown>) => {
@@ -431,7 +431,9 @@ FROM generate_series($1::bigint, $1::bigint + 19999) u, generate_series(1, 3) s`
 				const params = {user_ids: Array.from({length: size}, (_, i) => BASE + BigInt(i * 3))};
 				const plan = buildCandidatePlan(meta, params as never);
 				expect(plan.candidates.kind).toBe('ranges');
-				const fragments = planFragments(plan.candidates);
+				const groups = planFragmentGroups(plan.candidates);
+				expect(groups.length).toBe(1);
+				const fragments = groups[0]!;
 				expect(fragments.predicate).not.toContain('unnest');
 				const pushed = await explain(
 					`SELECT kv.row_key, kv.row_data FROM ${PERF} kv WHERE kv.table_name = $1${fragments.predicate} AND (kv.expires_at IS NULL OR kv.expires_at > now())`,
@@ -443,9 +445,22 @@ FROM generate_series($1::bigint, $1::bigint + 19999) u, generate_series(1, 3) s`
 			}
 			expect(deltas.join('\n')).toBe('');
 			const overCap = buildCandidatePlan(meta, {
-				user_ids: Array.from({length: 257}, (_, i) => BASE + BigInt(i * 3)),
+				user_ids: Array.from({length: 700}, (_, i) => BASE + BigInt(i * 3)),
 			} as never);
-			expect(overCap.candidates.kind).toBe('scan');
+			expect(overCap.candidates.kind).toBe('rangeGroups');
+			expect(overCap.exact).toBe(true);
+			const overCapGroups = planFragmentGroups(overCap.candidates);
+			expect(overCapGroups.map((group) => group.params.length)).toEqual([512, 512, 376]);
+			let chunkedMs = 0;
+			for (const fragments of overCapGroups) {
+				const pushed = await explain(
+					`SELECT kv.row_key, kv.row_data FROM ${PERF} kv WHERE kv.table_name = $1${fragments.predicate} AND (kv.expires_at IS NULL OR kv.expires_at > now())`,
+					['push_subscriptions', ...fragments.params],
+				);
+				expect(pushed.plan).not.toContain('Seq Scan');
+				chunkedMs += pushed.ms;
+			}
+			console.log(`n=700 chunked=${chunkedMs.toFixed(2)}ms scan=${scan.ms.toFixed(2)}ms`);
 			expect(SEP).toBe('\u001f');
 		}, 300_000);
 	});