Examples with VectorCursor org.apache.druid.segment.vector.VectorCursor used on opensource projects

Example 1 with VectorCursor

use of org.apache.druid.segment.vector.VectorCursor in project druid by druid-io.

the class BaseFilterTest method selectCountUsingVectorizedFilteredAggregator.

private long selectCountUsingVectorizedFilteredAggregator(final DimFilter dimFilter) {
    Preconditions.checkState(makeFilter(dimFilter).canVectorizeMatcher(adapter), "Cannot vectorize filter: %s", dimFilter);
    try (final VectorCursor cursor = makeVectorCursor(null)) {
        final FilteredAggregatorFactory aggregatorFactory = new FilteredAggregatorFactory(new CountAggregatorFactory("count"), maybeOptimize(dimFilter));
        final VectorAggregator aggregator = aggregatorFactory.factorizeVector(cursor.getColumnSelectorFactory());
        final ByteBuffer buf = ByteBuffer.allocate(aggregatorFactory.getMaxIntermediateSizeWithNulls() * 2);
        // Use two slots: one for each form of aggregate.
        aggregator.init(buf, 0);
        aggregator.init(buf, aggregatorFactory.getMaxIntermediateSizeWithNulls());
        for (; !cursor.isDone(); cursor.advance()) {
            aggregator.aggregate(buf, 0, 0, cursor.getCurrentVectorSize());
            final int[] positions = new int[cursor.getCurrentVectorSize()];
            Arrays.fill(positions, aggregatorFactory.getMaxIntermediateSizeWithNulls());
            final int[] allRows = new int[cursor.getCurrentVectorSize()];
            for (int i = 0; i < allRows.length; i++) {
                allRows[i] = i;
            }
            aggregator.aggregate(buf, cursor.getCurrentVectorSize(), positions, allRows, 0);
        }
        final long val1 = (long) aggregator.get(buf, 0);
        final long val2 = (long) aggregator.get(buf, aggregatorFactory.getMaxIntermediateSizeWithNulls());
        if (val1 != val2) {
            throw new ISE("Oh no, val1[%d] != val2[%d]", val1, val2);
        }
        return val1;
    }
}

Example 2 with VectorCursor

use of org.apache.druid.segment.vector.VectorCursor in project druid by druid-io.

the class VectorGroupByEngineIteratorTest method testCreateOneGrouperAndCloseItWhenClose.

@Test
public void testCreateOneGrouperAndCloseItWhenClose() throws IOException {
    final Interval interval = TestIndex.DATA_INTERVAL;
    final AggregatorFactory factory = new DoubleSumAggregatorFactory("index", "index");
    final GroupByQuery query = GroupByQuery.builder().setDataSource(QueryRunnerTestHelper.DATA_SOURCE).setGranularity(QueryRunnerTestHelper.DAY_GRAN).setInterval(interval).setDimensions(new DefaultDimensionSpec("market", null, null)).setAggregatorSpecs(factory).build();
    final StorageAdapter storageAdapter = new QueryableIndexStorageAdapter(TestIndex.getMMappedTestIndex());
    final ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[4096]);
    final VectorCursor cursor = storageAdapter.makeVectorCursor(Filters.toFilter(query.getDimFilter()), interval, query.getVirtualColumns(), false, QueryContexts.getVectorSize(query), null);
    final List<GroupByVectorColumnSelector> dimensions = query.getDimensions().stream().map(dimensionSpec -> ColumnProcessors.makeVectorProcessor(dimensionSpec, GroupByVectorColumnProcessorFactory.instance(), cursor.getColumnSelectorFactory())).collect(Collectors.toList());
    final MutableObject<VectorGrouper> grouperCaptor = new MutableObject<>();
    final VectorGroupByEngineIterator iterator = new VectorGroupByEngineIterator(query, new GroupByQueryConfig(), storageAdapter, cursor, interval, dimensions, byteBuffer, null) {

        @Override
        VectorGrouper makeGrouper() {
            grouperCaptor.setValue(Mockito.spy(super.makeGrouper()));
            return grouperCaptor.getValue();
        }
    };
    iterator.close();
    Mockito.verify(grouperCaptor.getValue()).close();
}

Example 3 with VectorCursor

use of org.apache.druid.segment.vector.VectorCursor in project druid by druid-io.

the class ExpressionVectorSelectorsTest method sanityTestVectorizedExpressionSelectors.

