All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.druid.query.groupby.orderby.DefaultLimitSpec Maven / Gradle / Ivy

There is a newer version: 0.12.3
Show newest version
/*
 * 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.groupby.orderby;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Ordering;
import com.google.common.primitives.Ints;
import com.google.common.primitives.Longs;
import com.metamx.common.ISE;
import com.metamx.common.guava.Sequence;
import com.metamx.common.guava.Sequences;
import io.druid.data.input.Row;
import io.druid.query.aggregation.AggregatorFactory;
import io.druid.query.aggregation.PostAggregator;
import io.druid.query.dimension.DimensionSpec;

import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Map;

/**
 */
public class DefaultLimitSpec implements LimitSpec
{
  private static final byte CACHE_KEY = 0x1;

  private final List columns;
  private final int limit;

  @JsonCreator
  public DefaultLimitSpec(
      @JsonProperty("columns") List columns,
      @JsonProperty("limit") Integer limit
  )
  {
    this.columns = (columns == null) ? ImmutableList.of() : columns;
    this.limit = (limit == null) ? Integer.MAX_VALUE : limit;

    Preconditions.checkArgument(this.limit > 0, "limit[%s] must be >0", limit);
  }

  @JsonProperty
  public List getColumns()
  {
    return columns;
  }

  @JsonProperty
  public int getLimit()
  {
    return limit;
  }

  @Override
  public Function, Sequence> build(
      List dimensions, List aggs, List postAggs
  )
  {
    if (columns.isEmpty()) {
      return new LimitingFn(limit);
    }

    // Materialize the Comparator first for fast-fail error checking.
    final Ordering ordering = makeComparator(dimensions, aggs, postAggs);

    if (limit == Integer.MAX_VALUE) {
      return new SortingFn(ordering);
    } else {
      return new TopNFunction(ordering, limit);
    }
  }

  @Override
  public LimitSpec merge(LimitSpec other)
  {
    return this;
  }

  private Ordering makeComparator(
      List dimensions, List aggs, List postAggs
  )
  {
    Ordering ordering = new Ordering()
    {
      @Override
      public int compare(Row left, Row right)
      {
        return Longs.compare(left.getTimestampFromEpoch(), right.getTimestampFromEpoch());
      }
    };

    Map dimensionsMap = Maps.newHashMap();
    for (DimensionSpec spec : dimensions) {
      dimensionsMap.put(spec.getOutputName(), spec);
    }

    Map aggregatorsMap = Maps.newHashMap();
    for (final AggregatorFactory agg : aggs) {
      aggregatorsMap.put(agg.getName(), agg);
    }

    Map postAggregatorsMap = Maps.newHashMap();
    for (PostAggregator postAgg : postAggs) {
      postAggregatorsMap.put(postAgg.getName(), postAgg);
    }

    for (OrderByColumnSpec columnSpec : columns) {
      String columnName = columnSpec.getDimension();
      Ordering nextOrdering = null;

      if (postAggregatorsMap.containsKey(columnName)) {
        nextOrdering = metricOrdering(columnName, postAggregatorsMap.get(columnName).getComparator());
      } else if (aggregatorsMap.containsKey(columnName)) {
        nextOrdering = metricOrdering(columnName, aggregatorsMap.get(columnName).getComparator());
      } else if (dimensionsMap.containsKey(columnName)) {
        nextOrdering = dimensionOrdering(columnName);
      }

      if (nextOrdering == null) {
        throw new ISE("Unknown column in order clause[%s]", columnSpec);
      }

      switch (columnSpec.getDirection()) {
        case DESCENDING:
          nextOrdering = nextOrdering.reverse();
      }

      ordering = ordering.compound(nextOrdering);
    }

    return ordering;
  }

  private Ordering metricOrdering(final String column, final Comparator comparator)
  {
    return new Ordering()
    {
      @SuppressWarnings("unchecked")
      @Override
      public int compare(Row left, Row right)
      {
        return comparator.compare(left.getRaw(column), right.getRaw(column));
      }
    };
  }

