spb/metrika Public

1//2//  ExpressionEvaluator.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation11import ZQData12import ZQParser1314/// Vectorized evaluation of a `ZQExpression` against the working dataset.15///16/// Missing-value semantics: any arithmetic or comparison touching a17/// missing operand yields missing; undefined operations (division by18/// zero, log of a non-positive number) yield missing. The `missing(x)`19/// function itself never returns missing. In `if` qualifiers a missing20/// condition excludes the observation.21struct ExpressionEvaluator {22    struct EvaluationError: Error, CustomStringConvertible {23        let message: String24        var description: String { message }25    }2627    /// A column vector or a broadcastable scalar.28    enum Value {29        case numeric(values: [Double], missing: [Bool])30        case strings([String?])31        case numericScalar(Double)32        case stringScalar(String)33        case missingScalar34    }3536    let frame: ZQDataFrame3738    func evaluate(_ expression: ZQExpression) throws -> Value {39        switch expression {40        case .number(let value):41            return .numericScalar(value)42        case .string(let value):43            return .stringScalar(value)44        case .missing:45            return .missingScalar46        case .variable(let name):47            guard let column = frame.column(named: name) else {48                throw EvaluationError(message: "variable '\(name)' not found")49            }50            switch column.data {51            case .float64(let values, let missing):52                return .numeric(values: values, missing: missing)53            case .string(let values):54                return .strings(values)55            }56        case .unary(let op, let operand):57            return try evaluateUnary(op: op, operand: operand)58        case .binary(let op, let lhs, let rhs):59            return try evaluateBinary(op: op, lhs: lhs, rhs: rhs)60        case .call(let name, let arguments):61            return try evaluateCall(name: name, arguments: arguments)62        }63    }6465    /// Evaluates to a full-length numeric column (broadcasting scalars).66    func evaluateNumericColumn(67        _ expression: ZQExpression68    ) throws -> (values: [Double], missing: [Bool]) {69        switch try evaluate(expression) {70        case .numeric(let values, let missing):71            return (values, missing)72        case .numericScalar(let value):73            return (74                [Double](repeating: value, count: frame.rowCount),75                [Bool](repeating: false, count: frame.rowCount)76            )77        case .missingScalar:78            return (79                [Double](repeating: .nan, count: frame.rowCount),80                [Bool](repeating: true, count: frame.rowCount)81            )82        case .strings, .stringScalar:83            throw EvaluationError(message: "type mismatch: expected a numeric expression")84        }85    }8687    /// Evaluates an `if` qualifier into a keep-mask. Missing → excluded.88    func evaluateCondition(_ expression: ZQExpression) throws -> [Bool] {89        let (values, missing) = try evaluateNumericColumn(expression)90        return (0..<values.count).map { !missing[$0] && values[$0] != 0 }91    }9293    // MARK: - Operators9495    private func evaluateUnary(op: String, operand: ZQExpression) throws -> Value {96        let value = try evaluate(operand)97        switch op {98        case "-":99            return try mapNumeric(value) { -$0 }100        case "!":101            return try mapNumeric(value) { $0 == 0 ? 1 : 0 }102        default:103            throw EvaluationError(message: "unknown unary operator '\(op)'")104        }105    }106107    private func evaluateBinary(108        op: String, lhs: ZQExpression, rhs: ZQExpression109    ) throws -> Value {110        let left = try evaluate(lhs)111        let right = try evaluate(rhs)112113        // String equality against a literal or another string column.114        if isString(left) || isString(right) {115            guard op == "==" || op == "!=" else {116                throw EvaluationError(117                    message: "operator '\(op)' is not defined for strings"118                )119            }120            return try compareStrings(left, right, equal: op == "==")121        }122123        let combine: (Double, Double) -> Double124        switch op {125        case "+": combine = (+)126        case "-": combine = (-)127        case "*": combine = (*)128        case "/": combine = { $1 == 0 ? .nan : $0 / $1 }129        case "^": combine = { pow($0, $1) }130        case "==": combine = { $0 == $1 ? 1 : 0 }131        case "!=": combine = { $0 != $1 ? 1 : 0 }132        case "<": combine = { $0 < $1 ? 1 : 0 }133        case "<=": combine = { $0 <= $1 ? 1 : 0 }134        case ">": combine = { $0 > $1 ? 1 : 0 }135        case ">=": combine = { $0 >= $1 ? 1 : 0 }136        case "&": combine = { ($0 != 0 && $1 != 0) ? 1 : 0 }137        case "|": combine = { ($0 != 0 || $1 != 0) ? 1 : 0 }138        default:139            throw EvaluationError(message: "unknown operator '\(op)'")140        }141        return try zipNumeric(left, right, combine)142    }143144    private func evaluateCall(name: String, arguments: [ZQExpression]) throws -> Value {145        // missing(x): 1/0 indicator, itself never missing.146        if name == "missing" || name == "mi" {147            guard arguments.count == 1 else {148                throw EvaluationError(message: "missing() takes exactly one argument")149            }150            switch try evaluate(arguments[0]) {151            case .numeric(_, let missing):152                return .numeric(153                    values: missing.map { $0 ? 1.0 : 0.0 },154                    missing: [Bool](repeating: false, count: missing.count)155                )156            case .strings(let values):157                return .numeric(158                    values: values.map { $0 == nil ? 1.0 : 0.0 },159                    missing: [Bool](repeating: false, count: values.count)160                )161            case .missingScalar:162                return .numericScalar(1)163            case .numericScalar, .stringScalar:164                return .numericScalar(0)165            }166        }167168        let unaryFunctions: [String: @Sendable (Double) -> Double] = [169            "ln": { $0 > 0 ? Foundation.log($0) : .nan },170            "log": { $0 > 0 ? Foundation.log($0) : .nan },171            "log10": { $0 > 0 ? Foundation.log10($0) : .nan },172            "exp": Foundation.exp,173            "sqrt": { $0 >= 0 ? $0.squareRoot() : .nan },174            "abs": abs,175            "floor": { $0.rounded(.down) },176            "ceil": { $0.rounded(.up) },177            "round": { $0.rounded(.toNearestOrAwayFromZero) },178            "int": { $0.rounded(.towardZero) },179        ]180        if let function = unaryFunctions[name] {181            guard arguments.count == 1 else {182                throw EvaluationError(message: "\(name)() takes exactly one argument")183            }184            return try mapNumeric(try evaluate(arguments[0]), function)185        }186187        if name == "max" || name == "min" {188            guard arguments.count >= 2 else {189                throw EvaluationError(message: "\(name)() takes at least two arguments")190            }191            var result = try evaluate(arguments[0])192            for argument in arguments.dropFirst() {193                let next = try evaluate(argument)194                result = try zipNumeric(result, next, name == "max" ? max : min)195            }196            return result197        }198199        throw EvaluationError(message: "unknown function '\(name)()'")200    }201202    // MARK: - Helpers203204    private func isString(_ value: Value) -> Bool {205        switch value {206        case .strings, .stringScalar: return true207        default: return false208        }209    }210211    private func compareStrings(212        _ lhs: Value, _ rhs: Value, equal: Bool213    ) throws -> Value {214        func stringAt(_ value: Value, _ index: Int) -> String?? {215            switch value {216            case .strings(let values): return values[index]217            case .stringScalar(let value): return value218            default: return Optional<String?>.none219            }220        }221        let count: Int222        switch (lhs, rhs) {223        case (.strings(let values), _), (_, .strings(let values)):224            count = values.count225        default:226            // scalar vs scalar227            if case .stringScalar(let a) = lhs, case .stringScalar(let b) = rhs {228                return .numericScalar((a == b) == equal ? 1 : 0)229            }230            throw EvaluationError(message: "type mismatch in string comparison")231        }232        var values = [Double](repeating: 0, count: count)233        let missing = [Bool](repeating: false, count: count)234        for i in 0..<count {235            guard let a = stringAt(lhs, i), let b = stringAt(rhs, i) else {236                throw EvaluationError(message: "type mismatch in string comparison")237            }238            switch (a, b) {239            case (nil, nil), (nil, _), (_, nil):240                // Stata treats "" as the string missing value; comparisons241                // against it still evaluate.242                let equalValues = (a ?? "") == (b ?? "")243                values[i] = equalValues == equal ? 1 : 0244            default:245                values[i] = (a! == b!) == equal ? 1 : 0246            }247        }248        return .numeric(values: values, missing: missing)249    }250251    private func mapNumeric(252        _ value: Value, _ transform: (Double) -> Double253    ) throws -> Value {254        switch value {255        case .numeric(let values, var missing):256            var result = [Double](repeating: .nan, count: values.count)257            for i in 0..<values.count where !missing[i] {258                result[i] = transform(values[i])259                if result[i].isNaN { missing[i] = true }260            }261            return .numeric(values: result, missing: missing)262        case .numericScalar(let value):263            let result = transform(value)264            return result.isNaN ? .missingScalar : .numericScalar(result)265        case .missingScalar:266            return .missingScalar267        case .strings, .stringScalar:268            throw EvaluationError(message: "type mismatch: expected a numeric operand")269        }270    }271272    private func zipNumeric(273        _ lhs: Value, _ rhs: Value, _ combine: (Double, Double) -> Double274    ) throws -> Value {275        switch (lhs, rhs) {276        case (.numericScalar(let a), .numericScalar(let b)):277            let result = combine(a, b)278            return result.isNaN ? .missingScalar : .numericScalar(result)279        case (.missingScalar, _), (_, .missingScalar):280            if case .numeric(let values, _) = lhs {281                return .numeric(282                    values: [Double](repeating: .nan, count: values.count),283                    missing: [Bool](repeating: true, count: values.count)284                )285            }286            if case .numeric(let values, _) = rhs {287                return .numeric(288                    values: [Double](repeating: .nan, count: values.count),289                    missing: [Bool](repeating: true, count: values.count)290                )291            }292            return .missingScalar293        default:294            let (leftValues, leftMissing) = try broadcast(lhs)295            let (rightValues, rightMissing) = try broadcast(rhs)296            guard leftValues.count == rightValues.count else {297                throw EvaluationError(message: "operand length mismatch")298            }299            var values = [Double](repeating: .nan, count: leftValues.count)300            var missing = [Bool](repeating: true, count: leftValues.count)301            for i in 0..<leftValues.count where !leftMissing[i] && !rightMissing[i] {302                let result = combine(leftValues[i], rightValues[i])303                if !result.isNaN {304                    values[i] = result305                    missing[i] = false306                }307            }308            return .numeric(values: values, missing: missing)309        }310    }311312    private func broadcast(313        _ value: Value314    ) throws -> (values: [Double], missing: [Bool]) {315        switch value {316        case .numeric(let values, let missing):317            return (values, missing)318        case .numericScalar(let scalar):319            return (320                [Double](repeating: scalar, count: frame.rowCount),321                [Bool](repeating: false, count: frame.rowCount)322            )323        case .missingScalar:324            return (325                [Double](repeating: .nan, count: frame.rowCount),326                [Bool](repeating: true, count: frame.rowCount)327            )328        case .strings, .stringScalar:329            throw EvaluationError(message: "type mismatch: expected a numeric operand")330        }331    }332}333