diff --git a/modules/experimental/src/ludf/index.ts b/modules/experimental/src/ludf/index.ts index 3378902852..34e2f72f30 100644 --- a/modules/experimental/src/ludf/index.ts +++ b/modules/experimental/src/ludf/index.ts @@ -39,6 +39,8 @@ export type { } from './lu-global-aggregation-query'; 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 {and, column, literal, LuExpression, not, or, parameter} from './lu-expression'; export type { LuExpressionBinaryOperator, @@ -51,3 +53,4 @@ export type {LuDataFrameQueryParameters} from './lu-query-compiler'; export {CompiledLuDataFrameGroupedAggregation} from './lu-group-aggregation-compiler'; export {CompiledLuDataFrameAggregation} from './lu-global-aggregation-compiler'; export {CompiledLuDataFrameHistogram} from './lu-histogram-compiler'; +export {CompiledLuDataFrameSort} from './lu-sort-compiler'; diff --git a/modules/experimental/src/ludf/lu-data-frame-query.ts b/modules/experimental/src/ludf/lu-data-frame-query.ts index cb0038a750..c6f939a307 100644 --- a/modules/experimental/src/ludf/lu-data-frame-query.ts +++ b/modules/experimental/src/ludf/lu-data-frame-query.ts @@ -22,6 +22,7 @@ import { type CompiledLuDataFrameQuery, type LuDataFrameQueryParameters } from './lu-query-compiler'; +import {LuDataFrameSortQuery, type LuDataFrameSortOptions} from './lu-sort-query'; /** Portable scalar storage formats supported by computed dataframe columns. */ export type LuDataFrameDerivedColumnFormat = 'float32' | 'sint32' | 'uint32'; @@ -178,6 +179,23 @@ export class LuDataFrameQuery< return new LuDataFrameHistogramQuery(this, column, options); } + /** Plans stable source-batch scalar sorting without allocating GPU resources or reading rows. */ + sortBy>( + column: Column, + options: LuDataFrameSortOptions = {} + ): LuDataFrameSortQuery { + return new LuDataFrameSortQuery(this, column, options); + } + + /** Plans descending stable top-K selection independently within every source record batch. */ + topK>( + column: Column, + limit: number, + options: LuDataFrameSortOptions = {} + ): LuDataFrameSortQuery { + return new LuDataFrameSortQuery(this, column, options, limit, 'descending'); + } + /** 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 c95874d56f..5258d5a858 100644 --- a/modules/experimental/src/ludf/lu-data-frame.ts +++ b/modules/experimental/src/ludf/lu-data-frame.ts @@ -33,6 +33,7 @@ import type { LuDataFrameGroupByQuery } from './lu-group-by-query'; import type {LuDataFrameHistogramOptions, LuDataFrameHistogramQuery} from './lu-histogram-query'; +import type {LuDataFrameSortOptions, LuDataFrameSortQuery} from './lu-sort-query'; /** Whether a dataframe borrows its source resources or releases them after its final view. */ export type LuDataFrameOwnership = 'borrowed' | 'owned'; @@ -236,6 +237,32 @@ export class LuDataFrame { ); } + /** Plans stable scalar ordering independently within every existing source record batch. */ + sortBy>( + column: Column, + options: LuDataFrameSortOptions = {} + ): LuDataFrameSortQuery { + this.assertAvailable(); + return new LuDataFrameQuery(this, [], this.columnNames).sortBy( + column, + options + ); + } + + /** Plans descending stable top-K selection without concatenating source record batches. */ + topK>( + column: Column, + limit: number, + options: LuDataFrameSortOptions = {} + ): LuDataFrameSortQuery { + this.assertAvailable(); + return new LuDataFrameQuery(this, [], this.columnNames).topK( + column, + limit, + options + ); + } + /** * Returns an independent borrowed projection without mutating or destroying source columns. * diff --git a/modules/experimental/src/ludf/lu-query-compiler.ts b/modules/experimental/src/ludf/lu-query-compiler.ts index 7e968609f2..a1ab03d819 100644 --- a/modules/experimental/src/ludf/lu-query-compiler.ts +++ b/modules/experimental/src/ludf/lu-query-compiler.ts @@ -94,6 +94,8 @@ export type LuDataFrameQueryExtensionContext = { validity: Readonly>; dictionaries: Readonly>; selectionMask: GraphVectorView<'uint32'>; + rowIndices: GraphVectorView<'uint32'>; + selectedCounts: GraphVectorView<'uint32'>; }; /** Result resources contributed by one graph-native extension. @internal */ @@ -342,7 +344,9 @@ export function compileLuDataFrameQuery< table: rowTable, validity, dictionaries, - selectionMask: maskView + selectionMask: maskView, + rowIndices: rowIndexView, + selectedCounts: countView }); } diff --git a/modules/experimental/src/ludf/lu-sort-compiler.ts b/modules/experimental/src/ludf/lu-sort-compiler.ts new file mode 100644 index 0000000000..a40fbb3b86 --- /dev/null +++ b/modules/experimental/src/ludf/lu-sort-compiler.ts @@ -0,0 +1,459 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors + +import type {GPUTypeMap} from '@luma.gl/tables'; +import {GPUBatchSort} from '../gpu-primitives/gpu-batch-sort'; +import { + type GPUCommandGraph, + type GraphBufferUse, + type GraphDataView +} from '../gpu-primitives/gpu-command-graph'; +import { + createTransientVectorView, + getViewElementOffset +} from '../gpu-primitives/graph-data-view-utils'; +import { + LU_ANALYTICS_WORKGROUP_SIZE, + addLuAnalyticsComputePass, + getLuAnalyticsShaderType, + getLuAnalyticsVector, + validateLuAnalyticsSource, + type LuAnalyticsScalarFormat +} 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 { + CompiledLuDataFrameQuery, + compileLuDataFrameQuery, + type CompiledLuDataFrameQueryProps, + type LuDataFrameQueryExtensionContext, + type LuDataFrameQueryExtensionResult, + type LuDataFrameQueryParameters +} from './lu-query-compiler'; +import type {LuDataFrameNormalizedSortOptions} from './lu-sort-query'; + +const MAXIMUM_UINT32 = 0xffffffff; + +type LuSortChunk = { + input: GraphDataView; + selection: GraphDataView<'uint32'>; + validity?: GraphDataView<'uint32'>; + encodedKeys: GraphDataView<'uint32'>; + localIndices: GraphDataView<'uint32'>; + sortedKeys: GraphDataView<'uint32'>; + sortedIndices: GraphDataView<'uint32'>; +}; + +/** Stable, source-batch-preserving numeric row ordering with optional per-batch top-K limits. */ +export class CompiledLuDataFrameSort< + T extends GPUTypeMap = GPUTypeMap +> extends CompiledLuDataFrameQuery { + /** Selected numeric source or derived column used to produce the row permutation. */ + readonly sortColumn: keyof T & string; + /** Final order of valid numeric source values. */ + readonly direction: LuDataFrameNormalizedSortOptions['direction']; + /** Placement of explicitly null source rows, outside the numeric/NaN groups. */ + readonly nulls: LuDataFrameNormalizedSortOptions['nulls']; + /** Placement of NaN source rows within the non-null group. */ + readonly nans: LuDataFrameNormalizedSortOptions['nans']; + /** Stable sort algorithm requested independently for every source batch. */ + readonly algorithm: LuDataFrameNormalizedSortOptions['algorithm']; + /** Optional maximum number of retained rows in each original source batch. */ + readonly limit?: number; + + /** @internal */ + constructor( + props: CompiledLuDataFrameQueryProps, + column: keyof T & string, + options: LuDataFrameNormalizedSortOptions + ) { + super(props); + this.sortColumn = column; + this.direction = options.direction; + this.nulls = options.nulls; + this.nans = options.nans; + this.algorithm = options.algorithm; + this.limit = options.limit; + } +} + +/** Adds stable numeric ordering to one reusable filter/derived graph without repacking batches. */ +export function compileLuDataFrameSort( + source: LuDataFrame, + predicates: readonly LuExpression[], + selectedColumns: readonly (keyof Selection & string)[], + derivedColumns: readonly LuDataFrameDerivedColumn[], + column: keyof Selection & string, + options: LuDataFrameNormalizedSortOptions, + graph: GPUCommandGraph +): CompiledLuDataFrameSort { + validateLuAnalyticsSource(source, [column]); + validateLuSortSourceOffsets(source); + return compileLuDataFrameQuery>( + source, + predicates, + selectedColumns, + graph, + derivedColumns, + { + allowEmptyPredicates: true, + prepare: context => addLuBatchSortToGraph(context, column, options) + } + ); +} + +/** Ensures stable source-row identifiers remain exactly representable by uint32 output vectors. */ +function validateLuSortSourceOffsets(source: LuDataFrame): void { + let sourceOffset = 0; + for (const batch of source.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 sort source-row identities must fit uint32'); + } + sourceOffset += batch.numRows; + } +} + +/** Composes one stable full-width numeric sort and one stable selected/null/NaN class sort. */ +function addLuBatchSortToGraph( + context: LuDataFrameQueryExtensionContext, + column: keyof Selection & string, + options: LuDataFrameNormalizedSortOptions +): LuDataFrameQueryExtensionResult> { + const graph = context.graph; + const id = `${context.queryId}-sort`; + const vector = getLuAnalyticsVector(context, column); + const input = graph.importGPUVector(`${id}-input`, vector); + const field = context.table.schema.fields.find(candidate => candidate.name === column); + const validityVector = field?.nullable ? context.validity[column] : undefined; + if (field?.nullable && !validityVector && input.length > 0) { + throw new Error(`LuDataFrame nullable sort column "${column}" requires GPU validity`); + } + const validity = validityVector + ? graph.importGPUVector(`${id}-source-validity`, validityVector) + : undefined; + + const encodedKeys = createTransientVectorView(graph, `${id}-encoded-keys`, context.selectionMask); + const localIndices = createTransientVectorView( + graph, + `${id}-local-indices`, + context.selectionMask + ); + const sortedKeys = createTransientVectorView(graph, `${id}-sorted-keys`, context.selectionMask); + const sortedIndices = createTransientVectorView( + graph, + `${id}-sorted-indices`, + context.selectionMask + ); + + for (const [batchIndex, values] of input.data.entries()) { + if (values.length === 0) { + continue; + } + addLuSortEncodeKeysPass(graph, `${id}-encode-batch-${batchIndex}`, { + input: values, + selection: context.selectionMask.data[batchIndex], + validity: validity?.data[batchIndex], + encodedKeys: encodedKeys.data[batchIndex], + localIndices: localIndices.data[batchIndex], + sortedKeys: sortedKeys.data[batchIndex], + sortedIndices: sortedIndices.data[batchIndex] + }); + } + + new GPUBatchSort({ + id: `${id}-numeric`, + keys: encodedKeys, + values: localIndices, + outputKeys: sortedKeys, + outputValues: sortedIndices, + direction: options.direction, + algorithm: options.algorithm + }).addToGraph(graph); + + for (const [batchIndex, values] of input.data.entries()) { + if (values.length === 0) { + continue; + } + addLuSortEncodeClassesPass( + graph, + `${id}-classify-batch-${batchIndex}`, + { + input: values, + selection: context.selectionMask.data[batchIndex], + validity: validity?.data[batchIndex], + encodedKeys: encodedKeys.data[batchIndex], + localIndices: localIndices.data[batchIndex], + sortedKeys: sortedKeys.data[batchIndex], + sortedIndices: sortedIndices.data[batchIndex] + }, + options + ); + } + + new GPUBatchSort({ + id: `${id}-classes`, + keys: sortedKeys, + values: sortedIndices, + outputKeys: encodedKeys, + outputValues: localIndices, + direction: 'ascending', + algorithm: options.algorithm + }).addToGraph(graph); + + let sourceOffset = 0; + for (const [batchIndex, batch] of context.table.batches.entries()) { + if (batch.numRows > 0) { + addLuSortPublishPermutationPass(graph, `${id}-publish-batch-${batchIndex}`, { + sortedClasses: encodedKeys.data[batchIndex], + sortedIndices: localIndices.data[batchIndex], + selection: context.selectionMask.data[batchIndex], + outputIndices: context.rowIndices.data[batchIndex], + selectedCount: context.selectedCounts.data[batchIndex], + sourceOffset: batch.sourceInfo?.sourceRowIndexOffset ?? sourceOffset, + limit: options.limit ?? MAXIMUM_UINT32 + }); + } + sourceOffset += batch.numRows; + } + + return { + table: context.table, + validity: context.validity, + dictionaries: context.dictionaries, + createCompiled: props => new CompiledLuDataFrameSort(props, column, options) + }; +} + +/** Encodes signed/floating order into the complete unsigned domain and preserves stable ties. */ +function addLuSortEncodeKeysPass( + graph: GPUCommandGraph, + id: string, + chunk: LuSortChunk +): void { + const shaderType = getLuAnalyticsShaderType(chunk.input.format); + const nullable = Boolean(chunk.validity); + const validityBinding = nullable + ? '@group(0) @binding(1) var validityValues: array;' + : ''; + const keyBindingIndex = nullable ? 2 : 1; + const validityOffset = chunk.validity ? getViewElementOffset(chunk.validity) : 0; + const isNull = nullable ? 'validityValues[VALIDITY_OFFSET + index] == 0u' : 'false'; + const nanExpression = + chunk.input.format === 'float32' + ? '(bitcast(value) & 0x7fffffffu) > 0x7f800000u' + : 'false'; + const key = getLuSortNumericKeyExpression(chunk.input.format); + const source = /* wgsl */ ` +const ELEMENT_COUNT: u32 = ${chunk.input.length}u; +const INPUT_OFFSET: u32 = ${getViewElementOffset(chunk.input)}u; +const VALIDITY_OFFSET: u32 = ${validityOffset}u; +const KEY_OFFSET: u32 = ${getViewElementOffset(chunk.encodedKeys)}u; +const INDEX_OFFSET: u32 = ${getViewElementOffset(chunk.localIndices)}u; +@group(0) @binding(0) var inputValues: array<${shaderType}>; +${validityBinding} +@group(0) @binding(${keyBindingIndex}) var outputKeys: array; +@group(0) @binding(${keyBindingIndex + 1}) var outputIndices: 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 value = inputValues[INPUT_OFFSET + index]; + let isNull = ${isNull}; + let isNaN = ${nanExpression}; + outputKeys[KEY_OFFSET + index] = select(${key}, 0u, isNull || isNaN); + outputIndices[INDEX_OFFSET + index] = index; + } +}`; + const resources: GraphBufferUse[] = [{buffer: chunk.input, usage: 'storage-read'}]; + const bindings: Record = {inputValues: chunk.input}; + if (chunk.validity) { + resources.push({buffer: chunk.validity, usage: 'storage-read'}); + bindings['validityValues'] = chunk.validity; + } + resources.push( + {buffer: chunk.encodedKeys, usage: 'storage-write'}, + {buffer: chunk.localIndices, usage: 'storage-write'} + ); + bindings['outputKeys'] = chunk.encodedKeys; + bindings['outputIndices'] = chunk.localIndices; + addLuAnalyticsComputePass(graph, { + id, + source, + resources, + bindings, + length: chunk.input.length + }); +} + +/** Maps numeric storage formats into monotonic unsigned keys without dropping any value bits. */ +function getLuSortNumericKeyExpression(format: LuAnalyticsScalarFormat): string { + switch (format) { + case 'uint32': + return 'value'; + case 'sint32': + return 'bitcast(value) ^ 0x80000000u'; + case 'float32': + return 'select(bitcast(value), 0u, value == 0.0) ^ select(0x80000000u, 0xffffffffu, (bitcast(value) & 0x80000000u) != 0u && value != 0.0)'; + } +} + +/** Stably ranks selected numeric, NaN, null, and rejected rows after the numeric permutation. */ +function addLuSortEncodeClassesPass( + graph: GPUCommandGraph, + id: string, + chunk: LuSortChunk, + options: LuDataFrameNormalizedSortOptions +): void { + const shaderType = getLuAnalyticsShaderType(chunk.input.format); + const nullable = Boolean(chunk.validity); + const validityBinding = nullable + ? '@group(0) @binding(3) var validityValues: array;' + : ''; + const outputBindingIndex = nullable ? 4 : 3; + const validityOffset = chunk.validity ? getViewElementOffset(chunk.validity) : 0; + const isNull = nullable ? 'validityValues[VALIDITY_OFFSET + localIndex] == 0u' : 'false'; + const nanExpression = + chunk.input.format === 'float32' + ? '(bitcast(value) & 0x7fffffffu) > 0x7f800000u' + : 'false'; + const nullRank = options.nulls === 'first' ? 0 : 2; + const nanRank = + options.nulls === 'first' + ? options.nans === 'first' + ? 1 + : 2 + : options.nans === 'first' + ? 0 + : 1; + const numericRank = + options.nulls === 'first' + ? options.nans === 'first' + ? 2 + : 1 + : options.nans === 'first' + ? 1 + : 0; + const source = /* wgsl */ ` +const ELEMENT_COUNT: u32 = ${chunk.input.length}u; +const INPUT_OFFSET: u32 = ${getViewElementOffset(chunk.input)}u; +const SELECTION_OFFSET: u32 = ${getViewElementOffset(chunk.selection)}u; +const INDEX_OFFSET: u32 = ${getViewElementOffset(chunk.sortedIndices)}u; +const VALIDITY_OFFSET: u32 = ${validityOffset}u; +const OUTPUT_OFFSET: u32 = ${getViewElementOffset(chunk.sortedKeys)}u; +@group(0) @binding(0) var inputValues: array<${shaderType}>; +@group(0) @binding(1) var selectionMask: array; +@group(0) @binding(2) var sortedIndices: array; +${validityBinding} +@group(0) @binding(${outputBindingIndex}) var outputClasses: 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 localIndex = sortedIndices[INDEX_OFFSET + index]; + let value = inputValues[INPUT_OFFSET + localIndex]; + let isNull = ${isNull}; + let isNaN = ${nanExpression}; + var rank = ${numericRank}u; + if (isNaN) { rank = ${nanRank}u; } + if (isNull) { rank = ${nullRank}u; } + if (selectionMask[SELECTION_OFFSET + localIndex] == 0u) { rank = 3u; } + outputClasses[OUTPUT_OFFSET + index] = rank; + } +}`; + const resources: GraphBufferUse[] = [ + {buffer: chunk.input, usage: 'storage-read'}, + {buffer: chunk.selection, usage: 'storage-read'}, + {buffer: chunk.sortedIndices, usage: 'storage-read'} + ]; + const bindings: Record = { + inputValues: chunk.input, + selectionMask: chunk.selection, + sortedIndices: chunk.sortedIndices + }; + if (chunk.validity) { + resources.push({buffer: chunk.validity, usage: 'storage-read'}); + bindings['validityValues'] = chunk.validity; + } + resources.push({buffer: chunk.sortedKeys, usage: 'storage-write'}); + bindings['outputClasses'] = chunk.sortedKeys; + addLuAnalyticsComputePass(graph, { + id, + source, + resources, + bindings, + length: chunk.input.length + }); +} + +/** Publishes stable source identities and keeps masks/counts consistent with each batch's top-K. */ +function addLuSortPublishPermutationPass( + graph: GPUCommandGraph, + id: string, + props: { + sortedClasses: GraphDataView<'uint32'>; + sortedIndices: GraphDataView<'uint32'>; + selection: GraphDataView<'uint32'>; + outputIndices: GraphDataView<'uint32'>; + selectedCount: GraphDataView<'uint32'>; + sourceOffset: number; + limit: number; + } +): void { + const source = /* wgsl */ ` +const ELEMENT_COUNT: u32 = ${props.sortedIndices.length}u; +const CLASS_OFFSET: u32 = ${getViewElementOffset(props.sortedClasses)}u; +const LOCAL_INDEX_OFFSET: u32 = ${getViewElementOffset(props.sortedIndices)}u; +const MASK_OFFSET: u32 = ${getViewElementOffset(props.selection)}u; +const OUTPUT_OFFSET: u32 = ${getViewElementOffset(props.outputIndices)}u; +const COUNT_OFFSET: u32 = ${getViewElementOffset(props.selectedCount)}u; +const SOURCE_OFFSET: u32 = ${props.sourceOffset}u; +const ROW_LIMIT: u32 = ${props.limit}u; +@group(0) @binding(0) var sortedClasses: array; +@group(0) @binding(1) var sortedIndices: array; +@group(0) @binding(2) var selectionMask: array; +@group(0) @binding(3) var outputIndices: array; +@group(0) @binding(4) var selectedCounts: 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 localIndex = sortedIndices[LOCAL_INDEX_OFFSET + index]; + let selected = sortedClasses[CLASS_OFFSET + index] != 3u && index < ROW_LIMIT; + selectionMask[MASK_OFFSET + localIndex] = select(0u, 1u, selected); + outputIndices[OUTPUT_OFFSET + index] = select(0u, SOURCE_OFFSET + localIndex, selected); + if (index == 0u) { + selectedCounts[COUNT_OFFSET] = min(selectedCounts[COUNT_OFFSET], ROW_LIMIT); + } + } +}`; + addLuAnalyticsComputePass(graph, { + id, + source, + resources: [ + {buffer: props.sortedClasses, usage: 'storage-read'}, + {buffer: props.sortedIndices, usage: 'storage-read'}, + {buffer: props.selection, usage: 'storage-write'}, + {buffer: props.outputIndices, usage: 'storage-write'}, + {buffer: props.selectedCount, usage: 'storage-read-write'} + ], + bindings: { + sortedClasses: props.sortedClasses, + sortedIndices: props.sortedIndices, + selectionMask: props.selection, + outputIndices: props.outputIndices, + selectedCounts: props.selectedCount + }, + length: props.sortedIndices.length + }); +} diff --git a/modules/experimental/src/ludf/lu-sort-query.ts b/modules/experimental/src/ludf/lu-sort-query.ts new file mode 100644 index 0000000000..57fd20ab53 --- /dev/null +++ b/modules/experimental/src/ludf/lu-sort-query.ts @@ -0,0 +1,143 @@ +// 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 {LuDataFrameQuery} from './lu-data-frame-query'; +import { + getLuDataFrameAnalyticColumnFormat, + type LuDataFrameScalarColumnNames +} from './lu-global-aggregation-query'; +import type {LuDataFrameQueryParameters} from './lu-query-compiler'; +import {compileLuDataFrameSort, type CompiledLuDataFrameSort} from './lu-sort-compiler'; + +const MAXIMUM_UINT32 = 0xffffffff; + +/** Stable per-batch scalar ordering, explicit null/NaN placement, and GPU sort implementation. */ +export type LuDataFrameSortOptions = Readonly<{ + /** Ascending or descending finite/infinite numeric ordering. */ + direction?: 'ascending' | 'descending'; + /** Absolute placement of explicitly invalid rows among accepted source rows. */ + nulls?: 'first' | 'last'; + /** Placement of NaN among valid rows, independently of source null ordering. */ + nans?: 'first' | 'last'; + /** Optional existing stable GPU sorting implementation. */ + algorithm?: 'auto' | 'bitonic' | 'radix'; +}>; + +/** Closed, immutable sort controls consumed directly by the graph-native sort compiler. @internal */ +export type LuDataFrameNormalizedSortOptions = Readonly<{ + direction: 'ascending' | 'descending'; + nulls: 'first' | 'last'; + nans: 'first' | 'last'; + algorithm: 'auto' | 'bitonic' | 'radix'; + limit?: number; +}>; + +/** + * Immutable stable scalar sorting or top-K plan. + * + * Each existing source batch is sorted independently: no source columns are copied, concatenated, + * or repacked. Filtered rows always remain after accepted values, null placement is absolute, and + * NaN placement applies only within the non-null portion of each batch. + */ +export class LuDataFrameSortQuery< + Logical extends GPUTypeMap, + SelectedColumns extends keyof Logical & string, + Column extends LuDataFrameScalarColumnNames, + Source extends GPUTypeMap = Logical +> { + /** Original filtered, projected, or derived dataframe query. */ + readonly query: LuDataFrameQuery; + /** Selected scalar column supplying stable numeric order. */ + readonly column: Column; + /** Immutable, closed ordering controls and optional per-batch result limit. */ + readonly options: LuDataFrameNormalizedSortOptions; + + /** Validates stable sorting metadata without allocating or retaining GPU resources. @internal */ + constructor( + query: LuDataFrameQuery, + column: Column, + options: LuDataFrameSortOptions = {}, + limit?: number, + defaultDirection: 'ascending' | 'descending' = 'ascending' + ) { + if (!query.selectedColumns.includes(column)) { + throw new Error(`LuDataFrame sort column "${column}" is not selected`); + } + if (!getLuDataFrameAnalyticColumnFormat(query, column)) { + throw new Error(`LuDataFrame sort column "${column}" requires scalar GPU data`); + } + if (query.source.table.gpuConstants[column]) { + throw new Error(`LuDataFrame sort column "${column}" cannot be a constant`); + } + + this.query = query; + this.column = column; + this.options = normalizeLuDataFrameSortOptions(options, limit, defaultDirection); + Object.freeze(this); + } + + /** Restricts each source batch independently while preserving this plan's ordering controls. */ + topK(limit: number): LuDataFrameSortQuery { + return new LuDataFrameSortQuery(this.query, this.column, this.options, limit); + } + + /** Adds stable source-row sorting and optional per-batch limiting to one reusable GPU graph. */ + compile( + graph: GPUCommandGraph + ): CompiledLuDataFrameSort> { + return compileLuDataFrameSort>( + this.query.source, + this.query.predicates, + this.query.selectedColumns, + this.query.derivedColumns, + this.column, + this.options, + graph + ); + } +} + +/** Clones closed ordering controls and rejects limits that cannot be represented by GPU counts. */ +function normalizeLuDataFrameSortOptions( + options: LuDataFrameSortOptions, + limit: number | undefined, + defaultDirection: 'ascending' | 'descending' +): LuDataFrameNormalizedSortOptions { + if (!options || typeof options !== 'object') { + throw new Error('LuDataFrame sorting options must be an object'); + } + + const direction = options.direction ?? defaultDirection; + const nulls = options.nulls ?? 'last'; + const nans = options.nans ?? 'last'; + const algorithm = options.algorithm ?? 'auto'; + if (direction !== 'ascending' && direction !== 'descending') { + throw new Error('LuDataFrame sort direction must be ascending or descending'); + } + if (nulls !== 'first' && nulls !== 'last') { + throw new Error('LuDataFrame sort nulls must be first or last'); + } + if (nans !== 'first' && nans !== 'last') { + throw new Error('LuDataFrame sort NaNs must be first or last'); + } + if (algorithm !== 'auto' && algorithm !== 'bitonic' && algorithm !== 'radix') { + throw new Error('LuDataFrame sort algorithm must be auto, bitonic, or radix'); + } + if ( + limit !== undefined && + (!Number.isSafeInteger(limit) || limit < 0 || limit > MAXIMUM_UINT32) + ) { + throw new Error('LuDataFrame top-K limits require a nonnegative uint32 count'); + } + + return Object.freeze({ + direction, + nulls, + nans, + algorithm, + ...(limit === undefined ? {} : {limit}) + }); +} diff --git a/modules/experimental/test/index.ts b/modules/experimental/test/index.ts index fed4dc005a..a02c9a5f76 100644 --- a/modules/experimental/test/index.ts +++ b/modules/experimental/test/index.ts @@ -44,6 +44,7 @@ import './ludf/lu-data-frame-query.spec'; import './ludf/lu-derived-columns.spec'; import './ludf/lu-group-aggregation.spec'; import './ludf/lu-global-aggregation.spec'; +import './ludf/lu-sort.spec'; import './luraster'; import './luxfilter'; import './luproj/luproj.spec'; diff --git a/modules/experimental/test/ludf/lu-sort.node.spec.ts b/modules/experimental/test/ludf/lu-sort.node.spec.ts new file mode 100644 index 0000000000..27668d6262 --- /dev/null +++ b/modules/experimental/test/ludf/lu-sort.node.spec.ts @@ -0,0 +1,346 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors + +import {Buffer} from '@luma.gl/core'; +import {GPUCommandGraph} from '@luma.gl/experimental'; +import { + column, + CompiledLuDataFrameSort, + literal, + LuDataFrame, + LuDataFrameSortQuery, + parameter, + type LuDataFrameQueryParameters, + type LuDataFrameSortOptions +} from '@luma.gl/experimental/ludf'; +import { + GPUConstant, + 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 SortSourceColumns = { + score: 'float32'; + signed: 'sint32'; + category: 'uint32'; + coordinates: 'float32x2'; +}; + +type SortSourceFixture = { + device: NullDevice; + table: GPUTable; + buffers: Buffer[]; +}; + +describe('LuDataFrame immutable stable scalar sorting', () => { + test('plans ascending source-batch ordering without GPU allocation or resource retention', () => { + const fixture = createSortSourceFixture([2, 0, 3]); + const source = new LuDataFrame({table: fixture.table, 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 sorted = source.sortBy('score'); + + expect(sorted).toBeInstanceOf(LuDataFrameSortQuery); + expect(sorted.query.source).toBe(source); + expect(sorted.column).toBe('score'); + expect(sorted.options).toEqual({ + direction: 'ascending', + nulls: 'last', + nans: 'last', + algorithm: 'auto' + }); + expect(Object.isFrozen(sorted)).toBe(true); + expect(Object.isFrozen(sorted.options)).toBe(true); + expect(source.batches.map(batch => batch.numRows)).toEqual([2, 0, 3]); + expect(source.sourceInfo.map(info => info?.sourceRowIndexOffset)).toEqual([40, 42, 42]); + 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('defaults direct top-K to descending while sorted-plan limiting preserves its direction', () => { + const fixture = createSortSourceFixture([2, 0, 3]); + const source = new LuDataFrame({table: fixture.table}); + + const largest = source.topK('score', 2); + const smallest = source.sortBy('score').topK(2); + const explicit = source.topK('signed', 0, { + direction: 'ascending', + nulls: 'first', + nans: 'first', + algorithm: 'radix' + }); + const unchanged = source.sortBy('category', {direction: 'descending', algorithm: 'bitonic'}); + const limited = unchanged.topK(0xffffffff); + + expect(largest.options).toEqual({ + direction: 'descending', + nulls: 'last', + nans: 'last', + algorithm: 'auto', + limit: 2 + }); + expect(smallest.options).toEqual({ + direction: 'ascending', + nulls: 'last', + nans: 'last', + algorithm: 'auto', + limit: 2 + }); + expect(explicit.options).toEqual({ + direction: 'ascending', + nulls: 'first', + nans: 'first', + algorithm: 'radix', + limit: 0 + }); + expect(unchanged.options).not.toHaveProperty('limit'); + expect(limited.options).toEqual({...unchanged.options, limit: 0xffffffff}); + + source.destroy(); + fixture.table.destroy(); + }); + + test('clones immutable options and keeps precise projected and derived source output types', () => { + const fixture = createSortSourceFixture([2]); + const source = new LuDataFrame({table: fixture.table}); + const options: {direction: 'ascending' | 'descending'; nulls: 'first' | 'last'} = { + direction: 'descending', + nulls: 'first' + }; + const sorted = source.sortBy('score', options); + options.direction = 'ascending'; + options.nulls = 'last'; + + const projected = source + .filter(column('score').greaterThan(literal(0))) + .select(['score', 'signed']) + .sortBy('signed'); + const derived = source + .withColumn('adjustedScore', column('score').add(literal(2))) + .select(['score', 'adjustedScore']) + .topK('adjustedScore', 3); + + expect(sorted.options.direction).toBe('descending'); + expect(sorted.options.nulls).toBe('first'); + expectTypeOf(projected.compile).returns.toEqualTypeOf< + CompiledLuDataFrameSort<{score: 'float32'; signed: 'sint32'}> + >(); + expectTypeOf(derived.compile).returns.toEqualTypeOf< + CompiledLuDataFrameSort<{score: 'float32'; adjustedScore: 'float32'}> + >(); + + source.destroy(); + fixture.table.destroy(); + }); + + test('preserves interaction parameters, hidden dependencies, and empty source-batch topology', () => { + for (const batchLengths of [[], [0], [2, 0, 3]] as const) { + const fixture = createSortSourceFixture(batchLengths, {nullableScore: true}); + const source = new LuDataFrame({table: fixture.table}); + const filtered = source.filter(column('score').greaterThan(parameter('minimumScore', 10))); + const adjusted = filtered.withColumn('adjustedScore', column('score').add(literal(2))); + const sorted = adjusted.select(['signed', 'adjustedScore']).topK('adjustedScore', 2, { + nulls: 'first', + nans: 'first' + }); + + expect(sorted.query.predicates[0]).toBe(filtered.predicates[0]); + expect(sorted.query.derivedColumns.map(({name}) => name)).toEqual(['adjustedScore']); + expect(sorted.query.selectedColumns).toEqual(['signed', 'adjustedScore']); + expect(sorted.query.source.batches.map(batch => batch.numRows)).toEqual(batchLengths); + expect(sorted.options.direction).toBe('descending'); + + source.destroy(); + fixture.table.destroy(); + } + }); + + test('rejects unknown, hidden, vector-valued, and constant sort keys before GPU work', () => { + const fixture = createSortSourceFixture([2]); + const source = new LuDataFrame({table: fixture.table}); + const createBuffer = vi.spyOn(fixture.device, 'createBuffer'); + + expect(() => + // @ts-expect-error Sort keys must be selected 32-bit scalar GPU columns. + source.sortBy('coordinates') + ).toThrow(/scalar|column/i); + expect(() => + // @ts-expect-error Sort keys must exist in the source dataframe. + source.sortBy('missing') + ).toThrow(/selected|column/i); + const selected = source.filter(column('score').greaterThan(literal(0))).select(['score']); + expect(() => + // @ts-expect-error Hidden source columns cannot supply sort keys. + selected.sortBy('signed') + ).toThrow(/selected|column/i); + + const constant = new GPUConstant({format: 'uint32', value: Uint32Array.of(7)}); + const constantTable = new GPUTable({ + batches: fixture.table.batches, + constants: {tier: constant} + }); + const constantSource = new LuDataFrame({table: constantTable}); + expect(() => constantSource.sortBy('tier')).toThrow(/constant/i); + expect(createBuffer).not.toHaveBeenCalled(); + + createBuffer.mockRestore(); + constantSource.destroy(); + source.destroy(); + constantTable.destroy(); + }); + + test('rejects strided scalar sort keys before allocating or planning GPU output buffers', () => { + const fixture = createSortSourceFixture([2, 0, 3], {stridedScore: true}); + Object.defineProperty(fixture.device, 'type', {value: 'webgpu'}); + const source = new LuDataFrame({table: fixture.table}); + const graph = new GPUCommandGraph(fixture.device, { + id: 'ludf-strided-sort-key' + }); + const createBuffer = vi.spyOn(fixture.device, 'createBuffer'); + const addComputePass = vi.spyOn(graph, 'addComputePass'); + const score = source.table.gpuVectors['score']; + + expect(score.byteStride).toBe(8); + expect(score.rowByteLength).toBe(4); + expect(score.stride).toBe(1); + expect(score.data.map(chunk => chunk.byteStride)).toEqual([8, 8, 8]); + expect(() => source.sortBy('score').compile(graph)).toThrow(/packed|stride|aligned/i); + expect(createBuffer).not.toHaveBeenCalled(); + expect(addComputePass).not.toHaveBeenCalled(); + + createBuffer.mockRestore(); + addComputePass.mockRestore(); + source.destroy(); + fixture.table.destroy(); + }); + + test('rejects non-closed ordering options and non-uint32 per-batch top-K limits', () => { + const fixture = createSortSourceFixture([2]); + const source = new LuDataFrame({table: fixture.table}); + const createBuffer = vi.spyOn(fixture.device, 'createBuffer'); + + for (const invalid of [ + {direction: 'sideways'}, + {nulls: 'middle'}, + {nans: 'ignore'}, + {algorithm: 'quicksort'} + ]) { + expect(() => source.sortBy('score', invalid as LuDataFrameSortOptions)).toThrow( + /direction|null|nan|algorithm/i + ); + } + expect(() => source.sortBy('score', null as unknown as LuDataFrameSortOptions)).toThrow( + /options/i + ); + for (const limit of [-1, 0.5, Number.NaN, Number.POSITIVE_INFINITY, 0x1_0000_0000]) { + expect(() => source.topK('score', limit)).toThrow(/limit|uint32/i); + expect(() => source.sortBy('score').topK(limit)).toThrow(/limit|uint32/i); + } + expect(createBuffer).not.toHaveBeenCalled(); + + createBuffer.mockRestore(); + source.destroy(); + fixture.table.destroy(); + }); + + test('rejects new sorting and top-K plans after the source dataframe was destroyed', () => { + const fixture = createSortSourceFixture([2]); + const source = new LuDataFrame({table: fixture.table}); + source.destroy(); + + expect(() => source.sortBy('score')).toThrow(/destroyed/i); + expect(() => source.topK('score', 2)).toThrow(/destroyed/i); + fixture.table.destroy(); + }); +}); + +function createSortSourceFixture( + batchLengths: readonly number[], + options: {nullableScore?: boolean; stridedScore?: boolean} = {} +): SortSourceFixture { + const device = new NullDevice({id: 'ludf-stable-sort-node-device'}); + const buffers: Buffer[] = []; + const fields: GPUField[] = [ + {name: 'score', format: 'float32', nullable: options.nullableScore ?? false}, + {name: 'signed', format: 'sint32', nullable: false}, + {name: 'category', format: 'uint32', nullable: false}, + {name: 'coordinates', format: 'float32x2', nullable: false} + ]; + let sourceRowIndexOffset = 40; + const batches = batchLengths.map((length, sourceBatchIndex) => { + const sourceInfo: GPURecordBatchSourceInfo = { + sourceBatchIndex, + sourceRowIndexOffset, + sourceRowCount: length + }; + sourceRowIndexOffset += length; + + return new GPURecordBatch({ + gpuData: { + score: makeSortSourceData( + device, + buffers, + length, + 'float32', + options.stridedScore ? {byteStride: 8, rowByteLength: 4, stride: 1} : undefined + ), + signed: makeSortSourceData(device, buffers, length, 'sint32'), + category: makeSortSourceData(device, buffers, length, 'uint32'), + coordinates: makeSortSourceData(device, buffers, length, 'float32x2') + }, + fields, + numRows: length, + sourceInfo + }); + }); + + const table = + batches.length > 0 + ? new GPUTable({batches}) + : new GPUTable({ + schema: {fields, metadata: new Map()}, + bufferLayout: [ + {name: 'score', format: 'float32'}, + {name: 'signed', format: 'sint32'}, + {name: 'category', format: 'uint32'}, + {name: 'coordinates', format: 'float32x2'} + ] + }); + return {device, table, buffers}; +} + +function makeSortSourceData( + device: NullDevice, + buffers: Buffer[], + length: number, + format: Format, + layout?: {byteStride: number; rowByteLength: number; stride: number} +): GPUData { + const byteStride = + layout?.byteStride ?? (format === 'float32x2' ? 8 : Uint32Array.BYTES_PER_ELEMENT); + const buffer = device.createBuffer({ + byteLength: Math.max(length, 1) * byteStride, + usage: Buffer.STORAGE | Buffer.COPY_SRC | Buffer.COPY_DST + }); + buffers.push(buffer); + return new GPUData({buffer, format, length, ownsBuffer: true, ...layout}); +} diff --git a/modules/experimental/test/ludf/lu-sort.spec.ts b/modules/experimental/test/ludf/lu-sort.spec.ts new file mode 100644 index 0000000000..4d61035540 --- /dev/null +++ b/modules/experimental/test/ludf/lu-sort.spec.ts @@ -0,0 +1,546 @@ +// 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 LuSortSourceSchema = { + score: 'float32'; + signed: 'sint32'; + category: 'uint32'; +}; + +type LuSortFixture = { + frame: LuDataFrame; + sourceBuffers: Buffer[]; +}; + +test('LuDataFrame stably sorts nullable floating-point batches without materializing source rows', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuSortFixture(device); + const createBufferSpy = vi.spyOn(device, 'createBuffer'); + const submitSpy = vi.spyOn(device, 'submit'); + const query = fixture.frame.sortBy('score'); + + testContext.equal( + createBufferSpy.mock.calls.length, + 0, + 'immutable sort planning allocates no GPU buffers' + ); + testContext.equal( + submitSpy.mock.calls.length, + 0, + 'immutable sort planning submits no GPU commands' + ); + + const graph = new GPUCommandGraph(device, { + id: 'ludf-floating-point-stable-sort' + }); + const compiled = query.compile(graph); + + try { + testContext.equal( + compiled.sortColumn, + 'score', + 'compiled sorting exposes its selected source column' + ); + testContext.equal(compiled.direction, 'ascending', 'sortBy defaults to ascending order'); + testContext.equal(compiled.nulls, 'last', 'null placement defaults to the final selected rows'); + testContext.equal(compiled.nans, 'last', 'NaNs default to the end of the non-null values'); + testContext.deepEqual( + compiled.table.batches.map(batch => batch.numRows), + [6, 0, 7], + 'sorting retains source record batches and does not fabricate a materialized table' + ); + testContext.deepEqual( + compiled.dictionaries.category, + {values: ['economy', 'standard', 'premium'], ordered: false}, + 'sorting keeps adapter-owned categorical dictionary metadata' + ); + + fixture.frame.destroy(); + testContext.ok( + fixture.sourceBuffers.every(buffer => !buffer.destroyed), + 'compiled sorts retain their owned source lease' + ); + + const commandEncoder = device.createCommandEncoder({id: 'ludf-float-sort-encode'}); + compiled.encode(commandEncoder); + testContext.equal( + submitSpy.mock.calls.length, + 0, + 'sorting only records caller-owned graph commands' + ); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + await readLuSortedSourceRows(compiled.rowIndices, compiled.selectedCounts), + [[43, 40, 41, 44, 42, 45], [], [52, 47, 50, 51, 46, 48, 49]], + 'stable per-batch float sorting handles infinities, signed zeros, NaNs, nulls, and duplicate keys' + ); + testContext.deepEqual( + await readLuSortChunks(compiled.selectedCounts), + [[6], [0], [7]], + 'sorting retains each source batch independently, including explicit empty batches' + ); + testContext.deepEqual( + await readLuSortChunks(compiled.selectionMask), + [[1, 1, 1, 1, 1, 1], [], [1, 1, 1, 1, 1, 1, 1]], + 'plain sorting includes explicit null and NaN rows while retaining source-row selection masks' + ); + + compiled.destroy(); + testContext.ok( + fixture.sourceBuffers.every(buffer => buffer.destroyed), + 'owned source buffers release after the compiled sorting lease is destroyed' + ); + } finally { + compiled.destroy(); + fixture.frame.destroy(); + createBufferSpy.mockRestore(); + submitSpy.mockRestore(); + } + + testContext.end(); +}); + +test('LuDataFrame orders nulls, NaNs, signed integers, and full-width unsigned values explicitly', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuSortFixture(device); + const nullableGraph = new GPUCommandGraph(device, { + id: 'ludf-null-nan-sort' + }); + const nullable = fixture.frame + .sortBy('score', {nulls: 'first', nans: 'first'}) + .compile(nullableGraph); + const descendingGraph = new GPUCommandGraph(device, { + id: 'ludf-descending-sort' + }); + const descending = fixture.frame + .sortBy('score', {direction: 'descending'}) + .compile(descendingGraph); + const signedGraph = new GPUCommandGraph(device, { + id: 'ludf-signed-sort' + }); + const signed = fixture.frame.sortBy('signed', {algorithm: 'radix'}).compile(signedGraph); + const unsignedGraph = new GPUCommandGraph(device, { + id: 'ludf-unsigned-sort' + }); + const unsigned = fixture.frame.sortBy('category').compile(unsignedGraph); + + try { + const commandEncoder = device.createCommandEncoder({id: 'ludf-all-sort-formats-encode'}); + nullable.encode(commandEncoder); + descending.encode(commandEncoder); + signed.encode(commandEncoder); + unsigned.encode(commandEncoder); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + await readLuSortedSourceRows(nullable.rowIndices, nullable.selectedCounts), + [[45, 42, 43, 40, 41, 44], [], [49, 52, 47, 50, 51, 46, 48]], + 'nulls sort outermost while NaNs sort first among remaining non-null values' + ); + testContext.deepEqual( + await readLuSortedSourceRows(descending.rowIndices, descending.selectedCounts), + [[44, 40, 41, 43, 42, 45], [], [46, 48, 50, 51, 47, 52, 49]], + 'descending floating order keeps signed-zero and duplicate source-row ties stable' + ); + testContext.deepEqual( + await readLuSortedSourceRows(signed.rowIndices, signed.selectedCounts), + [[41, 42, 43, 45, 40, 44], [], [48, 46, 50, 47, 52, 51, 49]], + 'signed bit transforms preserve full int32 ordering and stable duplicate ties' + ); + testContext.deepEqual( + await readLuSortedSourceRows(unsigned.rowIndices, unsigned.selectedCounts), + [[41, 44, 42, 43, 40, 45], [], [47, 50, 46, 48, 49, 51, 52]], + 'unsigned sorting supports the full uint32 domain without reserving null sentinels' + ); + } finally { + nullable.destroy(); + descending.destroy(); + signed.destroy(); + unsigned.destroy(); + fixture.frame.destroy(); + } + + testContext.end(); +}); + +test('LuDataFrame publishes filtered per-batch top-K rows and encoder-ordered parameter updates', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuSortFixture(device); + const graph = new GPUCommandGraph(device, { + id: 'ludf-filtered-top-k' + }); + const compiled = fixture.frame + .filter(column('score').greaterThan(parameter('minimumScore', -1))) + .topK('score', 2) + .compile(graph); + + const firstCounts = compiled.selectedCounts.data.map((_, batchIndex) => + device.createBuffer({ + id: `ludf-top-k-count-${batchIndex}`, + byteLength: Uint32Array.BYTES_PER_ELEMENT, + usage: Buffer.COPY_SRC | Buffer.COPY_DST + }) + ); + const firstMasks = compiled.selectionMask.data.map((chunk, batchIndex) => + chunk.length > 0 + ? device.createBuffer({ + id: `ludf-top-k-mask-${batchIndex}`, + byteLength: chunk.length * Uint32Array.BYTES_PER_ELEMENT, + usage: Buffer.COPY_SRC | Buffer.COPY_DST + }) + : undefined + ); + + try { + testContext.equal(compiled.limit, 2, 'compiled top-K exposes its per-batch limit'); + testContext.equal(compiled.direction, 'descending', 'top-K defaults to largest values first'); + + const commandEncoder = device.createCommandEncoder({id: 'ludf-top-k-two-encodes'}); + compiled.encode(commandEncoder, {minimumScore: -1}); + for (const [batchIndex, count] of compiled.selectedCounts.data.entries()) { + commandEncoder.copyBufferToBuffer({ + sourceBuffer: getLuSortBuffer(count), + destinationBuffer: firstCounts[batchIndex], + size: Uint32Array.BYTES_PER_ELEMENT + }); + const maskSnapshot = firstMasks[batchIndex]; + const sourceMask = compiled.selectionMask.data[batchIndex]; + if (maskSnapshot && sourceMask.length > 0) { + commandEncoder.copyBufferToBuffer({ + sourceBuffer: getLuSortBuffer(sourceMask), + destinationBuffer: maskSnapshot, + size: sourceMask.length * Uint32Array.BYTES_PER_ELEMENT + }); + } + } + compiled.encode(commandEncoder, {minimumScore: 2}); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + await Promise.all(firstCounts.map(async buffer => (await readLuSortBuffer(buffer, 1))[0])), + [2, 0, 2], + 'initial top-K clamps each preserved batch count independently' + ); + testContext.deepEqual( + await Promise.all( + firstMasks.map((buffer, batchIndex) => + buffer + ? readLuSortBuffer(buffer, compiled.selectionMask.data[batchIndex].length) + : Promise.resolve([]) + ) + ), + [[1, 0, 0, 0, 1, 0], [], [1, 0, 1, 0, 0, 0, 0]], + 'top-K source masks include exactly the selected sorted source rows' + ); + testContext.deepEqual( + await readLuSortedSourceRows(compiled.rowIndices, compiled.selectedCounts), + [[44], [], [46, 48]], + 'the same compiled graph updates ordered IDs after a stricter filter parameter' + ); + testContext.deepEqual( + await readLuSortChunks(compiled.selectionMask), + [[0, 0, 0, 0, 1, 0], [], [1, 0, 1, 0, 0, 0, 0]], + 'second-encode masks remain coherent with sorted IDs and clamped batch counts' + ); + } finally { + for (const buffer of firstCounts) buffer.destroy(); + for (const buffer of firstMasks) buffer?.destroy(); + compiled.destroy(); + fixture.frame.destroy(); + } + + testContext.end(); +}); + +test('LuDataFrame supports zero, oversized, and preordered per-batch top-K limits', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuSortFixture(device); + const emptyGraph = new GPUCommandGraph(device, { + id: 'ludf-zero-top-k' + }); + const empty = fixture.frame.topK('signed', 0).compile(emptyGraph); + const oversizedGraph = new GPUCommandGraph(device, { + id: 'ludf-oversized-top-k' + }); + const oversized = fixture.frame.topK('signed', 100).compile(oversizedGraph); + const orderedGraph = new GPUCommandGraph(device, { + id: 'ludf-ordered-top-k' + }); + const ordered = fixture.frame + .sortBy('signed', {direction: 'ascending'}) + .topK(2) + .compile(orderedGraph); + + try { + const commandEncoder = device.createCommandEncoder({id: 'ludf-top-k-limits-encode'}); + empty.encode(commandEncoder); + oversized.encode(commandEncoder); + ordered.encode(commandEncoder); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + await readLuSortChunks(empty.selectedCounts), + [[0], [0], [0]], + 'a zero top-K limit rejects every row while retaining every batch' + ); + testContext.deepEqual( + await readLuSortChunks(empty.selectionMask), + [[0, 0, 0, 0, 0, 0], [], [0, 0, 0, 0, 0, 0, 0]], + 'a zero top-K limit clears every source-row selection mask' + ); + testContext.deepEqual( + await readLuSortChunks(oversized.selectedCounts), + [[6], [0], [7]], + 'an oversized per-batch limit retains all selected rows without padding' + ); + testContext.equal( + ordered.direction, + 'ascending', + 'sorted-plan top-K preserves explicit direction' + ); + testContext.deepEqual( + await readLuSortedSourceRows(ordered.rowIndices, ordered.selectedCounts), + [[41, 42], [], [48, 46]], + 'top-K applied to an existing ascending plan returns the smallest stable rows per batch' + ); + } finally { + empty.destroy(); + oversized.destroy(); + ordered.destroy(); + fixture.frame.destroy(); + } + + testContext.end(); +}); + +test('LuDataFrame sorts nullable derived columns and schema-only source tables', async testContext => { + const device = await getWebGPUTestDevice(); + if (!device) { + testContext.comment('WebGPU is not available'); + testContext.end(); + return; + } + + const fixture = createLuSortFixture(device); + const derivedGraph = new GPUCommandGraph(device, { + id: 'ludf-derived-stable-sort' + }); + const derived = fixture.frame + .withColumn('shiftedScore', column('score').add(literal(1)), {format: 'float32'}) + .sortBy('shiftedScore', {nulls: 'first'}) + .compile(derivedGraph); + + const emptyFrame = new LuDataFrame({ + table: new GPUTable({ + schema: { + fields: [ + {name: 'score', format: 'float32', nullable: true}, + {name: 'signed', format: 'sint32', nullable: false}, + {name: 'category', format: 'uint32', nullable: false} + ], + metadata: new Map([['dataset', 'empty-sort']]) + }, + bufferLayout: [ + {name: 'score', format: 'float32', byteStride: 4}, + {name: 'signed', format: 'sint32', byteStride: 4}, + {name: 'category', format: 'uint32', byteStride: 4} + ] + }), + ownership: 'owned' + }); + const emptyGraph = new GPUCommandGraph(device, { + id: 'ludf-schema-only-sort' + }); + const empty = emptyFrame.sortBy('score').compile(emptyGraph); + + try { + const commandEncoder = device.createCommandEncoder({id: 'ludf-derived-and-empty-sort'}); + derived.encode(commandEncoder); + empty.encode(commandEncoder); + device.submit(commandEncoder.finish()); + + testContext.deepEqual( + await readLuSortedSourceRows(derived.rowIndices, derived.selectedCounts), + [[45, 43, 40, 41, 44, 42], [], [52, 47, 50, 51, 46, 48, 49]], + 'nullable derived keys preserve null-first ordering, NaN-last placement, and stable ties' + ); + testContext.deepEqual( + empty.table.batches, + [], + 'schema-only sorts do not invent source batches' + ); + testContext.deepEqual( + empty.selectedCounts.data, + [], + 'schema-only sorts allocate no batch counts' + ); + testContext.equal( + empty.table.schema.metadata.get('dataset'), + 'empty-sort', + 'schema-only sorted projections retain source metadata' + ); + } finally { + derived.destroy(); + empty.destroy(); + emptyFrame.destroy(); + fixture.frame.destroy(); + } + + testContext.end(); +}); + +function createLuSortFixture(device: Device): LuSortFixture { + const sourceBuffers: Buffer[] = []; + const scores = [ + Float32Array.from([-0, 0, Number.NaN, Number.NEGATIVE_INFINITY, Number.POSITIVE_INFINITY, 5]), + new Float32Array(0), + Float32Array.from([3, -2, 3, Number.NaN, -0, 0, Number.NEGATIVE_INFINITY]) + ]; + const scoreValidity = [ + Uint32Array.from([1, 1, 1, 1, 1, 0]), + new Uint32Array(0), + Uint32Array.from([1, 1, 1, 1, 1, 1, 1]) + ]; + const signed = [ + Int32Array.from([0x7fffffff, -0x80000000, -3, 0, 0x7fffffff, 7]), + new Int32Array(0), + Int32Array.from([-1, 0, -0x80000000, 0x7fffffff, -1, 8, 7]) + ]; + const categories = [ + Uint32Array.from([0xffffffff, 0, 2, 2, 1, 0xffffffff]), + new Uint32Array(0), + Uint32Array.from([1, 0, 1, 2, 0, 2, 0xffffffff]) + ]; + const validityChunks: GPUData<'uint32'>[] = []; + let sourceRowIndexOffset = 40; + + const batches = scores.map((values, batchIndex) => { + const batch = new GPURecordBatch({ + gpuData: { + score: createLuSortData(device, sourceBuffers, values, 'float32'), + signed: createLuSortData(device, sourceBuffers, signed[batchIndex], 'sint32'), + category: createLuSortData(device, sourceBuffers, categories[batchIndex], 'uint32') + }, + fields: [ + {name: 'score', format: 'float32', nullable: true}, + {name: 'signed', format: 'sint32', nullable: false}, + {name: 'category', format: 'uint32', nullable: false} + ], + sourceInfo: { + sourceBatchIndex: batchIndex + 4, + sourceRowIndexOffset, + sourceRowCount: values.length + } + }); + sourceRowIndexOffset += values.length; + validityChunks.push( + createLuSortData(device, sourceBuffers, scoreValidity[batchIndex], 'uint32') + ); + return batch; + }); + + return { + frame: new LuDataFrame({ + table: new GPUTable({batches}), + validity: { + score: new GPUVector<'uint32'>({ + type: 'data', + name: 'ludf-sort-score-validity', + format: 'uint32', + data: validityChunks, + ownsData: true + }) + }, + dictionaries: { + category: {values: ['economy', 'standard', 'premium'], ordered: false} + }, + ownership: 'owned' + }), + sourceBuffers + }; +} + +function createLuSortData( + 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 getLuSortBuffer(data: GPUData): Buffer { + return data.buffer instanceof Buffer ? data.buffer : data.buffer.buffer; +} + +async function readLuSortBuffer(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 readLuSortChunks(vector: GPUVector<'uint32'>): Promise { + return Promise.all( + vector.data.map(chunk => readLuSortBuffer(getLuSortBuffer(chunk), chunk.length)) + ); +} + +async function readLuSortedSourceRows( + rowIndices: GPUVector<'uint32'>, + selectedCounts: GPUVector<'uint32'> +): Promise { + const counts = await readLuSortChunks(selectedCounts); + return Promise.all( + rowIndices.data.map((chunk, batchIndex) => + readLuSortBuffer(getLuSortBuffer(chunk), counts[batchIndex][0] ?? 0) + ) + ); +}