Fix dimension selectors with extractionFns on missing columns. (#4717)

* Fix dimension selectors with extractionFns on missing columns.

This patch properly applies the requested extractionFn to missing columns.
It's important when the extractionFn maps null to something other than null.

* Extract helper method.

* Change contracts of VirtualColumns and VirtualColumn methods based on review comments.

* Remove unused import.

* Remove unused method.

* Adjust helper function.

* Adjustments

extensions-contrib/virtual-columns/src/main/java/io/druid/segment/MapVirtualColumn.java

@@ -165,25 +165,25 @@ public class MapVirtualColumn implements VirtualColumn
   public DimensionSelector makeDimensionSelector(DimensionSpec dimensionSpec, ColumnSelectorFactory factory)
   {
     // Could probably do something useful here if the column name is dot-style. But for now just return nothing.
-    return null;
+    return dimensionSpec.decorate(DimensionSelectorUtils.constantSelector(null, dimensionSpec.getExtractionFn()));
   }
 
   @Override
   public FloatColumnSelector makeFloatColumnSelector(String columnName, ColumnSelectorFactory factory)
   {
-    return null;
+    return ZeroFloatColumnSelector.instance();
   }
 
   @Override
   public LongColumnSelector makeLongColumnSelector(String columnName, ColumnSelectorFactory factory)
   {
-    return null;
+    return ZeroLongColumnSelector.instance();
   }
 
   @Override
   public DoubleColumnSelector makeDoubleColumnSelector(String columnName, ColumnSelectorFactory factory)
   {
-    return null;
+    return ZeroDoubleColumnSelector.instance();
   }
 
   @Override

processing/src/main/java/io/druid/query/search/SearchQueryRunner.java

@@ -44,7 +44,6 @@ import io.druid.segment.DimensionSelector;
 import io.druid.segment.DoubleColumnSelector;
 import io.druid.segment.FloatColumnSelector;
 import io.druid.segment.LongColumnSelector;
-import io.druid.segment.NullDimensionSelector;
 import io.druid.segment.Segment;
 import io.druid.segment.column.ColumnCapabilities;
 import io.druid.segment.column.ValueType;
@@ -132,7 +131,7 @@ public class SearchQueryRunner implements QueryRunner<Result<SearchResultValue>>
         final Object2IntRBTreeMap<SearchHit> set
     )
     {
-      if (selector != null && !(selector instanceof NullDimensionSelector)) {
+      if (selector != null && !isNilSelector(selector)) {
         final IndexedInts vals = selector.getRow();
         for (int i = 0; i < vals.size(); ++i) {
           final String dimVal = selector.lookupName(vals.get(i));
@@ -147,6 +146,13 @@ public class SearchQueryRunner implements QueryRunner<Result<SearchResultValue>>
     }
   }
 
+  private static boolean isNilSelector(final DimensionSelector selector)
+  {
+    return selector.nameLookupPossibleInAdvance()
+           && selector.getValueCardinality() == 1
+           && selector.lookupName(0) == null;
+  }
+
   public static class LongSearchColumnSelectorStrategy implements SearchColumnSelectorStrategy<LongColumnSelector>
   {
     @Override

processing/src/main/java/io/druid/segment/ConstantDimensionSelector.java

@@ -0,0 +1,121 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Metamarkets licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package io.druid.segment;
+
+import com.google.common.base.Predicate;
+import com.google.common.base.Strings;
+import io.druid.query.filter.ValueMatcher;
+import io.druid.query.monomorphicprocessing.RuntimeShapeInspector;
+import io.druid.segment.data.IndexedInts;
+import io.druid.segment.data.ZeroIndexedInts;
+import io.druid.segment.filter.BooleanValueMatcher;
+import io.druid.segment.historical.SingleValueHistoricalDimensionSelector;
+
+import javax.annotation.Nullable;
+import java.util.Objects;
+
+public class ConstantDimensionSelector implements SingleValueHistoricalDimensionSelector, IdLookup
+{
+  private final String value;
+
+  public ConstantDimensionSelector(final String value)
+  {
+    if (Strings.isNullOrEmpty(value)) {
+      // There's an optimized implementation for nulls that callers should use instead.
+      throw new IllegalArgumentException("Use NullDimensionSelector or DimensionSelectorUtils.constantSelector");
+    }
+
+    this.value = value;
+  }
+
+  @Override
+  public IndexedInts getRow()
+  {
+    return ZeroIndexedInts.instance();
+  }
+
+  @Override
+  public int getRowValue()
+  {
+    return 0;
+  }
+
+  @Override
+  public int getRowValue(int offset)
+  {
+    return 0;
+  }
+
+  @Override
+  public IndexedInts getRow(int offset)
+  {
+    return getRow();
+  }
+
+  @Override
+  public ValueMatcher makeValueMatcher(String matchValue)
+  {
+    return BooleanValueMatcher.of(Objects.equals(value, matchValue));
+  }
+
+  @Override
+  public ValueMatcher makeValueMatcher(Predicate<String> predicate)
+  {
+    return BooleanValueMatcher.of(predicate.apply(value));
+  }
+
+  @Override
+  public int getValueCardinality()
+  {
+    return 1;
+  }
+
+  @Override
+  public String lookupName(int id)
+  {
+    assert id == 0 : "id = " + id;
+    return value;
+  }
+
+  @Override
+  public boolean nameLookupPossibleInAdvance()
+  {
+    return true;
+  }
+
+  @Nullable
+  @Override
+  public IdLookup idLookup()
+  {
+    return this;
+  }
+
+  @Override
+  public int lookupId(String name)
+  {
+    return value.equals(name) ? 0 : -1;
+  }
+
+  @Override
+  public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+  {
+    inspector.visit("value", value);
+  }
+}

processing/src/main/java/io/druid/segment/DimensionSelectorUtils.java

@@ -21,12 +21,15 @@ package io.druid.segment;
 
 import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
+import com.google.common.base.Strings;
 import io.druid.java.util.common.IAE;
+import io.druid.query.extraction.ExtractionFn;
 import io.druid.query.filter.ValueMatcher;
 import io.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import io.druid.segment.data.IndexedInts;
 import io.druid.segment.filter.BooleanValueMatcher;
 
+import javax.annotation.Nullable;
 import java.util.BitSet;
 import java.util.Objects;
 
@@ -246,4 +249,25 @@ public final class DimensionSelectorUtils
     }
     return valueIds;
   }
+
+  public static DimensionSelector constantSelector(@Nullable final String value)
+  {
+    if (Strings.isNullOrEmpty(value)) {
+      return NullDimensionSelector.instance();
+    } else {
+      return new ConstantDimensionSelector(value);
+    }
+  }
+
+  public static DimensionSelector constantSelector(
+      @Nullable final String value,
+      @Nullable final ExtractionFn extractionFn
+  )
+  {
+    if (extractionFn == null) {
+      return constantSelector(value);
+    } else {
+      return constantSelector(extractionFn.apply(value));
+    }
+  }
 }

