Merge pull request #5 from tomasko126/fix/fix-adding-multiple-fields

fix: Update value of expression column names key in the JSON object with new column names
oskar11120 · Oct 24, 2024 · d7136f6 · d7136f6
2 parents 4e31ef1 + 5269afc
commit d7136f6
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 0 deletions.
diff --git a/Apache.Druid.Querying.Tests.Unit/DataSerializerShould_ApplyTo.cs b/Apache.Druid.Querying.Tests.Unit/DataSerializerShould_ApplyTo.cs
@@ -30,6 +30,7 @@ private sealed record Data(
     private sealed record Aggregations(
         DateTimeOffset Time,
         bool Bool);
+
     [Test]
     public void ExpressionAggregation()
     {

diff --git a/Apache.Druid.Querying.Tests.Unit/QueryShould_MapToRightJson.cs b/Apache.Druid.Querying.Tests.Unit/QueryShould_MapToRightJson.cs
@@ -123,6 +123,42 @@ public void Granularity_()
             var eight = query.Granularity("T1M", "utc", t);
             AssertMatch(eight);
         }
+
+        private sealed record Activity(
+            [property: DataSourceTimeColumn] DateTimeOffset Timestamp,
+            int Duration,
+            int DomainID,
+            int UserID
+        );
+        private sealed record ActivityDimensions(int DomainID);
+        private sealed record ActivityAggregations(List<long> UserIds, int Duration);
+
+        [Test]
+        public void ExpressionAggregationWithCombine()
+        {
+            var query = new Query<Activity>
+                    .GroupBy<ActivityDimensions>
+                    .WithNoVirtualColumns
+                    .WithAggregations<ActivityAggregations>()
+                .Dimensions(type => new ActivityDimensions(type.Default(activity => activity.DomainID)))
+                .Aggregations(type => new ActivityAggregations(
+                    type.Expression<List<long>, string>(
+                        "ARRAY<LONG>[]",
+                        "__acc",
+                        data => $"array_set_add(__acc, {data.UserID})",
+                        data => $"array_set_add_all(__acc, {data.Duration})",
+                        null,
+                        null,
+                        data => "ARRAY<LONG>[]",
+                        true,
+                        true,
+                        false,
+                        1024 * 10
+                    ),
+                    type.Sum(activity => activity.Duration))
+                );
+            AssertMatch(query);
+        }
 
         private record IotMeasurementTimestamps(DateTimeOffset Normal, DateTimeOffset Processed);
         [Test]

diff --git a/...Unit/__snapshots__/QueryShould_MapToRightJson.ExpressionAggregationWithCombine_query.snap b/...Unit/__snapshots__/QueryShould_MapToRightJson.ExpressionAggregationWithCombine_query.snap
@@ -0,0 +1,34 @@
+{
+  "queryType": "groupBy",
+  "dimensions": [
+    {
+      "type": "default",
+      "outputName": "DomainID",
+      "dimension": "DomainID"
+    }
+  ],
+  "aggregations": [
+    {
+      "type": "expression",
+      "name": "UserIds",
+      "initialValue": "ARRAY\u003CLONG\u003E[]",
+      "accumulatorIdentifier": "__acc",
+      "fold": "array_set_add(__acc, \u0022UserID\u0022)",
+      "combine": "array_set_add_all(__acc, \u0022Duration\u0022)",
+      "initialValueCombine": "ARRAY\u003CLONG\u003E[]",
+      "fields": [
+        "\u0022UserID\u0022",
+        "\u0022Duration\u0022"
+      ],
+      "isNullUnlessAggregated": true,
+      "shouldAggregateNullInputs": true,
+      "shouldCombineAggregateNullInputs": false,
+      "maxSizeBytes": 10240
+    },
+    {
+      "type": "longSum",
+      "name": "Duration",
+      "fieldName": "Duration"
+    }
+  ]
+}
diff --git a/Apache.Druid.Querying/Internal/Sections/SectionFactoryJsonMapper.cs b/Apache.Druid.Querying/Internal/Sections/SectionFactoryJsonMapper.cs
@@ -40,8 +40,16 @@ void MapCallParam(ElementFactoryCall.Parameter.Any param, string? callResultMemb
 
                     var (value, columnNames) = DruidExpression.Map(expression.Value, context.ColumnNames, context.DataSerializerOptions);
                     result.Add(expression.Name, value);
+
                     if (options.ExpressionColumnNamesKey is string existing)
+                    {
+                        if (result.Remove(existing, out var node))
+                        {
+                            var existingColumnNames = node.Deserialize<IEnumerable<string>>()!;
+                            columnNames = existingColumnNames.Concat(columnNames).Distinct().ToArray();
+                        }
                         result.Add(existing, JsonSerializer.SerializeToNode(columnNames, context.QuerySerializerOptions));
+                    }
                 },
                 (filterFactory, result) =>
                 {