Line data    Source code

       1              : import { ArrowTableFormatter } from '../arrow_formatter.js';
       2              : import { DataFrame, generateTableName } from '../data_frame.js';
       3              : import { createArrowFieldIndex } from '../computation_state.js';
       4              : import { TableAggregationTask, TableAggregation, ORDINAL_COLUMN, STRING_COLUMN, LIST_COLUMN, ColumnGroup, SKIPPED_COLUMN, ROWNUMBER_COLUMN } from '../computation_types.js';
       5              : import { LoggableException, LoggerLike, stringifyError } from '../../platform/logger/logger.js';
       6              : import { SQLFrame } from '../sql/sqlframe_builder.js';
       7              : import { ComputeQueryExecution, executeDirectly } from '../computation_tasks.js';
       8              : import { isIntegerType, getSmallIntegerDomainSize } from '../arrow_columns.js';
       9              : 
      10            1 : const LOG_CTX = "compute";
      11              : 
      12            1 : export async function computeTableAggregates(task: TableAggregationTask, logger: LoggerLike, executeQuery: ComputeQueryExecution = executeDirectly): Promise<[TableAggregation, ColumnGroup[]]> {
      13            1 :     const [sql, columnEntries, countStarColumn] = buildTableAggregationSQL(task);
      14              : 
      15            1 :     try {
      16            1 :         logger.info("Aggregating table", {
      17            1 :             "version": task.tableVersion.toString(),
      18            1 :             "table": task.tableId.toString(),
      19            1 :         }, LOG_CTX);
      20            1 :         const summaryStart = performance.now();
      21            1 :         const tableName = generateTableName("__tbl_agg");
      22            1 :         const transformedDataFrame = await executeQuery(sql, clientContext => DataFrame.fromSQL(task.inputDataFrame.database, sql, tableName, undefined, clientContext));
      23            1 :         const summaryEnd = performance.now();
      24            1 :         logger.info("Aggregated table", {
      25            1 :             "version": task.tableVersion.toString(),
      26            1 :             "table": task.tableId.toString(),
      27            1 :             "duration": Math.floor(summaryEnd - summaryStart).toString()
      28            1 :         }, LOG_CTX);
      29              : 
      30            1 :         const statsTable = await transformedDataFrame.readTable();
      31            1 :         const statsTableFields = createArrowFieldIndex(statsTable);
      32            1 :         const statsTableFormatter = new ArrowTableFormatter(statsTable.schema, statsTable.batches, logger);
      33              : 
      34            1 :         for (let i = 0; i < columnEntries.length; ++i) {
      35            1 :             const entry = columnEntries[i];
      36            1 :             if (entry.type != ORDINAL_COLUMN || !isIntegerType(entry.value.inputFieldType.typeId)) {
      37            1 :                 continue;
      38              :             }
      39            1 :             const minValue = statsTable.getChild(entry.value.statsFields!.minAggregateFieldName!)?.get(0);
      40            1 :             const maxValue = statsTable.getChild(entry.value.statsFields!.maxAggregateFieldName!)?.get(0);
      41            1 :             const domainSize = getSmallIntegerDomainSize(minValue, maxValue, entry.value.binCount);
      42            1 :             if (domainSize != null) {
      43            1 :                 columnEntries[i] = {
      44            1 :                     type: ORDINAL_COLUMN,
      45            1 :                     value: { ...entry.value, binCount: domainSize },
      46            1 :                 };
      47              :             }
      48            1 :         }
      49              : 
      50              :         const summary: TableAggregation = {
      51            1 :             table: statsTable,
      52            1 :             tableFormatter: statsTableFormatter,
      53            1 :             tableFieldsByName: statsTableFields,
      54            1 :             dataFrame: transformedDataFrame,
      55            1 :             countStarFieldName: countStarColumn,
      56              :         };
      57            1 :         return [summary, columnEntries];
      58              : 
      59            1 :     } catch (error: any) {
      60            0 :         if (error instanceof LoggableException) {
      61            0 :             throw error;
      62              :         } else {
      63            0 :             throw new LoggableException("computing aggregate failed", {
      64            0 :                 "table": task.tableId.toString(),
      65            0 :                 "error": stringifyError(error)
      66            0 :             }, LOG_CTX);
      67              :         }
      68            0 :     }
      69            1 : }
      70              : 
      71            1 : function buildTableAggregationSQL(task: TableAggregationTask): [string, ColumnGroup[], string] {
      72            1 :     const countColumn = `_count`;
      73            1 :     const aggregates: { fieldName?: string; outputAlias: string; func: "min" | "max" | "count" | "count_star"; distinct?: boolean }[] = [];
      74              : 
      75            1 :     aggregates.push({ outputAlias: `_count`, func: "count_star" });
      76              : 
      77            1 :     const updatedEntries: ColumnGroup[] = [];
      78            1 :     for (let i = 0; i < task.columnEntries.length; ++i) {
      79            1 :         const entry = task.columnEntries[i];
      80            1 :         switch (entry.type) {
      81            1 :             case SKIPPED_COLUMN:
      82            1 :             case ROWNUMBER_COLUMN:
      83            1 :                 updatedEntries.push(entry);
      84            1 :                 break;
      85            1 :             case ORDINAL_COLUMN: {
      86            1 :                 const countAggregateColumn = `_${i}_count`;
      87            1 :                 const minAggregateColumn = `_${i}_min`;
      88            1 :                 const maxAggregateColumn = `_${i}_max`;
      89            1 :                 aggregates.push({ fieldName: entry.value.inputFieldName, outputAlias: countAggregateColumn, func: "count" });
      90            1 :                 aggregates.push({ fieldName: entry.value.inputFieldName, outputAlias: minAggregateColumn, func: "min" });
      91            1 :                 aggregates.push({ fieldName: entry.value.inputFieldName, outputAlias: maxAggregateColumn, func: "max" });
      92            1 :                 updatedEntries.push({
      93            1 :                     type: ORDINAL_COLUMN,
      94            1 :                     value: {
      95            1 :                         ...entry.value,
      96            1 :                         statsFields: {
      97            1 :                             countFieldName: countAggregateColumn,
      98            1 :                             distinctCountFieldName: null,
      99            1 :                             minAggregateFieldName: minAggregateColumn,
     100            1 :                             maxAggregateFieldName: maxAggregateColumn,
     101            1 :                         }
     102            1 :                     }
     103            1 :                 });
     104            1 :                 break;
     105            1 :             }
     106            1 :             case STRING_COLUMN: {
     107            1 :                 const countAggregateColumn = `_${i}_count`;
     108            1 :                 const countDistinctAggregateColumn = `_${i}_countd`;
     109            1 :                 aggregates.push({ fieldName: entry.value.inputFieldName, outputAlias: countAggregateColumn, func: "count" });
     110            1 :                 aggregates.push({ fieldName: entry.value.inputFieldName, outputAlias: countDistinctAggregateColumn, func: "count", distinct: true });
     111            1 :                 updatedEntries.push({
     112            1 :                     type: STRING_COLUMN,
     113            1 :                     value: {
     114            1 :                         ...entry.value,
     115            1 :                         statsFields: {
     116            1 :                             countFieldName: countAggregateColumn,
     117            1 :                             distinctCountFieldName: countDistinctAggregateColumn,
     118            1 :                             minAggregateFieldName: null,
     119            1 :                             maxAggregateFieldName: null,
     120            1 :                         }
     121            1 :                     }
     122            1 :                 });
     123            1 :                 break;
     124            1 :             }
     125            1 :             case LIST_COLUMN:
     126              :                 // No table-level aggregates for LIST columns. In particular
     127              :                 // count(DISTINCT array) over float32 embeddings hashes every full
     128              :                 // vector at high cost, and no summary consumes it. Leaving
     129              :                 // statsFields null and passing the entry through unchanged.
     130            1 :                 updatedEntries.push(entry);
     131              :                 break;
     132            1 :         }
     133            1 :     }
     134              : 
     135            1 :     const sql = SQLFrame.from(task.inputDataFrame.tableName)
     136            1 :         .groupBy({
     137            1 :             keys: [],
     138            1 :             aggregates,
     139            1 :         })
     140            1 :         .toSQL();
     141            1 :     return [sql, updatedEntries, countColumn];
     142            1 : }