ambers 0.4.0

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

"""Stress test: Very wide files (3 million columns).

Verifies ambers can read .sav files far beyond ReadStat's 16 MB allocation
limit (pyreadstat issue #79). Tests 3,000,000 Float64 columns × 1 row.

Usage: .venv/Scripts/python tests/test_wide_columns.py
"""

import os
import sys
import time
import tempfile

import pytest
import polars as pl
pyreadstat = pytest.importorskip("pyreadstat")
import ambers

PASSED = 0
FAILED = 0


def check(name, condition, detail=""):
    global PASSED, FAILED
    if condition:
        PASSED += 1
        print(f"  PASS: {name}")
    else:
        FAILED += 1
        print(f"  FAIL: {name} — {detail}")


def _test_wide_columns(n_cols):
    """Test reading a file with n_cols Float64 columns."""
    print(f"\n--- {n_cols:,} columns × 1 row ---")

    # Generate column names and data
    print(f"  Generating {n_cols:,} columns...")
    t0 = time.perf_counter()
    col_names = [f"V{i}" for i in range(1, n_cols + 1)]
    # Use a dict of lists for Polars DataFrame
    data = {name: [float(i)] for i, name in enumerate(col_names, 1)}
    df_write = pl.DataFrame(data)
    t_gen = time.perf_counter() - t0
    print(f"  DataFrame generated in {t_gen:.1f}s ({df_write.width:,} cols)")

    with tempfile.NamedTemporaryFile(suffix=".sav", delete=False) as f:
        path = f.name

    try:
        # Write with pyreadstat
        print(f"  Writing .sav file...")
        t0 = time.perf_counter()
        pyreadstat.write_sav(df_write, path)
        t_write = time.perf_counter() - t0
        file_size = os.path.getsize(path)
        print(f"  Written in {t_write:.1f}s ({file_size / 1024 / 1024:.1f} MB)")

        # Read with ambers
        print(f"  Reading with ambers...")
        t0 = time.perf_counter()
        sav = ambers.read_sav(path)
        df_ambers, meta = sav.data, sav.meta
        t_read = time.perf_counter() - t0
        print(f"  Read in {t_read:.1f}s")

        # Verify column count
        check(
            f"ambers column count = {n_cols:,}",
            df_ambers.width == n_cols,
            f"got {df_ambers.width:,}",
        )

        # Verify row count
        check("ambers row count = 1", df_ambers.height == 1, f"got {df_ambers.height}")

        # Spot-check first, middle, and last columns
        check(
            "first column (V1) = 1.0",
            df_ambers["V1"][0] == 1.0,
            f"got {df_ambers['V1'][0]}",
        )

        mid = n_cols // 2
        mid_name = f"V{mid}"
        check(
            f"middle column ({mid_name}) = {float(mid)}",
            df_ambers[mid_name][0] == float(mid),
            f"got {df_ambers[mid_name][0]}",
        )

        last_name = f"V{n_cols}"
        check(
            f"last column ({last_name}) = {float(n_cols)}",
            df_ambers[last_name][0] == float(n_cols),
            f"got {df_ambers[last_name][0]}",
        )

        # Verify metadata
        check(
            "metadata column count",
            meta.number_columns == n_cols,
            f"got {meta.number_columns:,}",
        )

        print(f"\n  Summary: {t_write:.1f}s write, {t_read:.1f}s read, {file_size / 1024 / 1024:.1f} MB file")

    finally:
        os.unlink(path)


def main():
    global PASSED, FAILED

    print("=" * 60)
    print("WIDE COLUMN STRESS TEST")
    print("(pyreadstat issue #79 — ReadStat 16 MB limit)")
    print("=" * 60)

    # Warm up with a smaller test first
    _test_wide_columns(100_000)

    # The big one: 3 million columns
    _test_wide_columns(3_000_000)

    # Summary
    total = PASSED + FAILED
    print(f"\n{'=' * 60}")
    print(f"RESULTS: {PASSED}/{total} passed, {FAILED} failed")
    print(f"{'=' * 60}")

    if FAILED > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()