Split extend into pre/post-summarize for calculated columns vs measures

Calculated columns (non-aggregate expressions like strcat, iff) now get
their extend clause placed BEFORE summarize so they exist for group-by.
Calculated measures (expressions referencing aggregate aliases) remain
AFTER summarize where they can reference aggregated values.

Tracks all aggregate aliases in a separate set to avoid overwrites when
multiple measures use the same aggregate expression (e.g. two count(*)).

Adds test_pre_aggregated_calculations from PR review comment to verify
calculated columns used as group-by dimensions work correctly.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: Alison Gim

sqlalchemy_kusto/dialect_kql.py

@@ -142,13 +142,17 @@ def visit_select(
                 )
                 compiled_query_lines.append(f"| where {converted_where_clause}")
 
-        # Add summarize first if it exists
+        # Add pre-extend (calculated columns) before summarize
+        if "pre_extend" in projections_parts_dict:
+            compiled_query_lines.append(projections_parts_dict.pop("pre_extend"))
+
+        # Add summarize
         if "summarize" in projections_parts_dict:
             compiled_query_lines.append(projections_parts_dict.pop("summarize"))
 
-        # Then add extend after summarize
-        if "extend" in projections_parts_dict:
-            compiled_query_lines.append(projections_parts_dict.pop("extend"))
+        # Add post-extend (calculated measures) after summarize
+        if "post_extend" in projections_parts_dict:
+            compiled_query_lines.append(projections_parts_dict.pop("post_extend"))
 
         # Add remaining parts (project, sort)
         for statement_part in projections_parts_dict.values():
@@ -204,7 +208,8 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s
         group_by_cols = select._group_by_clauses
         order_by_cols = select._order_by_clauses
         summarize_statement = ""
-        extend_statement = ""
+        pre_extend_statement = ""
+        post_extend_statement = ""
         project_statement = ""
         has_aggregates = False
         # The following is the logic
@@ -219,8 +224,10 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s
         #                N---> Add to projection
         if columns is not None:
             summarize_columns = set()
-            extend_columns = set()
+            pre_extend_columns = []
+            post_extend_columns = []
             projection_columns = []
+            all_agg_aliases: set[str] = set()
             for column in [c for c in columns if c.name != "*"]:
                 column_name, column_alias = self._extract_column_name_and_alias(column)
                 column_alias = self._escape_and_quote_columns(column_alias, True)
@@ -232,13 +239,24 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s
                     summarize_columns.add(
                         self._build_column_projection(kql_agg, column_alias)
                     )
+                    if column_alias:
+                        all_agg_aliases.add(column_alias)
                 # No group by clause
                 # Do the columns have aliases ?
                 # Add additional and to handle case where : SELECT column_name as column_name
                 elif column_alias and column_alias != column_name:
