Add dimension extraction functionality to SearchQuery

processing/src/main/java/io/druid/query/Druids.java

@@ -17,17 +17,22 @@
 
 package io.druid.query;
 
+import com.google.common.base.Function;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import io.druid.granularity.QueryGranularity;
 import io.druid.query.aggregation.AggregatorFactory;
 import io.druid.query.aggregation.PostAggregator;
+import io.druid.query.datasourcemetadata.DataSourceMetadataQuery;
+import io.druid.query.dimension.DefaultDimensionSpec;
+import io.druid.query.dimension.DimensionSpec;
 import io.druid.query.filter.AndDimFilter;
 import io.druid.query.filter.DimFilter;
 import io.druid.query.filter.NoopDimFilter;
 import io.druid.query.filter.NotDimFilter;
 import io.druid.query.filter.OrDimFilter;
 import io.druid.query.filter.SelectorDimFilter;
-import io.druid.query.datasourcemetadata.DataSourceMetadataQuery;
 import io.druid.query.metadata.metadata.ColumnIncluderator;
 import io.druid.query.metadata.metadata.SegmentMetadataQuery;
 import io.druid.query.search.SearchResultValue;
@@ -44,13 +49,24 @@
 import org.joda.time.DateTime;
 import org.joda.time.Interval;
 
