|
| 1 | +/** |
| 2 | + * Date tag filters in real PostgreSQL: each range selects exactly the rows the calendar-day |
| 3 | + * comparison does, and a selective one is served by the slot index. |
| 4 | + */ |
| 5 | + |
| 6 | +import { document, embedding } from '@sim/db/schema' |
| 7 | +import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure' |
| 8 | +import { withUtcTimestamps } from '@sim/db/timestamps' |
| 9 | +import { generateId } from '@sim/utils/id' |
| 10 | +import { type SQL, sql } from 'drizzle-orm' |
| 11 | +import { drizzle } from 'drizzle-orm/postgres-js' |
| 12 | +import postgres from 'postgres' |
| 13 | +import { afterAll, beforeAll, describe, expect, it } from 'vitest' |
| 14 | +import { buildTagFilterCondition } from '@/lib/knowledge/documents/tag-filter' |
| 15 | +import { getStructuredTagFilters } from '@/lib/knowledge/search/tag-filters' |
| 16 | + |
| 17 | +const DAY = '2026-03-15' |
| 18 | +const NEXT_DAY = '2026-03-16' |
| 19 | + |
| 20 | +/** Every instant a day boundary can misplace, plus an untagged row. */ |
| 21 | +const TAG_VALUES = [ |
| 22 | + '2026-03-14 23:59:59.999999', |
| 23 | + '2026-03-15 00:00:00', |
| 24 | + '2026-03-15 12:00:00', |
| 25 | + '2026-03-15 23:59:59.999999', |
| 26 | + '2026-03-16 00:00:00', |
| 27 | + '2026-03-16 23:59:59.999999', |
| 28 | + '2026-03-17 00:00:00', |
| 29 | + null, |
| 30 | +] |
| 31 | + |
| 32 | +/** Session time zones on both sides of UTC, far enough to move a date across midnight. */ |
| 33 | +const TIME_ZONES = ['UTC', 'Pacific/Kiritimati', 'Pacific/Pago_Pago'] |
| 34 | + |
| 35 | +interface DateFilter { |
| 36 | + operator: string |
| 37 | + value: string |
| 38 | + valueTo?: string |
| 39 | +} |
| 40 | + |
| 41 | +const FILTERS: DateFilter[] = [ |
| 42 | + { operator: 'eq', value: DAY }, |
| 43 | + { operator: 'neq', value: DAY }, |
| 44 | + { operator: 'gt', value: DAY }, |
| 45 | + { operator: 'gte', value: DAY }, |
| 46 | + { operator: 'lt', value: DAY }, |
| 47 | + { operator: 'lte', value: DAY }, |
| 48 | + { operator: 'between', value: DAY, valueTo: DAY }, |
| 49 | + { operator: 'between', value: DAY, valueTo: NEXT_DAY }, |
| 50 | +] |
| 51 | + |
| 52 | +/** The calendar-day comparison date filters have always meant, written as a `::date` cast. */ |
| 53 | +function calendarDayCondition(column: SQL, { operator, value, valueTo }: DateFilter): SQL { |
| 54 | + switch (operator) { |
| 55 | + case 'neq': |
| 56 | + return sql`${column}::date != ${value}::date` |
| 57 | + case 'gt': |
| 58 | + return sql`${column}::date > ${value}::date` |
| 59 | + case 'gte': |
| 60 | + return sql`${column}::date >= ${value}::date` |
| 61 | + case 'lt': |
| 62 | + return sql`${column}::date < ${value}::date` |
| 63 | + case 'lte': |
| 64 | + return sql`${column}::date <= ${value}::date` |
| 65 | + case 'between': |
| 66 | + return sql`${column}::date >= ${value}::date AND ${column}::date <= ${valueTo}::date` |
| 67 | + default: |
| 68 | + return sql`${column}::date = ${value}::date` |
| 69 | + } |
| 70 | +} |
| 71 | + |
| 72 | +function searchCondition(filter: DateFilter): SQL { |
| 73 | + const [condition] = getStructuredTagFilters( |
| 74 | + [{ tagSlot: 'date1', fieldType: 'date', ...filter }], |
| 75 | + embedding |
| 76 | + ) |
| 77 | + return condition |
| 78 | +} |
| 79 | + |
| 80 | +function documentCondition(filter: DateFilter): SQL { |
| 81 | + const condition = buildTagFilterCondition({ tagSlot: 'date1', fieldType: 'date', ...filter }) |
| 82 | + if (!condition) throw new Error(`No document predicate for ${filter.operator}`) |
| 83 | + return condition |
| 84 | +} |
| 85 | + |
| 86 | +interface PlanNode { |
| 87 | + 'Index Name'?: string |
| 88 | + Plans?: PlanNode[] |
| 89 | +} |
| 90 | + |
| 91 | +function indexNames(node: PlanNode): string[] { |
| 92 | + return [ |
| 93 | + ...(node['Index Name'] ? [node['Index Name']] : []), |
| 94 | + ...(node.Plans ?? []).flatMap(indexNames), |
| 95 | + ] |
| 96 | +} |
| 97 | + |
| 98 | +describe('date tag filters in PostgreSQL', () => { |
| 99 | + const schemaName = `date_tag_filters_${generateId().replaceAll('-', '')}` |
| 100 | + const connection = postgres( |
| 101 | + readTestDatabaseUrl(), |
| 102 | + withUtcTimestamps({ |
| 103 | + max: 1, |
| 104 | + prepare: false, |
| 105 | + fetch_types: false, |
| 106 | + connection: { search_path: schemaName }, |
| 107 | + onnotice: () => {}, |
| 108 | + }) |
| 109 | + ) |
| 110 | + const db = drizzle(connection) |
| 111 | + |
| 112 | + /** Copies the shipped slot index definition, so the plan proves the migration's index. */ |
| 113 | + async function copyIndex(indexName: string) { |
| 114 | + const [{ indexdef }] = await connection<{ indexdef: string }[]>` |
| 115 | + SELECT indexdef FROM pg_indexes WHERE schemaname = 'public' AND indexname = ${indexName}` |
| 116 | + await connection.unsafe(indexdef.replace(' ON public.', ' ON ')) |
| 117 | + } |
| 118 | + |
| 119 | + async function matchingIds(table: typeof document | typeof embedding, condition: SQL) { |
| 120 | + const rows = await db.select({ id: table.id }).from(table).where(condition) |
| 121 | + return rows.map((row) => row.id).sort() |
| 122 | + } |
| 123 | + |
| 124 | + async function plannedIndexes(table: typeof document | typeof embedding, condition: SQL) { |
| 125 | + const query = db.select({ id: table.id }).from(table).where(condition).toSQL() |
| 126 | + const [row] = await connection.unsafe<{ 'QUERY PLAN': [{ Plan: PlanNode }] }[]>( |
| 127 | + `EXPLAIN (FORMAT JSON) ${query.sql}`, |
| 128 | + query.params as never[] |
| 129 | + ) |
| 130 | + return indexNames(row['QUERY PLAN'][0].Plan) |
| 131 | + } |
| 132 | + |
| 133 | + beforeAll(async () => { |
| 134 | + await connection`CREATE SCHEMA ${connection(schemaName)}` |
| 135 | + await connection`CREATE TABLE document (LIKE public.document INCLUDING DEFAULTS)` |
| 136 | + await connection`CREATE TABLE embedding (LIKE public.embedding INCLUDING DEFAULTS INCLUDING GENERATED)` |
| 137 | + }) |
| 138 | + |
| 139 | + afterAll(async () => { |
| 140 | + try { |
| 141 | + await connection`DROP SCHEMA ${connection(schemaName)} CASCADE` |
| 142 | + } finally { |
| 143 | + await connection.end() |
| 144 | + } |
| 145 | + }) |
| 146 | + |
| 147 | + describe('matching rows', () => { |
| 148 | + beforeAll(async () => { |
| 149 | + for (const [index, date1] of TAG_VALUES.entries()) { |
| 150 | + const id = `row-${index}` |
| 151 | + await connection` |
| 152 | + INSERT INTO document (id, knowledge_base_id, filename, file_url, mime_type, file_size, date1) |
| 153 | + VALUES (${id}, 'kb', 'file.txt', 'file-url', 'text/plain', 1, ${date1}::timestamp)` |
| 154 | + await connection` |
| 155 | + INSERT INTO embedding (id, knowledge_base_id, document_id, chunk_index, chunk_hash, content, |
| 156 | + content_length, token_count, start_offset, end_offset, date1) |
| 157 | + VALUES (${id}, 'kb', ${id}, 0, 'hash', 'chunk', 5, 1, 0, 5, ${date1}::timestamp)` |
| 158 | + } |
| 159 | + }) |
| 160 | + |
| 161 | + it.each(TIME_ZONES)( |
| 162 | + 'selects the rows of the calendar-day comparison for every operator under %s', |
| 163 | + async (timeZone) => { |
| 164 | + await connection.unsafe(`SET TIME ZONE '${timeZone}'`) |
| 165 | + try { |
| 166 | + for (const filter of FILTERS) { |
| 167 | + const expectedDocuments = await matchingIds( |
| 168 | + document, |
| 169 | + calendarDayCondition(sql`${document.date1}`, filter) |
| 170 | + ) |
| 171 | + const expectedChunks = await matchingIds( |
| 172 | + embedding, |
| 173 | + calendarDayCondition(sql`${embedding.date1}`, filter) |
| 174 | + ) |
| 175 | + expect(await matchingIds(document, documentCondition(filter)), filter.operator).toEqual( |
| 176 | + expectedDocuments |
| 177 | + ) |
| 178 | + expect(await matchingIds(embedding, searchCondition(filter)), filter.operator).toEqual( |
| 179 | + expectedChunks |
| 180 | + ) |
| 181 | + } |
| 182 | + } finally { |
| 183 | + await connection`SET TIME ZONE 'UTC'` |
| 184 | + } |
| 185 | + } |
| 186 | + ) |
| 187 | + |
| 188 | + it('narrows search to the day for an unknown operator or an unbounded between', async () => { |
| 189 | + const day = await matchingIds( |
| 190 | + embedding, |
| 191 | + calendarDayCondition(sql`${embedding.date1}`, { operator: 'eq', value: DAY }) |
| 192 | + ) |
| 193 | + expect(day).toEqual(['row-1', 'row-2', 'row-3']) |
| 194 | + expect(await matchingIds(embedding, searchCondition({ operator: 'on', value: DAY }))).toEqual( |
| 195 | + day |
| 196 | + ) |
| 197 | + expect( |
| 198 | + await matchingIds(embedding, searchCondition({ operator: 'between', value: DAY })) |
| 199 | + ).toEqual(day) |
| 200 | + }) |
| 201 | + }) |
| 202 | + |
| 203 | + describe('query plans', () => { |
| 204 | + beforeAll(async () => { |
| 205 | + await connection`TRUNCATE document, embedding` |
| 206 | + await connection` |
| 207 | + INSERT INTO document (id, knowledge_base_id, filename, file_url, mime_type, file_size, date1) |
| 208 | + SELECT 'doc-' || n, 'kb', 'file.txt', 'file-url', 'text/plain', 1, |
| 209 | + '2020-01-01'::timestamp + n * interval '1 hour' |
| 210 | + FROM generate_series(1, 50000) AS n` |
| 211 | + await connection` |
| 212 | + INSERT INTO embedding (id, knowledge_base_id, document_id, chunk_index, chunk_hash, content, |
| 213 | + content_length, token_count, start_offset, end_offset, date1) |
| 214 | + SELECT 'chunk-' || n, 'kb', 'doc-' || n, 0, 'hash', 'chunk', 5, 1, 0, 5, |
| 215 | + '2020-01-01'::timestamp + n * interval '1 hour' |
| 216 | + FROM generate_series(1, 50000) AS n` |
| 217 | + await copyIndex('doc_date1_idx') |
| 218 | + await copyIndex('emb_date1_idx') |
| 219 | + await connection`ANALYZE document` |
| 220 | + await connection`ANALYZE embedding` |
| 221 | + }) |
| 222 | + |
| 223 | + it.each([ |
| 224 | + { operator: 'eq', value: '2022-06-01' }, |
| 225 | + { operator: 'between', value: '2022-06-01', valueTo: '2022-06-03' }, |
| 226 | + { operator: 'gte', value: '2025-09-01' }, |
| 227 | + { operator: 'lt', value: '2020-01-03' }, |
| 228 | + ])('serves a selective $operator filter from the date slot index', async (filter) => { |
| 229 | + expect(await plannedIndexes(document, documentCondition(filter))).toContain('doc_date1_idx') |
| 230 | + expect(await plannedIndexes(embedding, searchCondition(filter))).toContain('emb_date1_idx') |
| 231 | + }) |
| 232 | + }) |
| 233 | +}) |
0 commit comments