Merge pull request #259 from diffix/edon/instrumentation

Add logs and instrumentation utils to query context

Author: edongashi

src/OpenDiffix.Core.Tests/Anonymizer.Tests.fs

@@ -26,20 +26,16 @@ let companyColumn = ColumnReference(2, StringType)
 let allAidColumns = ListExpr [ aidColumn; companyColumn ]
 
 let executionContext =
-  { ExecutionContext.Default with
-      QueryContext =
-        { QueryContext.Default with
-            AnonymizationParams =
-              {
-                TableSettings = Map.empty
-                Salt = [||]
-                Suppression = { LowThreshold = 2; LowMeanGap = 0.0; SD = 0. }
-                OutlierCount = { Lower = 1; Upper = 1 }
-                TopCount = { Lower = 1; Upper = 1 }
-                NoiseSD = 0.
-              }
-        }
-  }
+  (QueryContext.makeWithAnonParams
+    {
+      TableSettings = Map.empty
+      Salt = [||]
+      Suppression = { LowThreshold = 2; LowMeanGap = 0.0; SD = 0. }
+      OutlierCount = { Lower = 1; Upper = 1 }
+      TopCount = { Lower = 1; Upper = 1 }
+      NoiseSD = 0.
+    })
+  |> ExecutionContext.fromQueryContext
 
 let anonymizedAggregationContext =
   let threshold = { Lower = 2; Upper = 2 }
@@ -50,9 +46,7 @@ let anonymizedAggregationContext =
         TopCount = threshold
     }
 
-  { ExecutionContext.Default with
-      QueryContext = { QueryContext.Default with AnonymizationParams = anonParams }
-  }
+  QueryContext.makeWithAnonParams anonParams |> ExecutionContext.fromQueryContext
 
 let evaluateAggregator fn args =
   evaluateAggregator executionContext fn args

src/OpenDiffix.Core.Tests/Executor.Tests.fs

@@ -27,9 +27,8 @@ type Tests(db: DBFixture) =
     FunctionExpr(AggregateFunction(Count, { Distinct = true; OrderBy = [] }), [ expression ])
 
   let executionContext =
-    { ExecutionContext.Default with
-        QueryContext = { QueryContext.Default with DataProvider = db.DataProvider }
-    }
+    QueryContext.makeWithDataProvider db.DataProvider
+    |> ExecutionContext.fromQueryContext
 
   let execute plan =
     plan |> Executor.execute executionContext |> Seq.toList

src/OpenDiffix.Core.Tests/Expression.Tests.fs

@@ -535,7 +535,7 @@ let colRef2 = ColumnReference(2, RealType)
 let evaluate expr = Expression.evaluate testRow expr
 
 let evaluateAggregator fn args =
-  evaluateAggregator ExecutionContext.Default fn args testRows
+  evaluateAggregator (ExecutionContext.makeDefault ()) fn args testRows
 
 [<Fact>]
 let ``evaluate scalar expressions`` () =

src/OpenDiffix.Core/CommonTypes.fs

@@ -3,6 +3,8 @@
 module rec OpenDiffix.Core.CommonTypes
 
 open System
+open System.Diagnostics
+open System.Runtime.CompilerServices
 
 // ----------------------------------------------------------------
 // Values
@@ -165,19 +167,14 @@ type QueryContext =
   {
     AnonymizationParams: AnonymizationParams
     DataProvider: IDataProvider
+    Metadata: QueryMetadata
   }
-  static member Default =
-    {
-      AnonymizationParams = AnonymizationParams.Default
-      DataProvider =
-        { new IDataProvider with
-            member _.OpenTable(_table, _columnIndices) = failwith "No tables in data provider"
-            member _.GetSchema() = []
-            member _.Dispose() = ()
-        }
-    }
 
-type NoiseLayers = { BucketSeed: Hash }
+type NoiseLayers =
+  {
+    BucketSeed: Hash
+  }
+  static member Default = { BucketSeed = 0UL }
 
 type ExecutionContext =
   {
@@ -186,4 +183,215 @@ type ExecutionContext =
   }
   member this.AnonymizationParams = this.QueryContext.AnonymizationParams
   member this.DataProvider = this.QueryContext.DataProvider
