ambers 0.4.2

Pure Rust reader for SPSS .sav and .zsav files
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
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312

"""Tests for ambers.apply_missing().

Covers all SPSS user-defined missing value combinations:
discrete numeric (1-3 values), discrete string, range, range+discrete.

Run with:
    pytest tests/test_apply_missing.py -v
"""

import polars as pl
import pytest
from polars.testing import assert_frame_equal

import ambers as am


# ---------------------------------------------------------------------------
# Helper: build SpssMetadata with missing specs
# ---------------------------------------------------------------------------

def _meta_with_missing(specs: dict) -> am.SpssMetadata:
    return am.SpssMetadata(variable_missing_values=specs)


# ---------------------------------------------------------------------------
# Discrete numeric
# ---------------------------------------------------------------------------

class TestDiscreteNumeric:
    def test_single_value(self):
        df = pl.DataFrame({"Q1": [1.0, 2.0, 99.0, 3.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [99]}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"Q1": [1.0, 2.0, None, 3.0]})
        assert_frame_equal(result, expected)

    def test_two_values(self):
        df = pl.DataFrame({"Q1": [1.0, 98.0, 99.0, 3.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [98, 99]}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"Q1": [1.0, None, None, 3.0]})
        assert_frame_equal(result, expected)

    def test_three_values_max(self):
        df = pl.DataFrame({"Q1": [1.0, 97.0, 98.0, 99.0, 5.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [97, 98, 99]}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"Q1": [1.0, None, None, None, 5.0]})
        assert_frame_equal(result, expected)

    def test_float_values(self):
        df = pl.DataFrame({"score": [1.5, 9.99, 99.99, 50.0]})
        meta = _meta_with_missing({"score": {"type": "discrete", "values": [9.99, 99.99]}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"score": [1.5, None, None, 50.0]})
        assert_frame_equal(result, expected)

    def test_zero_value(self):
        df = pl.DataFrame({"Q1": [0.0, 1.0, 2.0, 3.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [0]}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"Q1": [None, 1.0, 2.0, 3.0]})
        assert_frame_equal(result, expected)

    def test_negative_values(self):
        df = pl.DataFrame({"Q1": [-9.0, -1.0, 1.0, 2.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [-1, -9]}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"Q1": [None, None, 1.0, 2.0]})
        assert_frame_equal(result, expected)


# ---------------------------------------------------------------------------
# Discrete string
# ---------------------------------------------------------------------------

class TestDiscreteString:
    def test_single_string(self):
        df = pl.DataFrame({"city": ["NYC", "NA", "LA"]})
        meta = _meta_with_missing({"city": {"type": "discrete", "values": ["NA"]}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"city": ["NYC", None, "LA"]})
        assert_frame_equal(result, expected)

    def test_multiple_strings(self):
        df = pl.DataFrame({"city": ["NYC", "N/A", "DK", "RF", "LA"]})
        meta = _meta_with_missing({"city": {"type": "discrete", "values": ["N/A", "DK", "RF"]}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"city": ["NYC", None, None, None, "LA"]})
        assert_frame_equal(result, expected)

    def test_string_with_spaces(self):
        df = pl.DataFrame({"note": ["hello", "  ", "world"]})
        meta = _meta_with_missing({"note": {"type": "discrete", "values": ["  "]}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"note": ["hello", None, "world"]})
        assert_frame_equal(result, expected)


# ---------------------------------------------------------------------------
# Range (numeric only)
# ---------------------------------------------------------------------------

