Fix KeyError When Chaining PARTITION on a Redefined/Transformed Term

pydough/conversion/hybrid_translator.py

@@ -1700,13 +1700,19 @@ def make_hybrid_tree(
                             successor_hybrid, node.child_access, child_ref_mapping
                         )
                         partition_child_idx = child_ref_mapping[0]
-                        for key_name in node.calc_terms:
+                        for key_name in sorted(node.calc_terms, key=str):
                             key = node.get_expr(key_name)
                             expr = self.make_hybrid_expr(
                                 successor_hybrid, key, child_ref_mapping, False
                             )
                             partition.add_key(key_name, expr)
-                            key_exprs.append(HybridRefExpr(key_name, expr.typ))
+                            # Note: uses `expr.name`, not `key_name`, since the
+                            # underlying CALCULATE may have internally renamed
+                            # `key_name` (e.g. to shadow a pre-existing term),
+                            # and `expr.name` reflects that actual,
+                            # possibly-renamed name within the child subtree.
+                            assert isinstance(expr, HybridChildRefExpr)
+                            key_exprs.append(HybridRefExpr(expr.name, expr.typ))
                         successor_hybrid.children[
                             partition_child_idx
                         ].subtree.agg_keys = key_exprs

tests/test_pipeline_tpch_custom.py

@@ -97,6 +97,7 @@
     bad_child_reuse_3,
     bad_child_reuse_4,
     bad_child_reuse_5,
+    chained_partition_transformed_key,
     cross_similar_partitions,
     customer_largest_order_deltas,
     customer_most_recent_orders,
@@ -6817,6 +6818,45 @@
             ),
             id="partition_key_name_collision",
         ),
+        pytest.param(
+            PyDoughPandasTest(
+                chained_partition_transformed_key,
+                "TPCH",
+                lambda: pd.DataFrame(
+                    {
+                        "name": [
+                            "algeria",
+                            "argentina",
+                            "brazil",
+                            "canada",
+                            "china",
+                            "egypt",
+                            "ethiopia",
+                            "france",
+                            "germany",
+                            "india",
+                            "indonesia",
+                            "iran",
+                            "iraq",
+                            "japan",
+                            "jordan",
+                            "kenya",
+                            "morocco",
+                            "mozambique",
+                            "peru",
+                            "romania",
+                            "russia",
+                            "saudi arabia",
+                            "united kingdom",
+                            "united states",
+                            "vietnam",
+                        ],
+                    }
+                ),
+                "chained_partition_transformed_key",
+            ),
+            id="chained_partition_transformed_key",
+        ),
     ],
 )
 def tpch_custom_pipeline_test_data(request) -> PyDoughPandasTest:

tests/test_plan_refsols/chained_partition_transformed_key.txt

@@ -0,0 +1,3 @@
+ROOT(columns=[('name', lower_n_name)], orderings=[(lower_n_name):asc_first])
+ AGGREGATE(keys={'lower_n_name': LOWER(n_name)}, aggregations={})
+  SCAN(table=tpch.NATION, columns={'n_name': n_name})

tests/test_pydough_functions/simple_pydough_functions.py

@@ -3965,3 +3965,19 @@ def partition_key_name_collision():
         )
         .WHERE(key <= 3)
     )
+
+
+def chained_partition_transformed_key():
+    """
+    Test for fixing the problem with chaining a second PARTITION on top of a
+    PARTITION whose key comes from a CALCULATE that redefines a term using
+    its own old value (e.g. `name=LOWER(name)`), used to cause an unhandled
+    `KeyError` during relational conversion.
+    """
+    return (
+        nations.CALCULATE(name=LOWER(name))
+        .PARTITION(name="a", by=name)
+        .PARTITION(name="b", by=name)
+        .CALCULATE(name)
+        .ORDER_BY(name.ASC())
+    )

tests/test_sql_refsols/chained_partition_transformed_key_ansi.sql

@@ -0,0 +1,10 @@
+WITH _t0 AS (
+  SELECT DISTINCT
+    LOWER(n_name) AS lower_n_name
+  FROM tpch.nation
+)
+SELECT
+  lower_n_name AS name
+FROM _t0
+ORDER BY
+  1

tests/test_sql_refsols/chained_partition_transformed_key_databricks.sql

@@ -0,0 +1,10 @@
+WITH _t0 AS (
+  SELECT DISTINCT
+    LOWER(n_name) AS lower_n_name
+  FROM tpch.nation
+)
+SELECT
+  lower_n_name AS name
+FROM _t0
+ORDER BY
+  1

tests/test_sql_refsols/chained_partition_transformed_key_duckdb.sql

@@ -0,0 +1,10 @@
+WITH _t0 AS (
+  SELECT DISTINCT
+    LOWER(n_name) AS lower_n_name
+  FROM tpch.nation
+)
+SELECT
+  lower_n_name AS name
+FROM _t0
+ORDER BY
+  1 NULLS FIRST

tests/test_sql_refsols/chained_partition_transformed_key_mysql.sql

@@ -0,0 +1,10 @@
+WITH _t0 AS (
+  SELECT DISTINCT
+    LOWER(n_name) AS lower_n_name
+  FROM tpch.NATION
+)
+SELECT
+  lower_n_name COLLATE utf8mb4_bin AS name
+FROM _t0
+ORDER BY
+  1

tests/test_sql_refsols/chained_partition_transformed_key_oracle.sql

@@ -0,0 +1,10 @@
+WITH "_T0" AS (
+  SELECT DISTINCT
+    LOWER(n_name) AS LOWER_N_NAME
+  FROM TPCH.NATION
+)
+SELECT
+  LOWER_N_NAME AS name
+FROM "_T0"
+ORDER BY
+  1 NULLS FIRST

tests/test_sql_refsols/chained_partition_transformed_key_postgres.sql

@@ -0,0 +1,10 @@
+WITH _t0 AS (
+  SELECT DISTINCT
+    LOWER(n_name) AS lower_n_name
+  FROM tpch.nation
+)
+SELECT
+  lower_n_name AS name
+FROM _t0
+ORDER BY
+  1 NULLS FIRST

tests/test_sql_refsols/chained_partition_transformed_key_snowflake.sql

@@ -0,0 +1,10 @@
+WITH _t0 AS (
+  SELECT DISTINCT
+    LOWER(n_name) AS lower_n_name
+  FROM tpch.nation
+)
+SELECT
+  lower_n_name AS name
+FROM _t0
+ORDER BY
+  1 NULLS FIRST

tests/test_sql_refsols/chained_partition_transformed_key_sqlite.sql

@@ -0,0 +1,10 @@
+WITH _t0 AS (
+  SELECT DISTINCT
+    LOWER(n_name) AS lower_n_name
+  FROM tpch.nation
+)
+SELECT
+  lower_n_name AS name
+FROM _t0
+ORDER BY
+  1

tests/test_sql_refsols/chained_partition_transformed_key_trino.sql

@@ -0,0 +1,10 @@
+WITH _t0 AS (
+  SELECT DISTINCT
+    LOWER(n_name) AS lower_n_name
+  FROM tpch.nation
+)
+SELECT
+  lower_n_name AS name
+FROM _t0
+ORDER BY
+  1 NULLS FIRST