@sixb/ducklake 0.0.1 → 0.1.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (38) hide show
  1. package/README.md +37 -0
  2. package/dist/ducklake-storage.d.ts +2 -1
  3. package/dist/ducklake-storage.d.ts.map +1 -1
  4. package/dist/index.js +598 -43
  5. package/dist/index.js.map +15 -13
  6. package/dist/internal/dataset-row-commit.d.ts +6 -0
  7. package/dist/internal/dataset-row-commit.d.ts.map +1 -1
  8. package/dist/internal/dataset-row-merge.d.ts +16 -0
  9. package/dist/internal/dataset-row-merge.d.ts.map +1 -0
  10. package/dist/internal/ducklake-catalog-schema.d.ts +4 -1
  11. package/dist/internal/ducklake-catalog-schema.d.ts.map +1 -1
  12. package/dist/internal/ducklake-current-dataset-definitions.d.ts.map +1 -1
  13. package/dist/internal/ducklake-dataset-catalog.d.ts +1 -0
  14. package/dist/internal/ducklake-dataset-catalog.d.ts.map +1 -1
  15. package/dist/internal/ducklake-merge-session.d.ts +21 -0
  16. package/dist/internal/ducklake-merge-session.d.ts.map +1 -0
  17. package/dist/internal/ducklake-snapshot-reader.d.ts +1 -0
  18. package/dist/internal/ducklake-snapshot-reader.d.ts.map +1 -1
  19. package/dist/internal/ducklake-write-coordinator.d.ts +5 -1
  20. package/dist/internal/ducklake-write-coordinator.d.ts.map +1 -1
  21. package/dist/internal/ducklake-write-session.d.ts.map +1 -1
  22. package/dist/internal/schema.d.ts +2 -0
  23. package/dist/internal/schema.d.ts.map +1 -1
  24. package/dist/internal/versions.d.ts +1 -0
  25. package/dist/internal/versions.d.ts.map +1 -1
  26. package/package.json +7 -4
  27. package/src/ducklake-storage.ts +6 -0
  28. package/src/internal/dataset-row-commit.ts +95 -2
  29. package/src/internal/dataset-row-merge.ts +268 -0
  30. package/src/internal/ducklake-catalog-schema.ts +90 -1
  31. package/src/internal/ducklake-current-dataset-definitions.ts +14 -5
  32. package/src/internal/ducklake-dataset-catalog.ts +43 -9
  33. package/src/internal/ducklake-merge-session.ts +187 -0
  34. package/src/internal/ducklake-snapshot-reader.ts +9 -2
  35. package/src/internal/ducklake-write-coordinator.ts +160 -20
  36. package/src/internal/ducklake-write-session.ts +28 -5
  37. package/src/internal/schema.ts +7 -0
  38. package/src/internal/versions.ts +17 -1
