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

"""
PySpark parity tests for Issue #336: WindowFunction comparison. Uses get_imports from fixture only.
"""

import os

import pytest

from sparkless.testing import get_imports

_imports = get_imports()
SparkSession = _imports.SparkSession
F = _imports.F
Window = _imports.Window


class TestWindowFunctionComparisonParity:
    """PySpark parity tests for WindowFunction comparison operators."""

    @pytest.mark.skipif(
        os.getenv("SPARKLESS_TEST_MODE") != "pyspark",
        reason="PySpark parity test - only run with SPARKLESS_TEST_MODE=pyspark",
    )
    def test_window_function_gt_comparison_parity(self):
        """Test WindowFunction > comparison matches PySpark."""
        spark = SparkSession.builder.appName("issue-336-parity").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A"},
                    {"Name": "Bob", "Type": "B"},
                ]
            )

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "GT-Zero",
                F.when(F.row_number().over(w) > 0, F.lit(True)).otherwise(F.lit(False)),
            )
            rows = result.collect()

            assert len(rows) == 2
            assert rows[0]["GT-Zero"] is True
            assert rows[1]["GT-Zero"] is True
        finally:
            spark.stop()

    @pytest.mark.skipif(
        os.getenv("SPARKLESS_TEST_MODE") != "pyspark",
        reason="PySpark parity test - only run with SPARKLESS_TEST_MODE=pyspark",
    )
    def test_window_function_eq_comparison_parity(self):
        """Test WindowFunction == comparison matches PySpark."""
        spark = SparkSession.builder.appName("issue-336-parity").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A"},
                    {"Name": "Bob", "Type": "B"},
                ]
            )

            w = Window().partitionBy("Type").orderBy("Type")
            result = df.withColumn(
                "EQ-One",
                F.when(F.row_number().over(w) == 1, F.lit("First")).otherwise(
                    F.lit("Other")
                ),
            )
            rows = result.collect()

            assert len(rows) == 2
            # Both should be "First" since each is rank 1 in their partition
            assert rows[0]["EQ-One"] == "First"
            assert rows[1]["EQ-One"] == "First"
        finally:
            spark.stop()

    @pytest.mark.skipif(
        os.getenv("SPARKLESS_TEST_MODE") != "pyspark",
        reason="PySpark parity test - only run with SPARKLESS_TEST_MODE=pyspark",
    )
    def test_window_function_comparison_in_filter_parity(self):
        """Test WindowFunction comparison in filter matches PySpark."""
        spark = SparkSession.builder.appName("issue-336-parity").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "Type": "A", "Score": 100},
                    {"Name": "Bob", "Type": "A", "Score": 90},
                    {"Name": "Charlie", "Type": "A", "Score": 80},
                ]
            )

            w = Window().partitionBy("Type").orderBy(F.col("Score").desc())
            result = df.filter(F.row_number().over(w) == 1).select(
                "Name", "Type", "Score"
            )
            rows = result.collect()

            assert len(rows) == 1
            assert rows[0]["Name"] == "Alice"
            assert rows[0]["Score"] == 100
        finally:
            spark.stop()