SPB Git

spb/metrika Public

Stata-class statistics, GPU-accelerated by Apple Silicon. Native Swift — no Electron, no Python runtime, no compromises.

macosswiftstatisticsgpu
Swift 92.4% HTML 3.3% R 3% Shell 1.3%

feat(data): native Stata .dta reader (117-119) and writer (118)

- DTAReader: formats 117/118/119, both byte orders, all numeric storage
  types widened to Float64 with Stata missing codes (., .a-.z) mapped to
  the validity mask; str# and strL (GSO) load as strings; value labels
  read past but not yet applied; pre-117 files rejected with a clear
  message
- DTAWriter: format 118 (UTF-8, LSF), doubles + str# up to 2045 bytes,
  real map offsets so Stata/haven/pandas can seek
- ZQDataStore routes .dta through the native path for use and save
- ZQColumnData equality now ignores value slots at missing positions
  (undefined by contract, often NaN — synthesized == failed on NaN != NaN)
- Fixtures: haven-written 117 + 118 files; tests assert bit-identical
  numeric loads vs the CSV, lossless write/read roundtrip; haven
  cross-reads Metrika-written files (verified)
- 61 tests green

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed 6 days ago (Aug 5, 2026) parent 76c5371

Showing 9 changed files with +760 and −0

