perf(kv): collate key columns in C and drop the duplicate index (#2162)

This commit is contained in:
Hampus authored and GitHub committed 2026-08-30 23:01:06 +02:00
1 parent 0da94965dc
commit f0e7c25e4c
3 files changed
+95 -10

No files matched your search

@@ -693,6 +693,19 @@ WHERE idx.indrelid = to_regclass($1)
}
}
async function postgresKvRowKeyIsCCollated(db: PostgresQueryable, kvTable: string): Promise<boolean> {
const result = await db.query<{c_collated: boolean}>(
`SELECT col.collname = 'C' AND col.collnamespace = 'pg_catalog'::regnamespace AS c_collated
FROM pg_attribute att
JOIN pg_collation col ON col.oid = att.attcollation
WHERE att.attrelid = to_regclass($1)
AND att.attname = 'row_key'
AND NOT att.attisdropped`,
[kvTable],
);
return result.rows[0]?.c_collated === true;
}
export async function ensurePostgresKvSchema(client: IPostgresClient): Promise<void> {
const kvTable = client.kvTable();
const table = quoteIdentifier(kvTable);
@@ -703,8 +716,8 @@ export async function ensurePostgresKvSchema(client: IPostgresClient): Promise<v
await db.query(`
CREATE TABLE IF NOT EXISTS ${table} (
table_name text NOT NULL,
partition_key text NOT NULL,
row_key text NOT NULL,
partition_key text COLLATE "C" NOT NULL,
row_key text COLLATE "C" NOT NULL,
row_data jsonb NOT NULL,
expires_at timestamptz,
updated_at timestamptz NOT NULL DEFAULT now(),
@@ -713,9 +726,11 @@ CREATE TABLE IF NOT EXISTS ${table} (
await db.query(
`CREATE INDEX IF NOT EXISTS ${quoteIdentifier(`${kvTable}_partition_row_idx`)} ON ${table} (table_name, partition_key, row_key)`,
);
await db.query(
`CREATE INDEX IF NOT EXISTS ${quoteIdentifier(`${kvTable}_row_key_c_idx`)} ON ${table} (table_name, row_key COLLATE "C")`,
);
if (!(await postgresKvRowKeyIsCCollated(db, kvTable))) {
await db.query(
`CREATE INDEX IF NOT EXISTS ${quoteIdentifier(`${kvTable}_row_key_c_idx`)} ON ${table} (table_name, row_key COLLATE "C")`,
);
}
await db.query(
`CREATE INDEX IF NOT EXISTS ${quoteIdentifier(`${kvTable}_expires_idx`)} ON ${table} (expires_at) WHERE expires_at IS NOT NULL`,
);
@@ -315,15 +315,52 @@ suite('postgres kv upgrade safety', () => {
) d`);
const schemaDiff = await raw.query<{n: string}>(`
SELECT count(*) AS n FROM (
(SELECT replace(indexdef, '${OLD}', 'KV') FROM pg_indexes WHERE tablename = '${OLD}'
(SELECT replace(indexdef, '${OLD}', 'KV') FROM pg_indexes
WHERE tablename = '${OLD}' AND indexname <> '${OLD}_row_key_c_idx'
EXCEPT SELECT replace(indexdef, '${NEW}', 'KV') FROM pg_indexes WHERE tablename = '${NEW}')
UNION ALL
(SELECT replace(indexdef, '${NEW}', 'KV') FROM pg_indexes WHERE tablename = '${NEW}'
EXCEPT SELECT replace(indexdef, '${OLD}', 'KV') FROM pg_indexes WHERE tablename = '${OLD}')
) d`);
const collations = await raw.query<{tablename: string; attname: string; collname: string}>(`
SELECT cls.relname AS tablename, att.attname, col.collname
FROM pg_attribute att
JOIN pg_class cls ON cls.oid = att.attrelid
JOIN pg_collation col ON col.oid = att.attcollation
WHERE att.attrelid IN ('${OLD}'::regclass, '${NEW}'::regclass)
AND att.attname IN ('partition_key', 'row_key')
ORDER BY cls.relname, att.attname`);
const cIndexes = await raw.query<{tablename: string}>(
`SELECT tablename FROM pg_indexes WHERE indexname IN ('${OLD}_row_key_c_idx', '${NEW}_row_key_c_idx')`,
);
console.log(`key diffs=${diff.rows[0]!.n} index-definition diffs=${schemaDiff.rows[0]!.n}`);
console.log(collations.rows.map((r) => `${r.tablename}.${r.attname}=${r.collname}`).join(' '));
expect(Number(diff.rows[0]!.n)).toBe(0);
expect(Number(schemaDiff.rows[0]!.n)).toBe(0);
expect(collations.rows.filter((r) => r.tablename === OLD).map((r) => r.collname)).toEqual(['default', 'default']);
expect(collations.rows.filter((r) => r.tablename === NEW).map((r) => r.collname)).toEqual(['C', 'C']);
expect(cIndexes.rows.map((r) => r.tablename)).toEqual([OLD]);
});
it('never recollates an existing table and keeps its C index', async () => {
const OLD = `${KV}_keep_old`;
await raw.query(`DROP TABLE IF EXISTS ${OLD}`);
const oldClient = new TableClient(raw, OLD);
await legacyEnsurePostgresKvSchema(oldClient);
await ensurePostgresKvSchema(oldClient);
await ensurePostgresKvSchema(oldClient);
const collations = await raw.query<{attname: string; collname: string}>(`
SELECT att.attname, col.collname
FROM pg_attribute att
JOIN pg_collation col ON col.oid = att.attcollation
WHERE att.attrelid = '${OLD}'::regclass AND att.attname IN ('partition_key', 'row_key')
ORDER BY att.attname`);
const indexes = await raw.query<{indexname: string}>(
`SELECT indexname FROM pg_indexes WHERE tablename = '${OLD}' ORDER BY indexname`,
);
console.log(`after boot on a legacy table: ${indexes.rows.map((r) => r.indexname).join(', ')}`);
expect(collations.rows.map((r) => r.collname)).toEqual(['default', 'default']);
expect(indexes.rows.map((r) => r.indexname)).toContain(`${OLD}_row_key_c_idx`);
});
it('survives three simultaneous boots (two api replicas and a worker)', async () => {