public static void sanityTestVectorizedExpressionSelectors(String expression, @Nullable ExpressionType outputType, QueryableIndex index, Closer closer, int rowsPerSegment) {
    final List<Object> results = new ArrayList<>(rowsPerSegment);
    final VirtualColumns virtualColumns = VirtualColumns.create(ImmutableList.of(new ExpressionVirtualColumn("v", expression, ExpressionType.toColumnType(outputType), TestExprMacroTable.INSTANCE)));
    final QueryableIndexStorageAdapter storageAdapter = new QueryableIndexStorageAdapter(index);
    VectorCursor cursor = storageAdapter.makeVectorCursor(null, index.getDataInterval(), virtualColumns, false, 512, null);
    ColumnCapabilities capabilities = virtualColumns.getColumnCapabilities(storageAdapter, "v");
    int rowCount = 0;
    if (capabilities.isDictionaryEncoded().isTrue()) {
        SingleValueDimensionVectorSelector selector = cursor.getColumnSelectorFactory().makeSingleValueDimensionSelector(DefaultDimensionSpec.of("v"));
        while (!cursor.isDone()) {
            int[] row = selector.getRowVector();
            for (int i = 0; i < selector.getCurrentVectorSize(); i++, rowCount++) {
                results.add(selector.lookupName(row[i]));
            }
            cursor.advance();
        }
    } else {
        VectorValueSelector selector = null;
        VectorObjectSelector objectSelector = null;
        if (outputType != null && outputType.isNumeric()) {
            selector = cursor.getColumnSelectorFactory().makeValueSelector("v");
        } else {
            objectSelector = cursor.getColumnSelectorFactory().makeObjectSelector("v");
        }
        while (!cursor.isDone()) {
            boolean[] nulls;
            switch(outputType.getType()) {
                case LONG:
                    nulls = selector.getNullVector();
                    long[] longs = selector.getLongVector();
                    for (int i = 0; i < selector.getCurrentVectorSize(); i++, rowCount++) {
                        results.add(nulls != null && nulls[i] ? null : longs[i]);
                    }
                    break;
                case DOUBLE:
                    // special case to test floats just to get coverage on getFloatVector
                    if ("float2".equals(expression)) {
                        nulls = selector.getNullVector();
                        float[] floats = selector.getFloatVector();
                        for (int i = 0; i < selector.getCurrentVectorSize(); i++, rowCount++) {
                            results.add(nulls != null && nulls[i] ? null : (double) floats[i]);
                        }
                    } else {
                        nulls = selector.getNullVector();
                        double[] doubles = selector.getDoubleVector();
                        for (int i = 0; i < selector.getCurrentVectorSize(); i++, rowCount++) {
                            results.add(nulls != null && nulls[i] ? null : doubles[i]);
                        }
                    }
                    break;
                case STRING:
                    Object[] objects = objectSelector.getObjectVector();
                    for (int i = 0; i < objectSelector.getCurrentVectorSize(); i++, rowCount++) {
                        results.add(objects[i]);
                    }
                    break;
            }
            cursor.advance();
        }
    }
    closer.register(cursor);
    Sequence<Cursor> cursors = new QueryableIndexStorageAdapter(index).makeCursors(null, index.getDataInterval(), virtualColumns, Granularities.ALL, false, null);
    int rowCountCursor = cursors.map(nonVectorized -> {
        final ColumnValueSelector nonSelector = nonVectorized.getColumnSelectorFactory().makeColumnValueSelector("v");
        int rows = 0;
        while (!nonVectorized.isDone()) {
            Assert.assertEquals(StringUtils.format("Failed at row %s", rows), nonSelector.getObject(), results.get(rows));
            rows++;
            nonVectorized.advance();
        }
        return rows;
    }).accumulate(0, (acc, in) -> acc + in);
    Assert.assertTrue(rowCountCursor > 0);
    Assert.assertEquals(rowCountCursor, rowCount);
}

Example 4 with VectorCursor

use of org.apache.druid.segment.vector.VectorCursor in project druid by druid-io.

the class TimeseriesQueryEngine method processVectorized.

