sadpandajoe commented on code in PR #45793:
URL: https://github.com/apache/superset/pull/45793#discussion_r4393778828
##########
superset/db_engine_specs/clickhouse.py:
##########
@@ -72,6 +73,16 @@ class ClickHouseQueryStoppedBeforeDispatchError(Exception):
"""
+class AggregateFunctionState(types.UserDefinedType):
+ cache_ok = True
+
+ def get_col_spec(self, **kw: Any) -> str:
+ return "AggregateFunction"
+
+ def column_expression(self, colexpr: ColumnElement) -> ColumnElement:
+ return func.finalizeAggregation(colexpr)
Review Comment:
This hook applies to every column that carries this type, including
calculated columns: `TableColumn.get_sqla_col()` passes the mapped type to
`literal_column(expression, type_=...)` too. A dataset column that kept the
reflected `AGGREGATEFUNCTION` type and already uses `finalizeAggregation(rtt)`
as its expression (a natural workaround for this bug) would now compile to
`finalizeAggregation((finalizeAggregation(rtt)))`, and ClickHouse rejects the
outer call because its argument is no longer an aggregate state. Should
explicit-expression columns skip the wrapper, or is there another way to keep
those datasets working?
##########
tests/unit_tests/db_engine_specs/test_clickhouse.py:
##########
@@ -1107,3 +1114,107 @@ def
test_legacy_clickhouse_engine_spec_has_no_cancel_support() -> None:
cursor_mock = Mock()
assert spec.get_cancel_query_id(cursor_mock, Query()) is None
assert spec.cancel_query(cursor_mock, Query(), "123") is False
+
+
[email protected](
+ "native_type",
+ [
+ "AggregateFunction(quantileTiming, UInt16)",
+ "AggregateFunction(uniqHLL12, String)",
+ "AGGREGATEFUNCTION",
+ ],
+)
+def test_aggregate_function_column_spec(native_type: str) -> None:
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ AggregateFunctionState,
+ ClickHouseConnectEngineSpec as spec,
+ )
+
+ assert_column_spec(
+ spec,
+ native_type,
+ AggregateFunctionState,
+ None,
+ GenericDataType.NUMERIC,
+ False,
+ )
+ assert AggregateFunctionState().compile() == "AggregateFunction"
+
+
+def test_select_star_finalizes_aggregate_function_columns() -> None:
+ from clickhouse_connect.cc_sqlalchemy.datatypes.base import
sqla_type_from_name
+ from clickhouse_connect.cc_sqlalchemy.dialect import ClickHouseDialect
+
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.sql.parse import Table
+ from superset.superset_typing import ResultSetColumnType
+
+ dialect = ClickHouseDialect()
+ database = Mock()
+ database.compile_sqla_query = lambda query, catalog, schema: str(
+ query.compile(dialect=dialect, compile_kwargs={"literal_binds": True})
+ )
+ cols: list[ResultSetColumnType] = [
+ {
+ "name": name,
+ "column_name": name,
+ "type": sqla_type_from_name(native_type),
+ "is_dttm": False,
+ }
+ for name, native_type in [
+ ("rows", "SimpleAggregateFunction(sum, UInt64)"),
+ ("rtt", "AggregateFunction(quantileTiming, UInt16)"),
+ ]
+ ]
+
+ sql = spec.select_star(
+ database=database,
+ table=Table("t"),
+ dialect=dialect,
+ limit=10,
+ show_cols=True,
+ indent=False,
+ latest_partition=False,
+ cols=cols,
+ )
+
+ assert sql == (
+ "SELECT `rows`, finalizeAggregation(`rtt`) AS `rtt` \nFROM `t`\n LIMIT
10"
+ )
+
+
+def test_dataset_query_finalizes_aggregate_function_columns(
+ app_context: None,
+) -> None:
+ from superset.connectors.sqla.models import SqlaTable, TableColumn
+ from superset.models.core import Database
+
+ table = SqlaTable(
+ database=Database(
+ database_name="ch",
sqlalchemy_uri="clickhousedb://localhost/default"
+ ),
+ schema="default",
+ table_name="t",
+ columns=[
+ TableColumn(column_name="ts", type="DATETIME", is_dttm=True),
+ TableColumn(column_name="rtt", type="AGGREGATEFUNCTION"),
+ ],
+ )
+
+ sqlaq = table.get_sqla_query(
+ columns=["ts", "rtt"],
+ metrics=None,
Review Comment:
The description says finalization applies only in the SELECT list, and that
existing merge-style metrics keep working, but the dataset test passes
`metrics=None` and orders by the plain `ts` column. If the wrapper ever leaked
into a `uniqMerge(rtt)` metric argument, or into a filter or `ORDER BY` on
`rtt`, it would feed a scalar where ClickHouse needs a state and nothing here
would fail. Could you add a compiled-query case that includes a merge metric
over `rtt` plus a filter and an `ORDER BY rtt`, asserting `finalizeAggregation`
shows up only in the projection (or not at all for the metric)?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