+import javax.annotation.Nullable;
 import java.util.List;
 import java.util.Map;
 
 /**
  */
 public class Druids
 {
+  public static final Function<String, DimensionSpec> DIMENSION_IDENTITY = new Function<String, DimensionSpec>()
+  {
+    @Nullable
+    @Override
+    public DimensionSpec apply(String input)
+    {
+      return new DefaultDimensionSpec(input, input);
+    }
+  };
+
   private Druids()
   {
     throw new AssertionError();
@@ -60,7 +76,7 @@ private Druids()
    * A Builder for AndDimFilter.
    *
    * Required: fields() must be called before build()
-   * 
+   *
    * Usage example:
    * <pre><code>
    *   AndDimFilter andDimFilter = Druids.newAndDimFilterBuilder()
@@ -104,9 +120,9 @@ public static AndDimFilterBuilder newAndDimFilterBuilder()
 
   /**
    * A Builder for OrDimFilter.
-   * 
+   *
    * Required: fields() must be called before build()
-   * 
+   *
    * Usage example:
    * <pre><code>
    *   OrDimFilter orDimFilter = Druids.newOrDimFilterBuilder()
@@ -159,9 +175,9 @@ public static OrDimFilterBuilder newOrDimFilterBuilder()
 
   /**
    * A Builder for NotDimFilter.
-   * 
+   *
    * Required: field() must be called before build()
-   * 
+   *
    * Usage example:
    * <pre><code>
    *   NotDimFilter notDimFilter = Druids.newNotDimFilterBuilder()
@@ -205,9 +221,9 @@ public static NotDimFilterBuilder newNotDimFilterBuilder()
 
   /**
    * A Builder for SelectorDimFilter.
-   * 
+   *
    * Required: dimension() and value() must be called before build()
-   * 
+   *
    * Usage example:
    * <pre><code>
    *   Selector selDimFilter = Druids.newSelectorDimFilterBuilder()
@@ -284,10 +300,10 @@ public static NoopDimFilterBuilder newNoopDimFilterBuilder()
 
   /**
    * A Builder for TimeseriesQuery.
-   * 
+   *
    * Required: dataSource(), intervals(), and aggregators() must be called before build()
    * Optional: filters(), granularity(), postAggregators(), and context() can be called before build()
-   * 
+   *
    * Usage example:
    * <pre><code>
    *   TimeseriesQuery query = Druids.newTimeseriesQueryBuilder()
@@ -482,11 +498,11 @@ public static TimeseriesQueryBuilder newTimeseriesQueryBuilder()
 
   /**
    * A Builder for SearchQuery.
-   * 
+   *
    * Required: dataSource(), intervals(), dimensions() and query() must be called before build()
-   * 
+   *
    * Optional: filters(), granularity(), and context() can be called before build()
-   * 
+   *
    * Usage example:
    * <pre><code>
    *   SearchQuery query = Druids.newSearchQueryBuilder()
@@ -506,7 +522,7 @@ public static class SearchQueryBuilder
     private QueryGranularity granularity;
     private int limit;
     private QuerySegmentSpec querySegmentSpec;
-    private List<String> dimensions;
+    private List<DimensionSpec> dimensions;
     private SearchQuerySpec querySpec;
     private Map<String, Object> context;
 
@@ -634,12 +650,24 @@ public SearchQueryBuilder intervals(List<Interval> l)
     }
 
     public SearchQueryBuilder dimensions(String d)
+    {
+      dimensions = ImmutableList.of(DIMENSION_IDENTITY.apply(d));
+      return this;
+    }
+
+    public SearchQueryBuilder dimensions(Iterable<String> d)
+    {
+      dimensions = ImmutableList.copyOf(Iterables.transform(d, DIMENSION_IDENTITY));
+      return this;
+    }
+
+    public SearchQueryBuilder dimensions(DimensionSpec d)
     {
       dimensions = Lists.newArrayList(d);
       return this;
     }
 
-    public SearchQueryBuilder dimensions(List<String> d)
+    public SearchQueryBuilder dimensions(List<DimensionSpec> d)
     {
       dimensions = d;
       return this;
@@ -677,9 +705,9 @@ public static SearchQueryBuilder newSearchQueryBuilder()
 
   /**
    * A Builder for TimeBoundaryQuery.
-   * 
+   *
    * Required: dataSource() must be called before build()
-   * 
+   *
    * Usage example:
    * <pre><code>
    *   TimeBoundaryQuery query = new MaxTimeQueryBuilder()
@@ -773,9 +801,9 @@ public static TimeBoundaryQueryBuilder newTimeBoundaryQueryBuilder()
 
   /**
    * A Builder for Result.
-   * 
+   *
    * Required: timestamp() and value() must be called before build()
-   * 
+   *
    * Usage example:
    * <pre><code>
    *   Result&lt;T&gt; result = Druids.newResultBuilder()
@@ -839,9 +867,9 @@ public static ResultBuilder<TimeBoundaryResultValue> newTimeBoundaryResultBuilde
 
   /**
    * A Builder for SegmentMetadataQuery.
-   * 
+   *
    * Required: dataSource(), intervals() must be called before build()
-   * 
+   *
    * Usage example:
    * <pre><code>
    *   SegmentMetadataQuery query = new SegmentMetadataQueryBuilder()
@@ -947,9 +975,9 @@ public static SegmentMetadataQueryBuilder newSegmentMetadataQueryBuilder()
 
   /**
    * A Builder for SelectQuery.
-   * 
+   *
    * Required: dataSource(), intervals() must be called before build()
-   * 
+   *
    * Usage example:
    * <pre><code>
    *   SelectQuery query = new SelectQueryBuilder()

processing/src/main/java/io/druid/query/extraction/ExtractionFn.java

@@ -30,7 +30,8 @@
     @JsonSubTypes.Type(name = "searchQuery", value = SearchQuerySpecDimExtractionFn.class),
     @JsonSubTypes.Type(name = "javascript", value = JavascriptExtractionFn.class),
     @JsonSubTypes.Type(name = "timeFormat", value = TimeFormatExtractionFn.class),
-    @JsonSubTypes.Type(name = "lookup", value = LookupExtractionFn.class)
+    @JsonSubTypes.Type(name = "lookup", value = LookupExtractionFn.class),
+    @JsonSubTypes.Type(name = "identity", value = IdentityExtractionFn.class)
 })
 /**
  * An ExtractionFn is a function that can be used to transform the values of a column (typically a dimension)

processing/src/main/java/io/druid/query/extraction/IdentityExtractionFn.java

@@ -0,0 +1,62 @@
+/*
+ * 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.query.extraction;
+
+import com.google.common.base.Strings;
+
+public class IdentityExtractionFn implements ExtractionFn
+{
+  private static final byte CACHE_TYPE_ID = 0x6;
+  @Override
+  public byte[] getCacheKey()
+  {
+    return new byte[]{CACHE_TYPE_ID};
+  }
+
+  @Override
+  public String apply(Object value)
+  {
+    return value == null ? null : Strings.emptyToNull(value.toString());
+  }
+
+  @Override
+  public String apply(String value)
+  {
+    return Strings.emptyToNull(value);
+  }
+
+  @Override
+  public String apply(long value)
+  {
+    return Long.toString(value);
+  }
+
+  @Override
+  public boolean preservesOrdering()
+  {
+    return true;
+  }
+
+  @Override
+  public ExtractionType getExtractionType()
+  {
+    return ExtractionType.ONE_TO_ONE;
+  }
+}

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

@@ -20,10 +20,10 @@
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.google.common.base.Function;
 import com.google.common.base.Functions;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Ordering;
-import com.google.common.collect.Sets;
 import com.google.common.primitives.Ints;
 import com.google.inject.Inject;
 import com.metamx.common.IAE;
@@ -32,19 +32,19 @@
 import com.metamx.common.guava.Sequence;
 import com.metamx.common.guava.Sequences;
 import com.metamx.common.guava.nary.BinaryFn;
-import com.metamx.common.StringUtils;
 import com.metamx.emitter.service.ServiceMetricEvent;
 import io.druid.collections.OrderedMergeSequence;
 import io.druid.query.CacheStrategy;
+import io.druid.query.DruidMetrics;
 import io.druid.query.IntervalChunkingQueryRunnerDecorator;
 import io.druid.query.Query;
-import io.druid.query.DruidMetrics;
 import io.druid.query.QueryRunner;
 import io.druid.query.QueryToolChest;
 import io.druid.query.Result;
 import io.druid.query.ResultGranularTimestampComparator;
 import io.druid.query.ResultMergeQueryRunner;
 import io.druid.query.aggregation.MetricManipulationFn;
+import io.druid.query.dimension.DimensionSpec;
 import io.druid.query.filter.DimFilter;
 import io.druid.query.search.search.SearchHit;
 import io.druid.query.search.search.SearchQuery;
@@ -53,9 +53,9 @@
 
 import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 /**
  */
@@ -152,16 +152,15 @@ public byte[] computeCacheKey(SearchQuery query)
         final byte[] querySpecBytes = query.getQuery().getCacheKey();
         final byte[] granularityBytes = query.getGranularity().cacheKey();
 
-        final Set<String> dimensions = Sets.newTreeSet();
-        if (query.getDimensions() != null) {
-          dimensions.addAll(query.getDimensions());
-        }
+        final Collection<DimensionSpec> dimensions = query.getDimensions() == null
+                                                     ? ImmutableList.<DimensionSpec>of()
+                                                     : query.getDimensions();
 
         final byte[][] dimensionsBytes = new byte[dimensions.size()][];
         int dimensionsBytesSize = 0;
         int index = 0;
-        for (String dimension : dimensions) {
-          dimensionsBytes[index] = StringUtils.toUtf8(dimension);
+        for (DimensionSpec dimension : dimensions) {
+          dimensionsBytes[index] = dimension.getCacheKey();
           dimensionsBytesSize += dimensionsBytes[index].length;
           ++index;
         }
@@ -215,7 +214,7 @@ public Result<SearchResultValue> apply(Object input)
             List<Object> result = (List<Object>) input;
 
             return new Result<>(
-                new DateTime(((Number)result.get(0)).longValue()),
+                new DateTime(((Number) result.get(0)).longValue()),
                 new SearchResultValue(
                     Lists.transform(
                         (List) result.get(1),