fix(api): page postgres kv full scans by keyset (#3465)

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

@@ -0,0 +1,334 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import {defineTable} from '@app/api/database/CassandraTableDsl';
+import type {PreparedQuery} from '@app/api/database/CassandraTypes';
+import {keyFromColumns, PostgresKvQueryExecutor} from '@app/api/database/PostgresKvQueryExecutor';
+import {dockerAvailable, startPostgresKv, stopPostgresKv} from '@app/api/test/DockerTestContainer';
+import type {IPostgresClient} from '@pkgs/postgres/src/Client';
+import {afterAll, beforeAll, beforeEach, describe, expect, it} from 'vitest';
+
+const KV_TABLE = 'kv_keyset_scan';
+
+type Row = Record<string, unknown>;
+
+interface UserRow {
+	user_id: bigint;
+	name: string;
+	kind: number;
+}
+
+interface BucketRow {
+	channel_id: bigint;
+	bucket: number;
+	note: string;
+}
+
+interface NameRow {
+	name: string;
+	value: number;
+}
+
+const KeysetUsers = defineTable<UserRow, 'user_id'>({
+	name: 'keyset_users',
+	columns: ['user_id', 'name', 'kind'],
+	primaryKey: ['user_id'],
+});
+
+const KeysetBuckets = defineTable<BucketRow, 'channel_id' | 'bucket', 'channel_id'>({
+	name: 'keyset_buckets',
+	columns: ['channel_id', 'bucket', 'note'],
+	primaryKey: ['channel_id', 'bucket'],
+	partitionKey: ['channel_id'],
+});
+
+const KeysetNames = defineTable<NameRow, 'name'>({
+	name: 'keyset_names',
+	columns: ['name', 'value'],
+	primaryKey: ['name'],
+});
+
+interface Statement {
+	text: string;
+	rows: number;
+}
+
+function recordingClient(client: IPostgresClient, statements: Array<Statement>): IPostgresClient {
+	return {
+		query: (async (text: string, values?: Array<unknown>, name?: string) => {
+			const result = await client.query(text, values, name);
+			statements.push({text, rows: result.rows.length});
+			return result;
+		}) as IPostgresClient['query'],
+		connect: () => client.connect(),
+		shutdown: () => client.shutdown(),
+		transaction: (fn) => client.transaction(fn),
+		kvTable: () => client.kvTable(),
+	};
+}
+
+function pageState(state: Row): string {
+	return Buffer.from(JSON.stringify(state)).toString('base64url');
+}
+
+function readPageState(state: string): Row {
+	return JSON.parse(Buffer.from(state, 'base64url').toString('utf8')) as Row;
+}
+
+function userId(index: number): bigint {
+	return 1_000_000_000_000_000_000n + BigInt(index) * 7_919n;
+}
+
+describe.skipIf(!dockerAvailable)('Postgres KV keyset scans', () => {
+	let raw: IPostgresClient;
+	let statements: Array<Statement>;
+	let executor: PostgresKvQueryExecutor;
+	let postgresContainer: string;
+
+	beforeAll(async () => {
+		const started = await startPostgresKv({namePrefix: 'fluxer-kvkeyset', kvTable: KV_TABLE, maxConnections: 4});
+		postgresContainer = started.container;
+		raw = started.client;
+	}, 900_000);
+
+	beforeEach(async () => {
+		await raw.query(`DELETE FROM ${KV_TABLE}`);
+		statements = [];
+		executor = new PostgresKvQueryExecutor(recordingClient(raw, statements));
+	});
+
+	afterAll(async () => {
+		await stopPostgresKv(postgresContainer);
+	});
+
+	async function seed(queries: ReadonlyArray<PreparedQuery>): Promise<void> {
+		await executor.executeBatch(queries.map((query) => ({query: query.cql, params: query.params, meta: query.kvMeta})));
+	}
+
+	async function seedUsers(count: number): Promise<Array<UserRow>> {
+		const rows = Array.from({length: count}, (_, index) => ({
+			user_id: userId((index * 37) % count),
+			name: `user-${index}`,
+			kind: index % 7,
+		}));
+		await seed(rows.map((row) => KeysetUsers.insert(row)));
+		return rows;
+	}
+
+	async function seedBuckets(channels: number, buckets: number): Promise<void> {
+		const queries: Array<PreparedQuery> = [];
+		for (let channel = 0; channel < channels; channel += 1) {
+			for (let bucket = 0; bucket < buckets; bucket += 1) {
+				queries.push(
+					KeysetBuckets.insert({
+						channel_id: userId((channel * 13) % channels),
+						bucket: bucket * 11 - 20,
+						note: `note-${(channel + bucket) % 5}`,
+					}),
+				);
+			}
+		}
+		await seed(queries);
+	}
+
+	async function oracle(query: PreparedQuery): Promise<Array<Row>> {
+		const rows = await executor.executeQuery<Row>(query);
+		statements.length = 0;
+		return rows;
+	}
+
+	async function walk(
+		query: PreparedQuery,
+		pageSize: number,
+		start: string | null = null,
+	): Promise<{pages: Array<Array<Row>>; states: Array<string>}> {
+		const pages: Array<Array<Row>> = [];
+		const states: Array<string> = [];
+		let state = start;
+		do {
+			const page = await executor.executePagedQuery<Row>(query, {pageSize, pageState: state});
+			pages.push(page.rows);
+			state = page.pageState;
+			if (state !== null) states.push(state);
+			expect(pages.length).toBeLessThan(10_000);
+		} while (state !== null);
+		return {pages, states};
+	}
+
+	function expectBounded(batch: number): void {
+		expect(statements.length).toBeGreaterThan(0);
+		for (const statement of statements) {
+			expect(statement.rows).toBeLessThanOrEqual(batch);
+			if (statement.text.includes('row_data FROM')) expect(statement.text).toMatch(/LIMIT \$2$/);
+		}
+	}
+
+	function byKey(columns: ReadonlyArray<string>, rows: ReadonlyArray<Row>): Array<Row> {
+		return [...rows].sort((left, right) => {
+			const l = keyFromColumns(columns, left);
+			const r = keyFromColumns(columns, right);
+			return l < r ? -1 : l > r ? 1 : 0;
+		});
+	}
+
+	it('walks a numeric single key table in the same order without a full load', async () => {
+		await seedUsers(250);
+		const query = KeysetUsers.select({}).bind({});
+		const expected = await oracle(query);
+		expect(expected).toHaveLength(250);
+		const {pages, states} = await walk(query, 40);
+		expect(pages.flat()).toEqual(expected);
+		expect(pages.slice(0, -1).every((page) => page.length === 40)).toBe(true);
+		expect(readPageState(states[0]!)).toEqual({
+			offset: 40,
+			after: keyFromColumns(['user_id'], expected[39]!),
+			keyed: true,
+		});
+		expectBounded(41);
+	});
+
+	it('walks a composite key table by row key without reloading it per page', async () => {
+		await seedBuckets(30, 7);
+		const query = KeysetBuckets.select({}).bind({});
+		const expected = await oracle(query);
+		expect(expected).toHaveLength(210);
+		const {pages} = await walk(query, 25);
+		const rows = pages.flat();
+		expect(byKey(['channel_id', 'bucket'], rows)).toEqual(byKey(['channel_id', 'bucket'], expected));
+		expect(new Set(rows.map((row) => keyFromColumns(['channel_id', 'bucket'], row))).size).toBe(210);
+		const channelRuns = rows.filter((row, index) => index === 0 || rows[index - 1]!.channel_id !== row.channel_id);
+		expect(channelRuns).toHaveLength(30);
+		expect(statements.filter((statement) => statement.text.includes('row_data FROM'))).toHaveLength(pages.length);
+		expectBounded(26);
+	});
+
+	it('walks a text key table in the same order', async () => {
+		await seed(
+			Array.from({length: 90}, (_, index) => KeysetNames.insert({name: `name-${(index * 31) % 90}-x`, value: index})),
+		);
+		const query = KeysetNames.select({}).bind({});
+		const expected = await oracle(query);
+		const {pages} = await walk(query, 16);
+		expect(pages.flat()).toEqual(expected);
+		expectBounded(17);
+	});
+
+	it('fills filtered pages from bounded batches on numeric and composite tables', async () => {
+		await seedUsers(300);
+		await seedBuckets(20, 9);
+		const users = KeysetUsers.select({where: KeysetUsers.where.eq('kind')}).bind({kind: 3});
+		const expectedUsers = await oracle(users);
+		const userWalk = await walk(users, 10);
+		expect(userWalk.pages.flat()).toEqual(expectedUsers);
+		expect(userWalk.pages.slice(0, -1).every((page) => page.length === 10)).toBe(true);
+		expectBounded(11);
+		const buckets = KeysetBuckets.select({where: KeysetBuckets.where.eq('note')}).bind({note: 'note-2'});
+		const expectedBuckets = await oracle(buckets);
+		const bucketWalk = await walk(buckets, 6);
+		expect(byKey(['channel_id', 'bucket'], bucketWalk.pages.flat())).toEqual(
+			byKey(['channel_id', 'bucket'], expectedBuckets),
+		);
+		expect(bucketWalk.pages.slice(0, -1).every((page) => page.length === 6)).toBe(true);
+		expectBounded(7);
+	});
+
+	it('skips expired rows', async () => {
+		await seedUsers(60);
+		await seedBuckets(6, 5);
+		await raw.query(
+			`UPDATE ${KV_TABLE} SET expires_at = now() - interval '1 minute' WHERE (row_data ->> 'kind')::int = 2 OR row_data ->> 'note' = 'note-1'`,
+		);
+		const users = KeysetUsers.select({}).bind({});
+		const expectedUsers = await oracle(users);
+		expect(expectedUsers.length).toBeLessThan(60);
+		expect((await walk(users, 7)).pages.flat()).toEqual(expectedUsers);
+		const buckets = KeysetBuckets.select({}).bind({});
+		const expectedBuckets = await oracle(buckets);
+		expect(expectedBuckets.length).toBeLessThan(30);
+		expect(byKey(['channel_id', 'bucket'], (await walk(buckets, 7)).pages.flat())).toEqual(
+			byKey(['channel_id', 'bucket'], expectedBuckets),
+		);
+		expectBounded(8);
+	});
+
+	it('returns one empty page for an empty table', async () => {
+		for (const query of [KeysetUsers.select({}).bind({}), KeysetBuckets.select({}).bind({})]) {
+			expect(await executor.executePagedQuery(query, {pageSize: 10})).toEqual({rows: [], pageState: null});
+		}
+		expectBounded(1);
+	});
+
+	it('never drops rows with non-numeric keys from a single key table', async () => {
+		await seedUsers(40);
+		await executor.executeQuery(KeysetUsers.insert({user_id: 'stray' as unknown as bigint, name: 'stray', kind: 0}));
+		const {pages} = await walk(KeysetUsers.select({}).bind({}), 9);
+		const names = pages.flat().map((row) => row.name);
+		expect(names).toHaveLength(41);
+		expect(new Set(names).size).toBe(41);
+		expect(names).toContain('stray');
+		expectBounded(10);
+	});
+
+	it('resumes every page state shape the previous release produced', async () => {
+		await seedUsers(120);
+		await seedBuckets(12, 6);
+		const users = KeysetUsers.select({}).bind({});
+		const expectedUsers = await oracle(users);
+		const userKey = keyFromColumns(['user_id'], expectedUsers[49]!);
+		for (const state of [{offset: 50, after: userKey, keyed: true}, {offset: 50, after: userKey}, {offset: 50}]) {
+			expect((await walk(users, 15, pageState(state))).pages.flat()).toEqual(expectedUsers.slice(50));
+		}
+		const buckets = KeysetBuckets.select({}).bind({});
+		const expectedBuckets = await oracle(buckets);
+		const bucketKey = keyFromColumns(['channel_id', 'bucket'], expectedBuckets[29]!);
+		statements.length = 0;
+		const resumed = await walk(buckets, 11, pageState({offset: 30, after: bucketKey}));
+		expect(byKey(['channel_id', 'bucket'], resumed.pages.flat())).toEqual(
+			byKey(['channel_id', 'bucket'], expectedBuckets.slice(30)),
+		);
+		expect(readPageState(resumed.states[0]!)).toMatchObject({byRowKey: true, since: bucketKey});
+		expectBounded(12);
+	});
+
+	it('streams limited scans in key order instead of loading the table', async () => {
+		const seeded = await seedUsers(200);
+		const sorted = [...seeded].sort((left, right) => (left.user_id < right.user_id ? -1 : 1));
+		const first = KeysetUsers.select({limit: 20}).bind({});
+		const next = KeysetUsers.select({where: KeysetUsers.where.tokenGt('user_id', 'last_user_id'), limit: 20});
+		const expectedFirst = await oracle(KeysetUsers.select({}).bind({}));
+		expect(await executor.executeQuery(first)).toEqual(expectedFirst.slice(0, 20));
+		expectBounded(21);
+		statements.length = 0;
+		const rows = await executor.executeQuery(next.bind({last_user_id: sorted[99]!.user_id}));
+		expect(rows).toEqual(expectedFirst.slice(100, 120));
+		expect(statements.filter((statement) => statement.text.includes('row_data FROM'))).toHaveLength(1);
+		expectBounded(21);
+		statements.length = 0;
+		expect(await executor.executeQuery(next.bind({last_user_id: sorted[195]!.user_id}))).toEqual(
+			expectedFirst.slice(196),
+		);
+		expectBounded(21);
+	});
+
+	it('probes for non-numeric keys through the primary key while other tables hold string keys', async () => {
+		await raw.query(
+			`INSERT INTO ${KV_TABLE} (table_name, partition_key, row_key, row_data) SELECT 'keyset_users', k, k, '{}'::jsonb FROM generate_series(1, 20000) g, LATERAL (SELECT '{"__fluxer_type":"bigint","value":"' || g || '"}' AS k) s`,
+		);
+		await raw.query(
+			`INSERT INTO ${KV_TABLE} (table_name, partition_key, row_key, row_data) SELECT 'keyset_strings', g::text, '"' || lpad(g::text, 8, '0') || '"', '{}'::jsonb FROM generate_series(1, 200000) g`,
+		);
+		await raw.query(`ANALYZE ${KV_TABLE}`);
+		await executor.executePagedQuery(KeysetUsers.select({}).bind({}), {pageSize: 10});
+		const probe = statements.find((statement) => statement.text.endsWith('AS found'));
+		expect(probe).toBeDefined();
+		const plan = await raw.query<{'QUERY PLAN': string}>(`EXPLAIN (COSTS OFF) ${probe!.text}`, ['keyset_users']);
+		expect(plan.rows.map((row) => row['QUERY PLAN']).join('\n')).not.toContain('Seq Scan');
+	});
+
+	it('keeps the sorted load when a limited scan asks for an order the key walk cannot give', async () => {
+		await seedUsers(30);
+		const descending = KeysetUsers.select({orderBy: {col: 'user_id', direction: 'DESC'}, limit: 5}).bind({});
+		const rows = await executor.executeQuery<Row>(descending);
+		const all = await oracle(KeysetUsers.select({}).bind({}));
+		expect(rows).toEqual(all.slice(-5).reverse());
+	});
+});

fluxer_api/src/api/database/PostgresKvQueryExecutor.ts

@@ -28,6 +28,19 @@ interface PageState {
 	offset: number;
 	after?: string;
 	keyed?: boolean;
+	byRowKey?: boolean;
+	since?: string;
+}
+
+interface ScanWalk {
+	numeric: boolean;
+	after: string | null;
+	since: string | null;
+}
+
+interface NumericFloor {
+	value: string;
+	inclusive: boolean;
 }
 
 interface PageEntry {
@@ -90,6 +103,12 @@ const EXPIRED_STORED_ROW = 'kv.expires_at IS NOT NULL AND kv.expires_at <= now()
 const MERGED_ROW_DATA = `CASE WHEN ${EXPIRED_STORED_ROW} THEN EXCLUDED.row_data ELSE kv.row_data || EXCLUDED.row_data END`;
 const KEPT_EXPIRES_AT = `CASE WHEN ${EXPIRED_STORED_ROW} THEN NULL ELSE kv.expires_at END`;
 const NO_EXPIRY = 'infinity';
+const BIGINT_ROW_KEY_PREFIX = `{"${ENCODED_TYPE_KEY}":"bigint","value":"`;
+const NON_NUMERIC_ROW_KEY_RANGES = [
+	`kv.row_key COLLATE "C" < '-'`,
+	`kv.row_key COLLATE "C" >= ':' AND kv.row_key COLLATE "C" < '${BIGINT_ROW_KEY_PREFIX}'`,
+	`kv.row_key COLLATE "C" >= '${BIGINT_ROW_KEY_PREFIX}:'`,
+];
 
 export async function postgresKvPassIsFresh(
 	client: IPostgresClient,
@@ -676,15 +695,18 @@ function logError(details: Record<string, unknown>, message: string): void {
 
 const fullScanLoggedAt = new Map<string, number>();
 
-function logFullScan(meta: KvQueryMeta): void {
+function logFullScan(meta: KvQueryMeta, walk?: 'numeric' | 'row_key'): void {
 	const shape = queryShape(meta);
-	const key = `${meta.table.name}|${meta.action}|${shape.summary}`;
+	const key = `${meta.table.name}|${meta.action}|${shape.summary}|${walk ?? 'load'}`;
 	const now = Date.now();
 	const last = fullScanLoggedAt.get(key);
 	if (last !== undefined && now - last < FULL_SCAN_LOG_INTERVAL_MS) return;
 	if (fullScanLoggedAt.size >= FULL_SCAN_LOG_KEY_LIMIT) fullScanLoggedAt.clear();
 	fullScanLoggedAt.set(key, now);
-	logWarn({table: meta.table.name, action: meta.action, where: shape.summary || 'none'}, 'Postgres KV full table scan');
+	logWarn(
+		{table: meta.table.name, action: meta.action, where: shape.summary || 'none', walk: walk ?? 'load'},
+		'Postgres KV full table scan',
+	);
 }
 
 function ttlExpiresAt(meta: KvQueryMeta, params: CassandraParams): Date | typeof NO_EXPIRY | null | undefined {
@@ -715,7 +737,13 @@ function decodePageState(pageState: string | null | undefined): PageState {
 		throw new Error('Invalid Postgres KV page state');
 	}
 	if (typeof decoded.after !== 'string') return {offset: decoded.offset};
-	return {offset: decoded.offset, after: decoded.after, keyed: decoded.keyed === true};
+	return {
+		offset: decoded.offset,
+		after: decoded.after,
+		keyed: decoded.keyed === true,
+		byRowKey: decoded.byRowKey === true,
+		...(typeof decoded.since === 'string' ? {since: decoded.since} : {}),
+	};
 }
 
 function pageableSelect(meta: KvQueryMeta, pageSize: number): boolean {
@@ -734,22 +762,30 @@ function numericKeyValue(value: unknown): string | null {
 	return null;
 }
 
-function numericScanPlan(meta: KvQueryMeta, plan: QueryPlan): boolean {
+function numericRowKeyValue(key: string): string | null {
-	return plan.exact && plan.candidates.kind === 'scan' && (meta.table.primaryKey as ReadonlyArray<string>).length === 1;
+	const values = decodeRowKey(key, 1);
+	return values === null ? null : numericKeyValue(values[0]);
 }
 
-function numericScanKeyed(meta: KvQueryMeta, plan: QueryPlan, entries: ReadonlyArray<PageEntry>): boolean {
+function singlePrimaryKey(meta: KvQueryMeta): boolean {
-	if (!numericScanPlan(meta, plan)) return false;
+	return (meta.table.primaryKey as ReadonlyArray<string>).length === 1;
-	const column = (meta.table.primaryKey as ReadonlyArray<string>)[0]!;
-	return entries.every((entry) => numericKeyValue(entry.row[column]) !== null);
 }
 
-function numericScanCursor(state: PageState): {rowKey: string; value: string} | null {
+function walksInPrimaryKeyOrder(meta: KvQueryMeta): boolean {
-	if (state.keyed !== true || state.after === undefined) return null;
+	if (!singlePrimaryKey(meta)) return false;
-	const values = decodeRowKey(state.after, 1);
+	if (meta.orderBy === undefined) return true;
-	if (values === null) return null;
+	return meta.orderBy.col === meta.table.primaryKey[0] && meta.orderBy.direction !== 'DESC';
-	const value = numericKeyValue(values[0]);
+}
-	return value === null ? null : {rowKey: state.after, value};
+
+function numericFloor(meta: KvQueryMeta, params: CassandraParams): NumericFloor | null {
+	const column = meta.table.primaryKey[0];
+	for (const clause of whereClauses(meta)) {
+		if (clause.col !== column) continue;
+		if (clause.kind !== 'gt' && clause.kind !== 'tokenGt' && clause.kind !== 'gte') continue;
+		const value = numericKeyValue(getParam(params, clause.param));
+		if (value !== null) return {value, inclusive: clause.kind === 'gte'};
+	}
+	return null;
 }
 
 function pageStart(meta: KvQueryMeta, entries: ReadonlyArray<PageEntry>, state: PageState): number {
@@ -1018,9 +1054,9 @@ export class PostgresKvQueryExecutor {
 		const meta = this.meta(query);
 		if (pageableSelect(meta, options.pageSize) && (state.after !== undefined || state.offset === 0)) {
 			const plan = buildCandidatePlan(meta, query.params);
-			const cursor = numericScanPlan(meta, plan) ? numericScanCursor(state) : null;
+			const page =
-			const page = cursor
+				plan.candidates.kind === 'scan'
-				? await this.numericScanPage(meta, query.params, cursor, state.offset, options.pageSize)
+					? await this.scanPage(meta, query.params, await this.scanWalk(meta, state), state.offset, options.pageSize)
 					: await this.sortedPage(meta, query.params, plan, state, options.pageSize);
 			return {rows: page.rows as Array<T>, pageState: page.pageState};
 		}
@@ -1033,29 +1069,111 @@ export class PostgresKvQueryExecutor {
 		};
 	}
 
-	private async numericScanPage(
+	private async scanWalk(meta: KvQueryMeta, state: PageState): Promise<ScanWalk> {
+		if (state.byRowKey) return {numeric: false, after: state.after ?? null, since: state.since ?? null};
+		const after = state.after ?? null;
+		const numericAfter = after === null || numericRowKeyValue(after) !== null;
+		if (state.keyed && after !== null && numericAfter && singlePrimaryKey(meta)) {
+			return {numeric: true, after, since: null};
+		}
+		const numeric = singlePrimaryKey(meta) && !(await this.hasNonNumericRows(meta, this.client));
+		if (numeric && numericAfter) return {numeric: true, after, since: null};
+		return {numeric: false, after: null, since: after};
+	}
+
+	private async hasNonNumericRows(meta: KvQueryMeta, db: PostgresQueryable): Promise<boolean> {
+		const probes = NON_NUMERIC_ROW_KEY_RANGES.map(
+			(range) =>
+				`(SELECT kv.row_key FROM ${this.table} kv WHERE kv.table_name = $1 AND ${range} AND (kv.expires_at IS NULL OR kv.expires_at > now()) AND ${NUMERIC_ROW_KEY} IS NULL ORDER BY kv.row_key COLLATE "C" LIMIT 1) IS NOT NULL`,
+		);
+		const result = await db.query<{found: boolean}>(`SELECT ${probes.join(' OR ')} AS found`, [meta.table.name]);
+		return result.rows[0]?.found === true;
+	}
+
+	private async scanPage(
 		meta: KvQueryMeta,
 		params: CassandraParams,
-		cursor: {rowKey: string; value: string},
+		walk: ScanWalk,
 		offset: number,
 		pageSize: number,
 	): Promise<{rows: Array<Row>; pageState: string | null}> {
-		logFullScan(meta);
+		const page = await this.walkScan(meta, params, walk, pageSize, this.client);
-		const result = await this.client.query<StoredRow>(
+		if (!page.more || page.after === null) return {rows: this.projected(meta, page.entries), pageState: null};
-			`SELECT kv.row_key, kv.row_data FROM ${this.table} kv WHERE kv.table_name = $1 AND ${NUMERIC_ROW_KEY} IS NOT NULL AND (${NUMERIC_ROW_KEY}, kv.row_key COLLATE "C") > ($2::numeric, $3) AND (kv.expires_at IS NULL OR kv.expires_at > now()) ORDER BY ${NUMERIC_ROW_KEY}, kv.row_key COLLATE "C" LIMIT $4`,
+		const position = {offset: offset + page.entries.length, after: page.after};
-			[meta.table.name, cursor.value, cursor.rowKey, pageSize + 1],
-		);
-		const entries = this.matchingEntries(meta, result.rows.slice(0, pageSize), params);
-		const last = entries[entries.length - 1];
 		return {
-			rows: this.projected(meta, entries),
+			rows: this.projected(meta, page.entries),
-			pageState:
+			pageState: encodePageState(
-				result.rows.length > pageSize && last
+				walk.numeric
-					? encodePageState({offset: offset + entries.length, after: last.key, keyed: true})
+					? {...position, keyed: true}
-					: null,
+					: {...position, byRowKey: true, ...(walk.since === null ? {} : {since: walk.since})},
+			),
 		};
 	}
 
+	private async walkScan(
+		meta: KvQueryMeta,
+		params: CassandraParams,
+		walk: ScanWalk,
+		want: number,
+		db: PostgresQueryable,
+	): Promise<{entries: Array<PageEntry>; after: string | null; more: boolean}> {
+		logFullScan(meta, walk.numeric ? 'numeric' : 'row_key');
+		const floor = walk.numeric ? numericFloor(meta, params) : null;
+		const since = walk.since === null ? null : decodeRowKey(walk.since, meta.table.primaryKey.length);
+		if (walk.since !== null && since === null) throw new Error('Invalid Postgres KV page state');
+		const limit = want + 1;
+		const entries: Array<PageEntry> = [];
+		let after = walk.after;
+		while (entries.length < want) {
+			const batch = await this.scanBatch(meta, walk.numeric, floor, after, limit, db);
+			for (let index = 0; index < batch.length; index += 1) {
+				const stored = batch[index]!;
+				after = stored.row_key;
+				const entry = this.matchingEntry(meta, stored, params);
+				if (entry && (since === null || compareRowToKeyValues(meta, entry.row, since) > 0)) entries.push(entry);
+				if (entries.length === want) {
+					return {entries, after, more: index < batch.length - 1 || batch.length === limit};
+				}
+			}
+			if (batch.length < limit) return {entries, after, more: false};
+		}
+		return {entries, after, more: false};
+	}
+
+	private async scanBatch(
+		meta: KvQueryMeta,
+		numeric: boolean,
+		floor: NumericFloor | null,
+		after: string | null,
+		limit: number,
+		db: PostgresQueryable,
+	): Promise<Array<StoredRow>> {
+		const params: Array<unknown> = [meta.table.name, limit];
+		let predicate = '';
+		if (numeric) {
+			predicate = ` AND ${NUMERIC_ROW_KEY} IS NOT NULL`;
+			if (floor) {
+				params.push(floor.value);
+				predicate += ` AND ${NUMERIC_ROW_KEY} ${floor.inclusive ? '>=' : '>'} $${params.length}::numeric`;
+			}
+			if (after !== null) {
+				const value = numericRowKeyValue(after);
+				if (value === null) throw new Error('Invalid Postgres KV page state');
+				params.push(value, after);
+				predicate += ` AND (${NUMERIC_ROW_KEY}, kv.row_key COLLATE "C") > ($${params.length - 1}::numeric, $${params.length})`;
+			}
+		} else if (after !== null) {
+			params.push(after);
+			predicate = ` AND kv.row_key COLLATE "C" > $${params.length}`;
+		}
+		const order = numeric ? `${NUMERIC_ROW_KEY}, kv.row_key COLLATE "C"` : 'kv.row_key COLLATE "C"';
+		const result = await db.query<StoredRow>(
+			`SELECT kv.row_key, kv.row_data FROM ${this.table} kv WHERE kv.table_name = $1${predicate} AND (kv.expires_at IS NULL OR kv.expires_at > now()) ORDER BY ${order} LIMIT $2`,
+			params,
+		);
+		return result.rows;
+	}
+
 	private async sortedPage(
 		meta: KvQueryMeta,
 		params: CassandraParams,
@@ -1071,12 +1189,9 @@ export class PostgresKvQueryExecutor {
 		const last = page[page.length - 1];
 		const nextOffset = start + page.length;
 		if (nextOffset >= entries.length || !last) return {rows: this.projected(meta, page), pageState: null};
-		const keyed = numericScanKeyed(meta, plan, entries);
 		return {
 			rows: this.projected(meta, page),
-			pageState: encodePageState(
+			pageState: encodePageState({offset: nextOffset, after: last.key}),
-				keyed ? {offset: nextOffset, after: last.key, keyed: true} : {offset: nextOffset, after: last.key},
-			),
 		};
 	}
 
@@ -1146,17 +1261,22 @@ export class PostgresKvQueryExecutor {
 		stored: ReadonlyArray<StoredRow>,
 		params: CassandraParams,
 	): Array<PageEntry> {
-		const required = queryShape(meta).requiredColumns;
 		const entries: Array<PageEntry> = [];
 		for (const entry of stored) {
-			const row = required ? decodeRowColumns(entry.row_data, required) : decodeRow(entry.row_data);
+			const matched = this.matchingEntry(meta, entry, params);
-			if (matchesWhere(row, meta.where as ReadonlyArray<WhereExpr<Row>> | undefined, params)) {
+			if (matched) entries.push(matched);
-				entries.push({key: entry.row_key, row});
-			}
 		}
 		return entries;
 	}
 
+	private matchingEntry(meta: KvQueryMeta, stored: StoredRow, params: CassandraParams): PageEntry | null {
+		const required = queryShape(meta).requiredColumns;
+		const row = required ? decodeRowColumns(stored.row_data, required) : decodeRow(stored.row_data);
+		return matchesWhere(row, meta.where as ReadonlyArray<WhereExpr<Row>> | undefined, params)
+			? {key: stored.row_key, row}
+			: null;
+	}
+
 	private matchingRows(meta: KvQueryMeta, stored: ReadonlyArray<StoredRow>, params: CassandraParams): Array<Row> {
 		return this.matchingEntries(meta, stored, params).map((entry) => entry.row);
 	}
@@ -1185,6 +1305,16 @@ export class PostgresKvQueryExecutor {
 		) {
 			return this.unorderedScan(meta, meta.limit, db);
 		}
+		if (
+			typeof meta.limit === 'number' &&
+			Number.isInteger(meta.limit) &&
+			plan.candidates.kind === 'scan' &&
+			walksInPrimaryKeyOrder(meta) &&
+			!(await this.hasNonNumericRows(meta, db))
+		) {
+			const walk = {numeric: true, after: null, since: null};
+			return this.projected(meta, (await this.walkScan(meta, params, walk, meta.limit, db)).entries);
+		}
 		let rows = this.matchingRows(meta, await this.candidates(meta, plan, db), params);
 		rows = sortRows(meta, rows);
 		if (typeof meta.limit === 'number') rows = rows.slice(0, meta.limit);

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

@@ -55,15 +55,15 @@ describe('Postgres KV unordered scan', () => {
 		expect(rows).toEqual([{user_id: '2'}, {user_id: '1'}]);
 	});
 
-	it('keeps the sorted full scan when order is asked for', async () => {
+	it('keeps the sorted full scan when the order cannot be walked by key', async () => {
 		const statements: Array<{text: string; values: Array<unknown> | undefined}> = [];
 		const executor = new PostgresKvQueryExecutor(recordingClient(statements));
 		const rows = await executor.executeQuery({
 			cql: '__ordered_probe',
 			params: {},
-			kvMeta: selectMeta({limit: 1}),
+			kvMeta: selectMeta({limit: 1, orderBy: {col: 'user_id', direction: 'DESC'}}),
 		});
 		expect(statements[0]!.text).not.toContain('LIMIT');
-		expect(rows).toEqual([{user_id: '1'}]);
+		expect(rows).toEqual([{user_id: '2'}]);
 	});
 });