support for null values in DimExtractionFn

processing/src/main/java/io/druid/query/topn/DimExtractionTopNAlgorithm.java

@@ -96,10 +96,7 @@ public class DimExtractionTopNAlgorithm extends BaseTopNAlgorithm<Aggregator[][]
         Aggregator[] theAggregators = rowSelector[dimIndex];
         if (theAggregators == null) {
           String key = query.getDimensionSpec().getDimExtractionFn().apply(dimSelector.lookupName(dimIndex));
-          if (key == null) {
+          // null keys are allowed
-            rowSelector[dimIndex] = EMPTY_ARRAY;
-            continue;
-          }
           theAggregators = aggregatesStore.get(key);
           if (theAggregators == null) {
             theAggregators = makeAggregators(cursor, query.getAggregatorSpecs());

processing/src/test/java/io/druid/query/topn/TopNQueryRunnerTest.java

@@ -23,6 +23,7 @@ import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
 import com.metamx.common.guava.Sequences;
 import io.druid.collections.StupidPool;
 import io.druid.query.BySegmentResultValueClass;
@@ -36,6 +37,7 @@ import io.druid.query.aggregation.MaxAggregatorFactory;
 import io.druid.query.aggregation.MinAggregatorFactory;
 import io.druid.query.aggregation.PostAggregator;
 import io.druid.query.dimension.ExtractionDimensionSpec;
+import io.druid.query.extraction.DimExtractionFn;
 import io.druid.query.extraction.RegexDimExtractionFn;
 import io.druid.query.filter.AndDimFilter;
 import io.druid.query.filter.DimFilter;
@@ -52,6 +54,7 @@ import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -1129,6 +1132,74 @@ public class TopNQueryRunnerTest
     TestHelper.assertExpectedResults(expectedResults, runner.run(query));
   }
 
+  @Test
+  public void testTopNDimExtractionNulls()
+  {
+    TopNQuery query = new TopNQueryBuilder()
+        .dataSource(QueryRunnerTestHelper.dataSource)
+        .granularity(QueryRunnerTestHelper.allGran)
+        .dimension(
+            new ExtractionDimensionSpec(
+                providerDimension, providerDimension, new DimExtractionFn()
+            {
+              @Override
+              public byte[] getCacheKey()
+              {
+                return new byte[]{(byte)0xFF};
+              }
+
+              @Override
+              public String apply(String dimValue)
+              {
+                return dimValue.equals("total_market") ? null : dimValue;
+              }
+            })
+        )
+        .metric("rows")
+        .threshold(4)
+        .intervals(QueryRunnerTestHelper.firstToThird)
+        .aggregators(QueryRunnerTestHelper.commonAggregators)
+        .postAggregators(Arrays.<PostAggregator>asList(QueryRunnerTestHelper.addRowsIndexConstant))
+        .build();
+
+    Map<String, Object> nullValue = Maps.newHashMap();
+    nullValue.put(providerDimension, null);
+    nullValue.putAll(ImmutableMap.of(
+                         "rows", 4L,
+                         "index", 5351.814697265625D,
+                         "addRowsIndexConstant", 5356.814697265625D,
+                         "uniques", QueryRunnerTestHelper.UNIQUES_2
+                     ));
+
+    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(
+                        providerDimension, "spot",
+                        "rows", 18L,
+                        "index", 2231.8768157958984D,
+                        "addRowsIndexConstant", 2250.8768157958984D,
+                        "uniques", QueryRunnerTestHelper.UNIQUES_9
+                    ),
+                    nullValue
+                    ,
+                    ImmutableMap.<String, Object>of(
+                        providerDimension, "upfront",
+                        "rows", 4L,
+                        "index", 4875.669677734375D,
+                        "addRowsIndexConstant", 4880.669677734375D,
+                        "uniques", QueryRunnerTestHelper.UNIQUES_2
+                    )
+                )
+            )
+        )
+    );
+
+    TestHelper.assertExpectedResults(expectedResults, runner.run(query));
+  }
+
   @Test
   public void testInvertedTopNQuery()
   {