Skip to content

Commit fa5f0db

Browse files
authored
perf(tables): keep the batch unique check on per-value GIN probes (#8978)
* perf(tables): keep the batch unique check on per-value GIN probes The batch unique check ORed every value into one query under LIMIT n. The LIMIT made the planner expect an early exit and walk the table's rows through a btree with every value as a filter, so large batches read the whole table inside the write transaction while holding the unique and schema locks. Split the values into UNION ALL branches of at most 25 ORed predicates with no LIMIT, so each value plans as its own GIN probe. Cast-compared columns (ttl), which no index serves, stay one pass. The predicates, null and in-batch handling, error messages, result shape, and lock order are unchanged. * chore(tables): fold the unique-check test helpers into one * fix(tables): coalesce a null equality cast to undefined * fix(tables): measure the unique-check scans by delta on a flushed GIN index
1 parent 4940b0a commit fa5f0db

4 files changed

Lines changed: 231 additions & 24 deletions

File tree

‎apps/sim/lib/table/column-types/comparison-sql.ts‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,20 @@ import { columnTypeOf } from '@/lib/table/column-types/registry'
33
import { validatedTimestampSql } from '@/lib/table/column-types/timestamp-sql'
44
import type { ColumnDefinition } from '@/lib/table/types'
55

6+
/**
7+
* The database cast an equality match on `column` compares through, or `undefined` when equality is
8+
* JSONB containment. Only a type with an equality projection compares by cast, which no index serves.
9+
*/
10+
export function equalityCast(column: ColumnDefinition): string | undefined {
11+
const definition = columnTypeOf(column)
12+
return definition.valueForEquality ? (definition.jsonbCast ?? undefined) : undefined
13+
}
14+
615
/** Casts stored text for equality, treating malformed explicit timestamps as absent instants. */
716
export function columnTextForEquality(cell: SQL, column: ColumnDefinition): SQL {
8-
const definition = columnTypeOf(column)
9-
if (!definition.valueForEquality || !definition.jsonbCast) return cell
10-
if (definition.timestampValidation) {
11-
return validatedTimestampSql(cell, definition.timestampValidation)
12-
}
13-
return sql`(${cell})::${sql.raw(definition.jsonbCast)}`
17+
const cast = equalityCast(column)
18+
if (!cast) return cell
19+
const { timestampValidation } = columnTypeOf(column)
20+
if (timestampValidation) return validatedTimestampSql(cell, timestampValidation)
21+
return sql`(${cell})::${sql.raw(cast)}`
1422
}

‎apps/sim/lib/table/rows/row-writes.integration.ts‎

Lines changed: 186 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ import { lockUniqueColumns, lockUniqueValues } from '@/lib/table/rows/unique-loc
4444
import { getTableById } from '@/lib/table/service'
4545
import { getOrCreateTableSnapshot } from '@/lib/table/snapshot-cache'
4646
import type { ColumnDefinition, JsonValue, RowData, TableDefinition } from '@/lib/table/types'
47+
import { checkBatchUniqueConstraintsDb } from '@/lib/table/validation'
4748

4849
const url = readTestDatabaseUrl()
4950
if (process.env.DATABASE_URL !== url) {
@@ -87,7 +88,7 @@ async function createTable(columns: ColumnDefinition[]): Promise<TableDefinition
8788

8889
async function seedRows(
8990
tableId: string,
90-
rows: Array<{ id: string; data: RowData; orderKey: string | null }>
91+
rows: Array<{ id: string; data: RowData; orderKey: string | null; position?: number }>
9192
) {
9293
await db.insert(userTableRows).values(rows.map((row) => ({ ...row, tableId, workspaceId })))
9394
}
@@ -481,6 +482,33 @@ describe('table row writes against real PostgreSQL', () => {
481482
expect(await storedCount(table.id, { email: 'dup@example.test' })).toBe(1)
482483
})
483484

485+
it('rejects the second of two concurrent batches over the value-lock cap that share a value', async () => {
486+
const table = await createTable(uniqueColumns)
487+
const batch = (prefix: string) => () =>
488+
batchInsertRows(
489+
{
490+
tableId: table.id,
491+
workspaceId,
492+
rows: [
493+
...Array.from({ length: 100 }, (_, i) => ({ email: `${prefix}-${i}@example.test` })),
494+
{ email: 'shared@example.test' },
495+
],
496+
secretProvenance: undefined,
497+
capabilityGovernedUserId: null,
498+
},
499+
table,
500+
'unique-race'
501+
)
502+
503+
const results = await raceUnderHeldDefinitionRow(table.id, [batch('left'), batch('right')], {
504+
onDefinitionRow: 1,
505+
onValueLock: 1,
506+
})
507+
508+
expect(fulfilled(results)).toBe(1)
509+
expect(await storedCount(table.id, { email: 'shared@example.test' })).toBe(1)
510+
})
511+
484512
it('rejects an import batch racing a single insert of the same value', async () => {
485513
const table = await createTable(uniqueColumns)
486514

@@ -700,6 +728,163 @@ describe('table row writes against real PostgreSQL', () => {
700728
})
701729
})
702730

731+
describe('batch unique check across many values', () => {
732+
/** Seeds `cells` at positions 0.., which the conflict messages report one-based. */
733+
const seedAt = (tableId: string, cells: RowData[]) =>
734+
seedRows(
735+
tableId,
736+
cells.map((data, position) => ({
737+
id: `${tableId}-${position}`,
738+
data,
739+
orderKey: `a${position}`,
740+
position,
741+
}))
742+
)
743+
744+
/**
745+
* Index and sequential scans this backend has counted on `user_table_rows`, by relation. The
746+
* counters are the backend's unflushed statistics, which can still hold scans from earlier
747+
* transactions on the same pooled connection, so callers diff two readings.
748+
*/
749+
async function rowScans(trx: DbTransaction): Promise<Map<string, number>> {
750+
const scans = await trx.execute<{ relname: string; scans: number }>(sql`
751+
SELECT c.relname, pg_stat_get_xact_numscans(c.oid)::int AS scans
752+
FROM pg_class c
753+
WHERE c.oid = 'user_table_rows'::regclass
754+
OR c.oid IN (SELECT indexrelid FROM pg_index WHERE indrelid = 'user_table_rows'::regclass)`)
755+
return new Map(scans.map((row) => [row.relname, row.scans]))
756+
}
757+
758+
/** Runs the check in a transaction and returns its result with the scans the check alone ran. */
759+
const checkWithScans = (table: TableDefinition, rows: RowData[]) =>
760+
db.transaction(async (trx) => {
761+
// A parallel worker's scans are not counted in the leader's transaction statistics.
762+
await trx.execute(sql`SET LOCAL max_parallel_workers_per_gather = 0`)
763+
const before = await rowScans(trx)
764+
const result = await checkBatchUniqueConstraintsDb(table.id, rows, table.schema, trx)
765+
const scans: Record<string, number> = {}
766+
for (const [relname, count] of await rowScans(trx)) {
767+
const ran = count - (before.get(relname) ?? 0)
768+
if (ran > 0) scans[relname] = ran
769+
}
770+
return { result, scans }
771+
})
772+
773+
it('reports stored and in-batch conflicts across every probe group', async () => {
774+
const table = await createTable([
775+
{ id: 'email', name: 'email', type: 'string', unique: true },
776+
{ id: 'meta', name: 'meta', type: 'json', unique: true },
777+
])
778+
await seedAt(table.id, [
779+
{ email: 'first@example.test' },
780+
{ email: 'middle@example.test' },
781+
{ email: 'last@example.test' },
782+
{ meta: { a: 1, b: 2 } },
783+
{ meta: [1, 2] },
784+
])
785+
const rows: RowData[] = Array.from({ length: 260 }, (_, i) => ({
786+
email: `new-${i}@example.test`,
787+
}))
788+
rows[0] = { email: 'new-0@example.test', meta: { b: 2, a: 1 } }
789+
rows[1] = { email: 'new-1@example.test', meta: { a: 1 } }
790+
rows[2] = { email: 'new-2@example.test', meta: [2, 1] }
791+
rows[3] = { email: 'first@example.test', meta: [1, 2] }
792+
rows[7] = { email: null }
793+
rows[8] = {}
794+
rows[130] = { email: 'middle@example.test' }
795+
rows[200] = { email: 'new-5@example.test' }
796+
rows[259] = { email: 'last@example.test' }
797+
798+
expect((await checkWithScans(table, rows)).result).toEqual({
799+
valid: false,
800+
errors: [
801+
{
802+
row: 0,
803+
errors: [
804+
'Column "meta" must be unique. Value "[object Object]" already exists in row 4',
805+
],
806+
},
807+
{
808+
row: 3,
809+
errors: [
810+
'Column "email" must be unique. Value "first@example.test" already exists in row 1',
811+
'Column "meta" must be unique. Value "1,2" already exists in row 5',
812+
],
813+
},
814+
{
815+
row: 130,
816+
errors: [
817+
'Column "email" must be unique. Value "middle@example.test" already exists in row 2',
818+
],
819+
},
820+
{
821+
row: 200,
822+
errors: [
823+
'Column "email" must be unique. Value "new-5@example.test" duplicates row 6 in batch',
824+
],
825+
},
826+
{
827+
row: 259,
828+
errors: [
829+
'Column "email" must be unique. Value "last@example.test" already exists in row 3',
830+
],
831+
},
832+
],
833+
})
834+
})
835+
836+
it('matches an expiration by instant in one pass over the table', async () => {
837+
const table = await createTable([
838+
{ id: 'expires', name: 'expires', type: 'ttl', unique: true },
839+
])
840+
await seedAt(table.id, [{ expires: '2024-01-31T00:00:00-00:00' }])
841+
const rows: RowData[] = Array.from({ length: 30 }, (_, i) => ({
842+
expires: `2025-01-${String(i + 1).padStart(2, '0')}T00:00:00-00:00`,
843+
}))
844+
rows[27] = { expires: '2024-01-31T01:00:00+01:00' }
845+
rows[28] = { expires: null }
846+
847+
const { result, scans } = await checkWithScans(table, rows)
848+
849+
expect(result).toEqual({
850+
valid: false,
851+
errors: [
852+
{
853+
row: 27,
854+
errors: [
855+
'Column "expires" must be unique. Value "2024-01-31T01:00:00+01:00" already exists in row 1',
856+
],
857+
},
858+
],
859+
})
860+
expect(Object.values(scans).reduce((sum, count) => sum + count, 0)).toBe(1)
861+
})
862+
863+
it('probes the containment index once per value on a large table', async () => {
864+
const table = await createTable([
865+
{ id: 'email', name: 'email', type: 'string', unique: true },
866+
])
867+
await control`INSERT INTO user_table_rows (id, table_id, workspace_id, data, position, order_key)
868+
SELECT ${table.id}::text || '-' || i, ${table.id}, ${workspaceId},
869+
jsonb_build_object('email', 'stored-' || i || '@example.test'), i, 'a' || i
870+
FROM generate_series(0, 39999) AS i`
871+
// A bulk insert leaves its entries in the GIN pending list, which a manual ANALYZE does not
872+
// merge and which the planner charges to every probe, so flush it for a deterministic plan.
873+
await control`SELECT gin_clean_pending_list('user_table_rows_tenant_data_gin_idx'::regclass)`
874+
await control`ANALYZE user_table_rows`
875+
const rows: RowData[] = Array.from({ length: 260 }, (_, i) => ({
876+
email: i % 10 === 0 ? `stored-${i * 100}@example.test` : `new-${i}@example.test`,
877+
}))
878+
879+
const { result, scans } = await checkWithScans(table, rows)
880+
881+
await control`DELETE FROM user_table_rows WHERE table_id = ${table.id}`
882+
883+
expect(result.errors).toHaveLength(26)
884+
expect(scans).toEqual({ user_table_rows_tenant_data_gin_idx: 260 })
885+
}, 60_000)
886+
})
887+
703888
describe('writers validate against the live schema, not their snapshot', () => {
704889
const columns: ColumnDefinition[] = [
705890
{ id: 'key', name: 'key', type: 'string', unique: true },

‎apps/sim/lib/table/sql.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import {
1919
SINGLE_SELECT_OPERATORS,
2020
SINGLE_SELECT_OPS,
2121
} from '@/lib/table/column-types'
22-
import { columnTextForEquality } from '@/lib/table/column-types/comparison-sql'
22+
import { columnTextForEquality, equalityCast } from '@/lib/table/column-types/comparison-sql'
2323
import { NAME_PATTERN } from '@/lib/table/constants'
2424
import { normalizeDateCellValue } from '@/lib/table/dates'
2525
import { TableQueryValidationError } from '@/lib/table/errors'
@@ -630,10 +630,10 @@ export function fieldPredicate(
630630
: value
631631

632632
const equalityClause = (operand: JsonValue): SQL => {
633-
const definition = column && columnTypeOf(column)
634-
if (column && operand !== null && definition?.valueForEquality && definition.jsonbCast) {
633+
const cast = column && equalityCast(column)
634+
if (column && operand !== null && cast) {
635635
const cell = columnTextForEquality(sql.raw(`${tableName}.data->>'${field}'`), column)
636-
return sql`COALESCE(${cell} = ${operand}::${sql.raw(definition.jsonbCast)}, false)`
636+
return sql`COALESCE(${cell} = ${operand}::${sql.raw(cast)}, false)`
637637
}
638638
return buildContainmentClause(tableName, field, operand)
639639
}

‎apps/sim/lib/table/validation.ts‎

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import { db } from '@sim/db'
66
import { userTableRows } from '@sim/db/schema'
7+
import { chunkArray } from '@sim/utils/helpers'
78
import { and, eq, or, type SQL, sql } from 'drizzle-orm'
89
import { NextResponse } from 'next/server'
910
import { canonicalJson } from '@/lib/api/cursor-binding'
@@ -19,6 +20,7 @@ import {
1920
validateColumnTypeLimits,
2021
validateTypeMetadata,
2122
} from '@/lib/table/column-types'
23+
import { equalityCast } from '@/lib/table/column-types/comparison-sql'
2224
import {
2325
getMaxRowSizeBytes,
2426
NAME_PATTERN,
@@ -537,6 +539,13 @@ async function withSeqscanOffOn(
537539
*/
538540
type UniqueCheckExecutor = Pick<typeof db, 'select' | 'execute'>
539541

542+
/**
543+
* Most values one branch of the batch conflict query ORs together. The branches join with UNION ALL
544+
* and carry no LIMIT, so each value stays a GIN probe: under a LIMIT the planner walks every row of
545+
* the table instead, and a much larger OR adds a scan of the table's whole entry in that index.
546+
*/
547+
const UNIQUE_PROBE_GROUP_SIZE = 25
548+
540549
/**
541550
* Checks unique constraints for a batch of rows using targeted database queries.
542551
* Validates both against existing database rows and within the batch itself.
@@ -615,8 +624,7 @@ export async function checkBatchUniqueConstraintsDb(
615624
throw new Error(`Invalid column id: ${columnId}`)
616625
}
617626

618-
const valueArray = Array.from(values)
619-
const valueConditions = valueArray.map((normalizedValue) => {
627+
const valueConditions = Array.from(values, (normalizedValue) => {
620628
// Reconstruct the original typed value from its normalized key. Both
621629
// directions go through JSON unconditionally: keying the write on the
622630
// value's RUNTIME type while keying the read on the column's DECLARED
@@ -627,20 +635,26 @@ export async function checkBatchUniqueConstraintsDb(
627635
return uniqueValuePredicate(USER_TABLE_ROWS_SQL_NAME, columnId, originalValue, column)
628636
})
629637

630-
const conflictingRows = await ex
631-
.select({
632-
id: userTableRows.id,
633-
data: userTableRows.data,
634-
position: userTableRows.position,
635-
})
636-
.from(userTableRows)
637-
.where(and(eq(userTableRows.tableId, tableId), or(...valueConditions)))
638-
.limit(valueArray.length) // We only need up to one conflict per value
638+
// A cast comparison is a filter over the table's rows whatever its shape, so it stays one
639+
// pass rather than one per group.
640+
const groups = equalityCast(column)
641+
? [valueConditions]
642+
: chunkArray(valueConditions, UNIQUE_PROBE_GROUP_SIZE)
643+
const conflictingRows = await ex.execute<{ data: RowData; position: number }>(
644+
sql.join(
645+
groups.map((group) =>
646+
ex
647+
.select({ data: userTableRows.data, position: userTableRows.position })
648+
.from(userTableRows)
649+
.where(and(eq(userTableRows.tableId, tableId), or(...group)))
650+
),
651+
sql` UNION ALL `
652+
)
653+
)
639654

640655
// Map conflicts back to batch rows
641656
for (const conflict of conflictingRows) {
642-
const conflictData = conflict.data as RowData
643-
const conflictValue = cellOf(conflictData, columnId)
657+
const conflictValue = cellOf(conflict.data, columnId)
644658
if (conflictValue === null || conflictValue === undefined) continue
645659
// Keyed like the batch, since stored jsonb comes back with its keys reordered.
646660
const normalizedConflictValue = uniqueValueKey(conflictValue, column)

0 commit comments

Comments
 (0)