spb/metrika Public

1//2//  DataFrame.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910/// Column payload. Numeric data is Float64 end-to-end (CLAUDE.md §5: all11/// CPU compute in Float64) with an explicit validity mask mirroring Arrow12/// semantics: `missing[i] == true` means observation `i` is missing, and13/// the corresponding value slot is undefined.14public enum ZQColumnData: Equatable, Sendable {15    case float64(values: [Double], missing: [Bool])16    case string([String?])1718    /// Value slots at missing positions are undefined (often NaN), so19    /// equality compares masks and only the valid values — synthesized20    /// 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 true30        case let (.string(left), .string(right)):31            return left == right32        default:33            return false34        }35    }3637    public var count: Int {38        switch self {39        case .float64(let values, _): return values.count40        case .string(let values): return values.count41        }42    }4344    public var missingCount: Int {45        switch self {46        case .float64(_, let missing): return missing.count { $0 }47        case .string(let values): return values.count { $0 == nil }48        }49    }5051    public var isNumeric: Bool {52        if case .float64 = self { return true }53        return false54    }55}5657/// A named column of a data frame.58public struct ZQColumn: Equatable, Sendable {59    public var name: String60    public var data: ZQColumnData6162    public init(name: String, data: ZQColumnData) {63        self.name = name64        self.data = data65    }6667    public var count: Int { data.count }6869    /// Numeric view: values with their validity mask, or nil for string70    /// columns.71    public var numeric: (values: [Double], missing: [Bool])? {72        if case .float64(let values, let missing) = data {73            return (values, missing)74        }75        return nil76    }77}7879public struct ZQDataError: Error, Equatable, Sendable, CustomStringConvertible {80    public var message: String81    public init(_ message: String) { self.message = message }82    public var description: String { message }83}8485/// In-memory data frame: the working dataset of a session. Materialized86/// from DuckDB query results; columnar, immutable value semantics.87public struct ZQDataFrame: Equatable, Sendable {88    public private(set) var columns: [ZQColumn]89    public let rowCount: Int9091    public init(columns: [ZQColumn] = []) throws(ZQDataError) {92        let counts = Set(columns.map(\.count))93        guard counts.count <= 1 else {94            throw ZQDataError("columns have inconsistent lengths: \(counts.sorted())")95        }96        self.columns = columns97        self.rowCount = counts.first ?? 098    }99100    public var isEmpty: Bool { columns.isEmpty || rowCount == 0 }101    public var columnNames: [String] { columns.map(\.name) }102103    public func column(named name: String) -> ZQColumn? {104        columns.first { $0.name == name }105    }106107    public func requireColumn(_ name: String) throws(ZQDataError) -> ZQColumn {108        guard let column = column(named: name) else {109            throw ZQDataError("variable '\(name)' not found")110        }111        return column112    }113114    public func requireNumeric(115        _ name: String116    ) throws(ZQDataError) -> (values: [Double], missing: [Bool]) {117        guard let numeric = try requireColumn(name).numeric else {118            throw ZQDataError("variable '\(name)' is a string variable")119        }120        return numeric121    }122123    /// Adds a new column; fails if the name already exists (Stata `generate`124    /// semantics — use `replaceColumn` for `replace`).125    public mutating func addColumn(_ column: ZQColumn) throws(ZQDataError) {126        guard self.column(named: column.name) == nil else {127            throw ZQDataError("variable '\(column.name)' already defined")128        }129        guard columns.isEmpty || column.count == rowCount else {130            throw ZQDataError(131                "column '\(column.name)' has \(column.count) rows, expected \(rowCount)"132            )133        }134        columns.append(column)135    }136137    /// Replaces an existing column in place (Stata `replace`).138    public mutating func replaceColumn(_ column: ZQColumn) throws(ZQDataError) {139        guard let index = columns.firstIndex(where: { $0.name == column.name }) else {140            throw ZQDataError("variable '\(column.name)' not found")141        }142        guard column.count == rowCount else {143            throw ZQDataError(144                "column '\(column.name)' has \(column.count) rows, expected \(rowCount)"145            )146        }147        columns[index] = column148    }149150    public mutating func dropColumns(_ names: [String]) throws(ZQDataError) {151        for name in names {152            guard columns.contains(where: { $0.name == name }) else {153                throw ZQDataError("variable '\(name)' not found")154            }155        }156        columns.removeAll { names.contains($0.name) }157    }158159    /// Row subset by keep-mask (`mask[i] == true` keeps observation `i`).160    public func filtered(by mask: [Bool]) throws(ZQDataError) -> ZQDataFrame {161        guard mask.count == rowCount else {162            throw ZQDataError("filter mask length \(mask.count) != \(rowCount) rows")163        }164        let kept = (0..<rowCount).filter { mask[$0] }165        let newColumns = columns.map { column -> ZQColumn in166            switch column.data {167            case .float64(let values, let missing):168                return ZQColumn(169                    name: column.name,170                    data: .float64(171                        values: kept.map { values[$0] },172                        missing: kept.map { missing[$0] }173                    )174                )175            case .string(let values):176                return ZQColumn(name: column.name, data: .string(kept.map { values[$0] }))177            }178        }179        // Lengths are uniform by construction.180        return try! ZQDataFrame(columns: newColumns)181    }182}183