Merge pull request #2349 from navis/dimensionspec-for-selectquery

Support dimension spec for select query

docs/content/querying/select-query.md

@@ -28,7 +28,7 @@ There are several main parts to a select query:
 |intervals|A JSON Object representing ISO-8601 Intervals. This defines the time ranges to run the query over.|yes|
 |descending|Whether to make descending ordered result. Default is `false`(ascending). When this is `true`, page identifier and offsets will be negative value.|no|
 |filter|See [Filters](../querying/filters.html)|no|
-|dimensions|A String array of dimensions to select. If left empty, all dimensions are returned.|no|
+|dimensions|A JSON list of dimensions to select; or see [DimensionSpec](../querying/dimensionspecs.html) for ways to extract dimensions. If left empty, all dimensions are returned.|no|
 |metrics|A String array of metrics to select. If left empty, all metrics are returned.|no|
 |pagingSpec|A JSON object indicating offsets into different scanned segments. Query results will return a `pagingIdentifiers` value that can be reused in the next query for pagination.|yes|
 |context|An additional JSON Object which can be used to specify certain flags.|no|

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

@@ -1077,7 +1077,7 @@ public class Druids
     private Map<String, Object> context;
     private DimFilter dimFilter;
     private QueryGranularity granularity;
-    private List<String> dimensions;
+    private List<DimensionSpec> dimensions;
     private List<String> metrics;
     private PagingSpec pagingSpec;
 
@@ -1101,9 +1101,7 @@ public class Druids
           descending,
           dimFilter,
           granularity,
-          dimensions,
+          dimensions, metrics, pagingSpec,
-          metrics,
-          pagingSpec,
           context
       );
     }
@@ -1192,12 +1190,18 @@ public class Druids
       return this;
     }
 
-    public SelectQueryBuilder dimensions(List<String> d)
+    public SelectQueryBuilder dimensionSpecs(List<DimensionSpec> d)
     {
       dimensions = d;
       return this;
     }
 
