Basic select API working #1674

sfc-gh-heshah · 2024-05-24T02:51:28Z

Columns are created with no reference to the parent Dataframe in the current Snowpark API
Following this pattern when generating the AST (invoking proto IR entity constructors directly)
ast parameter added to Column constructor to enable Column methods to create and push-down ASTs
ast property added to Column class to enable backwards compatibility with col._expression: Expression using SpColumnSqlExpr as a backup

Current test case output AST

client_version {
  major: 42
}
client_language {
  python_language {
    version {
      major: 3
      minor: 8
    }
  }
}
body {
  assign {
    uid: 1
    var_id {
      bitfield1: 1
    }
    expr {
      sp_table {
        table: "test_table"
      }
    }
  }
}
body {
  assign {
    uid: 2
    var_id {
      bitfield1: 2
    }
    expr {
      sp_dataframe_filter {
        df {
          sp_dataframe_ref {
            id {
              bitfield1: 1
            }
          }
        }
        condition {
          sp_column_sql_expr {
            sql: "STR LIKE \'%e%\'"
          }
        }
      }
    }
  }
}
body {
  assign {
    uid: 3
    var_id {
      bitfield1: 3
    }
    expr {
      sp_dataframe_select__columns {
        df {
          sp_dataframe_ref {
            id {
              bitfield1: 2
            }
          }
        }
        cols {
          sp_column {
            name: "\"A\""
          }
        }
        cols {
          sp_column_sql_expr {
            sql: "\"STR\""
          }
        }
        cols {
          sp_column {
            name: "\"B\""
          }
        }
        cols {
          sp_column_sql_expr {
            sql: "\"STR\""
          }
        }
        cols {
          sp_column_sql_expr {
            sql: "ADD(\"A\", \"B\")"
          }
        }
      }
    }
  }
}

sfc-gh-azwiegincew · 2024-05-24T14:47:09Z

src/snowflake/snowpark/column.py

@@ -79,6 +79,7 @@
    TimestampType,
 )
 from snowflake.snowpark.window import Window, WindowSpec
+import snowflake.snowpark._internal.proto.ast_pb2 as proto


This should be sorted.

sfc-gh-azwiegincew · 2024-05-24T14:59:24Z

tests/thin-client/steel-thread.py

-df = df.filter("STR LIKE '%e%'")
-df.show()
+df = session.table('test_table').filter("STR LIKE '%e%'")
+# col = df.col("STR") # TODO: determine whether this should create and assign statement or not


It should not. We have no visibility into the assignment of Python objects. We create assign statements in the AST that correspond to dataframe-typed values (df, col-expr, etc.).

sfc-gh-azwiegincew · 2024-05-24T15:00:38Z

tests/thin-client/steel-thread.py

-df = session.table('test_table')
-df = df.filter("STR LIKE '%e%'")
-df.show()
+df = session.table('test_table').filter("STR LIKE '%e%'")


Maybe add this as a separate "test" or use cmdline args to choose the test.

sfc-gh-azwiegincew · 2024-05-24T15:06:05Z

src/snowflake/snowpark/dataframe.py

