Merge pull request apache#1532 from metamx/fixTopNDimExtractionDouble…

…Apply Fix TopN dimension extractions being applied twice
mistercrunch · Jul 16, 2015 · c130820 · c130820
2 parents 3a0793a + 9092c66
commit c130820
Show file tree

Hide file tree

Showing 2 changed files with 129 additions and 5 deletions.
diff --git a/processing/src/main/java/io/druid/query/topn/TopNQueryQueryToolChest.java b/processing/src/main/java/io/druid/query/topn/TopNQueryQueryToolChest.java
@@ -238,9 +238,6 @@ public Function<Result<TopNResultValue>, Result<TopNResultValue>> makePostComput
       final TopNQuery query, final MetricManipulationFn fn
   )
   {
-    final ExtractionFn extractionFn = TopNQueryEngine.canApplyExtractionInPost(query)
-                                      ? query.getDimensionSpec().getExtractionFn()
-                                      : null;
     return new Function<Result<TopNResultValue>, Result<TopNResultValue>>()
     {
       private String dimension = query.getDimensionSpec().getOutputName();
@@ -284,8 +281,7 @@ public Map<String, Object> apply(DimensionAndMetricValueExtractor input)
                       values.put(name, fn.manipulate(aggregatorFactories[i], input.getMetric(name)));
                     }
 
-                    final Object dimValue = input.getDimensionValue(dimension);
-                    values.put(dimension, extractionFn == null ? dimValue : extractionFn.apply(dimValue));
+                    values.put(dimension, input.getDimensionValue(dimension));
 
                     return values;
                   }

diff --git a/processing/src/test/java/io/druid/query/topn/TopNQueryRunnerTest.java b/processing/src/test/java/io/druid/query/topn/TopNQueryRunnerTest.java
@@ -1610,6 +1610,134 @@ public void testTopNDimExtraction()
     assertExpectedResults(expectedResults, query);
   }
 
+
+
+  @Test
+  public void testTopNDimExtractionFastTopNOptimalWithReplaceMissing()
+  {
+    TopNQuery query = new TopNQueryBuilder()
+        .dataSource(QueryRunnerTestHelper.dataSource)
+        .granularity(QueryRunnerTestHelper.allGran)
+        .dimension(
+            new ExtractionDimensionSpec(
+                QueryRunnerTestHelper.marketDimension,
+                QueryRunnerTestHelper.marketDimension,
+                new LookupExtractionFn(
+                    new MapLookupExtractor(
+                        ImmutableMap.of(
+                            "spot", "2spot0",
+                            "total_market", "1total_market0",
+                            "upfront", "3upfront0"
+                        )
+                    ), false, "MISSING", true
+                ),
+                null
+            )
+        )
+        .metric("rows")
+        .threshold(4)
+        .intervals(QueryRunnerTestHelper.firstToThird)
+        .aggregators(QueryRunnerTestHelper.commonAggregators)
+        .postAggregators(Arrays.<PostAggregator>asList(QueryRunnerTestHelper.addRowsIndexConstant))
+        .build();
+
+    List<Result<TopNResultValue>> expectedResults = Arrays.asList(
+        new Result<>(
+            new DateTime("2011-04-01T00:00:00.000Z"),
+            new TopNResultValue(
+                Arrays.<Map<String, Object>>asList(
+                    ImmutableMap.<String, Object>of(
+                        QueryRunnerTestHelper.marketDimension, "2spot0",
+                        "rows", 18L,
+                        "index", 2231.8768157958984D,
+                        "addRowsIndexConstant", 2250.8768157958984D,
+                        "uniques", QueryRunnerTestHelper.UNIQUES_9
+                    ),
+                    ImmutableMap.<String, Object>of(
+                        QueryRunnerTestHelper.marketDimension, "1total_market0",
+                        "rows", 4L,
+                        "index", 5351.814697265625D,
+                        "addRowsIndexConstant", 5356.814697265625D,
+                        "uniques", QueryRunnerTestHelper.UNIQUES_2
+                    ),
+                    ImmutableMap.<String, Object>of(
+                        QueryRunnerTestHelper.marketDimension, "3upfront0",
+                        "rows", 4L,
+                        "index", 4875.669677734375D,
+                        "addRowsIndexConstant", 4880.669677734375D,
+                        "uniques", QueryRunnerTestHelper.UNIQUES_2
+                    )
+                )
+            )
+        )
+    );
+    assertExpectedResults(expectedResults, query);
+  }
+
+
+  @Test
+  public void testTopNDimExtractionFastTopNUnOptimalWithReplaceMissing()
+  {
+    TopNQuery query = new TopNQueryBuilder()
+        .dataSource(QueryRunnerTestHelper.dataSource)
+        .granularity(QueryRunnerTestHelper.allGran)
+        .dimension(
+            new ExtractionDimensionSpec(
+                QueryRunnerTestHelper.marketDimension,
+                QueryRunnerTestHelper.marketDimension,
+                new LookupExtractionFn(
+                    new MapLookupExtractor(
+                        ImmutableMap.of(
+                            "spot", "2spot0",
+                            "total_market", "1total_market0",
+                            "upfront", "3upfront0"
+                        )
+                    ), false, "MISSING", false
+                ),
+                null
+            )
+        )
+        .metric("rows")
+        .threshold(4)
+        .intervals(QueryRunnerTestHelper.firstToThird)
+        .aggregators(QueryRunnerTestHelper.commonAggregators)
+        .postAggregators(Arrays.<PostAggregator>asList(QueryRunnerTestHelper.addRowsIndexConstant))
+        .build();
+
+    List<Result<TopNResultValue>> expectedResults = Arrays.asList(
+        new Result<>(
+            new DateTime("2011-04-01T00:00:00.000Z"),
+            new TopNResultValue(
+                Arrays.<Map<String, Object>>asList(
+                    ImmutableMap.<String, Object>of(
+                        QueryRunnerTestHelper.marketDimension, "2spot0",
+                        "rows", 18L,
+                        "index", 2231.8768157958984D,
+                        "addRowsIndexConstant", 2250.8768157958984D,
+                        "uniques", QueryRunnerTestHelper.UNIQUES_9
+                    ),
+                    ImmutableMap.<String, Object>of(
+                        QueryRunnerTestHelper.marketDimension, "1total_market0",
+                        "rows", 4L,
+                        "index", 5351.814697265625D,
+                        "addRowsIndexConstant", 5356.814697265625D,
+                        "uniques", QueryRunnerTestHelper.UNIQUES_2
+                    ),
+                    ImmutableMap.<String, Object>of(
+                        QueryRunnerTestHelper.marketDimension, "3upfront0",
+                        "rows", 4L,
+                        "index", 4875.669677734375D,
+                        "addRowsIndexConstant", 4880.669677734375D,
+                        "uniques", QueryRunnerTestHelper.UNIQUES_2
+                    )
+                )
+            )
+        )
+    );
+    assertExpectedResults(expectedResults, query);
+  }
+
+
   @Test
   // Test a "direct" query
   public void testTopNDimExtractionFastTopNOptimal()