From 144a3520ffd0f94c48dae689a7b7a71014dcb5cc Mon Sep 17 00:00:00 2001 From: Hadia Ahmed Date: Fri, 9 Oct 2026 14:29:45 -0700 Subject: [PATCH 1/2] [run all] fix and test --- pydough/conversion/hybrid_translator.py | 8 +++- tests/test_pipeline_tpch_custom.py | 40 +++++++++++++++++++ .../chained_partition_transformed_key.txt | 3 ++ .../simple_pydough_functions.py | 16 ++++++++ ...chained_partition_transformed_key_ansi.sql | 10 +++++ ...d_partition_transformed_key_databricks.sql | 10 +++++ ...ained_partition_transformed_key_duckdb.sql | 10 +++++ ...hained_partition_transformed_key_mysql.sql | 10 +++++ ...ained_partition_transformed_key_oracle.sql | 10 +++++ ...ned_partition_transformed_key_postgres.sql | 10 +++++ ...ed_partition_transformed_key_snowflake.sql | 10 +++++ ...ained_partition_transformed_key_sqlite.sql | 10 +++++ ...hained_partition_transformed_key_trino.sql | 10 +++++ 13 files changed, 156 insertions(+), 1 deletion(-) create mode 100644 tests/test_plan_refsols/chained_partition_transformed_key.txt create mode 100644 tests/test_sql_refsols/chained_partition_transformed_key_ansi.sql create mode 100644 tests/test_sql_refsols/chained_partition_transformed_key_databricks.sql create mode 100644 tests/test_sql_refsols/chained_partition_transformed_key_duckdb.sql create mode 100644 tests/test_sql_refsols/chained_partition_transformed_key_mysql.sql create mode 100644 tests/test_sql_refsols/chained_partition_transformed_key_oracle.sql create mode 100644 tests/test_sql_refsols/chained_partition_transformed_key_postgres.sql create mode 100644 tests/test_sql_refsols/chained_partition_transformed_key_snowflake.sql create mode 100644 tests/test_sql_refsols/chained_partition_transformed_key_sqlite.sql create mode 100644 tests/test_sql_refsols/chained_partition_transformed_key_trino.sql diff --git a/pydough/conversion/hybrid_translator.py b/pydough/conversion/hybrid_translator.py index 81edb90e5..de3961631 100644 --- a/pydough/conversion/hybrid_translator.py +++ b/pydough/conversion/hybrid_translator.py @@ -1706,7 +1706,13 @@ def make_hybrid_tree( 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 diff --git a/tests/test_pipeline_tpch_custom.py b/tests/test_pipeline_tpch_custom.py index c2b4f9b07..04c9ae6ac 100644 --- a/tests/test_pipeline_tpch_custom.py +++ b/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: diff --git a/tests/test_plan_refsols/chained_partition_transformed_key.txt b/tests/test_plan_refsols/chained_partition_transformed_key.txt new file mode 100644 index 000000000..6a905cb0e --- /dev/null +++ b/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}) diff --git a/tests/test_pydough_functions/simple_pydough_functions.py b/tests/test_pydough_functions/simple_pydough_functions.py index 64fd47698..aa9ea4d41 100644 --- a/tests/test_pydough_functions/simple_pydough_functions.py +++ b/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()) + ) diff --git a/tests/test_sql_refsols/chained_partition_transformed_key_ansi.sql b/tests/test_sql_refsols/chained_partition_transformed_key_ansi.sql new file mode 100644 index 000000000..6bc142ce8 --- /dev/null +++ b/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 diff --git a/tests/test_sql_refsols/chained_partition_transformed_key_databricks.sql b/tests/test_sql_refsols/chained_partition_transformed_key_databricks.sql new file mode 100644 index 000000000..6bc142ce8 --- /dev/null +++ b/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 diff --git a/tests/test_sql_refsols/chained_partition_transformed_key_duckdb.sql b/tests/test_sql_refsols/chained_partition_transformed_key_duckdb.sql new file mode 100644 index 000000000..0612156a6 --- /dev/null +++ b/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 diff --git a/tests/test_sql_refsols/chained_partition_transformed_key_mysql.sql b/tests/test_sql_refsols/chained_partition_transformed_key_mysql.sql new file mode 100644 index 000000000..c809a4c18 --- /dev/null +++ b/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 diff --git a/tests/test_sql_refsols/chained_partition_transformed_key_oracle.sql b/tests/test_sql_refsols/chained_partition_transformed_key_oracle.sql new file mode 100644 index 000000000..a03439b34 --- /dev/null +++ b/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 diff --git a/tests/test_sql_refsols/chained_partition_transformed_key_postgres.sql b/tests/test_sql_refsols/chained_partition_transformed_key_postgres.sql new file mode 100644 index 000000000..0612156a6 --- /dev/null +++ b/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 diff --git a/tests/test_sql_refsols/chained_partition_transformed_key_snowflake.sql b/tests/test_sql_refsols/chained_partition_transformed_key_snowflake.sql new file mode 100644 index 000000000..0612156a6 --- /dev/null +++ b/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 diff --git a/tests/test_sql_refsols/chained_partition_transformed_key_sqlite.sql b/tests/test_sql_refsols/chained_partition_transformed_key_sqlite.sql new file mode 100644 index 000000000..6bc142ce8 --- /dev/null +++ b/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 diff --git a/tests/test_sql_refsols/chained_partition_transformed_key_trino.sql b/tests/test_sql_refsols/chained_partition_transformed_key_trino.sql new file mode 100644 index 000000000..0612156a6 --- /dev/null +++ b/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 From 3028a43b0729912f7c1eafc1b52aa5685c5a7277 Mon Sep 17 00:00:00 2001 From: Hadia Ahmed Date: Mon, 12 Oct 2026 15:38:20 -0700 Subject: [PATCH 2/2] [run ci] add sorted --- pydough/conversion/hybrid_translator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pydough/conversion/hybrid_translator.py b/pydough/conversion/hybrid_translator.py index de3961631..7e1f4ce10 100644 --- a/pydough/conversion/hybrid_translator.py +++ b/pydough/conversion/hybrid_translator.py @@ -1700,7 +1700,7 @@ 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