class TestRange:
    def test_basic_range(self):
        df = pl.DataFrame({"score": [1.0, 500.0, 900.0, 950.0, 999.0]})
        meta = _meta_with_missing({"score": {"type": "range", "low": 900, "high": 999}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"score": [1.0, 500.0, None, None, None]})
        assert_frame_equal(result, expected)

    def test_range_boundaries_inclusive(self):
        """Both low and high boundaries are inclusive (closed), matching SPSS."""
        df = pl.DataFrame({"score": [899.0, 900.0, 950.0, 999.0, 1000.0]})
        meta = _meta_with_missing({"score": {"type": "range", "low": 900, "high": 999}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"score": [899.0, None, None, None, 1000.0]})
        assert_frame_equal(result, expected)

    def test_range_fractional_bounds(self):
        df = pl.DataFrame({"x": [99.4, 99.5, 99.7, 99.9, 100.0]})
        meta = _meta_with_missing({"x": {"type": "range", "low": 99.5, "high": 99.9}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"x": [99.4, None, None, None, 100.0]})
        assert_frame_equal(result, expected)


# ---------------------------------------------------------------------------
# Range + discrete
# ---------------------------------------------------------------------------

class TestRangeDiscrete:
    def test_range_with_discrete_below(self):
        df = pl.DataFrame({"income": [0.0, 500.0, 999990.0, 999999.0, 50000.0]})
        meta = _meta_with_missing({
            "income": {"type": "range", "low": 999990, "high": 999999, "discrete": 0}
        })
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"income": [None, 500.0, None, None, 50000.0]})
        assert_frame_equal(result, expected)

    def test_range_with_discrete_above(self):
        df = pl.DataFrame({"score": [1.0, 900.0, 999.0, 9999.0, 500.0]})
        meta = _meta_with_missing({
            "score": {"type": "range", "low": 900, "high": 999, "discrete": 9999}
        })
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"score": [1.0, None, None, None, 500.0]})
        assert_frame_equal(result, expected)


# ---------------------------------------------------------------------------
# Multi-column
# ---------------------------------------------------------------------------

class TestMultiColumn:
    def test_different_spec_types(self):
        df = pl.DataFrame({
            "Q1": [1.0, 99.0, 3.0],
            "Q2": [500.0, 900.0, 999.0],
            "Q3": [0.0, 50.0, 999999.0],
        })
        meta = _meta_with_missing({
            "Q1": {"type": "discrete", "values": [99]},
            "Q2": {"type": "range", "low": 900, "high": 999},
            "Q3": {"type": "range", "low": 999990, "high": 999999, "discrete": 0},
        })
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({
            "Q1": [1.0, None, 3.0],
            "Q2": [500.0, None, None],
            "Q3": [None, 50.0, None],
        })
        assert_frame_equal(result, expected)

    def test_some_columns_without_specs(self):
        df = pl.DataFrame({
            "Q1": [1.0, 99.0, 3.0],
            "age": [25.0, 30.0, 35.0],
        })
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [99]}})
        result = am.apply_missing(df, meta)
        # Q1 has missing applied, age is untouched
        assert result["Q1"].to_list() == [1.0, None, 3.0]
        assert result["age"].to_list() == [25.0, 30.0, 35.0]


# ---------------------------------------------------------------------------
# Columns filter
# ---------------------------------------------------------------------------

