DataFrame (#11)

Author: Alexander-Ignition

Playgrounds/README.playground/Contents.swift

@@ -53,3 +53,26 @@ struct Contact: Codable {
 
 let select = try database.prepare("SELECT * FROM contacts;")
 let contacts = try select.array(Contact.self)
+/*:
+ ## DataFrame
+
+ The [DataFrame](https://developer.apple.com/documentation/tabulardata/dataframe) from the [TabularData](https://developer.apple.com/documentation/tabulardata) framework is supported.
+
+ It can help to print the table.
+ */
+let df = try database.prepare("SELECT * FROM contacts;").dataFrame()
+print(df)
+/*:
+ ```
+ ┏━━━┳━━━━━━━┳━━━━━━━━━━┓
+ ┃   ┃ id    ┃ name     ┃
+ ┃   ┃ <Int> ┃ <String> ┃
+ ┡━━━╇━━━━━━━╇━━━━━━━━━━┩
+ │ 0 │     1 │ Paul     │
+ │ 1 │     2 │ John     │
+ └───┴───────┴──────────┘
+ ```
+ ## License
+
+ MIT
+ */

README.md

@@ -53,3 +53,24 @@ struct Contact: Codable {
 let select = try database.prepare("SELECT * FROM contacts;")
 let contacts = try select.array(Contact.self)
 ```
+## DataFrame
+
+The [DataFrame](https://developer.apple.com/documentation/tabulardata/dataframe) from the [TabularData](https://developer.apple.com/documentation/tabulardata) framework is supported.
+
+It can help to print the table.
+```swift
+let df = try database.prepare("SELECT * FROM contacts;").dataFrame()
+print(df)
+```
+```
+┏━━━┳━━━━━━━┳━━━━━━━━━━┓
+┃   ┃ id    ┃ name     ┃
+┃   ┃ <Int> ┃ <String> ┃
+┡━━━╇━━━━━━━╇━━━━━━━━━━┩
+│ 0 │     1 │ Paul     │
+│ 1 │     2 │ John     │
+└───┴───────┴──────────┘
+```
+## License
+
+MIT

Sources/SQLyra/DataFrame.swift

@@ -0,0 +1,86 @@
+#if canImport(TabularData)
+
+import TabularData
+import Foundation
+
+// MARK: - PreparedStatement + DataFrame
+
+@available(macOS 12, iOS 15, tvOS 15, watchOS 8, *)
+extension PreparedStatement {
+    /// Creates a new data frame from a prepared statement.
+    ///
+    /// ```swift
+    /// let df = try db.prepare("SELECT * FROM contacts;").dataFrame()
+    /// print(df)
+    /// ```
+    /// ```
+    /// ┏━━━┳━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━┓
+    /// ┃   ┃ id    ┃ name     ┃ rating   ┃ image   ┃
+    /// ┃   ┃ <Int> ┃ <String> ┃ <Double> ┃ <Data>  ┃
+    /// ┡━━━╇━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━┩
+    /// │ 0 │     5 │ A        │      2,0 │ nil     │
+    /// │ 1 │     6 │ B        │      nil │ 3 bytes │
+    /// └───┴───────┴──────────┴──────────┴─────────┘
+    /// ```
+    /// - Parameters:
+    ///   - capacity: An integer that represents the number of elements the columns can initially store.
+    ///   - transformers: SQLite column value transformers.
+    /// - Returns: The data frame that can print a table.
+    /// - Throws: ``DatabaseError``
+    public func dataFrame(
+        capacity: Int = 0,
+        transformers: [String: ColumnValueTransformer] = ColumnValueTransformer.defaults
+    ) throws -> TabularData.DataFrame {
+
+        let valueTransformers: [ColumnValueTransformer] = (0..<columnCount).map { index in
+            columnDeclaration(at: index).flatMap { transformers[$0] } ?? .string
+        }
+        let columns: [TabularData.AnyColumn] = (0..<columnCount).map { index in
+            let name = columnName(at: index) ?? "N/A"
+            return valueTransformers[index].column(name, capacity)
+        }
+        var df = DataFrame(columns: columns)
+        var count = 0
+        while let row = try row() {
+            df.appendEmptyRow()
+            for index in (0..<columnCount) {
+                df.rows[count][index] = row[index].flatMap { valueTransformers[index].transform($0) }
+            }
+            count += 1
+        }
+        return df
+    }
+}
+
+@available(macOS 12, iOS 15, tvOS 15, watchOS 8, *)
+public struct ColumnValueTransformer: Sendable {
+    public static let string = ColumnValueTransformer(transform: \.string)
+    public static let int = ColumnValueTransformer(transform: \.int)
+    public static let double = ColumnValueTransformer(transform: \.double)
+    public static let blob = ColumnValueTransformer(transform: \.blob)
+
+    public static let defaults: [String: ColumnValueTransformer] = [
+        "INT": .int,
+        "INTEGER": .int,
+        "NUM": .double,
+        "REAL": .double,
+        "FLOAT": .double,
+        "TEXT": .string,
+        "BLOB": .blob,
+    ]
+
+    @usableFromInline
+    let column: @Sendable (_ name: String, _ capacity: Int) -> AnyColumn
+    @usableFromInline
+    let transform: @Sendable (PreparedStatement.Value) -> Any?
+
+    @inlinable
+    public init<T>(transform: @escaping @Sendable (PreparedStatement.Value) -> T?) {
+        self.column = { name, capacity in
+            TabularData.Column<T>(name: name, capacity: capacity).eraseToAnyColumn()
+        }
+        self.transform = transform
+    }
+}
+
+#endif  // TabularData

Sources/SQLyra/PreparedStatement.swift

@@ -149,6 +149,10 @@ extension PreparedStatement {
     public func columnName(at index: Int) -> String? {
         sqlite3_column_name(stmt, Int32(index)).string
     }
+
+    func columnDeclaration(at index: Int) -> String? {
+        sqlite3_column_decltype(stmt, Int32(index)).string
+    }
 }
 
 // MARK: - Result values from a Query

Tests/SQLyraTests/PreparedStatementTests.swift

@@ -93,4 +93,25 @@ struct PreparedStatementTests {
         ]
         #expect(contracts == expected)
     }
+
+    @available(macOS 12, iOS 15, tvOS 15, watchOS 8, *)
+    @Test func dataFrame() throws {
+        let insert = try db.prepare(Contact.insert)
+
+        try insert.bind(parameters: 5, "A").execute().reset()
+        try insert.bind(parameters: 6, "B").execute()
+
+        let df = try db.prepare("SELECT * FROM contacts;").dataFrame()
+        let expected = """
+            ┏━━━┳━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━┓
+            ┃   ┃ id    ┃ name     ┃ rating   ┃ image  ┃
+            ┃   ┃ <Int> ┃ <String> ┃ <Double> ┃ <Data> ┃
+            ┡━━━╇━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━┩
+            │ 0 │     5 │ A        │      nil │ nil    │
+            │ 1 │     6 │ B        │      nil │ nil    │
+            └───┴───────┴──────────┴──────────┴────────┘
+            2 rows, 4 columns
+            """
+        #expect(df.description == expected + "\n")
+    }
 }