Use custom struct for DataSketches ThetaSketch Aggregator (#56)

winsmith · Copilot · web-flow · commit 652865bd9ffd · 2025-04-10T17:42:41.000+02:00
* Use custom struct for DataSketches ThetaSketch Aggregator Also automatically set a size for the theta sketch in funnel queries. Fix #55 * Update Sources/DataTransferObjects/Query/CustomQuery+CompileDown.swift Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
diff --git a/Sources/DataTransferObjects/Query/Aggregator.swift b/Sources/DataTransferObjects/Query/Aggregator.swift
@@ -118,7 +118,7 @@ public indirect enum Aggregator: Codable, Hashable, Equatable {
     /// querying faster.
     ///
     /// https://druid.apache.org/docs/latest/development/extensions-core/datasketches-theta.html
-    case thetaSketch(GenericAggregator)
+    case thetaSketch(ThetaSketchAggregator)
 
     case quantilesDoublesSketch(QuantilesDoublesSketchAggregator)
 
@@ -199,7 +199,7 @@ public indirect enum Aggregator: Codable, Hashable, Equatable {
         case "stringAny":
             self = try .stringAny(GenericAggregator(from: decoder))
         case "thetaSketch":
-            self = try .thetaSketch(GenericAggregator(from: decoder))
+            self = try .thetaSketch(ThetaSketchAggregator(from: decoder))
         case "quantilesDoublesSketch":
             self = try .quantilesDoublesSketch(QuantilesDoublesSketchAggregator(from: decoder))
         case "filtered":
@@ -424,7 +424,7 @@ public enum AggregatorType: String, Codable, Hashable {
     case quantilesDoublesSketch
     case filtered
 
-    // JavaScript aggregator missing
+    // JavaScript aggregator missing on purpose
 }
 
 /// A filtered aggregator wraps any given aggregator, but only aggregates the values for which the given dimension filter matches.
@@ -450,6 +450,41 @@ public struct FilteredAggregator: Codable, Hashable {
     public let name: String?
 }
 
+/// DataSketches Theta Sketch Aggregator
+///
+/// https://druid.apache.org/docs/latest/development/extensions-core/datasketches-theta/
+public struct ThetaSketchAggregator: Codable, Hashable {
+    public init(name: String, fieldName: String, size: Int? = nil, shouldFinalize: Bool? = nil) {
+        type = .thetaSketch
+        self.name = name
+        self.fieldName = fieldName
+        self.size = size
+        self.shouldFinalize = shouldFinalize
+    }
+
+    public let type: AggregatorType
+
+    /// String representing the output column to store sketch values.
+    public let name: String
+
+    /// A string for the name of the aggregator used at ingestion time.
+    public let fieldName: String
+
+    /// The maximum number of entries sketch object retains.
+    ///
+    /// Must be a power of 2. Higher size means higher accuracy but more space to store sketches. After you index with a particular size, Druid persists the sketch in segments.
+    /// At query time you must use a size greater or equal to the ingested size. See the DataSketches site for details. The default is recommended for the majority of use cases.
+    ///
+    /// Defaults to 16384
+    public let size: Int?
+
+    /// Return the final double type representing the estimate rather than the intermediate sketch type itself. In addition to controlling the finalization of this aggregator, you can control whether
+    /// all aggregators are finalized with the query context parameters finalize and sqlFinalizeOuterSketches.
+    ///
+    /// Defaults to true
+    public let shouldFinalize: Bool?
+}
+
 public struct QuantilesDoublesSketchAggregator: Codable, Hashable {
     public init(
         name: String,
@@ -509,7 +544,7 @@ public struct UserCountAggregator: Codable, Hashable, PrecompilableAggregator {
     public let name: String?
 
     public func precompile() -> (aggregators: [Aggregator], postAggregators: [PostAggregator]) {
-        let aggregators = [Aggregator.thetaSketch(.init(type: .thetaSketch, name: name ?? "Users", fieldName: "clientUser"))]
+        let aggregators = [Aggregator.thetaSketch(.init(name: name ?? "Users", fieldName: "clientUser"))]
 
         return (aggregators: aggregators, postAggregators: [])
     }
diff --git a/Sources/DataTransferObjects/Query/CustomQuery+CompileDown.swift b/Sources/DataTransferObjects/Query/CustomQuery+CompileDown.swift
@@ -42,7 +42,10 @@ public extension CustomQuery {
 
         // Custom Query Types
         if query.queryType == .funnel {
-            query = try precompiledFunnelQuery()
+            // If a namespace is set, increase the accuracy of the funnel query (i.e. the size of the theta sketch).
+            // This helps increase accuracy for bigger customers who have their own namespace while hopefully keeping costs down in telemetry-signals.
+            // https://github.com/TelemetryDeck/SwiftDataTransferObjects/issues/55
+            query = try namespace == nil ? precompiledFunnelQuery() : precompiledFunnelQuery(accuracy: 65536)
         } else if query.queryType == .experiment {
             query = try precompiledExperimentQuery()
         }
diff --git a/Sources/DataTransferObjects/QueryGeneration/CustomQuery+Experiment.swift b/Sources/DataTransferObjects/QueryGeneration/CustomQuery+Experiment.swift
@@ -20,7 +20,6 @@ extension CustomQuery {
                 .filtered(.init(
                     filter: combined.1.filter ?? Filter.empty,
                     aggregator: .thetaSketch(.init(
-                        type: .thetaSketch,
                         name: combined.0,
                         fieldName: "clientUser"
                     ))
@@ -34,7 +33,6 @@ extension CustomQuery {
                 sample2.filter ?? Filter.empty,
             ])),
             aggregator: .thetaSketch(.init(
-                type: .thetaSketch,
                 name: "users",
                 fieldName: "clientUser"
             ))
diff --git a/Sources/DataTransferObjects/QueryGeneration/CustomQuery+Funnel.swift b/Sources/DataTransferObjects/QueryGeneration/CustomQuery+Funnel.swift
@@ -1,5 +1,5 @@
 extension CustomQuery {
-    func precompiledFunnelQuery() throws -> CustomQuery {
+    func precompiledFunnelQuery(accuracy: Int? = nil) throws -> CustomQuery {
         var query = self
 
         guard let steps = steps else { throw QueryGenerationError.keyMissing(reason: "Missing key 'steps'") }
@@ -16,9 +16,9 @@ extension CustomQuery {
             aggregations.append(.filtered(.init(
                 filter: step.filter ?? Filter.empty,
                 aggregator: .thetaSketch(.init(
-                    type: .thetaSketch,
                     name: "\(aggregationNamePrefix)\(index)",
-                    fieldName: "clientUser"
+                    fieldName: "clientUser",
+                    size: accuracy
                 ))
             )))
         }
diff --git a/Sources/DataTransferObjects/QueryGeneration/RetentionQueryGenerator.swift b/Sources/DataTransferObjects/QueryGeneration/RetentionQueryGenerator.swift
@@ -94,7 +94,6 @@ public enum RetentionQueryGenerator {
                 intervals: [.init(dateInterval: interval)]
             )),
             aggregator: .thetaSketch(.init(
-                type: .thetaSketch,
                 name: "_\(title(for: interval))",
                 fieldName: "clientUser"
             ))
diff --git a/Tests/QueryGenerationTests/CompileDownTests.swift b/Tests/QueryGenerationTests/CompileDownTests.swift
@@ -58,7 +58,7 @@ final class CompileDownTests: XCTestCase {
         // This should succeed because an app ID is provided
         let appID = UUID()
         let query = CustomQuery(queryType: .timeseries, appID: appID, baseFilters: .thisApp, relativeIntervals: relativeIntervals, granularity: .all)
-        let precompiledQuery = try query.precompile(organizationAppIDs: [appID1, appID2], isSuperOrg: false)
+        let precompiledQuery = try query.precompile(organizationAppIDs: [appID, appID1, appID2], isSuperOrg: false)
 
         XCTAssertEqual(
             precompiledQuery.filter,
diff --git a/Tests/QueryGenerationTests/ConvenienceAggregatorTests.swift b/Tests/QueryGenerationTests/ConvenienceAggregatorTests.swift
@@ -5,7 +5,7 @@ final class ConvenienceAggregatorTests: XCTestCase {
     func testUserCountQueryGetsPrecompiled() throws {
         let query = CustomQuery(queryType: .timeseries, aggregations: [.userCount(.init())])
         let precompiled = try query.precompile(organizationAppIDs: [UUID()], isSuperOrg: false)
-        let expectedAggregations: [Aggregator] = [.thetaSketch(.init(type: .thetaSketch, name: "Users", fieldName: "clientUser"))]
+        let expectedAggregations: [Aggregator] = [.thetaSketch(.init(name: "Users", fieldName: "clientUser"))]
         XCTAssertEqual(precompiled.aggregations, expectedAggregations)
     }
 
diff --git a/Tests/QueryGenerationTests/ExperimentQueryGenerationTests.swift b/Tests/QueryGenerationTests/ExperimentQueryGenerationTests.swift
@@ -31,7 +31,6 @@ final class ExperimentQueryGenerationTests: XCTestCase {
                 filter: .selector(.init(dimension: "type", value: "payScreenALaunched")),
                 aggregator: .thetaSketch(
                     .init(
-                        type: .thetaSketch,
                         name: "cohort_1",
                         fieldName: "clientUser"
                     )
@@ -41,7 +40,6 @@ final class ExperimentQueryGenerationTests: XCTestCase {
                 filter: .selector(.init(dimension: "type", value: "payScreenBLaunched")),
                 aggregator: .thetaSketch(
                     .init(
-                        type: .thetaSketch,
                         name: "cohort_2",
                         fieldName: "clientUser"
                     )
@@ -51,7 +49,6 @@ final class ExperimentQueryGenerationTests: XCTestCase {
                 filter: .selector(.init(dimension: "type", value: "paymentSucceeded")),
                 aggregator: .thetaSketch(
                     .init(
-                        type: .thetaSketch,
                         name: "success",
                         fieldName: "clientUser"
                     )
@@ -64,7 +61,6 @@ final class ExperimentQueryGenerationTests: XCTestCase {
                 ])),
                 aggregator: .thetaSketch(
                     .init(
-                        type: .thetaSketch,
                         name: "users",
                         fieldName: "clientUser"
                     )
diff --git a/Tests/QueryGenerationTests/FunnelQueryGenerationTests.swift b/Tests/QueryGenerationTests/FunnelQueryGenerationTests.swift
@@ -26,7 +26,6 @@ final class FunnelQueryGenerationTests: XCTestCase {
                 filter: .selector(.init(dimension: "type", value: "appLaunchedRegularly")),
                 aggregator: .thetaSketch(
                     .init(
-                        type: AggregatorType.thetaSketch,
                         name: "_funnel_step_0",
                         fieldName: "clientUser"
                     )
@@ -36,7 +35,6 @@ final class FunnelQueryGenerationTests: XCTestCase {
                 filter: .selector(.init(dimension: "type", value: "dataEntered")),
                 aggregator: .thetaSketch(
                     .init(
-                        type: AggregatorType.thetaSketch,
                         name: "_funnel_step_1",
                         fieldName: "clientUser"
                     )
@@ -46,7 +44,6 @@ final class FunnelQueryGenerationTests: XCTestCase {
                 filter: .selector(.init(dimension: "type", value: "paywallSeen")),
                 aggregator: .thetaSketch(
                     .init(
-                        type: AggregatorType.thetaSketch,
                         name: "_funnel_step_2",
                         fieldName: "clientUser"
                     )
@@ -56,7 +53,6 @@ final class FunnelQueryGenerationTests: XCTestCase {
                 filter: .selector(.init(dimension: "type", value: "conversion")),
                 aggregator: .thetaSketch(
                     .init(
-                        type: AggregatorType.thetaSketch,
                         name: "_funnel_step_3",
                         fieldName: "clientUser"
                     )
diff --git a/Tests/QueryGenerationTests/RetentionQueryGenerationTests.swift b/Tests/QueryGenerationTests/RetentionQueryGenerationTests.swift
@@ -36,7 +36,6 @@ final class RetentionQueryGenerationTests: XCTestCase {
                 )),
                 aggregator: .thetaSketch(
                     .init(
-                        type: AggregatorType.thetaSketch,
                         name: "_2022-08-01T00:00:00.000Z_2022-08-31T23:59:59.000Z",
                         fieldName: "clientUser"
                     )
@@ -54,7 +53,6 @@ final class RetentionQueryGenerationTests: XCTestCase {
                 )),
                 aggregator: .thetaSketch(
                     .init(
-                        type: .thetaSketch,
                         name: "_2022-09-01T00:00:00.000Z_2022-09-30T23:59:59.000Z",
                         fieldName: "clientUser"
                     )
diff --git a/Tests/QueryTests/AggregatorTests.swift b/Tests/QueryTests/AggregatorTests.swift
@@ -68,7 +68,7 @@ final class AggregatorTests: XCTestCase {
     func testThetaSketchAggregatorDecoding() throws {
         let decodedAggregators = try JSONDecoder.telemetryDecoder.decode([Aggregator].self, from: exampleDruidAggregatorsThetaSketch)
 
-        XCTAssertEqual(decodedAggregators, [Aggregator.thetaSketch(.init(type: .thetaSketch, name: "count", fieldName: "clientUser"))])
+        XCTAssertEqual(decodedAggregators, [Aggregator.thetaSketch(.init(name: "count", fieldName: "clientUser"))])
     }
 
     func testQuantilesDoublesSketchAggregator() throws {
@@ -134,7 +134,6 @@ final class AggregatorTests: XCTestCase {
                     ),
                     aggregator: .thetaSketch(
                         .init(
-                            type: .thetaSketch,
                             name: "newSessionBegan",
                             fieldName: "clientUser"
                         )
@@ -158,7 +157,6 @@ final class AggregatorTests: XCTestCase {
                     ),
                     aggregator: .thetaSketch(
                         .init(
-                            type: .thetaSketch,
                             name: "InsightShown",
                             fieldName: "clientUser"
                         )
@@ -180,7 +178,6 @@ final class AggregatorTests: XCTestCase {
                     ),
                     aggregator: .thetaSketch(
                         .init(
-                            type: .thetaSketch,
                             name: "newSessionBegan",
                             fieldName: "clientUser"
                         )
@@ -204,7 +201,6 @@ final class AggregatorTests: XCTestCase {
                     ),
                     aggregator: .thetaSketch(
                         .init(
-                            type: .thetaSketch,
                             name: "InsightShown",
                             fieldName: "clientUser"
                         )
diff --git a/Tests/QueryTests/HashingTests.swift b/Tests/QueryTests/HashingTests.swift
@@ -191,7 +191,7 @@ class HashingTests: XCTestCase {
             granularity: .day,
             aggregations: [
                 .count(.init(name: "signals")),
-                .thetaSketch(.init(type: .thetaSketch, name: "users", fieldName: "clientUser")),
+                .thetaSketch(.init(name: "users", fieldName: "clientUser")),
             ]
         )
 
@@ -203,7 +203,7 @@ class HashingTests: XCTestCase {
             granularity: .day,
             aggregations: [
                 .count(.init(name: "signals")),
-                .thetaSketch(.init(type: .thetaSketch, name: "users", fieldName: "clientUser")),
+                .thetaSketch(.init(name: "users", fieldName: "clientUser")),
             ]
         )
 
diff --git a/Tests/QueryTests/RetentionQueryTests.swift b/Tests/QueryTests/RetentionQueryTests.swift
@@ -22,7 +22,6 @@ class RetentionQueryTests: XCTestCase {
             )),
             aggregator: .thetaSketch(
                 .init(
-                    type: AggregatorType.thetaSketch,
                     name: "_august_clientUser_count",
                     fieldName: "clientUser"
                 )
@@ -40,7 +39,6 @@ class RetentionQueryTests: XCTestCase {
             )),
             aggregator: .thetaSketch(
                 .init(
-                    type: .thetaSketch,
                     name: "_september_clientUser_count",
                     fieldName: "clientUser"
                 )
@@ -58,7 +56,6 @@ class RetentionQueryTests: XCTestCase {
             )),
             aggregator: .thetaSketch(
                 .init(
-                    type: .thetaSketch,
                     name: "_october_clientUser_count",
                     fieldName: "clientUser"
                 )
@@ -76,7 +73,6 @@ class RetentionQueryTests: XCTestCase {
             )),
             aggregator: .thetaSketch(
                 .init(
-                    type: .thetaSketch,
                     name: "_november_clientUser_count",
                     fieldName: "clientUser"
                 )
diff --git a/Tests/QueryTests/VirtualColumnTests.swift b/Tests/QueryTests/VirtualColumnTests.swift
@@ -41,7 +41,7 @@ final class VirtualColumnTests: XCTestCase {
         ],
         granularity: .all,
         aggregations: [
-            .thetaSketch(.init(type: .thetaSketch, name: "count", fieldName: "clientUser"))
+            .thetaSketch(.init(name: "count", fieldName: "clientUser"))
         ],
         threshold: 10,
         metric: .numeric(.init(metric: "count")),

Original file line number	Diff line number	Diff line change
`@@ -5,7 +5,7 @@ final class ConvenienceAggregatorTests: XCTestCase {`
`5`	`5`	`func testUserCountQueryGetsPrecompiled() throws {`
`6`	`6`	`let query = CustomQuery(queryType: .timeseries, aggregations: [.userCount(.init())])`
`7`	`7`	`let precompiled = try query.precompile(organizationAppIDs: [UUID()], isSuperOrg: false)`
`8`		`- let expectedAggregations: [Aggregator] = [.thetaSketch(.init(type: .thetaSketch, name: "Users", fieldName: "clientUser"))]`
	`8`	`+ let expectedAggregations: [Aggregator] = [.thetaSketch(.init(name: "Users", fieldName: "clientUser"))]`
`9`	`9`	`XCTAssertEqual(precompiled.aggregations, expectedAggregations)`
`10`	`10`	`}`
`11`	`11`