class TestColumnsFilter:
    def test_columns_none_applies_all(self):
        df = pl.DataFrame({"Q1": [1.0, 99.0], "Q2": [98.0, 2.0]})
        meta = _meta_with_missing({
            "Q1": {"type": "discrete", "values": [99]},
            "Q2": {"type": "discrete", "values": [98]},
        })
        result = am.apply_missing(df, meta, columns=None)
        assert result["Q1"].to_list() == [1.0, None]
        assert result["Q2"].to_list() == [None, 2.0]

    def test_columns_filter_specific(self):
        df = pl.DataFrame({"Q1": [1.0, 99.0], "Q2": [98.0, 2.0]})
        meta = _meta_with_missing({
            "Q1": {"type": "discrete", "values": [99]},
            "Q2": {"type": "discrete", "values": [98]},
        })
        result = am.apply_missing(df, meta, columns=["Q1"])
        assert result["Q1"].to_list() == [1.0, None]
        assert result["Q2"].to_list() == [98.0, 2.0]  # NOT nullified

    def test_column_not_in_df_silent(self):
        """Column in columns= but not in df is silently skipped."""
        df = pl.DataFrame({"Q1": [1.0, 99.0]})
        meta = _meta_with_missing({
            "Q1": {"type": "discrete", "values": [99]},
            "NONEXISTENT": {"type": "discrete", "values": [1]},
        })
        result = am.apply_missing(df, meta, columns=["Q1", "NONEXISTENT"])
        assert result["Q1"].to_list() == [1.0, None]

    def test_column_no_spec_silent(self):
        """Column in columns= but no spec in metadata is silently skipped."""
        df = pl.DataFrame({"Q1": [1.0, 2.0], "Q2": [3.0, 4.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [99]}})
        result = am.apply_missing(df, meta, columns=["Q2"])
        # Q2 has no spec, so nothing changes
        assert_frame_equal(result, df)

    def test_exclude_skips_column(self):
        df = pl.DataFrame({"Q1": [1.0, 99.0], "Q2": [98.0, 2.0]})
        meta = _meta_with_missing({
            "Q1": {"type": "discrete", "values": [99]},
            "Q2": {"type": "discrete", "values": [98]},
        })
        result = am.apply_missing(df, meta, exclude=["Q2"])
        assert result["Q1"].to_list() == [1.0, None]
        assert result["Q2"].to_list() == [98.0, 2.0]  # NOT nullified

    def test_exclude_all_columns(self):
        df = pl.DataFrame({"Q1": [1.0, 99.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [99]}})
        result = am.apply_missing(df, meta, exclude=["Q1"])
        assert_frame_equal(result, df)

    def test_columns_and_exclude_mutually_exclusive(self):
        df = pl.DataFrame({"Q1": [1.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [99]}})
        with pytest.raises(ValueError, match="mutually exclusive"):
            am.apply_missing(df, meta, columns=["Q1"], exclude=["Q2"])


# ---------------------------------------------------------------------------
# Type handling
# ---------------------------------------------------------------------------

class TestTypeHandling:
    def test_dataframe_returns_dataframe(self):
        df = pl.DataFrame({"Q1": [1.0, 99.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [99]}})
        result = am.apply_missing(df, meta)
        assert isinstance(result, pl.DataFrame)

    def test_lazyframe_returns_lazyframe(self):
        lf = pl.DataFrame({"Q1": [1.0, 99.0]}).lazy()
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [99]}})
        result = am.apply_missing(lf, meta)
        assert isinstance(result, pl.LazyFrame)
        collected = result.collect()
        assert collected["Q1"].to_list() == [1.0, None]

    def test_existing_nulls_preserved(self):
        df = pl.DataFrame({"Q1": [1.0, None, 99.0, None, 3.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [99]}})
        result = am.apply_missing(df, meta)
        expected = pl.DataFrame({"Q1": [1.0, None, None, None, 3.0]})
        assert_frame_equal(result, expected)

    def test_non_matching_values_preserved(self):
        df = pl.DataFrame({"Q1": [1.0, 2.0, 3.0, 4.0, 5.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [99]}})
        result = am.apply_missing(df, meta)
        assert_frame_equal(result, df)

    def test_invalid_df_type_raises(self):
        with pytest.raises(TypeError):
            am.apply_missing([1, 2, 3], _meta_with_missing({}))


# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------

class TestEdgeCases:
    def test_no_missing_specs(self):
        df = pl.DataFrame({"Q1": [1.0, 2.0, 3.0]})
        meta = am.SpssMetadata()
        result = am.apply_missing(df, meta)
        assert_frame_equal(result, df)

    def test_empty_dataframe(self):
        df = pl.DataFrame({"Q1": pl.Series([], dtype=pl.Float64)})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [99]}})
        result = am.apply_missing(df, meta)
        assert result.height == 0

    def test_all_values_missing(self):
        df = pl.DataFrame({"Q1": [98.0, 99.0, 99.0]})
        meta = _meta_with_missing({"Q1": {"type": "discrete", "values": [98, 99]}})
        result = am.apply_missing(df, meta)
        assert result["Q1"].null_count() == 3