spb/metrika Public

1//2//  DTAWriter.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation1112/// Native writer for Stata .dta format 118 (Stata 14+, UTF-8, little13/// endian). Numeric columns store as double; string columns as str# up to14/// 2045 bytes. The map section carries real offsets, so Stata, haven, and15/// pandas can seek.16enum DTAWriter {1718    static func write(_ frame: ZQDataFrame, to url: URL) throws {19        try data(for: frame).write(to: url)20    }2122    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.count27        guard variableCount <= 32767 else {28            throw ZQDataError(".dta 118 supports at most 32,767 variables")29        }3031        // 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 = 136                for value in values {37                    if let value {38                        let bytes = value.utf8.count39                        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] = width48            }49        }5051        var out = Data()52        var mapOffsets = [UInt64](repeating: 0, count: 14)5354        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        }6667        /// str# data cell: exactly width bytes, null-padded, no terminator68        /// 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        }7475        // Header.76        mapOffsets[0] = 077        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>")8788        // Map: placeholder now, patched at the end.89        mapOffsets[1] = UInt64(out.count)90        append("<map>")91        let mapPayloadOffset = out.count92        out.append(Data(repeating: 0, count: 14 * 8))93        append("</map>")9495        // 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.isNumeric100                ? DTAFormat.typeDouble101                : UInt16(stringWidths[j])102            appendUInt(UInt64(code), 2)103        }104        append("</variable_types>")105106        // 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>")116117        // 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>")122123        // 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>")131132        // 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>")137138        mapOffsets[7] = UInt64(out.count)139        append("<variable_labels>")140        out.append(Data(repeating: 0, count: 321 * variableCount))141        append("</variable_labels>")142143        mapOffsets[8] = UInt64(out.count)144        append("<characteristics></characteristics>")145146        // 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>")162163        mapOffsets[10] = UInt64(out.count)164        append("<strls></strls>")165166        mapOffsets[11] = UInt64(out.count)167        append("<value_labels></value_labels>")168169        mapOffsets[12] = UInt64(out.count)170        append("</stata_dta>")171        mapOffsets[13] = UInt64(out.count)172173        // 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 out181    }182}183