diff --git a/modules/experimental/src/gpu-primitives/gpu-batch-hash-index.ts b/modules/experimental/src/gpu-primitives/gpu-batch-hash-index.ts new file mode 100644 index 0000000000..becf09a9a8 --- /dev/null +++ b/modules/experimental/src/gpu-primitives/gpu-batch-hash-index.ts @@ -0,0 +1,223 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// Copyright (c) vis.gl contributors + +import {GraphVectorView, type GPUCommandGraph, type GraphDataView} from './gpu-command-graph'; +import { + doGraphDataViewsOverlap, + validateMatchingVectorTopology, + validatePackedUint32View +} from './graph-data-view-utils'; +import { + addGPUHashIndexBuildBatchesToGraph, + GPU_HASH_INDEX_STATISTICS_LENGTH, + type GPUHashIndexBuildBatch, + type GPUHashIndexStats, + type GPUHashIndexView +} from './gpu-hash-index'; + +const MAXIMUM_UINT32 = 0xffffffff; + +/** Properties for rebuilding one hash index from ordered, preserved source batches. */ +export type GPUBatchHashIndexProps = { + /** Prefix for the shared initialization and ordered per-batch graph passes. */ + id?: string; + /** Ordered packed unsigned key chunks. */ + keys: GraphVectorView<'uint32'>; + /** Optional packed values with exactly the same ordered chunk topology as `keys`. */ + values?: GraphVectorView<'uint32'>; + /** Optional packed nonzero/zero validity with exactly the same source topology. */ + validity?: GraphVectorView<'uint32'>; + /** First generated value for each source chunk. Mutually exclusive with `values`. */ + firstValues?: readonly number[]; + /** Caller-owned power-of-two key table shared by all source chunks. */ + tableKeys: GraphDataView<'uint32'>; + /** Caller-owned values aligned with `tableKeys`. */ + tableValues: GraphDataView<'uint32'>; + /** Caller-owned six-row cumulative build-statistics block. */ + statistics: GraphDataView<'uint32'>; + /** Maximum slots examined per valid input row. Defaults to table capacity. */ + maxProbeCount?: number; +}; + +/** CPU-visible storage, bounded-work, and preserved source-topology facts. */ +export type GPUBatchHashIndexStats = GPUHashIndexStats & { + batchCount: number; + inputLength: number; +}; + +/** + * Rebuilds one packed unsigned-key hash index from ordered source chunks. + * + * Chunks retain their original GPU buffers and are processed in source order. Duplicate keys + * retain the globally earliest source row, including duplicates encountered in later chunks. + * Zero-validity rows are silently excluded; valid reserved keys increment the invalid statistic. + */ +export class GPUBatchHashIndex implements GPUHashIndexView { + readonly id: string; + readonly keys: GraphVectorView<'uint32'>; + readonly values?: GraphVectorView<'uint32'>; + readonly validity?: GraphVectorView<'uint32'>; + readonly firstValues: readonly number[]; + readonly tableKeys: GraphDataView<'uint32'>; + readonly tableValues: GraphDataView<'uint32'>; + readonly statistics: GraphDataView<'uint32'>; + readonly maxProbeCount: number; + readonly stats: GPUBatchHashIndexStats; + readonly updatePolicy = 'rebuild' as const; + + constructor(props: GPUBatchHashIndexProps) { + this.id = props.id ?? 'gpu-batch-hash-index'; + this.keys = props.keys; + this.values = props.values; + this.validity = props.validity; + this.tableKeys = props.tableKeys; + this.tableValues = props.tableValues; + this.statistics = props.statistics; + this.maxProbeCount = props.maxProbeCount ?? this.tableKeys.length; + + validateSourceVector(this.keys, `${this.id} keys`); + if (this.values) { + validateSourceVector(this.values, `${this.id} values`); + validateMatchingVectorTopology(this.keys, this.values, `${this.id} values`); + } + if (this.validity) { + validateSourceVector(this.validity, `${this.id} validity`); + validateMatchingVectorTopology(this.keys, this.validity, `${this.id} validity`); + } + if (this.values && props.firstValues !== undefined) { + throw new Error(`${this.id} values and firstValues are mutually exclusive`); + } + + let firstSourceRow = 0; + const defaultFirstValues = this.keys.data.map(chunk => { + const firstValue = firstSourceRow; + firstSourceRow += chunk.length; + return firstValue; + }); + this.firstValues = Object.freeze( + props.firstValues === undefined ? defaultFirstValues : [...props.firstValues] + ); + if (this.firstValues.length !== this.keys.data.length) { + throw new Error(`${this.id} firstValues must contain one value per source chunk`); + } + for (const [chunkIndex, chunk] of this.keys.data.entries()) { + const firstValue = this.firstValues[chunkIndex]; + if ( + !Number.isSafeInteger(firstValue) || + firstValue < 0 || + firstValue > MAXIMUM_UINT32 || + (chunk.length > 0 && firstValue + chunk.length - 1 > MAXIMUM_UINT32) + ) { + throw new Error(`${this.id} generated values must fit in uint32`); + } + } + + for (const [view, name] of [ + [this.tableKeys, 'tableKeys'], + [this.tableValues, 'tableValues'], + [this.statistics, 'statistics'] + ] as const) { + validatePackedUint32View(view, `${this.id} ${name}`); + } + if ( + !Number.isSafeInteger(this.tableKeys.length) || + this.tableKeys.length < 1 || + !Number.isInteger(Math.log2(this.tableKeys.length)) + ) { + throw new Error(`${this.id} table capacity must be a positive power of two`); + } + if (this.tableKeys.length > MAXIMUM_UINT32) { + throw new Error(`${this.id} table capacity must fit in uint32`); + } + if (this.tableValues.length !== this.tableKeys.length) { + throw new Error(`${this.id} table key and value capacities must match`); + } + if (this.statistics.length < GPU_HASH_INDEX_STATISTICS_LENGTH) { + throw new Error(`${this.id} statistics must contain six uint32 rows`); + } + if ( + !Number.isSafeInteger(this.maxProbeCount) || + this.maxProbeCount < 1 || + this.maxProbeCount > this.tableKeys.length + ) { + throw new Error(`${this.id} maxProbeCount must be an integer from one through capacity`); + } + if (this.keys.length * this.maxProbeCount > MAXIMUM_UINT32) { + throw new Error(`${this.id} aggregate probe count must fit in uint32 statistics`); + } + validateDisjointViews(this); + + this.stats = Object.freeze({ + capacity: this.tableKeys.length, + maxProbeCount: this.maxProbeCount, + tableByteLength: this.tableKeys.length * 8, + statisticsByteLength: GPU_HASH_INDEX_STATISTICS_LENGTH * 4, + outputByteLength: this.tableKeys.length * 8 + GPU_HASH_INDEX_STATISTICS_LENGTH * 4, + batchCount: this.keys.data.length, + inputLength: this.keys.length + }); + } + + /** Adds one shared clear and sequential chunk-local build/finalize passes to a graph. */ + addToGraph(graph: GPUCommandGraph): void { + const views = [ + ...this.keys.data, + ...(this.values?.data ?? []), + ...(this.validity?.data ?? []), + this.tableKeys, + this.tableValues, + this.statistics + ]; + if (views.some(view => view.buffer.graph !== graph)) { + throw new Error(`${this.id} views must belong to the target graph`); + } + + const batches: GPUHashIndexBuildBatch[] = this.keys.data.map((keys, chunkIndex) => ({ + keys, + ...(this.values ? {values: this.values.data[chunkIndex]} : {}), + ...(this.validity ? {validity: this.validity.data[chunkIndex]} : {}), + firstValue: this.firstValues[chunkIndex] + })); + addGPUHashIndexBuildBatchesToGraph(graph, this, batches); + } +} + +function validateSourceVector(vector: GraphVectorView<'uint32'>, name: string): void { + if (!(vector instanceof GraphVectorView) || vector.format !== 'uint32') { + throw new Error(`${name} must be a uint32 GraphVectorView`); + } + if ( + !Number.isSafeInteger(vector.length) || + vector.length < 0 || + vector.data.reduce((length, chunk) => length + chunk.length, 0) !== vector.length + ) { + throw new Error(`${name} length must equal its ordered source chunks`); + } + for (const [chunkIndex, chunk] of vector.data.entries()) { + validatePackedUint32View(chunk, `${name} chunk ${chunkIndex}`); + } +} + +function validateDisjointViews(index: GPUBatchHashIndex): void { + const inputs = [ + ...index.keys.data, + ...(index.values?.data ?? []), + ...(index.validity?.data ?? []) + ]; + const outputs = [index.tableKeys, index.tableValues, index.statistics]; + for (const input of inputs) { + for (const output of outputs) { + if (doGraphDataViewsOverlap(input, output)) { + throw new Error(`${index.id} input and output views must not overlap`); + } + } + } + for (let first = 0; first < outputs.length; first++) { + for (let second = first + 1; second < outputs.length; second++) { + if (doGraphDataViewsOverlap(outputs[first], outputs[second])) { + throw new Error(`${index.id} output views must not overlap`); + } + } + } +} diff --git a/modules/experimental/src/gpu-primitives/gpu-hash-index.ts b/modules/experimental/src/gpu-primitives/gpu-hash-index.ts index 5dfed78c8a..c68c2b4be6 100644 --- a/modules/experimental/src/gpu-primitives/gpu-hash-index.ts +++ b/modules/experimental/src/gpu-primitives/gpu-hash-index.ts @@ -35,6 +35,24 @@ export type GPUHashIndexView = { maxProbeCount: number; }; +/** One preserved packed source chunk contributed to a shared hash-index rebuild. @internal */ +export type GPUHashIndexBuildBatch = { + keys: GraphDataView<'uint32'>; + values?: GraphDataView<'uint32'>; + validity?: GraphDataView<'uint32'>; + firstValue: number; +}; + +type GPUHashIndexBuildTarget = GPUHashIndexView & { + id: string; + statistics: GraphDataView<'uint32'>; +}; + +type GPUHashIndexBuildPass = GPUHashIndexBuildBatch & { + id: string; + firstSourceRow: number; +}; + /** Properties for one packed fixed-capacity hash-index rebuild. */ export type GPUHashIndexProps = { id?: string; @@ -161,9 +179,45 @@ export class GPUHashIndex implements GPUHashIndexView { 'uint32', this.tableKeys.length ); + const batch: GPUHashIndexBuildPass = { + id: this.id, + keys: this.keys, + ...(this.values ? {values: this.values} : {}), + firstValue: this.firstValue, + firstSourceRow: 0 + }; addBuildInitializePass(graph, this, sourceRows); - if (this.keys.length > 0) addBuildPass(graph, this, sourceRows); - addBuildFinalizePass(graph, this, sourceRows); + if (this.keys.length > 0) addBuildPass(graph, this, sourceRows, batch); + addBuildFinalizePass(graph, this, sourceRows, batch); + } +} + +/** Rebuilds one shared index from ordered source chunks without concatenating their buffers. @internal */ +export function addGPUHashIndexBuildBatchesToGraph( + graph: GPUCommandGraph, + index: GPUHashIndexBuildTarget, + batches: readonly GPUHashIndexBuildBatch[] +): void { + const sourceRows = createTransientView( + graph, + `${index.id}-source-rows`, + 'uint32', + index.tableKeys.length + ); + addBuildInitializePass(graph, index, sourceRows); + + let firstSourceRow = 0; + for (const [batchIndex, batch] of batches.entries()) { + if (batch.keys.length > 0) { + const pass: GPUHashIndexBuildPass = { + ...batch, + id: `${index.id}-batch-${batchIndex}`, + firstSourceRow + }; + addBuildPass(graph, index, sourceRows, pass); + addBuildFinalizePass(graph, index, sourceRows, pass); + } + firstSourceRow += batch.keys.length; } } @@ -261,7 +315,7 @@ export class GPUHashIndexQuery { function addBuildInitializePass( graph: GPUCommandGraph, - index: GPUHashIndex, + index: GPUHashIndexBuildTarget, sourceRows: GraphDataView<'uint32'> ): void { const layout = getDispatchLayout( @@ -316,27 +370,38 @@ const STATISTICS_OFFSET: u32 = ${getViewElementOffset(index.statistics)}u; function addBuildPass( graph: GPUCommandGraph, - index: GPUHashIndex, - sourceRows: GraphDataView<'uint32'> + index: GPUHashIndexBuildTarget, + sourceRows: GraphDataView<'uint32'>, + batch: GPUHashIndexBuildPass ): void { const layout = getDispatchLayout( - index.keys.length, + batch.keys.length, graph.device.limits.maxComputeWorkgroupsPerDimension ); + const validityBinding = batch.validity + ? '@group(0) @binding(1) var inputValidity: array;' + : ''; + const firstTableBinding = batch.validity ? 2 : 1; + const validityOffset = batch.validity ? getViewElementOffset(batch.validity) : 0; + const rejectInvalid = batch.validity + ? `if (inputValidity[${validityOffset}u + inputIndex] == 0u) { return; }` + : ''; const source = /* wgsl */ ` -const ELEMENT_COUNT: u32 = ${index.keys.length}u; +const ELEMENT_COUNT: u32 = ${batch.keys.length}u; const CAPACITY_MASK: u32 = ${index.tableKeys.length - 1}u; const MAX_PROBES: u32 = ${index.maxProbeCount}u; const DISPATCH_X: u32 = ${layout.x}u; const DISPATCH_Y: u32 = ${layout.y}u; -const KEYS_OFFSET: u32 = ${getViewElementOffset(index.keys)}u; +const KEYS_OFFSET: u32 = ${getViewElementOffset(batch.keys)}u; const TABLE_KEYS_OFFSET: u32 = ${getViewElementOffset(index.tableKeys)}u; const SOURCE_ROWS_OFFSET: u32 = ${getViewElementOffset(sourceRows)}u; const STATISTICS_OFFSET: u32 = ${getViewElementOffset(index.statistics)}u; +const FIRST_SOURCE_ROW: u32 = ${batch.firstSourceRow}u; @group(0) @binding(0) var inputKeys: array; -@group(0) @binding(1) var tableKeys: array>; -@group(0) @binding(2) var sourceRows: array>; -@group(0) @binding(3) var statistics: array>; +${validityBinding} +@group(0) @binding(${firstTableBinding}) var tableKeys: array>; +@group(0) @binding(${firstTableBinding + 1}) var sourceRows: array>; +@group(0) @binding(${firstTableBinding + 2}) var statistics: array>; fn hashKey(key: u32) -> u32 { var value = key; @@ -352,6 +417,7 @@ fn hashKey(key: u32) -> u32 { let workgroupIndex = (workgroupId.z * DISPATCH_Y + workgroupId.y) * DISPATCH_X + workgroupId.x; let inputIndex = workgroupIndex * ${HASH_INDEX_WORKGROUP_SIZE}u + localIndex; if (inputIndex >= ELEMENT_COUNT) { return; } + ${rejectInvalid} let key = inputKeys[KEYS_OFFSET + inputIndex]; if (key == ${GPU_HASH_INDEX_EMPTY_KEY}u) { atomicAdd(&statistics[STATISTICS_OFFSET + 3u], 1u); @@ -383,7 +449,7 @@ fn hashKey(key: u32) -> u32 { break; } if (result.exchanged || result.old_value == key) { - atomicMin(&sourceRows[SOURCE_ROWS_OFFSET + slot], inputIndex); + atomicMin(&sourceRows[SOURCE_ROWS_OFFSET + slot], FIRST_SOURCE_ROW + inputIndex); inserted = result.exchanged; duplicate = !result.exchanged; break; @@ -400,16 +466,20 @@ fn hashKey(key: u32) -> u32 { } }`; addComputationPass(graph, { - id: `${index.id}-build`, + id: `${batch.id}-build`, source, resources: [ - {buffer: index.keys, usage: 'storage-read'}, + {buffer: batch.keys, usage: 'storage-read'}, + ...(batch.validity + ? ([{buffer: batch.validity, usage: 'storage-read'}] as GraphBufferUse[]) + : []), {buffer: index.tableKeys, usage: 'storage-read-write'}, {buffer: sourceRows, usage: 'storage-read-write'}, {buffer: index.statistics, usage: 'storage-read-write'} ], bindings: { - inputKeys: index.keys, + inputKeys: batch.keys, + ...(batch.validity ? {inputValidity: batch.validity} : {}), tableKeys: index.tableKeys, sourceRows, statistics: index.statistics @@ -420,22 +490,25 @@ fn hashKey(key: u32) -> u32 { function addBuildFinalizePass( graph: GPUCommandGraph, - index: GPUHashIndex, - sourceRows: GraphDataView<'uint32'> + index: GPUHashIndexBuildTarget, + sourceRows: GraphDataView<'uint32'>, + batch: GPUHashIndexBuildPass ): void { const layout = getDispatchLayout( index.tableKeys.length, graph.device.limits.maxComputeWorkgroupsPerDimension ); - const hasExplicitValues = Boolean(index.values && index.keys.length > 0); + const hasExplicitValues = Boolean(batch.values && batch.keys.length > 0); const valueBinding = hasExplicitValues ? '@group(0) @binding(3) var inputValues: array;' : ''; const valueExpression = hasExplicitValues - ? `inputValues[${getViewElementOffset(index.values!)}u + sourceRow]` - : `${index.firstValue}u + sourceRow`; + ? `inputValues[${getViewElementOffset(batch.values!)}u + localSourceRow]` + : `${batch.firstValue}u + localSourceRow`; const source = /* wgsl */ ` const CAPACITY: u32 = ${index.tableKeys.length}u; +const FIRST_SOURCE_ROW: u32 = ${batch.firstSourceRow}u; +const SOURCE_ROW_COUNT: u32 = ${batch.keys.length}u; const DISPATCH_X: u32 = ${layout.x}u; const DISPATCH_Y: u32 = ${layout.y}u; const TABLE_KEYS_OFFSET: u32 = ${getViewElementOffset(index.tableKeys)}u; @@ -453,24 +526,26 @@ ${valueBinding} let slot = workgroupIndex * ${HASH_INDEX_WORKGROUP_SIZE}u + localIndex; if (slot >= CAPACITY || tableKeys[TABLE_KEYS_OFFSET + slot] == ${GPU_HASH_INDEX_EMPTY_KEY}u) { return; } let sourceRow = sourceRows[SOURCE_ROWS_OFFSET + slot]; + if (sourceRow < FIRST_SOURCE_ROW || sourceRow - FIRST_SOURCE_ROW >= SOURCE_ROW_COUNT) { return; } + let localSourceRow = sourceRow - FIRST_SOURCE_ROW; tableValues[TABLE_VALUES_OFFSET + slot] = ${valueExpression}; }`; addComputationPass(graph, { - id: `${index.id}-finalize`, + id: `${batch.id}-finalize`, source, resources: [ {buffer: index.tableKeys, usage: 'storage-read'}, {buffer: sourceRows, usage: 'storage-read'}, {buffer: index.tableValues, usage: 'storage-write'}, ...(hasExplicitValues - ? ([{buffer: index.values!, usage: 'storage-read'}] as GraphBufferUse[]) + ? ([{buffer: batch.values!, usage: 'storage-read'}] as GraphBufferUse[]) : []) ], bindings: { tableKeys: index.tableKeys, sourceRows, tableValues: index.tableValues, - ...(hasExplicitValues ? {inputValues: index.values!} : {}) + ...(hasExplicitValues ? {inputValues: batch.values!} : {}) }, dispatchSize: layout }); diff --git a/modules/experimental/src/gpu-primitives/index.ts b/modules/experimental/src/gpu-primitives/index.ts index dbf7a74939..54d5892fc7 100644 --- a/modules/experimental/src/gpu-primitives/index.ts +++ b/modules/experimental/src/gpu-primitives/index.ts @@ -300,6 +300,9 @@ export type { GPUHashIndexView } from './gpu-hash-index'; +export {GPUBatchHashIndex} from './gpu-batch-hash-index'; +export type {GPUBatchHashIndexProps, GPUBatchHashIndexStats} from './gpu-batch-hash-index'; + export {GPUHashJoin} from './gpu-hash-join'; export type {GPUHashJoinProps, GPUHashJoinStats} from './gpu-hash-join'; export {GPUBatchHashJoin} from './gpu-batch-hash-join'; diff --git a/modules/experimental/src/ludf/index.ts b/modules/experimental/src/ludf/index.ts index 9e68584892..68b673d51d 100644 --- a/modules/experimental/src/ludf/index.ts +++ b/modules/experimental/src/ludf/index.ts @@ -41,6 +41,8 @@ export {LuDataFrameHistogramQuery} from './lu-histogram-query'; export type {LuDataFrameHistogramOptions} from './lu-histogram-query'; export {LuDataFrameSortQuery} from './lu-sort-query'; export type {LuDataFrameSortOptions} from './lu-sort-query'; +export {LuDataFrameJoinQuery, LuDataFrameLookupQuery} from './lu-join-query'; +export type {LuDataFrameJoinOptions, LuDataFrameLookupOptions} from './lu-join-query'; export {and, column, literal, LuExpression, not, or, parameter} from './lu-expression'; export type { LuExpressionBinaryOperator, @@ -54,3 +56,4 @@ export {CompiledLuDataFrameGroupedAggregation} from './lu-group-aggregation-comp export {CompiledLuDataFrameAggregation} from './lu-global-aggregation-compiler'; export {CompiledLuDataFrameHistogram} from './lu-histogram-compiler'; export {CompiledLuDataFrameSort} from './lu-sort-compiler'; +export {CompiledLuDataFrameJoin, CompiledLuDataFrameLookup} from './lu-join-compiler'; diff --git a/modules/experimental/src/ludf/lu-data-frame-query.ts b/modules/experimental/src/ludf/lu-data-frame-query.ts index bba98de28d..00ba9aab60 100644 --- a/modules/experimental/src/ludf/lu-data-frame-query.ts +++ b/modules/experimental/src/ludf/lu-data-frame-query.ts @@ -17,6 +17,12 @@ import { type LuDataFrameGroupByOptions } from './lu-group-by-query'; import {LuDataFrameHistogramQuery, type LuDataFrameHistogramOptions} from './lu-histogram-query'; +import { + LuDataFrameJoinQuery, + LuDataFrameLookupQuery, + type LuDataFrameJoinOptions, + type LuDataFrameLookupOptions +} from './lu-join-query'; import { compileLuDataFrameQuery, type CompiledLuDataFrameQuery, @@ -196,6 +202,30 @@ export class LuDataFrameQuery< return new LuDataFrameSortQuery(this, column, options, limit, 'descending'); } + /** Plans a stable, unique-right inner join without allocating, repacking, or retaining rows. */ + innerJoin< + Right extends GPUTypeMap, + LeftKey extends LuDataFrameColumnNamesOfFormat, + RightKey extends LuDataFrameColumnNamesOfFormat + >( + right: LuDataFrame, + options: LuDataFrameJoinOptions + ): LuDataFrameJoinQuery { + return new LuDataFrameJoinQuery(this, right, options); + } + + /** Plans a bounded, source-aligned unique-right lookup while preserving both batch topologies. */ + lookup< + Right extends GPUTypeMap, + LeftKey extends LuDataFrameColumnNamesOfFormat, + RightKey extends LuDataFrameColumnNamesOfFormat + >( + right: LuDataFrame, + options: LuDataFrameLookupOptions + ): LuDataFrameLookupQuery { + return new LuDataFrameLookupQuery(this, right, options); + } + /** Materializes reusable GPU graph passes and compiler-owned selection/index/count outputs. */ compile( graph: GPUCommandGraph diff --git a/modules/experimental/src/ludf/lu-data-frame.ts b/modules/experimental/src/ludf/lu-data-frame.ts index d2a6691d45..6415d7dd97 100644 --- a/modules/experimental/src/ludf/lu-data-frame.ts +++ b/modules/experimental/src/ludf/lu-data-frame.ts @@ -33,6 +33,12 @@ import type { LuDataFrameGroupByQuery } from './lu-group-by-query'; import type {LuDataFrameHistogramOptions, LuDataFrameHistogramQuery} from './lu-histogram-query'; +import type { + LuDataFrameJoinOptions, + LuDataFrameJoinQuery, + LuDataFrameLookupOptions, + LuDataFrameLookupQuery +} from './lu-join-query'; import type {LuDataFrameSortOptions, LuDataFrameSortQuery} from './lu-sort-query'; /** Whether a dataframe borrows its source resources or releases them after its final view. */ @@ -263,6 +269,38 @@ export class LuDataFrame { ); } + /** Plans a stable unique-right unsigned inner join without materializing either dataframe. */ + innerJoin< + Right extends GPUTypeMap, + LeftKey extends LuDataFrameColumnNamesOfFormat, + RightKey extends LuDataFrameColumnNamesOfFormat + >( + right: LuDataFrame, + options: LuDataFrameJoinOptions + ): LuDataFrameJoinQuery { + this.assertAvailable(); + return new LuDataFrameQuery(this, [], this.columnNames).innerJoin( + right, + options + ); + } + + /** Plans bounded, source-aligned unique-right lookups without flattening source batches. */ + lookup< + Right extends GPUTypeMap, + LeftKey extends LuDataFrameColumnNamesOfFormat, + RightKey extends LuDataFrameColumnNamesOfFormat + >( + right: LuDataFrame, + options: LuDataFrameLookupOptions + ): LuDataFrameLookupQuery { + this.assertAvailable(); + return new LuDataFrameQuery(this, [], this.columnNames).lookup( + right, + options + ); + } + /** * Returns an independent borrowed projection without mutating or destroying source columns. * diff --git a/modules/experimental/src/ludf/lu-join-compiler.ts b/modules/experimental/src/ludf/lu-join-compiler.ts new file mode 100644 index 0000000000..43077970bc --- /dev/null +++ b/modules/experimental/src/ludf/lu-join-compiler.ts @@ -0,0 +1,757 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// Copyright (c) vis.gl contributors + +import {Buffer, type Device} from '@luma.gl/core'; +import {GPUData, GPUVector, type GPUTable, type GPUTypeMap} from '@luma.gl/tables'; +import {GPUBatchHashIndex} from '../gpu-primitives/gpu-batch-hash-index'; +import { + type GPUCommandGraph, + type GraphBufferUse, + type GraphDataView, + type GraphVectorView +} from '../gpu-primitives/gpu-command-graph'; +import { + GPUHashIndexQuery, + GPU_HASH_INDEX_EMPTY_KEY, + GPU_HASH_INDEX_STATISTICS_LENGTH, + GPU_HASH_QUERY_STATISTICS_LENGTH +} from '../gpu-primitives/gpu-hash-index'; +import {GPUHashJoin} from '../gpu-primitives/gpu-hash-join'; +import {GPUScan} from '../gpu-primitives/gpu-scan'; +import { + createTransientVectorView, + createTransientView, + getViewElementOffset +} from '../gpu-primitives/graph-data-view-utils'; +import { + LU_ANALYTICS_WORKGROUP_SIZE, + addLuAnalyticsComputePass, + getLuAnalyticsVector, + validateLuAnalyticsSource +} from './lu-analytics-compiler-utils'; +import type {LuDataFrame} from './lu-data-frame'; +import type {LuDataFrameDerivedColumn} from './lu-data-frame-query'; +import type {LuExpression} from './lu-expression'; +import type {LuDataFrameNormalizedJoinOptions} from './lu-join-query'; +import { + CompiledLuDataFrameQuery, + compileLuDataFrameQuery, + type CompiledLuDataFrameQueryProps, + type LuDataFrameQueryExtensionContext, + type LuDataFrameQueryExtensionResult, + type LuDataFrameQueryParameters +} from './lu-query-compiler'; + +const UINT32_BYTE_LENGTH = Uint32Array.BYTES_PER_ELEMENT; +const MAXIMUM_UINT32 = 0xffffffff; + +type LuJoinIndexState = { + index: GPUBatchHashIndex; + indexStatistics: GPUVector<'uint32'>; + contractViolation: GPUVector<'uint32'>; + leftKeys: GraphVectorView<'uint32'>; + maskedLeftKeys: GraphVectorView<'uint32'>; +}; + +type LuJoinCommonResources = { + right: LuDataFrame; + rightRowIndices: GPUVector<'uint32'>; + indexStatistics: GPUVector<'uint32'>; + lookupStatistics: GPUVector<'uint32'>; + contractViolation: GPUVector<'uint32'>; +}; + +type LuJoinResources = LuJoinCommonResources & { + requiredCounts: GPUVector<'uint32'>; + overflows: GPUVector<'uint32'>; +}; + +type LuLookupResources = LuJoinCommonResources & { + matchMask: GPUVector<'uint32'>; + probeCounts: GPUVector<'uint32'>; +}; + +/** Shared right-source lease and explicit GPU diagnostics retained by bounded hash consumers. */ +abstract class CompiledLuDataFrameHashQuery< + Left extends GPUTypeMap, + Right extends GPUTypeMap +> extends CompiledLuDataFrameQuery { + /** Independently batched, retained right source table; source rows are never repacked. */ + readonly rightTable: GPUTable; + /** Stable right source-row IDs aligned with compacted pairs or original lookup rows. */ + readonly rightRowIndices: GPUVector<'uint32'>; + /** Six GPU words: unique, duplicate, overflow, invalid, total probes, and maximum probes. */ + readonly indexStatistics: GPUVector<'uint32'>; + /** Four GPU-resident lookup statistics for each original left source batch. */ + readonly lookupStatistics: GPUVector<'uint32'>; + /** One nonzero GPU word when right uniqueness, reserved-key, or index completeness fails. */ + readonly contractViolation: GPUVector<'uint32'>; + + private readonly retainedRight: LuDataFrame; + + /** @internal */ + constructor(props: CompiledLuDataFrameQueryProps, resources: LuJoinCommonResources) { + super(props); + this.retainedRight = resources.right; + this.rightTable = resources.right.table; + this.rightRowIndices = resources.rightRowIndices; + this.indexStatistics = resources.indexStatistics; + this.lookupStatistics = resources.lookupStatistics; + this.contractViolation = resources.contractViolation; + } + + /** Releases graph-owned work, both sets of owned outputs, and both retained source leases. */ + override destroy(): void { + super.destroy(); + this.retainedRight.destroy(); + } +} + +/** Stable, source-batch-preserving unique-right inner join with explicit bounded diagnostics. */ +export class CompiledLuDataFrameJoin< + Left extends GPUTypeMap = GPUTypeMap, + Right extends GPUTypeMap = GPUTypeMap +> extends CompiledLuDataFrameHashQuery { + /** Exact required pair count for each source batch, independent of publication capacity. */ + readonly requiredCounts: GPUVector<'uint32'>; + /** One source-index or per-batch publication overflow flag for each original left batch. */ + readonly overflows: GPUVector<'uint32'>; + + /** @internal */ + constructor(props: CompiledLuDataFrameQueryProps, resources: LuJoinResources) { + super(props, resources); + this.requiredCounts = resources.requiredCounts; + this.overflows = resources.overflows; + } +} + +/** Source-aligned bounded unique-right lookup preserving every original left GPU row. */ +export class CompiledLuDataFrameLookup< + Left extends GPUTypeMap = GPUTypeMap, + Right extends GPUTypeMap = GPUTypeMap +> extends CompiledLuDataFrameHashQuery { + /** Nonzero for source rows with one valid, selected, unique-right match. */ + readonly matchMask: GPUVector<'uint32'>; + /** Number of bounded hash probes performed independently for each source row. */ + readonly probeCounts: GPUVector<'uint32'>; + + /** @internal */ + constructor(props: CompiledLuDataFrameQueryProps, resources: LuLookupResources) { + super(props, resources); + this.matchMask = resources.matchMask; + this.probeCounts = resources.probeCounts; + } +} + +/** Compiles a nullable, filtered, bounded inner join against an independently batched right side. */ +export function compileLuDataFrameJoin< + Source extends GPUTypeMap, + Selection extends GPUTypeMap, + Right extends GPUTypeMap +>( + source: LuDataFrame, + predicates: readonly LuExpression[], + selectedColumns: readonly (keyof Selection & string)[], + derivedColumns: readonly LuDataFrameDerivedColumn[], + right: LuDataFrame, + options: LuDataFrameNormalizedJoinOptions, + graph: GPUCommandGraph +): CompiledLuDataFrameJoin { + validateLuJoinSources(source, right, options, graph); + const retainedRight = right.select(right.columnNames); + try { + return compileLuDataFrameQuery< + Source, + Selection, + Selection, + CompiledLuDataFrameJoin + >(source, predicates, selectedColumns, graph, derivedColumns, { + allowEmptyPredicates: true, + prepare: context => addLuInnerJoinToGraph(context, retainedRight, options) + }); + } catch (error) { + retainedRight.destroy(); + throw error; + } +} + +/** Compiles a nullable, filtered, source-aligned lookup without changing source batch boundaries. */ +export function compileLuDataFrameLookup< + Source extends GPUTypeMap, + Selection extends GPUTypeMap, + Right extends GPUTypeMap +>( + source: LuDataFrame, + predicates: readonly LuExpression[], + selectedColumns: readonly (keyof Selection & string)[], + derivedColumns: readonly LuDataFrameDerivedColumn[], + right: LuDataFrame, + options: LuDataFrameNormalizedJoinOptions, + graph: GPUCommandGraph +): CompiledLuDataFrameLookup { + validateLuJoinSources(source, right, options, graph); + const retainedRight = right.select(right.columnNames); + try { + return compileLuDataFrameQuery< + Source, + Selection, + Selection, + CompiledLuDataFrameLookup + >(source, predicates, selectedColumns, graph, derivedColumns, { + allowEmptyPredicates: true, + prepare: context => addLuLookupToGraph(context, retainedRight, options) + }); + } catch (error) { + retainedRight.destroy(); + throw error; + } +} + +/** Rejects missing validity, unsupported packed layouts, and unrepresentable IDs before GPU work. */ +function validateLuJoinSources( + source: LuDataFrame, + right: LuDataFrame, + options: LuDataFrameNormalizedJoinOptions, + graph: GPUCommandGraph +): void { + validateLuAnalyticsSource(source, [options.leftOn]); + validateLuAnalyticsSource(right, [options.rightOn]); + const indexByteLength = options.indexCapacity * UINT32_BYTE_LENGTH; + if ( + indexByteLength > graph.device.limits.maxBufferSize || + indexByteLength > graph.device.limits.maxStorageBufferBindingSize + ) { + throw new Error('LuDataFrame join index exceeds available GPU buffer capacity'); + } + for (const table of [source, right]) { + let sourceOffset = 0; + for (const batch of table.batches) { + const offset = batch.sourceInfo?.sourceRowIndexOffset ?? sourceOffset; + if ( + !Number.isSafeInteger(offset) || + offset < 0 || + offset > MAXIMUM_UINT32 || + offset + Math.max(batch.numRows - 1, 0) > MAXIMUM_UINT32 + ) { + throw new Error('LuDataFrame join source-row identifiers must fit uint32'); + } + if (table === source && batch.numRows * options.maxProbeCount > MAXIMUM_UINT32) { + throw new Error('LuDataFrame left join probe counts must fit uint32'); + } + sourceOffset += batch.numRows; + } + } +} + +/** Materializes owned pair diagnostics while sharing one chunk-preserving right index. */ +function addLuInnerJoinToGraph( + context: LuDataFrameQueryExtensionContext, + right: LuDataFrame, + options: LuDataFrameNormalizedJoinOptions +): LuDataFrameQueryExtensionResult> { + const ownedVectors: GPUVector[] = []; + const id = `${context.queryId}-inner-join`; + try { + const indexState = buildLuJoinIndex(context, right, options, ownedVectors, id); + const lengths = context.table.batches.map(batch => batch.numRows); + const rightRowIndices = createLuJoinOutputVector(context.graph.device, `${id}-right`, lengths); + ownedVectors.push(rightRowIndices); + const requiredCounts = createLuJoinOutputVector( + context.graph.device, + `${id}-required`, + lengths.map(() => 1) + ); + ownedVectors.push(requiredCounts); + const overflows = createLuJoinOutputVector( + context.graph.device, + `${id}-overflow`, + lengths.map(() => 1) + ); + ownedVectors.push(overflows); + const lookupStatistics = createLuJoinOutputVector( + context.graph.device, + `${id}-lookup-statistics`, + lengths.map(() => GPU_HASH_QUERY_STATISTICS_LENGTH) + ); + ownedVectors.push(lookupStatistics); + + const rightRows = context.graph.importGPUVector(`${id}-right-rows`, rightRowIndices); + const required = context.graph.importGPUVector(`${id}-required-counts`, requiredCounts); + const overflow = context.graph.importGPUVector(`${id}-overflows`, overflows); + const statistics = context.graph.importGPUVector(`${id}-statistics`, lookupStatistics); + const matches = createTransientVectorView( + context.graph, + `${id}-matches`, + context.selectionMask + ); + + let firstLeftRow = 0; + for (const [batchIndex, batch] of context.table.batches.entries()) { + const batchId = `${id}-batch-${batchIndex}`; + const capacity = Math.min(options.capacity ?? batch.numRows, batch.numRows); + const outputLeftRows = getLuJoinCapacityView( + context.graph, + context.rowIndices.data[batchIndex], + capacity + ); + const outputRightRows = getLuJoinCapacityView( + context.graph, + rightRows.data[batchIndex], + capacity + ); + + new GPUHashJoin({ + id: batchId, + index: indexState.index, + keys: indexState.maskedLeftKeys.data[batchIndex], + firstLeftRow: batch.sourceInfo?.sourceRowIndexOffset ?? firstLeftRow, + outputLeftRows, + outputRightRows, + count: required.data[batchIndex], + overflow: overflow.data[batchIndex], + statistics: statistics.data[batchIndex], + found: matches.data[batchIndex], + maxProbeCount: options.maxProbeCount + }).addToGraph(context.graph); + + const offsets = createTransientView( + context.graph, + `${batchId}-match-offsets`, + 'uint32', + batch.numRows + ); + new GPUScan({ + id: `${batchId}-published-offsets`, + input: matches.data[batchIndex], + output: offsets + }).addToGraph(context.graph); + addLuJoinPublishPass(context.graph, `${batchId}-publish`, { + matches: matches.data[batchIndex], + offsets, + selection: context.selectionMask.data[batchIndex], + leftRows: context.rowIndices.data[batchIndex], + rightRows: rightRows.data[batchIndex], + required: required.data[batchIndex], + published: context.selectedCounts.data[batchIndex], + capacity + }); + firstLeftRow += batch.numRows; + } + + const resources: LuJoinResources = { + right, + rightRowIndices, + requiredCounts, + overflows, + lookupStatistics, + indexStatistics: indexState.indexStatistics, + contractViolation: indexState.contractViolation + }; + return { + table: context.table, + validity: context.validity, + dictionaries: context.dictionaries, + ownedVectors, + createCompiled: props => new CompiledLuDataFrameJoin(props, resources) + }; + } catch (error) { + for (const vector of ownedVectors) { + vector.destroy(); + } + throw error; + } +} + +/** Publishes right row IDs, match flags, probes, and diagnostics for every left source row. */ +function addLuLookupToGraph( + context: LuDataFrameQueryExtensionContext, + right: LuDataFrame, + options: LuDataFrameNormalizedJoinOptions +): LuDataFrameQueryExtensionResult> { + const ownedVectors: GPUVector[] = []; + const id = `${context.queryId}-lookup`; + try { + const indexState = buildLuJoinIndex(context, right, options, ownedVectors, id); + const lengths = context.table.batches.map(batch => batch.numRows); + const rightRowIndices = createLuJoinOutputVector(context.graph.device, `${id}-right`, lengths); + ownedVectors.push(rightRowIndices); + const matchMask = createLuJoinOutputVector(context.graph.device, `${id}-matched`, lengths); + ownedVectors.push(matchMask); + const probeCounts = createLuJoinOutputVector(context.graph.device, `${id}-probes`, lengths); + ownedVectors.push(probeCounts); + const lookupStatistics = createLuJoinOutputVector( + context.graph.device, + `${id}-statistics`, + lengths.map(() => GPU_HASH_QUERY_STATISTICS_LENGTH) + ); + ownedVectors.push(lookupStatistics); + + const rightRows = context.graph.importGPUVector(`${id}-right-rows`, rightRowIndices); + const matches = context.graph.importGPUVector(`${id}-match-mask`, matchMask); + const probes = context.graph.importGPUVector(`${id}-probe-counts`, probeCounts); + const statistics = context.graph.importGPUVector(`${id}-query-statistics`, lookupStatistics); + + for (const [batchIndex, keys] of indexState.maskedLeftKeys.data.entries()) { + new GPUHashIndexQuery({ + id: `${id}-batch-${batchIndex}`, + index: indexState.index, + keys, + values: rightRows.data[batchIndex], + found: matches.data[batchIndex], + probes: probes.data[batchIndex], + statistics: statistics.data[batchIndex], + maxProbeCount: options.maxProbeCount + }).addToGraph(context.graph); + } + + const resources: LuLookupResources = { + right, + rightRowIndices, + matchMask, + probeCounts, + lookupStatistics, + indexStatistics: indexState.indexStatistics, + contractViolation: indexState.contractViolation + }; + return { + table: context.table, + validity: context.validity, + dictionaries: context.dictionaries, + ownedVectors, + createCompiled: props => new CompiledLuDataFrameLookup(props, resources) + }; + } catch (error) { + for (const vector of ownedVectors) { + vector.destroy(); + } + throw error; + } +} + +/** Builds one right index without concatenation and sanitizes filtered/nullable left source keys. */ +function buildLuJoinIndex( + context: LuDataFrameQueryExtensionContext, + right: LuDataFrame, + options: LuDataFrameNormalizedJoinOptions, + ownedVectors: GPUVector[], + id: string +): LuJoinIndexState { + const graph = context.graph; + const leftVector = getLuAnalyticsVector(context, options.leftOn); + if (leftVector.format !== 'uint32') { + throw new Error('LuDataFrame left join keys must use packed uint32 GPU data'); + } + const leftKeys = graph.importGPUVector(`${id}-left-keys`, leftVector as GPUVector<'uint32'>); + const rightField = right.schema.fields.find(field => field.name === options.rightOn); + const rightVector = right.table.gpuVectors[options.rightOn]; + if (rightField?.format !== 'uint32' && rightVector?.format !== 'uint32') { + throw new Error('LuDataFrame right join keys must use packed uint32 GPU data'); + } + const rightKeys = graph.importGPUVector( + `${id}-right-keys`, + rightVector + ? (rightVector as GPUVector<'uint32'>) + : new GPUVector({ + type: 'data', + name: options.rightOn, + format: 'uint32', + data: [], + ownsData: false + }) + ); + const rightValidity = rightField?.nullable + ? right.validity[options.rightOn as keyof Right & string] + : undefined; + if (rightField?.nullable && rightKeys.length > 0 && !rightValidity) { + throw new Error('LuDataFrame nullable right join keys require explicit GPU validity'); + } + + const indexStatistics = createLuJoinOutputVector(graph.device, `${id}-index-statistics`, [ + GPU_HASH_INDEX_STATISTICS_LENGTH + ]); + ownedVectors.push(indexStatistics); + const contractViolation = createLuJoinOutputVector(graph.device, `${id}-contract-violation`, [1]); + ownedVectors.push(contractViolation); + const statistics = graph.importGPUVector(`${id}-index-stats`, indexStatistics).data[0]; + const violation = graph.importGPUVector(`${id}-violation`, contractViolation).data[0]; + const tableKeys = createTransientView(graph, `${id}-table-keys`, 'uint32', options.indexCapacity); + const tableValues = createTransientView( + graph, + `${id}-table-values`, + 'uint32', + options.indexCapacity + ); + + let rightOffset = 0; + const firstValues = right.batches.map(batch => { + const firstValue = batch.sourceInfo?.sourceRowIndexOffset ?? rightOffset; + rightOffset += batch.numRows; + return firstValue; + }); + const index = new GPUBatchHashIndex({ + id: `${id}-right-index`, + keys: rightKeys, + ...(rightValidity + ? {validity: graph.importGPUVector(`${id}-right-validity`, rightValidity)} + : {}), + firstValues, + tableKeys, + tableValues, + statistics, + maxProbeCount: options.maxProbeCount + }); + index.addToGraph(graph); + addLuJoinContractPass(graph, `${id}-validate-contract`, statistics, violation); + + const leftField = context.table.schema.fields.find(field => field.name === options.leftOn); + const leftValidity = leftField?.nullable + ? context.validity[options.leftOn as keyof Left & string] + : undefined; + if (leftField?.nullable && leftKeys.length > 0 && !leftValidity) { + throw new Error('LuDataFrame nullable left join keys require explicit GPU validity'); + } + const validity = leftValidity + ? graph.importGPUVector(`${id}-left-validity`, leftValidity) + : undefined; + const maskedLeftKeys = createTransientVectorView( + graph, + `${id}-masked-keys`, + context.selectionMask + ); + for (const [batchIndex, keys] of leftKeys.data.entries()) { + if (keys.length > 0) { + addLuJoinPrepareKeysPass(graph, `${id}-prepare-batch-${batchIndex}`, { + input: keys, + selection: context.selectionMask.data[batchIndex], + validity: validity?.data[batchIndex], + violation, + output: maskedLeftKeys.data[batchIndex] + }); + } + } + return {index, indexStatistics, contractViolation, leftKeys, maskedLeftKeys}; +} + +/** Allocates one independently owned packed uint32 chunk for each caller-preserved source batch. */ +function createLuJoinOutputVector( + device: Device, + name: string, + lengths: readonly number[] +): GPUVector<'uint32'> { + const chunks: GPUData<'uint32'>[] = []; + try { + for (const [batchIndex, length] of lengths.entries()) { + const buffer = device.createBuffer({ + id: `${name}-batch-${batchIndex}`, + byteLength: Math.max(length, 1) * UINT32_BYTE_LENGTH, + usage: Buffer.STORAGE | Buffer.COPY_SRC | Buffer.COPY_DST | Buffer.VERTEX + }); + try { + chunks.push(new GPUData({buffer, format: 'uint32', length, ownsBuffer: true})); + } catch (error) { + buffer.destroy(); + throw error; + } + } + return new GPUVector({type: 'data', name, format: 'uint32', data: chunks, ownsData: true}); + } catch (error) { + for (const chunk of chunks) { + chunk.destroy(); + } + throw error; + } +} + +/** Uses the canonical imported graph handle while exposing a bounded logical pair-output prefix. */ +function getLuJoinCapacityView( + graph: GPUCommandGraph, + view: GraphDataView<'uint32'>, + capacity: number +): GraphDataView<'uint32'> { + return graph.createDataView(view.buffer, { + format: 'uint32', + length: capacity, + byteOffset: view.byteOffset + }); +} + +/** Publishes one strict GPU contract flag for duplicates, bounded index overflow, or reserved keys. */ +function addLuJoinContractPass( + graph: GPUCommandGraph, + id: string, + statistics: GraphDataView<'uint32'>, + violation: GraphDataView<'uint32'> +): void { + const source = /* wgsl */ ` +const STATISTICS_OFFSET: u32 = ${getViewElementOffset(statistics)}u; +const OUTPUT_OFFSET: u32 = ${getViewElementOffset(violation)}u; +@group(0) @binding(0) var indexStatistics: array; +@group(0) @binding(1) var contractViolation: array; + +@compute @workgroup_size(1) +fn main() { + let invalid = indexStatistics[STATISTICS_OFFSET + 1u] != 0u || + indexStatistics[STATISTICS_OFFSET + 2u] != 0u || + indexStatistics[STATISTICS_OFFSET + 3u] != 0u; + contractViolation[OUTPUT_OFFSET] = select(0u, 1u, invalid); +}`; + addLuAnalyticsComputePass(graph, { + id, + source, + resources: [ + {buffer: statistics, usage: 'storage-read'}, + {buffer: violation, usage: 'storage-write'} + ], + bindings: {indexStatistics: statistics, contractViolation: violation}, + length: 1 + }); +} + +/** Excludes filtered/null rows and invalid right indexes without repacking left source chunks. */ +function addLuJoinPrepareKeysPass( + graph: GPUCommandGraph, + id: string, + props: { + input: GraphDataView<'uint32'>; + selection: GraphDataView<'uint32'>; + validity?: GraphDataView<'uint32'>; + violation: GraphDataView<'uint32'>; + output: GraphDataView<'uint32'>; + } +): void { + const nullable = Boolean(props.validity); + const validityBinding = nullable + ? '@group(0) @binding(2) var validityMask: array;' + : ''; + const validityOffset = props.validity ? getViewElementOffset(props.validity) : 0; + const firstTailBinding = nullable ? 3 : 2; + const isValid = nullable ? 'validityMask[VALIDITY_OFFSET + index] != 0u' : 'true'; + const source = /* wgsl */ ` +const ELEMENT_COUNT: u32 = ${props.input.length}u; +const KEY_OFFSET: u32 = ${getViewElementOffset(props.input)}u; +const SELECTION_OFFSET: u32 = ${getViewElementOffset(props.selection)}u; +const VALIDITY_OFFSET: u32 = ${validityOffset}u; +const CONTRACT_OFFSET: u32 = ${getViewElementOffset(props.violation)}u; +const OUTPUT_OFFSET: u32 = ${getViewElementOffset(props.output)}u; +@group(0) @binding(0) var sourceKeys: array; +@group(0) @binding(1) var selectionMask: array; +${validityBinding} +@group(0) @binding(${firstTailBinding}) var contractViolation: array; +@group(0) @binding(${firstTailBinding + 1}) var preparedKeys: array; + +@compute @workgroup_size(${LU_ANALYTICS_WORKGROUP_SIZE}) +fn main(@builtin(global_invocation_id) globalId: vec3) { + let index = globalId.x; + if (index < ELEMENT_COUNT) { + let selected = selectionMask[SELECTION_OFFSET + index] != 0u; + let valid = ${isValid}; + let permitted = selected && valid && contractViolation[CONTRACT_OFFSET] == 0u; + preparedKeys[OUTPUT_OFFSET + index] = select( + ${GPU_HASH_INDEX_EMPTY_KEY}u, + sourceKeys[KEY_OFFSET + index], + permitted + ); + } +}`; + const resources: GraphBufferUse[] = [ + {buffer: props.input, usage: 'storage-read'}, + {buffer: props.selection, usage: 'storage-read'} + ]; + const bindings: Record = { + sourceKeys: props.input, + selectionMask: props.selection + }; + if (props.validity) { + resources.push({buffer: props.validity, usage: 'storage-read'}); + bindings['validityMask'] = props.validity; + } + resources.push( + {buffer: props.violation, usage: 'storage-read'}, + {buffer: props.output, usage: 'storage-write'} + ); + bindings['contractViolation'] = props.violation; + bindings['preparedKeys'] = props.output; + addLuAnalyticsComputePass(graph, { + id, + source, + resources, + bindings, + length: props.input.length + }); +} + +/** Keeps inherited masks/counts and stable left/right row prefixes coherent after bounded joins. */ +function addLuJoinPublishPass( + graph: GPUCommandGraph, + id: string, + props: { + matches: GraphDataView<'uint32'>; + offsets: GraphDataView<'uint32'>; + selection: GraphDataView<'uint32'>; + leftRows: GraphDataView<'uint32'>; + rightRows: GraphDataView<'uint32'>; + required: GraphDataView<'uint32'>; + published: GraphDataView<'uint32'>; + capacity: number; + } +): void { + const source = /* wgsl */ ` +const ELEMENT_COUNT: u32 = ${props.matches.length}u; +const MATCH_OFFSET: u32 = ${getViewElementOffset(props.matches)}u; +const OFFSET_OFFSET: u32 = ${getViewElementOffset(props.offsets)}u; +const SELECTION_OFFSET: u32 = ${getViewElementOffset(props.selection)}u; +const LEFT_OFFSET: u32 = ${getViewElementOffset(props.leftRows)}u; +const RIGHT_OFFSET: u32 = ${getViewElementOffset(props.rightRows)}u; +const REQUIRED_OFFSET: u32 = ${getViewElementOffset(props.required)}u; +const PUBLISHED_OFFSET: u32 = ${getViewElementOffset(props.published)}u; +const OUTPUT_CAPACITY: u32 = ${props.capacity}u; +@group(0) @binding(0) var matchedRows: array; +@group(0) @binding(1) var matchedOffsets: array; +@group(0) @binding(2) var selectionMask: array; +@group(0) @binding(3) var outputLeftRows: array; +@group(0) @binding(4) var outputRightRows: array; +@group(0) @binding(5) var requiredCounts: array; +@group(0) @binding(6) var selectedCounts: array; + +@compute @workgroup_size(${LU_ANALYTICS_WORKGROUP_SIZE}) +fn main(@builtin(global_invocation_id) globalId: vec3) { + let index = globalId.x; + let published = min(requiredCounts[REQUIRED_OFFSET], OUTPUT_CAPACITY); + if (index == 0u) { + selectedCounts[PUBLISHED_OFFSET] = published; + } + if (index < ELEMENT_COUNT) { + let selected = matchedRows[MATCH_OFFSET + index] != 0u && + matchedOffsets[OFFSET_OFFSET + index] < OUTPUT_CAPACITY; + selectionMask[SELECTION_OFFSET + index] = select(0u, 1u, selected); + if (index >= published) { + outputLeftRows[LEFT_OFFSET + index] = 0u; + outputRightRows[RIGHT_OFFSET + index] = 0u; + } + } +}`; + addLuAnalyticsComputePass(graph, { + id, + source, + resources: [ + {buffer: props.matches, usage: 'storage-read'}, + {buffer: props.offsets, usage: 'storage-read'}, + {buffer: props.selection, usage: 'storage-write'}, + {buffer: props.leftRows, usage: 'storage-read-write'}, + {buffer: props.rightRows, usage: 'storage-read-write'}, + {buffer: props.required, usage: 'storage-read'}, + {buffer: props.published, usage: 'storage-write'} + ], + bindings: { + matchedRows: props.matches, + matchedOffsets: props.offsets, + selectionMask: props.selection, + outputLeftRows: props.leftRows, + outputRightRows: props.rightRows, + requiredCounts: props.required, + selectedCounts: props.published + }, + length: Math.max(props.matches.length, 1) + }); +} diff --git a/modules/experimental/src/ludf/lu-join-query.ts b/modules/experimental/src/ludf/lu-join-query.ts new file mode 100644 index 0000000000..88631c2271 --- /dev/null +++ b/modules/experimental/src/ludf/lu-join-query.ts @@ -0,0 +1,301 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// Copyright (c) vis.gl contributors + +import type {GPUTypeMap} from '@luma.gl/tables'; +import type {GPUCommandGraph} from '../gpu-primitives/gpu-command-graph'; +import type {LuDataFrame, LuDataFrameDictionary} from './lu-data-frame'; +import type {LuDataFrameQuery} from './lu-data-frame-query'; +import {getLuDataFrameAnalyticColumnFormat} from './lu-global-aggregation-query'; +import type {LuDataFrameColumnNamesOfFormat} from './lu-group-by-query'; +import { + compileLuDataFrameJoin, + compileLuDataFrameLookup, + type CompiledLuDataFrameJoin, + type CompiledLuDataFrameLookup +} from './lu-join-compiler'; +import type {LuDataFrameQueryParameters} from './lu-query-compiler'; + +const MAXIMUM_UINT32 = 0xffffffff; +const MAXIMUM_HASH_CAPACITY = 0x80000000; + +/** Unique-right inner-join columns, per-batch output capacity, and bounded hash-index controls. */ +export type LuDataFrameJoinOptions< + LeftKey extends string = string, + RightKey extends string = string +> = Readonly<{ + /** Selected unsigned key column on the left dataframe. */ + leftOn: LeftKey; + /** Unsigned unique-key column on the right dataframe. */ + rightOn: RightKey; + /** Optional maximum number of published pairs in each original left source batch. */ + capacity?: number; + /** Optional power-of-two slot count for the reusable GPU hash index. */ + indexCapacity?: number; + /** Optional bounded linear-probe count for index construction and left lookups. */ + maxProbeCount?: number; +}>; + +/** Source-aligned unique-right lookup columns and bounded GPU hash-index controls. */ +export type LuDataFrameLookupOptions< + LeftKey extends string = string, + RightKey extends string = string +> = Readonly<{ + /** Selected unsigned key column on the left dataframe. */ + leftOn: LeftKey; + /** Unsigned unique-key column on the right dataframe. */ + rightOn: RightKey; + /** Optional power-of-two slot count for the reusable GPU hash index. */ + indexCapacity?: number; + /** Optional bounded linear-probe count for index construction and left lookups. */ + maxProbeCount?: number; +}>; + +/** Validated, immutable unique-right index and per-batch join controls. @internal */ +export type LuDataFrameNormalizedJoinOptions< + LeftKey extends string = string, + RightKey extends string = string +> = Readonly<{ + leftOn: LeftKey; + rightOn: RightKey; + indexCapacity: number; + maxProbeCount: number; + capacity?: number; +}>; + +/** Immutable bounded, stable, unique-right unsigned inner join without source materialization. */ +export class LuDataFrameJoinQuery< + Logical extends GPUTypeMap, + SelectedColumns extends keyof Logical & string, + Right extends GPUTypeMap, + LeftKey extends LuDataFrameColumnNamesOfFormat, + RightKey extends LuDataFrameColumnNamesOfFormat, + Source extends GPUTypeMap = Logical +> { + /** Filtered, projected, or derived left dataframe plan. */ + readonly query: LuDataFrameQuery; + /** Borrowed right dataframe, retained only during explicit graph compilation. */ + readonly right: LuDataFrame; + /** Immutable, fully validated key names and fixed GPU index/publication bounds. */ + readonly options: LuDataFrameNormalizedJoinOptions; + + /** Validates both key schemas and bounded capacities without touching GPU resources. @internal */ + constructor( + query: LuDataFrameQuery, + right: LuDataFrame, + options: LuDataFrameJoinOptions + ) { + this.query = query; + this.right = right; + this.options = normalizeLuDataFrameJoinOptions(query, right, options, true); + Object.freeze(this); + } + + /** Adds one chunk-preserving unique-right index and bounded joins to the caller-owned graph. */ + compile( + graph: GPUCommandGraph + ): CompiledLuDataFrameJoin, Right> { + return compileLuDataFrameJoin, Right>( + this.query.source, + this.query.predicates, + this.query.selectedColumns, + this.query.derivedColumns, + this.right, + this.options, + graph + ); + } +} + +/** Immutable bounded, source-aligned unique-right unsigned lookup without row compaction. */ +export class LuDataFrameLookupQuery< + Logical extends GPUTypeMap, + SelectedColumns extends keyof Logical & string, + Right extends GPUTypeMap, + LeftKey extends LuDataFrameColumnNamesOfFormat, + RightKey extends LuDataFrameColumnNamesOfFormat, + Source extends GPUTypeMap = Logical +> { + /** Filtered, projected, or derived left dataframe plan. */ + readonly query: LuDataFrameQuery; + /** Borrowed right dataframe, retained only during explicit graph compilation. */ + readonly right: LuDataFrame; + /** Immutable, fully validated key names and fixed GPU hash-index bounds. */ + readonly options: LuDataFrameNormalizedJoinOptions; + + /** Validates both unsigned key schemas without retaining either dataframe. @internal */ + constructor( + query: LuDataFrameQuery, + right: LuDataFrame, + options: LuDataFrameLookupOptions + ) { + this.query = query; + this.right = right; + this.options = normalizeLuDataFrameJoinOptions(query, right, options, false); + Object.freeze(this); + } + + /** Adds one chunk-preserving right index and bounded source-aligned left lookup graph passes. */ + compile( + graph: GPUCommandGraph + ): CompiledLuDataFrameLookup, Right> { + return compileLuDataFrameLookup, Right>( + this.query.source, + this.query.predicates, + this.query.selectedColumns, + this.query.derivedColumns, + this.right, + this.options, + graph + ); + } +} + +/** Validates two existing unsigned keys and normalizes safely bounded right-index controls. */ +function normalizeLuDataFrameJoinOptions< + Logical extends GPUTypeMap, + SelectedColumns extends keyof Logical & string, + Right extends GPUTypeMap, + LeftKey extends LuDataFrameColumnNamesOfFormat, + RightKey extends LuDataFrameColumnNamesOfFormat, + Source extends GPUTypeMap +>( + query: LuDataFrameQuery, + right: LuDataFrame, + options: LuDataFrameJoinOptions | LuDataFrameLookupOptions, + allowCapacity: boolean +): LuDataFrameNormalizedJoinOptions { + if (!options || typeof options !== 'object') { + throw new Error('LuDataFrame joins require explicit leftOn and rightOn key columns'); + } + if (!right || typeof right !== 'object') { + throw new Error('LuDataFrame joins require an existing right dataframe'); + } + if (!query.selectedColumns.includes(options.leftOn)) { + throw new Error(`LuDataFrame left join key "${options.leftOn}" is not selected`); + } + if (getLuDataFrameAnalyticColumnFormat(query, options.leftOn) !== 'uint32') { + throw new Error(`LuDataFrame left join key "${options.leftOn}" must be uint32`); + } + if (query.source.table.gpuConstants[options.leftOn]) { + throw new Error(`LuDataFrame left join key "${options.leftOn}" cannot be constant`); + } + if (!right.columnNames.includes(options.rightOn)) { + throw new Error(`LuDataFrame right join key "${options.rightOn}" does not exist`); + } + const rightField = right.schema.fields.find(field => field.name === options.rightOn); + const rightFormat = right.table.gpuColumns[options.rightOn]?.format ?? rightField?.format; + if (rightFormat !== 'uint32') { + throw new Error(`LuDataFrame right join key "${options.rightOn}" must be uint32`); + } + if (right.table.gpuConstants[options.rightOn]) { + throw new Error(`LuDataFrame right join key "${options.rightOn}" cannot be constant`); + } + assertCompatibleLuDataFrameJoinDictionaries(query, right, options.leftOn, options.rightOn); + if (!Number.isSafeInteger(right.numRows) || right.numRows < 0 || right.numRows > MAXIMUM_UINT32) { + throw new Error('LuDataFrame right join row counts must fit uint32'); + } + + const indexCapacity = options.indexCapacity ?? getDefaultLuDataFrameJoinCapacity(right.numRows); + if ( + !Number.isSafeInteger(indexCapacity) || + indexCapacity < 1 || + indexCapacity > MAXIMUM_HASH_CAPACITY || + !Number.isInteger(Math.log2(indexCapacity)) + ) { + throw new Error('LuDataFrame joins require a positive power-of-two uint32 index capacity'); + } + const maximumSafeProbeCount = Math.max( + 1, + Math.floor(MAXIMUM_UINT32 / Math.max(right.numRows, 1)) + ); + const maxProbeCount = options.maxProbeCount ?? Math.min(indexCapacity, maximumSafeProbeCount); + if ( + !Number.isSafeInteger(maxProbeCount) || + maxProbeCount < 1 || + maxProbeCount > indexCapacity || + right.numRows * maxProbeCount > MAXIMUM_UINT32 + ) { + throw new Error('LuDataFrame joins require a safely bounded uint32 probe count'); + } + + if ('capacity' in options) { + const capacity = options.capacity; + if (!allowCapacity) { + throw new Error('LuDataFrame lookups do not accept a compacted output capacity'); + } + if ( + capacity !== undefined && + (!Number.isSafeInteger(capacity) || capacity < 0 || capacity > MAXIMUM_UINT32) + ) { + throw new Error('LuDataFrame joins require a nonnegative uint32 output capacity'); + } + } + + return Object.freeze({ + leftOn: options.leftOn, + rightOn: options.rightOn, + indexCapacity, + maxProbeCount, + ...(allowCapacity && 'capacity' in options && options.capacity !== undefined + ? {capacity: options.capacity} + : {}) + }); +} + +/** Prevents incompatible dictionary codebooks from silently joining different logical labels. */ +function assertCompatibleLuDataFrameJoinDictionaries< + Logical extends GPUTypeMap, + SelectedColumns extends keyof Logical & string, + Right extends GPUTypeMap, + Source extends GPUTypeMap +>( + query: LuDataFrameQuery, + right: LuDataFrame, + leftOn: string, + rightOn: keyof Right & string +): void { + const leftDictionary = ( + query.source.dictionaries as Readonly> + )[leftOn]; + const rightDictionary = right.dictionaries[rightOn]; + if (!leftDictionary && !rightDictionary) { + return; + } + if (!leftDictionary || !rightDictionary) { + throw new Error('LuDataFrame join dictionaries must exist on both key columns'); + } + + const leftMetadata = leftDictionary as Readonly<{ + values: readonly unknown[]; + ordered?: boolean; + }>; + const rightMetadata = rightDictionary as Readonly<{ + values: readonly unknown[]; + ordered?: boolean; + }>; + const leftValues = Array.isArray(leftDictionary) ? leftDictionary : leftMetadata.values; + const rightValues = Array.isArray(rightDictionary) ? rightDictionary : rightMetadata.values; + const leftOrdered = Array.isArray(leftDictionary) ? false : Boolean(leftMetadata.ordered); + const rightOrdered = Array.isArray(rightDictionary) ? false : Boolean(rightMetadata.ordered); + if ( + leftOrdered !== rightOrdered || + leftValues.length !== rightValues.length || + leftValues.some((value, index) => !Object.is(value, rightValues[index])) + ) { + throw new Error('LuDataFrame join dictionaries must use identical labels and ordering'); + } +} + +/** Chooses a bounded power-of-two index with a maximum default half-full load factor. */ +function getDefaultLuDataFrameJoinCapacity(rightRowCount: number): number { + const requiredCapacity = Math.max(1, rightRowCount * 2); + let capacity = 1; + while (capacity < requiredCapacity) { + capacity *= 2; + if (capacity > MAXIMUM_HASH_CAPACITY) { + throw new Error('LuDataFrame default join index exceeds uint32 hash capacity'); + } + } + return capacity; +} diff --git a/modules/experimental/test/gpu-primitives/gpu-batch-hash-index.node.spec.ts b/modules/experimental/test/gpu-primitives/gpu-batch-hash-index.node.spec.ts new file mode 100644 index 0000000000..4ab485720d --- /dev/null +++ b/modules/experimental/test/gpu-primitives/gpu-batch-hash-index.node.spec.ts @@ -0,0 +1,278 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// Copyright (c) vis.gl contributors + +import {Buffer} from '@luma.gl/core'; +import { + GPUBatchHashIndex, + GPUCommandGraph, + GraphVectorView, + type GPUBatchHashIndexProps, + type GraphDataView +} from '@luma.gl/experimental'; +import {NullDevice} from '@luma.gl/test-utils'; +import {describe, expect, test, vi} from 'vitest'; + +describe('GPUBatchHashIndex planning', () => { + test('preserves source batches, empty chunks, offsets, and allocation-free graph planning', () => { + const fixture = createGraphFixture(); + const createBuffer = vi.spyOn(fixture.device, 'createBuffer'); + const addComputePass = vi.spyOn(fixture.graph, 'addComputePass'); + + try { + const index = new GPUBatchHashIndex({ + ...createIndexProps(fixture.graph, [2, 0, 3]), + validity: createVector(fixture.graph, 'validity', [2, 0, 3]), + firstValues: [40, 100, 500] + }); + expect(index.firstValues).toEqual([40, 100, 500]); + expect(Object.isFrozen(index.firstValues)).toBe(true); + expect(index.stats).toEqual({ + capacity: 8, + maxProbeCount: 8, + tableByteLength: 64, + statisticsByteLength: 24, + outputByteLength: 88, + batchCount: 3, + inputLength: 5 + }); + expect(index.updatePolicy).toBe('rebuild'); + expect(createBuffer).not.toHaveBeenCalled(); + + index.addToGraph(fixture.graph); + + expect(addComputePass.mock.calls.map(([pass]) => pass.id)).toEqual([ + 'node-batch-index-initialize', + 'node-batch-index-batch-0-build', + 'node-batch-index-batch-0-finalize', + 'node-batch-index-batch-2-build', + 'node-batch-index-batch-2-finalize' + ]); + expect(createBuffer).not.toHaveBeenCalled(); + } finally { + createBuffer.mockRestore(); + addComputePass.mockRestore(); + fixture.device.destroy(); + } + }); + + test('generates contiguous batch offsets and accepts explicit aligned payload vectors', () => { + const fixture = createGraphFixture(); + + try { + const props = createIndexProps(fixture.graph, [2, 0, 3]); + const values = createVector(fixture.graph, 'values', [2, 0, 3]); + const index = new GPUBatchHashIndex({...props, values}); + + expect(index.values).toBe(values); + expect(index.firstValues).toEqual([0, 2, 2]); + expect(() => new GPUBatchHashIndex({...props, values, firstValues: [0, 2, 2]})).toThrow( + /values and firstValues are mutually exclusive/ + ); + } finally { + fixture.device.destroy(); + } + }); + + test('rejects mismatched ordered topology and unsupported scalar layouts', () => { + const fixture = createGraphFixture(); + + try { + const props = createIndexProps(fixture.graph, [2, 0, 3]); + expect( + () => + new GPUBatchHashIndex({ + ...props, + values: createVector(fixture.graph, 'mismatched-values', [1, 0, 4]) + }) + ).toThrow(/same chunk topology/); + expect( + () => + new GPUBatchHashIndex({ + ...props, + validity: createVector(fixture.graph, 'mismatched-validity', [2, 3]) + }) + ).toThrow(/same chunk topology/); + expect( + () => + new GPUBatchHashIndex({ + ...props, + keys: createVector(fixture.graph, 'incorrect-length', [2, 0, 3], {length: 4}) + }) + ).toThrow(/length must equal its ordered source chunks/); + expect( + () => + new GPUBatchHashIndex({ + ...props, + keys: createVector(fixture.graph, 'strided-keys', [2, 0, 3], {byteStride: 8}) + }) + ).toThrow(/packed, uint32-aligned uint32/); + } finally { + fixture.device.destroy(); + } + }); + + test('rejects ambiguous offsets, uint32 source identities, and cumulative probe overflow', () => { + const fixture = createGraphFixture(); + + try { + const props = createIndexProps(fixture.graph, [2, 0, 3]); + expect(() => new GPUBatchHashIndex({...props, firstValues: [0, 2]})).toThrow( + /one value per source chunk/ + ); + expect(() => new GPUBatchHashIndex({...props, firstValues: [0xffffffff, 2, 2]})).toThrow( + /generated values must fit in uint32/ + ); + expect(() => new GPUBatchHashIndex({...props, firstValues: [-1, 2, 2]})).toThrow( + /generated values must fit in uint32/ + ); + expect(() => new GPUBatchHashIndex({...props, maxProbeCount: 9})).toThrow( + /one through capacity/ + ); + + const oversizedKeys = createVector(fixture.graph, 'oversized-keys', [0x40000000, 0x40000000]); + expect( + () => new GPUBatchHashIndex({...props, keys: oversizedKeys, maxProbeCount: 2}) + ).toThrow(/aggregate probe count must fit in uint32/); + } finally { + fixture.device.destroy(); + } + }); + + test('rejects shared output ranges and source/output aliases', () => { + const fixture = createGraphFixture(); + + try { + const props = createIndexProps(fixture.graph, [2, 0, 3]); + expect(() => new GPUBatchHashIndex({...props, tableValues: props.tableKeys})).toThrow( + /output views must not overlap/ + ); + expect( + () => + new GPUBatchHashIndex({ + ...props, + statistics: fixture.graph.createDataView(props.keys.data[0].buffer, { + format: 'uint32', + length: 2 + }) + }) + ).toThrow(/statistics must contain six uint32 rows/); + + const sourceAlias = fixture.graph.createDataView(props.tableKeys.buffer, { + format: 'uint32', + length: 2 + }); + const aliasedKeys = createVector(fixture.graph, 'alias-keys', [2, 0, 3], { + firstChunk: sourceAlias + }); + expect(() => new GPUBatchHashIndex({...props, keys: aliasedKeys})).toThrow( + /input and output views must not overlap/ + ); + } finally { + fixture.device.destroy(); + } + }); + + test('clears empty topologies exactly once without importing empty input bindings', () => { + for (const lengths of [[], [0, 0]] as readonly number[][]) { + const fixture = createGraphFixture(); + const addComputePass = vi.spyOn(fixture.graph, 'addComputePass'); + const createBuffer = vi.spyOn(fixture.device, 'createBuffer'); + + try { + new GPUBatchHashIndex(createIndexProps(fixture.graph, lengths)).addToGraph(fixture.graph); + expect(addComputePass.mock.calls.map(([pass]) => pass.id)).toEqual([ + 'node-batch-index-initialize' + ]); + expect(createBuffer).not.toHaveBeenCalled(); + } finally { + addComputePass.mockRestore(); + createBuffer.mockRestore(); + fixture.device.destroy(); + } + } + }); + + test('rejects views owned by a different command graph before adding compute passes', () => { + const fixture = createGraphFixture(); + const other = createGraphFixture(); + const addComputePass = vi.spyOn(fixture.graph, 'addComputePass'); + + try { + const props = createIndexProps(fixture.graph, [2, 0, 3]); + const index = new GPUBatchHashIndex({ + ...props, + validity: createVector(other.graph, 'external-validity', [2, 0, 3]) + }); + expect(() => index.addToGraph(fixture.graph)).toThrow( + /views must belong to the target graph/ + ); + expect(addComputePass).not.toHaveBeenCalled(); + } finally { + addComputePass.mockRestore(); + fixture.device.destroy(); + other.device.destroy(); + } + }); +}); + +function createGraphFixture(): {device: NullDevice; graph: GPUCommandGraph} { + const device = new NullDevice({id: 'batch-hash-index-node-device'}); + Object.defineProperty(device, 'type', {value: 'webgpu'}); + device.limits.maxComputeWorkgroupsPerDimension = 65_535; + device.limits.maxBufferSize = Number.MAX_SAFE_INTEGER; + return {device, graph: new GPUCommandGraph(device, {id: 'batch-hash-index-node-graph'})}; +} + +function createIndexProps( + graph: GPUCommandGraph, + chunkLengths: readonly number[] +): GPUBatchHashIndexProps { + return { + id: 'node-batch-index', + keys: createVector(graph, 'keys', chunkLengths), + tableKeys: createView(graph, 'table-keys', 8), + tableValues: createView(graph, 'table-values', 8), + statistics: createView(graph, 'statistics', 6) + }; +} + +function createVector( + graph: GPUCommandGraph, + id: string, + chunkLengths: readonly number[], + options: {length?: number; byteStride?: number; firstChunk?: GraphDataView<'uint32'>} = {} +): GraphVectorView<'uint32'> { + const sourceLength = chunkLengths.reduce((length, chunkLength) => length + chunkLength, 0); + const length = options.length ?? sourceLength; + const byteStride = options.byteStride ?? Uint32Array.BYTES_PER_ELEMENT; + return new GraphVectorView({ + id, + name: id, + format: 'uint32', + length, + valueLength: length, + stride: 1, + byteStride, + rowByteLength: Uint32Array.BYTES_PER_ELEMENT, + data: chunkLengths.map((chunkLength, chunkIndex) => + chunkIndex === 0 && options.firstChunk + ? options.firstChunk + : createView(graph, `${id}-chunk-${chunkIndex}`, chunkLength, byteStride) + ) + }); +} + +function createView( + graph: GPUCommandGraph, + id: string, + length: number, + byteStride = Uint32Array.BYTES_PER_ELEMENT +): GraphDataView<'uint32'> { + const buffer = graph.createTransientBuffer({ + id, + byteLength: Math.max(length, 1) * byteStride, + usage: Buffer.STORAGE + }); + return graph.createDataView(buffer, {format: 'uint32', length, byteStride}); +} diff --git a/modules/experimental/test/gpu-primitives/gpu-batch-hash-index.spec.ts b/modules/experimental/test/gpu-primitives/gpu-batch-hash-index.spec.ts new file mode 100644 index 0000000000..840db132bf --- /dev/null +++ b/modules/experimental/test/gpu-primitives/gpu-batch-hash-index.spec.ts @@ -0,0 +1,304 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// Copyright (c) vis.gl contributors + +import test from 'test/utils/vitest-tape'; +import {Buffer, type Device} from '@luma.gl/core'; +import { + GPUBatchHashIndex, + GPUCommandGraph, + GPUHashIndexQuery, + GPU_HASH_INDEX_EMPTY_KEY, + GraphVectorView, + type GraphDataView +} from '@luma.gl/experimental'; +import {getWebGPUTestDevice} from '@luma.gl/test-utils'; + +test('GPUBatchHashIndex preserves nullable batches, earliest duplicates, and source offsets', async testCase => { + const device = await getWebGPUTestDevice(); + if (!device) { + testCase.comment('WebGPU is not available'); + testCase.end(); + return; + } + + const result = await runBatchHashIndex(device, { + keys: [[7, 10], [], [10, GPU_HASH_INDEX_EMPTY_KEY, 20, 30, GPU_HASH_INDEX_EMPTY_KEY]], + validity: [[1, 1], [], [1, 1, 0, 1, 0]], + firstValues: [40, 100, 500], + queryKeys: [7, 10, 20, 30, GPU_HASH_INDEX_EMPTY_KEY, 99], + capacity: 8, + encodingCount: 2 + }); + + testCase.deepEqual( + result.values, + [40, 41, GPU_HASH_INDEX_EMPTY_KEY, 503, GPU_HASH_INDEX_EMPTY_KEY, GPU_HASH_INDEX_EMPTY_KEY], + 'generated values retain discontinuous batch offsets and globally earliest duplicates' + ); + testCase.deepEqual(result.found, [1, 1, 0, 1, 0, 0], 'nullable keys are excluded from lookups'); + testCase.deepEqual( + result.buildStatistics.slice(0, 4), + [3, 1, 0, 1], + 'counts valid duplicates and reserved keys while silently skipping invalid/null rows' + ); + testCase.equal( + result.tableKeys.filter(key => key !== GPU_HASH_INDEX_EMPTY_KEY).length, + 3, + 'one shared table slot is retained per distinct valid key' + ); + testCase.ok(result.buildStatistics[4] >= 4, 'cumulative probes include every non-null valid row'); + testCase.ok(result.buildStatistics[5] <= 8, 'all chunks obey the common probe bound'); + testCase.deepEqual(result.queryStatistics.slice(0, 2), [3, 3], 're-encoding resets diagnostics'); + testCase.end(); +}); + +test('GPUBatchHashIndex resolves explicit payloads across offset chunk views', async testCase => { + const device = await getWebGPUTestDevice(); + if (!device) { + testCase.comment('WebGPU is not available'); + testCase.end(); + return; + } + + const result = await runBatchHashIndex(device, { + keys: [[5, 8], [], [8, 13, 21]], + values: [[500, 800], [], [801, 1300, 2100]], + queryKeys: [5, 8, 13, 21], + capacity: 8, + sourceByteOffset: 8 + }); + + testCase.deepEqual( + result.values, + [500, 800, 1300, 2100], + 'later duplicate chunks cannot overwrite the globally first explicit payload' + ); + testCase.deepEqual(result.found, [1, 1, 1, 1], 'finds keys originating in distinct GPU buffers'); + testCase.deepEqual( + result.buildStatistics.slice(0, 4), + [4, 1, 0, 0], + 'accumulates chunk statistics' + ); + testCase.end(); +}); + +test('GPUBatchHashIndex accumulates bounded overflow without resetting between chunks', async testCase => { + const device = await getWebGPUTestDevice(); + if (!device) { + testCase.comment('WebGPU is not available'); + testCase.end(); + return; + } + + const result = await runBatchHashIndex(device, { + keys: [[1, 2], [], [3, 4]], + firstValues: [10, 99, 30], + queryKeys: [1, 2, 3, 4], + capacity: 2 + }); + + testCase.deepEqual( + result.buildStatistics.slice(0, 4), + [2, 0, 2, 0], + 'later distinct keys report fixed-capacity overflow without clearing earlier rows' + ); + testCase.deepEqual(result.found, [1, 1, 0, 0], 'only first-batch keys remain available'); + testCase.deepEqual( + result.values, + [10, 11, GPU_HASH_INDEX_EMPTY_KEY, GPU_HASH_INDEX_EMPTY_KEY], + 'surviving keys preserve original first-source offsets' + ); + testCase.end(); +}); + +test('GPUBatchHashIndex clears zero-chunk and empty-chunk index topologies', async testCase => { + const device = await getWebGPUTestDevice(); + if (!device) { + testCase.comment('WebGPU is not available'); + testCase.end(); + return; + } + + for (const keys of [[], [[], []]] as readonly (readonly number[])[][]) { + const result = await runBatchHashIndex(device, {keys, queryKeys: [1], capacity: 4}); + testCase.deepEqual( + result.tableKeys, + [ + GPU_HASH_INDEX_EMPTY_KEY, + GPU_HASH_INDEX_EMPTY_KEY, + GPU_HASH_INDEX_EMPTY_KEY, + GPU_HASH_INDEX_EMPTY_KEY + ], + 'empty source topologies initialize every shared table slot' + ); + testCase.deepEqual( + result.buildStatistics, + [0, 0, 0, 0, 0, 0], + 'empty chunks add no statistics' + ); + testCase.deepEqual(result.found, [0], 'empty tables do not publish matches'); + } + testCase.end(); +}); + +type BatchHashIndexFixture = { + keys: readonly (readonly number[])[]; + values?: readonly (readonly number[])[]; + validity?: readonly (readonly number[])[]; + firstValues?: readonly number[]; + queryKeys: readonly number[]; + capacity: number; + encodingCount?: number; + sourceByteOffset?: number; +}; + +async function runBatchHashIndex(device: Device, fixture: BatchHashIndexFixture) { + const graph = new GPUCommandGraph(device, {id: 'batch-hash-index-browser'}); + const resources: Buffer[] = []; + const keys = createImportedVector( + graph, + device, + resources, + 'input-keys', + fixture.keys, + fixture.sourceByteOffset + ); + const values = fixture.values + ? createImportedVector( + graph, + device, + resources, + 'input-values', + fixture.values, + fixture.sourceByteOffset + ) + : undefined; + const validity = fixture.validity + ? createImportedVector( + graph, + device, + resources, + 'input-validity', + fixture.validity, + fixture.sourceByteOffset + ) + : undefined; + const queryKeys = createImportedView(graph, device, resources, 'query-keys', fixture.queryKeys); + const tableKeys = createOutputView(graph, device, resources, 'table-keys', fixture.capacity); + const tableValues = createOutputView(graph, device, resources, 'table-values', fixture.capacity); + const buildStatistics = createOutputView(graph, device, resources, 'build-statistics', 6); + const outputValues = createOutputView(graph, device, resources, 'query-values', queryKeys.length); + const found = createOutputView(graph, device, resources, 'query-found', queryKeys.length); + const probes = createOutputView(graph, device, resources, 'query-probes', queryKeys.length); + const queryStatistics = createOutputView(graph, device, resources, 'query-statistics', 4); + + const index = new GPUBatchHashIndex({ + id: 'browser-batch-index', + keys, + ...(values ? {values} : {}), + ...(validity ? {validity} : {}), + ...(fixture.firstValues ? {firstValues: fixture.firstValues} : {}), + tableKeys: tableKeys.view, + tableValues: tableValues.view, + statistics: buildStatistics.view + }); + index.addToGraph(graph); + new GPUHashIndexQuery({ + id: 'browser-batch-query', + index, + keys: queryKeys, + values: outputValues.view, + found: found.view, + probes: probes.view, + statistics: queryStatistics.view + }).addToGraph(graph); + + const compiled = graph.compile(); + try { + for (let encoding = 0; encoding < (fixture.encodingCount ?? 1); encoding++) { + const commandEncoder = device.createCommandEncoder({id: `batch-hash-index-${encoding}`}); + compiled.encode(commandEncoder, {parameters: undefined}); + device.submit(commandEncoder.finish()); + } + return { + tableKeys: await readUint32(tableKeys.buffer, fixture.capacity), + buildStatistics: await readUint32(buildStatistics.buffer, 6), + values: await readUint32(outputValues.buffer, queryKeys.length), + found: await readUint32(found.buffer, queryKeys.length), + probes: await readUint32(probes.buffer, queryKeys.length), + queryStatistics: await readUint32(queryStatistics.buffer, 4) + }; + } finally { + compiled.destroy(); + for (const resource of resources) resource.destroy(); + } +} + +function createImportedVector( + graph: GPUCommandGraph, + device: Device, + resources: Buffer[], + id: string, + chunks: readonly (readonly number[])[], + byteOffset = 0 +): GraphVectorView<'uint32'> { + const length = chunks.reduce((total, chunk) => total + chunk.length, 0); + return new GraphVectorView({ + id, + name: id, + format: 'uint32', + length, + valueLength: length, + stride: 1, + byteStride: Uint32Array.BYTES_PER_ELEMENT, + rowByteLength: Uint32Array.BYTES_PER_ELEMENT, + data: chunks.map((chunk, chunkIndex) => + createImportedView(graph, device, resources, `${id}-chunk-${chunkIndex}`, chunk, byteOffset) + ) + }); +} + +function createImportedView( + graph: GPUCommandGraph, + device: Device, + resources: Buffer[], + id: string, + values: readonly number[], + byteOffset = 0 +): GraphDataView<'uint32'> { + const prefixLength = byteOffset / Uint32Array.BYTES_PER_ELEMENT; + const data = new Uint32Array(prefixLength + Math.max(values.length, 1)); + data.set(values, prefixLength); + const buffer = device.createBuffer({data, usage: Buffer.STORAGE | Buffer.COPY_DST}); + resources.push(buffer); + const handle = graph.importBuffer( + {id, byteLength: buffer.byteLength, usage: buffer.usage}, + buffer + ); + return graph.createDataView(handle, {format: 'uint32', length: values.length, byteOffset}); +} + +function createOutputView( + graph: GPUCommandGraph, + device: Device, + resources: Buffer[], + id: string, + length: number +): {buffer: Buffer; view: GraphDataView<'uint32'>} { + const buffer = device.createBuffer({ + byteLength: Math.max(length, 1) * Uint32Array.BYTES_PER_ELEMENT, + usage: Buffer.STORAGE | Buffer.COPY_SRC + }); + resources.push(buffer); + const handle = graph.importBuffer( + {id, byteLength: buffer.byteLength, usage: buffer.usage}, + buffer + ); + return {buffer, view: graph.createDataView(handle, {format: 'uint32', length})}; +} + +async function readUint32(buffer: Buffer, length: number): Promise { + const bytes = await buffer.readAsync(); + return Array.from(new Uint32Array(bytes.buffer, bytes.byteOffset, length)); +} diff --git a/modules/experimental/test/index.ts b/modules/experimental/test/index.ts index 216c030db0..23e1213fbe 100644 --- a/modules/experimental/test/index.ts +++ b/modules/experimental/test/index.ts @@ -39,6 +39,7 @@ import './ludf/lu-derived-columns.spec'; import './ludf/lu-group-aggregation.spec'; import './ludf/lu-global-aggregation.spec'; import './ludf/lu-sort.spec'; +import './ludf/lu-join.spec'; import './luxfilter'; import './luproj/luproj.spec'; import './luproj/projection-benchmark.spec'; diff --git a/modules/experimental/test/ludf/lu-join.node.spec.ts b/modules/experimental/test/ludf/lu-join.node.spec.ts new file mode 100644 index 0000000000..c537c9fd24 --- /dev/null +++ b/modules/experimental/test/ludf/lu-join.node.spec.ts @@ -0,0 +1,403 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// Copyright (c) vis.gl contributors + +import {Buffer} from '@luma.gl/core'; +import { + column, + CompiledLuDataFrameJoin, + CompiledLuDataFrameLookup, + literal, + LuDataFrame, + LuDataFrameJoinQuery, + LuDataFrameLookupQuery, + parameter, + type LuDataFrameJoinOptions, + type LuDataFrameLookupOptions +} from '@luma.gl/experimental/ludf'; +import { + GPUConstant, + GPUData, + GPURecordBatch, + GPUTable, + type GPUField, + type GPURecordBatchSourceInfo, + type GPUTypeMap +} from '@luma.gl/tables'; +import {NullDevice} from '@luma.gl/test-utils'; +import {describe, expect, expectTypeOf, test, vi} from 'vitest'; + +type LeftJoinColumns = { + key: 'uint32'; + amount: 'float32'; + signed: 'sint32'; +}; + +type RightJoinColumns = { + lookupKey: 'uint32'; + value: 'float32'; + signed: 'sint32'; +}; + +type JoinSourceFixture = { + device: NullDevice; + table: GPUTable; + buffers: Buffer[]; +}; + +const LEFT_FIELDS: GPUField[] = [ + {name: 'key', format: 'uint32', nullable: false}, + {name: 'amount', format: 'float32', nullable: false}, + {name: 'signed', format: 'sint32', nullable: false} +]; + +const RIGHT_FIELDS: GPUField[] = [ + {name: 'lookupKey', format: 'uint32', nullable: false}, + {name: 'value', format: 'float32', nullable: false}, + {name: 'signed', format: 'sint32', nullable: false} +]; + +describe('LuDataFrame immutable unique-right joins and bounded lookups', () => { + test('plans mismatched preserved batch topologies without GPU work or source retention', () => { + const leftFixture = createJoinSourceFixture('left', [3, 0, 5], LEFT_FIELDS, [100, 400, 800]); + const rightFixture = createJoinSourceFixture('right', [2, 0, 3], RIGHT_FIELDS, [500, 750, 900]); + const left = new LuDataFrame({table: leftFixture.table, ownership: 'owned'}); + const right = new LuDataFrame({table: rightFixture.table, ownership: 'owned'}); + const createLeftBuffer = vi.spyOn(leftFixture.device, 'createBuffer'); + const createRightBuffer = vi.spyOn(rightFixture.device, 'createBuffer'); + const submitLeft = vi.spyOn(leftFixture.device, 'submit'); + const submitRight = vi.spyOn(rightFixture.device, 'submit'); + const selectLeft = vi.spyOn(leftFixture.table, 'select'); + const selectRight = vi.spyOn(rightFixture.table, 'select'); + + const joined = left.innerJoin(right, {leftOn: 'key', rightOn: 'lookupKey'}); + const lookup = left.lookup(right, {leftOn: 'key', rightOn: 'lookupKey'}); + + expect(joined).toBeInstanceOf(LuDataFrameJoinQuery); + expect(lookup).toBeInstanceOf(LuDataFrameLookupQuery); + expect(joined.query.source).toBe(left); + expect(joined.right).toBe(right); + expect(joined.options).toEqual({ + leftOn: 'key', + rightOn: 'lookupKey', + indexCapacity: 16, + maxProbeCount: 16 + }); + expect(lookup.options).toEqual(joined.options); + expect(Object.isFrozen(joined)).toBe(true); + expect(Object.isFrozen(joined.options)).toBe(true); + expect(Object.isFrozen(lookup)).toBe(true); + expect(Object.isFrozen(lookup.options)).toBe(true); + expect(left.batches.map(batch => batch.numRows)).toEqual([3, 0, 5]); + expect(right.batches.map(batch => batch.numRows)).toEqual([2, 0, 3]); + expect(left.sourceInfo.map(info => info?.sourceRowIndexOffset)).toEqual([100, 400, 800]); + expect(right.sourceInfo.map(info => info?.sourceRowIndexOffset)).toEqual([500, 750, 900]); + expect(createLeftBuffer).not.toHaveBeenCalled(); + expect(createRightBuffer).not.toHaveBeenCalled(); + expect(submitLeft).not.toHaveBeenCalled(); + expect(submitRight).not.toHaveBeenCalled(); + expect(selectLeft).not.toHaveBeenCalled(); + expect(selectRight).not.toHaveBeenCalled(); + + left.destroy(); + right.destroy(); + expect(leftFixture.buffers.every(buffer => buffer.destroyed)).toBe(true); + expect(rightFixture.buffers.every(buffer => buffer.destroyed)).toBe(true); + + createLeftBuffer.mockRestore(); + createRightBuffer.mockRestore(); + submitLeft.mockRestore(); + submitRight.mockRestore(); + selectLeft.mockRestore(); + selectRight.mockRestore(); + }); + + test('retains precise left/right schemas through filtered, projected, and derived join plans', () => { + const leftFixture = createJoinSourceFixture('left', [3, 0, 5], LEFT_FIELDS); + const rightFixture = createJoinSourceFixture('right', [2, 0, 3], RIGHT_FIELDS); + const left = new LuDataFrame({table: leftFixture.table}); + const right = new LuDataFrame({table: rightFixture.table}); + const filtered = left + .filter(column('amount').greaterThan(parameter('minimumAmount', 5))) + .select(['key', 'amount']); + const joined = filtered.innerJoin(right, {leftOn: 'key', rightOn: 'lookupKey', capacity: 2}); + const derived = left + .withColumn('shiftedKey', column('key').add(literal(1)), {format: 'uint32'}) + .select(['shiftedKey', 'amount']) + .lookup(right, {leftOn: 'shiftedKey', rightOn: 'lookupKey'}); + + expect(joined.query.predicates[0]).toBe(filtered.predicates[0]); + expect(joined.query.selectedColumns).toEqual(['key', 'amount']); + expect(joined.options.capacity).toBe(2); + expect(derived.query.derivedColumns.map(({name}) => name)).toEqual(['shiftedKey']); + expectTypeOf(joined.compile).returns.toEqualTypeOf< + CompiledLuDataFrameJoin<{key: 'uint32'; amount: 'float32'}, RightJoinColumns> + >(); + expectTypeOf(derived.compile).returns.toEqualTypeOf< + CompiledLuDataFrameLookup<{shiftedKey: 'uint32'; amount: 'float32'}, RightJoinColumns> + >(); + + left.destroy(); + right.destroy(); + leftFixture.table.destroy(); + rightFixture.table.destroy(); + }); + + test('clones bounded options and safely clamps default cumulative hash probing', () => { + const leftFixture = createJoinSourceFixture('left', [3], LEFT_FIELDS); + const rightFixture = createJoinSourceFixture('right', [70_000], RIGHT_FIELDS); + const left = new LuDataFrame({table: leftFixture.table}); + const right = new LuDataFrame({table: rightFixture.table}); + const mutableOptions: LuDataFrameJoinOptions<'key', 'lookupKey'> = { + leftOn: 'key', + rightOn: 'lookupKey', + capacity: 0, + indexCapacity: 262_144 + }; + + const joined = left.innerJoin(right, mutableOptions); + expect(joined.options.indexCapacity).toBe(262_144); + expect(joined.options.maxProbeCount).toBe(Math.floor(0xffffffff / 70_000)); + expect(joined.options.capacity).toBe(0); + expect(Object.isFrozen(joined.options)).toBe(true); + + const explicit = left.lookup(right, { + leftOn: 'key', + rightOn: 'lookupKey', + indexCapacity: 131_072, + maxProbeCount: 8 + }); + expect(explicit.options.maxProbeCount).toBe(8); + + left.destroy(); + right.destroy(); + leftFixture.table.destroy(); + rightFixture.table.destroy(); + }); + + test('supports zero-row schema-only sources and independently empty preserved batches', () => { + for (const [leftLengths, rightLengths] of [ + [[], []], + [[0, 0], []], + [[], [0, 0]], + [ + [2, 0], + [0, 3] + ] + ] as const) { + const leftFixture = createJoinSourceFixture('left', leftLengths, LEFT_FIELDS); + const rightFixture = createJoinSourceFixture('right', rightLengths, RIGHT_FIELDS); + const left = new LuDataFrame({table: leftFixture.table}); + const right = new LuDataFrame({table: rightFixture.table}); + + const joined = left.innerJoin(right, {leftOn: 'key', rightOn: 'lookupKey'}); + const lookup = left.lookup(right, {leftOn: 'key', rightOn: 'lookupKey'}); + + expect(joined.query.source.batches.map(batch => batch.numRows)).toEqual(leftLengths); + expect(joined.right.batches.map(batch => batch.numRows)).toEqual(rightLengths); + expect(lookup.options.indexCapacity).toBe(right.numRows === 0 ? 1 : 8); + + left.destroy(); + right.destroy(); + leftFixture.table.destroy(); + rightFixture.table.destroy(); + } + }); + + test('rejects unknown, hidden, signed, floating, or constant keys without GPU allocation', () => { + const leftFixture = createJoinSourceFixture('left', [3], LEFT_FIELDS); + const rightFixture = createJoinSourceFixture('right', [2], RIGHT_FIELDS); + const left = new LuDataFrame({table: leftFixture.table}); + const right = new LuDataFrame({table: rightFixture.table}); + const createBuffer = vi.spyOn(leftFixture.device, 'createBuffer'); + + expect(() => + // @ts-expect-error Join keys must be selected unsigned 32-bit scalar columns. + left.innerJoin(right, {leftOn: 'amount', rightOn: 'lookupKey'}) + ).toThrow(/uint32/i); + expect(() => + // @ts-expect-error Right join keys must be unsigned 32-bit scalar columns. + left.lookup(right, {leftOn: 'key', rightOn: 'signed'}) + ).toThrow(/uint32/i); + expect(() => + // @ts-expect-error Right join keys must exist in the right schema. + left.innerJoin(right, {leftOn: 'key', rightOn: 'missing'}) + ).toThrow(/right|exist/i); + const selected = left.filter(column('amount').greaterThan(literal(0))).select(['amount']); + expect(() => + // @ts-expect-error Projected-away source keys cannot participate in joins. + selected.innerJoin(right, {leftOn: 'key', rightOn: 'lookupKey'}) + ).toThrow(/selected/i); + + const constant = new GPUConstant({format: 'uint32', value: Uint32Array.of(7)}); + const constantLeftTable = new GPUTable({ + batches: leftFixture.table.batches, + constants: {constantKey: constant} + }); + const constantLeft = new LuDataFrame({table: constantLeftTable}); + expect(() => + constantLeft.innerJoin(right, {leftOn: 'constantKey', rightOn: 'lookupKey'}) + ).toThrow(/constant/i); + + const constantRightTable = new GPUTable({ + batches: rightFixture.table.batches, + constants: {constantKey: constant} + }); + const constantRight = new LuDataFrame({table: constantRightTable}); + expect(() => left.lookup(constantRight, {leftOn: 'key', rightOn: 'constantKey'})).toThrow( + /constant/i + ); + expect(createBuffer).not.toHaveBeenCalled(); + + createBuffer.mockRestore(); + constantLeft.destroy(); + constantRight.destroy(); + left.destroy(); + right.destroy(); + constantLeftTable.destroy(); + constantRightTable.destroy(); + }); + + test('rejects incompatible dictionary labels, ordering, and one-sided categorical encoding', () => { + const leftFixture = createJoinSourceFixture('left', [3], LEFT_FIELDS); + const rightFixture = createJoinSourceFixture('right', [2], RIGHT_FIELDS); + const left = new LuDataFrame({ + table: leftFixture.table, + dictionaries: {key: {values: ['economy', 'premium'], ordered: true}} + }); + const compatible = new LuDataFrame({ + table: rightFixture.table, + dictionaries: {lookupKey: {values: ['economy', 'premium'], ordered: true}} + }); + const reordered = new LuDataFrame({ + table: rightFixture.table, + dictionaries: {lookupKey: {values: ['premium', 'economy'], ordered: true}} + }); + const unordered = new LuDataFrame({ + table: rightFixture.table, + dictionaries: {lookupKey: {values: ['economy', 'premium'], ordered: false}} + }); + const raw = new LuDataFrame({table: rightFixture.table}); + + expect(left.innerJoin(compatible, {leftOn: 'key', rightOn: 'lookupKey'})).toBeInstanceOf( + LuDataFrameJoinQuery + ); + expect(() => left.lookup(reordered, {leftOn: 'key', rightOn: 'lookupKey'})).toThrow( + /dictionary|dictionaries|labels/i + ); + expect(() => left.lookup(unordered, {leftOn: 'key', rightOn: 'lookupKey'})).toThrow( + /dictionary|dictionaries|ordering/i + ); + expect(() => left.innerJoin(raw, {leftOn: 'key', rightOn: 'lookupKey'})).toThrow( + /dictionary|dictionaries|both/i + ); + + left.destroy(); + compatible.destroy(); + reordered.destroy(); + unordered.destroy(); + raw.destroy(); + leftFixture.table.destroy(); + rightFixture.table.destroy(); + }); + + test('rejects unsafe index capacities, unbounded probes, and invalid per-batch output limits', () => { + const leftFixture = createJoinSourceFixture('left', [3], LEFT_FIELDS); + const rightFixture = createJoinSourceFixture('right', [2], RIGHT_FIELDS); + const left = new LuDataFrame({table: leftFixture.table}); + const right = new LuDataFrame({table: rightFixture.table}); + const createBuffer = vi.spyOn(leftFixture.device, 'createBuffer'); + + for (const indexCapacity of [0, -1, 3, 1.5, Number.NaN, 0x1_0000_0000]) { + expect(() => + left.innerJoin(right, {leftOn: 'key', rightOn: 'lookupKey', indexCapacity}) + ).toThrow(/capacity|power/i); + } + for (const maxProbeCount of [0, -1, 1.5, Number.NaN, 9]) { + expect(() => + left.lookup(right, {leftOn: 'key', rightOn: 'lookupKey', indexCapacity: 8, maxProbeCount}) + ).toThrow(/probe/i); + } + for (const capacity of [-1, 1.5, Number.NaN, 0x1_0000_0000]) { + expect(() => left.innerJoin(right, {leftOn: 'key', rightOn: 'lookupKey', capacity})).toThrow( + /capacity|uint32/i + ); + } + const lookupWithCapacity = { + leftOn: 'key', + rightOn: 'lookupKey', + capacity: 1 + } as LuDataFrameLookupOptions<'key', 'lookupKey'>; + expect(() => left.lookup(right, lookupWithCapacity)).toThrow(/capacity/i); + expect(createBuffer).not.toHaveBeenCalled(); + + createBuffer.mockRestore(); + left.destroy(); + right.destroy(); + leftFixture.table.destroy(); + rightFixture.table.destroy(); + }); + + test('rejects new join and lookup plans after the left dataframe was explicitly destroyed', () => { + const leftFixture = createJoinSourceFixture('left', [3], LEFT_FIELDS); + const rightFixture = createJoinSourceFixture('right', [2], RIGHT_FIELDS); + const left = new LuDataFrame({table: leftFixture.table}); + const right = new LuDataFrame({table: rightFixture.table}); + left.destroy(); + + expect(() => left.innerJoin(right, {leftOn: 'key', rightOn: 'lookupKey'})).toThrow( + /destroyed/i + ); + expect(() => left.lookup(right, {leftOn: 'key', rightOn: 'lookupKey'})).toThrow(/destroyed/i); + + right.destroy(); + leftFixture.table.destroy(); + rightFixture.table.destroy(); + }); +}); + +function createJoinSourceFixture( + side: 'left' | 'right', + batchLengths: readonly number[], + fields: readonly GPUField[], + sourceOffsets?: readonly number[] +): JoinSourceFixture { + const device = new NullDevice({id: `ludf-${side}-join-node-device`}); + const buffers: Buffer[] = []; + let defaultOffset = side === 'left' ? 100 : 500; + const batches = batchLengths.map((length, sourceBatchIndex) => { + const sourceInfo: GPURecordBatchSourceInfo = { + sourceBatchIndex, + sourceRowIndexOffset: sourceOffsets?.[sourceBatchIndex] ?? defaultOffset, + sourceRowCount: length + }; + defaultOffset += length; + const gpuData: Record = {}; + for (const field of fields) { + const format = field.format; + if (format !== 'float32' && format !== 'sint32' && format !== 'uint32') { + throw new Error('Join fixtures require scalar formats'); + } + const buffer = device.createBuffer({ + byteLength: Math.max(length, 1) * Uint32Array.BYTES_PER_ELEMENT, + usage: Buffer.STORAGE | Buffer.COPY_SRC | Buffer.COPY_DST + }); + buffers.push(buffer); + gpuData[field.name] = new GPUData({buffer, format, length, ownsBuffer: true}); + } + return new GPURecordBatch({ + gpuData, + fields: [...fields], + numRows: length, + sourceInfo + }); + }); + + const table = + batches.length > 0 + ? new GPUTable({batches}) + : new GPUTable({ + schema: {fields: [...fields], metadata: new Map()} + }); + return {device, table, buffers}; +} diff --git a/modules/experimental/test/ludf/lu-join.spec.ts b/modules/experimental/test/ludf/lu-join.spec.ts new file mode 100644 index 0000000000..87800e1012 --- /dev/null +++ b/modules/experimental/test/ludf/lu-join.spec.ts @@ -0,0 +1,555 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// Copyright (c) vis.gl contributors + +import {Buffer, type Device} from '@luma.gl/core'; +import {GPUCommandGraph} from '@luma.gl/experimental'; +import { + LuDataFrame, + column, + parameter, + type LuDataFrameQueryParameters +} from '@luma.gl/experimental/ludf'; +import {GPUData, GPURecordBatch, GPUTable, GPUVector} from '@luma.gl/tables'; +import {getWebGPUTestDevice} from '@luma.gl/test-utils'; +import test from 'test/utils/vitest-tape'; +import {vi} from 'vitest'; + +const MISSING_JOIN_ROW = 0xffffffff; + +type LuJoinLeftSchema = {key: 'uint32'; fare: 'float32'}; +type LuJoinRightSchema = {lookupKey: 'uint32'; weight: 'float32'}; + +type LuJoinFixture = { + left: LuDataFrame; + right: LuDataFrame; + leftBuffers: Buffer[]; + rightBuffers: Buffer[]; +}; + +type LuJoinFixtureOptions = { + duplicateRight?: boolean; + invalidRight?: boolean; + noRightBatches?: boolean; +}; + +test('LuDataFrame inner joins preserve mismatched nullable batches, stable row identities, and both source leases', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuJoinFixture(device); + const createBufferSpy = vi.spyOn(device, 'createBuffer'); + const submitSpy = vi.spyOn(device, 'submit'); + const query = fixture.left.innerJoin(fixture.right, {leftOn: 'key', rightOn: 'lookupKey'}); + + testContext.equal(createBufferSpy.mock.calls.length, 0, 'join planning allocates no GPU buffers'); + testContext.equal(submitSpy.mock.calls.length, 0, 'join planning submits no GPU work'); + + const graph = new GPUCommandGraph(device, { + id: 'ludf-mismatched-batch-inner-join' + }); + const compiled = query.compile(graph); + + try { + testContext.deepEqual( + compiled.table.batches.map(batch => batch.numRows), + [3, 0, 5], + 'joined output retains the original left record batches' + ); + testContext.deepEqual( + compiled.rightTable.batches.map(batch => batch.numRows), + [2, 0, 3], + 'the retained right side preserves its independent source topology' + ); + + fixture.left.destroy(); + fixture.right.destroy(); + testContext.ok( + [...fixture.leftBuffers, ...fixture.rightBuffers].every(buffer => !buffer.destroyed), + 'compiled joins retain both owned source leases' + ); + + const commandEncoder = device.createCommandEncoder({id: 'ludf-inner-join-encode'}); + compiled.encode(commandEncoder); + testContext.equal( + submitSpy.mock.calls.length, + 0, + 'encoding leaves submission application-owned' + ); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + await readLuJoinChunks(compiled.requiredCounts), + [[2], [0], [3]], + 'required matches are counted independently per left batch' + ); + testContext.deepEqual( + await readLuJoinChunks(compiled.selectedCounts), + [[2], [0], [3]], + 'published counts remain coherent with inherited stable row indices' + ); + testContext.deepEqual(await readLuJoinChunks(compiled.overflows), [[0], [0], [0]]); + testContext.deepEqual( + await readLuJoinPublishedRows(compiled.rowIndices, compiled.selectedCounts), + [[100, 102], [], [800, 802, 803]], + 'left source identifiers retain discontinuous sourceInfo offsets and stable match order' + ); + testContext.deepEqual( + await readLuJoinPublishedRows(compiled.rightRowIndices, compiled.selectedCounts), + [[501, 500], [], [900, 901, 500]], + 'right identities resolve across multiple batches without concatenating or repacking' + ); + testContext.deepEqual( + (await readLuJoinChunks(compiled.indexStatistics))[0].slice(0, 4), + [4, 0, 0, 0], + 'ordinary nullable right rows are skipped without becoming reserved-key violations' + ); + testContext.deepEqual(await readLuJoinChunks(compiled.contractViolation), [[0]]); + + compiled.destroy(); + testContext.ok( + [...fixture.leftBuffers, ...fixture.rightBuffers].every(buffer => buffer.destroyed), + 'both owned sources are destroyed only after the compiled join is released' + ); + } finally { + compiled.destroy(); + fixture.left.destroy(); + fixture.right.destroy(); + createBufferSpy.mockRestore(); + submitSpy.mockRestore(); + } + + testContext.end(); +}); + +test('LuDataFrame bounded lookups keep source-aligned matches, nullable keys, and missing markers on the GPU', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuJoinFixture(device); + const graph = new GPUCommandGraph(device, { + id: 'ludf-source-aligned-lookup' + }); + const compiled = fixture.left + .lookup(fixture.right, {leftOn: 'key', rightOn: 'lookupKey'}) + .compile(graph); + + try { + const commandEncoder = device.createCommandEncoder({id: 'ludf-lookup-encode'}); + compiled.encode(commandEncoder); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + await readLuJoinChunks(compiled.matchMask), + [[1, 0, 1], [], [1, 0, 1, 1, 0]], + 'missing and explicitly null left keys never report a match' + ); + testContext.deepEqual( + await readLuJoinChunks(compiled.rightRowIndices), + [[501, MISSING_JOIN_ROW, 500], [], [900, MISSING_JOIN_ROW, 901, 500, MISSING_JOIN_ROW]], + 'bounded lookups retain source-aligned stable right identities and explicit missing markers' + ); + testContext.deepEqual( + compiled.probeCounts.data.map(chunk => chunk.length), + [3, 0, 5], + 'probe counts preserve the original left chunk topology' + ); + testContext.deepEqual( + compiled.lookupStatistics.data.map(chunk => chunk.length), + [4, 4, 4], + 'lookup diagnostics retain one four-word statistics block per left batch' + ); + testContext.deepEqual(await readLuJoinChunks(compiled.contractViolation), [[0]]); + } finally { + compiled.destroy(); + fixture.left.destroy(); + fixture.right.destroy(); + } + + testContext.end(); +}); + +test('LuDataFrame reuses filtered joins across two ordered encodings without reading source rows', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuJoinFixture(device); + const graph = new GPUCommandGraph(device, { + id: 'ludf-parameterized-inner-join' + }); + const compiled = fixture.left + .filter(column('fare').greaterThan(parameter('minimumFare', 0))) + .innerJoin(fixture.right, {leftOn: 'key', rightOn: 'lookupKey'}) + .compile(graph); + const firstCounts = compiled.selectedCounts.data.map((_, batchIndex) => + device.createBuffer({ + id: `ludf-first-join-count-${batchIndex}`, + byteLength: Uint32Array.BYTES_PER_ELEMENT, + usage: Buffer.COPY_SRC | Buffer.COPY_DST + }) + ); + + try { + const commandEncoder = device.createCommandEncoder({id: 'ludf-join-two-encodes'}); + compiled.encode(commandEncoder, {minimumFare: 10}); + for (const [batchIndex, count] of compiled.selectedCounts.data.entries()) { + commandEncoder.copyBufferToBuffer({ + sourceBuffer: getLuJoinBuffer(count), + destinationBuffer: firstCounts[batchIndex], + size: Uint32Array.BYTES_PER_ELEMENT + }); + } + compiled.encode(commandEncoder, {minimumFare: 35}); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + await Promise.all(firstCounts.map(buffer => readLuJoinBuffer(buffer, 1))), + [[2], [0], [2]], + 'the first encoder-ordered parameter update retains matching filtered rows per batch' + ); + testContext.deepEqual( + await readLuJoinChunks(compiled.selectedCounts), + [[0], [0], [1]], + 'the second update reuses the same index and graph with a stricter source predicate' + ); + testContext.deepEqual( + await readLuJoinPublishedRows(compiled.rowIndices, compiled.selectedCounts), + [[], [], [802]], + 'reused joins publish only the final matching stable left row' + ); + testContext.deepEqual( + await readLuJoinPublishedRows(compiled.rightRowIndices, compiled.selectedCounts), + [[], [], [901]], + 'reused joins retain the corresponding stable right source identity' + ); + } finally { + for (const buffer of firstCounts) buffer.destroy(); + compiled.destroy(); + fixture.left.destroy(); + fixture.right.destroy(); + } + + testContext.end(); +}); + +test('LuDataFrame reports bounded join overflow and suppresses duplicate, reserved-key, and incomplete right indexes', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const boundedFixture = createLuJoinFixture(device); + const duplicateFixture = createLuJoinFixture(device, {duplicateRight: true}); + const invalidFixture = createLuJoinFixture(device, {invalidRight: true}); + const incompleteFixture = createLuJoinFixture(device); + + const bounded = boundedFixture.left + .innerJoin(boundedFixture.right, {leftOn: 'key', rightOn: 'lookupKey', capacity: 1}) + .compile(new GPUCommandGraph(device, {id: 'ludf-bounded-join'})); + const duplicate = duplicateFixture.left + .innerJoin(duplicateFixture.right, {leftOn: 'key', rightOn: 'lookupKey'}) + .compile(new GPUCommandGraph(device, {id: 'ludf-duplicate-join'})); + const invalid = invalidFixture.left + .innerJoin(invalidFixture.right, {leftOn: 'key', rightOn: 'lookupKey'}) + .compile(new GPUCommandGraph(device, {id: 'ludf-invalid-join'})); + const incomplete = incompleteFixture.left + .innerJoin(incompleteFixture.right, { + leftOn: 'key', + rightOn: 'lookupKey', + indexCapacity: 2, + maxProbeCount: 2 + }) + .compile(new GPUCommandGraph(device, {id: 'ludf-incomplete-join'})); + + try { + const commandEncoder = device.createCommandEncoder({id: 'ludf-join-contracts'}); + bounded.encode(commandEncoder); + duplicate.encode(commandEncoder); + invalid.encode(commandEncoder); + incomplete.encode(commandEncoder); + device.submit(commandEncoder.finish()); + + testContext.deepEqual(await readLuJoinChunks(bounded.requiredCounts), [[2], [0], [3]]); + testContext.deepEqual(await readLuJoinChunks(bounded.selectedCounts), [[1], [0], [1]]); + testContext.deepEqual(await readLuJoinChunks(bounded.overflows), [[1], [0], [1]]); + testContext.deepEqual( + await readLuJoinPublishedRows(bounded.rowIndices, bounded.selectedCounts), + [[100], [], [800]], + 'bounded batches publish their earliest stable matching left row' + ); + testContext.deepEqual( + await readLuJoinPublishedRows(bounded.rightRowIndices, bounded.selectedCounts), + [[501], [], [900]], + 'bounded partner outputs remain aligned with the published left rows' + ); + + for (const [compiled, statisticIndex, label] of [ + [duplicate, 1, 'duplicate right keys'], + [invalid, 3, 'reserved valid right keys'], + [incomplete, 2, 'incomplete right indexes'] + ] as const) { + testContext.ok( + (await readLuJoinChunks(compiled.indexStatistics))[0][statisticIndex] > 0, + `${label} remain visible in GPU-resident index diagnostics` + ); + testContext.deepEqual( + await readLuJoinChunks(compiled.contractViolation), + [[1]], + `${label} raise an explicit GPU contract violation` + ); + testContext.deepEqual( + await readLuJoinChunks(compiled.selectedCounts), + [[0], [0], [0]], + `${label} never publish potentially incorrect join matches` + ); + } + } finally { + bounded.destroy(); + duplicate.destroy(); + invalid.destroy(); + incomplete.destroy(); + for (const fixture of [boundedFixture, duplicateFixture, invalidFixture, incompleteFixture]) { + fixture.left.destroy(); + fixture.right.destroy(); + } + } + + testContext.end(); +}); + +test('LuDataFrame joins and lookups preserve empty left chunks against a schema-only right source', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuJoinFixture(device, {noRightBatches: true}); + const join = fixture.left + .innerJoin(fixture.right, {leftOn: 'key', rightOn: 'lookupKey'}) + .compile(new GPUCommandGraph(device, {id: 'ludf-empty-right-join'})); + const lookup = fixture.left + .lookup(fixture.right, {leftOn: 'key', rightOn: 'lookupKey'}) + .compile(new GPUCommandGraph(device, {id: 'ludf-empty-right-lookup'})); + + try { + const commandEncoder = device.createCommandEncoder({id: 'ludf-empty-right-queries'}); + join.encode(commandEncoder); + lookup.encode(commandEncoder); + device.submit(commandEncoder.finish()); + + testContext.deepEqual(join.rightTable.batches, [], 'no right batches are fabricated'); + testContext.deepEqual(await readLuJoinChunks(join.selectedCounts), [[0], [0], [0]]); + testContext.deepEqual(await readLuJoinChunks(join.requiredCounts), [[0], [0], [0]]); + testContext.deepEqual( + await readLuJoinChunks(lookup.matchMask), + [[0, 0, 0], [], [0, 0, 0, 0, 0]], + 'empty right indexes leave every preserved left row unmatched' + ); + testContext.deepEqual( + (await readLuJoinChunks(join.indexStatistics))[0].slice(0, 4), + [0, 0, 0, 0], + 'empty right sources clear index statistics without reading data back' + ); + } finally { + join.destroy(); + lookup.destroy(); + fixture.left.destroy(); + fixture.right.destroy(); + } + + testContext.end(); +}); + +function createLuJoinFixture(device: Device, options: LuJoinFixtureOptions = {}): LuJoinFixture { + const leftBuffers: Buffer[] = []; + const rightBuffers: Buffer[] = []; + const leftKeys = [ + Uint32Array.from([70, 1, 20]), + new Uint32Array(0), + Uint32Array.from([40, 70, 90, 20, MISSING_JOIN_ROW]) + ]; + const fares = [ + Float32Array.from([30, 2, 12]), + new Float32Array(0), + Float32Array.from([6, 40, 50, 11, 99]) + ]; + const leftValidity = [ + Uint32Array.from([1, 1, 1]), + new Uint32Array(0), + Uint32Array.from([1, 0, 1, 1, 0]) + ]; + const leftOffsets = [100, 400, 800]; + const leftValidityData: GPUData<'uint32'>[] = []; + const leftBatches = leftKeys.map((keys, batchIndex) => { + leftValidityData.push( + createLuJoinData(device, leftBuffers, leftValidity[batchIndex], 'uint32') + ); + return new GPURecordBatch({ + gpuData: { + key: createLuJoinData(device, leftBuffers, keys, 'uint32'), + fare: createLuJoinData(device, leftBuffers, fares[batchIndex], 'float32') + }, + fields: [ + {name: 'key', format: 'uint32', nullable: true}, + {name: 'fare', format: 'float32', nullable: false} + ], + sourceInfo: { + sourceBatchIndex: batchIndex, + sourceRowIndexOffset: leftOffsets[batchIndex], + sourceRowCount: keys.length + } + }); + }); + + const rightKeys = options.noRightBatches + ? [] + : [ + Uint32Array.from([20, 70]), + new Uint32Array(0), + Uint32Array.from([options.duplicateRight ? 20 : 40, 90, MISSING_JOIN_ROW]) + ]; + const rightValidity = options.noRightBatches + ? [] + : [ + Uint32Array.from([1, 1]), + new Uint32Array(0), + Uint32Array.from([1, 1, options.invalidRight ? 1 : 0]) + ]; + const rightOffsets = [500, 750, 900]; + const rightValidityData: GPUData<'uint32'>[] = []; + const rightBatches = rightKeys.map((keys, batchIndex) => { + rightValidityData.push( + createLuJoinData(device, rightBuffers, rightValidity[batchIndex], 'uint32') + ); + return new GPURecordBatch({ + gpuData: { + lookupKey: createLuJoinData(device, rightBuffers, keys, 'uint32'), + weight: createLuJoinData( + device, + rightBuffers, + Float32Array.from(keys, (_, index) => batchIndex * 10 + index), + 'float32' + ) + }, + fields: [ + {name: 'lookupKey', format: 'uint32', nullable: true}, + {name: 'weight', format: 'float32', nullable: false} + ], + sourceInfo: { + sourceBatchIndex: batchIndex + 10, + sourceRowIndexOffset: rightOffsets[batchIndex], + sourceRowCount: keys.length + } + }); + }); + + const rightTable = + rightBatches.length > 0 + ? new GPUTable({batches: rightBatches}) + : new GPUTable({ + schema: { + fields: [ + {name: 'lookupKey', format: 'uint32', nullable: true}, + {name: 'weight', format: 'float32', nullable: false} + ] + }, + bufferLayout: [ + {name: 'lookupKey', format: 'uint32', byteStride: 4}, + {name: 'weight', format: 'float32', byteStride: 4} + ] + }); + + return { + left: new LuDataFrame({ + table: new GPUTable({batches: leftBatches}), + validity: { + key: new GPUVector<'uint32'>({ + type: 'data', + name: 'ludf-left-join-validity', + format: 'uint32', + data: leftValidityData, + ownsData: true + }) + }, + ownership: 'owned' + }), + right: new LuDataFrame({ + table: rightTable, + ...(rightValidityData.length > 0 + ? { + validity: { + lookupKey: new GPUVector<'uint32'>({ + type: 'data', + name: 'ludf-right-join-validity', + format: 'uint32', + data: rightValidityData, + ownsData: true + }) + } + } + : {}), + ownership: 'owned' + }), + leftBuffers, + rightBuffers + }; +} + +function createLuJoinData( + device: Device, + sourceBuffers: Buffer[], + values: Float32Array | Uint32Array, + format: Format +): GPUData { + const buffer = device.createBuffer({ + byteLength: Math.max(values.byteLength, Uint32Array.BYTES_PER_ELEMENT), + usage: Buffer.STORAGE | Buffer.VERTEX | Buffer.COPY_SRC | Buffer.COPY_DST, + ...(values.byteLength > 0 ? {data: values} : {}) + }); + sourceBuffers.push(buffer); + return new GPUData({buffer, format, length: values.length, ownsBuffer: true}); +} + +function getLuJoinBuffer(data: GPUData): Buffer { + return data.buffer instanceof Buffer ? data.buffer : data.buffer.buffer; +} + +async function readLuJoinBuffer(buffer: Buffer, length: number): Promise { + if (length === 0) return []; + const bytes = await buffer.readAsync(0, length * Uint32Array.BYTES_PER_ELEMENT); + return Array.from(new Uint32Array(bytes.buffer, bytes.byteOffset, length)); +} + +async function readLuJoinChunks(vector: GPUVector<'uint32'>): Promise { + return Promise.all( + vector.data.map(chunk => readLuJoinBuffer(getLuJoinBuffer(chunk), chunk.length)) + ); +} + +async function readLuJoinPublishedRows( + rows: GPUVector<'uint32'>, + counts: GPUVector<'uint32'> +): Promise { + const publishedCounts = await readLuJoinChunks(counts); + return Promise.all( + rows.data.map((chunk, batchIndex) => + readLuJoinBuffer(getLuJoinBuffer(chunk), publishedCounts[batchIndex][0] ?? 0) + ) + ); +}