Typescript, adding custom methods to op

[aleixmorgadas]

I have faced a challenge when adding custom functions using addAggregateFunction, for example, and make that method available so that Typescript doesn't complain.

This is the way I did it:

File: arquero-extensions.ts

import { addAggregateFunction } from 'arquero';

const meanArray = {
  create: () => ({
    init: (state: any) => {
      state.sum = 0;
      state.counter = 0;
    },
    add: (state: any, ...values) => {
      for (const val of values) {
        if (val != null) {
          state.sum += val.reduce((acc: number, currentValue: number) => {
            return acc + currentValue;
          }, 0);
          state.counter += val.length;
        }
      }
      return state;
    },
    rem: (state: any, ...values) => {
      for (const val of values) {
        if (val != null) {
          state.sum -= val.reduce((acc: number, currentValue: number) => {
            return acc + currentValue;
          }, 0);
          state.counter -= val.length;
        }
      }
      return state;
    },
    value: (state: any) =>
      state.counter ? state.sum / state.counter : undefined,
  }),
  param: [1, 0], // 1 field input, 0 extra parameters
};

addAggregateFunction('meanArray', meanArray, { override: true });

File: types/arquero-extensions.d.ts

import 'arquero';
import { opApi } from 'arquero/dist/types/op/op-api';

interface MeanArray {
  meanArray: (...fields: any[]) => Op;
}

type Op = MeanArray & opApi;

declare module 'arquero' {
  declare const op: Op;
}

File: tsconfig.json

{
  "compilerOptions": {
    "types": ["./types/arquero-extensions.d.ts"]
  }
}

Example of usage:

import './arquero-extensions';
import { ColumnTable, op, table } from 'arquero';

df
  .rollup({
    mean: (d) =>
      op.meanArray([
        d['A'],
        d['B'],
        d['C'],
      ]),
  })
  .print();

Test:

import { describe, expect, it } from 'vitest';
import { op, table } from 'arquero';
import './arquero-extensions';

describe('arquero-extensions', () => {
  const df = table({
    A: [1, 2, 3],
    B: [1, 1, 1],
    C: [3, 4, 5],
  });

  it('meanArray', () => {
    const test = df.rollup({
      mean: (d) => op.meanArray([d.A, d.B, d.C]),
    });

    expect(+test.get('mean', 0).toFixed(2)).toBe(2.33);
  });
});