From 73291513577c589b01b31faa229ed04150afa468 Mon Sep 17 00:00:00 2001 From: Evgeny Ivanov <14168452+e10v@users.noreply.github.com> Date: Sun, 12 Jul 2026 00:20:54 +0300 Subject: [PATCH] perf: precompute centered products for fallback aggregates --- src/tea_tasting/backends/ibis.py | 45 +++++++++++++------------ src/tea_tasting/backends/narwhals.py | 14 ++++---- src/tea_tasting/backends/sql.py | 50 ++++++++++++---------------- 3 files changed, 52 insertions(+), 57 deletions(-) diff --git a/src/tea_tasting/backends/ibis.py b/src/tea_tasting/backends/ibis.py index caa51b3..0b82b2a 100644 --- a/src/tea_tasting/backends/ibis.py +++ b/src/tea_tasting/backends/ibis.py @@ -25,9 +25,8 @@ import tea_tasting.aggr # noqa: TC004 -_CENTERED = "__centered__{}__" -_CENTERED_LEFT = "__centered__left__{}__{}__" -_CENTERED_RIGHT = "__centered_right__{}__{}__" +_CENTERED_SQUARE = "__centered_square__{}__" +_CENTERED_PRODUCT = "__centered_product__{}__{}__" class IbisTable(BaseTable): # noqa: D101 @@ -176,7 +175,7 @@ def _aggregate( fallback_var_cols = () if ibis_table.has_var else aggr_cols.var_cols fallback_cov_cols = () if ibis_table.has_cov else aggr_cols.cov_cols if len(fallback_var_cols) > 0 or len(fallback_cov_cols) > 0: - data = _add_centered_cols( + data = _add_centered_product_cols( data=data, aggr_cols=aggr_cols, group_col=group_col, @@ -195,7 +194,7 @@ def _aggregate( return _get_aggregates(result, aggr_cols, group_col) -def _add_centered_cols( +def _add_centered_product_cols( data: ibis.expr.types.Table, aggr_cols: tea_tasting.aggr.AggrCols, group_col: str | None, @@ -218,18 +217,25 @@ def _add_centered_cols( for col in var_cols: col_expr = data[col].cast("float") centered_col_expr = col_expr - _mean_over(data, col_expr, group_col) - exprs.append(centered_col_expr.name(_CENTERED.format(col))) + exprs.append( + (centered_col_expr * centered_col_expr) + .name(_CENTERED_SQUARE.format(col)), + ) for left, right in cov_cols: valid = data[left].notnull() & data[right].notnull() left_expr = data[left].cast("float") right_expr = data[right].cast("float") - exprs.extend(( - (left_expr - _mean_over(data, valid.ifelse(left_expr, null), group_col)) - .name(_CENTERED_LEFT.format(left, right)), - (right_expr - _mean_over(data, valid.ifelse(right_expr, null), group_col)) - .name(_CENTERED_RIGHT.format(left, right)), - )) + centered_left = ( + left_expr - _mean_over(data, valid.ifelse(left_expr, null), group_col) + ) + centered_right = ( + right_expr - _mean_over(data, valid.ifelse(right_expr, null), group_col) + ) + exprs.append( + (centered_left * centered_right) + .name(_CENTERED_PRODUCT.format(left, right)), + ) return data.select(*exprs) @@ -278,9 +284,7 @@ def _sample_var( ) -> ibis.expr.types.Value: if has_var: return data[col].cast("float").var(how="sample") - return _fallback_sample_aggr( - data[_CENTERED.format(col)] * data[_CENTERED.format(col)], - ) + return _fallback_sample_aggr(data[_CENTERED_SQUARE.format(col)]) def _sample_cov( @@ -292,16 +296,13 @@ def _sample_cov( ) -> ibis.expr.types.Value: if has_cov: return data[left].cast("float").cov(data[right].cast("float"), how="sample") - return _fallback_sample_aggr( - data[_CENTERED_LEFT.format(left, right)] * - data[_CENTERED_RIGHT.format(left, right)], - ) + return _fallback_sample_aggr(data[_CENTERED_PRODUCT.format(left, right)]) def _fallback_sample_aggr( - centered_expr: ibis.expr.types.Column, + product_expr: ibis.expr.types.Column, ) -> ibis.expr.types.Value: import ibis # noqa: PLC0415 - count = centered_expr.count() - return (count > 1).ifelse(centered_expr.sum() / (count - 1), ibis.null()) # ty:ignore[unsupported-operator] + count = product_expr.count() + return (count > 1).ifelse(product_expr.sum() / (count - 1), ibis.null()) # ty:ignore[unsupported-operator] diff --git a/src/tea_tasting/backends/narwhals.py b/src/tea_tasting/backends/narwhals.py index d8a3a39..748a81e 100644 --- a/src/tea_tasting/backends/narwhals.py +++ b/src/tea_tasting/backends/narwhals.py @@ -30,7 +30,7 @@ type CovarianceSupport = Literal["full", "ungrouped_only", "none", "auto"] -_CROSS_PRODUCT = "__cross_product__{}__{}__" +_CENTERED_PRODUCT = "__centered_product__{}__{}__" _COV_COUNT = "__cov_count__{}__{}__" _COV_SUM = "__cov_sum__{}__{}__" @@ -185,7 +185,7 @@ def _aggregate( narwhals_frame.cov == "ungrouped_only" and group_col is None ) if not has_cov and len(aggr_cols.cov_cols) > 0: - frame = _add_cross_product_cols(frame, aggr_cols, group_col) + frame = _add_centered_product_cols(frame, aggr_cols, group_col) exprs = _aggr_exprs(aggr_cols, has_cov=has_cov) frame = ( @@ -200,7 +200,7 @@ def _aggregate( return _get_aggregates(rows, aggr_cols, group_col) -def _add_cross_product_cols( +def _add_centered_product_cols( frame: nw.LazyFrame, aggr_cols: tea_tasting.aggr.AggrCols, group_col: str | None, @@ -221,7 +221,7 @@ def _add_cross_product_cols( centered_cols.append(valid_col - mean) exprs.append( (centered_cols[0] * centered_cols[1]).alias( - _CROSS_PRODUCT.format(left, right), + _CENTERED_PRODUCT.format(left, right), ), ) return frame.select(*exprs) @@ -248,10 +248,10 @@ def _aggr_exprs( ) else: for left, right in aggr_cols.cov_cols: - cross_product = nw.col(_CROSS_PRODUCT.format(left, right)) + centered_product = nw.col(_CENTERED_PRODUCT.format(left, right)) exprs.extend(( - cross_product.sum().alias(_COV_SUM.format(left, right)), - cross_product.count().alias(_COV_COUNT.format(left, right)), + centered_product.sum().alias(_COV_SUM.format(left, right)), + centered_product.count().alias(_COV_COUNT.format(left, right)), )) return exprs diff --git a/src/tea_tasting/backends/sql.py b/src/tea_tasting/backends/sql.py index 8f2a1d6..978f305 100644 --- a/src/tea_tasting/backends/sql.py +++ b/src/tea_tasting/backends/sql.py @@ -137,9 +137,8 @@ ) | {"singlestore"} _SUBQUERY = "__subquery__" -_CENTERED = "__centered__{}__" -_CENTERED_LEFT = "__centered__left__{}__{}__" -_CENTERED_RIGHT = "__centered_right__{}__{}__" +_CENTERED_SQUARE = "__centered_square__{}__" +_CENTERED_PRODUCT = "__centered_product__{}__{}__" class SQLQuery(BaseTable): # noqa: D101 @@ -318,7 +317,7 @@ def _aggregate( fallback_var_cols = () if sql_query.var else aggr_cols.var_cols fallback_cov_cols = () if sql_query.cov else aggr_cols.cov_cols if len(fallback_var_cols) > 0 or len(fallback_cov_cols) > 0: - query = _add_centered_cols( + query = _add_centered_product_cols( query=query, aggr_cols=aggr_cols, group_col=group_col, @@ -341,7 +340,7 @@ def _aggregate( return _get_aggregates(data, aggr_cols, group_col) -def _add_centered_cols( +def _add_centered_product_cols( query: sqlglot.exp.Subquery, aggr_cols: tea_tasting.aggr.AggrCols, group_col: str | None, @@ -358,13 +357,14 @@ def _add_centered_cols( keep_cols.add(group_col) exprs: list[sqlglot.exp.Expression] = [_col(col) for col in keep_cols] - exprs.extend( - sqlglot.exp.alias_( - _float(col) - _mean_over(_float(col), group_col), - _CENTERED.format(col), + for col in var_cols: + centered = _float(col) - _mean_over(_float(col), group_col) + exprs.append( + sqlglot.exp.alias_( + centered * centered.copy(), + _CENTERED_SQUARE.format(col), + ), ) - for col in var_cols - ) for left, right in cov_cols: valid = ( _col(left).is_(sqlglot.exp.Null()).not_() & @@ -374,16 +374,14 @@ def _add_centered_cols( right_float = _float(right) left_valid = _if_valid(valid, left_float.copy()) right_valid = _if_valid(valid, right_float.copy()) - exprs.extend(( - sqlglot.exp.alias_( - left_float - _mean_over(left_valid, group_col), - _CENTERED_LEFT.format(left, right), - ), + centered_left = left_float - _mean_over(left_valid, group_col) + centered_right = right_float - _mean_over(right_valid, group_col) + exprs.append( sqlglot.exp.alias_( - right_float - _mean_over(right_valid, group_col), - _CENTERED_RIGHT.format(left, right), + centered_left * centered_right, + _CENTERED_PRODUCT.format(left, right), ), - )) + ) return sqlglot.exp.select(*exprs).from_(query).subquery(_SUBQUERY) @@ -422,8 +420,7 @@ def _sample_var(col: str, *, var: bool | str) -> sqlglot.exp.Expression: return sqlglot.exp.Variance(this=_float(col)) if isinstance(var, str): return sqlglot.exp.Anonymous(this=var, expressions=[_float(col)]) - centered = _col(_CENTERED.format(col)) - return _fallback_sample_aggr(centered * centered) + return _fallback_sample_aggr(_col(_CENTERED_SQUARE.format(col))) def _sample_cov( @@ -439,20 +436,17 @@ def _sample_cov( this=cov, expressions=[_float(left), _float(right)], ) - return _fallback_sample_aggr( - _col(_CENTERED_LEFT.format(left, right)) * - _col(_CENTERED_RIGHT.format(left, right)), - ) + return _fallback_sample_aggr(_col(_CENTERED_PRODUCT.format(left, right))) def _fallback_sample_aggr( - centered_expr: sqlglot.exp.Expression, + product_expr: sqlglot.exp.Expression, ) -> sqlglot.exp.Expression: - count = sqlglot.exp.Count(this=centered_expr.copy()) + count = sqlglot.exp.Count(this=product_expr.copy()) one = sqlglot.exp.Literal.number(1) return sqlglot.exp.Case().when( count > one, - sqlglot.exp.Sum(this=centered_expr) / (count.copy() - one.copy()), + sqlglot.exp.Sum(this=product_expr) / (count.copy() - one.copy()), ).else_(sqlglot.exp.Null())