-                    extend_columns.add(
-                        self._build_column_projection(column_name, column_alias, True)
+                    extend_entry = self._build_column_projection(
+                        column_name, column_alias, True
                     )
+                    escaped_expr = self._escape_and_quote_columns(column_name)
+
+                    # Calculated measures (aggregate-dependent) go after summarize;
+                    # calculated columns (no aggregate dependency) go before summarize
+                    if any(alias in escaped_expr for alias in all_agg_aliases):
+                        post_extend_columns.append(extend_entry)
+                    else:
+                        pre_extend_columns.append(extend_entry)
+
                 if column_alias:
                     projection_columns.append(
                         self._escape_and_quote_columns(column_alias, True)
@@ -257,8 +275,10 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s
                     summarize_statement = (
                         f"{summarize_statement} by {', '.join(by_columns)}"
                     )
-            if extend_columns:
-                extend_statement = f"| extend {', '.join(sorted(extend_columns))}"
+            if pre_extend_columns:
+                pre_extend_statement = f"| extend {', '.join(pre_extend_columns)}"
+            if post_extend_columns:
+                post_extend_statement = f"| extend {', '.join(post_extend_columns)}"
             project_statement = (
                 f"| project {', '.join(projection_columns)}"
                 if projection_columns
@@ -269,7 +289,8 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s
             f"| order by {', '.join(unwrapped_order_by)}" if unwrapped_order_by else ""
         )
         return {
-            "extend": extend_statement,
+            "pre_extend": pre_extend_statement,
+            "post_extend": post_extend_statement,
             "summarize": summarize_statement,
             "project": project_statement,
             "sort": sort_statement,

tests/unit/test_dialect_kql.py

@@ -173,12 +173,12 @@ def test_group_by_text():
     query_compiled = str(
         query.compile(engine, compile_kwargs={"literal_binds": True})
     ).replace("\n", "")
-    # raw query text from query
+    # raw query text from query - order matches column appearance
     query_expected = (
         '["ActiveUsersLastMonth"]'
+        '| extend ["EventInfo_Time"] = ["EventInfo_Time"] / time(1d), '
+        '["ActiveUserMetric"] = ["ActiveUsers"]'
         '| summarize   by ["EventInfo_Time"] / time(1d)'
-        '| extend ["ActiveUserMetric"] = ["ActiveUsers"], '
-        '["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)'
         '| project ["EventInfo_Time"], ["ActiveUserMetric"]'
         '| order by ["ActiveUserMetric"] desc'
     )
@@ -202,12 +202,13 @@ def test_function_text(f: str, expected: str):
     query_compiled = str(
         query.compile(engine, compile_kwargs={"literal_binds": True})
     ).replace("\n", "")
+    # Order matches column appearance in select
     query_expected = (
         '["ActiveUsersLastMonth"]'
-        '| extend ["ActiveUserMetric"] = ["ActiveUsers"], '
-        '["EventInfo_Time"] = '
+        '| extend ["EventInfo_Time"] = '
         + expected
-        + '| project ["EventInfo_Time"], ["ActiveUserMetric"]'
+        + ', ["ActiveUserMetric"] = ["ActiveUsers"]'
+        '| project ["EventInfo_Time"], ["ActiveUserMetric"]'
     )
     assert query_compiled == query_expected
 
@@ -226,8 +227,8 @@ def test_group_by_text_vaccine_dataset():
     ).replace("\n", "")
     query_expected = (
         'database("superset").["CovidVaccineData"]'
-        '| summarize   by ["country_name"]'
         '| extend ["country_name"] = ["country_name"]'
+        '| summarize   by ["country_name"]'
         '| project ["country_name"]'
         '| order by ["country_name"] asc'
     )
@@ -328,8 +329,8 @@ def test_distinct_count_by_text():
     # raw query text from query
     query_expected = (
         '["ActiveUsersLastMonth"]'
-        '| summarize ["DistinctUsers"] = dcount(["ActiveUsers"])  by ["EventInfo_Time"] / time(1d)'
         '| extend ["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)'
+        '| summarize ["DistinctUsers"] = dcount(["ActiveUsers"])  by ["EventInfo_Time"] / time(1d)'
         '| project ["EventInfo_Time"], ["DistinctUsers"]'
         '| order by ["ActiveUserMetric"] desc'
     )
@@ -354,8 +355,8 @@ def test_distinct_count_alt_by_text():
     # raw query text from query
     query_expected = (
         '["ActiveUsersLastMonth"]'
-        '| summarize ["DistinctUsers"] = dcount(["ActiveUsers"])  by ["EventInfo_Time"] / time(1d)'
         '| extend ["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)'
+        '| summarize ["DistinctUsers"] = dcount(["ActiveUsers"])  by ["EventInfo_Time"] / time(1d)'
         '| project ["EventInfo_Time"], ["DistinctUsers"]'
         '| order by ["ActiveUserMetric"] desc'
     )
@@ -571,6 +572,37 @@ def test_calculated_measure_with_adhoc_measure_and_constant():
     assert query_compiled == query_expected
 
 
+def test_pre_aggregated_calculations():
+    """Test from PR #48 review: calculated column before aggregation.
+
+    A calculated column (strcat) is created and used as a group-by dimension.
+    The extend must appear before summarize so the column exists for grouping.
+    """
+    app_col = literal_column("App")
+    ns_col = literal_column("Namespace")
+    id_col = literal_column("_id")
+
+    app_namespace = sa.func.strcat(app_col, ns_col).label("App_Namespace")
+
+    query = (
+        select(app_namespace, sa.func.count(id_col).label("TotalLogs"))
+        .select_from(text("Logs"))
+        .group_by(app_namespace)
+    )
+
+    query_compiled = str(
+        query.compile(engine, compile_kwargs={"literal_binds": True})
+    ).replace("\n", "")
+
+    query_expected = (
+        '["Logs"]'
+        '| extend ["App_Namespace"] = strcat(App, Namespace)'
+        '| summarize ["TotalLogs"] = count(["_id"])  by ["App_Namespace"]'
+        '| project ["App_Namespace"], ["TotalLogs"]'
+    )
+    assert query_compiled == query_expected
+
+
 @pytest.mark.parametrize(
     ("query_table_name", "expected_table_name"),
     [