processing/src/main/java/io/druid/segment/QueryableIndexStorageAdapter.java

@@ -404,7 +404,10 @@ public class QueryableIndexStorageAdapter implements StorageAdapter
                 public Cursor apply(final Interval inputInterval)
                 {
                   final long timeStart = Math.max(interval.getStartMillis(), inputInterval.getStartMillis());
-                  final long timeEnd = Math.min(interval.getEndMillis(), gran.increment(inputInterval.getStart()).getMillis());
+                  final long timeEnd = Math.min(
+                      interval.getEndMillis(),
+                      gran.increment(inputInterval.getStart()).getMillis()
+                  );
 
                   if (descending) {
                     for (; baseOffset.withinBounds(); baseOffset.increment()) {
@@ -503,7 +506,7 @@ public class QueryableIndexStorageAdapter implements StorageAdapter
 
                       final Column columnDesc = index.getColumn(dimension);
                       if (columnDesc == null) {
-                        return NullDimensionSelector.instance();
+                        return DimensionSelectorUtils.constantSelector(null, extractionFn);
                       }
 
                       if (dimension.equals(Column.TIME_COLUMN_NAME)) {
@@ -534,7 +537,7 @@ public class QueryableIndexStorageAdapter implements StorageAdapter
 
                       final DictionaryEncodedColumn<String> column = cachedColumn;
                       if (column == null) {
-                        return NullDimensionSelector.instance();
+                        return DimensionSelectorUtils.constantSelector(null, extractionFn);
                       } else {
                         return column.makeDimensionSelector(this, extractionFn);
                       }
@@ -883,7 +886,7 @@ public class QueryableIndexStorageAdapter implements StorageAdapter
                     return new QueryableIndexBaseCursor<FilteredOffset>()
                     {
                       private Offset baseOffset;
-                      
+
                       {
                         cursorOffset = new FilteredOffset(this, descending, postFilter, bitmapIndexSelector);
                         reset();

processing/src/main/java/io/druid/segment/VirtualColumn.java

@@ -26,7 +26,6 @@ import io.druid.query.dimension.DimensionSpec;
 import io.druid.segment.column.ColumnCapabilities;
 import io.druid.segment.virtual.ExpressionVirtualColumn;
 
-import javax.annotation.Nullable;
 import java.util.List;
 
 /**
@@ -69,9 +68,8 @@ public interface VirtualColumn extends Cacheable
    * @param dimensionSpec the dimensionSpec this column was referenced with
    * @param factory       column selector factory
    *
-   * @return the selector, or null if we can't make a selector
+   * @return the selector, must not be null
    */
-  @Nullable
   DimensionSelector makeDimensionSelector(DimensionSpec dimensionSpec, ColumnSelectorFactory factory);
 
   /**
@@ -81,9 +79,8 @@ public interface VirtualColumn extends Cacheable
    * @param columnName the name this virtual column was referenced with
    * @param factory    column selector factory
    *
-   * @return the selector, or null if we can't make a selector
+   * @return the selector, must not be null
    */
-  @Nullable
   FloatColumnSelector makeFloatColumnSelector(String columnName, ColumnSelectorFactory factory);
 
   /**
@@ -93,9 +90,8 @@ public interface VirtualColumn extends Cacheable
    * @param columnName the name this virtual column was referenced with
    * @param factory    column selector factory
    *
-   * @return the selector, or null if we can't make a selector
+   * @return the selector, must not be null
    */
-  @Nullable
   LongColumnSelector makeLongColumnSelector(String columnName, ColumnSelectorFactory factory);
 
   /**
@@ -105,9 +101,8 @@ public interface VirtualColumn extends Cacheable
    * @param columnName the name this virtual column was referenced with
    * @param factory    column selector factory
    *
-   * @return the selector, or null if we can't make a selector
+   * @return the selector, must not be null
    */
-  @Nullable
   DoubleColumnSelector makeDoubleColumnSelector(String columnName, ColumnSelectorFactory factory);
 
   /**

processing/src/main/java/io/druid/segment/VirtualColumns.java

@@ -127,6 +127,13 @@ public class VirtualColumns implements Cacheable
   private final Map<String, VirtualColumn> withDotSupport;
   private final Map<String, VirtualColumn> withoutDotSupport;
 
+  /**
+   * Returns true if a virtual column exists with a particular columnName.
+   *
+   * @param columnName the column name
+   *
+   * @return true or false
+   */
   public boolean exists(String columnName)
   {
     return getVirtualColumn(columnName) != null;
@@ -142,11 +149,21 @@ public class VirtualColumns implements Cacheable
     return withDotSupport.get(baseColumnName);
   }
 
+  /**
+   * Create an object selector.
+   *
+   * @param columnName column mame
+   * @param factory    base column selector factory
+   *
+   * @return selector
+   *
+   * @throws IllegalArgumentException if the virtual column does not exist (see {@link #exists(String)}
+   */
   public ObjectColumnSelector makeObjectColumnSelector(String columnName, ColumnSelectorFactory factory)
   {
     final VirtualColumn virtualColumn = getVirtualColumn(columnName);
     if (virtualColumn == null) {
-      return null;
+      throw new IAE("No such virtual column[%s]", columnName);
     } else {
       return Preconditions.checkNotNull(
           virtualColumn.makeObjectColumnSelector(columnName, factory),
@@ -157,39 +174,82 @@ public class VirtualColumns implements Cacheable
     }
   }
 
+  /**
+   * Create a dimension (string) selector.
+   *
+   * @param dimensionSpec the dimensionSpec for this selector
+   * @param factory       base column selector factory
+   *
+   * @return selector
+   *
+   * @throws IllegalArgumentException if the virtual column does not exist (see {@link #exists(String)}
+   */
   public DimensionSelector makeDimensionSelector(DimensionSpec dimensionSpec, ColumnSelectorFactory factory)
   {
     final VirtualColumn virtualColumn = getVirtualColumn(dimensionSpec.getDimension());
     if (virtualColumn == null) {
-      return dimensionSpec.decorate(NullDimensionSelector.instance());
+      throw new IAE("No such virtual column[%s]", dimensionSpec.getDimension());
     } else {
       final DimensionSelector selector = virtualColumn.makeDimensionSelector(dimensionSpec, factory);
-      return selector == null ? dimensionSpec.decorate(NullDimensionSelector.instance()) : selector;
+      Preconditions.checkNotNull(selector, "selector");
+      return selector;
     }
   }
 
+  /**
+   * Create a float selector.
+   *
+   * @param columnName column mame
+   * @param factory    base column selector factory
+   *
+   * @return selector
+   *
+   * @throws IllegalArgumentException if the virtual column does not exist (see {@link #exists(String)}
+   */
   public FloatColumnSelector makeFloatColumnSelector(String columnName, ColumnSelectorFactory factory)
   {
     final VirtualColumn virtualColumn = getVirtualColumn(columnName);
     if (virtualColumn == null) {
-      return ZeroFloatColumnSelector.instance();
+      throw new IAE("No such virtual column[%s]", columnName);
     } else {
       final FloatColumnSelector selector = virtualColumn.makeFloatColumnSelector(columnName, factory);
-      return selector == null ? ZeroFloatColumnSelector.instance() : selector;
+      Preconditions.checkNotNull(selector, "selector");
+      return selector;
     }
   }
 
+  /**
+   * Create a long selector.
+   *
+   * @param columnName column mame
+   * @param factory    base column selector factory
+   *
+   * @return selector
+   *
+   * @throws IllegalArgumentException if the virtual column does not exist (see {@link #exists(String)}
+   */
   public LongColumnSelector makeLongColumnSelector(String columnName, ColumnSelectorFactory factory)
   {
     final VirtualColumn virtualColumn = getVirtualColumn(columnName);
     if (virtualColumn == null) {
-      return ZeroLongColumnSelector.instance();
+      throw new IAE("No such virtual column[%s]", columnName);
     } else {
       final LongColumnSelector selector = virtualColumn.makeLongColumnSelector(columnName, factory);
-      return selector == null ? ZeroLongColumnSelector.instance() : selector;
+      Preconditions.checkNotNull(selector, "selector");
+      return selector;
     }
   }
 
+  /**
+   * Create a double selector.
+   *
+   * @param columnName column mame
+   * @param factory    base column selector factory
+   *
+   * @return selector
+   *
+   * @throws IllegalArgumentException if the virtual column does not exist (see {@link #exists(String)}
+   */
   public DoubleColumnSelector makeDoubleColumnSelector(
       String columnName,
       ColumnSelectorFactory factory
@@ -197,10 +257,11 @@ public class VirtualColumns implements Cacheable
   {
     final VirtualColumn virtualColumn = getVirtualColumn(columnName);
     if (virtualColumn == null) {
-      return ZeroDoubleColumnSelector.instance();
+      throw new IAE("No such virtual column[%s]", columnName);
     } else {
       final DoubleColumnSelector selector = virtualColumn.makeDoubleColumnSelector(columnName, factory);
-      return selector == null ? ZeroDoubleColumnSelector.instance() : selector;
+      Preconditions.checkNotNull(selector, "selector");
+      return selector;
     }
   }
 

processing/src/main/java/io/druid/segment/incremental/IncrementalIndexStorageAdapter.java

@@ -38,6 +38,7 @@ import io.druid.segment.Cursor;
 import io.druid.segment.DimensionHandler;
 import io.druid.segment.DimensionIndexer;
 import io.druid.segment.DimensionSelector;
+import io.druid.segment.DimensionSelectorUtils;
 import io.druid.segment.DoubleColumnSelector;
 import io.druid.segment.DoubleWrappingDimensionSelector;
 import io.druid.segment.FloatColumnSelector;
@@ -45,7 +46,6 @@ import io.druid.segment.FloatWrappingDimensionSelector;
 import io.druid.segment.LongColumnSelector;
 import io.druid.segment.LongWrappingDimensionSelector;
 import io.druid.segment.Metadata;
-import io.druid.segment.NullDimensionSelector;
 import io.druid.segment.ObjectColumnSelector;
 import io.druid.segment.SingleScanTimeDimSelector;
 import io.druid.segment.StorageAdapter;
@@ -418,7 +418,7 @@ public class IncrementalIndexStorageAdapter implements StorageAdapter
                   // not a dimension, column may be a metric
                   ColumnCapabilities capabilities = getColumnCapabilities(dimension);
                   if (capabilities == null) {
-                    return NullDimensionSelector.instance();
+                    return DimensionSelectorUtils.constantSelector(null, extractionFn);
                   }
                   if (capabilities.getType() == ValueType.LONG) {
                     return new LongWrappingDimensionSelector(makeLongColumnSelector(dimension), extractionFn);
@@ -431,7 +431,7 @@ public class IncrementalIndexStorageAdapter implements StorageAdapter
                   }
 
                   // if we can't wrap the base column, just return a column of all nulls
-                  return NullDimensionSelector.instance();
+                  return DimensionSelectorUtils.constantSelector(null, extractionFn);
                 } else {
                   final DimensionIndexer indexer = dimensionDesc.getIndexer();
                   return indexer.makeDimensionSelector(dimensionSpec, currEntry, dimensionDesc);

processing/src/test/java/io/druid/query/groupby/GroupByQueryRunnerTest.java

@@ -94,6 +94,7 @@ import io.druid.query.extraction.ExtractionFn;
 import io.druid.query.extraction.JavaScriptExtractionFn;
 import io.druid.query.extraction.MapLookupExtractor;
 import io.druid.query.extraction.RegexDimExtractionFn;
+import io.druid.query.extraction.StringFormatExtractionFn;
 import io.druid.query.extraction.StrlenExtractionFn;
 import io.druid.query.extraction.TimeFormatExtractionFn;
 import io.druid.query.filter.AndDimFilter;
@@ -143,6 +144,7 @@ import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -453,6 +455,36 @@ public class GroupByQueryRunnerTest
     TestHelper.assertExpectedObjects(expectedResults, results, "");
   }
 
+  @Test
+  public void testGroupByOnMissingColumn()
+  {
+    GroupByQuery query = GroupByQuery
+        .builder()
+        .setDataSource(QueryRunnerTestHelper.dataSource)
+        .setQuerySegmentSpec(QueryRunnerTestHelper.firstToThird)
+        .setDimensions(
+            Lists.<DimensionSpec>newArrayList(
+                new DefaultDimensionSpec("nonexistent0", "alias0"),
+                new ExtractionDimensionSpec("nonexistent1", "alias1", new StringFormatExtractionFn("foo"))
+            )
+        )
+        .setAggregatorSpecs(Collections.singletonList(QueryRunnerTestHelper.rowsCount))
+        .setGranularity(QueryRunnerTestHelper.allGran)
+        .build();
+
+    List<Row> expectedResults = Collections.singletonList(
+        GroupByQueryRunnerTestHelper.createExpectedRow(
+            "2011-04-01",
+            "alias0", null,
+            "alias1", "foo",
+            "rows", 26L
+        )
+    );
+
+    Iterable<Row> results = GroupByQueryRunnerTestHelper.runQuery(factory, runner, query);
+    TestHelper.assertExpectedObjects(expectedResults, results, "");
+  }
+
   @Test
   public void testGroupByWithStringPostAggregator()
   {

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

@@ -80,6 +80,7 @@ import io.druid.query.extraction.ExtractionFn;
 import io.druid.query.extraction.JavaScriptExtractionFn;
 import io.druid.query.extraction.MapLookupExtractor;
 import io.druid.query.extraction.RegexDimExtractionFn;
+import io.druid.query.extraction.StringFormatExtractionFn;
 import io.druid.query.extraction.StrlenExtractionFn;
 import io.druid.query.extraction.TimeFormatExtractionFn;
 import io.druid.query.filter.AndDimFilter;
@@ -377,6 +378,61 @@ public class TopNQueryRunnerTest
     );
   }
 
+  @Test
+  public void testTopNOnMissingColumn()
+  {
+    TopNQuery query = new TopNQueryBuilder()
+        .dataSource(QueryRunnerTestHelper.dataSource)
+        .granularity(QueryRunnerTestHelper.allGran)
+        .dimension(new DefaultDimensionSpec("nonexistentColumn", "alias"))
+        .metric("rows")
+        .threshold(4)
+        .intervals(QueryRunnerTestHelper.fullOnInterval)
+        .aggregators(Collections.singletonList(new CountAggregatorFactory("rows")))
+        .build();
+
+    final HashMap<String, Object> resultMap = new HashMap<>();
+    resultMap.put("alias", null);
+    resultMap.put("rows", 1209L);
+
+    List<Result<TopNResultValue>> expectedResults = Collections.singletonList(
+        new Result<>(
+            DateTimes.of("2011-01-12T00:00:00.000Z"),
+            new TopNResultValue(Collections.<Map<String, Object>>singletonList(resultMap))
+        )
+    );
+    assertExpectedResults(expectedResults, query);
+  }
+
+  @Test
+  public void testTopNOnMissingColumnWithExtractionFn()
+  {
+    TopNQuery query = new TopNQueryBuilder()
+        .dataSource(QueryRunnerTestHelper.dataSource)
+        .granularity(QueryRunnerTestHelper.allGran)
+        .dimension(new ExtractionDimensionSpec("nonexistentColumn", "alias", new StringFormatExtractionFn("theValue")))
+        .metric("rows")
+        .threshold(4)
+        .intervals(QueryRunnerTestHelper.fullOnInterval)
+        .aggregators(Collections.singletonList(new CountAggregatorFactory("rows")))
+        .build();
+
+    List<Result<TopNResultValue>> expectedResults = Collections.singletonList(
+        new Result<>(
+            DateTimes.of("2011-01-12T00:00:00.000Z"),
+            new TopNResultValue(
+                Collections.<Map<String, Object>>singletonList(
+                    ImmutableMap.<String, Object>builder()
+                        .put("alias", "theValue")
+                        .put("rows", 1209L)
+                        .build()
+                )
+            )
+        )
+    );
+    assertExpectedResults(expectedResults, query);
+  }
+
   @Test
   public void testFullOnTopNOverPostAggs()
   {

processing/src/test/java/io/druid/segment/ConstantDimensionSelectorTest.java

@@ -0,0 +1,97 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Metamarkets licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package io.druid.segment;
+
+import io.druid.query.extraction.StringFormatExtractionFn;
+import io.druid.query.extraction.SubstringDimExtractionFn;
+import io.druid.segment.data.IndexedInts;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.Iterator;
+
+public class ConstantDimensionSelectorTest
+{
+  private final DimensionSelector NULL_SELECTOR = DimensionSelectorUtils.constantSelector(null);
+  private final DimensionSelector CONST_SELECTOR = DimensionSelectorUtils.constantSelector("billy");
+  private final DimensionSelector NULL_EXTRACTION_SELECTOR = DimensionSelectorUtils.constantSelector(
+      null,
+      new StringFormatExtractionFn("billy")
+  );
+  private final DimensionSelector CONST_EXTRACTION_SELECTOR = DimensionSelectorUtils.constantSelector(
+      "billybilly",
+      new SubstringDimExtractionFn(0, 5)
+  );
+
+  @Test
+  public void testGetRow() throws Exception
+  {
+    IndexedInts row = NULL_SELECTOR.getRow();
+    Assert.assertEquals(1, row.size());
+    Assert.assertEquals(0, row.get(0));
+
+    Iterator<Integer> iter = row.iterator();
+    Assert.assertEquals(true, iter.hasNext());
+    Assert.assertEquals(0, iter.next().intValue());
+    Assert.assertEquals(false, iter.hasNext());
+  }
+
+  @Test
+  public void testGetValueCardinality() throws Exception
+  {
+    Assert.assertEquals(1, NULL_SELECTOR.getValueCardinality());
+    Assert.assertEquals(1, CONST_SELECTOR.getValueCardinality());
+    Assert.assertEquals(1, NULL_EXTRACTION_SELECTOR.getValueCardinality());
+    Assert.assertEquals(1, CONST_EXTRACTION_SELECTOR.getValueCardinality());
+  }
+
+  @Test
+  public void testLookupName() throws Exception
+  {
+    Assert.assertEquals(null, NULL_SELECTOR.lookupName(0));
+    Assert.assertEquals("billy", CONST_SELECTOR.lookupName(0));
+    Assert.assertEquals("billy", NULL_EXTRACTION_SELECTOR.lookupName(0));
+    Assert.assertEquals("billy", CONST_EXTRACTION_SELECTOR.lookupName(0));
+  }
+
+  @Test
+  public void testLookupId() throws Exception
+  {
+    Assert.assertEquals(0, NULL_SELECTOR.idLookup().lookupId(null));
+    Assert.assertEquals(0, NULL_SELECTOR.idLookup().lookupId(""));
+    Assert.assertEquals(-1, NULL_SELECTOR.idLookup().lookupId("billy"));
+    Assert.assertEquals(-1, NULL_SELECTOR.idLookup().lookupId("bob"));
+
+    Assert.assertEquals(-1, CONST_SELECTOR.idLookup().lookupId(null));
+    Assert.assertEquals(-1, CONST_SELECTOR.idLookup().lookupId(""));
+    Assert.assertEquals(0, CONST_SELECTOR.idLookup().lookupId("billy"));
+    Assert.assertEquals(-1, CONST_SELECTOR.idLookup().lookupId("bob"));
+
+    Assert.assertEquals(-1, NULL_EXTRACTION_SELECTOR.idLookup().lookupId(null));
+    Assert.assertEquals(-1, NULL_EXTRACTION_SELECTOR.idLookup().lookupId(""));
+    Assert.assertEquals(0, NULL_EXTRACTION_SELECTOR.idLookup().lookupId("billy"));
+    Assert.assertEquals(-1, NULL_EXTRACTION_SELECTOR.idLookup().lookupId("bob"));
+
+    Assert.assertEquals(-1, CONST_EXTRACTION_SELECTOR.idLookup().lookupId(null));
+    Assert.assertEquals(-1, CONST_EXTRACTION_SELECTOR.idLookup().lookupId(""));
+    Assert.assertEquals(0, CONST_EXTRACTION_SELECTOR.idLookup().lookupId("billy"));
+    Assert.assertEquals(-1, CONST_EXTRACTION_SELECTOR.idLookup().lookupId("bob"));
+  }
+}

processing/src/test/java/io/druid/segment/NullDimensionSelectorTest.java

@@ -1,65 +0,0 @@
-/*
- * Licensed to Metamarkets Group Inc. (Metamarkets) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. Metamarkets licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-package io.druid.segment;
-
-import io.druid.segment.data.IndexedInts;
-import org.junit.Assert;
-import org.junit.Test;
-
-import java.util.Iterator;
-
-public class NullDimensionSelectorTest
-{
-
-  private final NullDimensionSelector selector = NullDimensionSelector.instance();
-
-  @Test
-  public void testGetRow() throws Exception
-  {
-    IndexedInts row = selector.getRow();
-    Assert.assertEquals(1, row.size());
-    Assert.assertEquals(0, row.get(0));
-
-    Iterator<Integer> iter = row.iterator();
-    Assert.assertEquals(true, iter.hasNext());
-    Assert.assertEquals(0, iter.next().intValue());
-    Assert.assertEquals(false, iter.hasNext());
-  }
-
-  @Test
-  public void testGetValueCardinality() throws Exception
-  {
-    Assert.assertEquals(1, selector.getValueCardinality());
-  }
-
-  @Test
-  public void testLookupName() throws Exception
-  {
-    Assert.assertEquals(null, selector.lookupName(0));
-  }
-
-  @Test
-  public void testLookupId() throws Exception
-  {
-    Assert.assertEquals(0, selector.idLookup().lookupId(null));
-    Assert.assertEquals(0, selector.idLookup().lookupId(""));
-    Assert.assertEquals(-1, selector.idLookup().lookupId("billy"));
-  }
-}

processing/src/test/java/io/druid/segment/virtual/VirtualColumnsTest.java

@@ -65,6 +65,28 @@ public class VirtualColumnsTest
   @Rule
   public ExpectedException expectedException = ExpectedException.none();
 
+  @Test
+  public void testExists()
+  {
+    final VirtualColumns virtualColumns = makeVirtualColumns();
+
+    Assert.assertTrue(virtualColumns.exists("expr"));
+    Assert.assertTrue(virtualColumns.exists("foo"));
+    Assert.assertTrue(virtualColumns.exists("foo.5"));
+    Assert.assertFalse(virtualColumns.exists("bar"));
+  }
+
+  @Test
+  public void testNonExistentSelector()
+  {
+    final VirtualColumns virtualColumns = makeVirtualColumns();
+
+    expectedException.expect(IllegalArgumentException.class);
+    expectedException.expectMessage("No such virtual column[bar]");
+
+    virtualColumns.makeObjectColumnSelector("bar", null);
+  }
+
   @Test
   public void testMakeSelectors()
   {
@@ -406,7 +428,8 @@ public class VirtualColumnsTest
     public DoubleColumnSelector makeDoubleColumnSelector(String columnName, ColumnSelectorFactory factory)
     {
       final ColumnValueSelector selector = makeLongColumnSelector(columnName, factory);
-      return new TestDoubleColumnSelector() {
+      return new TestDoubleColumnSelector()
+      {
 
         @Override
         public double getDouble()