Add Aggregate Function geomean (mosaic-sql)

docs/api/sql/aggregate-functions.md

@@ -99,6 +99,12 @@ Create an aggregate function that calculates the sum of the input _expression_.
 
 Create an aggregate function that calculates the product of the input _expression_.
 
+## geomean
+
+`geomean(expression)`
+
+Create an aggregate function that calculates the geometric mean of the input _expression_.
+
 ## median
 
 `median(expression)`

packages/core/src/preagg/sufficient-statistics.js

@@ -1,4 +1,4 @@
-import { AggregateNode, and, argmax, argmin, count, div, ExprNode, isNotNull, max, min, mul, pow, regrAvgX, regrAvgY, regrCount, sql, sqrt, sub, sum } from '@uwdata/mosaic-sql';
+import { AggregateNode, and, argmax, argmin, count, div, exp, ExprNode, isNotNull, ln, max, min, mul, pow, regrAvgX, regrAvgY, regrCount, sql, sqrt, sub, sum } from '@uwdata/mosaic-sql';
 import { fnv_hash } from '../util/hash.js';
 
 /**
@@ -18,6 +18,8 @@ export function sufficientStatistics(node, preagg, avg) {
       return sumExpr(preagg, node);
     case 'avg':
       return avgExpr(preagg, node);
+    case 'geomean':
+      return geomeanExpr(preagg, node);
     case 'arg_max':
       return argmaxExpr(preagg, node);
     case 'arg_min':
@@ -155,6 +157,24 @@ function avgExpr(preagg, node) {
   return div(sum(mul(as, name)), expr);
 }
 
+/**
+ * Generate an expression for calculating geometric means over data dimensions.
+ * This method uses log-based computations to ensure numerical stability. The
+ * geomean calculation uses two sufficient statistics: the sum of log values
+ * and the count of non-null values. As a side effect, this method adds columns
+ * for these statistics to the input *preagg* object.
+ * @param {Record<string, ExprNode>} preagg A map of columns (such as
+ *  sufficient statistics) to pre-aggregate.
+ * @param {AggregateNode} node The originating aggregate function call.
+ * @returns {ExprNode} An aggregate expression over pre-aggregated dimensions.
+ */
+function geomeanExpr(preagg, node) {
+  const x = node.args[0];
+  const expr = addStat(preagg, sum(ln(x)), node);
+  const { expr: n } = countExpr(preagg, node);
+  return exp(div(sum(expr), n));
+}
+
 /**
  * Generate an expression for calculating argmax over data dimensions.
  * As a side effect, this method adds a column to the input *preagg* object

packages/core/test/preaggregator.test.js

@@ -1,5 +1,5 @@
 import { describe, it, expect } from 'vitest';
-import { Query, add, argmax, argmin, avg, corr, count, covarPop, covariance, gt, isNotDistinct, literal, loadObjects, max, min, product, regrAvgX, regrAvgY, regrCount, regrIntercept, regrR2, regrSXX, regrSXY, regrSYY, regrSlope, stddev, stddevPop, sum, varPop, variance } from '@uwdata/mosaic-sql';
+import { Query, add, argmax, argmin, avg, corr, count, covarPop, covariance, geomean, gt, isNotDistinct, literal, loadObjects, max, min, product, regrAvgX, regrAvgY, regrCount, regrIntercept, regrR2, regrSXX, regrSXY, regrSYY, regrSlope, stddev, stddevPop, sum, varPop, variance } from '@uwdata/mosaic-sql';
 import { Coordinator, Selection } from '../src/index.js';
 import { nodeConnector } from './util/node-connector.js';
 import { TestClient } from './util/test-client.js';
@@ -61,6 +61,13 @@ describe('PreAggregator', () => {
     expect(await run(avg('x'))).toStrictEqual([3.5, true]);
   });
 
+  it('supports geomean aggregate', async () => {
+    const [result, optimized] = await run(geomean('x'));
+
+    expect(result).toBeCloseTo(Math.sqrt(12), 10);
+    expect(optimized).toBe(true);
+  });
+
   it('supports min aggregate', async () => {
     expect(await run(min('x'))).toStrictEqual([3, true]);
   });

packages/spec/src/config/transforms.js

@@ -18,6 +18,7 @@ export function transformNames(overrides = []) {
     'dateMonthDay',
     'dateDay',
     'first',
+    'geomean',
     'geojson',
     'last',
     'max',

packages/spec/src/spec/Transform.ts

@@ -229,6 +229,14 @@ export interface First extends AggregateOptions, WindowOptions {
   first: Arg1;
 }
 
+/* A geometric mean aggregate transform. */
+export interface Geomean extends AggregateOptions, WindowOptions {
+  /**
+   * Compute the geometric mean value of the given column.
+   */
+  geomean: Arg1;
+}
+
 /* A last aggregate transform. */
 export interface Last extends AggregateOptions, WindowOptions {
   /**

packages/sql/src/functions/aggregate.js

@@ -98,6 +98,15 @@ export function first(expr) {
   return aggFn('first', expr);
 }
 
+/**
+ * Compute a geomean aggregate.
+ * @param {import('../types.js').ExprValue} expr The expression to aggregate.
+ * @returns {AggregateNode} A SQL aggregate function call.
+ */
+export function geomean(expr) {
+  return aggFn('geomean', expr);
+}
+
 /**
  * Compute a sample kurtosis aggregate.
  * @param {import('../types.js').ExprValue} expr The expression to aggregate.

packages/sql/src/index.js

@@ -24,7 +24,7 @@ export { VerbatimNode } from './ast/verbatim.js';
 export { WindowClauseNode, WindowDefNode, WindowFrameNode, WindowFunctionNode, WindowNode } from './ast/window.js';
 export { WithClauseNode } from './ast/with.js';
 
-export { argmax, argmin, arrayAgg, avg, corr, count, covariance, covarPop, entropy, first, kurtosis, mad, max, median, min, mode, last, product, quantile, regrAvgX, regrAvgY, regrCount, regrIntercept, regrR2, regrSXX, regrSXY, regrSYY, regrSlope, skewness, stddev, stddevPop, stringAgg, sum, variance, varPop } from './functions/aggregate.js';
+export { argmax, argmin, arrayAgg, avg, corr, count, covariance, covarPop, entropy, first, geomean, kurtosis, mad, max, median, min, mode, last, product, quantile, regrAvgX, regrAvgY, regrCount, regrIntercept, regrR2, regrSXX, regrSXY, regrSYY, regrSlope, skewness, stddev, stddevPop, stringAgg, sum, variance, varPop } from './functions/aggregate.js';
 export { cond } from './functions/case.js';
 export { cast, float32, float64, int32 } from './functions/cast.js';
 export { column } from './functions/column.js';

packages/sql/test/aggregate.test.js

@@ -1,6 +1,6 @@
 import { expect, describe, it } from 'vitest';
 import { columns } from './util/columns.js';
-import { argmax, argmin, arrayAgg, avg, column, corr, count, covariance, covarPop, entropy, first, gt, kurtosis, last, mad, max, median, min, mode, product, quantile, regrAvgX, regrAvgY, regrCount, regrIntercept, regrR2, regrSlope, regrSXX, regrSXY, regrSYY, skewness, stddev, stddevPop, stringAgg, sum, variance, varPop } from '../src/index.js';
+import { argmax, argmin, arrayAgg, avg, column, corr, count, covariance, covarPop, entropy, first, geomean, gt, kurtosis, last, mad, max, median, min, mode, product, quantile, regrAvgX, regrAvgY, regrCount, regrIntercept, regrR2, regrSlope, regrSXX, regrSXY, regrSYY, skewness, stddev, stddevPop, stringAgg, sum, variance, varPop } from '../src/index.js';
 
 describe('Aggregate functions', () => {
   it('include accessible metadata', () => {
@@ -63,6 +63,10 @@ describe('Aggregate functions', () => {
     expect(String(first('foo'))).toBe('first("foo")');
   });
 
+  it('include geomean', () => {
+    expect(String(geomean('foo'))).toBe('geomean("foo")');
+  });
+
   it('include kurtosis', () => {
     expect(String(kurtosis('foo'))).toBe('kurtosis("foo")');
   });

packages/vgplot/src/api.js

@@ -28,6 +28,7 @@ export {
   covarPop,
   entropy,
   first,
+  geomean,
   kurtosis,
   mad,
   max,