+    public SelectQueryBuilder dimensions(List<String> d)
+    {
+      dimensions = DefaultDimensionSpec.toSpec(d);
+      return this;
+    }
+
     public SelectQueryBuilder metrics(List<String> m)
     {
       metrics = m;

processing/src/main/java/io/druid/query/dimension/DefaultDimensionSpec.java

@@ -21,16 +21,42 @@ package io.druid.query.dimension;
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.base.Function;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
 import com.metamx.common.StringUtils;
 import io.druid.query.extraction.ExtractionFn;
 import io.druid.segment.DimensionSelector;
 
 import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.List;
 
 /**
  */
 public class DefaultDimensionSpec implements DimensionSpec
 {
+  public static List<DimensionSpec> toSpec(String... dimensionNames)
+  {
+    return toSpec(Arrays.asList(dimensionNames));
+  }
+
+  public static List<DimensionSpec> toSpec(Iterable<String> dimensionNames)
+  {
+    return Lists.newArrayList(
+        Iterables.transform(
+            dimensionNames, new Function<String, DimensionSpec>()
+            {
+              @Override
+              public DimensionSpec apply(String input)
+              {
+                return new DefaultDimensionSpec(input, input);
+              }
+            }
+        )
+    );
+  }
+
   private static final byte CACHE_TYPE_ID = 0x0;
   private final String dimension;
   private final String outputName;
@@ -103,7 +129,8 @@ public class DefaultDimensionSpec implements DimensionSpec
   public boolean equals(Object o)
   {
     if (this == o) return true;
-    if (o == null || getClass() != o.getClass()) return false;
+    // LegacyDimensionSpec can be equal to DefaultDimensionSpec
+    if (!(o instanceof DefaultDimensionSpec)) return false;
 
     DefaultDimensionSpec that = (DefaultDimensionSpec) o;
 

processing/src/main/java/io/druid/query/dimension/ExtractionDimensionSpec.java

@@ -57,6 +57,11 @@ public class ExtractionDimensionSpec implements DimensionSpec
     this.outputName = outputName == null ? dimension : outputName;
   }
 
+  public ExtractionDimensionSpec(String dimension, String outputName, ExtractionFn extractionFn)
+  {
+    this(dimension, outputName, extractionFn, null);
+  }
+
   @Override
   @JsonProperty
   public String getDimension()

processing/src/main/java/io/druid/query/select/PagingSpec.java

@@ -90,6 +90,36 @@ public class PagingSpec
     return queryCacheKey.array();
   }
 
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (!(o instanceof PagingSpec)) {
+      return false;
+    }
+
+    PagingSpec that = (PagingSpec) o;
+
+    if (threshold != that.threshold) {
+      return false;
+    }
+    if (!pagingIdentifiers.equals(that.pagingIdentifiers)) {
+      return false;
+    }
+
+    return true;
+  }
+
+  @Override
+  public int hashCode()
+  {
+    int result = pagingIdentifiers.hashCode();
+    result = 31 * result + threshold;
+    return result;
+  }
+
   @Override
   public String toString()
   {

processing/src/main/java/io/druid/query/select/SelectQuery.java

@@ -28,6 +28,7 @@ import io.druid.query.BaseQuery;
 import io.druid.query.DataSource;
 import io.druid.query.Query;
 import io.druid.query.Result;
+import io.druid.query.dimension.DimensionSpec;
 import io.druid.query.filter.DimFilter;
 import io.druid.query.spec.QuerySegmentSpec;
 
@@ -41,7 +42,7 @@ public class SelectQuery extends BaseQuery<Result<SelectResultValue>>
 {
   private final DimFilter dimFilter;
   private final QueryGranularity granularity;
-  private final List<String> dimensions;
+  private final List<DimensionSpec> dimensions;
   private final List<String> metrics;
   private final PagingSpec pagingSpec;
 
@@ -52,7 +53,7 @@ public class SelectQuery extends BaseQuery<Result<SelectResultValue>>
       @JsonProperty("descending") boolean descending,
       @JsonProperty("filter") DimFilter dimFilter,
       @JsonProperty("granularity") QueryGranularity granularity,
-      @JsonProperty("dimensions") List<String> dimensions,
+      @JsonProperty("dimensions") List<DimensionSpec> dimensions,
       @JsonProperty("metrics") List<String> metrics,
       @JsonProperty("pagingSpec") PagingSpec pagingSpec,
       @JsonProperty("context") Map<String, Object> context
@@ -104,7 +105,7 @@ public class SelectQuery extends BaseQuery<Result<SelectResultValue>>
   }
 
   @JsonProperty
-  public List<String> getDimensions()
+  public List<DimensionSpec> getDimensions()
   {
     return dimensions;
   }

processing/src/main/java/io/druid/query/select/SelectQueryEngine.java

@@ -27,6 +27,7 @@ import com.metamx.common.guava.Sequence;
 import io.druid.query.QueryRunnerHelper;
 import io.druid.query.Result;
 import io.druid.query.dimension.DefaultDimensionSpec;
+import io.druid.query.dimension.DimensionSpec;
 import io.druid.segment.Cursor;
 import io.druid.segment.DimensionSelector;
 import io.druid.segment.LongColumnSelector;
@@ -55,9 +56,9 @@ public class SelectQueryEngine
       );
     }
 
-    final Iterable<String> dims;
+    final Iterable<DimensionSpec> dims;
     if (query.getDimensions() == null || query.getDimensions().isEmpty()) {
-      dims = adapter.getAvailableDimensions();
+      dims = DefaultDimensionSpec.toSpec(adapter.getAvailableDimensions());
     } else {
       dims = query.getDimensions();
     }
@@ -89,10 +90,9 @@ public class SelectQueryEngine
             final LongColumnSelector timestampColumnSelector = cursor.makeLongColumnSelector(Column.TIME_COLUMN_NAME);
 
             final Map<String, DimensionSelector> dimSelectors = Maps.newHashMap();
