Example #1
0
  private static GroupByQueryRunnerFactory getGroupByQueryRunnerFactory() {
    ObjectMapper mapper = new DefaultObjectMapper();
    GroupByQueryConfig config = new GroupByQueryConfig();
    config.setMaxIntermediateRows(10000);

    Supplier<GroupByQueryConfig> configSupplier = Suppliers.ofInstance(config);
    GroupByQueryEngine engine = new GroupByQueryEngine(configSupplier, Utils.getBufferPool());

    GroupByQueryRunnerFactory factory =
        new GroupByQueryRunnerFactory(
            engine,
            Utils.NOOP_QUERYWATCHER,
            configSupplier,
            new GroupByQueryQueryToolChest(
                configSupplier,
                mapper,
                engine,
                Utils.getBufferPool(),
                Utils.NoopIntervalChunkingQueryRunnerDecorator()),
            Utils.getBufferPool());
    return factory;
  }
  public static Pair<Grouper<RowBasedKey>, Accumulator<Grouper<RowBasedKey>, Row>>
      createGrouperAccumulatorPair(
          final GroupByQuery query,
          final boolean isInputRaw,
          final GroupByQueryConfig config,
          final ByteBuffer buffer,
          final int concurrencyHint,
          final LimitedTemporaryStorage temporaryStorage,
          final ObjectMapper spillMapper,
          final AggregatorFactory[] aggregatorFactories) {
    // concurrencyHint >= 1 for concurrent groupers, -1 for single-threaded
    Preconditions.checkArgument(
        concurrencyHint >= 1 || concurrencyHint == -1, "invalid concurrencyHint");

    final GroupByQueryConfig querySpecificConfig = config.withOverrides(query);
    final DateTime fudgeTimestamp = GroupByStrategyV2.getUniversalTimestamp(query);
    final Grouper.KeySerdeFactory<RowBasedKey> keySerdeFactory =
        new RowBasedKeySerdeFactory(
            fudgeTimestamp,
            query.getDimensions().size(),
            querySpecificConfig.getMaxMergingDictionarySize()
                / (concurrencyHint == -1 ? 1 : concurrencyHint));
    final RowBasedColumnSelectorFactory columnSelectorFactory = new RowBasedColumnSelectorFactory();
    final Grouper<RowBasedKey> grouper;
    if (concurrencyHint == -1) {
      grouper =
          new SpillingGrouper<>(
              buffer,
              keySerdeFactory,
              columnSelectorFactory,
              aggregatorFactories,
              querySpecificConfig.getBufferGrouperMaxSize(),
              querySpecificConfig.getBufferGrouperMaxLoadFactor(),
              querySpecificConfig.getBufferGrouperInitialBuckets(),
              temporaryStorage,
              spillMapper,
              true);
    } else {
      grouper =
          new ConcurrentGrouper<>(
              buffer,
              keySerdeFactory,
              columnSelectorFactory,
              aggregatorFactories,
              querySpecificConfig.getBufferGrouperMaxSize(),
              querySpecificConfig.getBufferGrouperMaxLoadFactor(),
              querySpecificConfig.getBufferGrouperInitialBuckets(),
              temporaryStorage,
              spillMapper,
              concurrencyHint);
    }

    final DimensionSelector[] dimensionSelectors;
    if (isInputRaw) {
      dimensionSelectors = new DimensionSelector[query.getDimensions().size()];
      for (int i = 0; i < dimensionSelectors.length; i++) {
        dimensionSelectors[i] =
            columnSelectorFactory.makeDimensionSelector(query.getDimensions().get(i));
      }
    } else {
      dimensionSelectors = null;
    }

    final Accumulator<Grouper<RowBasedKey>, Row> accumulator =
        new Accumulator<Grouper<RowBasedKey>, Row>() {
          @Override
          public Grouper<RowBasedKey> accumulate(
              final Grouper<RowBasedKey> theGrouper, final Row row) {
            if (Thread.interrupted()) {
              throw new QueryInterruptedException(new InterruptedException());
            }

            if (theGrouper == null) {
              // Pass-through null returns without doing more work.
              return null;
            }

            long timestamp = row.getTimestampFromEpoch();
            if (isInputRaw) {
              if (query.getGranularity() instanceof AllGranularity) {
                timestamp = query.getIntervals().get(0).getStartMillis();
              } else {
                timestamp = query.getGranularity().truncate(timestamp);
              }
            }

            columnSelectorFactory.setRow(row);
            final String[] dimensions = new String[query.getDimensions().size()];
            for (int i = 0; i < dimensions.length; i++) {
              final String value;
              if (isInputRaw) {
                IndexedInts index = dimensionSelectors[i].getRow();
                value = index.size() == 0 ? "" : dimensionSelectors[i].lookupName(index.get(0));
              } else {
                value = (String) row.getRaw(query.getDimensions().get(i).getOutputName());
              }
              dimensions[i] = Strings.nullToEmpty(value);
            }

            final boolean didAggregate =
                theGrouper.aggregate(new RowBasedKey(timestamp, dimensions));
            if (!didAggregate) {
              // null return means grouping resources were exhausted.
              return null;
            }
            columnSelectorFactory.setRow(null);

            return theGrouper;
          }
        };

    return new Pair<>(grouper, accumulator);
  }