@@ -0,0 +1,268 @@
1
+ import type { DatasetDefinition } from "@sixb/core"
2
+ import { getDatasetPrimaryKeyColumns, LakeStorageError } from "@sixb/core/lake-storage"
3
+ import type { DuckLakeStorageOptions } from "../types"
4
+ import { type ApplyDatasetRowsResult, inspectUniqueKeyedRelation } from "./dataset-row-commit"
5
+ import { getBigIntLike } from "./duckdb-row"
6
+ import type { DuckDbQueryRuntime } from "./duckdb-runtime"
7
+ import { encodeDatasetTableName } from "./names"
8
+ import { datasetSchemaColumnNamesSql } from "./schema"
9
+ import { qualifiedTableName, quoteIdentifier, quoteSqlString } from "./sql"
10
+
11
+ export interface ApplyDatasetMergeFromRelationInput {
12
+ readonly options: DuckLakeStorageOptions
13
+ readonly runtime: DuckDbQueryRuntime
14
+ readonly dataset: DatasetDefinition
15
+ readonly stagingTableName: string
16
+ readonly sequenceColumnName: string
17
+ readonly kindColumnName: string
18
+ readonly previousRowCount?: number
19
+ readonly validatedPrimaryKeyColumns?: readonly string[]
20
+ }
21
+
22
+ interface MergeEffectCounts {
23
+ readonly inserts: number
24
+ readonly updates: number
25
+ readonly deletes: number
26
+ }
27
+
28
+ export async function applyDatasetMergeFromRelation(
29
+ input: ApplyDatasetMergeFromRelationInput
30
+ ): Promise<ApplyDatasetRowsResult> {
31
+ const primaryKeyColumns = getDatasetPrimaryKeyColumns(input.dataset)
32
+ if (primaryKeyColumns === null) {
33
+ throw new LakeStorageError(
34
+ `[SixbDuckLake] Dataset '${input.dataset.id}' must define a primaryKey before it can be merged.`
35
+ )
36
+ }
37
+
38
+ const stagingTable = quoteIdentifier(input.stagingTableName)
39
+ const targetTable = qualifiedTableName(input.options, encodeDatasetTableName(input.dataset.id))
40
+ const sourceRowCount = await countRows(input.runtime, stagingTable)
41
+ const previousRowCount = await currentBaselineRowCount(input, targetTable, primaryKeyColumns)
42
+
43
+ return withMaterializedFinalChanges(input, primaryKeyColumns, async (finalChanges) => {
44
+ const effects = await countMergeEffects(input, targetTable, finalChanges, primaryKeyColumns)
45
+ const effectiveChangeCount = effects.inserts + effects.updates + effects.deletes
46
+
47
+ if (effectiveChangeCount === 0) {
48
+ return {
49
+ dataChangeExpected: false,
50
+ sourceRowCount,
51
+ resultingRowCount: { kind: "exact", value: previousRowCount },
52
+ }
53
+ }
54
+
55
+ const keyMatchSql = primaryKeyMatchSql(primaryKeyColumns, "target", "source")
56
+ const rowDifferenceSql = datasetRowDifferenceSql(input.dataset, "target", "source")
57
+ const kindColumn = `source.${quoteIdentifier(input.kindColumnName)}`
58
+
59
+ await input.runtime.run(`
60
+ DELETE FROM ${targetTable} AS target
61
+ USING ${finalChanges} AS source
62
+ WHERE ${keyMatchSql}
63
+ AND (
64
+ ${kindColumn} = ${quoteSqlString("delete")}
65
+ OR (
66
+ ${kindColumn} = ${quoteSqlString("upsert")}
67
+ AND (${rowDifferenceSql})
68
+ )
69
+ )
70
+ `)
71
+
72
+ const columnsSql = datasetSchemaColumnNamesSql(input.dataset.schema)
73
+ const selectedColumnsSql = input.dataset.schema.columns
74
+ .map((column) => `source.${quoteIdentifier(column.name)}`)
75
+ .join(", ")
76
+ const firstSequenceColumnName = firstSequenceName(input.sequenceColumnName)
77
+ await input.runtime.run(`
78
+ INSERT INTO ${targetTable} (${columnsSql})
79
+ SELECT ${selectedColumnsSql}
80
+ FROM ${finalChanges} AS source
81
+ WHERE ${kindColumn} = ${quoteSqlString("upsert")}
82
+ AND NOT EXISTS (
83
+ SELECT 1
84
+ FROM ${targetTable} AS target
85
+ WHERE ${keyMatchSql}
86
+ )
87
+ ORDER BY source.${quoteIdentifier(firstSequenceColumnName)}
88
+ `)
89
+
90
+ return {
91
+ dataChangeExpected: true,
92
+ sourceRowCount,
93
+ resultingRowCount: {
94
+ kind: "exact",
95
+ value: previousRowCount + effects.inserts - effects.deletes,
96
+ },
97
+ }
98
+ })
99
+ }
100
+
101
+ async function currentBaselineRowCount(
102
+ input: ApplyDatasetMergeFromRelationInput,
103
+ targetTable: string,
104
+ primaryKeyColumns: readonly string[]
105
+ ): Promise<number> {
106
+ // Sixb writes mark keyed versions only after the write path has enforced uniqueness. Older or
107
+ // external versions carry no marker and keep the full defensive audit.
108
+ if (sameColumns(input.validatedPrimaryKeyColumns, primaryKeyColumns)) {
109
+ return input.previousRowCount ?? countRows(input.runtime, targetTable)
110
+ }
111
+
112
+ return inspectUniqueKeyedRelation({
113
+ runtime: input.runtime,
114
+ dataset: input.dataset,
115
+ relationSql: targetTable,
116
+ context: "current baseline",
117
+ })
118
+ }
119
+
120
+ async function withMaterializedFinalChanges<T>(
121
+ input: ApplyDatasetMergeFromRelationInput,
122
+ primaryKeyColumns: readonly string[],
123
+ run: (finalChangesTable: string) => Promise<T>
124
+ ): Promise<T> {
125
+ const table = quoteIdentifier(`${input.stagingTableName}_final`)
126
+ await input.runtime.run(`
127
+ CREATE TEMP TABLE ${table} AS
128
+ ${finalChangesSelectSql(input, primaryKeyColumns)}
129
+ `)
130
+
131
+ let outcome:
132
+ | { readonly kind: "success"; readonly value: T }
133
+ | { readonly kind: "error"; readonly error: unknown }
134
+ try {
135
+ outcome = { kind: "success", value: await run(table) }
136
+ } catch (error) {
137
+ outcome = { kind: "error", error }
138
+ }
139
+
140
+ try {
141
+ await input.runtime.run(`DROP TABLE IF EXISTS ${table}`)
142
+ } catch (cleanupError) {
143
+ if (outcome.kind === "success") {
144
+ throw cleanupError
145
+ }
146
+ }
147
+
148
+ if (outcome.kind === "error") {
149
+ throw outcome.error
150
+ }
151
+
152
+ return outcome.value
153
+ }
154
+
155
+ function finalChangesSelectSql(
156
+ input: ApplyDatasetMergeFromRelationInput,
157
+ primaryKeyColumns: readonly string[]
158
+ ): string {
159
+ const sequenceColumn = quoteIdentifier(input.sequenceColumnName)
160
+ const keyColumnsSql = primaryKeyColumns.map((column) => quoteIdentifier(column)).join(", ")
161
+ const rankColumn = quoteIdentifier(rankName(input.sequenceColumnName))
162
+ const firstSequenceColumn = quoteIdentifier(firstSequenceName(input.sequenceColumnName))
163
+
164
+ return `
165
+ SELECT * EXCLUDE (${rankColumn})
166
+ FROM (
167
+ SELECT
168
+ *,
169
+ min(${sequenceColumn}) OVER (PARTITION BY ${keyColumnsSql}) AS ${firstSequenceColumn},
170
+ row_number() OVER (
171
+ PARTITION BY ${keyColumnsSql}
172
+ ORDER BY ${sequenceColumn} DESC
173
+ ) AS ${rankColumn}
174
+ FROM ${quoteIdentifier(input.stagingTableName)}
175
+ ) ranked_changes
176
+ WHERE ${rankColumn} = 1
177
+ `
178
+ }
179
+
180
+ async function countMergeEffects(
181
+ input: ApplyDatasetMergeFromRelationInput,
182
+ targetTable: string,
183
+ finalChanges: string,
184
+ primaryKeyColumns: readonly string[]
185
+ ): Promise<MergeEffectCounts> {
186
+ const matchColumn = primaryKeyColumns[0]
187
+ if (matchColumn === undefined) {
188
+ throw new LakeStorageError(
189
+ `[SixbDuckLake] Dataset '${input.dataset.id}' primaryKey must contain at least one column.`
190
+ )
191
+ }
192
+
193
+ const keyMatchSql = primaryKeyMatchSql(primaryKeyColumns, "target", "source")
194
+ const targetMatch = `target.${quoteIdentifier(matchColumn)} IS NOT NULL`
195
+ const kindColumn = `source.${quoteIdentifier(input.kindColumnName)}`
196
+ const rowDifferenceSql = datasetRowDifferenceSql(input.dataset, "target", "source")
197
+ const [row] = await input.runtime.query(`
198
+ SELECT
199
+ count(*) FILTER (
200
+ WHERE ${kindColumn} = ${quoteSqlString("upsert")} AND NOT (${targetMatch})
201
+ ) AS insert_count,
202
+ count(*) FILTER (
203
+ WHERE ${kindColumn} = ${quoteSqlString("upsert")}
204
+ AND ${targetMatch}
205
+ AND (${rowDifferenceSql})
206
+ ) AS update_count,
207
+ count(*) FILTER (
208
+ WHERE ${kindColumn} = ${quoteSqlString("delete")} AND ${targetMatch}
209
+ ) AS delete_count
210
+ FROM ${finalChanges} AS source
211
+ LEFT JOIN ${targetTable} AS target ON ${keyMatchSql}
212
+ `)
213
+
214
+ return {
215
+ inserts: row === undefined ? 0 : Number(getBigIntLike(row, "insert_count")),
216
+ updates: row === undefined ? 0 : Number(getBigIntLike(row, "update_count")),
217
+ deletes: row === undefined ? 0 : Number(getBigIntLike(row, "delete_count")),
218
+ }
219
+ }
220
+
221
+ function primaryKeyMatchSql(
222
+ primaryKeyColumns: readonly string[],
223
+ leftAlias: string,
224
+ rightAlias: string
225
+ ): string {
226
+ return primaryKeyColumns
227
+ .map(
228
+ (column) =>
229
+ `${leftAlias}.${quoteIdentifier(column)} = ${rightAlias}.${quoteIdentifier(column)}`
230
+ )
231
+ .join(" AND ")
232
+ }
233
+
234
+ function datasetRowDifferenceSql(
235
+ dataset: DatasetDefinition,
236
+ leftAlias: string,
237
+ rightAlias: string
238
+ ): string {
239
+ return dataset.schema.columns
240
+ .map(
241
+ (column) =>
242
+ `${leftAlias}.${quoteIdentifier(column.name)} IS DISTINCT FROM ${rightAlias}.${quoteIdentifier(
243
+ column.name
244
+ )}`
245
+ )
246
+ .join(" OR ")
247
+ }
248
+
249
+ async function countRows(runtime: DuckDbQueryRuntime, relationSql: string): Promise<number> {
250
+ const [row] = await runtime.query(`SELECT count(*) AS row_count FROM ${relationSql}`)
251
+ return row === undefined ? 0 : Number(getBigIntLike(row, "row_count"))
252
+ }
253
+
254
+ function rankName(sequenceColumnName: string): string {
255
+ return `${sequenceColumnName}_rank`
256
+ }
257
+
258
+ function firstSequenceName(sequenceColumnName: string): string {
259
+ return `${sequenceColumnName}_first`
260
+ }
261
+
262
+ function sameColumns(left: readonly string[] | undefined, right: readonly string[]): boolean {
263
+ return (
264
+ left !== undefined &&
265
+ left.length === right.length &&
266
+ left.every((value, index) => value === right[index])
267
+ )
268
+ }
@@ -1,4 +1,4 @@
1
- import type { DatasetColumnDefinition, DatasetSchema } from "@sixb/core"
1
+ import type { DatasetColumnDefinition, DatasetPrimaryKey, DatasetSchema } from "@sixb/core"
2
2
  import { LakeStorageError } from "@sixb/core/lake-storage"
