From 1d424691d6ca19f44c1c069344cde4deb2426817 Mon Sep 17 00:00:00 2001 From: Afroz Alam Date: Fri, 26 Jan 2024 23:13:08 +0000 Subject: [PATCH] SNOW-945984: Fix sql simplifier when non select statement is used with limit (#1216) * fix sql simplifier when non select statement is used with limit * changelog updates --- CHANGELOG.md | 3 +++ .../snowpark/_internal/analyzer/select_statement.py | 2 ++ tests/integ/test_simplifier_suite.py | 7 +++++++ 3 files changed, 12 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 61378e71652..6c7a2a20f9a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,12 +34,15 @@ For instance, `df.select("a", seq1().alias("b")).select("a", "b").sort("a")` won't flatten the sort clause anymore. - a window or sequence-dependent data generator column is used after `DataFrame.limit()`. For instance, `df.limit(10).select(row_number().over())` won't flatten the limit and select in the generated SQL. - Fixed a bug that aliasing a DataFrame column raises an error when the DataFame is copied from another DataFrame with an aliased column. For instance, + ```python df = df.select(col("a").alias("b")) df = copy(df) df.select(col("b").alias("c")) # threw an error. Now it's fixed. ``` + - Fixed a bug in `Session.create_dataframe` that the non-nullable field in schema is not respected for boolean type. Note that this fix is only effective when the user to have the privilege to create a temp table. +- Fixed a bug in sql simplifier where non-select statements in `session.sql` dropped sql query when used with `limit()`. - Fixed a bug that raised an exception when session parameter `ERROR_ON_NONDETERMINISTIC_UPDATE` is true. ### Behavior Changes (API Compatible) diff --git a/src/snowflake/snowpark/_internal/analyzer/select_statement.py b/src/snowflake/snowpark/_internal/analyzer/select_statement.py index 9eb359be295..50e7819ee81 100644 --- a/src/snowflake/snowpark/_internal/analyzer/select_statement.py +++ b/src/snowflake/snowpark/_internal/analyzer/select_statement.py @@ -811,6 +811,8 @@ def limit(self, n: int, *, offset: int = 0) -> "SelectStatement": new.limit_ = min(self.limit_, n) if self.limit_ else n new.offset = offset or self.offset new.column_states = self.column_states + new.pre_actions = new.from_.pre_actions + new.post_actions = new.from_.post_actions return new diff --git a/tests/integ/test_simplifier_suite.py b/tests/integ/test_simplifier_suite.py index e7813b9af31..b93d71cbc85 100644 --- a/tests/integ/test_simplifier_suite.py +++ b/tests/integ/test_simplifier_suite.py @@ -735,6 +735,13 @@ def test_limit(session, simplifier_table): == f"SELECT * FROM (select * from {simplifier_table}) LIMIT 10" ) + # test for non-select sql statement + temp_table_name = Utils.random_table_name() + query = f"create or replace temporary table {temp_table_name} (bar string)" + df = session.sql(query).limit(1) + assert df.queries["queries"][-2] == query + assert df.collect()[0][0] == f"Table {temp_table_name} successfully created." + def test_filter_order_limit_together(session, simplifier_table): df = session.table(simplifier_table)