processing/src/main/java/org/apache/druid/query/topn/BaseTopNAlgorithm.java - druid - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF 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 org.apache.druid.query.topn;

 import com.google.common.annotations.VisibleForTesting;
 import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.common.Pair;
 import org.apache.druid.query.aggregation.Aggregator;
 import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.query.aggregation.BufferAggregator;
 import org.apache.druid.segment.Cursor;
 import org.apache.druid.segment.DimensionSelector;
 import org.apache.druid.segment.IdLookup;
 import org.apache.druid.segment.StorageAdapter;
 import org.apache.druid.segment.column.ColumnCapabilities;

 import javax.annotation.Nullable;
 import java.util.Arrays;
 import java.util.Comparator;
 import java.util.List;

 /**
  */
 public abstract class BaseTopNAlgorithm<DimValSelector, DimValAggregateStore, Parameters extends TopNParams>
     implements TopNAlgorithm<DimValSelector, Parameters>
 {
   public static Aggregator[] makeAggregators(Cursor cursor, List<AggregatorFactory> aggregatorSpecs)
   {
     Aggregator[] aggregators = new Aggregator[aggregatorSpecs.size()];
     int aggregatorIndex = 0;
     for (AggregatorFactory spec : aggregatorSpecs) {
       aggregators[aggregatorIndex] = spec.factorize(cursor.getColumnSelectorFactory());
       ++aggregatorIndex;
     }
     return aggregators;
   }

   protected static BufferAggregator[] makeBufferAggregators(Cursor cursor, List<AggregatorFactory> aggregatorSpecs)
   {
     BufferAggregator[] aggregators = new BufferAggregator[aggregatorSpecs.size()];
     int aggregatorIndex = 0;
     for (AggregatorFactory spec : aggregatorSpecs) {
       aggregators[aggregatorIndex] = spec.factorizeBuffered(cursor.getColumnSelectorFactory());
       ++aggregatorIndex;
     }
     return aggregators;
   }

   protected final StorageAdapter storageAdapter;

   protected BaseTopNAlgorithm(StorageAdapter storageAdapter)
   {
     this.storageAdapter = storageAdapter;
   }

   @Override
   public void run(
       Parameters params,
       TopNResultBuilder resultBuilder,
       DimValSelector dimValSelector,
       @Nullable TopNQueryMetrics queryMetrics
   )
   {
     if (params.getCardinality() != TopNParams.CARDINALITY_UNKNOWN) {
       runWithCardinalityKnown(params, resultBuilder, dimValSelector, queryMetrics);
     } else {
       runWithCardinalityUnknown(params, resultBuilder, queryMetrics);
     }
   }

   private void runWithCardinalityKnown(
       Parameters params,
       TopNResultBuilder resultBuilder,
       DimValSelector dimValSelector,
       @Nullable TopNQueryMetrics queryMetrics
   )
   {
     if (queryMetrics != null) {
       queryMetrics.startRecordingScanTime();
     }
     boolean hasDimValSelector = (dimValSelector != null);

     int cardinality = params.getCardinality();
     int numProcessed = 0;
     long processedRows = 0;
     while (numProcessed < cardinality) {
       final int numToProcess;
       int maxNumToProcess = Math.min(params.getNumValuesPerPass(), cardinality - numProcessed);

       DimValSelector theDimValSelector;
       if (!hasDimValSelector) {
         numToProcess = maxNumToProcess;
         theDimValSelector = makeDimValSelector(params, numProcessed, numToProcess);
       } else {
         //skip invalid, calculate length to have enough valid value to process or hit the end.
         numToProcess = computeNewLength(dimValSelector, numProcessed, maxNumToProcess);
         theDimValSelector = updateDimValSelector(dimValSelector, numProcessed, numToProcess);
       }

       DimValAggregateStore aggregatesStore = makeDimValAggregateStore(params);

       processedRows = scanAndAggregate(params, theDimValSelector, aggregatesStore);

       updateResults(params, theDimValSelector, aggregatesStore, resultBuilder);

       closeAggregators(aggregatesStore);

       numProcessed += numToProcess;
       params.getCursor().reset();
     }
     if (queryMetrics != null) {
       queryMetrics.addProcessedRows(processedRows);
       queryMetrics.stopRecordingScanTime();
     }
   }

   /**
    * This function currently handles TopNs on long and float columns, which do not provide cardinality or an ID lookup.
    * When cardinality is unknown, process everything in one pass.
    * Existing implementations of makeDimValSelector() require cardinality as well, so the DimValSelector is not used.
    *
    * @param params        TopN parameters from run()
    * @param resultBuilder Result builder from run()
    */
   private void runWithCardinalityUnknown(
       Parameters params,
       TopNResultBuilder resultBuilder,
       @Nullable TopNQueryMetrics queryMetrics
   )
   {
     DimValAggregateStore aggregatesStore = makeDimValAggregateStore(params);
     if (queryMetrics != null) {
       queryMetrics.startRecordingScanTime();
     }
     long processedRows = scanAndAggregate(params, null, aggregatesStore);
     updateResults(params, null, aggregatesStore, resultBuilder);
     closeAggregators(aggregatesStore);
     params.getCursor().reset();
     if (queryMetrics != null) {
       queryMetrics.addProcessedRows(processedRows);
       queryMetrics.stopRecordingScanTime();
     }
   }

   protected abstract DimValSelector makeDimValSelector(Parameters params, int numProcessed, int numToProcess);

   /**
    * Skip invalid value, calculate length to have enough valid value to process or hit the end.
    *
    * @param dimValSelector the dim value selector which record value is valid or invalid.
    * @param numProcessed   the start position to process
    * @param numToProcess   the number of valid value to process
    *
    * @return the length between which have enough valid value to process or hit the end.
    */
   protected int computeNewLength(DimValSelector dimValSelector, int numProcessed, int numToProcess)
   {
     return numToProcess;
   }

   protected abstract DimValSelector updateDimValSelector(
       DimValSelector dimValSelector,
       int numProcessed,
       int numToProcess
   );

   protected abstract DimValAggregateStore makeDimValAggregateStore(Parameters params);

   /**
    * Returns the number of processed rows (i. e. after postFilters are applied inside the cursor being processed).
    */
   protected abstract long scanAndAggregate(
       Parameters params,
       DimValSelector dimValSelector,
       DimValAggregateStore dimValAggregateStore
   );

   protected abstract void updateResults(
       Parameters params,
       DimValSelector dimValSelector,
       DimValAggregateStore dimValAggregateStore,
       TopNResultBuilder resultBuilder
   );

   protected abstract void closeAggregators(
       DimValAggregateStore dimValAggregateStore
   );

   public static class AggregatorArrayProvider extends BaseArrayProvider<Aggregator[][]>
   {
     Aggregator[][] expansionAggs;
     int cardinality;

     public AggregatorArrayProvider(
         DimensionSelector dimSelector,
         TopNQuery query,
         int cardinality,
         StorageAdapter storageAdapter
     )
     {
       super(dimSelector, query, storageAdapter);

       this.expansionAggs = new Aggregator[cardinality][];
       this.cardinality = cardinality;
     }

     @Override
     public Aggregator[][] build()
     {
       Pair<Integer, Integer> startEnd = computeStartEnd(cardinality);

       Arrays.fill(expansionAggs, 0, startEnd.lhs, EMPTY_ARRAY);
       Arrays.fill(expansionAggs, startEnd.lhs, startEnd.rhs, null);
       Arrays.fill(expansionAggs, startEnd.rhs, expansionAggs.length, EMPTY_ARRAY);

       return expansionAggs;
     }
   }

   protected abstract static class BaseArrayProvider<T> implements TopNMetricSpecBuilder<T>
   {
     private volatile String previousStop;
     private volatile boolean ignoreAfterThreshold;
     private volatile int ignoreFirstN;
     private volatile int keepOnlyN;

     private final IdLookup idLookup;
     private final TopNQuery query;
     private final StorageAdapter storageAdapter;

     public BaseArrayProvider(
         DimensionSelector dimSelector,
         TopNQuery query,
         StorageAdapter storageAdapter
     )
     {
       this.idLookup = dimSelector.idLookup();
       this.query = query;
       this.storageAdapter = storageAdapter;

       previousStop = null;
       ignoreAfterThreshold = false;
       ignoreFirstN = 0;
       keepOnlyN = dimSelector.getValueCardinality();

       if (keepOnlyN < 0) {
         throw new IAE("Cannot operate on a dimension with no dictionary");
       }
     }

     @Override
     public void skipTo(String previousStop)
     {
       ColumnCapabilities capabilities = storageAdapter.getColumnCapabilities(query.getDimensionSpec().getDimension());
       if (capabilities != null && capabilities.areDictionaryValuesSorted().isTrue()) {
         this.previousStop = previousStop;
       }
     }

     @Override
     public void ignoreAfterThreshold()
     {
       ignoreAfterThreshold = true;
     }

     @Override
     public void ignoreFirstN(int n)
     {
       ignoreFirstN = n;
     }

     @Override
     public void keepOnlyN(int n)
     {
       keepOnlyN = n;
     }

     @VisibleForTesting
     public Pair<Integer, Integer> computeStartEnd(int cardinality)
     {
       int startIndex = ignoreFirstN;

       if (previousStop != null) {
         if (idLookup == null) {
           throw new UnsupportedOperationException("Only DimensionSelectors which support idLookup() are supported yet");
         }
         int lookupId = idLookup.lookupId(previousStop) + 1;
         if (lookupId < 0) {
           lookupId *= -1;
         }
         if (lookupId > ignoreFirstN + keepOnlyN) {
           startIndex = ignoreFirstN + keepOnlyN;
         } else {
           startIndex = Math.max(lookupId, startIndex);
         }
       }

       int endIndex = Math.min(ignoreFirstN + keepOnlyN, cardinality);

       if (ignoreAfterThreshold &&
           query.getDimensionsFilter() == null &&
           !storageAdapter.hasBuiltInFilters() &&
           query.getIntervals().stream().anyMatch(interval -> interval.contains(storageAdapter.getInterval()))) {
         endIndex = Math.min(endIndex, startIndex + query.getThreshold());
       }

       return Pair.of(startIndex, endIndex);
     }
   }

   public static TopNResultBuilder makeResultBuilder(TopNParams params, TopNQuery query)
   {
     final Comparator comparator = query.getTopNMetricSpec()
                                        .getComparator(query.getAggregatorSpecs(), query.getPostAggregatorSpecs());
     return query.getTopNMetricSpec().getResultBuilder(
         params.getCursor().getTime(),
         query.getDimensionSpec(),
         query.getThreshold(),
         comparator,
         query.getAggregatorSpecs(),
         query.getPostAggregatorSpecs()
     );
   }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF 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 org.apache.druid.query.topn;

	import com.google.common.annotations.VisibleForTesting;
	import org.apache.druid.java.util.common.IAE;
	import org.apache.druid.java.util.common.Pair;
	import org.apache.druid.query.aggregation.Aggregator;
	import org.apache.druid.query.aggregation.AggregatorFactory;
	import org.apache.druid.query.aggregation.BufferAggregator;
	import org.apache.druid.segment.Cursor;
	import org.apache.druid.segment.DimensionSelector;
	import org.apache.druid.segment.IdLookup;
	import org.apache.druid.segment.StorageAdapter;
	import org.apache.druid.segment.column.ColumnCapabilities;

	import javax.annotation.Nullable;
	import java.util.Arrays;
	import java.util.Comparator;
	import java.util.List;

	/**
	*/
	public abstract class BaseTopNAlgorithm<DimValSelector, DimValAggregateStore, Parameters extends TopNParams>
	implements TopNAlgorithm<DimValSelector, Parameters>
	{
	public static Aggregator[] makeAggregators(Cursor cursor, List<AggregatorFactory> aggregatorSpecs)
	{
	Aggregator[] aggregators = new Aggregator[aggregatorSpecs.size()];
	int aggregatorIndex = 0;
	for (AggregatorFactory spec : aggregatorSpecs) {
	aggregators[aggregatorIndex] = spec.factorize(cursor.getColumnSelectorFactory());
	++aggregatorIndex;
	}
	return aggregators;
	}

	protected static BufferAggregator[] makeBufferAggregators(Cursor cursor, List<AggregatorFactory> aggregatorSpecs)
	{
	BufferAggregator[] aggregators = new BufferAggregator[aggregatorSpecs.size()];
	int aggregatorIndex = 0;
	for (AggregatorFactory spec : aggregatorSpecs) {
	aggregators[aggregatorIndex] = spec.factorizeBuffered(cursor.getColumnSelectorFactory());
	++aggregatorIndex;
	}
	return aggregators;
	}

	protected final StorageAdapter storageAdapter;

	protected BaseTopNAlgorithm(StorageAdapter storageAdapter)
	{
	this.storageAdapter = storageAdapter;
	}

	@Override
	public void run(
	Parameters params,
	TopNResultBuilder resultBuilder,
	DimValSelector dimValSelector,
	@Nullable TopNQueryMetrics queryMetrics
	)
	{
	if (params.getCardinality() != TopNParams.CARDINALITY_UNKNOWN) {
	runWithCardinalityKnown(params, resultBuilder, dimValSelector, queryMetrics);
	} else {
	runWithCardinalityUnknown(params, resultBuilder, queryMetrics);
	}
	}

	private void runWithCardinalityKnown(
	Parameters params,
	TopNResultBuilder resultBuilder,
	DimValSelector dimValSelector,
	@Nullable TopNQueryMetrics queryMetrics
	)
	{
	if (queryMetrics != null) {
	queryMetrics.startRecordingScanTime();
	}
	boolean hasDimValSelector = (dimValSelector != null);

	int cardinality = params.getCardinality();
	int numProcessed = 0;
	long processedRows = 0;
	while (numProcessed < cardinality) {
	final int numToProcess;
	int maxNumToProcess = Math.min(params.getNumValuesPerPass(), cardinality - numProcessed);

	DimValSelector theDimValSelector;
	if (!hasDimValSelector) {
	numToProcess = maxNumToProcess;
	theDimValSelector = makeDimValSelector(params, numProcessed, numToProcess);
	} else {
	//skip invalid, calculate length to have enough valid value to process or hit the end.
	numToProcess = computeNewLength(dimValSelector, numProcessed, maxNumToProcess);
	theDimValSelector = updateDimValSelector(dimValSelector, numProcessed, numToProcess);
	}

	DimValAggregateStore aggregatesStore = makeDimValAggregateStore(params);

	processedRows = scanAndAggregate(params, theDimValSelector, aggregatesStore);

	updateResults(params, theDimValSelector, aggregatesStore, resultBuilder);

	closeAggregators(aggregatesStore);

	numProcessed += numToProcess;
	params.getCursor().reset();
	}
	if (queryMetrics != null) {
	queryMetrics.addProcessedRows(processedRows);
	queryMetrics.stopRecordingScanTime();
	}
	}

	/**
	* This function currently handles TopNs on long and float columns, which do not provide cardinality or an ID lookup.
	* When cardinality is unknown, process everything in one pass.
	* Existing implementations of makeDimValSelector() require cardinality as well, so the DimValSelector is not used.
	*
	* @param params TopN parameters from run()
	* @param resultBuilder Result builder from run()
	*/
	private void runWithCardinalityUnknown(
	Parameters params,
	TopNResultBuilder resultBuilder,
	@Nullable TopNQueryMetrics queryMetrics
	)
	{
	DimValAggregateStore aggregatesStore = makeDimValAggregateStore(params);
	if (queryMetrics != null) {
	queryMetrics.startRecordingScanTime();
	}
	long processedRows = scanAndAggregate(params, null, aggregatesStore);
	updateResults(params, null, aggregatesStore, resultBuilder);
	closeAggregators(aggregatesStore);
	params.getCursor().reset();
	if (queryMetrics != null) {
	queryMetrics.addProcessedRows(processedRows);
	queryMetrics.stopRecordingScanTime();
	}
	}

	protected abstract DimValSelector makeDimValSelector(Parameters params, int numProcessed, int numToProcess);

	/**
	* Skip invalid value, calculate length to have enough valid value to process or hit the end.
	*
	* @param dimValSelector the dim value selector which record value is valid or invalid.
	* @param numProcessed the start position to process
	* @param numToProcess the number of valid value to process
	*
	* @return the length between which have enough valid value to process or hit the end.
	*/
	protected int computeNewLength(DimValSelector dimValSelector, int numProcessed, int numToProcess)
	{
	return numToProcess;
	}

	protected abstract DimValSelector updateDimValSelector(
	DimValSelector dimValSelector,
	int numProcessed,
	int numToProcess
	);

	protected abstract DimValAggregateStore makeDimValAggregateStore(Parameters params);

	/**
	* Returns the number of processed rows (i. e. after postFilters are applied inside the cursor being processed).
	*/
	protected abstract long scanAndAggregate(
	Parameters params,
	DimValSelector dimValSelector,
	DimValAggregateStore dimValAggregateStore
	);

	protected abstract void updateResults(
	Parameters params,
	DimValSelector dimValSelector,
	DimValAggregateStore dimValAggregateStore,
	TopNResultBuilder resultBuilder
	);

	protected abstract void closeAggregators(
	DimValAggregateStore dimValAggregateStore
	);

	public static class AggregatorArrayProvider extends BaseArrayProvider<Aggregator[][]>
	{
	Aggregator[][] expansionAggs;
	int cardinality;

	public AggregatorArrayProvider(
	DimensionSelector dimSelector,
	TopNQuery query,
	int cardinality,
	StorageAdapter storageAdapter
	)
	{
	super(dimSelector, query, storageAdapter);

	this.expansionAggs = new Aggregator[cardinality][];
	this.cardinality = cardinality;
	}

	@Override
	public Aggregator[][] build()
	{
	Pair<Integer, Integer> startEnd = computeStartEnd(cardinality);

	Arrays.fill(expansionAggs, 0, startEnd.lhs, EMPTY_ARRAY);
	Arrays.fill(expansionAggs, startEnd.lhs, startEnd.rhs, null);
	Arrays.fill(expansionAggs, startEnd.rhs, expansionAggs.length, EMPTY_ARRAY);

	return expansionAggs;
	}
	}

	protected abstract static class BaseArrayProvider<T> implements TopNMetricSpecBuilder<T>
	{
	private volatile String previousStop;
	private volatile boolean ignoreAfterThreshold;
	private volatile int ignoreFirstN;
	private volatile int keepOnlyN;

	private final IdLookup idLookup;
	private final TopNQuery query;
	private final StorageAdapter storageAdapter;

	public BaseArrayProvider(
	DimensionSelector dimSelector,
	TopNQuery query,
	StorageAdapter storageAdapter
	)
	{
	this.idLookup = dimSelector.idLookup();
	this.query = query;
	this.storageAdapter = storageAdapter;

	previousStop = null;
	ignoreAfterThreshold = false;
	ignoreFirstN = 0;
	keepOnlyN = dimSelector.getValueCardinality();

	if (keepOnlyN < 0) {
	throw new IAE("Cannot operate on a dimension with no dictionary");
	}
	}

	@Override
	public void skipTo(String previousStop)
	{
	ColumnCapabilities capabilities = storageAdapter.getColumnCapabilities(query.getDimensionSpec().getDimension());
	if (capabilities != null && capabilities.areDictionaryValuesSorted().isTrue()) {
	this.previousStop = previousStop;
	}
	}

	@Override
	public void ignoreAfterThreshold()
	{
	ignoreAfterThreshold = true;
	}

	@Override
	public void ignoreFirstN(int n)
	{
	ignoreFirstN = n;
	}

	@Override
	public void keepOnlyN(int n)
	{
	keepOnlyN = n;
	}

	@VisibleForTesting
	public Pair<Integer, Integer> computeStartEnd(int cardinality)
	{
	int startIndex = ignoreFirstN;

	if (previousStop != null) {
	if (idLookup == null) {
	throw new UnsupportedOperationException("Only DimensionSelectors which support idLookup() are supported yet");
	}
	int lookupId = idLookup.lookupId(previousStop) + 1;
	if (lookupId < 0) {
	lookupId *= -1;
	}
	if (lookupId > ignoreFirstN + keepOnlyN) {
	startIndex = ignoreFirstN + keepOnlyN;
	} else {
	startIndex = Math.max(lookupId, startIndex);
	}
	}

	int endIndex = Math.min(ignoreFirstN + keepOnlyN, cardinality);

	if (ignoreAfterThreshold &&
	query.getDimensionsFilter() == null &&
	!storageAdapter.hasBuiltInFilters() &&
	query.getIntervals().stream().anyMatch(interval -> interval.contains(storageAdapter.getInterval()))) {
	endIndex = Math.min(endIndex, startIndex + query.getThreshold());
	}

	return Pair.of(startIndex, endIndex);
	}
	}

	public static TopNResultBuilder makeResultBuilder(TopNParams params, TopNQuery query)
	{
	final Comparator comparator = query.getTopNMetricSpec()
	.getComparator(query.getAggregatorSpecs(), query.getPostAggregatorSpecs());
	return query.getTopNMetricSpec().getResultBuilder(
	params.getCursor().getTime(),
	query.getDimensionSpec(),
	query.getThreshold(),
	comparator,
	query.getAggregatorSpecs(),
	query.getPostAggregatorSpecs()
	);
	}
	}