Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

SNOW-1738538 Fix multi ast eval validation in decoding expectation test ast #2497

Merged
merged 3 commits into from
Oct 24, 2024
Merged

SNOW-1738538 Fix multi ast eval validation in decoding expectation test ast #2497

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c714710
SNOW-1738538 Fix multi-ast eval validation in decoding expectation te…
sfc-gh-evandenberg Oct 23, 2024
3f1a1ce
Update from review comment
sfc-gh-evandenberg Oct 23, 2024
f2af085
Update
sfc-gh-evandenberg Oct 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 23 additions & 3 deletions src/snowflake/snowpark/_internal/ast_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1342,12 +1342,32 @@ def ClearTempTables(message: proto.Request) -> None:

def base64_str_to_request(base64_str: str) -> proto.Request:
sfc-gh-evandenberg marked this conversation as resolved.
Show resolved Hide resolved
message = proto.Request()
message.ParseFromString(base64.b64decode(base64_str.strip()))
message.ParseFromString(base64.b64decode(base64_str))
return message


def base64_str_to_textproto(base64_str: str) -> str:
request = base64_str_to_request(base64_str)
def merge_requests(requests: List[proto.Request]) -> proto.Request:
"""Merge list of requests into a single request through accumulating the request body segments in same order."""
request = proto.Request()

# Copy the client_version, etc as part of first message.
request.CopyFrom(requests[0])

for next_request in requests[1:]:
for next_stmt in next_request.body:
stmt = request.body.add()
stmt.CopyFrom(next_stmt)

return request


def base64_lines_to_request(base64_lines: str) -> proto.Request:
messages = [base64_str_to_request(s) for s in base64_lines.split("\n")]
return merge_requests(messages)


def base64_lines_to_textproto(base64_str: str) -> str:
request = base64_lines_to_request(base64_str)

# Force a fixed python version to avoid unnecessary diffs
request.client_language.python_language.version.major = 3
Expand Down
200 changes: 200 additions & 0 deletions tests/ast/data/DataFrame.collect.test
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,206 @@ body {
}
}
}
body {
assign {
expr {
sp_dataframe_collect {
case_sensitive: true
id {
bitfield1: 1
}
src {
file: "SRC_POSITION_TEST_MODE"
start_line: 29
}
}
}
symbol {
}
uid: 4
var_id {
bitfield1: 4
}
}
}
body {
eval {
uid: 5
var_id {
bitfield1: 4
}
}
}
body {
assign {
expr {
sp_dataframe_collect {
block: true
case_sensitive: true
id {
bitfield1: 1
}
src {
file: "SRC_POSITION_TEST_MODE"
start_line: 31
}
statement_params {
_1: "SF_PARTNER"
_2: "FAKE_PARTNER"
}
}
}
symbol {
}
uid: 6
var_id {
bitfield1: 6
}
}
}
body {
eval {
uid: 7
var_id {
bitfield1: 6
}
}
}
body {
assign {
expr {
sp_dataframe_collect {
id {
bitfield1: 1
}
log_on_exception: true
src {
file: "SRC_POSITION_TEST_MODE"
start_line: 33
}
statement_params {
_1: "SF_PARTNER"
_2: "FAKE_PARTNER"
}
}
}
symbol {
}
uid: 8
var_id {
bitfield1: 8
}
}
}
body {
eval {
uid: 9
var_id {
bitfield1: 8
}
}
}
body {
assign {
expr {
sp_dataframe_collect {
case_sensitive: true
id {
bitfield1: 1
}
no_wait: true
src {
file: "SRC_POSITION_TEST_MODE"
start_line: 35
}
}
}
symbol {
}
uid: 10
var_id {
bitfield1: 10
}
}
}
body {
eval {
uid: 11
var_id {
bitfield1: 10
}
}
}
body {
assign {
expr {
sp_dataframe_collect {
case_sensitive: true
id {
bitfield1: 1
}
no_wait: true
src {
file: "SRC_POSITION_TEST_MODE"
start_line: 37
}
statement_params {
_1: "SF_PARTNER"
_2: "FAKE_PARTNER"
}
}
}
symbol {
}
uid: 12
var_id {
bitfield1: 12
}
}
}
body {
eval {
uid: 13
var_id {
bitfield1: 12
}
}
}
body {
assign {
expr {
sp_dataframe_collect {
id {
bitfield1: 1
}
log_on_exception: true
no_wait: true
src {
file: "SRC_POSITION_TEST_MODE"
start_line: 39
}
statement_params {
_1: "SF_PARTNER"
_2: "FAKE_PARTNER"
}
}
}
symbol {
}
uid: 14
var_id {
bitfield1: 14
}
}
}
body {
eval {
uid: 15
var_id {
bitfield1: 14
}
}
}
client_ast_version: 1
client_language {
python_language {
Expand Down
33 changes: 33 additions & 0 deletions tests/ast/data/DataFrame.count.test
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,39 @@ body {
}
}
}
body {
assign {
expr {
sp_dataframe_count {
id {
bitfield1: 1
}
src {
file: "SRC_POSITION_TEST_MODE"
start_line: 31
}
statement_params {
_1: "SF_PARTNER"
_2: "FAKE_PARTNER"
}
}
}
symbol {
}
uid: 6
var_id {
bitfield1: 6
}
}
}
body {
eval {
uid: 7
var_id {
bitfield1: 6
}
}
}
client_ast_version: 1
client_language {
python_language {
Expand Down
33 changes: 33 additions & 0 deletions tests/ast/data/DataFrame.count2.test
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,39 @@ body {
}
}
}
body {
assign {
expr {
sp_dataframe_count {
id {
bitfield1: 1
}
src {
file: "SRC_POSITION_TEST_MODE"
start_line: 34
}
statement_params {
_1: "SF_PARTNER"
_2: "FAKE_PARTNER"
}
}
}
symbol {
}
uid: 6
var_id {
bitfield1: 6
}
}
}
body {
eval {
uid: 7
var_id {
bitfield1: 6
}
}
}
client_ast_version: 1
client_language {
python_language {
Expand Down
Loading
Loading