private Sequence<Result<TimeseriesResultValue>> processVectorized(final TimeseriesQuery query, final StorageAdapter adapter, @Nullable final Filter filter, final Interval queryInterval, final Granularity gran, final boolean descending) {
    final boolean skipEmptyBuckets = query.isSkipEmptyBuckets();
    final List<AggregatorFactory> aggregatorSpecs = query.getAggregatorSpecs();
    final VectorCursor cursor = adapter.makeVectorCursor(filter, queryInterval, query.getVirtualColumns(), descending, QueryContexts.getVectorSize(query), null);
    if (cursor == null) {
        return Sequences.empty();
    }
    final Closer closer = Closer.create();
    closer.register(cursor);
    try {
        final VectorCursorGranularizer granularizer = VectorCursorGranularizer.create(adapter, cursor, gran, queryInterval);
        if (granularizer == null) {
            return Sequences.empty();
        }
        final VectorColumnSelectorFactory columnSelectorFactory = cursor.getColumnSelectorFactory();
        final AggregatorAdapters aggregators = closer.register(AggregatorAdapters.factorizeVector(columnSelectorFactory, query.getAggregatorSpecs()));
        final ResourceHolder<ByteBuffer> bufferHolder = closer.register(bufferPool.take());
        final ByteBuffer buffer = bufferHolder.get();
        if (aggregators.spaceNeeded() > buffer.remaining()) {
            throw new ISE("Not enough space for aggregators, needed [%,d] bytes but have only [%,d].", aggregators.spaceNeeded(), buffer.remaining());
        }
        return Sequences.withBaggage(Sequences.simple(granularizer.getBucketIterable()).map(bucketInterval -> {
            // Whether or not the current bucket is empty
            boolean emptyBucket = true;
            while (!cursor.isDone()) {
                granularizer.setCurrentOffsets(bucketInterval);
                if (granularizer.getEndOffset() > granularizer.getStartOffset()) {
                    if (emptyBucket) {
                        aggregators.init(buffer, 0);
                    }
                    aggregators.aggregateVector(buffer, 0, granularizer.getStartOffset(), granularizer.getEndOffset());
                    emptyBucket = false;
                }
                if (!granularizer.advanceCursorWithinBucket()) {
                    break;
                }
            }
            if (emptyBucket && skipEmptyBuckets) {
                // Return null, will get filtered out later by the Objects::nonNull filter.
                return null;
            }
            final TimeseriesResultBuilder bob = new TimeseriesResultBuilder(gran.toDateTime(bucketInterval.getStartMillis()));
            if (emptyBucket) {
                aggregators.init(buffer, 0);
            }
            for (int i = 0; i < aggregatorSpecs.size(); i++) {
                bob.addMetric(aggregatorSpecs.get(i).getName(), aggregators.get(buffer, 0, i));
            }
            return bob.build();
        }).filter(Objects::nonNull), closer);
    } catch (Throwable t1) {
        try {
            closer.close();
        } catch (Throwable t2) {
            t1.addSuppressed(t2);
        }
        throw t1;
    }
}

Example 5 with VectorCursor

use of org.apache.druid.segment.vector.VectorCursor in project druid by druid-io.

the class ToObjectVectorColumnProcessorFactoryTest method testRead.

@Test
public void testRead() {
    try (final VectorCursor cursor = makeCursor()) {
        final Supplier<Object[]> qualitySupplier = ColumnProcessors.makeVectorProcessor("quality", ToObjectVectorColumnProcessorFactory.INSTANCE, cursor.getColumnSelectorFactory());
        final Supplier<Object[]> qualityLongSupplier = ColumnProcessors.makeVectorProcessor("qualityLong", ToObjectVectorColumnProcessorFactory.INSTANCE, cursor.getColumnSelectorFactory());
        final Supplier<Object[]> qualityFloatSupplier = ColumnProcessors.makeVectorProcessor("qualityFloat", ToObjectVectorColumnProcessorFactory.INSTANCE, cursor.getColumnSelectorFactory());
        final Supplier<Object[]> qualityDoubleSupplier = ColumnProcessors.makeVectorProcessor("qualityDouble", ToObjectVectorColumnProcessorFactory.INSTANCE, cursor.getColumnSelectorFactory());
        final Supplier<Object[]> placementSupplier = ColumnProcessors.makeVectorProcessor("placement", ToObjectVectorColumnProcessorFactory.INSTANCE, cursor.getColumnSelectorFactory());
        final Supplier<Object[]> qualityUniquesSupplier = ColumnProcessors.makeVectorProcessor("quality_uniques", ToObjectVectorColumnProcessorFactory.INSTANCE, cursor.getColumnSelectorFactory());
    }
}