diff --git a/modules/experimental/src/ludf/index.ts b/modules/experimental/src/ludf/index.ts index a765455f86..9b5b7bde68 100644 --- a/modules/experimental/src/ludf/index.ts +++ b/modules/experimental/src/ludf/index.ts @@ -19,6 +19,16 @@ export type { LuDataFrameDerivedColumnFormatForExpression, LuDataFrameDerivedColumnOptions } from './lu-data-frame-query'; +export {LuDataFrameGroupByQuery, LuDataFrameGroupedAggregationQuery} from './lu-group-by-query'; +export type { + LuDataFrameAggregationDefinition, + LuDataFrameAggregationDefinitions, + LuDataFrameAggregationOperation, + LuDataFrameAggregationValue, + LuDataFrameColumnNamesOfFormat, + LuDataFrameGroupByOptions, + LuDataFrameGroupedAggregationResult +} from './lu-group-by-query'; export {and, column, literal, LuExpression, not, or, parameter} from './lu-expression'; export type { LuExpressionBinaryOperator, @@ -28,3 +38,4 @@ export type { } from './lu-expression'; export {CompiledLuDataFrameQuery} from './lu-query-compiler'; export type {LuDataFrameQueryParameters} from './lu-query-compiler'; +export {CompiledLuDataFrameGroupedAggregation} from './lu-group-aggregation-compiler'; diff --git a/modules/experimental/src/ludf/lu-data-frame-query.ts b/modules/experimental/src/ludf/lu-data-frame-query.ts index 2257ac01e8..7c7162cb50 100644 --- a/modules/experimental/src/ludf/lu-data-frame-query.ts +++ b/modules/experimental/src/ludf/lu-data-frame-query.ts @@ -6,6 +6,11 @@ import type {GPUTypeMap} from '@luma.gl/tables'; import type {GPUCommandGraph} from '../gpu-primitives/gpu-command-graph'; import type {LuDataFrame} from './lu-data-frame'; import {getLuExpressionColumnNames, LuExpression, type LuExpressionNode} from './lu-expression'; +import { + LuDataFrameGroupByQuery, + type LuDataFrameColumnNamesOfFormat, + type LuDataFrameGroupByOptions +} from './lu-group-by-query'; import { compileLuDataFrameQuery, type CompiledLuDataFrameQuery, @@ -144,6 +149,14 @@ export class LuDataFrameQuery< ); } + /** Plans dense categorical grouping without allocating GPU resources or reading source values. */ + groupBy>( + key: Key, + options: LuDataFrameGroupByOptions = {} + ): LuDataFrameGroupByQuery { + return new LuDataFrameGroupByQuery(this, key, 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 f4c818cb8e..1aba87bac4 100644 --- a/modules/experimental/src/ludf/lu-data-frame.ts +++ b/modules/experimental/src/ludf/lu-data-frame.ts @@ -22,6 +22,11 @@ import { type LuDataFrameDerivedColumnOptions } from './lu-data-frame-query'; import type {LuExpression} from './lu-expression'; +import type { + LuDataFrameColumnNamesOfFormat, + LuDataFrameGroupByOptions, + LuDataFrameGroupByQuery +} from './lu-group-by-query'; /** Whether a dataframe borrows its source resources or releases them after its final view. */ export type LuDataFrameOwnership = 'borrowed' | 'owned'; @@ -191,6 +196,18 @@ export class LuDataFrame { return query.withColumn(name, expression, options); } + /** Plans dense unsigned grouping without allocating GPU resources or retaining source leases. */ + groupBy>( + key: Key, + options: LuDataFrameGroupByOptions = {} + ): LuDataFrameGroupByQuery { + this.assertAvailable(); + return new LuDataFrameQuery(this, [], this.columnNames).groupBy( + key, + options + ); + } + /** * Returns an independent borrowed projection without mutating or destroying source columns. * diff --git a/modules/experimental/src/ludf/lu-expression-shader.ts b/modules/experimental/src/ludf/lu-expression-shader.ts index d39a9a18a6..d43a22d200 100644 --- a/modules/experimental/src/ludf/lu-expression-shader.ts +++ b/modules/experimental/src/ludf/lu-expression-shader.ts @@ -73,9 +73,10 @@ export function makeLuQueryExpressionShaderPlan( source: LuDataFrame, predicates: readonly LuExpression[], derivedColumns: readonly LuDataFrameDerivedColumn[] = [], - selectedColumns: readonly string[] = [] + selectedColumns: readonly string[] = [], + allowEmptyPredicates = false ): LuQueryExpressionShaderPlan { - if (predicates.length === 0 && derivedColumns.length === 0) { + if (predicates.length === 0 && derivedColumns.length === 0 && !allowEmptyPredicates) { throw new Error('LuDataFrame filtering requires at least one predicate'); } diff --git a/modules/experimental/src/ludf/lu-group-aggregation-compiler.ts b/modules/experimental/src/ludf/lu-group-aggregation-compiler.ts new file mode 100644 index 0000000000..6226771b75 --- /dev/null +++ b/modules/experimental/src/ludf/lu-group-aggregation-compiler.ts @@ -0,0 +1,563 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors + +import {Buffer, type Binding, type Device} from '@luma.gl/core'; +import {Computation} from '@luma.gl/engine'; +import { + GPUData, + GPURecordBatch, + GPUTable, + GPUVector, + type GPUField, + type GPUTypeMap +} from '@luma.gl/tables'; +import { + GraphVectorView, + type GPUCommandGraph, + type GraphBufferUse, + type GraphDataView +} from '../gpu-primitives/gpu-command-graph'; +import {GPUGroupAggregation} from '../gpu-primitives/gpu-group-aggregation'; +import {GPUMask} from '../gpu-primitives/gpu-mask'; +import { + createTransientVectorView, + getViewBinding, + getViewElementOffset +} from '../gpu-primitives/graph-data-view-utils'; +import type {LuDataFrame, LuDataFrameDictionaries, LuDataFrameValidity} from './lu-data-frame'; +import type {LuDataFrameDerivedColumn} from './lu-data-frame-query'; +import type {LuExpression} from './lu-expression'; +import type {LuDataFrameAggregationDefinition} from './lu-group-by-query'; +import { + CompiledLuDataFrameQuery, + compileLuDataFrameQuery, + type CompiledLuDataFrameQueryProps, + type LuDataFrameQueryExtensionContext, + type LuDataFrameQueryExtensionResult, + type LuDataFrameQueryParameters +} from './lu-query-compiler'; + +const LU_GROUP_WORKGROUP_SIZE = 256; +const UINT32_BYTE_LENGTH = Uint32Array.BYTES_PER_ELEMENT; +const MAXIMUM_UINT32 = 0xffffffff; + +type LuGroupedMetricState = { + values: GraphVectorView<'float32'>; + mask: GraphVectorView<'uint32'>; + validity: GPUVector<'uint32'>; +}; + +/** Dense grouped GPU statistics retaining their original source selection and row identities. */ +export class CompiledLuDataFrameGroupedAggregation< + T extends GPUTypeMap = GPUTypeMap +> extends CompiledLuDataFrameQuery { + /** Number of dense categorical group rows represented by the grouped result table. */ + readonly groupCount: number; + + /** @internal */ + constructor(props: CompiledLuDataFrameQueryProps, groupCount: number) { + super(props); + this.groupCount = groupCount; + } +} + +/** Adds dense categorical aggregation to source-row work before the shared graph is frozen. */ +export function compileLuDataFrameGroupedAggregation< + Source extends GPUTypeMap, + Selection extends GPUTypeMap, + Result extends GPUTypeMap +>( + source: LuDataFrame, + predicates: readonly LuExpression[], + selectedColumns: readonly (keyof Selection & string)[], + derivedColumns: readonly LuDataFrameDerivedColumn[], + key: keyof Selection & string, + groupCount: number, + definitions: readonly LuDataFrameAggregationDefinition[], + graph: GPUCommandGraph +): CompiledLuDataFrameGroupedAggregation { + validateLuGroupingCapacity(source, groupCount, graph); + validateLuGroupingSourceColumns(source, key, definitions); + return compileLuDataFrameQuery< + Source, + Selection, + Result, + CompiledLuDataFrameGroupedAggregation + >(source, predicates, selectedColumns, graph, derivedColumns, { + allowEmptyPredicates: true, + prepare: context => + addLuGroupedAggregationToGraph(context, key, groupCount, definitions) + }); +} + +/** Rejects unsupported constants and explicitly unknown source nullability before GPU allocation. */ +function validateLuGroupingSourceColumns( + source: LuDataFrame, + key: string, + definitions: readonly LuDataFrameAggregationDefinition[] +): void { + const columnNames = new Set([ + key, + ...definitions.flatMap(definition => (definition.column ? [definition.column] : [])) + ]); + for (const name of columnNames) { + if (source.table.gpuConstants[name]) { + throw new Error(`LuDataFrame grouping column "${name}" must contain GPU vector data`); + } + const field = source.schema.fields.find(candidate => candidate.name === name); + if (field?.nullable && source.numRows > 0 && !source.validity[name as keyof Source & string]) { + throw new Error(`LuDataFrame nullable grouping column "${name}" requires GPU validity`); + } + } +} + +/** Rejects overflow and unsupported output dimensions before retaining or allocating GPU state. */ +function validateLuGroupingCapacity( + source: LuDataFrame, + groupCount: number, + graph: GPUCommandGraph +): void { + if (!Number.isSafeInteger(groupCount) || groupCount <= 0 || groupCount > MAXIMUM_UINT32) { + throw new Error('LuDataFrame grouping requires a positive uint32 group count'); + } + if (source.numRows > MAXIMUM_UINT32) { + throw new Error('LuDataFrame group counts cannot represent more than uint32 source rows'); + } + if (groupCount > graph.device.limits.maxComputeWorkgroupsPerDimension * LU_GROUP_WORKGROUP_SIZE) { + throw new Error('LuDataFrame group count exceeds the supported dispatch capacity'); + } + const byteLength = groupCount * UINT32_BYTE_LENGTH; + if ( + byteLength > graph.device.limits.maxStorageBufferBindingSize || + byteLength > graph.device.limits.maxBufferSize + ) { + throw new Error('LuDataFrame group count exceeds the available GPU buffer capacity'); + } +} + +/** Creates grouped buffers, null-aware masks, and one cross-batch primitive per requested metric. */ +function addLuGroupedAggregationToGraph( + context: LuDataFrameQueryExtensionContext, + key: keyof Selection & string, + groupCount: number, + definitions: readonly LuDataFrameAggregationDefinition[] +): LuDataFrameQueryExtensionResult> { + const graph = context.graph; + const prefix = `${context.queryId}-group`; + const ownedVectors: GPUVector[] = []; + let resultTable: GPUTable | undefined; + + try { + const sourceKeys = getLuGroupingVector(context, key, 'uint32'); + const keys = graph.importGPUVector(`${prefix}-keys`, sourceKeys); + const keyValidity = getLuGroupingValidity(context, key, keys); + const baseMask = keyValidity + ? combineLuGroupingMasks(graph, `${prefix}-key-validity`, context.selectionMask, keyValidity) + : context.selectionMask; + + const outputVectors = new Map | GPUVector<'float32'>>(); + const groupKeys = createLuGroupedOutputVector( + graph.device, + `${prefix}-output-keys`, + groupCount, + 'uint32' + ); + ownedVectors.push(groupKeys); + outputVectors.set(key, groupKeys); + const groupKeyView = graph.importGPUVector(`${prefix}-output-key-vector`, groupKeys).data[0]; + addLuGroupIdentityPass(graph, `${prefix}-output-key-identity`, groupKeyView); + + const metricStates = new Map(); + const validity: Record> = {}; + + for (const [definitionIndex, definition] of definitions.entries()) { + const metricId = `${prefix}-metric-${definitionIndex}`; + if (definition.operation === 'count') { + const output = createLuGroupedOutputVector(graph.device, metricId, groupCount, 'uint32'); + ownedVectors.push(output); + outputVectors.set(definition.name, output); + new GPUGroupAggregation({ + id: metricId, + keys, + mask: baseMask, + output: graph.importGPUVector(`${metricId}-output`, output).data[0], + operation: 'count' + }).addToGraph(graph); + continue; + } + + const columnName = definition.column; + if (!columnName) { + throw new Error('LuDataFrame grouped statistics require a numeric value column'); + } + let state = metricStates.get(columnName); + if (!state) { + state = createLuGroupedMetricState( + context, + keys, + baseMask, + columnName, + groupCount, + `${prefix}-values-${metricStates.size}`, + ownedVectors + ); + metricStates.set(columnName, state); + } + + const output = createLuGroupedOutputVector(graph.device, metricId, groupCount, 'float32'); + ownedVectors.push(output); + outputVectors.set(definition.name, output); + validity[definition.name] = state.validity; + new GPUGroupAggregation({ + id: metricId, + keys, + values: state.values, + mask: state.mask, + output: graph.importGPUVector(`${metricId}-output`, output).data[0], + operation: definition.operation + }).addToGraph(graph); + } + + resultTable = createLuGroupedResultTable( + context.table, + key, + definitions, + outputVectors + ); + const dictionary = context.dictionaries[key]; + const dictionaries = Object.freeze(dictionary ? {[key]: dictionary} : {}) as Readonly< + LuDataFrameDictionaries + >; + + return { + table: resultTable, + validity: Object.freeze(validity) as Readonly>, + dictionaries, + ownedTables: [resultTable], + ownedVectors, + createCompiled: props => new CompiledLuDataFrameGroupedAggregation(props, groupCount) + }; + } catch (error) { + resultTable?.destroy(); + for (const vector of ownedVectors) { + vector.destroy(); + } + throw error; + } +} + +/** Resolves one batch-aligned grouped input, including schema-only sources without GPU chunks. */ +function getLuGroupingVector( + context: LuDataFrameQueryExtensionContext, + name: string, + format: Format +): GPUVector { + if (context.table.gpuConstants[name]) { + throw new Error(`LuDataFrame grouping column "${name}" must contain GPU vector data`); + } + const field = context.table.schema.fields.find(candidate => candidate.name === name); + const vector = context.table.gpuVectors[name]; + if ( + !field || + (vector && vector.format !== format) || + (!vector && context.table.batches.length > 0) + ) { + throw new Error(`LuDataFrame grouping column "${name}" requires ${format} GPU vector data`); + } + if (!vector) { + if (field.format !== format) { + throw new Error(`LuDataFrame grouping column "${name}" requires ${format} GPU vector data`); + } + return new GPUVector({type: 'data', name, format, data: [], ownsData: false}); + } + return vector as GPUVector; +} + +/** Imports nullable source/derived sidecars and rejects explicitly unknown nonempty validity. */ +function getLuGroupingValidity( + context: LuDataFrameQueryExtensionContext, + name: string, + template: GraphVectorView +): GraphVectorView<'uint32'> | undefined { + const field = context.table.schema.fields.find(candidate => candidate.name === name); + if (!field?.nullable) { + return undefined; + } + const validity = context.validity[name as keyof Selection & string]; + if (!validity) { + if (template.length === 0) { + return undefined; + } + throw new Error(`LuDataFrame nullable grouping column "${name}" requires GPU validity`); + } + return context.graph.importGPUVector(`${context.queryId}-group-validity-${name}`, validity); +} + +/** Intersects two source-row-aligned masks into graph-owned chunk-preserving scratch. */ +function combineLuGroupingMasks( + graph: GPUCommandGraph, + id: string, + first: GraphVectorView<'uint32'>, + second: GraphVectorView<'uint32'> +): GraphVectorView<'uint32'> { + const output = createTransientVectorView(graph, id, first); + new GPUMask({id: `${id}-compose`, inputs: [first, second], output}).addToGraph(graph); + return output; +} + +/** Shares null/NaN-filtered source masks and group validity across every metric on one column. */ +function createLuGroupedMetricState( + context: LuDataFrameQueryExtensionContext, + keys: GraphVectorView<'uint32'>, + baseMask: GraphVectorView<'uint32'>, + name: string, + groupCount: number, + id: string, + ownedVectors: GPUVector[] +): LuGroupedMetricState { + const values = context.graph.importGPUVector( + `${id}-input`, + getLuGroupingVector(context, name, 'float32') + ); + const valueValidity = getLuGroupingValidity(context, name, values); + const validRows = valueValidity + ? combineLuGroupingMasks(context.graph, `${id}-validity-mask`, baseMask, valueValidity) + : baseMask; + const finiteRows = createTransientVectorView(context.graph, `${id}-finite-mask`, validRows); + addLuGroupFiniteMaskPasses(context.graph, `${id}-finite`, values, validRows, finiteRows); + + const validity = createLuGroupedOutputVector( + context.graph.device, + `${id}-group-validity`, + groupCount, + 'uint32' + ); + ownedVectors.push(validity); + const output = context.graph.importGPUVector(`${id}-group-validity-vector`, validity).data[0]; + new GPUGroupAggregation({ + id: `${id}-accepted-count`, + keys, + mask: finiteRows, + output, + operation: 'count' + }).addToGraph(context.graph); + addLuNormalizeGroupValidityPass(context.graph, `${id}-normalize-validity`, output); + return {values, mask: finiteRows, validity}; +} + +/** Creates one owned dense, renderer-compatible result vector with one physical GPU chunk. */ +function createLuGroupedOutputVector( + device: Device, + name: string, + length: number, + format: Format +): GPUVector { + const buffer = device.createBuffer({ + id: name, + byteLength: Math.max(length, 1) * UINT32_BYTE_LENGTH, + usage: Buffer.STORAGE | Buffer.VERTEX | Buffer.COPY_SRC | Buffer.COPY_DST + }); + try { + const data = new GPUData({buffer, format, length, ownsBuffer: true}); + return new GPUVector({type: 'data', name, format, data: [data], ownsData: true}); + } catch (error) { + buffer.destroy(); + throw error; + } +} + +/** Produces one explicit dense categorical batch with independently borrowed output wrappers. */ +function createLuGroupedResultTable( + source: GPUTable, + key: string, + definitions: readonly LuDataFrameAggregationDefinition[], + vectors: ReadonlyMap | GPUVector<'float32'>> +): GPUTable { + const sourceKey = source.schema.fields.find(field => field.name === key); + if (!sourceKey) { + throw new Error('LuDataFrame grouped result requires an existing key field'); + } + const fields: GPUField[] = [ + { + name: key, + format: 'uint32', + nullable: false, + ...(sourceKey.metadata ? {metadata: new Map(sourceKey.metadata)} : {metadata: new Map()}) + }, + ...definitions.map(definition => ({ + name: definition.name, + format: definition.operation === 'count' ? ('uint32' as const) : ('float32' as const), + nullable: definition.operation !== 'count', + metadata: new Map() + })) + ]; + const gpuData: Record = {}; + for (const field of fields) { + const sourceData = vectors.get(field.name)?.data[0]; + if (!sourceData) { + throw new Error('LuDataFrame grouped result is missing an output vector'); + } + gpuData[field.name] = new GPUData({ + buffer: sourceData.buffer, + format: sourceData.format, + length: sourceData.length, + ownsBuffer: false + }); + } + const batch = new GPURecordBatch({ + gpuData, + fields, + metadata: new Map(source.schema.metadata) + }); + try { + return new GPUTable({batches: [batch]}); + } catch (error) { + batch.destroy(); + throw error; + } +} + +/** Initializes dense category IDs entirely on the GPU without reading or materializing source rows. */ +function addLuGroupIdentityPass( + graph: GPUCommandGraph, + id: string, + output: GraphDataView<'uint32'> +): void { + const source = /* wgsl */ ` +const ELEMENT_COUNT: u32 = ${output.length}u; +const OUTPUT_OFFSET: u32 = ${getViewElementOffset(output)}u; +@group(0) @binding(0) var outputValues: array; + +@compute @workgroup_size(${LU_GROUP_WORKGROUP_SIZE}) +fn main(@builtin(global_invocation_id) globalId: vec3) { + if (globalId.x < ELEMENT_COUNT) { + outputValues[OUTPUT_OFFSET + globalId.x] = globalId.x; + } +}`; + addLuGroupingComputePass(graph, { + id, + source, + resources: [{buffer: output, usage: 'storage-write'}], + bindings: {outputValues: output}, + length: output.length + }); +} + +/** Rejects null, NaN, and infinite contributions without flattening source record batches. */ +function addLuGroupFiniteMaskPasses( + graph: GPUCommandGraph, + id: string, + values: GraphVectorView<'float32'>, + input: GraphVectorView<'uint32'>, + output: GraphVectorView<'uint32'> +): void { + for (const [chunkIndex, mask] of output.data.entries()) { + if (mask.length === 0) { + continue; + } + const value = values.data[chunkIndex]; + const sourceMask = input.data[chunkIndex]; + const source = /* wgsl */ ` +const ELEMENT_COUNT: u32 = ${mask.length}u; +const VALUE_OFFSET: u32 = ${getViewElementOffset(value)}u; +const INPUT_OFFSET: u32 = ${getViewElementOffset(sourceMask)}u; +const OUTPUT_OFFSET: u32 = ${getViewElementOffset(mask)}u; +@group(0) @binding(0) var inputValues: array; +@group(0) @binding(1) var inputMask: array; +@group(0) @binding(2) var outputMask: array; + +@compute @workgroup_size(${LU_GROUP_WORKGROUP_SIZE}) +fn main(@builtin(global_invocation_id) globalId: vec3) { + if (globalId.x < ELEMENT_COUNT) { + let value = inputValues[VALUE_OFFSET + globalId.x]; + let finite = value == value && abs(value) <= 3.402823466e+38; + outputMask[OUTPUT_OFFSET + globalId.x] = select( + 0u, + 1u, + inputMask[INPUT_OFFSET + globalId.x] != 0u && finite + ); + } +}`; + addLuGroupingComputePass(graph, { + id: `${id}-chunk-${chunkIndex}`, + source, + resources: [ + {buffer: value, usage: 'storage-read'}, + {buffer: sourceMask, usage: 'storage-read'}, + {buffer: mask, usage: 'storage-write'} + ], + bindings: {inputValues: value, inputMask: sourceMask, outputMask: mask}, + length: mask.length + }); + } +} + +/** Converts accepted floating-point contribution counts into canonical 0/1 validity flags. */ +function addLuNormalizeGroupValidityPass( + graph: GPUCommandGraph, + id: string, + validity: GraphDataView<'uint32'> +): void { + const source = /* wgsl */ ` +const ELEMENT_COUNT: u32 = ${validity.length}u; +const OUTPUT_OFFSET: u32 = ${getViewElementOffset(validity)}u; +@group(0) @binding(0) var outputValidity: array; + +@compute @workgroup_size(${LU_GROUP_WORKGROUP_SIZE}) +fn main(@builtin(global_invocation_id) globalId: vec3) { + if (globalId.x < ELEMENT_COUNT) { + let offset = OUTPUT_OFFSET + globalId.x; + outputValidity[offset] = select(0u, 1u, outputValidity[offset] != 0u); + } +}`; + addLuGroupingComputePass(graph, { + id, + source, + resources: [{buffer: validity, usage: 'storage-read-write'}], + bindings: {outputValidity: validity}, + length: validity.length + }); +} + +/** Owns one small grouping computation while preserving typed graph buffer hazard declarations. */ +function addLuGroupingComputePass( + graph: GPUCommandGraph, + props: { + id: string; + source: string; + resources: readonly GraphBufferUse[]; + bindings: Readonly>; + length: number; + } +): void { + graph.addComputePass({ + id: props.id, + resources: [...props.resources], + compile: ({device}) => { + const computation = new Computation(device, { + id: props.id, + source: props.source, + shaderLayout: { + bindings: Object.keys(props.bindings).map((name, location) => ({ + name, + type: 'storage' as const, + group: 0, + location + })) + } + }); + return { + encode: ({computePass, getBuffer}) => { + const resolvedBindings: Record = {}; + for (const [name, view] of Object.entries(props.bindings)) { + resolvedBindings[name] = getViewBinding(view, getBuffer); + } + computation.setBindings(resolvedBindings); + computation.dispatch(computePass, Math.ceil(props.length / LU_GROUP_WORKGROUP_SIZE)); + }, + destroy: () => computation.destroy() + }; + } + }); +} diff --git a/modules/experimental/src/ludf/lu-group-by-query.ts b/modules/experimental/src/ludf/lu-group-by-query.ts new file mode 100644 index 0000000000..427471d457 --- /dev/null +++ b/modules/experimental/src/ludf/lu-group-by-query.ts @@ -0,0 +1,264 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors + +import type {GPUTypeMap} from '@luma.gl/tables'; +import type {GPUCommandGraph} from '../gpu-primitives/gpu-command-graph'; +import type {LuDataFrameDictionary} from './lu-data-frame'; +import type {LuDataFrameQuery} from './lu-data-frame-query'; +import {getLuExpressionColumnNames} from './lu-expression'; +import { + compileLuDataFrameGroupedAggregation, + type CompiledLuDataFrameGroupedAggregation +} from './lu-group-aggregation-compiler'; +import type {LuDataFrameQueryParameters} from './lu-query-compiler'; + +/** Scalar statistics supported by dense GPU-resident categorical grouping. */ +export type LuDataFrameAggregationOperation = 'count' | 'sum' | 'min' | 'max' | 'mean'; + +/** One normalized, immutable GPU group operation consumed by the graph compiler. */ +export type LuDataFrameAggregationDefinition = Readonly<{ + /** Output dataframe field name. */ + name: string; + /** Dense unsigned count or floating-point summary statistic. */ + operation: LuDataFrameAggregationOperation; + /** Floating-point input field for sum, minimum, maximum, or mean. */ + column?: string; +}>; + +/** Selects logical dataframe columns having one supported scalar storage format. */ +export type LuDataFrameColumnNamesOfFormat< + T extends GPUTypeMap, + SelectedColumns extends keyof T & string, + Format extends 'uint32' | 'float32' +> = { + [Name in SelectedColumns]: T[Name] extends Format ? Name : never; +}[SelectedColumns]; + +/** Optional dense key-domain size when source categorical labels cannot establish it. */ +export type LuDataFrameGroupByOptions = Readonly<{groupCount?: number}>; + +/** One named count or a float32 statistic referencing a selected numerical column. */ +export type LuDataFrameAggregationValue< + T extends GPUTypeMap, + SelectedColumns extends keyof T & string +> = + | 'count' + | Readonly<{sum: LuDataFrameColumnNamesOfFormat}> + | Readonly<{min: LuDataFrameColumnNamesOfFormat}> + | Readonly<{max: LuDataFrameColumnNamesOfFormat}> + | Readonly<{mean: LuDataFrameColumnNamesOfFormat}>; + +/** Caller-defined dense grouped output names and their corresponding statistics. */ +export type LuDataFrameAggregationDefinitions< + T extends GPUTypeMap, + SelectedColumns extends keyof T & string = keyof T & string +> = Readonly>>; + +/** Exact GPU result formats for the group-key column and every requested statistic. */ +export type LuDataFrameGroupedAggregationResult< + Key extends string, + Definitions extends Readonly> +> = Record & { + [Name in keyof Definitions & string]: Definitions[Name] extends 'count' ? 'uint32' : 'float32'; +}; + +/** + * Immutable dense-key grouping plan that never allocates GPU resources. + * + * Dictionary-backed category labels establish the complete dense key domain. Other unsigned key + * columns require an explicit positive `groupCount` so query planning never needs CPU readback. + */ +export class LuDataFrameGroupByQuery< + Logical extends GPUTypeMap, + SelectedColumns extends keyof Logical & string, + Key extends LuDataFrameColumnNamesOfFormat, + Source extends GPUTypeMap = Logical +> { + /** Complete immutable source/filter/derived query retained without acquiring a resource lease. */ + readonly query: LuDataFrameQuery; + /** Existing unsigned logical column providing dense categorical row keys. */ + readonly key: Key; + /** Number of output groups inferred from labels or supplied explicitly by the caller. */ + readonly groupCount: number; + + /** Validates grouping metadata entirely on the CPU. @internal */ + constructor( + query: LuDataFrameQuery, + key: Key, + options: LuDataFrameGroupByOptions = {} + ) { + if (!query.selectedColumns.includes(key)) { + throw new Error(`LuDataFrame group key "${key}" is not selected`); + } + if (getLuDataFrameQueryColumnFormats(query).get(key) !== 'uint32') { + throw new Error(`LuDataFrame group key "${key}" requires a uint32 column`); + } + + const dictionary = (query.source.dictionaries as Record)[key]; + const dictionaryCount = getLuDataFrameDictionaryLength(dictionary); + if (options.groupCount !== undefined && dictionaryCount !== undefined) { + if (options.groupCount !== dictionaryCount) { + throw new Error('LuDataFrame group count must match the category dictionary'); + } + } + const groupCount = options.groupCount ?? dictionaryCount; + if ( + !Number.isSafeInteger(groupCount) || + !groupCount || + groupCount < 1 || + groupCount > 0xffffffff + ) { + throw new Error('LuDataFrame grouping requires a positive uint32 group count'); + } + + this.query = query; + this.key = key; + this.groupCount = groupCount; + Object.freeze(this); + } + + /** Plans named dense count/sum/min/max/mean outputs without touching GPU resources. */ + aggregate>( + definitions: Definitions + ): LuDataFrameGroupedAggregationQuery { + return new LuDataFrameGroupedAggregationQuery< + Logical, + SelectedColumns, + Key, + Definitions, + Source + >(this.query, this.key, this.groupCount, definitions); + } +} + +/** Immutable dense aggregation query lowered to existing graph-native GPU group primitives. */ +export class LuDataFrameGroupedAggregationQuery< + Logical extends GPUTypeMap, + SelectedColumns extends keyof Logical & string, + Key extends LuDataFrameColumnNamesOfFormat, + Definitions extends LuDataFrameAggregationDefinitions, + Source extends GPUTypeMap = Logical +> { + /** Complete source query preserving immutable filters, projections, and derived expressions. */ + readonly query: LuDataFrameQuery; + /** Dense unsigned group key exposed as the first output dataframe column. */ + readonly key: Key; + /** Complete category domain and grouped output row count. */ + readonly groupCount: number; + /** Named operations normalized in the caller's original object insertion order. */ + readonly definitions: readonly LuDataFrameAggregationDefinition[]; + + /** Validates all supported operations, names, and selected float32 metric inputs. @internal */ + constructor( + query: LuDataFrameQuery, + key: Key, + groupCount: number, + definitions: Definitions + ) { + this.query = query; + this.key = key; + this.groupCount = groupCount; + this.definitions = Object.freeze(normalizeLuDataFrameAggregations(query, key, definitions)); + Object.freeze(this); + } + + /** Encodes reusable grouping, filtering, and derived work into one caller-owned command graph. */ + compile( + graph: GPUCommandGraph + ): CompiledLuDataFrameGroupedAggregation> { + return compileLuDataFrameGroupedAggregation< + Source, + Pick, + LuDataFrameGroupedAggregationResult + >( + this.query.source, + this.query.predicates, + this.query.selectedColumns, + this.query.derivedColumns, + this.key, + this.groupCount, + this.definitions, + graph + ); + } +} + +/** Normalizes a closed set of operation shapes without accepting arbitrary shader expressions. */ +function normalizeLuDataFrameAggregations< + Logical extends GPUTypeMap, + SelectedColumns extends keyof Logical & string, + Source extends GPUTypeMap +>( + query: LuDataFrameQuery, + key: string, + definitions: LuDataFrameAggregationDefinitions +): LuDataFrameAggregationDefinition[] { + const entries = Object.entries(definitions); + if (entries.length === 0) { + throw new Error('LuDataFrame grouping requires at least one aggregation'); + } + + const formats = getLuDataFrameQueryColumnFormats(query); + const normalized: LuDataFrameAggregationDefinition[] = []; + for (const [name, value] of entries) { + if (name.length === 0 || name === key) { + throw new Error('LuDataFrame aggregation names must be distinct from the group key'); + } + if (value === 'count') { + normalized.push(Object.freeze({name, operation: 'count'})); + continue; + } + if (!value || typeof value !== 'object') { + throw new Error(`LuDataFrame aggregation "${name}" has an unsupported operation`); + } + + const operations = Object.entries(value); + if (operations.length !== 1) { + throw new Error(`LuDataFrame aggregation "${name}" requires exactly one operation`); + } + const [operation, column] = operations[0]; + if (operation !== 'sum' && operation !== 'min' && operation !== 'max' && operation !== 'mean') { + throw new Error(`LuDataFrame aggregation "${name}" has an unsupported operation`); + } + if (typeof column !== 'string' || !query.selectedColumns.includes(column as SelectedColumns)) { + throw new Error(`LuDataFrame aggregation "${name}" requires a selected input column`); + } + if (formats.get(column) !== 'float32') { + throw new Error(`LuDataFrame aggregation "${name}" requires a float32 input column`); + } + normalized.push(Object.freeze({name, operation, column})); + } + + return normalized; +} + +/** Resolves original and derived logical scalar formats using canonical GPU column metadata. */ +function getLuDataFrameQueryColumnFormats< + Logical extends GPUTypeMap, + SelectedColumns extends keyof Logical & string +>(query: LuDataFrameQuery): Map { + const formats = new Map( + query.source.schema.fields.map( + field => + [field.name, query.source.table.gpuColumns[field.name]?.format ?? field.format] as const + ) + ); + + for (const definition of query.derivedColumns) { + const firstReference = getLuExpressionColumnNames(definition.expression)[0]; + formats.set(definition.name, definition.format ?? formats.get(firstReference) ?? 'float32'); + } + + return formats; +} + +/** Reads explicit adapter-owned categorical metadata without introducing an Apache Arrow import. */ +function getLuDataFrameDictionaryLength(dictionary?: LuDataFrameDictionary): number | undefined { + if (!dictionary) { + return undefined; + } + return Array.isArray(dictionary) + ? dictionary.length + : (dictionary as {values: readonly unknown[]}).values.length; +} diff --git a/modules/experimental/src/ludf/lu-query-compiler.ts b/modules/experimental/src/ludf/lu-query-compiler.ts index 8e93f688ac..7e968609f2 100644 --- a/modules/experimental/src/ludf/lu-query-compiler.ts +++ b/modules/experimental/src/ludf/lu-query-compiler.ts @@ -72,7 +72,8 @@ type LuQueryDerivedView = { validity?: GraphVectorView<'uint32'>; }; -type CompiledLuDataFrameQueryProps = { +/** Internal ownership and graph state transferred exactly once to a compiled dataframe query. */ +export type CompiledLuDataFrameQueryProps = { table: GPUTable; validity: Readonly>; dictionaries: Readonly>; @@ -81,10 +82,45 @@ type CompiledLuDataFrameQueryProps = { selectedCounts: GPUVector<'uint32'>; graph: CompiledGPUCommandGraph; sourceViews: readonly Pick[]; - ownedTable?: GPUTable; + ownedTables?: readonly Pick[]; ownedVectors?: readonly GPUVector[]; }; +/** Source-row GPU outputs available to one graph contribution before graph compilation. @internal */ +export type LuDataFrameQueryExtensionContext = { + graph: GPUCommandGraph; + queryId: string; + table: GPUTable; + validity: Readonly>; + dictionaries: Readonly>; + selectionMask: GraphVectorView<'uint32'>; +}; + +/** Result resources contributed by one graph-native extension. @internal */ +export type LuDataFrameQueryExtensionResult< + T extends GPUTypeMap, + Compiled extends CompiledLuDataFrameQuery = CompiledLuDataFrameQuery +> = { + table: GPUTable; + validity: Readonly>; + dictionaries: Readonly>; + ownedTables?: readonly Pick[]; + ownedVectors?: readonly GPUVector[]; + createCompiled: (props: CompiledLuDataFrameQueryProps) => Compiled; +}; + +/** Declares downstream GPU work after row filtering but before the graph is frozen. @internal */ +export type LuDataFrameQueryCompilationExtension< + Row extends GPUTypeMap, + Result extends GPUTypeMap, + Compiled extends CompiledLuDataFrameQuery = CompiledLuDataFrameQuery +> = { + allowEmptyPredicates?: boolean; + prepare: ( + context: LuDataFrameQueryExtensionContext + ) => LuDataFrameQueryExtensionResult; +}; + /** * Reusable GPU dataframe query with source-aligned masks and stable per-batch selected row IDs. * @@ -107,7 +143,7 @@ export class CompiledLuDataFrameQuery { private readonly graph: CompiledGPUCommandGraph; private readonly sourceViews: readonly Pick[]; - private readonly ownedTable?: GPUTable; + private readonly ownedTables: readonly Pick[]; private readonly ownedVectors: readonly GPUVector[]; private destroyed = false; @@ -121,7 +157,7 @@ export class CompiledLuDataFrameQuery { this.selectedCounts = props.selectedCounts; this.graph = props.graph; this.sourceViews = props.sourceViews; - this.ownedTable = props.ownedTable; + this.ownedTables = props.ownedTables ?? []; this.ownedVectors = props.ownedVectors ?? []; } @@ -146,7 +182,9 @@ export class CompiledLuDataFrameQuery { this.selectionMask.destroy(); this.rowIndices.destroy(); this.selectedCounts.destroy(); - this.ownedTable?.destroy(); + for (const table of this.ownedTables) { + table.destroy(); + } for (const vector of this.ownedVectors) { vector.destroy(); } @@ -157,16 +195,23 @@ export class CompiledLuDataFrameQuery { } /** Compiles immutable dataframe predicates into source-batch-preserving WebGPU command work. */ -export function compileLuDataFrameQuery( +export function compileLuDataFrameQuery< + Source extends GPUTypeMap, + Row extends GPUTypeMap, + Result extends GPUTypeMap = Row, + Compiled extends CompiledLuDataFrameQuery = CompiledLuDataFrameQuery +>( source: LuDataFrame, predicates: readonly LuExpression[], - selectedColumns: readonly (keyof Result & string)[], + selectedColumns: readonly (keyof Row & string)[], graph: GPUCommandGraph, - derivedColumns: readonly LuDataFrameDerivedColumn[] = [] -): CompiledLuDataFrameQuery { + derivedColumns: readonly LuDataFrameDerivedColumn[] = [], + extension?: LuDataFrameQueryCompilationExtension +): Compiled { const retainedSource = source.select(source.columnNames); let selectedSource: LuDataFrame | undefined; - let ownedTable: GPUTable | undefined; + let ownedTable: GPUTable | undefined; + let extensionResult: LuDataFrameQueryExtensionResult | undefined; const derivedOutputs: LuQueryDerivedOutput[] = []; let selectionMask: GPUVector<'uint32'> | undefined; let rowIndices: GPUVector<'uint32'> | undefined; @@ -182,7 +227,8 @@ export function compileLuDataFrameQuery 0) { - ownedTable = createLuQueryDerivedTable( + ownedTable = createLuQueryDerivedTable( selectedSource.table, selectedColumns, derivedOutputs ); } - const validity = selectLuQueryValidity(selectedSource, derivedOutputs); + const validity = selectLuQueryValidity(selectedSource, derivedOutputs); const dictionaries = Object.freeze({ ...selectedSource.dictionaries - }) as Readonly>; + }) as Readonly>; const queryId = `${graph.id}-ludf-query`; const sourceViews = importLuQuerySourceViews(graph, retainedSource, plan, queryId); @@ -288,26 +334,53 @@ export function compileLuDataFrameQuery); + if (extension) { + extensionResult = extension.prepare({ + graph, + queryId, + table: rowTable, + validity, + dictionaries, + selectionMask: maskView + }); + } + compiledGraph = graph.compile(); - return new CompiledLuDataFrameQuery({ - table: ownedTable ?? (selectedSource.table as GPUTable), - validity, - dictionaries, + const props: CompiledLuDataFrameQueryProps = { + table: extensionResult?.table ?? (rowTable as unknown as GPUTable), + validity: + extensionResult?.validity ?? (validity as unknown as Readonly>), + dictionaries: + extensionResult?.dictionaries ?? + (dictionaries as unknown as Readonly>), selectionMask, rowIndices, selectedCounts, graph: compiledGraph, sourceViews: [selectedSource, retainedSource], - ...(ownedTable ? {ownedTable} : {}), - ownedVectors: derivedOutputs.flatMap(output => - output.validity ? [output.values, output.validity] : [output.values] - ) - }); + ownedTables: [...(ownedTable ? [ownedTable] : []), ...(extensionResult?.ownedTables ?? [])], + ownedVectors: [ + ...derivedOutputs.flatMap(output => + output.validity ? [output.values, output.validity] : [output.values] + ), + ...(extensionResult?.ownedVectors ?? []) + ] + }; + return extensionResult + ? extensionResult.createCompiled(props) + : (new CompiledLuDataFrameQuery(props) as Compiled); } catch (error) { compiledGraph?.destroy(); selectionMask?.destroy(); rowIndices?.destroy(); selectedCounts?.destroy(); + for (const table of extensionResult?.ownedTables ?? []) { + table.destroy(); + } + for (const vector of extensionResult?.ownedVectors ?? []) { + vector.destroy(); + } ownedTable?.destroy(); for (const output of derivedOutputs) { output.values.destroy(); diff --git a/modules/experimental/test/index.ts b/modules/experimental/test/index.ts index e934ab3d31..a21b3e2ddd 100644 --- a/modules/experimental/test/index.ts +++ b/modules/experimental/test/index.ts @@ -42,6 +42,7 @@ import './geospatial/geospatial-projection-distance.spec'; import './ludf/lu-data-frame.spec'; import './ludf/lu-data-frame-query.spec'; import './ludf/lu-derived-columns.spec'; +import './ludf/lu-group-aggregation.spec'; import './luraster'; import './luxfilter'; import './luproj/luproj.spec'; diff --git a/modules/experimental/test/ludf/lu-group-aggregation.node.spec.ts b/modules/experimental/test/ludf/lu-group-aggregation.node.spec.ts new file mode 100644 index 0000000000..3455d3f2aa --- /dev/null +++ b/modules/experimental/test/ludf/lu-group-aggregation.node.spec.ts @@ -0,0 +1,341 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors + +import {Buffer} from '@luma.gl/core'; +import { + column, + CompiledLuDataFrameGroupedAggregation, + literal, + LuDataFrame, + LuDataFrameGroupByQuery, + LuDataFrameGroupedAggregationQuery, + parameter, + type LuDataFrameAggregationDefinitions, + type LuDataFrameGroupedAggregationResult +} from '@luma.gl/experimental/ludf'; +import { + GPUData, + GPURecordBatch, + GPUTable, + type GPUField, + type GPURecordBatchSourceInfo +} from '@luma.gl/tables'; +import {NullDevice} from '@luma.gl/test-utils'; +import {describe, expect, expectTypeOf, test, vi} from 'vitest'; + +type GroupSourceColumns = { + category: 'uint32'; + fare: 'float32'; + distance: 'sint32'; + otherGroup: 'uint32'; +}; + +type GroupSourceFixture = { + device: NullDevice; + table: GPUTable; + buffers: Buffer[]; +}; + +const GROUP_LABELS = ['economy', 'standard', 'premium', 'unused'] as const; + +describe('LuDataFrame immutable dense grouped-aggregation planning', () => { + test('plans all supported statistics without allocating GPU resources or retaining source leases', () => { + const fixture = createGroupSourceFixture([2, 0, 3]); + const source = new LuDataFrame({ + table: fixture.table, + dictionaries: {category: {values: GROUP_LABELS, ordered: true}}, + ownership: 'owned' + }); + const createBuffer = vi.spyOn(fixture.device, 'createBuffer'); + const createCommandEncoder = vi.spyOn(fixture.device, 'createCommandEncoder'); + const submit = vi.spyOn(fixture.device, 'submit'); + const tableSelect = vi.spyOn(fixture.table, 'select'); + + const grouped = source.groupBy('category'); + const aggregated = grouped.aggregate({ + count: 'count', + totalFare: {sum: 'fare'}, + minimumFare: {min: 'fare'}, + maximumFare: {max: 'fare'}, + averageFare: {mean: 'fare'} + }); + + expect(grouped).toBeInstanceOf(LuDataFrameGroupByQuery); + expect(aggregated).toBeInstanceOf(LuDataFrameGroupedAggregationQuery); + expect(grouped.key).toBe('category'); + expect(grouped.groupCount).toBe(4); + expect(aggregated.key).toBe('category'); + expect(aggregated.groupCount).toBe(4); + expect(aggregated.query.source).toBe(source); + expect(aggregated.query.predicates).toEqual([]); + expect(aggregated.definitions).toEqual([ + {name: 'count', operation: 'count'}, + {name: 'totalFare', operation: 'sum', column: 'fare'}, + {name: 'minimumFare', operation: 'min', column: 'fare'}, + {name: 'maximumFare', operation: 'max', column: 'fare'}, + {name: 'averageFare', operation: 'mean', column: 'fare'} + ]); + expect(Object.isFrozen(grouped)).toBe(true); + expect(Object.isFrozen(aggregated)).toBe(true); + expect(Object.isFrozen(aggregated.definitions)).toBe(true); + expect(Object.isFrozen(aggregated.definitions[0])).toBe(true); + expect(source.batches.map(batch => batch.numRows)).toEqual([2, 0, 3]); + expect(createBuffer).not.toHaveBeenCalled(); + expect(createCommandEncoder).not.toHaveBeenCalled(); + expect(submit).not.toHaveBeenCalled(); + expect(tableSelect).not.toHaveBeenCalled(); + + source.destroy(); + expect(fixture.buffers.every(buffer => buffer.destroyed)).toBe(true); + + createBuffer.mockRestore(); + createCommandEncoder.mockRestore(); + submit.mockRestore(); + tableSelect.mockRestore(); + }); + + test('preserves exact typed uint32 count and float32 statistic result formats', () => { + const fixture = createGroupSourceFixture([2]); + const source = new LuDataFrame({ + table: fixture.table, + dictionaries: {category: GROUP_LABELS} + }); + const grouped = source.groupBy('category').aggregate({ + count: 'count', + totalFare: {sum: 'fare'}, + minimumFare: {min: 'fare'}, + maximumFare: {max: 'fare'}, + averageFare: {mean: 'fare'} + }); + + expectTypeOf(grouped.compile).returns.toEqualTypeOf< + CompiledLuDataFrameGroupedAggregation< + Record<'category', 'uint32'> & { + count: 'uint32'; + totalFare: 'float32'; + minimumFare: 'float32'; + maximumFare: 'float32'; + averageFare: 'float32'; + } + > + >(); + expectTypeOf< + LuDataFrameGroupedAggregationResult< + 'category', + {accepted: 'count'; averageFare: {mean: 'fare'}} + > + >().toEqualTypeOf< + Record<'category', 'uint32'> & {accepted: 'uint32'; averageFare: 'float32'} + >(); + + source.destroy(); + fixture.table.destroy(); + }); + + test('preserves immutable filtered, parameterized, projected, and derived source plans', () => { + const fixture = createGroupSourceFixture([2, 0, 3]); + const source = new LuDataFrame({ + table: fixture.table, + dictionaries: {category: GROUP_LABELS} + }); + const filtered = source + .filter(column('fare').greaterThan(parameter('minimumFare', 10))) + .select(['category', 'fare']); + const first = filtered.groupBy('category').aggregate({accepted: 'count'}); + const second = filtered.groupBy('category').aggregate({totalFare: {sum: 'fare'}}); + const derived = filtered + .withColumn('adjustedFare', column('fare').add(literal(2))) + .groupBy('category') + .aggregate({averageAdjustedFare: {mean: 'adjustedFare'}}); + + expect(first.definitions).toEqual([{name: 'accepted', operation: 'count'}]); + expect(second.definitions).toEqual([{name: 'totalFare', operation: 'sum', column: 'fare'}]); + expect(derived.definitions).toEqual([ + {name: 'averageAdjustedFare', operation: 'mean', column: 'adjustedFare'} + ]); + expect(derived.query.derivedColumns.map(({name}) => name)).toEqual(['adjustedFare']); + expect(derived.query.predicates[0]).toBe(filtered.predicates[0]); + expect(first.query.columnNames).toEqual(['category', 'fare']); + expect(filtered.columnNames).toEqual(['category', 'fare']); + expect(source.columnNames).toEqual(['category', 'fare', 'distance', 'otherGroup']); + + source.destroy(); + fixture.table.destroy(); + }); + + test('infers dictionary domains and requires matching explicit unsigned group counts', () => { + const fixture = createGroupSourceFixture([2]); + const dictionarySource = new LuDataFrame({ + table: fixture.table, + dictionaries: {category: {values: GROUP_LABELS, ordered: false}} + }); + const plainSource = new LuDataFrame({table: fixture.table}); + + expect(dictionarySource.groupBy('category').groupCount).toBe(4); + expect(dictionarySource.groupBy('category', {groupCount: 4}).groupCount).toBe(4); + expect(() => dictionarySource.groupBy('category', {groupCount: 3})).toThrow( + /count|dictionary/i + ); + expect(plainSource.groupBy('otherGroup', {groupCount: 3}).groupCount).toBe(3); + expect(() => plainSource.groupBy('otherGroup')).toThrow(/count|positive/i); + + for (const groupCount of [0, -1, 1.5, Number.NaN, 0x1_0000_0000]) { + expect(() => plainSource.groupBy('otherGroup', {groupCount})).toThrow(/count|positive/i); + } + + dictionarySource.destroy(); + plainSource.destroy(); + fixture.table.destroy(); + }); + + test('rejects unsupported keys, hidden inputs, invalid operations, and non-float32 metrics', () => { + const fixture = createGroupSourceFixture([2]); + const source = new LuDataFrame({ + table: fixture.table, + dictionaries: {category: GROUP_LABELS} + }); + const grouped = source.groupBy('category'); + const createBuffer = vi.spyOn(fixture.device, 'createBuffer'); + + expect(() => + // @ts-expect-error Dense grouping keys must have uint32 GPU storage. + source.groupBy('fare', {groupCount: 2}) + ).toThrow(/key|uint32/i); + expect(() => + // @ts-expect-error Dense grouping keys must exist in the selected source schema. + source.groupBy('missing', {groupCount: 2}) + ).toThrow(/key|selected/i); + expect(() => + source + .filter(column('fare').greaterThan(literal(1))) + .select(['fare']) + // @ts-expect-error Group keys must remain selected in the current query projection. + .groupBy('category') + ).toThrow(/key|selected/i); + expect(() => grouped.aggregate({})).toThrow(/aggregation/i); + expect(() => grouped.aggregate({category: 'count'})).toThrow(/name|key/i); + expect(() => grouped.aggregate({'': 'count'})).toThrow(/name|key/i); + expect(() => + // @ts-expect-error Grouped numerical statistics require float32 metrics. + grouped.aggregate({wrongType: {sum: 'distance'}}) + ).toThrow(/float32/i); + expect(() => + // @ts-expect-error Grouped numerical metrics must be selected source fields. + grouped.aggregate({missingMetric: {sum: 'missing'}}) + ).toThrow(/selected|column/i); + const invalidOperation = { + broken: {median: 'fare'} + } as unknown as LuDataFrameAggregationDefinitions; + expect(() => grouped.aggregate(invalidOperation)).toThrow(/operation/i); + const multipleOperations = { + broken: {sum: 'fare', max: 'fare'} + } as unknown as LuDataFrameAggregationDefinitions; + expect(() => grouped.aggregate(multipleOperations)).toThrow(/one|operation/i); + expect(createBuffer).not.toHaveBeenCalled(); + + createBuffer.mockRestore(); + source.destroy(); + fixture.table.destroy(); + }); + + test('preserves nullable metadata and schema-only or empty source batch plans', () => { + for (const batchLengths of [[], [0], [2, 0, 3]] as const) { + const fixture = createGroupSourceFixture(batchLengths, {nullable: true}); + const source = new LuDataFrame({ + table: fixture.table, + dictionaries: {category: GROUP_LABELS} + }); + const grouped = source.groupBy('category').aggregate({totalFare: {sum: 'fare'}}); + + expect(grouped.groupCount).toBe(4); + expect(grouped.query.source.batches.map(batch => batch.numRows)).toEqual(batchLengths); + expect( + grouped.query.source.schema.fields.find(field => field.name === 'category')?.nullable + ).toBe(true); + expect( + grouped.query.source.schema.fields.find(field => field.name === 'fare')?.nullable + ).toBe(true); + expect(grouped.query.source.validity.category).toBeUndefined(); + expect(grouped.query.source.validity.fare).toBeUndefined(); + + source.destroy(); + fixture.table.destroy(); + } + }); + + test('rejects new group plans after the source dataframe was explicitly destroyed', () => { + const fixture = createGroupSourceFixture([2]); + const source = new LuDataFrame({ + table: fixture.table, + dictionaries: {category: GROUP_LABELS} + }); + source.destroy(); + + expect(() => source.groupBy('category')).toThrow(/destroyed/i); + fixture.table.destroy(); + }); +}); + +function createGroupSourceFixture( + batchLengths: readonly number[], + options: {nullable?: boolean} = {} +): GroupSourceFixture { + const device = new NullDevice({id: 'ludf-group-aggregation-node-device'}); + const buffers: Buffer[] = []; + const fields: GPUField[] = [ + {name: 'category', format: 'uint32', nullable: options.nullable ?? false}, + {name: 'fare', format: 'float32', nullable: options.nullable ?? false}, + {name: 'distance', format: 'sint32', nullable: false}, + {name: 'otherGroup', format: 'uint32', nullable: false} + ]; + let sourceRowIndexOffset = 40; + const batches = batchLengths.map((length, sourceBatchIndex) => { + const sourceInfo: GPURecordBatchSourceInfo = { + sourceBatchIndex, + sourceRowIndexOffset, + sourceRowCount: length + }; + sourceRowIndexOffset += length; + + return new GPURecordBatch({ + gpuData: { + category: makeGroupSourceData(device, buffers, length, 'uint32'), + fare: makeGroupSourceData(device, buffers, length, 'float32'), + distance: makeGroupSourceData(device, buffers, length, 'sint32'), + otherGroup: makeGroupSourceData(device, buffers, length, 'uint32') + }, + fields, + numRows: length, + sourceInfo + }); + }); + + const table = + batches.length > 0 + ? new GPUTable({batches}) + : new GPUTable({ + schema: {fields, metadata: new Map()}, + bufferLayout: [ + {name: 'category', format: 'uint32'}, + {name: 'fare', format: 'float32'}, + {name: 'distance', format: 'sint32'}, + {name: 'otherGroup', format: 'uint32'} + ] + }); + return {device, table, buffers}; +} + +function makeGroupSourceData( + device: NullDevice, + buffers: Buffer[], + length: number, + format: Format +): GPUData { + 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); + return new GPUData({buffer, format, length, ownsBuffer: true}); +} diff --git a/modules/experimental/test/ludf/lu-group-aggregation.spec.ts b/modules/experimental/test/ludf/lu-group-aggregation.spec.ts new file mode 100644 index 0000000000..ef92412996 --- /dev/null +++ b/modules/experimental/test/ludf/lu-group-aggregation.spec.ts @@ -0,0 +1,617 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors + +import {Buffer, type Device} from '@luma.gl/core'; +import {GPUCommandGraph} from '@luma.gl/experimental'; +import { + LuDataFrame, + column, + literal, + 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'; + +type LuGroupedSourceSchema = { + category: 'uint32'; + fare: 'float32'; + distance: 'sint32'; +}; + +type LuGroupedFixture = { + frame: LuDataFrame; + sourceBuffers: Buffer[]; +}; + +test('LuDataFrame groups nullable source batches into explicit dense GPU statistics', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuGroupedFixture(device); + const createBufferSpy = vi.spyOn(device, 'createBuffer'); + const submitSpy = vi.spyOn(device, 'submit'); + + const query = fixture.frame.groupBy('category').aggregate({ + count: 'count', + totalFare: {sum: 'fare'}, + minimumFare: {min: 'fare'}, + maximumFare: {max: 'fare'}, + averageFare: {mean: 'fare'} + }); + + testContext.equal( + createBufferSpy.mock.calls.length, + 0, + 'grouped aggregation planning never allocates GPU storage' + ); + testContext.equal( + submitSpy.mock.calls.length, + 0, + 'grouped aggregation planning never submits GPU work' + ); + + const graph = new GPUCommandGraph(device, { + id: 'ludf-nullable-dense-group-aggregation' + }); + const compiled = query.compile(graph); + + try { + testContext.equal(compiled.groupCount, 4, 'dictionary labels determine the dense group domain'); + testContext.deepEqual( + compiled.table.schema.fields.map(field => field.name), + ['category', 'count', 'totalFare', 'minimumFare', 'maximumFare', 'averageFare'], + 'grouped schema retains the dense key and requested metric ordering' + ); + testContext.deepEqual( + compiled.table.batches.map(batch => batch.numRows), + [4], + 'all original source batches contribute to one dense GPU-owned result batch' + ); + testContext.deepEqual( + compiled.dictionaries.category, + {values: ['economy', 'standard', 'premium', 'unused'], ordered: false}, + 'grouped categorical keys retain their adapter-owned dictionary labels' + ); + testContext.deepEqual( + compiled.selectionMask.data.map(chunk => chunk.length), + [2, 0, 3], + 'source selection masks preserve every original record-batch boundary' + ); + + fixture.frame.destroy(); + testContext.ok( + fixture.sourceBuffers.every(buffer => !buffer.destroyed), + 'compiled grouped queries retain their owned source lease' + ); + + const commandEncoder = device.createCommandEncoder({id: 'ludf-dense-group-encode'}); + compiled.encode(commandEncoder); + testContext.equal( + submitSpy.mock.calls.length, + 0, + 'grouped aggregation encodes work into the caller-owned command encoder' + ); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + await readUint32VectorChunks(compiled.table.gpuVectors.category), + [[0, 1, 2, 3]], + 'dense grouped rows publish stable unsigned categorical keys' + ); + testContext.deepEqual( + await readUint32VectorChunks(compiled.table.gpuVectors.count), + [[2, 2, 0, 0]], + 'group counts reject null keys but include rows with null metric values' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.totalFare), + [[40, 20, 0, 0]], + 'group sums exclude null keys and values across all source batches' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.minimumFare), + [[10, 20, NaN, NaN]], + 'group minimums publish NaN payloads for empty categories' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.maximumFare), + [[30, 20, NaN, NaN]], + 'group maximums publish NaN payloads for empty categories' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.averageFare), + [[20, 20, NaN, NaN]], + 'group means divide only finite, non-null source values' + ); + + for (const metricName of ['totalFare', 'minimumFare', 'maximumFare', 'averageFare'] as const) { + const validity = compiled.validity[metricName]; + if (!validity) { + throw new Error(`Expected explicit GPU group validity for ${metricName}`); + } + testContext.deepEqual( + await readUint32VectorChunks(validity), + [[1, 1, 0, 0]], + `${metricName} distinguishes populated groups from empty or entirely null groups` + ); + } + testContext.equal(compiled.validity.count, undefined, 'dense count columns are never nullable'); + + const groupedBuffers = Object.values(compiled.table.gpuVectors).flatMap(vector => + vector.data.map(getGPUDataBuffer) + ); + const validityBuffers = Array.from( + new Set( + Object.values(compiled.validity).flatMap(vector => + vector ? vector.data.map(getGPUDataBuffer) : [] + ) + ) + ); + compiled.destroy(); + testContext.ok( + groupedBuffers.every(buffer => buffer.destroyed), + 'grouped results release every owned dense value buffer' + ); + testContext.ok( + validityBuffers.every(buffer => buffer.destroyed), + 'grouped results release shared statistic-validity buffers exactly once' + ); + testContext.ok( + fixture.sourceBuffers.every(buffer => buffer.destroyed), + 'owned source buffers survive until the final grouped lease is destroyed' + ); + } finally { + compiled.destroy(); + fixture.frame.destroy(); + createBufferSpy.mockRestore(); + submitSpy.mockRestore(); + } + + testContext.end(); +}); + +test('LuDataFrame reuses filtered grouped aggregations with encoder-ordered parameters', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuGroupedFixture(device); + const graph = new GPUCommandGraph(device, { + id: 'ludf-parameterized-group-aggregation' + }); + const compiled = fixture.frame + .filter(column('fare').greaterThan(parameter('minimumFare', 0))) + .groupBy('category') + .aggregate({count: 'count', totalFare: {sum: 'fare'}}) + .compile(graph); + + const firstCount = device.createBuffer({ + id: 'ludf-first-group-count', + byteLength: 4 * Uint32Array.BYTES_PER_ELEMENT, + usage: Buffer.COPY_DST | Buffer.COPY_SRC + }); + const firstTotal = device.createBuffer({ + id: 'ludf-first-group-total', + byteLength: 4 * Float32Array.BYTES_PER_ELEMENT, + usage: Buffer.COPY_DST | Buffer.COPY_SRC + }); + + try { + const commandEncoder = device.createCommandEncoder({id: 'ludf-two-group-encodes'}); + compiled.encode(commandEncoder, {minimumFare: 15}); + commandEncoder.copyBufferToBuffer({ + sourceBuffer: getGPUDataBuffer(compiled.table.gpuVectors.count.data[0]), + destinationBuffer: firstCount, + size: 4 * Uint32Array.BYTES_PER_ELEMENT + }); + commandEncoder.copyBufferToBuffer({ + sourceBuffer: getGPUDataBuffer(compiled.table.gpuVectors.totalFare.data[0]), + destinationBuffer: firstTotal, + size: 4 * Float32Array.BYTES_PER_ELEMENT + }); + compiled.encode(commandEncoder, {minimumFare: 25}); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + await readUint32Buffer(firstCount, 4), + [1, 1, 0, 0], + 'the first encoding preserves its own filtered dense group counts' + ); + testContext.deepEqual( + await readFloat32Buffer(firstTotal, 4), + [30, 20, 0, 0], + 'the first encoding snapshots grouped sums before the parameter changes' + ); + testContext.deepEqual( + await readUint32VectorChunks(compiled.table.gpuVectors.count), + [[1, 0, 0, 0]], + 'the second encoding reuses the graph with a stricter selection threshold' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.totalFare), + [[30, 0, 0, 0]], + 'grouped statistic kernels observe encoder-ordered parameter uploads' + ); + + const validity = compiled.validity.totalFare; + if (!validity) { + throw new Error('Expected nullable grouped total validity'); + } + testContext.deepEqual( + await readUint32VectorChunks(validity), + [[1, 0, 0, 0]], + 'group validity is recomputed when dynamic filter parameters change' + ); + } finally { + firstCount.destroy(); + firstTotal.destroy(); + compiled.destroy(); + fixture.frame.destroy(); + } + + testContext.end(); +}); + +test('LuDataFrame groups chained nullable derived values without materializing hidden sources', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuGroupedFixture(device); + const graph = new GPUCommandGraph(device, { + id: 'ludf-derived-group-aggregation' + }); + const compiled = fixture.frame + .withColumn('doubleFare', column('fare').multiply(literal(2)), {format: 'float32'}) + .groupBy('category') + .aggregate({count: 'count', totalFare: {sum: 'doubleFare'}, averageFare: {mean: 'doubleFare'}}) + .compile(graph); + + try { + const commandEncoder = device.createCommandEncoder({id: 'ludf-derived-group-encode'}); + compiled.encode(commandEncoder); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + compiled.table.schema.fields.map(field => field.name), + ['category', 'count', 'totalFare', 'averageFare'], + 'derived grouping publishes only the key and requested aggregate aliases' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.totalFare), + [[80, 40, 0, 0]], + 'GPU grouping consumes nullable derived values across preserved source batches' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.averageFare), + [[40, 40, NaN, NaN]], + 'derived group means retain empty-group payload semantics' + ); + + const validity = compiled.validity.totalFare; + if (!validity) { + throw new Error('Expected derived group validity'); + } + testContext.deepEqual( + await readUint32VectorChunks(validity), + [[1, 1, 0, 0]], + 'nullable derived source sidecars propagate into dense grouped validity' + ); + } finally { + compiled.destroy(); + fixture.frame.destroy(); + } + + testContext.end(); +}); + +test('LuDataFrame excludes nonfinite metric values without changing categorical row counts', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuGroupedFixture(device); + const fare = fixture.frame.table.gpuVectors.fare; + getGPUDataBuffer(fare.data[0]).write(Float32Array.from([10, Number.NaN])); + getGPUDataBuffer(fare.data[2]).write(Float32Array.from([Number.POSITIVE_INFINITY, 99, 50])); + + const graph = new GPUCommandGraph(device, { + id: 'ludf-nonfinite-group-aggregation' + }); + const compiled = fixture.frame + .groupBy('category') + .aggregate({ + count: 'count', + totalFare: {sum: 'fare'}, + minimumFare: {min: 'fare'}, + averageFare: {mean: 'fare'} + }) + .compile(graph); + + try { + const commandEncoder = device.createCommandEncoder({id: 'ludf-nonfinite-group-encode'}); + compiled.encode(commandEncoder); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + await readUint32VectorChunks(compiled.table.gpuVectors.count), + [[2, 2, 0, 0]], + 'categorical count excludes only invalid group keys, not NaN or infinite metric values' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.totalFare), + [[10, 0, 0, 0]], + 'floating-point group sums discard NaN, infinity, and explicit null values' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.minimumFare), + [[10, NaN, NaN, NaN]], + 'groups with only nonfinite contributions retain invalid minimum payloads' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.averageFare), + [[10, NaN, NaN, NaN]], + 'floating-point means divide only finite and explicitly valid contributions' + ); + + for (const name of ['totalFare', 'minimumFare', 'averageFare'] as const) { + const validity = compiled.validity[name]; + if (!validity) { + throw new Error(`Expected explicit finite-value group validity for ${name}`); + } + testContext.deepEqual( + await readUint32VectorChunks(validity), + [[1, 0, 0, 0]], + `${name} excludes categories populated only by NaN, infinity, or null values` + ); + } + } finally { + compiled.destroy(); + fixture.frame.destroy(); + } + + testContext.end(); +}); + +test('LuDataFrame initializes dictionary groups when source tables have no record batches', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const frame = new LuDataFrame({ + table: new GPUTable({ + schema: { + fields: [ + {name: 'category', format: 'uint32', nullable: false}, + {name: 'fare', format: 'float32', nullable: true}, + {name: 'distance', format: 'sint32', nullable: false} + ], + metadata: new Map([['dataset', 'empty-groups']]) + }, + bufferLayout: [ + {name: 'category', format: 'uint32', byteStride: 4}, + {name: 'fare', format: 'float32', byteStride: 4}, + {name: 'distance', format: 'sint32', byteStride: 4} + ] + }), + dictionaries: { + category: {values: ['economy', 'standard', 'premium'], ordered: false} + }, + ownership: 'owned' + }); + const graph = new GPUCommandGraph(device, { + id: 'ludf-schema-only-group-aggregation' + }); + const compiled = frame + .groupBy('category') + .aggregate({count: 'count', totalFare: {sum: 'fare'}, averageFare: {mean: 'fare'}}) + .compile(graph); + + try { + const commandEncoder = device.createCommandEncoder({id: 'ludf-empty-group-encode'}); + compiled.encode(commandEncoder); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + compiled.table.batches.map(batch => batch.numRows), + [3], + 'schema-only sources still publish one dense categorical result batch' + ); + testContext.deepEqual( + await readUint32VectorChunks(compiled.table.gpuVectors.category), + [[0, 1, 2]], + 'empty source tables initialize every dictionary group key' + ); + testContext.deepEqual( + await readUint32VectorChunks(compiled.table.gpuVectors.count), + [[0, 0, 0]], + 'empty source tables publish deterministic zero group counts' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.totalFare), + [[0, 0, 0]], + 'empty source groups retain zero sum payloads' + ); + testContext.deepEqual( + await readFloat32VectorChunks(compiled.table.gpuVectors.averageFare), + [[NaN, NaN, NaN]], + 'empty source groups retain NaN mean payloads' + ); + for (const name of ['totalFare', 'averageFare'] as const) { + const validity = compiled.validity[name]; + if (!validity) { + throw new Error(`Expected explicit empty group validity for ${name}`); + } + testContext.deepEqual( + await readUint32VectorChunks(validity), + [[0, 0, 0]], + `${name} marks every schema-only source group invalid` + ); + } + testContext.deepEqual( + compiled.selectedCounts.data, + [], + 'schema-only grouping does not invent source selection batches' + ); + testContext.equal( + compiled.table.schema.metadata.get('dataset'), + 'empty-groups', + 'dense grouped schemas retain independent source metadata' + ); + } finally { + compiled.destroy(); + frame.destroy(); + } + + testContext.end(); +}); + +function createLuGroupedFixture(device: Device): LuGroupedFixture { + const sourceBuffers: Buffer[] = []; + const categoryValues = [ + Uint32Array.from([0, 1]), + new Uint32Array(0), + Uint32Array.from([0, 2, 1]) + ]; + const categoryValidityValues = [ + Uint32Array.from([1, 1]), + new Uint32Array(0), + Uint32Array.from([1, 0, 1]) + ]; + const fareValues = [ + Float32Array.from([10, 20]), + new Float32Array(0), + Float32Array.from([30, 99, 50]) + ]; + const fareValidityValues = [ + Uint32Array.from([1, 1]), + new Uint32Array(0), + Uint32Array.from([1, 1, 0]) + ]; + const distanceValues = [Int32Array.from([-2, 3]), new Int32Array(0), Int32Array.from([-1, 4, 8])]; + const categoryValidityChunks: GPUData<'uint32'>[] = []; + const fareValidityChunks: GPUData<'uint32'>[] = []; + let sourceRowIndexOffset = 40; + + const batches = categoryValues.map((values, batchIndex) => { + const batch = new GPURecordBatch({ + gpuData: { + category: createLuGroupedData(device, sourceBuffers, values, 'uint32'), + fare: createLuGroupedData(device, sourceBuffers, fareValues[batchIndex], 'float32'), + distance: createLuGroupedData(device, sourceBuffers, distanceValues[batchIndex], 'sint32') + }, + fields: [ + {name: 'category', format: 'uint32', nullable: true}, + {name: 'fare', format: 'float32', nullable: true}, + {name: 'distance', format: 'sint32', nullable: false} + ], + sourceInfo: { + sourceBatchIndex: batchIndex + 4, + sourceRowIndexOffset, + sourceRowCount: values.length + } + }); + sourceRowIndexOffset += values.length; + categoryValidityChunks.push( + createLuGroupedData(device, sourceBuffers, categoryValidityValues[batchIndex], 'uint32') + ); + fareValidityChunks.push( + createLuGroupedData(device, sourceBuffers, fareValidityValues[batchIndex], 'uint32') + ); + return batch; + }); + + return { + frame: new LuDataFrame({ + table: new GPUTable({batches}), + validity: { + category: new GPUVector<'uint32'>({ + type: 'data', + name: 'ludf-group-category-validity', + format: 'uint32', + data: categoryValidityChunks, + ownsData: true + }), + fare: new GPUVector<'uint32'>({ + type: 'data', + name: 'ludf-group-fare-validity', + format: 'uint32', + data: fareValidityChunks, + ownsData: true + }) + }, + dictionaries: { + category: {values: ['economy', 'standard', 'premium', 'unused'], ordered: false} + }, + ownership: 'owned' + }), + sourceBuffers + }; +} + +function createLuGroupedData( + device: Device, + sourceBuffers: Buffer[], + values: Float32Array | Int32Array | 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 getGPUDataBuffer(data: GPUData): Buffer { + return data.buffer instanceof Buffer ? data.buffer : data.buffer.buffer; +} + +async function readUint32Buffer(buffer: Buffer, length: number): Promise { + if (length === 0) { + return []; + } + const values = await buffer.readAsync(0, length * Uint32Array.BYTES_PER_ELEMENT); + return Array.from(new Uint32Array(values.buffer, values.byteOffset, length)); +} + +async function readFloat32Buffer(buffer: Buffer, length: number): Promise { + if (length === 0) { + return []; + } + const values = await buffer.readAsync(0, length * Float32Array.BYTES_PER_ELEMENT); + return Array.from(new Float32Array(values.buffer, values.byteOffset, length)); +} + +async function readUint32VectorChunks(vector: GPUVector): Promise { + return Promise.all( + vector.data.map(chunk => readUint32Buffer(getGPUDataBuffer(chunk), chunk.length)) + ); +} + +async function readFloat32VectorChunks(vector: GPUVector): Promise { + return Promise.all( + vector.data.map(chunk => readFloat32Buffer(getGPUDataBuffer(chunk), chunk.length)) + ); +}