modified MetrikaKit/Sources/ZQData/DataFrame.swift +19 −0
@@ -15,6 +15,25 @@ public enum ZQColumnData: Equatable, Sendable {
15 15 case float64(values: [Double], missing: [Bool])
16 16 case string([String?])
17 17
18 + /// Value slots at missing positions are undefined (often NaN), so
19 + /// equality compares masks and only the valid values — synthesized
20 + /// equality would fail on NaN ≠ NaN.
21 + public static func == (lhs: ZQColumnData, rhs: ZQColumnData) -> Bool {
22 + switch (lhs, rhs) {
23 + case let (.float64(leftValues, leftMissing), .float64(rightValues, rightMissing)):
24 + guard leftMissing == rightMissing,
25 + leftValues.count == rightValues.count else { return false }
26 + for i in 0..<leftValues.count where !leftMissing[i] {
27 + if leftValues[i] != rightValues[i] { return false }
28 + }
29 + return true
30 + case let (.string(left), .string(right)):
31 + return left == right
32 + default:
33 + return false
34 + }
35 + }
36 +
18 37 public var count: Int {
19 38 switch self {
20 39 case .float64(let values, _): return values.count
modified MetrikaKit/Sources/ZQData/DuckDBStore.swift +10 −0
@@ -32,6 +32,11 @@ public actor ZQDataStore {
32 32 guard FileManager.default.fileExists(atPath: path) else {
33 33 throw ZQDataError("file not found: \(path)")
34 34 }
35 + // Stata files use the native reader (CLAUDE.md §6).
36 + if ["dta"].contains(url.pathExtension.lowercased()) {
37 + return try DTAReader.read(contentsOf: url)
38 + }
39 +
35 40 let escaped = path.replacingOccurrences(of: "'", with: "''")
36 41 let reader: String
37 42 switch url.pathExtension.lowercased() {
@@ -106,6 +111,11 @@ public actor ZQDataStore {
106 111 /// Writes a data frame to disk via DuckDB `COPY`. Format inferred from
107 112 /// the extension (.parquet or .csv).
108 113 public func save(_ frame: ZQDataFrame, to url: URL) throws {
114 + if url.pathExtension.lowercased() == "dta" {
115 + try DTAWriter.write(frame, to: url)
116 + return
117 + }
118 +
109 119 let format: String
110 120 switch url.pathExtension.lowercased() {
111 121 case "parquet", "pq": format = "PARQUET"
added MetrikaKit/Sources/ZQData/Stata/DTAFormat.swift +157 −0
@@ -0,0 +1,157 @@
1 +//
2 +// DTAFormat.swift
3 +// Metrika
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Contact: contact@spboucher.ai
7 +// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
8 +//
9 +
10 +import Foundation
11 +
12 +/// Shared constants and byte-cursor plumbing for the native Stata .dta
13 +/// reader/writer (formats 117–119, CLAUDE.md §6).
14 +enum DTAFormat {
15 + /// Variable type codes (dta 117+).
16 + static let typeStrL: UInt16 = 32768
17 + static let typeDouble: UInt16 = 65526
18 + static let typeFloat: UInt16 = 65527
19 + static let typeLong: UInt16 = 65528
20 + static let typeInt: UInt16 = 65529
21 + static let typeByte: UInt16 = 65530
22 + static let maxFixedStringLength = 2045
23 +
24 + /// Smallest value treated as missing, per numeric storage type.
25 + /// Stata reserves the top of each range for `.` and `.a`–`.z`; the
26 + /// distinctions collapse into one missing state here.
27 + static let doubleMissingThreshold = 8.988465674311579e307
28 + static let floatMissingThreshold: Float = 1.701e38
29 + static let longMissingThreshold: Int32 = 2_147_483_621
30 + static let intMissingThreshold: Int16 = 32741
31 + static let byteMissingThreshold: Int8 = 101
32 +
33 + /// Canonical `.` (system missing) for doubles: +2¹⁰²³.
34 + static let doubleMissingValue = Double(bitPattern: 0x7FE0_0000_0000_0000)
35 +
36 + /// Per-variable field widths that differ across format versions.
37 + struct Layout {
38 + let release: Int
39 + let variableCountBytes: Int // K
40 + let observationCountBytes: Int// N
41 + let nameBytes: Int
42 + let formatBytes: Int
43 + let labelNameBytes: Int
44 + let variableLabelBytes: Int
45 + let sortEntryBytes: Int
46 + let datasetLabelLengthBytes: Int
47 + let strLVarBytes: Int // v in a data-section strL cell
48 + let strLObsBytes: Int // o in a data-section strL cell
49 + let gsoObsBytes: Int // o in a GSO record
50 +
51 + init?(release: Int) {
52 + self.release = release
53 + switch release {
54 + case 117:
55 + variableCountBytes = 2
56 + observationCountBytes = 4
57 + nameBytes = 33
58 + formatBytes = 49
59 + labelNameBytes = 33
60 + variableLabelBytes = 81
61 + sortEntryBytes = 2
62 + datasetLabelLengthBytes = 1
63 + strLVarBytes = 4
64 + strLObsBytes = 4
65 + gsoObsBytes = 4
66 + case 118:
67 + variableCountBytes = 2
68 + observationCountBytes = 8
69 + nameBytes = 129
70 + formatBytes = 57
71 + labelNameBytes = 129
72 + variableLabelBytes = 321
73 + sortEntryBytes = 2
74 + datasetLabelLengthBytes = 2
75 + strLVarBytes = 2
76 + strLObsBytes = 6
77 + gsoObsBytes = 8
78 + case 119:
79 + variableCountBytes = 4
80 + observationCountBytes = 8
81 + nameBytes = 129
82 + formatBytes = 57
83 + labelNameBytes = 129
84 + variableLabelBytes = 321
85 + sortEntryBytes = 4
86 + datasetLabelLengthBytes = 2
87 + strLVarBytes = 4
88 + strLObsBytes = 6
89 + gsoObsBytes = 8
90 + default:
91 + return nil
92 + }
93 + }
94 + }
95 +}
96 +
97 +/// Sequential binary cursor with endian-aware integer reads.
98 +struct DTACursor {
99 + let data: Data
100 + var offset: Int = 0
101 + var bigEndian = false
102 +
103 + init(data: Data) { self.data = data }
104 +
105 + var remaining: Int { data.count - offset }
106 +
107 + mutating func readBytes(_ count: Int) throws(ZQDataError) -> Data {
108 + guard count >= 0, offset + count <= data.count else {
109 + throw ZQDataError("corrupt .dta: unexpected end of file at offset \(offset)")
110 + }
111 + let slice = data.subdata(in: offset..<(offset + count))
112 + offset += count
113 + return slice
114 + }
115 +
116 + /// Consumes an exact ASCII tag or fails with its name.
117 + mutating func expect(_ tag: String) throws(ZQDataError) {
118 + let expected = Data(tag.utf8)
119 + let actual = try readBytes(expected.count)
120 + guard actual == expected else {
121 + throw ZQDataError(
122 + "corrupt .dta: expected '\(tag)' at offset \(offset - expected.count)"
123 + )
124 + }
125 + }
126 +
127 + mutating func readUInt(_ byteCount: Int) throws(ZQDataError) -> UInt64 {
128 + let bytes = try readBytes(byteCount)
129 + var value: UInt64 = 0
130 + if bigEndian {
131 + for byte in bytes { value = value << 8 | UInt64(byte) }
132 + } else {
133 + for byte in bytes.reversed() { value = value << 8 | UInt64(byte) }
134 + }
135 + return value
136 + }
137 +
138 + mutating func readDouble() throws(ZQDataError) -> Double {
139 + Double(bitPattern: try readUInt(8))
140 + }
141 +
142 + mutating func readFloat() throws(ZQDataError) -> Float {
143 + Float(bitPattern: UInt32(truncatingIfNeeded: try readUInt(4)))
144 + }
145 +
146 + /// Fixed-width, null-padded text field.
147 + mutating func readPaddedString(
148 + _ width: Int, release: Int
149 + ) throws(ZQDataError) -> String {
150 + let bytes = try readBytes(width)
151 + let trimmed = bytes.prefix { $0 != 0 }
152 + if release >= 118 {
153 + return String(data: trimmed, encoding: .utf8) ?? ""
154 + }
155 + return String(data: trimmed, encoding: .isoLatin1) ?? ""
156 + }
157 +}
added MetrikaKit/Sources/ZQData/Stata/DTAReader.swift +267 −0
@@ -0,0 +1,267 @@
1 +//
2 +// DTAReader.swift
3 +// Metrika
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Contact: contact@spboucher.ai
7 +// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
8 +//
9 +
10 +import Foundation
11 +
12 +/// Native reader for Stata .dta formats 117 (Stata 13), 118 (Stata 14–18)
13 +/// and 119 (>32k variables). Numeric storage types widen to Float64 with
14 +/// Stata missing codes (`.`, `.a`–`.z`) mapped to the validity mask; str#
15 +/// and strL variables load as strings. Value labels are not yet applied.
16 +enum DTAReader {
17 +
18 + static func read(contentsOf url: URL) throws -> ZQDataFrame {
19 + let data = try Data(contentsOf: url)
20 + return try read(data: data)
21 + }
22 +
23 + static func read(data: Data) throws -> ZQDataFrame {
24 + var cursor = DTACursor(data: data)
25 +
26 + guard data.starts(with: Data("<stata_dta>".utf8)) else {
27 + // Pre-117 files begin with a raw version byte (0x66–0x73).
28 + if let first = data.first, (0x66...0x76).contains(first) {
29 + throw ZQDataError(
30 + ".dta formats older than 117 (Stata 13) are not supported — re-save with a modern Stata"
31 + )
32 + }
33 + throw ZQDataError("not a Stata .dta file")
34 + }
35 +
36 + // Header.
37 + try cursor.expect("<stata_dta><header><release>")
38 + let releaseText = try cursor.readPaddedString(3, release: 117)
39 + guard let release = Int(releaseText) else {
40 + throw ZQDataError("corrupt .dta: bad release '\(releaseText)'")
41 + }
42 + guard let layout = DTAFormat.Layout(release: release) else {
43 + throw ZQDataError(
44 + ".dta format \(release) is not supported (117–119 are)"
45 + )
46 + }
47 + try cursor.expect("</release><byteorder>")
48 + let byteOrder = try cursor.readPaddedString(3, release: 117)
49 + cursor.bigEndian = byteOrder == "MSF"
50 + try cursor.expect("</byteorder><K>")
51 + let variableCount = Int(try cursor.readUInt(layout.variableCountBytes))
52 + try cursor.expect("</K><N>")
53 + let observationCount = Int(try cursor.readUInt(layout.observationCountBytes))
54 + try cursor.expect("</N><label>")
55 + let labelLength = Int(try cursor.readUInt(layout.datasetLabelLengthBytes))
56 + _ = try cursor.readBytes(labelLength)
57 + try cursor.expect("</label><timestamp>")
58 + let timestampLength = Int(try cursor.readUInt(1))
59 + _ = try cursor.readBytes(timestampLength)
60 + try cursor.expect("</timestamp></header>")
61 +
62 + // Map: 14 file offsets. Trusted for jumping to <strls> later; the
63 + // sections up to <data> are read sequentially.
64 + try cursor.expect("<map>")
65 + var map: [UInt64] = []
66 + for _ in 0..<14 { map.append(try cursor.readUInt(8)) }
67 + try cursor.expect("</map>")
68 +
69 + // Variable descriptors.
70 + try cursor.expect("<variable_types>")
71 + var types: [UInt16] = []
72 + for _ in 0..<variableCount {
73 + types.append(UInt16(truncatingIfNeeded: try cursor.readUInt(2)))
74 + }
75 + try cursor.expect("</variable_types>")
76 +
77 + try cursor.expect("<varnames>")
78 + var names: [String] = []
79 + for _ in 0..<variableCount {
80 + names.append(try cursor.readPaddedString(layout.nameBytes, release: release))
81 + }
82 + try cursor.expect("</varnames>")
83 +
84 + try cursor.expect("<sortlist>")
85 + _ = try cursor.readBytes(layout.sortEntryBytes * (variableCount + 1))
86 + try cursor.expect("</sortlist>")
87 +
88 + try cursor.expect("<formats>")
89 + _ = try cursor.readBytes(layout.formatBytes * variableCount)
90 + try cursor.expect("</formats>")
91 +
92 + try cursor.expect("<value_label_names>")
93 + _ = try cursor.readBytes(layout.labelNameBytes * variableCount)
94 + try cursor.expect("</value_label_names>")
95 +
96 + try cursor.expect("<variable_labels>")
97 + _ = try cursor.readBytes(layout.variableLabelBytes * variableCount)
98 + try cursor.expect("</variable_labels>")
99 +
100 + // Characteristics: <ch> length-prefixed blobs until the closer.
101 + try cursor.expect("<characteristics>")
102 + while cursor.remaining >= 4 {
103 + let peek = try DTACursor(data: data).peekBytes(at: cursor.offset, count: 4)
104 + if peek == Data("</ch".utf8) { break }
105 + try cursor.expect("<ch>")
106 + let length = Int(try cursor.readUInt(4))
107 + _ = try cursor.readBytes(length)
108 + try cursor.expect("</ch>")
109 + }
110 + try cursor.expect("</characteristics>")
111 +
112 + // strLs are stored after <data>; when any variable is strL, decode
113 + // that section first (via the map) so cells can resolve.
114 + var strLTable: [UInt64: String] = [:]
115 + if types.contains(DTAFormat.typeStrL) {
116 + strLTable = try readStrLs(
117 + data: data, at: Int(map[10]), layout: layout,
118 + bigEndian: cursor.bigEndian, release: release
119 + )
120 + }
121 +
122 + // Data.
123 + try cursor.expect("<data>")
124 + var numericValues = [[Double]](
125 + repeating: [Double](repeating: 0, count: observationCount),
126 + count: variableCount
127 + )
128 + var numericMissing = [[Bool]](
129 + repeating: [Bool](repeating: false, count: observationCount),
130 + count: variableCount
131 + )
132 + var stringValues = [[String?]](repeating: [], count: variableCount)
133 + for j in 0..<variableCount where isStringType(types[j]) {
134 + stringValues[j] = [String?](repeating: nil, count: observationCount)
135 + }
136 +
137 + for row in 0..<observationCount {
138 + for j in 0..<variableCount {
139 + let type = types[j]
140 + switch type {
141 + case DTAFormat.typeDouble:
142 + let value = try cursor.readDouble()
143 + if value.isNaN || value > DTAFormat.doubleMissingThreshold {
144 + numericMissing[j][row] = true
145 + numericValues[j][row] = .nan
146 + } else {
147 + numericValues[j][row] = value
148 + }
149 + case DTAFormat.typeFloat:
150 + let value = try cursor.readFloat()
151 + if value.isNaN || value > DTAFormat.floatMissingThreshold {
152 + numericMissing[j][row] = true
153 + numericValues[j][row] = .nan
154 + } else {
155 + numericValues[j][row] = Double(value)
156 + }
157 + case DTAFormat.typeLong:
158 + let value = Int32(truncatingIfNeeded: try cursor.readUInt(4))
159 + if value >= DTAFormat.longMissingThreshold {
160 + numericMissing[j][row] = true
161 + numericValues[j][row] = .nan
162 + } else {
163 + numericValues[j][row] = Double(value)
164 + }
165 + case DTAFormat.typeInt:
166 + let value = Int16(truncatingIfNeeded: try cursor.readUInt(2))
167 + if value >= DTAFormat.intMissingThreshold {
168 + numericMissing[j][row] = true
169 + numericValues[j][row] = .nan
170 + } else {
171 + numericValues[j][row] = Double(value)
172 + }
173 + case DTAFormat.typeByte:
174 + let value = Int8(truncatingIfNeeded: try cursor.readUInt(1))
175 + if value >= DTAFormat.byteMissingThreshold {
176 + numericMissing[j][row] = true
177 + numericValues[j][row] = .nan
178 + } else {
179 + numericValues[j][row] = Double(value)
180 + }
181 + case DTAFormat.typeStrL:
182 + let v = try cursor.readUInt(layout.strLVarBytes)
183 + let o = try cursor.readUInt(layout.strLObsBytes)
184 + if v == 0 && o == 0 {
185 + stringValues[j][row] = nil
186 + } else {
187 + stringValues[j][row] = strLTable[strLKey(v: v, o: o)] ?? ""
188 + }
189 + case 1...UInt16(DTAFormat.maxFixedStringLength):
190 + let text = try cursor.readPaddedString(Int(type), release: release)
191 + stringValues[j][row] = text.isEmpty ? nil : text
192 + default:
193 + throw ZQDataError("corrupt .dta: unknown variable type \(type)")
194 + }
195 + }
196 + }
197 + try cursor.expect("</data>")
198 +
199 + // Assemble columns.
200 + var columns: [ZQColumn] = []
201 + for j in 0..<variableCount {
202 + if isStringType(types[j]) {
203 + columns.append(ZQColumn(name: names[j], data: .string(stringValues[j])))
204 + } else {
205 + columns.append(ZQColumn(
206 + name: names[j],
207 + data: .float64(values: numericValues[j], missing: numericMissing[j])
208 + ))
209 + }
210 + }
211 + return try ZQDataFrame(columns: columns)
212 + }
213 +
214 + // MARK: - Helpers
215 +
216 + private static func isStringType(_ type: UInt16) -> Bool {
217 + type == DTAFormat.typeStrL
218 + || (1...UInt16(DTAFormat.maxFixedStringLength)).contains(type)
219 + }
220 +
221 + private static func strLKey(v: UInt64, o: UInt64) -> UInt64 {
222 + // (v, o) packed; v < 2^16 in 118 and < 2^32 in 117/119, o < 2^48.
223 + (v << 48) | o
224 + }
225 +
226 + /// Parses the <strls> section: consecutive GSO records.
227 + private static func readStrLs(
228 + data: Data, at offset: Int, layout: DTAFormat.Layout,
229 + bigEndian: Bool, release: Int
230 + ) throws -> [UInt64: String] {
231 + var cursor = DTACursor(data: data)
232 + cursor.offset = offset
233 + cursor.bigEndian = bigEndian
234 + try cursor.expect("<strls>")
235 +
236 + var table: [UInt64: String] = [:]
237 + while cursor.remaining >= 3 {
238 + let marker = try cursor.peekBytes(at: cursor.offset, count: 3)
239 + if marker != Data("GSO".utf8) { break }
240 + _ = try cursor.readBytes(3)
241 + let v = try cursor.readUInt(4)
242 + let o = try cursor.readUInt(layout.gsoObsBytes)
243 + let kind = try cursor.readUInt(1) // 129 binary, 130 ASCII/UTF-8
244 + let length = Int(try cursor.readUInt(4))
245 + var payload = try cursor.readBytes(length)
246 + if kind == 130, payload.last == 0 {
247 + payload = payload.dropLast()
248 + }
249 + let text = String(data: payload, encoding: .utf8)
250 + ?? String(data: payload, encoding: .isoLatin1)
251 + ?? ""
252 + table[strLKey(v: v, o: o)] = text
253 + }
254 + try cursor.expect("</strls>")
255 + return table
256 + }
257 +}
258 +
259 +extension DTACursor {
260 + /// Non-consuming read at an absolute offset.
261 + func peekBytes(at position: Int, count: Int) throws(ZQDataError) -> Data {
262 + guard position >= 0, position + count <= data.count else {
263 + throw ZQDataError("corrupt .dta: unexpected end of file at offset \(position)")
264 + }
265 + return data.subdata(in: position..<(position + count))
266 + }
267 +}
added MetrikaKit/Sources/ZQData/Stata/DTAWriter.swift +182 −0
@@ -0,0 +1,182 @@
1 +//
2 +// DTAWriter.swift
3 +// Metrika
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Contact: contact@spboucher.ai
7 +// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
8 +//
9 +
10 +import Foundation
11 +
12 +/// Native writer for Stata .dta format 118 (Stata 14+, UTF-8, little
13 +/// endian). Numeric columns store as double; string columns as str# up to
14 +/// 2045 bytes. The map section carries real offsets, so Stata, haven, and
15 +/// pandas can seek.
16 +enum DTAWriter {
17 +
18 + static func write(_ frame: ZQDataFrame, to url: URL) throws {
19 + try data(for: frame).write(to: url)
20 + }
21 +
22 + static func data(for frame: ZQDataFrame) throws -> Data {
23 + guard !frame.columns.isEmpty else {
24 + throw ZQDataError("nothing to save: dataset has no variables")
25 + }
26 + let variableCount = frame.columns.count
27 + guard variableCount <= 32767 else {
28 + throw ZQDataError(".dta 118 supports at most 32,767 variables")
29 + }
30 +
31 + // str# width per string column (UTF-8 bytes, Stata counts bytes).
32 + var stringWidths = [Int](repeating: 0, count: variableCount)
33 + for (j, column) in frame.columns.enumerated() {
34 + if case .string(let values) = column.data {
35 + var width = 1
36 + for value in values {
37 + if let value {
38 + let bytes = value.utf8.count
39 + guard bytes <= DTAFormat.maxFixedStringLength else {
40 + throw ZQDataError(
41 + "string variable '\(column.name)' exceeds 2045 bytes — strL writing not supported yet"
42 + )
43 + }
44 + width = max(width, bytes)
45 + }
46 + }
47 + stringWidths[j] = width
48 + }
49 + }
50 +
51 + var out = Data()
52 + var mapOffsets = [UInt64](repeating: 0, count: 14)
53 +
54 + func append(_ text: String) { out.append(Data(text.utf8)) }
55 + func appendUInt(_ value: UInt64, _ byteCount: Int) {
56 + for shift in 0..<byteCount {
57 + out.append(UInt8(truncatingIfNeeded: value >> (8 * shift)))
58 + }
59 + }
60 + /// Null-terminated field (names, formats): content ≤ width−1 bytes.
61 + func appendPadded(_ text: String, _ width: Int) {
62 + var bytes = Data(text.utf8).prefix(width - 1)
63 + bytes.append(contentsOf: [UInt8](repeating: 0, count: width - bytes.count))
64 + out.append(bytes)
65 + }
66 +
67 + /// str# data cell: exactly width bytes, null-padded, no terminator
68 + /// required when the content fills the field.
69 + func appendCell(_ text: String, _ width: Int) {
70 + var bytes = Data(text.utf8).prefix(width)
71 + bytes.append(contentsOf: [UInt8](repeating: 0, count: width - bytes.count))
72 + out.append(bytes)
73 + }
74 +
75 + // Header.
76 + mapOffsets[0] = 0
77 + append("<stata_dta><header><release>118</release>")
78 + append("<byteorder>LSF</byteorder><K>")
79 + appendUInt(UInt64(variableCount), 2)
80 + append("</K><N>")
81 + appendUInt(UInt64(frame.rowCount), 8)
82 + append("</N><label>")
83 + appendUInt(0, 2)
84 + append("</label><timestamp>")
85 + appendUInt(0, 1)
86 + append("</timestamp></header>")
87 +
88 + // Map: placeholder now, patched at the end.
89 + mapOffsets[1] = UInt64(out.count)
90 + append("<map>")
91 + let mapPayloadOffset = out.count
92 + out.append(Data(repeating: 0, count: 14 * 8))
93 + append("</map>")
94 +
95 + // Variable types.
96 + mapOffsets[2] = UInt64(out.count)
97 + append("<variable_types>")
98 + for (j, column) in frame.columns.enumerated() {
99 + let code: UInt16 = column.data.isNumeric
100 + ? DTAFormat.typeDouble
101 + : UInt16(stringWidths[j])
102 + appendUInt(UInt64(code), 2)
103 + }
104 + append("</variable_types>")
105 +
106 + // Names (129-byte fields).
107 + mapOffsets[3] = UInt64(out.count)
108 + append("<varnames>")
109 + for column in frame.columns {
110 + guard column.name.utf8.count <= 128 else {
111 + throw ZQDataError("variable name '\(column.name)' exceeds 128 bytes")
112 + }
113 + appendPadded(column.name, 129)
114 + }
115 + append("</varnames>")
116 +
117 + // Sort list: unsorted.
118 + mapOffsets[4] = UInt64(out.count)
119 + append("<sortlist>")
120 + out.append(Data(repeating: 0, count: 2 * (variableCount + 1)))
121 + append("</sortlist>")
122 +
123 + // Display formats (57-byte fields).
124 + mapOffsets[5] = UInt64(out.count)
125 + append("<formats>")
126 + for (j, column) in frame.columns.enumerated() {
127 + let format = column.data.isNumeric ? "%10.0g" : "%\(stringWidths[j])s"
128 + appendPadded(format, 57)
129 + }
130 + append("</formats>")
131 +
132 + // No value labels or variable labels.
133 + mapOffsets[6] = UInt64(out.count)
134 + append("<value_label_names>")
135 + out.append(Data(repeating: 0, count: 129 * variableCount))
136 + append("</value_label_names>")
137 +
138 + mapOffsets[7] = UInt64(out.count)
139 + append("<variable_labels>")
140 + out.append(Data(repeating: 0, count: 321 * variableCount))
141 + append("</variable_labels>")
142 +
143 + mapOffsets[8] = UInt64(out.count)
144 + append("<characteristics></characteristics>")
145 +
146 + // Data, row-major.
147 + mapOffsets[9] = UInt64(out.count)
148 + append("<data>")
149 + for row in 0..<frame.rowCount {
150 + for (j, column) in frame.columns.enumerated() {
151 + switch column.data {
152 + case .float64(let values, let missing):
153 + let value = missing[row] ? DTAFormat.doubleMissingValue : values[row]
154 + appendUInt(value.bitPattern, 8)
155 + case .string(let values):
156 + // Stata's string missing is the empty string.
157 + appendCell(values[row] ?? "", stringWidths[j])
158 + }
159 + }
160 + }
161 + append("</data>")
162 +
163 + mapOffsets[10] = UInt64(out.count)
164 + append("<strls></strls>")
165 +
166 + mapOffsets[11] = UInt64(out.count)
167 + append("<value_labels></value_labels>")
168 +
169 + mapOffsets[12] = UInt64(out.count)
170 + append("</stata_dta>")
171 + mapOffsets[13] = UInt64(out.count)
172 +
173 + // Patch the map.
174 + for (index, offset) in mapOffsets.enumerated() {
175 + for shift in 0..<8 {
176 + out[mapPayloadOffset + index * 8 + shift] =
177 + UInt8(truncatingIfNeeded: offset >> (8 * shift))
178 + }
179 + }
180 + return out
181 + }
182 +}
added MetrikaKit/Tests/MetrikaKitTests/DTATests.swift +111 −0
@@ -0,0 +1,111 @@
1 +//
2 +// DTATests.swift
3 +// Metrika
4 +//
5 +// Author: Simon-Pierre Boucher
6 +// Contact: contact@spboucher.ai
7 +// Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
8 +//
9 +
10 +import Foundation
11 +import Testing
12 +import ZQData
13 +
14 +/// Native .dta reader/writer tests: haven-written fixtures (formats 117
15 +/// and 118) must load identically to the CSV, and a write/read roundtrip
16 +/// must be lossless. All loads go through the public ZQDataStore path.
17 +@Suite("Stata .dta reader/writer")
18 +struct DTATests {
19 + let store: ZQDataStore
20 + let csvFrame: ZQDataFrame
21 + let fixturesDirectory: URL
22 +
23 + init() async throws {
24 + let fixtures = try Fixtures()
25 + self.store = try ZQDataStore()
26 + self.csvFrame = try await store.load(contentsOf: fixtures.datasetURL)
27 + self.fixturesDirectory = fixtures.datasetURL.deletingLastPathComponent()
28 + }
29 +
30 + private func expectMatchesCSV(_ frame: ZQDataFrame) throws {
31 + #expect(frame.rowCount == 60)
32 + #expect(frame.columnNames == [
33 + "revenue", "price", "region", "firm_id", "purchase", "orders", "firm_name",
34 + ])
35 +
36 + // Numeric columns: bit-identical to the CSV load (both sides parse
37 + // the same decimal text into the nearest double).
38 + for name in csvFrame.columnNames {
39 + let expected = try csvFrame.requireNumeric(name)
40 + let actual = try frame.requireNumeric(name)
41 + #expect(actual.missing == expected.missing, "missing mask of \(name)")
42 + for i in 0..<60 where !expected.missing[i] {
43 + #expect(actual.values[i] == expected.values[i], "\(name)[\(i)]")
44 + }
45 + }
46 +
47 + // String column: haven wrote firm_<id>, with row 5 blank (missing).
48 + guard case .string(let names)? = frame.column(named: "firm_name")?.data else {
49 + Issue.record("firm_name is not a string column")
50 + return
51 + }
52 + #expect(names[0] == "firm_1")
53 + #expect(names[4] == nil) // "" is Stata's string missing
54 + #expect(names[59] == "firm_12")
55 + }
56 +
57 + @Test("reads haven-written format 118")
58 + func readsFormat118() async throws {
59 + let url = fixturesDirectory.appendingPathComponent("regression_v118.dta")
60 + try expectMatchesCSV(try await store.load(contentsOf: url))
61 + }
62 +
63 + @Test("reads haven-written format 117")
64 + func readsFormat117() async throws {
65 + let url = fixturesDirectory.appendingPathComponent("regression_v117.dta")
66 + try expectMatchesCSV(try await store.load(contentsOf: url))
67 + }
68 +
69 + @Test("write/read roundtrip is lossless")
70 + func roundtrip() async throws {
71 + let url = fixturesDirectory.appendingPathComponent("regression_v118.dta")
72 + let original = try await store.load(contentsOf: url)
73 +
74 + let out = FileManager.default.temporaryDirectory
75 + .appendingPathComponent("metrika_roundtrip_\(UUID().uuidString).dta")
76 + defer { try? FileManager.default.removeItem(at: out) }
77 + try await store.save(original, to: out)
78 + let reloaded = try await store.load(contentsOf: out)
79 +
80 + #expect(reloaded == original)
81 + }
82 +
83 + @Test("rejects pre-117 files with a clear message")
84 + func rejectsOldFormats() async throws {
85 + // Format 115 files start with the raw version byte 0x73.
86 + let url = try temporaryFile(Data([0x73, 0x02, 0x01, 0x00]))
87 + defer { try? FileManager.default.removeItem(at: url) }
88 + do {
89 + _ = try await store.load(contentsOf: url)
90 + Issue.record("expected a format error")
91 + } catch {
92 + #expect("\(error)".contains("older than 117"))
93 + }
94 + }
95 +
96 + @Test("rejects non-dta bytes")
97 + func rejectsGarbage() async throws {
98 + let url = try temporaryFile(Data("not a dta file at all".utf8))
99 + defer { try? FileManager.default.removeItem(at: url) }
100 + await #expect(throws: (any Error).self) {
101 + _ = try await store.load(contentsOf: url)
102 + }
103 + }
104 +
105 + private func temporaryFile(_ data: Data) throws -> URL {
106 + let url = FileManager.default.temporaryDirectory
107 + .appendingPathComponent("metrika_dta_\(UUID().uuidString).dta")
108 + try data.write(to: url)
109 + return url
110 + }
111 +}
added MetrikaKit/Tests/MetrikaKitTests/Fixtures/regression_v117.dta +0 −0

Binary file not shown.

added MetrikaKit/Tests/MetrikaKitTests/Fixtures/regression_v118.dta +0 −0

Binary file not shown.

modified Tests/Fixtures/generate.R +14 −0
@@ -233,6 +233,20 @@ emit("dist_qt_0p975_df12", qt(0.975, 12))
233 233 emit("dist_qt_0p995_df4", qt(0.995, 4))
234 234 emit("dist_pnorm_1p64", pnorm(1.64))
235 235
236 +# ------------------------------------------------------------- .dta files
237 +# haven-written fixtures for the native DTA reader: same dataset plus a
238 +# string column, in formats 118 (version 14) and 117 (version 13).
239 +if (requireNamespace("haven", quietly = TRUE)) {
240 + dta <- data
241 + dta$firm_name <- paste0("firm_", data$firm_id)
242 + dta$firm_name[5] <- "" # string missing
243 + haven::write_dta(dta, file.path(out_dir, "regression_v118.dta"), version = 14)
244 + haven::write_dta(dta, file.path(out_dir, "regression_v117.dta"), version = 13)
245 + cat("wrote .dta fixtures (117, 118)\n")
246 +} else {
247 + cat("haven not installed — skipping .dta fixtures\n")
248 +}
249 +
236 250 writeLines(lines, sink_path)
237 251 cat("wrote", csv_path, "\n")
238 252 cat("wrote", sink_path, "with", length(lines), "expected values\n")
239 253