-  static member Default = { QueryContext = QueryContext.Default; NoiseLayers = { BucketSeed = 0UL } }
+  member this.Metadata = this.QueryContext.Metadata
+
+// ----------------------------------------------------------------
+// Constants
+// ----------------------------------------------------------------
+
+let NULL_TYPE = UnknownType "null_type"
+let MISSING_TYPE = UnknownType "missing_type"
+let MIXED_TYPE = UnknownType "mixed_type"
+
+// ----------------------------------------------------------------
+// Functions
+// ----------------------------------------------------------------
+
+module Row =
+  let equalityComparer = LanguagePrimitives.FastGenericEqualityComparer<Row>
+
+module ExpressionType =
+  /// Resolves the common type from a list of types.
+  let commonType types =
+    types
+    |> List.distinct
+    |> function
+      | [] -> MISSING_TYPE
+      | [ t ] -> t
+      | _ -> MIXED_TYPE
+
+module Function =
+  let fromString name =
+    match name with
+    | "count" -> AggregateFunction(Count, AggregateOptions.Default)
+    | "sum" -> AggregateFunction(Sum, AggregateOptions.Default)
+    | "diffix_count" -> AggregateFunction(DiffixCount, AggregateOptions.Default)
+    | "diffix_low_count" -> AggregateFunction(DiffixLowCount, AggregateOptions.Default)
+    | "+" -> ScalarFunction Add
+    | "-" -> ScalarFunction Subtract
+    | "*" -> ScalarFunction Multiply
+    | "/" -> ScalarFunction Divide
+    | "%" -> ScalarFunction Modulo
+    | "round" -> ScalarFunction Round
+    | "ceil" -> ScalarFunction Ceil
+    | "floor" -> ScalarFunction Floor
+    | "round_by" -> ScalarFunction RoundBy
+    | "ceil_by" -> ScalarFunction CeilBy
+    | "floor_by" -> ScalarFunction FloorBy
+    | "abs" -> ScalarFunction Abs
+    | "length" -> ScalarFunction Length
+    | "lower" -> ScalarFunction Lower
+    | "upper" -> ScalarFunction Upper
+    | "substring" -> ScalarFunction Substring
+    | "||" -> ScalarFunction Concat
+    | "width_bucket" -> ScalarFunction WidthBucket
+    | "cast" -> ScalarFunction Cast
+    | other -> failwith $"Unknown function `{other}`"
+
+module Table =
+  /// Finds a column along with its index. The index is zero-based.
+  let tryFindColumn table columnName =
+    table.Columns
+    |> List.indexed
+    |> List.tryFind (fun (_index, column) -> String.equalsI column.Name columnName)
+
+  /// Finds a column along with its index. The index is zero-based. Fails if column not found.
+  let findColumn table columnName =
+    columnName
+    |> tryFindColumn table
+    |> function
+      | Some column -> column
+      | None -> failwith $"Could not find column `{columnName}` in table `{table.Name}`"
+
+module Schema =
+  /// Finds a table by name in the schema.
+  let tryFindTable schema tableName =
+    schema |> List.tryFind (fun table -> String.equalsI table.Name tableName)
+
+  /// Finds a table by name in the schema. Fails if table not found.
+  let findTable schema tableName =
+    tableName
+    |> tryFindTable schema
+    |> function
+      | Some table -> table
+      | None -> failwith $"Could not find table `{tableName}`."
+
+module AnonymizationParams =
+  /// Returns whether the given column in the table is an AID column.
+  let isAidColumn anonParams tableName columnName =
+    anonParams.TableSettings
+    |> Map.tryFind tableName
+    |> function
+      | Some tableSettings -> tableSettings.AidColumns |> List.exists (String.equalsI columnName)
+      | None -> false
+
+module QueryContext =
+  let private defaultDataProvider =
+    { new IDataProvider with
+        member _.OpenTable(_table, _columnIndices) = failwith "No tables in data provider"
+        member _.GetSchema() = []
+        member _.Dispose() = ()
+    }
+
+  let make anonParams dataProvider =
+    {
+      AnonymizationParams = anonParams
+      DataProvider = dataProvider
+      Metadata = QueryMetadata(fun _msg -> ())
+    }
+
+  let makeDefault () =
+    make AnonymizationParams.Default defaultDataProvider
+
+  let makeWithAnonParams anonParams = make anonParams defaultDataProvider
+
+  let makeWithDataProvider dataProvider =
+    make AnonymizationParams.Default dataProvider
+
+  let withLogger logger queryContext =
+    { queryContext with Metadata = QueryMetadata(logger) }
+
+module ExecutionContext =
+  let fromQueryContext queryContext =
+    { QueryContext = queryContext; NoiseLayers = NoiseLayers.Default }
+
+  let makeDefault () =
+    fromQueryContext (QueryContext.makeDefault ())
+
+// ----------------------------------------------------------------
+// Logging & Instrumentation
+// ----------------------------------------------------------------
+
+// Events are relative to init time, expressed in ticks (unit of 100ns).
+type Ticks = int64
+
+type LogLevel =
+  | DebugLevel
+  | InfoLevel
+  | WarningLevel
+
+type LogMessage = { Timestamp: Ticks; Level: LogLevel; Message: string }
+
+module Ticks =
+  let private ticksPerMillisecond = float TimeSpan.TicksPerMillisecond
+
+  let toTimestamp (t: Ticks) = TimeSpan.FromTicks(t).ToString("c")
+
+  let toDuration (t: Ticks) =
+    let ms = (float t / ticksPerMillisecond)
+
+    if ms >= 1000.0 then
+      (ms / 1000.0).ToString("N3") + "s"
+    else
+      ms.ToString("N3") + "ms"
+
+module LogMessage =
+  let private levelToString =
+    function
+    | DebugLevel -> "[DBG]"
+    | InfoLevel -> "[INF]"
+    | WarningLevel -> "[WRN]"
+
+  let toString (message: LogMessage) : string =
+    $"{Ticks.toTimestamp message.Timestamp} {levelToString message.Level} {message.Message}"
+
+type LoggerCallback = LogMessage -> unit
+
+type QueryMetadata(logger: LoggerCallback) =
+  let globalTimer = Stopwatch.StartNew()
+  let measurements = Collections.Generic.Dictionary<string, Ticks>()
+  let counters = Collections.Generic.Dictionary<string, int>()
+
+  let makeMessage level message =
+    { Timestamp = globalTimer.Elapsed.Ticks; Level = level; Message = message }
+
+  [<Conditional("DEBUG")>]
+  member this.LogDebug(message: string) : unit = logger (makeMessage DebugLevel message)
+
+  member this.Log(message: string) : unit = logger (makeMessage InfoLevel message)
+
+  member this.LogWarning(message: string) : unit =
+    logger (makeMessage WarningLevel message)
+
+  member this.MeasureScope([<CallerMemberName>] ?event: string) : IDisposable =
+    let event = event.Value
+    let stopwatch = Stopwatch.StartNew()
+
+    { new IDisposable with
+        member _.Dispose() =
+          let total = stopwatch.Elapsed.Ticks + (Dictionary.getOrDefault event 0L measurements)
+          stopwatch.Reset()
+          measurements.[event] <- total
+    }
+
+  [<Conditional("DEBUG")>]
+  member this.CountDebug([<CallerMemberName>] ?event: string) : unit = this.Count(event.Value)
+
+  member this.Count([<CallerMemberName>] ?event: string) : unit =
+    let event = event.Value
+    let currentCount = Dictionary.getOrDefault event 0 counters
+    counters.[event] <- currentCount + 1
+
+  override this.ToString() =
+    let builder = Text.StringBuilder()
+
+    if measurements.Count > 0 || counters.Count > 0 then
+      builder.AppendLine("<Metadata>") |> ignore
+
+      measurements
+      |> Seq.iter (fun pair -> builder.AppendLine($"{pair.Key}: {Ticks.toDuration pair.Value}") |> ignore)
+
+      counters
+      |> Seq.iter (fun pair -> builder.AppendLine($"{pair.Key}: {pair.Value}") |> ignore)
+
+    builder.ToString()

src/OpenDiffix.Core/OpenDiffix.Core.fsproj

@@ -3,8 +3,8 @@
     <TargetFramework>net5.0</TargetFramework>
   </PropertyGroup>
   <ItemGroup>
-    <Compile Include="CommonTypes.fs" />
     <Compile Include="Utils.fs" />
+    <Compile Include="CommonTypes.fs" />
     <Compile Include="Value.fs" />
     <Compile Include="Expression.fs" />
     <Compile Include="Anonymizer.fs" />

src/OpenDiffix.Core/QueryEngine.fs

@@ -6,7 +6,7 @@ let rec private extractColumns query =
   query.TargetList
   |> List.filter TargetEntry.isRegular
   |> List.map (fun column -> //
-    { Name = column.Alias; Type = Expression.typeOf (column.Expression) }
+    { Name = column.Alias; Type = Expression.typeOf column.Expression }
   )
 
 // ----------------------------------------------------------------
@@ -16,6 +16,8 @@ let rec private extractColumns query =
 type QueryResult = { Columns: Column list; Rows: Row list }
 
 let run queryContext statement : QueryResult =
+  use _measurer = queryContext.Metadata.MeasureScope()
+
   let query, executionContext =
     statement
     |> Parser.parse