-            for (String dim : dims) {
+            for (DimensionSpec dim : dims) {
-              // switching to using DimensionSpec for select would allow the use of extractionFn here.
+              final DimensionSelector dimSelector = cursor.makeDimensionSelector(dim);
-              final DimensionSelector dimSelector = cursor.makeDimensionSelector(new DefaultDimensionSpec(dim, dim));
+              dimSelectors.put(dim.getOutputName(), dimSelector);
-              dimSelectors.put(dim, dimSelector);
             }
 
             final Map<String, ObjectColumnSelector> metSelectors = Maps.newHashMap();

processing/src/main/java/io/druid/query/select/SelectQueryQueryToolChest.java

@@ -40,11 +40,13 @@ 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 org.joda.time.DateTime;
 
 import java.nio.ByteBuffer;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -138,16 +140,16 @@ public class SelectQueryQueryToolChest extends QueryToolChest<Result<SelectResul
         final byte[] filterBytes = dimFilter == null ? new byte[]{} : dimFilter.getCacheKey();
         final byte[] granularityBytes = query.getGranularity().cacheKey();
 
-        final Set<String> dimensions = Sets.newTreeSet();
+        List<DimensionSpec> dimensionSpecs = query.getDimensions();
-        if (query.getDimensions() != null) {
+        if (dimensionSpecs == null) {
-          dimensions.addAll(query.getDimensions());
+          dimensionSpecs = Collections.emptyList();
         }
 
-        final byte[][] dimensionsBytes = new byte[dimensions.size()][];
+        final byte[][] dimensionsBytes = new byte[dimensionSpecs.size()][];
         int dimensionsBytesSize = 0;
         int index = 0;
-        for (String dimension : dimensions) {
+        for (DimensionSpec dimension : dimensionSpecs) {
-          dimensionsBytes[index] = StringUtils.toUtf8(dimension);
+          dimensionsBytes[index] = dimension.getCacheKey();
           dimensionsBytesSize += dimensionsBytes[index].length;
           ++index;
         }

processing/src/test/java/io/druid/query/select/SelectQueryRunnerTest.java

@@ -32,6 +32,11 @@ import io.druid.query.QueryRunner;
 import io.druid.query.QueryRunnerTestHelper;
 import io.druid.query.Result;
 import io.druid.query.TableDataSource;
+import io.druid.query.dimension.DefaultDimensionSpec;
+import io.druid.query.dimension.DimensionSpec;
+import io.druid.query.dimension.ExtractionDimensionSpec;
+import io.druid.query.extraction.LookupExtractionFn;
+import io.druid.query.extraction.MapLookupExtractor;
 import io.druid.query.filter.AndDimFilter;
 import io.druid.query.filter.DimFilter;
 import io.druid.query.filter.SelectorDimFilter;
@@ -130,7 +135,7 @@ public class SelectQueryRunnerTest
         descending,
         null,
         QueryRunnerTestHelper.allGran,
-        Arrays.<String>asList(),
+        DefaultDimensionSpec.toSpec(Arrays.<String>asList()),
         Arrays.<String>asList(),
         new PagingSpec(null, 3),
         null
@@ -150,6 +155,134 @@ public class SelectQueryRunnerTest
     verify(expectedResults, results);
   }
 
+  @Test
+  public void testFullOnSelectWithDimensionSpec()
+  {
+    Map<String, String> map = new HashMap<>();
+    map.put("automotive", "automotive0");
+    map.put("business", "business0");
+    map.put("entertainment", "entertainment0");
+    map.put("health", "health0");
+    map.put("mezzanine", "mezzanine0");
+    map.put("news", "news0");
+    map.put("premium", "premium0");
+    map.put("technology", "technology0");
+    map.put("travel", "travel0");
+
+    SelectQuery query = new SelectQuery(
+        new TableDataSource(QueryRunnerTestHelper.dataSource),
+        QueryRunnerTestHelper.fullOnInterval,
+        descending,
+        null,
+        QueryRunnerTestHelper.allGran,
+        Arrays.<DimensionSpec>asList(
+            new DefaultDimensionSpec(QueryRunnerTestHelper.marketDimension, "mar"),
+            new ExtractionDimensionSpec(
+                QueryRunnerTestHelper.qualityDimension,
+                "qual",
+                new LookupExtractionFn(new MapLookupExtractor(map, true), false, null, true, false)
+            ),
+            new DefaultDimensionSpec(QueryRunnerTestHelper.placementDimension, "place")
+        ), Lists.<String>newArrayList(), new PagingSpec(null, 3),
+        null
+    );
+    HashMap<String, Object> context = new HashMap<String, Object>();
+    Iterable<Result<SelectResultValue>> results = Sequences.toList(
+        runner.run(query, context),
+        Lists.<Result<SelectResultValue>>newArrayList()
+    );
+
+    List<Result<SelectResultValue>> expectedResultsAsc = Arrays.asList(
+        new Result<SelectResultValue>(
+            new DateTime("2011-01-12T00:00:00.000Z"),
+            new SelectResultValue(
+                ImmutableMap.of(QueryRunnerTestHelper.segmentId, 2),
+                Arrays.asList(
+                    new EventHolder(
+                        QueryRunnerTestHelper.segmentId,
+                        0,
+                        new ImmutableMap.Builder<String, Object>()
+                            .put(EventHolder.timestampKey, new DateTime("2011-01-12T00:00:00.000Z"))
+                            .put("mar", "spot")
+                            .put("qual", "automotive0")
+                            .put("place", "preferred")
+                            .put(QueryRunnerTestHelper.indexMetric, 100.000000F)
+                            .build()
+                    ),
+                    new EventHolder(
+                        QueryRunnerTestHelper.segmentId,
+                        1,
+                        new ImmutableMap.Builder<String, Object>()
+                            .put(EventHolder.timestampKey, new DateTime("2011-01-12T00:00:00.000Z"))
+                            .put("mar", "spot")
+                            .put("qual", "business0")
+                            .put("place", "preferred")
+                            .put(QueryRunnerTestHelper.indexMetric, 100.000000F)
+                            .build()
+                    ),
+                    new EventHolder(
+                        QueryRunnerTestHelper.segmentId,
+                        2,
+                        new ImmutableMap.Builder<String, Object>()
+                            .put(EventHolder.timestampKey, new DateTime("2011-01-12T00:00:00.000Z"))
+                            .put("mar", "spot")
+                            .put("qual", "entertainment0")
+                            .put("place", "preferred")
+                            .put(QueryRunnerTestHelper.indexMetric, 100.000000F)
+                            .build()
+                    )
+                )
+            )
+        )
+    );
+
+    List<Result<SelectResultValue>> expectedResultsDsc = Arrays.asList(
+        new Result<SelectResultValue>(
+            new DateTime("2011-01-12T00:00:00.000Z"),
+            new SelectResultValue(
+                ImmutableMap.of(QueryRunnerTestHelper.segmentId, -3),
+                Arrays.asList(
+                    new EventHolder(
+                        QueryRunnerTestHelper.segmentId,
+                        -1,
+                        new ImmutableMap.Builder<String, Object>()
+                            .put(EventHolder.timestampKey, new DateTime("2011-04-15T00:00:00.000Z"))
+                            .put("mar", "upfront")
+                            .put("qual", "premium0")
+                            .put("place", "preferred")
+                            .put(QueryRunnerTestHelper.indexMetric, 780.27197265625F)
+                            .build()
+                    ),
+                    new EventHolder(
+                        QueryRunnerTestHelper.segmentId,
+                        -2,
+                        new ImmutableMap.Builder<String, Object>()
+                            .put(EventHolder.timestampKey, new DateTime("2011-04-15T00:00:00.000Z"))
+                            .put("mar", "upfront")
+                            .put("qual", "mezzanine0")
+                            .put("place", "preferred")
+                            .put(QueryRunnerTestHelper.indexMetric, 962.731201171875F)
+                            .build()
+                    ),
+                    new EventHolder(
+                        QueryRunnerTestHelper.segmentId,
+                        -3,
+                        new ImmutableMap.Builder<String, Object>()
+                            .put(EventHolder.timestampKey, new DateTime("2011-04-15T00:00:00.000Z"))
+                            .put("mar", "total_market")
+                            .put("qual", "premium0")
+                            .put("place", "preferred")
+                            .put(QueryRunnerTestHelper.indexMetric, 1029.0570068359375F)
+                            .build()
+                    )
+                )
+            )
+        )
+    );
+
+    verify(descending ? expectedResultsDsc : expectedResultsAsc, results);
+  }
+
   @Test
   public void testSelectWithDimsAndMets()
   {
@@ -159,7 +292,7 @@ public class SelectQueryRunnerTest
         descending,
         null,
         QueryRunnerTestHelper.allGran,
-        Arrays.asList(QueryRunnerTestHelper.marketDimension),
+        DefaultDimensionSpec.toSpec(Arrays.asList(QueryRunnerTestHelper.marketDimension)),
         Arrays.asList(QueryRunnerTestHelper.indexMetric),
         new PagingSpec(null, 3),
         null
@@ -198,7 +331,7 @@ public class SelectQueryRunnerTest
         descending,
         null,
         QueryRunnerTestHelper.allGran,
-        Arrays.asList(QueryRunnerTestHelper.qualityDimension),
+        DefaultDimensionSpec.toSpec(Arrays.asList(QueryRunnerTestHelper.qualityDimension)),
         Arrays.asList(QueryRunnerTestHelper.indexMetric),
         new PagingSpec(toPagingIdentifier(3, descending), 3),
         null
@@ -236,7 +369,7 @@ public class SelectQueryRunnerTest
           descending,
           new SelectorDimFilter(QueryRunnerTestHelper.marketDimension, "spot"),
           QueryRunnerTestHelper.dayGran,
-          Lists.<String>newArrayList(QueryRunnerTestHelper.qualityDimension),
+          DefaultDimensionSpec.toSpec(Lists.<String>newArrayList(QueryRunnerTestHelper.qualityDimension)),
           Lists.<String>newArrayList(QueryRunnerTestHelper.indexMetric),
           new PagingSpec(toPagingIdentifier(param[0], descending), param[1]),
           null
@@ -299,15 +432,13 @@ public class SelectQueryRunnerTest
         I_0112_0114,
         descending,
         new AndDimFilter(
-            Arrays.<DimFilter>asList(
+                Arrays.<DimFilter>asList(
-                new SelectorDimFilter(QueryRunnerTestHelper.marketDimension, "spot"),
+                    new SelectorDimFilter(QueryRunnerTestHelper.marketDimension, "spot"),
-                new SelectorDimFilter(QueryRunnerTestHelper.marketDimension, "foo")
+                    new SelectorDimFilter(QueryRunnerTestHelper.marketDimension, "foo")
-            )
+                )
-        ),
+            ),
         QueryRunnerTestHelper.allGran,
-        Lists.<String>newArrayList(),
+        DefaultDimensionSpec.toSpec(Lists.<String>newArrayList()), Lists.<String>newArrayList(), new PagingSpec(null, 3),
-        Lists.<String>newArrayList(),
-        new PagingSpec(null, 3),
         null
     );
 
@@ -338,7 +469,7 @@ public class SelectQueryRunnerTest
         descending,
         null,
         QueryRunnerTestHelper.allGran,
-        Lists.<String>newArrayList("foo"),
+        DefaultDimensionSpec.toSpec(Lists.<String>newArrayList("foo")),
         Lists.<String>newArrayList("foo2"),
         new PagingSpec(null, 3),
         null

processing/src/test/java/io/druid/query/select/SelectQuerySpecTest.java

@@ -0,0 +1,82 @@
+/*
+ * 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.select;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import io.druid.jackson.DefaultObjectMapper;
+import io.druid.query.QueryRunnerTestHelper;
+import io.druid.query.TableDataSource;
+import io.druid.query.dimension.DefaultDimensionSpec;
+import io.druid.query.spec.LegacySegmentSpec;
+import org.joda.time.Interval;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.Arrays;
+
+/**
+ */
+public class SelectQuerySpecTest
+{
+  private static final ObjectMapper jsonMapper = new DefaultObjectMapper();
+
+  @Test
+  public void testSerializationLegacyString() throws Exception
+  {
+    String legacy =
+        "{\"queryType\":\"select\",\"dataSource\":{\"type\":\"table\",\"name\":\"testing\"},"
+        + "\"intervals\":{\"type\":\"LegacySegmentSpec\",\"intervals\":[\"2011-01-12T00:00:00.000Z/2011-01-14T00:00:00.000Z\"]},"
+        + "\"descending\":true,"
+        + "\"filter\":null,"
+        + "\"granularity\":{\"type\":\"all\"},"
+        + "\"dimensions\":[\"market\",\"quality\"],"
+        + "\"metrics\":[\"index\"],"
+        + "\"pagingSpec\":{\"pagingIdentifiers\":{},\"threshold\":3},"
+        + "\"context\":null}";
+
+    String current =
+        "{\"queryType\":\"select\",\"dataSource\":{\"type\":\"table\",\"name\":\"testing\"},"
+        + "\"intervals\":{\"type\":\"LegacySegmentSpec\",\"intervals\":[\"2011-01-12T00:00:00.000Z/2011-01-14T00:00:00.000Z\"]},"
+        + "\"descending\":true,"
+        + "\"filter\":null,"
+        + "\"granularity\":{\"type\":\"all\"},"
+        + "\"dimensions\":[{\"type\":\"default\",\"dimension\":\"market\",\"outputName\":\"market\"},{\"type\":\"default\",\"dimension\":\"quality\",\"outputName\":\"quality\"}],"
+        + "\"metrics\":[\"index\"],"
+        + "\"pagingSpec\":{\"pagingIdentifiers\":{},\"threshold\":3},"
+        + "\"context\":null}";
+
+    SelectQuery query = new SelectQuery(
+        new TableDataSource(QueryRunnerTestHelper.dataSource),
+        new LegacySegmentSpec(new Interval("2011-01-12/2011-01-14")),
+        true,
+        null,
+        QueryRunnerTestHelper.allGran,
+        DefaultDimensionSpec.toSpec(Arrays.<String>asList("market", "quality")),
+        Arrays.<String>asList("index"),
+        new PagingSpec(null, 3),
+        null
+    );
+
+    String actual = jsonMapper.writeValueAsString(query);
+    Assert.assertEquals(current, actual);
+    Assert.assertEquals(query, jsonMapper.readValue(actual, SelectQuery.class));
+    Assert.assertEquals(query, jsonMapper.readValue(legacy, SelectQuery.class));
+  }
+}