[Local Testing] SNOW-783554 Support selecting window expressions #1049

sfc-gh-stan · 2023-09-13T22:46:06Z

Make non-rank related functions with unspecified window frame use cumulative window by default
Support aggregating window expressions
Support more rank related functions/Refactor to allow patching for rank related functions (This will be addressed in a separate PR).

tests/integ/scala/test_window_frame_suite.py

src/snowflake/snowpark/mock/plan.py

src/snowflake/snowpark/mock/functions.py

src/snowflake/snowpark/mock/plan.py

sfc-gh-aling · 2023-09-21T20:29:36Z

src/snowflake/snowpark/mock/functions.py

@@ -137,7 +139,7 @@ def mock_avg(column: ColumnEmulator) -> ColumnEmulator:
            cnt += 1
    # round to 5 according to snowflake spec
    ret = (
-        ColumnEmulator(data=[round((ret / cnt), 5)])
+        ColumnEmulator(data=[round((ret / cnt), 3)])


there are other places in this module where we round to 5 (I did this previously)
so do we also need to update other mock implementations?

Yea, I will address these altogether in a separate PR.

src/snowflake/snowpark/mock/plan.py

sfc-gh-yixie · 2023-09-28T17:21:10Z

I tried this

from snowflake.snowpark import Session, Window
from snowflake.snowpark.functions import min as min_
from snowflake.snowpark.mock.connection import MockServerConnection

session = Session(MockServerConnection())

df = session.create_dataframe([[1, 2], [3, 4]], schema=["a", "b"])
df.select(min_("b").over(Window.partition_by("a").order_by("b"))).show()
df.select(min_("b").over(Window.partition_by("a").order_by("b").range_between(0, 1))).show()

Two problems:

The first show outputs column name min("B") OVER (PARTITION BY "A" ORDER BY "B" A.... I guess Snowflake uses all upper case.
the 2nd doesn't work. AttributeError: 'numpy.ndarray' object has no attribute 'loc'.

src/snowflake/snowpark/mock/functions.py

sfc-gh-stan · 2023-10-11T17:40:26Z

src/snowflake/snowpark/mock/analyzer.py

@@ -285,7 +308,7 @@ def analyze(
                self.session._conn._telemetry_client.send_function_usage_telemetry(
                    expr.api_call_source, TelemetryField.FUNC_CAT_USAGE.value
                )
-            func_name = expr.name.upper() if parse_local_name else expr.name


This change always converts function names to uppercase. cc: @sfc-gh-yixie

sfc-gh-stan · 2023-10-11T17:44:47Z

src/snowflake/snowpark/mock/snowflake_data_type.py

@@ -270,6 +270,8 @@ def add_date_and_number(


 class ColumnEmulator(pd.Series):
+    _metadata = ["sf_type"]


Adding this to preserve sf_type after index operations like reset_index or sort_index.

I don't see this is used anyhere?

This is needed by pandas internally to preserve added property in subclasses: https://pandas.pydata.org/docs/development/extending.html#define-original-properties

shall we add a comment here?

sfc-gh-stan · 2023-10-11T17:46:53Z

tests/integ/scala/test_window_frame_suite.py

@@ -251,20 +265,30 @@ def test_range_between_should_accept_at_most_one_order_by_expression_when_unboun
        df.select(
            min_("key").over(window.range_between(Window.unboundedPreceding, 1))
        ).collect()
-    assert "Cumulative window frame unsupported for function MIN" in str(ex_info)
+        if not local_testing_mode:
+            assert "Cumulative window frame unsupported for function MIN" in str(


The error message doesn't make sense. This error is caused by window.range_between(Window.unboundedPreceding, 1), you cannot specify range between for non-cumulative window frames. As an evidence, changing the expression to window.range_between(Window.unboundedPreceding, 0) runs regardless of the number of order expressions.

sfc-gh-stan · 2023-10-11T17:47:47Z

tests/integ/scala/test_window_frame_suite.py

@@ -283,19 +307,28 @@ def test_range_between_should_accept_numeric_values_only_when_bounded(session):
        df.select(
            min_("value").over(window.range_between(Window.unboundedPreceding, 1))
        ).collect()
-    assert "Cumulative window frame unsupported for function MIN" in str(ex_info)
+        if not local_testing_mode:


Same issue as above, the error message is confusing and the actual failure is not caused by numeric values.

We should rewrite these two tests for both live connection and local testing.

sfc-gh-stan · 2023-10-11T17:50:46Z

I tried this
from snowflake.snowpark import Session, Window
from snowflake.snowpark.functions import min as min_
from snowflake.snowpark.mock.connection import MockServerConnection

session = Session(MockServerConnection())

df = session.create_dataframe([[1, 2], [3, 4]], schema=["a", "b"])
df.select(min_("b").over(Window.partition_by("a").order_by("b"))).show()
df.select(min_("b").over(Window.partition_by("a").order_by("b").range_between(0, 1))).show()
Two problems:

The first show outputs column name min("B") OVER (PARTITION BY "A" ORDER BY "B" A.... I guess Snowflake uses all upper case.

the 2nd doesn't work. AttributeError: 'numpy.ndarray' object has no attribute 'loc'.

Fixed.

src/snowflake/snowpark/mock/plan.py

tests/integ/scala/test_window_spec_suite.py

src/snowflake/snowpark/mock/plan.py

sfc-gh-aling · 2023-10-17T23:18:16Z

src/snowflake/snowpark/mock/plan.py

+            )  # dtype=object prevents implicit converting None to Nan
+            res_col.index = res_index
+            return res_col.sort_index()
+        elif isinstance(window_function, FirstValue):


nit: it feels that FirstValue and LastValue calculation are sharing lots of logic
except for the w.iloc[0] vs w.iloc[len(w) - 1] and for cur_idx in range(len(w)): vs for cur_idx in range(len(w) - 1, -1, -1):.

This is very true. I am planning to refactor this part in a separate PR so all window function definitions can be found in mock/functions.py.

let's create a JIRA or log it in an existing jira so that we don't forget?

src/snowflake/snowpark/mock/plan.py

sfc-gh-stan commented Sep 14, 2023

View reviewed changes

tests/integ/scala/test_window_frame_suite.py Outdated Show resolved Hide resolved

sfc-gh-stan marked this pull request as ready for review September 14, 2023 19:23

sfc-gh-stan requested a review from a team as a code owner September 14, 2023 19:23

sfc-gh-stan requested review from sfc-gh-mkeller and sfc-gh-sfan and removed request for a team September 14, 2023 19:23

sfc-gh-stan commented Sep 14, 2023

View reviewed changes

src/snowflake/snowpark/mock/plan.py Outdated Show resolved Hide resolved

sfc-gh-stan commented Sep 15, 2023

View reviewed changes

src/snowflake/snowpark/mock/functions.py Show resolved Hide resolved

sfc-gh-stan commented Sep 15, 2023

View reviewed changes

src/snowflake/snowpark/mock/plan.py Outdated Show resolved Hide resolved

sfc-gh-stan force-pushed the local/support-window-functions branch from be640d9 to d07b157 Compare September 19, 2023 17:26

sfc-gh-aling reviewed Sep 21, 2023

View reviewed changes

src/snowflake/snowpark/mock/plan.py Show resolved Hide resolved

sfc-gh-aling reviewed Sep 21, 2023

View reviewed changes

src/snowflake/snowpark/mock/plan.py Show resolved Hide resolved

sfc-gh-aling reviewed Sep 21, 2023

View reviewed changes

src/snowflake/snowpark/mock/plan.py Outdated Show resolved Hide resolved

sfc-gh-aling reviewed Sep 21, 2023

View reviewed changes

src/snowflake/snowpark/mock/plan.py Show resolved Hide resolved

sfc-gh-stan force-pushed the local/support-window-functions branch from d07b157 to 40a2e7e Compare September 25, 2023 13:38

sfc-gh-stan requested a review from sfc-gh-yixie September 26, 2023 17:45

sfc-gh-yixie reviewed Sep 28, 2023

View reviewed changes

src/snowflake/snowpark/mock/functions.py Show resolved Hide resolved

sfc-gh-stan force-pushed the local/support-window-functions branch from 4934593 to 56b2034 Compare October 5, 2023 19:11

sfc-gh-stan added 7 commits October 11, 2023 11:22

Add changes

a0fb50d

Clean up

340994f

Support FirstValue and LastValue

49e8982

Address comments

fc8d026

Enable null inputs test

800db15

Fix bugs

79ac67c

Fix bugs and add tests

71ca36c

sfc-gh-stan force-pushed the local/support-window-functions branch from 56b2034 to 71ca36c Compare October 11, 2023 17:38

sfc-gh-stan commented Oct 11, 2023

View reviewed changes

Clean up

5fc9943

sfc-gh-stan commented Oct 11, 2023

View reviewed changes

sfc-gh-stan requested review from sfc-gh-yixie, sfc-gh-aling and sfc-gh-jdu and removed request for sfc-gh-mkeller October 11, 2023 17:50

sfc-gh-stan commented Oct 11, 2023

View reviewed changes

src/snowflake/snowpark/mock/plan.py Outdated Show resolved Hide resolved

Fix range frame and handle sorting better

6b40905

sfc-gh-stan mentioned this pull request Oct 12, 2023

[Local Testing] SNOW-904259 Enable DataFrame.drop_duplicates tests #1085

Closed

sfc-gh-stan commented Oct 12, 2023

View reviewed changes

tests/integ/scala/test_window_spec_suite.py Outdated Show resolved Hide resolved

sfc-gh-aling reviewed Oct 13, 2023

View reviewed changes

Address comments

c0f0d8e

sfc-gh-stan requested a review from sfc-gh-aling October 16, 2023 19:00

sfc-gh-aling reviewed Oct 17, 2023

View reviewed changes

src/snowflake/snowpark/mock/plan.py Outdated Show resolved Hide resolved

sfc-gh-aling reviewed Oct 17, 2023

View reviewed changes

src/snowflake/snowpark/mock/plan.py Outdated Show resolved Hide resolved

sfc-gh-aling reviewed Oct 17, 2023

View reviewed changes

src/snowflake/snowpark/mock/plan.py Outdated Show resolved Hide resolved

sfc-gh-aling reviewed Oct 17, 2023

View reviewed changes

src/snowflake/snowpark/mock/plan.py Outdated Show resolved Hide resolved

Address comments

413a5ca

sfc-gh-stan requested a review from sfc-gh-aling October 17, 2023 23:59

sfc-gh-aling approved these changes Oct 18, 2023

View reviewed changes

sfc-gh-stan merged commit 2ad4e3e into dev/local-testing Oct 18, 2023
8 of 39 checks passed

sfc-gh-stan deleted the local/support-window-functions branch October 18, 2023 22:22

github-actions bot locked and limited conversation to collaborators Oct 18, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Local Testing] SNOW-783554 Support selecting window expressions #1049

[Local Testing] SNOW-783554 Support selecting window expressions #1049

sfc-gh-stan commented Sep 13, 2023 •

edited

Loading

sfc-gh-aling Sep 21, 2023

sfc-gh-stan Sep 25, 2023

sfc-gh-yixie commented Sep 28, 2023

sfc-gh-stan Oct 11, 2023 •

edited

Loading

sfc-gh-stan Oct 11, 2023

sfc-gh-aling Oct 17, 2023

sfc-gh-stan Oct 17, 2023

sfc-gh-aling Oct 18, 2023

sfc-gh-stan Oct 11, 2023

sfc-gh-stan Oct 11, 2023

sfc-gh-stan Oct 11, 2023

sfc-gh-stan commented Oct 11, 2023

sfc-gh-aling Oct 17, 2023

sfc-gh-stan Oct 17, 2023

sfc-gh-aling Oct 18, 2023

		@@ -270,6 +270,8 @@ def add_date_and_number(


		class ColumnEmulator(pd.Series):
		_metadata = ["sf_type"]

[Local Testing] SNOW-783554 Support selecting window expressions #1049

[Local Testing] SNOW-783554 Support selecting window expressions #1049

Conversation

sfc-gh-stan commented Sep 13, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-yixie commented Sep 28, 2023

sfc-gh-stan Oct 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-stan commented Oct 11, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-stan commented Sep 13, 2023 •

edited

Loading

sfc-gh-stan Oct 11, 2023 •

edited

Loading