Line data    Source code

       1              : import * as arrow from 'apache-arrow';
       2              : 
       3              : import { DataFrame, generateTableName } from '../data_frame.js';
       4              : import { TableAggregation, ORDINAL_COLUMN, STRING_COLUMN, LIST_COLUMN, ColumnGroup, SKIPPED_COLUMN, SystemColumnComputationTask, 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 { getToNumericFn, isIntegerType, createUniqueColumnName } from '../arrow_columns.js';
       9              : 
      10            1 : const LOG_CTX = "compute";
      11              : 
      12            1 : export async function computeSystemColumns(task: SystemColumnComputationTask, logger: LoggerLike, executeQuery: ComputeQueryExecution = executeDirectly): Promise<[arrow.Table, DataFrame, ColumnGroup[]]> {
      13            1 :     try {
      14            1 :         const [sqlFrame, columnGroups] = buildSystemColumnSQLFrame(task.inputTable.schema, task.columnEntries, task.inputDataFrame.tableName, task.rowNumberFieldName, task.tableAggregate);
      15              : 
      16            1 :         logger.info("Computing system columns", {
      17            1 :             "version": task.tableVersion.toString(),
      18            1 :         }, LOG_CTX);
      19            1 :         const transformStart = performance.now();
      20            1 :         const tableName = generateTableName("__syscols");
      21            1 :         const sql = sqlFrame.toSQL();
      22            1 :         const transformed = await executeQuery(sql, clientContext => DataFrame.fromSQL(task.inputDataFrame.database, sql, tableName, undefined, clientContext));
      23            1 :         const transformEnd = performance.now();
      24            1 :         const transformedTable = await transformed.readTable({ orderBy: [{ field: task.rowNumberFieldName }] });
      25            1 :         logger.info("Computed system columns", {
      26            1 :             "version": task.tableVersion.toString(),
      27            1 :             "duration": Math.floor(transformEnd - transformStart).toString(),
      28            1 :             "sql": sql,
      29            1 :         }, LOG_CTX);
      30              : 
      31            1 :         let rowNumColumnName: string | null = null;
      32            1 :         for (let i = 0; i < columnGroups.length; ++i) {
      33            1 :             const group = columnGroups[i];
      34            1 :             switch (group.type) {
      35            1 :                 case ROWNUMBER_COLUMN:
      36            1 :                     rowNumColumnName = group.value.rowNumberFieldName;
      37              :                     break;
      38            1 :             }
      39            1 :         }
      40            1 :         if (!rowNumColumnName) {
      41            0 :             throw new LoggableException("Missing rownum column group", {}, LOG_CTX);
      42              :         }
      43            1 :         return [transformedTable, transformed, columnGroups];
      44              : 
      45            1 :     } catch (error: any) {
      46            0 :         if (error instanceof LoggableException) {
      47            0 :             throw error;
      48              :         } else {
      49            0 :             throw new LoggableException("Failed to precompute column", { "error": stringifyError(error) }, LOG_CTX);
      50              :         }
      51            0 :     }
      52            1 : }
      53              : 
      54            1 : function buildSystemColumnSQLFrame(schema: arrow.Schema, columns: ColumnGroup[], inputTableName: string, rowNumberFieldName: string, tableAggregate: TableAggregation): [SQLFrame, ColumnGroup[]] {
      55            1 :     let fieldNames = new Set<string>([rowNumberFieldName]);
      56            1 :     for (const field of schema.fields) {
      57            1 :         fieldNames.add(field.name);
      58              :     }
      59              : 
      60              :     const rowNumberGridColumn: ColumnGroup = {
      61            1 :         type: ROWNUMBER_COLUMN,
      62            1 :         value: {
      63            1 :             rowNumberFieldName: rowNumberFieldName
      64            1 :         }
      65              :     };
      66            1 :     let gridColumns: ColumnGroup[] = [
      67            1 :         rowNumberGridColumn,
      68            1 :         ...columns
      69            1 :     ];
      70              : 
      71            1 :     let frame = SQLFrame.from(inputTableName);
      72              : 
      73            1 :     for (let i = 1; i <= columns.length; ++i) {
      74            1 :         let column = gridColumns[i];
      75            1 :         switch (column.type) {
      76            0 :             case ROWNUMBER_COLUMN:
      77            1 :             case SKIPPED_COLUMN:
      78            1 :                 break;
      79            1 :             case ORDINAL_COLUMN: {
      80            1 :                 const binFieldName = createUniqueColumnName(`_${i}_bin`, fieldNames);
      81            1 :                 frame = frame.binning({
      82            1 :                     fieldName: column.value.inputFieldName,
      83            1 :                     statsTable: tableAggregate.dataFrame.tableName,
      84            1 :                     statsMinField: column.value.statsFields!.minAggregateFieldName!,
      85            1 :                     statsMaxField: column.value.statsFields!.maxAggregateFieldName!,
      86            1 :                     binCount: column.value.binCount,
      87            1 :                     minBinWidth: isIntegerType(column.value.inputFieldType.typeId) ? 1 : undefined,
      88            1 :                     outputAlias: binFieldName,
      89            1 :                     toNumericFn: getToNumericFn(column.value.inputFieldType.typeId),
      90            1 :                 });
      91            1 :                 gridColumns[i] = {
      92            1 :                     type: ORDINAL_COLUMN,
      93            1 :                     value: {
      94            1 :                         ...column.value,
      95            1 :                         binFieldName: binFieldName,
      96            1 :                     }
      97            1 :                 };
      98            1 :                 break;
      99            1 :             }
     100            1 :             case STRING_COLUMN: {
     101            1 :                 const valueFieldName = createUniqueColumnName(`_${i}_id`, fieldNames);
     102            1 :                 frame = frame.valueIdentifier(column.value.inputFieldName, valueFieldName);
     103            1 :                 gridColumns[i] = {
     104            1 :                     type: STRING_COLUMN,
     105            1 :                     value: {
     106            1 :                         ...column.value,
     107            1 :                         valueIdFieldName: valueFieldName
     108            1 :                     }
     109            1 :                 };
     110            1 :                 break;
     111            1 :             }
     112            1 :             case LIST_COLUMN:
     113              :                 // No value identifier for LIST columns: a dense_rank() over float32
     114              :                 // embedding arrays sorts the whole table by full-vector comparison at
     115              :                 // high cost, and nothing consumes the id (no summary is rendered).
     116              :                 // Leaving valueIdFieldName null; the layout guards on `!= null`.
     117              :                 break;
     118            1 :         }
     119            1 :     }
     120              : 
     121            1 :     return [frame, gridColumns];
     122            1 : }