@@ -517,6 +517,7 @@ def __init__(
                             referenced in subsequent dataframe expressions.
        """
        self._session = session
+        self._ast_expr = ast_stmt.expr


There's no reason to retain the AST expression.

sfc-gh-azwiegincew · 2024-05-24T15:15:56Z

src/snowflake/snowpark/dataframe.py

+        stmt = self._session._ast_batch.assign()
+        ast = stmt.expr
+        ast.sp_dataframe_select__columns.df.sp_dataframe_ref.id.bitfield1 = self._ast_id
+        ast.sp_dataframe_select__columns.variadic = False


variadic is true for some forms – the logic seems to be in parse_positional_args_to_list. We have to capture this detail to correctly represent the unparsed form in query history.

sfc-gh-azwiegincew · 2024-05-24T15:16:37Z

src/snowflake/snowpark/dataframe.py

-                names.append(Column(e)._named())
+                col = Column(e)
+                names.append(col._named())
+                ast.sp_dataframe_select__columns.cols.append(col.ast)


This is inefficient. Since the user passes in a string in this case, you're free to construct a new column expr from scratch without the inefficient copying.

This will naturally get refactored due to the previous comment though.

If this were to result in cleaner code, feel free to reimplement the logic, since in Phase 1 we will remove a lot of the existing complexity (we're preserving it for now so that this codebase can be used for Phase 0).

More info at https://protobuf.dev/reference/python/python-generated/#repeated-message-fields

sfc-gh-azwiegincew · 2024-05-24T15:26:39Z

src/snowflake/snowpark/column.py

                    )
+                    # TODO: figure out why basic aliasing is breaking above


What's the problem?

I misinterpreted the goal of this case in the Column constructor. When two strings are provided to the constructor the first is meant to use an alias for the Dataframe created prior to using the Column constructor.

Just to clarify, is the SpColumnAlias IR entity meant to capture statements of the form select colA as A?

Yes. In that case you would have two nested expressions.

sfc-gh-azwiegincew · 2024-05-24T15:26:45Z

src/snowflake/snowpark/column.py

    ) -> None:
+        self._ast = proto.SpColumnExpr() if ast is None else ast


This logic seems wrong. You should defer the construction of the AST to the specific cases below.

sfc-gh-azwiegincew · 2024-05-24T15:28:11Z

src/snowflake/snowpark/column.py

        elif isinstance(expr1, Expression):
            self._expression = expr1
        else:  # pragma: no cover
            raise TypeError("Column constructor only accepts str or expression.")

+    @property
+    def ast(self) -> proto.SpColumnExpr:


This shouldn't be necessary. This logic should happen at construction time for the appropriate case. We shouldn't be introducing imperative mutations with a hard-to-follow control flow in an API without mutable state like this one. This is why I'm suggesting having clear logic in the constructor that immediately constructs the required value without mutations spread across multiple branches and functions.

Would it be okay to still store a constructed _ast member and retrieve it later?

I'm facing an issue with deciding how to handle expressions of the form df.select(df.col("A"), col("B")), where the provided Column instances are constructed before df.select has an opportunity to construct the necessary SpColumnExpr instance and populate the AST fields.

More generally any Dataframe API that accepts previously constructed Column instances will face this issue, hence the current solution. I agree with the point you made for line 1137 though as that is a clear opportunity to avoid copying from a potentially complex message into the Session AstBatch.

Yes, I think you're right and we have to store a pre-constructed value in _ast. This property shouldn't be necessary though. Maybe a helper function to copy _ast into whatever builder is passed in?

Sounds good, would the goal of the helper function be to replace protobuf.message.CopyFrom(col._ast) and append(col._ast) when constructing the AST one level up from the constructor?

I was thinking as long as _ast is ensured to be an instance of SpColumnExpr a helper function might not be necessary, unless there is another purpose I'm missing, please let me know!

Yes. Definitely better without a helper if the code seems simple.

sfc-gh-azwiegincew · 2024-05-24T15:35:57Z

tests/thin-client/steel-thread.py

+df = session.table('test_table').filter("STR LIKE '%e%'")
+# col = df.col("STR") # TODO: determine whether this should create and assign statement or not
+df2 = df.select("A", df.col("STR"), col("B"), df.STR, col("A") + col("B"))
+print(session._ast_batch._request)


Could you attach the output to the PR description please? I wanted to double-check the result.

For sure, just added above

Thanks. I don't think it's right to map df.col() to SpColumnSqlExpr. It should be the same as the others. It's possible that AST construct was checked in by mistake. I don't see any Column APIs that include SQL.

I agree, currently I was using the SpColumnSqlExpr as a backup for APIs which didn't generate the AST yet.

There is however a Column API that allows SQL text through an UnresolvedAttribute Expression subclass I believe:

snowpark-python/src/snowflake/snowpark/column.py

Line 771 in cfe432b

def _expr(cls, e: str) -> "Column":

In the AST, our goal is to capture the original API call. This appears to be an implementation detail which we will be removing in stage 1.

sfc-gh-heshah · 2024-05-25T01:20:24Z

AST output from tests/thin-client/select_outputs.py after changes:

df = session.table("test_table").filter("STR LIKE '%e%'")
df2 = df.select("A", df.col("STR"), col("B"), df.STR)

client_version {
  major: 42
}
client_language {
  python_language {
    version {
      major: 3
      minor: 8
    }
  }
}
body {
  assign {
    uid: 1
    var_id {
      bitfield1: 1
    }
    expr {
      sp_table {
        table: "test_table"
      }
    }
  }
}
body {
  assign {
    uid: 2
    var_id {
      bitfield1: 2
    }
    expr {
      sp_dataframe_filter {
        df {
          sp_dataframe_ref {
            id {
              bitfield1: 1
            }
          }
        }
        condition {
          sp_column_sql_expr {
            sql: "STR LIKE \'%e%\'"
          }
        }
      }
    }
  }
}
body {
  assign {
    uid: 3
    var_id {
      bitfield1: 3
    }
    expr {
      sp_dataframe_select__columns {
        df {
          sp_dataframe_ref {
            id {
              bitfield1: 2
            }
          }
        }
        cols {
          sp_column {
            name: "\"A\""
          }
        }
        cols {
          sp_dataframe_col {
            df {
              sp_dataframe_ref {
                id {
                  bitfield1: 2
                }
              }
            }
            col_name: "\"STR\""
          }
        }
        cols {
          sp_column {
            name: "\"B\""
          }
        }
        cols {
          sp_dataframe_col {
            df {
              sp_dataframe_ref {
                id {
                  bitfield1: 2
                }
              }
            }
            col_name: "\"STR\""
          }
        }
        variadic: true
      }
    }
  }
}

sfc-gh-azwiegincew · 2024-05-28T15:38:48Z

src/snowflake/snowpark/_internal/ast.py

 class AstBatch:
    def __init__(self, session):
        self._session = session
        self._id_gen = itertools.count(start=1)
        self._init_batch()
        # TODO: extended version from the branch snowpark-ir.

-    def assign(self, symbol = None):
+    def assign(self, symbol=None):


This is inconsistent with the rest of the code.

sfc-gh-azwiegincew · 2024-05-28T15:40:26Z

src/snowflake/snowpark/column.py

    ) -> None:
+        if ast is not None:


This should be done without a guard. It's better to rely on a known slot being set to None rather than the absence of the key in the dictionary.

sfc-gh-azwiegincew · 2024-05-28T15:41:36Z

src/snowflake/snowpark/column.py

            if expr1 == "*":
                self._expression = Star([])
+                # TODO: Determine whether SpColumnSqlExpr is correct entity here


This seems correct.

sfc-gh-azwiegincew · 2024-05-28T15:48:13Z

src/snowflake/snowpark/dataframe.py

+        stmt = self._session._ast_batch.assign()
+        ast = stmt.expr
+        ast.sp_dataframe_select__columns.df.sp_dataframe_ref.id.bitfield1 = self._ast_id
+        ast.sp_dataframe_select__columns.variadic = (len(cols) > 1 or isinstance(cols[0], (list, tuple, set)))


This is incorrect. If a single list is passed, we want variadic = False.

sfc-gh-azwiegincew

LG!

sfc-gh-azwiegincew · 2024-05-28T18:00:22Z

src/snowflake/snowpark/column.py

-                self._ast.sp_column.name = col_name
+                self._expression = UnresolvedAttribute(quote_name(expr1))
+
+            # some repitition here, but _expression logic will be eliminated eventually


Basic select API working

cfe432b

sfc-gh-heshah requested a review from a team as a code owner May 24, 2024 02:51

sfc-gh-heshah requested review from sfc-gh-yixie, sfc-gh-yuwang, sfc-gh-jrose and sfc-gh-azwiegincew and removed request for a team May 24, 2024 02:51

sfc-gh-azwiegincew requested changes May 24, 2024

View reviewed changes

sfc-gh-heshah added 2 commits May 24, 2024 11:23

tox:fix-lint changes

36468c3

Address PR comments

70fd615

sfc-gh-azwiegincew requested changes May 28, 2024

View reviewed changes

Minor PR changes

fa4e382

sfc-gh-azwiegincew approved these changes May 28, 2024

View reviewed changes

misspelling

adfe610

sfc-gh-heshah merged commit 8de5d28 into server-side-snowpark May 28, 2024
11 of 13 checks passed

sfc-gh-heshah deleted the server-side-snowpark-colexpr branch May 28, 2024 18:21

github-actions bot locked and limited conversation to collaborators May 28, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Basic select API working #1674

Basic select API working #1674

sfc-gh-heshah commented May 24, 2024 •

edited

Loading

sfc-gh-azwiegincew May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-heshah May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-heshah May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-heshah May 25, 2024

sfc-gh-azwiegincew May 25, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-heshah May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-heshah May 24, 2024

sfc-gh-azwiegincew May 24, 2024

sfc-gh-heshah commented May 25, 2024 •

edited

Loading

sfc-gh-azwiegincew May 28, 2024

sfc-gh-azwiegincew May 28, 2024

sfc-gh-azwiegincew May 28, 2024

sfc-gh-azwiegincew May 28, 2024

sfc-gh-azwiegincew left a comment

sfc-gh-azwiegincew May 28, 2024

		) -> None:
		self._ast = proto.SpColumnExpr() if ast is None else ast

Basic select API working #1674

Basic select API working #1674

Conversation

sfc-gh-heshah commented May 24, 2024 • 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

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

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-heshah commented May 25, 2024 • 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

sfc-gh-azwiegincew left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sfc-gh-heshah commented May 24, 2024 •

edited

Loading

sfc-gh-heshah commented May 25, 2024 •

edited

Loading