feat(experimental): add stable luDF sorting and top-K

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';

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 extends LuDataFrameScalarColumnNames<Logical, SelectedColumns>>(
+    column: Column,
+    options: LuDataFrameSortOptions = {}
+  ): LuDataFrameSortQuery<Logical, SelectedColumns, Column, Source> {
+    return new LuDataFrameSortQuery(this, column, options);
+  }
+
+  /** Plans descending stable top-K selection independently within every source record batch. */
+  topK<Column extends LuDataFrameScalarColumnNames<Logical, SelectedColumns>>(
+    column: Column,
+    limit: number,
+    options: LuDataFrameSortOptions = {}
+  ): LuDataFrameSortQuery<Logical, SelectedColumns, Column, Source> {
+    return new LuDataFrameSortQuery(this, column, options, limit, 'descending');
+  }
+
   /** Materializes reusable GPU graph passes and compiler-owned selection/index/count outputs. */
   compile(
     graph: GPUCommandGraph<LuDataFrameQueryParameters>

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<T extends GPUTypeMap = GPUTypeMap> {
     );
   }
 
+  /** Plans stable scalar ordering independently within every existing source record batch. */
+  sortBy<Column extends LuDataFrameScalarColumnNames<T, keyof T & string>>(
+    column: Column,
+    options: LuDataFrameSortOptions = {}
+  ): LuDataFrameSortQuery<T, keyof T & string, Column, T> {
+    this.assertAvailable();
+    return new LuDataFrameQuery<T, keyof T & string, T>(this, [], this.columnNames).sortBy(
+      column,
+      options
+    );
+  }
+
+  /** Plans descending stable top-K selection without concatenating source record batches. */
+  topK<Column extends LuDataFrameScalarColumnNames<T, keyof T & string>>(
+    column: Column,
+    limit: number,
+    options: LuDataFrameSortOptions = {}
+  ): LuDataFrameSortQuery<T, keyof T & string, Column, T> {
+    this.assertAvailable();
+    return new LuDataFrameQuery<T, keyof T & string, T>(this, [], this.columnNames).topK(
+      column,
+      limit,
+      options
+    );
+  }
+
   /**
    * Returns an independent borrowed projection without mutating or destroying source columns.
    *

modules/experimental/src/ludf/lu-query-compiler.ts

@@ -94,6 +94,8 @@ export type LuDataFrameQueryExtensionContext<T extends GPUTypeMap> = {
   validity: Readonly<LuDataFrameValidity<T>>;
   dictionaries: Readonly<LuDataFrameDictionaries<T>>;
   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
       });
     }