HyukjinKwon commented on code in PR #39225: URL: https://github.com/apache/spark/pull/39225#discussion_r1057438509
########## python/pyspark/sql/connect/window.py: ########## @@ -218,27 +218,201 @@ class Window: @staticmethod def partitionBy(*cols: Union["ColumnOrName", List["ColumnOrName"]]) -> "WindowSpec": + """ + Examples + -------- + >>> from pyspark.sql.connect.window import Window + >>> from pyspark.sql.connect.functions import row_number + >>> df = spark.createDataFrame( + ... [(1, "a"), (1, "a"), (2, "a"), (1, "b"), (2, "b"), (3, "b")], ["id", "category"]) + >>> df.show() + +---+--------+ + | id|category| + +---+--------+ + | 1| a| + | 1| a| + | 2| a| + | 1| b| + | 2| b| + | 3| b| + +---+--------+ + + Show row number order by ``id`` in partition ``category``. + + >>> window = Window.partitionBy("category").orderBy("id") + >>> df.withColumn("row_number", row_number().over(window)).show() + +---+--------+----------+ + | id|category|row_number| + +---+--------+----------+ + | 1| a| 1| + | 1| a| 2| + | 2| a| 3| + | 1| b| 1| + | 2| b| 2| + | 3| b| 3| + +---+--------+----------+ + """ return Window._spec.partitionBy(*cols) partitionBy.__doc__ = PySparkWindow.partitionBy.__doc__ @staticmethod def orderBy(*cols: Union["ColumnOrName", List["ColumnOrName"]]) -> "WindowSpec": + """ + Examples + -------- + >>> from pyspark.sql.connect.window import Window + >>> from pyspark.sql.connect.functions import row_number + >>> df = spark.createDataFrame( + ... [(1, "a"), (1, "a"), (2, "a"), (1, "b"), (2, "b"), (3, "b")], ["id", "category"]) + >>> df.show() + +---+--------+ + | id|category| + +---+--------+ + | 1| a| + | 1| a| + | 2| a| + | 1| b| + | 2| b| + | 3| b| + +---+--------+ + + Show row number order by ``category`` in partition ``id``. + + >>> window = Window.partitionBy("id").orderBy("category") + >>> df.withColumn("row_number", row_number().over(window)).show() + +---+--------+----------+ + | id|category|row_number| + +---+--------+----------+ + | 1| a| 1| + | 1| a| 2| + | 1| b| 3| + | 2| a| 1| + | 2| b| 2| + | 3| b| 1| + +---+--------+----------+ + """ return Window._spec.orderBy(*cols) orderBy.__doc__ = PySparkWindow.orderBy.__doc__ @staticmethod def rowsBetween(start: int, end: int) -> "WindowSpec": + """ + Examples + -------- + >>> from pyspark.sql.connect.window import Window + >>> from pyspark.sql.connect import functions as func + >>> df = spark.createDataFrame( + ... [(1, "a"), (1, "a"), (2, "a"), (1, "b"), (2, "b"), (3, "b")], ["id", "category"]) + >>> df.show() + +---+--------+ + | id|category| + +---+--------+ + | 1| a| + | 1| a| + | 2| a| + | 1| b| + | 2| b| + | 3| b| + +---+--------+ + + Calculate sum of ``id`` in the range from currentRow to currentRow + 1 + in partition ``category`` + + >>> window = Window.partitionBy("category").orderBy("id").rowsBetween(Window.currentRow, 1) + >>> df.withColumn("sum", func.sum("id").over(window)).sort("id", "category", "sum").show() + +---+--------+---+ + | id|category|sum| + +---+--------+---+ + | 1| a| 2| + | 1| a| 3| + | 1| b| 3| + | 2| a| 2| + | 2| b| 5| + | 3| b| 3| + +---+--------+---+ + + """ return Window._spec.rowsBetween(start, end) rowsBetween.__doc__ = PySparkWindow.rowsBetween.__doc__ @staticmethod def rangeBetween(start: int, end: int) -> "WindowSpec": + """ + >>> from pyspark.sql.connect.window import Window + >>> from pyspark.sql.connect import functions as func + >>> df = spark.createDataFrame( + ... [(1, "a"), (1, "a"), (2, "a"), (1, "b"), (2, "b"), (3, "b")], ["id", "category"]) + >>> df.show() + +---+--------+ + | id|category| + +---+--------+ + | 1| a| + | 1| a| + | 2| a| + | 1| b| + | 2| b| + | 3| b| + +---+--------+ + + Calculate sum of ``id`` in the range from ``id`` of currentRow to ``id`` of currentRow + 1 + in partition ``category`` + + >>> window = Window.partitionBy("category").orderBy("id").rangeBetween(Window.currentRow, 1) + >>> df.withColumn("sum", func.sum("id").over(window)).sort("id", "category").show() + +---+--------+---+ + | id|category|sum| + +---+--------+---+ + | 1| a| 4| + | 1| a| 4| + | 1| b| 3| + | 2| a| 2| + | 2| b| 5| + | 3| b| 3| + +---+--------+---+ + + """ return Window._spec.rangeBetween(start, end) rangeBetween.__doc__ = PySparkWindow.rangeBetween.__doc__ Window.__doc__ = Window.__doc__ + + +def _test() -> None: + import os + import sys + import doctest + from pyspark import SparkContext + from pyspark.sql import SparkSession + from pyspark.testing.connectutils import should_test_connect, connect_requirement_message + import pyspark.sql.connect.window + + if should_test_connect: + globs = pyspark.sql.window.__dict__.copy() + sc = SparkContext("local[4]", "sql.connect.window tests") + globs["_spark"] = SparkSession(sc) Review Comment: yeah looks pretty good -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org