spb/metrika Public

1//2//  DuckDBStore.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import DuckDB11import Foundation1213/// DuckDB-backed data store (CLAUDE.md §6). File loading goes through14/// DuckDB readers (`read_parquet`, `read_csv`) and query results are15/// materialized into `ZQDataFrame` columns: every numeric database type16/// widens to Float64, text stays as strings.17public actor ZQDataStore {18    private let database: Database19    private let connection: Connection2021    public init() throws {22        self.database = try Database(store: .inMemory)23        self.connection = try database.connect()24    }2526    // MARK: - Loading2728    /// Loads a dataset file into a data frame. Format is inferred from the29    /// file extension: .parquet, .csv, .tsv, .json, .arrow.30    public func load(contentsOf url: URL) throws -> ZQDataFrame {31        let path = url.path32        guard FileManager.default.fileExists(atPath: path) else {33            throw ZQDataError("file not found: \(path)")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        }3940        let escaped = path.replacingOccurrences(of: "'", with: "''")41        let reader: String42        switch url.pathExtension.lowercased() {43        case "parquet", "pq":44            reader = "read_parquet('\(escaped)')"45        case "csv", "tsv", "txt":46            reader = "read_csv('\(escaped)')"47        case "json", "ndjson", "jsonl":48            reader = "read_json('\(escaped)')"49        case "arrow", "feather", "ipc":50            reader = "read_ipc('\(escaped)')"51        case let ext:52            throw ZQDataError("unsupported file format '.\(ext)'")53        }54        return try dataFrame(fromQuery: "SELECT * FROM \(reader)")55    }5657    /// Runs an arbitrary SQL query and materializes the result.58    ///59    /// Two-phase: a `LIMIT 0` probe through duckdb-swift discovers names60    /// and types, then the full query — with every column cast to DOUBLE61    /// or VARCHAR in SQL — extracts through the C-API bulk reader.62    /// duckdb-swift's per-element access rescans every chunk per read63    /// (O(rows × chunks)); the bulk path is a straight memcpy per numeric64    /// column, ~200× faster at 10M rows.65    public func dataFrame(fromQuery sql: String) throws -> ZQDataFrame {66        let probe: ResultSet67        do {68            probe = try connection.query("SELECT * FROM (\(sql)) __q LIMIT 0")69        } catch {70            throw ZQDataError("query failed: \(error)")71        }7273        var specs: [DuckDBFastReader.ColumnSpec] = []74        var selections: [String] = []75        for index in 0..<probe.columnCount {76            let name = probe.columnName(at: index)77            let quoted = "\"\(name.replacingOccurrences(of: "\"", with: "\"\""))\""78            switch probe.column(at: index).underlyingDatabaseType {79            case .double, .float, .decimal,80                 .tinyint, .smallint, .integer, .bigint,81                 .utinyint, .usmallint, .uinteger, .ubigint:82                specs.append(.init(name: name, isNumeric: true))83                selections.append("CAST(\(quoted) AS DOUBLE) AS \(quoted)")84            case .boolean:85                specs.append(.init(name: name, isNumeric: true))86                selections.append("CAST(CAST(\(quoted) AS INTEGER) AS DOUBLE) AS \(quoted)")87            default:88                // Text, dates, and everything else: string rendering.89                specs.append(.init(name: name, isNumeric: false))90                selections.append("CAST(\(quoted) AS VARCHAR) AS \(quoted)")91            }92        }93        guard !specs.isEmpty else { return try ZQDataFrame() }9495        let castSQL = "SELECT \(selections.joined(separator: ", ")) FROM (\(sql)) __q"96        let reader = try DuckDBFastReader()97        return try reader.read(sql: castSQL, columns: specs)98    }99100    // MARK: - Saving101102    /// Writes a data frame to disk via DuckDB `COPY`. Format inferred from103    /// the extension (.parquet or .csv).104    public func save(_ frame: ZQDataFrame, to url: URL) throws {105        if url.pathExtension.lowercased() == "dta" {106            try DTAWriter.write(frame, to: url)107            return108        }109110        let format: String111        switch url.pathExtension.lowercased() {112        case "parquet", "pq": format = "PARQUET"113        case "csv": format = "CSV, HEADER"114        case let ext: throw ZQDataError("unsupported save format '.\(ext)'")115        }116117        try register(frame, as: "__metrika_save")118        defer { _ = try? connection.query("DROP TABLE IF EXISTS __metrika_save") }119        let escaped = url.path.replacingOccurrences(of: "'", with: "''")120        _ = try connection.query(121            "COPY __metrika_save TO '\(escaped)' (FORMAT \(format))"122        )123    }124125    /// Registers a data frame as a DuckDB table (used by save and by SQL126    /// pushdown of `gen`/`collapse` expressions).127    private func register(_ frame: ZQDataFrame, as table: String) throws {128        _ = try? connection.query("DROP TABLE IF EXISTS \(table)")129        let declarations = frame.columns.map { column in130            let type = column.data.isNumeric ? "DOUBLE" : "VARCHAR"131            return "\"\(column.name)\" \(type)"132        }133        guard !declarations.isEmpty else {134            throw ZQDataError("nothing to save: dataset has no variables")135        }136        _ = try connection.query(137            "CREATE TABLE \(table) (\(declarations.joined(separator: ", ")))"138        )139140        let appender = try Appender(connection: connection, table: table)141        for row in 0..<frame.rowCount {142            for column in frame.columns {143                switch column.data {144                case .float64(let values, let missing):145                    try appender.append(missing[row] ? nil : values[row])146                case .string(let values):147                    try appender.append(values[row])148                }149            }150            try appender.endRow()151        }152        try appender.flush()153    }154}155