import { DataTypeConfig , FieldType } from '@terascope/types' ; import { BigMap } from '@terascope/utils' ; import { Column , KeyAggregation , ValueAggregation , valueAggMap , FieldAgg , KeyAggFn , keyAggMap , makeUniqueKeyAgg } from '../column' ; import { getCommonTupleType , isNumberLike , Vector } from '../vector' ; import { Builder } from '../builder' ; import { getBuilderForField , getMaxColumnSize } from './utils' ; import { FieldArg , flattenStringArg , freezeArray , getFieldsFromArg } from '../core' ; import { columnsToDataTypeConfig , makeKeyForRow } from '../data-frame/utils' ; /**
 * A deferred execution frame dedicated to running a aggregations.
 *
 * This is different from a DataFrame for a few reasons:
 *  - GroupBy and aggregations have to run at the same time in-order to get the correctly results.
 *  - The operations are added to an instruction set and in one optimized execution.
 *  - All methods in the AggregationFrame will mutate the execution
 *    instructions instead of return a new instance with the applied changes.
*/ export class AggregationFrame < T extends Record < string , any > > { /**
     * The name of the Frame
    */ name ? : string ; /**
     * The list of columns
    */ columns : readonly Column < any , keyof T > [ ] ; /**
     * An array of the column names
    */ fields : readonly ( keyof T ) [ ] ; /**
     * Metadata about the Frame
    */ readonly metadata : Record < string , any > ; internal protected _sortFields ? : ( readonly ( keyof T ) [ ] ) | ( readonly string [ ] ) ; /**
     * The number of records to limit the result by
    */ protected _limit ? : number ; /**
     * The Aggregations to run
    */ protected readonly _aggregations = new Map < keyof T , AggObject > ( ) ; /**
     * Group By fields
    */ protected _groupByFields : readonly ( keyof T ) [ ] ; /**
     * The field to sort by
    */ protected _selectFields ? : readonly ( keyof T ) [ ] ; /**
     * Use this to cache the the column index needed, this
     * should speed things up
    */ protected _fieldToColumnIndexCache ? : Map < ( keyof T ) , number > ; constructor ( columns , options ) { this . columns = freezeArray ( columns ) ; this . fields = Object . freeze ( this . columns . map ( ( col ) => col . name ) ) ; this . name = options ?. name ; this . metadata = options ?. metadata ? { ... options . metadata } : { } ; this . _groupByFields = [ ] ; } /**
     * Get the number of records in the AggregationFrame
    */ get size ( ) : number { return getMaxColumnSize ( this . columns ) ; } /**
     * Generate the DataType config from the columns.
    */ get config ( ) : DataTypeConfig { return columnsToDataTypeConfig ( this . columns ) ; } get id ( ) : string { throw new Error ( `${ this . constructor . name }.id is currently unsupported` ) ; } /**
     * GroupBy fields
    */ groupBy ( ... fieldArg ) { this . _groupByFields = Object . freeze ( this . _groupByFields . concat ( Array . from ( getFieldsFromArg ( this . fields , fieldArg ) ) ) ) ; return this ; } note param field param as [ ValueAggregation . avg ] ( field ) ; [ ValueAggregation . avg ] < A extends string > ( field , as ) ; [ ValueAggregation . avg ] < A extends string > ( field , as ? ) { const { name } = this . _ensureColumn ( field , as ) ; this . _ensureNumericLike ( name , ValueAggregation . avg ) ; const aggObject = this . _aggregations . get ( name ) ?? { } ; aggObject . value = ValueAggregation . avg ; this . _aggregations . set ( name , aggObject ) ; return this as AggregationFrame < T | WithAlias < T , A , number > > ; } note param field param as [ ValueAggregation . sum ] ( field ) ; [ ValueAggregation . sum ] < A extends string > ( field , as ) ; [ ValueAggregation . sum ] < A extends string > ( field , as ? ) { const { name } = this . _ensureColumn ( field , as ) ; this . _ensureNumericLike ( name , ValueAggregation . sum ) ; const aggObject = this . _aggregations . get ( name ) ?? { } ; aggObject . value = ValueAggregation . sum ; this . _aggregations . set ( name , aggObject ) ; return this as AggregationFrame < T | WithAlias < T , A , number > > ; } note param field param as [ ValueAggregation . min ] ( field ) ; [ ValueAggregation . min ] < A extends string > ( field , as ) ; [ ValueAggregation . min ] < A extends string > ( field , as ? ) { const { name } = this . _ensureColumn ( field , as ) ; this . _ensureNumericLike ( name , ValueAggregation . min ) ; const aggObject = this . _aggregations . get ( name ) ?? { } ; aggObject . value = ValueAggregation . min ; this . _aggregations . set ( name , aggObject ) ; return this as AggregationFrame < T | WithAlias < T , A , number > > ; } note param field param as [ ValueAggregation . max ] ( field ) ; [ ValueAggregation . max ] < A extends string > ( field , as ) ; [ ValueAggregation . max ] < A extends string > ( field , as ? ) { const { name } = this . _ensureColumn ( field , as ) ; this . _ensureNumericLike ( name , ValueAggregation . max ) ; const aggObject = this . _aggregations . get ( name ) ?? { } ; aggObject . value = ValueAggregation . max ; this . _aggregations . set ( name , aggObject ) ; return this as AggregationFrame < T | WithAlias < T , A , number > > ; } param field param as [ ValueAggregation . count ] ( field ) ; [ ValueAggregation . count ] < A extends string > ( field , as ) ; [ ValueAggregation . count ] < A extends string > ( field , as ? ) { const { name } = this . _ensureColumn ( field , as ) ; const aggObject = this . _aggregations . get ( name ) ?? { } ; aggObject . value = ValueAggregation . count ; this . _aggregations . set ( name , aggObject ) ; return this as AggregationFrame < T | WithAlias < T , A , number > > ; } note param field [ KeyAggregation . hourly ] ( field ) { const { name , type } = this . _ensureColumn ( field ) ; if ( type !== FieldType . Date ) { throw new Error ( `${ KeyAggregation . hourly } requires a ${ FieldType . Date } field type` ) ; } const aggObject = this . _aggregations . get ( name ) ?? { } ; aggObject . key = KeyAggregation . hourly ; this . _aggregations . set ( name , aggObject ) ; return this ; } note param field [ KeyAggregation . daily ] ( field ) { const { name , type } = this . _ensureColumn ( field ) ; if ( type !== FieldType . Date ) { throw new Error ( `${ KeyAggregation . daily } requires a ${ FieldType . Date } field type` ) ; } const aggObject = this . _aggregations . get ( name ) ?? { } ; aggObject . key = KeyAggregation . daily ; this . _aggregations . set ( name , aggObject ) ; return this ; } note param field [ KeyAggregation . monthly ] ( field ) { const { name , type } = this . _ensureColumn ( field ) ; if ( type !== FieldType . Date ) { throw new Error ( `${ KeyAggregation . monthly } requires a ${ FieldType . Date } field type` ) ; } const aggObject = this . _aggregations . get ( name ) ?? { } ; aggObject . key = KeyAggregation . monthly ; this . _aggregations . set ( name , aggObject ) ; return this ; } note param field [ KeyAggregation . yearly ] ( field ) { const { name , type } = this . _ensureColumn ( field ) ; if ( type !== FieldType . Date ) { throw new Error ( `${ KeyAggregation . yearly } requires a ${ FieldType . Date } field type` ) ; } const aggObject = this . _aggregations . get ( name ) ?? { } ; aggObject . key = KeyAggregation . yearly ; this . _aggregations . set ( name , aggObject ) ; return this ; } /**
     * Order the rows by fields, format of is `field:asc` or `field:desc`.
     * Defaults to `asc` if none specified
    */ orderBy ( ... fieldArgs ) ; orderBy ( ... fieldArgs ) ; orderBy ( ... fieldArgs ) { const fields = flattenStringArg ( fieldArgs ) ; if ( fields . size > 1 ) { throw new Error ( 'AggregationFrame.orderBy can only works with one field currently' ) ; } this . _sortFields = Object . freeze ( [ ... fields ] ) ; return this ; } see orderBy sort ( ... fieldArgs ) ; sort ( ... fieldArgs ) ; sort ( ... fieldArgs ) { return this . orderBy ( ... fieldArgs ) ; } /**
     * Limit the number of results being returned
    */ limit ( num ) { this . _limit = num ; return this ; } /**
     * After the aggregations run, return only these selected fields
    */ select < K extends keyof T > ( ... fieldArg ) { this . _selectFields = Object . freeze ( [ ... getFieldsFromArg ( this . fields , fieldArg ) ] ) ; return this as any ; } /**
     * Get a column by name
    */ getColumn < P extends keyof T > ( field ) { if ( this . _fieldToColumnIndexCache ?. has ( field ) ) { return this . getColumnAt < P > ( this . _fieldToColumnIndexCache . get ( field ) ! ) ; } const index = this . columns . findIndex ( ( col ) => col . name === field ) ; if ( ! this . _fieldToColumnIndexCache ) { this . _fieldToColumnIndexCache = new Map ( [ [ field , index ] ] ) ; } else { this . _fieldToColumnIndexCache . set ( field , index ) ; } return this . getColumnAt < P > ( index ) ; } /**
     * Get a column by name or throw if not found
    */ getColumnOrThrow < P extends keyof T > ( field ) { const column = this . getColumn ( field ) ; if ( ! column ) { throw new Error ( `Unknown column ${ field } in${ this . name ? ` ${ this . name }` : '' } ${ this . constructor . name }` ) ; } return column ; } /**
     * Get a column by index
    */ getColumnAt < P extends keyof T > ( index ) { return this . columns [ index ] as Column < any , P > | undefined ; } /**
     * Rename an existing column
    */ rename < K extends keyof T , R extends string > ( name , renameTo ) { this . columns = Object . freeze ( this . columns . map ( ( col ) : Column < any , any > => { if ( col . name !== name ) return col ; return col . rename ( renameTo ) ; } ) ) ; this . _fieldToColumnIndexCache ?. clear ( ) ; this . fields = Object . freeze ( this . fields . map ( ( field ) => { if ( field === name ) return renameTo ; return field ; } ) ) ; this . _sortFields = this . _sortFields ? Object . freeze ( ( this . _sortFields as string [ ] ) . map ( ( field ) => { if ( field === name ) return renameTo ; if ( field . startsWith ( `${ name }:` ) ) { return field . replace ( `${ name }:` , `${ renameTo }:` ) ; } return field ; } ) as any [ ] ) : this . _sortFields ; this . _selectFields = this . _selectFields ? Object . freeze ( this . _selectFields . map ( ( field ) => { if ( field === name ) return renameTo ; return field ; } ) ) : this . _selectFields ; this . _groupByFields = Object . freeze ( this . _groupByFields . map ( ( field ) => { if ( field === name ) return renameTo ; return field ; } ) ) ; const agg = this . _aggregations . get ( name ) ; if ( agg ) { this . _aggregations . delete ( name ) ; this . _aggregations . set ( renameTo , agg ) ; } return this as any ; } /**
     * Assign new columns, if given a column already exists,
     * the column will replace the existing one.
    */ assign < R extends Record < string , unknown > = Record < string , any > > ( columns ) { const newColumns = columns . filter ( ( col ) => { if ( this . getColumn ( col . name ) ) return false ; return true ; } ) ; this . columns = Object . freeze ( this . columns . map ( ( col ) => { const replaceCol = columns . find ( ( c ) => c . name === col . name ) ; if ( replaceCol ) return replaceCol ; return col ; } ) . concat ( newColumns ) ) as readonly Column < any > [ ] ; this . _fieldToColumnIndexCache ?. clear ( ) ; this . fields = Object . freeze ( this . fields . concat ( newColumns . map ( ( col ) => col . name ) ) ) ; return this as any ; } private _ensureColumn ( field , as ? ) { const col = this . getColumnOrThrow ( field ) ; if ( as ) { const columns = [ ] ; for ( const c of this . columns ) { columns . push ( c ) ; if ( c === col ) { columns . push ( c . rename ( as ) ) ; this . fields = Object . freeze ( this . fields . concat ( as ) ) ; } } this . columns = Object . freeze ( columns ) ; this . _fieldToColumnIndexCache ?. clear ( ) ; return { name : as , type : col . config . type as FieldType } ; } return { name : col . name , type : col . config . type as FieldType } ; } private _ensureNumericLike ( field , ctx ) { const column = this . getColumnOrThrow ( field ) ; let fieldType ; if ( column . config . type === FieldType . Tuple ) { fieldType = getCommonTupleType ( field as string , column . vector . childConfig ) ; } else { fieldType = column . config . type as FieldType ; } if ( ! isNumberLike ( fieldType ) ) { throw new Error ( `${ ctx } requires a numeric field type` ) ; } } todo async execute ( ) { if ( ! this . _aggregations . size && ! this . _groupByFields . length ) { throw new Error ( 'No aggregations specified to run' ) ; } const buckets = this . _generateBuckets ( this . _aggregationBuilders ( ) ) ; const builders = this . _generateBuilders ( buckets . size ) ; for ( const [ key , bucket ] of buckets ) { this . _buildBucket ( builders , bucket ) ; buckets . delete ( key ) ; } this . columns = Object . freeze ( [ ... builders ] . map ( ( [ name , builder ] ) => { const column = this . getColumnOrThrow ( name ) ; return column . fork ( builder . toVector ( ) ) ; } ) ) ; this . _fieldToColumnIndexCache ?. clear ( ) ; return this ; } /**
     * Reset the Aggregations
    */ reset ( ) { this . _limit = undefined ; this . _sortFields = undefined ; this . _selectFields = undefined ; this . _groupByFields = Object . freeze ( [ ] ) ; this . _aggregations . clear ( ) ; return this ; } private _generateBuckets ( { keyAggs , fieldAggMakers } ) { const buckets = new BigMap < string , Bucket < T > > ( ) ; const getFieldAggs = makeGetFieldAggs ( buckets ) ; for ( let i = 0 ; i < this . size ; i ++ ) { const res = makeKeyForRow ( keyAggs , i ) ; if ( res ) { fieldAggMakers . forEach ( makeProcessFieldAgg ( getFieldAggs ( res . key , i ) , i ) ) ; } } return buckets ; } private _generateBuilders ( size ) { const builders = new Map < keyof T , Builder < any > > ( ) ; for ( const col of this . columns ) { if ( ! this . _selectFields ?. length || this . _selectFields . includes ( col . name ) ) { const agg = this . _aggregations . get ( col . name ) ; const builder = getBuilderForField ( col , size , agg ?. key , agg ?. value ) ; builders . set ( col . name , builder ) ; } } return builders ; } private _buildBucket ( builders , bucket ) { if ( bucket [ 0 ] . adjustsSelectedRow ) { return this . _buildBucketWithAdjustedRowIndex ( builders , bucket ) ; } const [ fieldAggs , startIndex ] = bucket ; for ( const [ field , builder ] of builders ) { const agg = fieldAggs . get ( field ) ; if ( agg != null ) { builder . append ( agg . flush ( ) . value ) ; } else { const value = this . getColumnOrThrow ( field ) . vector . get ( startIndex ) ; const writeIndex = builder . currentIndex ++ ; if ( value != null ) { builder . data . set ( writeIndex , value ) ; } } } } /**
     * this is slower version of the aggregation
     * that will select a specific row that shows more
     * correct information, like for min and max
    */ private _buildBucketWithAdjustedRowIndex ( builders , [ fieldAggs , startIndex ] ) { let useIndex = startIndex ; const remainingFields = [ ] ; for ( const [ field , builder ] of builders ) { const agg = fieldAggs . get ( field ) ; if ( agg != null ) { const res = agg . flush ( ) ; if ( res . index != null && res . index > useIndex ) { useIndex = res . index ; } builder . append ( res . value ) ; } else { remainingFields . push ( field ) ; } } for ( const field of remainingFields ) { const builder = builders . get ( field ) ! ; const value = this . getColumnOrThrow ( field ) . vector . get ( useIndex ) ; const writeIndex = builder . currentIndex ++ ; if ( value != null ) { builder . data . set ( writeIndex , value ) ; } } } private _aggregationBuilders ( ) { const fieldAggMakers = new Map < keyof T , FieldAggMaker > ( ) ; const keyAggs = new Map < keyof T , KeyAggFn > ( ) ; for ( const col of this . columns ) { const agg = this . _aggregations . get ( col . name ) ; const isInGroupBy = this . _groupByFields . includes ( col . name ) ; if ( isInGroupBy ) { keyAggs . set ( col . name , makeUniqueKeyAgg ( col . vector ) ) ; } if ( agg ) { if ( agg . value ) { fieldAggMakers . set ( col . name , curryFieldAgg ( agg . value , col . vector ) ) ; } if ( agg . key ) { if ( isInGroupBy ) { throw new Error ( `Invalid to combination of groupBy and ${ agg . key } for field ${ col . name }` ) ; } keyAggs . set ( col . name , keyAggMap [ agg . key ] ( col . vector ) ) ; } } } return { fieldAggMakers , keyAggs } ; } } type AggObject = { key ? : KeyAggregation ; value ? : ValueAggregation ; } ; type WithAlias < T extends Record < string , unknown > , A extends string , V > = { [ P in ( keyof T ) | A ] : V ; } ; type FieldAggMaker = [ fieldAgg : ( ) => FieldAgg , vector : Vector < any > ] ; type FieldAggsMap < T extends Record < string , unknown > > = Map < keyof T , FieldAgg > & { adjustsSelectedRow : boolean ; } ; type Bucket < T extends Record < string , unknown > > = [ fieldAggs : FieldAggsMap < T > , startIndex : number ] ; interface AggBuilders < T extends Record < string , any > > { fieldAggMakers : Map < keyof T , FieldAggMaker > ; keyAggs : Map < keyof T , KeyAggFn > ; } /**
 * AggregationFrame options
*/ export interface AggregationFrameOptions { name ? : string ; metadata ? : Record < string , any > ; } function curryFieldAgg ( valueAgg , vector ) { return [ ( ) => valueAggMap [ valueAgg ] ( vector ) , vector ] ; } function makeGetFieldAggs < T extends Record < string , any > > ( buckets ) { return function getFieldAggs ( key , index ) : FieldAggsMap < T > { const fieldAggRes = buckets . get ( key ) ; if ( ! fieldAggRes ) { const fieldAggs = new Map ( ) as FieldAggsMap < T > ; fieldAggs . adjustsSelectedRow = false ; buckets . set ( key , [ fieldAggs , index ] ) ; return fieldAggs ; } return fieldAggRes [ 0 ] ; } ; } function makeProcessFieldAgg < T extends Record < string , any > > ( fieldAggs , index ) { return function processFieldAgg ( maker , field ) { let agg = fieldAggs . get ( field ) ; if ( ! agg ) { agg = maker [ 0 ] ( ) ; fieldAggs . set ( field , agg ) ; } if ( agg . adjustsSelectedRow ) { fieldAggs . adjustsSelectedRow = true ; } agg . push ( maker [ 1 ] . get ( index ) , index ) ; } ; }