spb/metrika Public

1//2//  DTAReader.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation1112/// Native reader for Stata .dta formats 117 (Stata 13), 118 (Stata 14–18)13/// and 119 (>32k variables). Numeric storage types widen to Float64 with14/// Stata missing codes (`.`, `.a`–`.z`) mapped to the validity mask; str#15/// and strL variables load as strings. Value labels are not yet applied.16enum DTAReader {1718    static func read(contentsOf url: URL) throws -> ZQDataFrame {19        let data = try Data(contentsOf: url)20        return try read(data: data)21    }2223    static func read(data: Data) throws -> ZQDataFrame {24        var cursor = DTACursor(data: data)2526        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        }3536        // 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>")6162        // Map: 14 file offsets. Trusted for jumping to <strls> later; the63        // 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>")6869        // 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>")7677        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>")8384        try cursor.expect("<sortlist>")85        _ = try cursor.readBytes(layout.sortEntryBytes * (variableCount + 1))86        try cursor.expect("</sortlist>")8788        try cursor.expect("<formats>")89        _ = try cursor.readBytes(layout.formatBytes * variableCount)90        try cursor.expect("</formats>")9192        try cursor.expect("<value_label_names>")93        _ = try cursor.readBytes(layout.labelNameBytes * variableCount)94        try cursor.expect("</value_label_names>")9596        try cursor.expect("<variable_labels>")97        _ = try cursor.readBytes(layout.variableLabelBytes * variableCount)98        try cursor.expect("</variable_labels>")99100        // 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>")111112        // strLs are stored after <data>; when any variable is strL, decode113        // 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: release119            )120        }121122        // Data.123        try cursor.expect("<data>")124        var numericValues = [[Double]](125            repeating: [Double](repeating: 0, count: observationCount),126            count: variableCount127        )128        var numericMissing = [[Bool]](129            repeating: [Bool](repeating: false, count: observationCount),130            count: variableCount131        )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        }136137        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] = true145                        numericValues[j][row] = .nan146                    } else {147                        numericValues[j][row] = value148                    }149                case DTAFormat.typeFloat:150                    let value = try cursor.readFloat()151                    if value.isNaN || value > DTAFormat.floatMissingThreshold {152                        numericMissing[j][row] = true153                        numericValues[j][row] = .nan154                    } 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] = true161                        numericValues[j][row] = .nan162                    } 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] = true169                        numericValues[j][row] = .nan170                    } 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] = true177                        numericValues[j][row] = .nan178                    } 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] = nil186                    } 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 : text192                default:193                    throw ZQDataError("corrupt .dta: unknown variable type \(type)")194                }195            }196        }197        try cursor.expect("</data>")198199        // 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    }213214    // MARK: - Helpers215216    private static func isStringType(_ type: UInt16) -> Bool {217        type == DTAFormat.typeStrL218            || (1...UInt16(DTAFormat.maxFixedStringLength)).contains(type)219    }220221    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) | o224    }225226    /// 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: Int230    ) throws -> [UInt64: String] {231        var cursor = DTACursor(data: data)232        cursor.offset = offset233        cursor.bigEndian = bigEndian234        try cursor.expect("<strls>")235236        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-8244            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)] = text253        }254        try cursor.expect("</strls>")255        return table256    }257}258259extension 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}268