3
3
  import { duckDbTypeToDatasetColumnType } from "./schema"
4
4
 
@@ -10,8 +10,12 @@ export interface DuckLakeCatalogColumn {
10
10
  readonly columnType: string
11
11
  readonly nullsAllowed: boolean
12
12
  readonly parentColumnId?: bigint
13
+ readonly comment?: string
13
14
  }
14
15
 
16
+ const SIXB_PRIMARY_KEY_COMMENT_NAMESPACE = "sixb:primary-key:"
17
+ const SIXB_PRIMARY_KEY_COMMENT_PREFIX = `${SIXB_PRIMARY_KEY_COMMENT_NAMESPACE}v1:`
18
+
15
19
  const FILE_REF_STRUCT_CHILDREN = [
16
20
  { name: "blobId", type: "string" },
17
21
  { name: "digest", type: "string" },
@@ -46,6 +50,91 @@ export function duckLakeCatalogColumnsToDatasetSchema(
46
50
  }
47
51
  }
48
52
 
53
+ export function duckLakePrimaryKeyColumnComment(ordinal: number): string {
54
+ return `${SIXB_PRIMARY_KEY_COMMENT_PREFIX}${ordinal}`
55
+ }
56
+
57
+ export function duckLakeCatalogColumnsToDatasetPrimaryKey(
58
+ tableName: string,
59
+ rows: readonly DuckLakeCatalogColumn[]
60
+ ): DatasetPrimaryKey | undefined {
61
+ const keyedColumns: { readonly column: DuckLakeCatalogColumn; readonly ordinal: number }[] = []
62
+ const seenOrdinals = new Set<number>()
63
+
64
+ for (const column of rows) {
65
+ const comment = column.comment
66
+ if (comment === undefined || !comment.startsWith(SIXB_PRIMARY_KEY_COMMENT_NAMESPACE)) {
67
+ continue
68
+ }
69
+
70
+ if (
71
+ column.parentColumnId !== undefined ||
72
+ !comment.startsWith(SIXB_PRIMARY_KEY_COMMENT_PREFIX)
73
+ ) {
74
+ throwInvalidPrimaryKeyMetadata(tableName, column, comment)
75
+ }
76
+
77
+ const ordinalText = comment.slice(SIXB_PRIMARY_KEY_COMMENT_PREFIX.length)
78
+ if (!/^(0|[1-9]\d*)$/.test(ordinalText)) {
79
+ throwInvalidPrimaryKeyMetadata(tableName, column, comment)
80
+ }
81
+
82
+ const ordinal = Number(ordinalText)
83
+ if (!Number.isSafeInteger(ordinal)) {
84
+ throwInvalidPrimaryKeyMetadata(tableName, column, comment)
85
+ }
86
+ if (seenOrdinals.has(ordinal)) {
87
+ throw new LakeStorageError(
88
+ `[SixbDuckLake] Dataset table '${tableName}' has duplicate primary-key ordinal '${ordinal}'.`
89
+ )
90
+ }
91
+ seenOrdinals.add(ordinal)
92
+
93
+ if (duckDbTypeToDatasetColumnType(column.columnType) !== "string") {
94
+ throw new LakeStorageError(
95
+ `[SixbDuckLake] Dataset table '${tableName}' primary-key column '${column.columnName}' must map to type 'string'.`
96
+ )
97
+ }
98
+ if (column.nullsAllowed) {
99
+ throw new LakeStorageError(
100
+ `[SixbDuckLake] Dataset table '${tableName}' primary-key column '${column.columnName}' must not be nullable.`
101
+ )
102
+ }
103
+
104
+ keyedColumns.push({ column, ordinal })
105
+ }
106
+
107
+ if (keyedColumns.length === 0) {
108
+ return undefined
109
+ }
110
+
111
+ keyedColumns.sort((left, right) => left.ordinal - right.ordinal)
112
+ for (const [expectedOrdinal, keyedColumn] of keyedColumns.entries()) {
113
+ if (keyedColumn.ordinal !== expectedOrdinal) {
114
+ throw new LakeStorageError(
115
+ `[SixbDuckLake] Dataset table '${tableName}' primary-key ordinals must be contiguous from 0.`
116
+ )
117
+ }
118
+ }
119
+
120
+ const columnNames = keyedColumns.map(({ column }) => column.columnName)
121
+ if (columnNames.length === 1) {
122
+ return columnNames[0]
123
+ }
124
+
125
+ return columnNames as [string, string, ...string[]]
126
+ }
127
+
128
+ function throwInvalidPrimaryKeyMetadata(
129
+ tableName: string,
130
+ column: DuckLakeCatalogColumn,
131
+ comment: string
132
+ ): never {
133
+ throw new LakeStorageError(
134
+ `[SixbDuckLake] Dataset table '${tableName}' has malformed primary-key metadata '${comment}' on column '${column.columnName}'.`
135
+ )
136
+ }
137
+
49
138
  function duckLakeCatalogColumnToDefinition(
50
139
  tableName: string,
51
140
  column: DuckLakeCatalogColumn,
@@ -11,6 +11,7 @@ import {
11
11
  import type { DuckDbQueryRuntime } from "./duckdb-runtime"
12
12
  import {
13
13
  type DuckLakeCatalogColumn,
14
+ duckLakeCatalogColumnsToDatasetPrimaryKey,
14
15
  duckLakeCatalogColumnsToDatasetSchema,
15
16
  } from "./ducklake-catalog-schema"
16
17
  import { decodeDatasetTableName, encodeDatasetTableName } from "./names"
@@ -56,14 +57,14 @@ export async function readCurrentDatasetDefinitions(
56
57
  const definitions = new Map<string, DatasetDefinition>()
57
58
 
58
59
  for (const table of tables) {
60
+ const columns = columnsByTableName.get(table.tableName) ?? []
61
+ const primaryKey = duckLakeCatalogColumnsToDatasetPrimaryKey(table.tableName, columns)
59
62
  const partitionBy = partitionByTableName.get(table.tableName) ?? []
60
63
  definitions.set(table.datasetId, {
61
64
  kind: "dataset",
62
65
  id: table.datasetId,
63
- schema: duckLakeCatalogColumnsToDatasetSchema(
64
- table.tableName,
65
- columnsByTableName.get(table.tableName) ?? []
66
- ),
66
+ schema: duckLakeCatalogColumnsToDatasetSchema(table.tableName, columns),
67
+ ...(primaryKey !== undefined ? { primaryKey } : {}),
67
68
  ...(partitionBy.length > 0 ? { partitionBy } : {}),
68
69
  ...(table.comment !== undefined ? { description: table.comment } : {}),
69
70
  })
@@ -121,6 +122,7 @@ async function readCurrentDatasetColumns(
121
122
 
122
123
  const ducklakeTable = duckLakeMetadataTableName(input.options, "ducklake_table")
123
124
  const ducklakeColumn = duckLakeMetadataTableName(input.options, "ducklake_column")
125
+ const ducklakeColumnTag = duckLakeMetadataTableName(input.options, "ducklake_column_tag")
124
126
  const rows = await input.runtime.query(`
125
127
  SELECT
126
128
  table_meta.table_name,
@@ -129,11 +131,17 @@ async function readCurrentDatasetColumns(
129
131
  column_meta.column_name,
130
132
  column_meta.column_type,
131
133
  CAST(column_meta.nulls_allowed AS BOOLEAN) AS nulls_allowed,
132
- column_meta.parent_column
134
+ column_meta.parent_column,
135
+ column_tag."value" AS comment
133
136
  FROM ${ducklakeTable} table_meta
134
137
  JOIN ${ducklakeColumn} column_meta
135
138
  ON column_meta.table_id = table_meta.table_id
136
139
  AND column_meta.end_snapshot IS NULL
140
+ LEFT JOIN ${ducklakeColumnTag} column_tag
141
+ ON column_tag.table_id = column_meta.table_id
142
+ AND column_tag.column_id = column_meta.column_id
143
+ AND column_tag."key" = 'comment'
144
+ AND column_tag.end_snapshot IS NULL
137
145
  WHERE table_meta.end_snapshot IS NULL
138
146
  AND table_meta.table_name IN (${quoteSqlStringList(tableNames)})
139
147
  ORDER BY table_meta.table_name, column_meta.column_order
@@ -147,6 +155,7 @@ async function readCurrentDatasetColumns(
147
155
  columnType: getString(row, "column_type"),
148
156
  nullsAllowed: getBoolean(row, "nulls_allowed"),
149
157
  parentColumnId: getOptionalBigIntLike(row, "parent_column"),
158
+ comment: getOptionalString(row, "comment"),
150
159
  }))
151
160
  }
152
161
 
@@ -8,6 +8,7 @@ import type { DuckDbQueryRuntime } from "./duckdb-runtime"
8
8
  import {
9
9
  type DuckLakeCatalogColumn,
10
10
  duckLakeCatalogColumnsToDatasetSchema,
11
+ duckLakePrimaryKeyColumnComment,
11
12
  } from "./ducklake-catalog-schema"
12
13
  import type { DuckLakeConnectionManager } from "./ducklake-connection-manager"
13
14
  import { readCurrentDatasetDefinitions } from "./ducklake-current-dataset-definitions"
@@ -59,19 +60,32 @@ export class DuckLakeDatasetCatalog {
59
60
  const qualifiedName = qualifiedTableName(this.options, tableName)
60
61
 
61
62
  // Creating a dataset can be several catalog DDL statements. Keep them in
62
- // one runtime slot so writes cannot observe a half-created table.
63
+ // one transaction and runtime slot so nothing can observe a partial definition.
63
64
  await this.connections.withExclusiveAttached(async (runtime) => {
64
- await runtime.run(
65
- `CREATE TABLE ${qualifiedName} (${datasetSchemaToDuckDbColumnsSql(definition.schema)})`
66
- )
67
-
68
- if (definition.description !== undefined) {
65
+ await runtime.run("BEGIN TRANSACTION")
66
+ let committed = false
67
+ try {
69
68
  await runtime.run(
70
- `COMMENT ON TABLE ${qualifiedName} IS ${quoteSqlString(definition.description)}`
69
+ `CREATE TABLE ${qualifiedName} (${datasetSchemaToDuckDbColumnsSql(definition.schema)})`
71
70
  )
72
- }
73
71
 
74
- await this.applyPartitionBy(runtime, qualifiedName, definition)
72
+ if (definition.description !== undefined) {
73
+ await runtime.run(
74
+ `COMMENT ON TABLE ${qualifiedName} IS ${quoteSqlString(definition.description)}`
75
+ )
76
+ }
77
+
78
+ await this.applyPrimaryKey(runtime, qualifiedName, definition)
79
+ await this.applyPartitionBy(runtime, qualifiedName, definition)
80
+
81
+ await runtime.run("COMMIT")
82
+ committed = true
83
+ } catch (error) {
84
+ if (!committed) {
85
+ await this.rollbackTransaction(runtime)
86
+ }
87
+ throw error
88
+ }
75
89
  })
76
90
  this.connections.markLocalCatalogChanged()
77
91
 
@@ -289,6 +303,26 @@ export class DuckLakeDatasetCatalog {
289
303
  }
290
304
  }
291
305
 
306
+ private async applyPrimaryKey(
307
+ runtime: DuckDbQueryRuntime,
308
+ qualifiedName: string,
309
+ definition: DatasetDefinition
310
+ ): Promise<void> {
311
+ if (definition.primaryKey === undefined) {
312
+ return
313
+ }
314
+
315
+ const columnNames =
316
+ typeof definition.primaryKey === "string" ? [definition.primaryKey] : definition.primaryKey
317
+ for (const [ordinal, columnName] of columnNames.entries()) {
318
+ await runtime.run(
319
+ `COMMENT ON COLUMN ${qualifiedName}.${quoteIdentifier(columnName)} IS ${quoteSqlString(
320
+ duckLakePrimaryKeyColumnComment(ordinal)
321
+ )}`
322
+ )
323
+ }
324
+ }
325
+
292
326
  private async setSchemaEvolutionCommitMetadata(
293
327
  runtime: DuckDbQueryRuntime,
294
328
  datasetId: string,
@@ -0,0 +1,187 @@
1
+ import { randomUUID } from "node:crypto"
2
+ import type { DatasetRow, MergeChange } from "@sixb/core"
3
+ import type {
4
+ BeginDatasetMergeInput,
5
+ CommitDatasetMergeInput,
6
+ DatasetMergeCommitResult,
7
+ LakeMergeSession,
8
+ } from "@sixb/core/lake-storage"
9
+ import {
10
+ cloneDatasetMergeChange,
11
+ getDatasetMergeChangeValidationError,
12
+ LakeStorageError,
13
+ } from "@sixb/core/lake-storage"
14
+ import type { DuckDbRuntime } from "./duckdb-runtime"
15
+ import { appendDatasetRow } from "./row-appender"
16
+ import { datasetSchemaToDuckDbNullableColumnsSql } from "./schema"
17
+ import { quoteIdentifier } from "./sql"
18
+
19
+ export interface DuckLakeCommitMergeInput {
20
+ readonly merge: BeginDatasetMergeInput
21
+ readonly baseVersionId: string | null
22
+ readonly commit?: CommitDatasetMergeInput
23
+ readonly stagingTableName: string
24
+ readonly sequenceColumnName: string
25
+ readonly kindColumnName: string
26
+ readonly changesWritten: number
27
+ }
28
+
29
+ type DuckLakeCommitMerge = (input: DuckLakeCommitMergeInput) => Promise<DatasetMergeCommitResult>
30
+
31
+ const MERGE_CHANGE_BATCH_SIZE = 1000
32
+
33
+ interface StagedMergeChange {
34
+ readonly sequence: bigint
35
+ readonly change: MergeChange<DatasetRow, DatasetRow>
36
+ }
37
+
38
+ interface CreateDuckLakeMergeSessionInput {
39
+ readonly commitMerge: DuckLakeCommitMerge
40
+ readonly runtime: DuckDbRuntime
41
+ readonly merge: BeginDatasetMergeInput
42
+ readonly baseVersionId: string | null
43
+ }
44
+
45
+ export async function createDuckLakeMergeSession(
46
+ input: CreateDuckLakeMergeSessionInput
47
+ ): Promise<LakeMergeSession> {
48
+ const suffix = randomUUID().replaceAll("-", "")
49
+ const stagingTableName = `sixb_merge_${suffix}`
50
+ const sequenceColumnName = `sixb_merge_sequence_${suffix}`
51
+ const kindColumnName = `sixb_merge_kind_${suffix}`
52
+
53
+ await input.runtime.run(`
54
+ CREATE TEMP TABLE ${quoteIdentifier(stagingTableName)} (
55
+ ${quoteIdentifier(sequenceColumnName)} UBIGINT NOT NULL,
56
+ ${quoteIdentifier(kindColumnName)} VARCHAR NOT NULL,
57
+ ${datasetSchemaToDuckDbNullableColumnsSql(input.merge.dataset.schema)}
58
+ )
59
+ `)
60
+
61
+ return new DuckLakeMergeSession(
62
+ input.commitMerge,
63
+ input.runtime,
64
+ input.merge,
65
+ input.baseVersionId,
66
+ stagingTableName,
67
+ sequenceColumnName,
68
+ kindColumnName
69
+ )
70
+ }
71
+
72
+ class DuckLakeMergeSession implements LakeMergeSession {
73
+ private closed = false
74
+ private cleanedUp = false
75
+ private nextSequence = 0n
76
+ private changesWritten = 0
77
+
78
+ constructor(
79
+ private readonly commitMerge: DuckLakeCommitMerge,
80
+ private readonly runtime: DuckDbRuntime,
81
+ private readonly input: BeginDatasetMergeInput,
82
+ private readonly baseVersionId: string | null,
83
+ private readonly stagingTableName: string,
84
+ private readonly sequenceColumnName: string,
85
+ private readonly kindColumnName: string
86
+ ) {}
87
+
88
+ async writeChanges(
89
+ changes:
90
+ | Iterable<MergeChange<DatasetRow, DatasetRow>>
91
+ | AsyncIterable<MergeChange<DatasetRow, DatasetRow>>
92
+ ): Promise<void> {
93
+ this.assertOpen()
94
+
95
+ let batch: StagedMergeChange[] = []
96
+ for await (const change of changes) {
97
+ const validationError = getDatasetMergeChangeValidationError(change, this.input.dataset)
98
+ if (validationError) {
99
+ throw new LakeStorageError(`[SixbDuckLake] ${validationError}`)
100
+ }
101
+
102
+ batch.push({
103
+ sequence: this.nextSequence,
104
+ change: cloneDatasetMergeChange(change),
105
+ })
106
+ this.nextSequence += 1n
107
+
108
+ if (batch.length >= MERGE_CHANGE_BATCH_SIZE) {
109
+ await this.appendBatchToStagingTable(batch)
110
+ batch = []
111
+ }
112
+ }
113
+
114
+ if (batch.length > 0) {
115
+ await this.appendBatchToStagingTable(batch)
116
+ }
117
+ }
118
+
119
+ private async appendBatchToStagingTable(batch: readonly StagedMergeChange[]): Promise<void> {
120
+ await this.runtime.withAppender(this.stagingTableName, (appender) => {
121
+ for (const staged of batch) {
122
+ appender.appendBigInt(staged.sequence)
123
+ appender.appendVarchar(staged.change.kind)
124
+ appendDatasetRow(
125
+ appender,
126
+ this.input.dataset.schema,
127
+ staged.change.kind === "upsert" ? staged.change.row : staged.change.key
128
+ )
129
+ }
130
+ })
131
+ this.changesWritten += batch.length
132
+ }
133
+
134
+ async commit(input?: CommitDatasetMergeInput): Promise<DatasetMergeCommitResult> {
135
+ this.assertOpen()
136
+ this.closed = true
137
+
138
+ let committed = false
139
+ try {
140
+ const result = await this.commitMerge({
141
+ merge: this.input,
142
+ baseVersionId: this.baseVersionId,
143
+ commit: input,
144
+ stagingTableName: this.stagingTableName,
145
+ sequenceColumnName: this.sequenceColumnName,
146
+ kindColumnName: this.kindColumnName,
147
+ changesWritten: this.changesWritten,
148
+ })
149
+ committed = true
150
+ await this.cleanup()
151
+ return result
152
+ } catch (error) {
153
+ if (!committed) {
154
+ await this.cleanup()
155
+ }
156
+ throw error
157
+ }
158
+ }
159
+
160
+ async abort(): Promise<void> {
161
+ if (this.closed) {
162
+ return
163
+ }
164
+
165
+ this.closed = true
166
+ await this.cleanup()
167
+ }
168
+
169
+ private assertOpen(): void {
170
+ if (this.closed) {
171
+ throw new LakeStorageError("[SixbDuckLake] Merge session is already closed.")
172
+ }
173
+ }
174
+
175
+ private async cleanup(): Promise<void> {
176
+ if (this.cleanedUp) {
177
+ return
178
+ }
179
+
180
+ this.cleanedUp = true
181
+ try {
182
+ await this.runtime.run(`DROP TABLE IF EXISTS ${quoteIdentifier(this.stagingTableName)}`)
183
+ } catch {
184
+ // Preserve the commit/abort outcome. Temp tables disappear when the storage runtime closes.
185
+ }
186
+ }
187
+ }