Skip to content

Commit

Permalink
fix doctest
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-aalam committed Aug 28, 2024
1 parent 783f973 commit eb68c91
Showing 1 changed file with 13 additions and 14 deletions.
27 changes: 13 additions & 14 deletions src/snowflake/snowpark/_internal/compiler/large_query_breakdown.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,20 +77,19 @@ class LargeQueryBreakdown:
Example:
For a data pipeline with a large query plan created like so:
>>> base_df = session.sql("select 1 as A, 2 as B")
>>> df1 = base_df.with_column("A", F.col("A") + F.lit(1))
>>> x = 100
>>> for i in range(x):
>>> df1 = df1.with_column("A", F.col("A") + F.lit(i))
>>> df1 = df1.group_by(F.col("A")).agg(F.sum(F.col("B")).alias("B"))
>>> df2 = base_df.with_column("B", F.col("B") + F.lit(1))
>>> for i in range(x):
>>> df2 = df2.with_column("B", F.col("B") + F.lit(i))
>>> df2 = df2.group_by(F.col("B")).agg(F.sum(F.col("A")).alias("A"))
>>> union_df = df1.union_all(df2)
>>> final_df = union_df.with_column("A", F.col("A") + F.lit(1))
base_df = session.sql("select 1 as A, 2 as B")
df1 = base_df.with_column("A", F.col("A") + F.lit(1))
df2 = base_df.with_column("B", F.col("B") + F.lit(1))
for i in range(100):
df1 = df1.with_column("A", F.col("A") + F.lit(i))
df2 = df2.with_column("B", F.col("B") + F.lit(i))
df1 = df1.group_by(F.col("A")).agg(F.sum(F.col("B")).alias("B"))
df2 = df2.group_by(F.col("B")).agg(F.sum(F.col("A")).alias("A"))
union_df = df1.union_all(df2)
final_df = union_df.with_column("A", F.col("A") + F.lit(1))
The corresponding query plan has the following structure:
Expand Down

0 comments on commit eb68c91

Please sign in to comment.