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

"""Write benchmark: ambers.write_sav vs pyreadstat.write_sav.

Compares write speed across test files for both bytecode (.sav) and
zlib (.zsav) compression modes.

Usage: python tests/bench_write.py
"""

import gc
import os
import tempfile
import time

import ambers
import pyreadstat

RUNS = 5
TEST_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "test_data")

FILES = [
    ("test_1", "test_1_small.sav"),
    ("test_2", "test_2_medium.sav"),
    ("test_3", "test_3_medium.sav"),
    ("test_4", "test_4_small.sav"),
    ("test_5", "test_5_small.sav"),
    ("test_6", "test_6_large.sav"),
]


def fmt_size(path):
    size_mb = os.path.getsize(path) / 1024 / 1024
    return f"{size_mb / 1024:.1f} GB" if size_mb >= 1024 else f"{size_mb:.0f} MB"


def fmt_ratio(a_best, p_best):
    if a_best <= p_best:
        return f"\033[32m{p_best / a_best:.1f}x faster\033[0m"
    else:
        return f"\033[31m{a_best / p_best:.1f}x SLOWER\033[0m"


def prepare_pyreadstat_args(meta):
    """Extract metadata fields into pyreadstat.write_sav kwargs."""
    kwargs = {}

    if meta.file_label:
        kwargs["file_label"] = meta.file_label

    if meta.variable_labels:
        kwargs["column_labels"] = meta.variable_labels

    if meta.variable_value_labels:
        kwargs["variable_value_labels"] = meta.variable_value_labels

    if meta.variable_measures:
        kwargs["variable_measure"] = meta.variable_measures

    if meta.notes:
        kwargs["note"] = meta.notes[0] if len(meta.notes) == 1 else meta.notes

    return kwargs


def bench_sav(label, path):
    """Benchmark bytecode .sav writing."""
    size_str = fmt_size(path)

    # Read with ambers
    sav = ambers.read_sav(path)
    df_polars, meta = sav.data, sav.meta
    rows, cols = df_polars.height, df_polars.width

    # Pre-convert to pandas for pyreadstat (exclude conversion from timing)
    df_pandas = df_polars.to_pandas()
    pyr_kwargs = prepare_pyreadstat_args(meta)

    with tempfile.TemporaryDirectory() as tmp:
        am_path = os.path.join(tmp, "ambers.sav")
        pyr_path = os.path.join(tmp, "pyreadstat.sav")

        # Warmup
        ambers.write_sav(df_polars, am_path, meta=meta)
        pyreadstat.write_sav(df_pandas, pyr_path, row_compress=True, **pyr_kwargs)

        # Bench ambers
        a_times = []
        for _ in range(RUNS):
            gc.collect()
            t0 = time.perf_counter()
            ambers.write_sav(df_polars, am_path, meta=meta)
            a_times.append(time.perf_counter() - t0)

        # Bench pyreadstat
        p_times = []
        for _ in range(RUNS):
            gc.collect()
            t0 = time.perf_counter()
            pyreadstat.write_sav(df_pandas, pyr_path, row_compress=True, **pyr_kwargs)
            p_times.append(time.perf_counter() - t0)

    a_best = min(a_times)
    p_best = min(p_times)
    ratio = fmt_ratio(a_best, p_best)

    # Output sizes
    am_sz = os.path.getsize(am_path) if os.path.exists(am_path) else 0
    pyr_sz = os.path.getsize(pyr_path) if os.path.exists(pyr_path) else 0

    print(
        f"  .sav   {label:<8} {size_str:>8} {rows:>8,} x {cols:<5} | "
        f"ambers {a_best:.3f}s  pyr {p_best:.3f}s  -> {ratio}"
    )
    return a_best, p_best


def bench_zsav(label, path):
    """Benchmark zlib .zsav writing."""
    size_str = fmt_size(path)

    # Read with ambers
    sav = ambers.read_sav(path)
    df_polars, meta = sav.data, sav.meta
    rows, cols = df_polars.height, df_polars.width

    # Pre-convert to pandas for pyreadstat
    df_pandas = df_polars.to_pandas()
    pyr_kwargs = prepare_pyreadstat_args(meta)

    with tempfile.TemporaryDirectory() as tmp:
        am_path = os.path.join(tmp, "ambers.zsav")
        pyr_path = os.path.join(tmp, "pyreadstat.zsav")

        # Warmup
        ambers.write_sav(df_polars, am_path, meta=meta)
        pyreadstat.write_sav(df_pandas, pyr_path, compress=True, **pyr_kwargs)

        # Bench ambers
        a_times = []
        for _ in range(RUNS):
            gc.collect()
            t0 = time.perf_counter()
            ambers.write_sav(df_polars, am_path, meta=meta)
            a_times.append(time.perf_counter() - t0)

        # Bench pyreadstat
        p_times = []
        for _ in range(RUNS):
            gc.collect()
            t0 = time.perf_counter()
            pyreadstat.write_sav(df_pandas, pyr_path, compress=True, **pyr_kwargs)
            p_times.append(time.perf_counter() - t0)

    a_best = min(a_times)
    p_best = min(p_times)
    ratio = fmt_ratio(a_best, p_best)

    print(
        f"  .zsav  {label:<8} {size_str:>8} {rows:>8,} x {cols:<5} | "
        f"ambers {a_best:.3f}s  pyr {p_best:.3f}s  -> {ratio}"
    )
    return a_best, p_best


def main():
    print(f"Write Benchmark: ambers vs pyreadstat ({RUNS} runs, best time)")
    print("=" * 80)

    print(f"\n{'  Mode':<8} {'File':<8} {'Size':>8} {'Shape':>16} | {'Results'}")
    print("-" * 80)

    for label, filename in FILES:
        path = os.path.join(TEST_DIR, filename)
        if not os.path.exists(path):
            print(f"  {label}: SKIPPED (not found)")
            continue

        bench_sav(label, path)
        bench_zsav(label, path)
        print()

    print("Done.")


if __name__ == "__main__":
    main()