robin-sparkless 4.4.0

PySpark-like DataFrame API in Rust on Polars; no JVM.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164

"""
PySpark parity tests for DataFrame window operations.

Tests validate that Sparkless window operations behave identically to PySpark.
"""

from tests.tools.parity_base import ParityTestBase
from sparkless.testing import get_imports
from sparkless.testing import Mode, get_mode


def _is_pyspark_mode() -> bool:
    """Check if running in PySpark mode."""
    backend = get_mode()
    return backend == Mode.PYSPARK


class TestWindowOperationsParity(ParityTestBase):
    """Test DataFrame window operations parity with PySpark."""

    def test_row_number(self, spark):
        """Test row_number window function matches PySpark behavior."""
        imports = get_imports()
        F = imports.F
        Window = imports.Window
        expected = self.load_expected("windows", "row_number")

        df = spark.createDataFrame(expected["input_data"])
        window_spec = Window.partitionBy("dept").orderBy("salary")
        result = df.withColumn("row_num", F.row_number().over(window_spec))

        self.assert_parity(result, expected)

    def test_rank(self, spark):
        """Test rank window function matches PySpark behavior."""
        imports = get_imports()
        F = imports.F
        Window = imports.Window
        expected = self.load_expected("windows", "rank")

        df = spark.createDataFrame(expected["input_data"])
        window_spec = Window.partitionBy("dept").orderBy("salary")
        result = df.withColumn("rank", F.rank().over(window_spec))

        self.assert_parity(result, expected)

    def test_dense_rank(self, spark):
        """Test dense_rank window function matches PySpark behavior."""
        imports = get_imports()
        F = imports.F
        Window = imports.Window
        expected = self.load_expected("windows", "dense_rank")

        df = spark.createDataFrame(expected["input_data"])
        window_spec = Window.partitionBy("dept").orderBy("salary")
        result = df.withColumn("dense_rank", F.dense_rank().over(window_spec))

        self.assert_parity(result, expected)

    def test_sum_over_window(self, spark):
        """Test sum over window matches PySpark behavior."""
        imports = get_imports()
        F = imports.F
        Window = imports.Window
        expected = self.load_expected("windows", "sum_over_window")

        df = spark.createDataFrame(expected["input_data"])
        # Match PySpark fixture: total salary per department, repeated on each
        # row for that department (not a running total).
        window_spec = Window.partitionBy("department")
        result = df.withColumn("dept_total", F.sum("salary").over(window_spec))

        self.assert_parity(result, expected)

    def test_lag(self, spark):
        """Test lag window function matches PySpark behavior."""
        imports = get_imports()
        F = imports.F
        Window = imports.Window
        expected = self.load_expected("windows", "lag")

        df = spark.createDataFrame(expected["input_data"])
        window_spec = Window.partitionBy("dept").orderBy("salary")
        result = df.withColumn("lag_salary", F.lag("salary", 1).over(window_spec))

        self.assert_parity(result, expected)

    def test_lead(self, spark):
        """Test lead window function matches PySpark behavior."""
        imports = get_imports()
        F = imports.F
        Window = imports.Window
        expected = self.load_expected("windows", "lead")

        df = spark.createDataFrame(expected["input_data"])
        window_spec = Window.partitionBy("dept").orderBy("salary")
        result = df.withColumn("lead_salary", F.lead("salary", 1).over(window_spec))

        self.assert_parity(result, expected)

    def test_cume_dist(self, spark):
        """Test cume_dist window function matches PySpark behavior."""
        imports = get_imports()
        F = imports.F
        Window = imports.Window
        expected = self.load_expected("windows", "cume_dist")

        df = spark.createDataFrame(expected["input_data"])
        window_spec = Window.partitionBy("dept").orderBy("salary")
        result = df.withColumn("cume_dist", F.cume_dist().over(window_spec))

        self.assert_parity(result, expected)

    def test_first_value(self, spark):
        """Test first_value window function matches PySpark behavior."""
        imports = get_imports()
        F = imports.F
        Window = imports.Window
        expected = self.load_expected("windows", "first_value")

        df = spark.createDataFrame(expected["input_data"])
        window_spec = Window.partitionBy("dept").orderBy("salary")
        result = df.withColumn("first_salary", F.first("salary").over(window_spec))

        self.assert_parity(result, expected)

    def test_last_value(self, spark):
        """Test last_value window function matches PySpark behavior."""
        imports = get_imports()
        F = imports.F
        Window = imports.Window
        expected = self.load_expected("windows", "last_value")

        df = spark.createDataFrame(expected["input_data"])
        window_spec = Window.partitionBy("dept").orderBy("salary")
        result = df.withColumn("last_salary", F.last("salary").over(window_spec))

        self.assert_parity(result, expected)

    def test_percent_rank(self, spark):
        """Test percent_rank window function matches PySpark behavior."""
        imports = get_imports()
        F = imports.F
        Window = imports.Window
        expected = self.load_expected("windows", "percent_rank")

        df = spark.createDataFrame(expected["input_data"])
        window_spec = Window.partitionBy("dept").orderBy("salary")
        result = df.withColumn("percent_rank", F.percent_rank().over(window_spec))

        self.assert_parity(result, expected)

    def test_ntile(self, spark):
        """Test ntile window function matches PySpark behavior."""
        imports = get_imports()
        F = imports.F
        Window = imports.Window
        expected = self.load_expected("windows", "ntile")

        df = spark.createDataFrame(expected["input_data"])
        window_spec = Window.partitionBy("dept").orderBy("salary")
        result = df.withColumn("ntile", F.ntile(2).over(window_spec))

        self.assert_parity(result, expected)