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

"""
PySpark parity tests for Issue #339: Column subscript notation for struct access.

Tests that Column subscript notation (e.g., F.col("StructVal")["E1"]) matches PySpark behavior.
Uses get_imports from fixture only; same logic for both backends.
"""

from sparkless.testing import get_imports

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


class TestColumnSubscriptParity:
    """PySpark parity tests for Column subscript notation."""

    def test_column_subscript_single_field_parity(self):
        """Test that Column subscript notation matches PySpark for single field."""
        spark = SparkSession.builder.appName("issue-339-parity").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "StructVal": {"E1": 1, "E2": 2}},
                    {"Name": "Bob", "StructVal": {"E1": 3, "E2": 4}},
                ]
            )

            # Test subscript notation
            result = df.withColumn("Extract-E1", F.col("StructVal")["E1"])
            rows = result.collect()

            assert len(rows) == 2
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            bob_row = next(row for row in rows if row["Name"] == "Bob")

            assert alice_row["Extract-E1"] == 1
            assert bob_row["Extract-E1"] == 3
        finally:
            spark.stop()

    def test_column_subscript_multiple_fields_parity(self):
        """Test that Column subscript notation matches PySpark for multiple fields."""
        spark = SparkSession.builder.appName("issue-339-parity").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "StructVal": {"E1": 1, "E2": 2}},
                    {"Name": "Bob", "StructVal": {"E1": 3, "E2": 4}},
                ]
            )

            result = df.withColumn("Extract-E1", F.col("StructVal")["E1"]).withColumn(
                "Extract-E2", F.col("StructVal")["E2"]
            )
            rows = result.collect()

            assert len(rows) == 2
            alice_row = next(row for row in rows if row["Name"] == "Alice")
            bob_row = next(row for row in rows if row["Name"] == "Bob")

            assert alice_row["Extract-E1"] == 1
            assert alice_row["Extract-E2"] == 2
            assert bob_row["Extract-E1"] == 3
            assert bob_row["Extract-E2"] == 4
        finally:
            spark.stop()

    def test_column_subscript_equals_dot_notation_parity(self):
        """Test that subscript notation produces same results as dot notation in PySpark."""
        spark = SparkSession.builder.appName("issue-339-parity").getOrCreate()
        try:
            df = spark.createDataFrame(
                [
                    {"Name": "Alice", "StructVal": {"E1": 1, "E2": 2}},
                    {"Name": "Bob", "StructVal": {"E1": 3, "E2": 4}},
                ]
            )

            result_subscript = df.withColumn("Extract-E1", F.col("StructVal")["E1"])
            result_dot = df.withColumn("Extract-E1", F.col("StructVal.E1"))

            rows_subscript = result_subscript.collect()
            rows_dot = result_dot.collect()

            assert len(rows_subscript) == len(rows_dot) == 2

            for sub_row, dot_row in zip(rows_subscript, rows_dot):
                assert sub_row["Extract-E1"] == dot_row["Extract-E1"]
        finally:
            spark.stop()