  private Ordering dimensionOrdering(final String dimension)
  {
    return Ordering.natural()
                   .nullsFirst()
                   .onResultOf(
                       new Function()
                       {
                         @Override
                         public String apply(Row input)
                         {
                           // Multi-value dimensions have all been flattened at this point;
                           final List dimList = input.getDimension(dimension);
                           return dimList.isEmpty() ? null : dimList.get(0);
                         }
                       }
                   );
  }

  @Override
  public String toString()
  {
    return "DefaultLimitSpec{" +
           "columns='" + columns + '\'' +
           ", limit=" + limit +
           '}';
  }

  private static class LimitingFn implements Function, Sequence>
  {
    private int limit;

    public LimitingFn(int limit)
    {
      this.limit = limit;
    }

    @Override
    public Sequence apply(
        Sequence input
    )
    {
      return Sequences.limit(input, limit);
    }

    @Override
    public boolean equals(Object o)
    {
      if (this == o) {
        return true;
      }
      if (o == null || getClass() != o.getClass()) {
        return false;
      }

      LimitingFn that = (LimitingFn) o;

      if (limit != that.limit) {
        return false;
      }

      return true;
    }

    @Override
    public int hashCode()
    {
      return limit;
    }
  }

  private static class SortingFn implements Function, Sequence>
  {
    private final Ordering ordering;

    public SortingFn(Ordering ordering) {this.ordering = ordering;}

    @Override
    public Sequence apply(@Nullable Sequence input)
    {
      return Sequences.sort(input, ordering);
    }

    @Override
    public boolean equals(Object o)
    {
      if (this == o) {
        return true;
      }
      if (o == null || getClass() != o.getClass()) {
        return false;
      }

      SortingFn sortingFn = (SortingFn) o;

      if (ordering != null ? !ordering.equals(sortingFn.ordering) : sortingFn.ordering != null) {
        return false;
      }

      return true;
    }

    @Override
    public int hashCode()
    {
      return ordering != null ? ordering.hashCode() : 0;
    }
  }

  private static class TopNFunction implements Function, Sequence>
  {
    private final TopNSorter sorter;
    private final int limit;

    public TopNFunction(Ordering ordering, int limit)
    {
      this.limit = limit;

      this.sorter = new TopNSorter<>(ordering);
    }

    @Override
    public Sequence apply(
        Sequence input
    )
    {
      final ArrayList materializedList = Sequences.toList(input, Lists.newArrayList());
      return Sequences.simple(sorter.toTopN(materializedList, limit));
    }

    @Override
    public boolean equals(Object o)
    {
      if (this == o) {
        return true;
      }
      if (o == null || getClass() != o.getClass()) {
        return false;
      }

      TopNFunction that = (TopNFunction) o;

      if (limit != that.limit) {
        return false;
      }
      if (sorter != null ? !sorter.equals(that.sorter) : that.sorter != null) {
        return false;
      }

      return true;
    }

    @Override
    public int hashCode()
    {
      int result = sorter != null ? sorter.hashCode() : 0;
      result = 31 * result + limit;
      return result;
    }
  }

  @Override
  public boolean equals(Object o)
  {
    if (this == o) {
      return true;
    }
    if (o == null || getClass() != o.getClass()) {
      return false;
    }

    DefaultLimitSpec that = (DefaultLimitSpec) o;

    if (limit != that.limit) {
      return false;
    }
    if (columns != null ? !columns.equals(that.columns) : that.columns != null) {
      return false;
    }

    return true;
  }

  @Override
  public int hashCode()
  {
    int result = columns != null ? columns.hashCode() : 0;
    result = 31 * result + limit;
    return result;
  }

  @Override
  public byte[] getCacheKey()
  {
    final byte[][] columnBytes = new byte[columns.size()][];
    int columnsBytesSize = 0;
    int index = 0;
    for (OrderByColumnSpec column : columns) {
      columnBytes[index] = column.getCacheKey();
      columnsBytesSize += columnBytes[index].length;
      ++index;
    }

    ByteBuffer buffer = ByteBuffer.allocate(1 + columnsBytesSize + 4)
                                  .put(CACHE_KEY);
    for (byte[] columnByte : columnBytes) {
      buffer.put(columnByte);
    }
    buffer.put(Ints.toByteArray(limit));
    return buffer.array();
  }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy