robin-sparkless 4.0.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

"""
Test for issue #138: Column reference error after drop(). Uses get_spark_imports from fixture only.
"""

from datetime import datetime

from tests.fixtures.spark_imports import get_spark_imports

_imports = get_spark_imports()
SparkSession = _imports.SparkSession
F = _imports.F
col = F.col
to_timestamp = F.to_timestamp
regexp_replace = F.regexp_replace


class TestIssue138ColumnDropReference:
    """Test cases for issue #138: column drop reference errors."""

    def test_drop_column_after_transform(self):
        """Test that dropping a column after transformation works correctly.

        This is the exact scenario from issue #138.
        """
        spark = SparkSession.builder.appName("BugRepro").getOrCreate()
        try:
            data = [("inv1", "2024-01-15T10:30:00", 100)]
            df = spark.createDataFrame(
                data, ["inventory_id", "snapshot_date", "quantity_on_hand"]
            )

            transformed = (
                df.withColumn(
                    "snapshot_date_parsed",
                    to_timestamp(
                        regexp_replace(col("snapshot_date"), r"\.\d+", "").cast(
                            "string"
                        ),
                        "yyyy-MM-dd'T'HH:mm:ss",
                    ),
                ).drop("snapshot_date")  # Drop original column
            )

            # This should not fail with AttributeError
            result = transformed.select("inventory_id", "snapshot_date_parsed")
            count = result.count()
            assert count == 1, f"Expected count 1, got {count}"

            # Verify the data is correct
            rows = result.collect()
            assert len(rows) == 1
            assert rows[0]["inventory_id"] == "inv1"
            assert isinstance(rows[0]["snapshot_date_parsed"], datetime)

        finally:
            spark.stop()

    def test_drop_multiple_columns_after_transform(self):
        """Test dropping multiple columns after transformation."""
        spark = SparkSession.builder.appName("BugRepro").getOrCreate()
        try:
            data = [("inv1", "2024-01-15T10:30:00", 100, "temp")]
            df = spark.createDataFrame(
                data,
                ["inventory_id", "snapshot_date", "quantity_on_hand", "temp_col"],
            )

            transformed = (
                df.withColumn(
                    "snapshot_date_parsed",
                    to_timestamp(
                        regexp_replace(col("snapshot_date"), r"\.\d+", "").cast(
                            "string"
                        ),
                        "yyyy-MM-dd'T'HH:mm:ss",
                    ),
                )
                .drop("snapshot_date")
                .drop("temp_col")
            )

            # Should work without AttributeError
            result = transformed.select(
                "inventory_id", "snapshot_date_parsed", "quantity_on_hand"
            )
            count = result.count()
            assert count == 1

            rows = result.collect()
            assert len(rows) == 1
            assert "snapshot_date" not in rows[0]
            assert "temp_col" not in rows[0]

        finally:
            spark.stop()

    def test_drop_then_select(self):
        """Test selecting columns after dropping some columns."""
        spark = SparkSession.builder.appName("BugRepro").getOrCreate()
        try:
            data = [("inv1", "2024-01-15T10:30:00", 100)]
            df = spark.createDataFrame(
                data, ["inventory_id", "snapshot_date", "quantity_on_hand"]
            )

            transformed = df.withColumn(
                "snapshot_date_parsed",
                to_timestamp(
                    regexp_replace(col("snapshot_date"), r"\.\d+", "").cast("string"),
                    "yyyy-MM-dd'T'HH:mm:ss",
                ),
            ).drop("snapshot_date")

            # Select should work even after drop
            result = transformed.select(
                "inventory_id", "snapshot_date_parsed", "quantity_on_hand"
            )
            count = result.count()
            assert count == 1

        finally:
            spark.stop()

    def test_drop_then_filter(self):
        """Test filtering after dropping columns."""
        spark = SparkSession.builder.appName("BugRepro").getOrCreate()
        try:
            data = [("inv1", "2024-01-15T10:30:00", 100)]
            df = spark.createDataFrame(
                data, ["inventory_id", "snapshot_date", "quantity_on_hand"]
            )

            transformed = df.withColumn(
                "snapshot_date_parsed",
                to_timestamp(
                    regexp_replace(col("snapshot_date"), r"\.\d+", "").cast("string"),
                    "yyyy-MM-dd'T'HH:mm:ss",
                ),
            ).drop("snapshot_date")

            # Filter should work without referencing dropped column
            result = transformed.filter(col("inventory_id") == "inv1")
            count = result.count()
            assert count == 1

        finally:
            spark.stop()