use of io.crate.operation.collect.InputCollectExpression in project crate by crate.
the class IndexWriterProjectorUnitTest method testNullPKValue.
@Test
public void testNullPKValue() throws Throwable {
InputCollectExpression sourceInput = new InputCollectExpression(0);
List<CollectExpression<Row, ?>> collectExpressions = Collections.<CollectExpression<Row, ?>>singletonList(sourceInput);
final IndexWriterProjector indexWriter = new IndexWriterProjector(clusterService, TestingHelpers.getFunctions(), new IndexNameExpressionResolver(Settings.EMPTY), Settings.EMPTY, mock(TransportBulkCreateIndicesAction.class), mock(BulkRequestExecutor.class), () -> "foo", mock(BulkRetryCoordinatorPool.class, Answers.RETURNS_DEEP_STUBS.get()), rawSourceReference, ImmutableList.of(ID_IDENT), Arrays.<Symbol>asList(new InputColumn(1)), null, null, sourceInput, collectExpressions, 20, null, null, false, false, UUID.randomUUID());
RowN rowN = new RowN(new Object[] { new BytesRef("{\"y\": \"x\"}"), null });
BatchIterator batchIterator = RowsBatchIterator.newInstance(Collections.singletonList(rowN), rowN.numColumns());
batchIterator = indexWriter.apply(batchIterator);
TestingBatchConsumer testingBatchConsumer = new TestingBatchConsumer();
testingBatchConsumer.accept(batchIterator, null);
expectedException.expect(IllegalArgumentException.class);
expectedException.expectMessage("A primary key value must not be NULL");
testingBatchConsumer.getResult();
}
use of io.crate.operation.collect.InputCollectExpression in project crate by crate.
the class GroupingIntegerCollectorBenchmark method createGroupBySumCollector.
private GroupingCollector createGroupBySumCollector(Functions functions) {
InputCollectExpression keyInput = new InputCollectExpression(0);
List<Input<?>> keyInputs = Arrays.<Input<?>>asList(keyInput);
CollectExpression[] collectExpressions = new CollectExpression[] { keyInput };
FunctionIdent functionIdent = new FunctionIdent(SumAggregation.NAME, Arrays.asList(DataTypes.INTEGER));
FunctionInfo functionInfo = new FunctionInfo(functionIdent, DataTypes.INTEGER, FunctionInfo.Type.AGGREGATE);
AggregationFunction sumAgg = (AggregationFunction) functions.get(functionIdent);
Aggregation aggregation = Aggregation.finalAggregation(functionInfo, Arrays.asList(new InputColumn(0)), Aggregation.Step.ITER);
Aggregator[] aggregators = new Aggregator[] { new Aggregator(RAM_ACCOUNTING_CONTEXT, aggregation, sumAgg, new Input[] { keyInput }) };
return GroupingCollector.singleKey(collectExpressions, aggregators, RAM_ACCOUNTING_CONTEXT, keyInputs.get(0), DataTypes.INTEGER);
}
use of io.crate.operation.collect.InputCollectExpression in project crate by crate.
the class AggregationTest method executeAggregation.
public Object[][] executeAggregation(String name, DataType dataType, Object[][] data, List<DataType> argumentTypes) throws Exception {
FunctionIdent fi;
InputCollectExpression[] inputs;
if (dataType != null) {
fi = new FunctionIdent(name, argumentTypes);
inputs = new InputCollectExpression[argumentTypes.size()];
for (int i = 0; i < argumentTypes.size(); i++) {
inputs[i] = new InputCollectExpression(i);
}
} else {
fi = new FunctionIdent(name, ImmutableList.<DataType>of());
inputs = new InputCollectExpression[0];
}
AggregationFunction impl = (AggregationFunction) functions.get(fi);
Object state = impl.newState(ramAccountingContext);
ArrayBucket bucket = new ArrayBucket(data);
for (Row row : bucket) {
for (InputCollectExpression i : inputs) {
i.setNextRow(row);
}
state = impl.iterate(ramAccountingContext, state, inputs);
}
state = impl.terminatePartial(ramAccountingContext, state);
return new Object[][] { { state } };
}
use of io.crate.operation.collect.InputCollectExpression in project crate by crate.
the class GroupingBytesRefCollectorBenchmark method createGroupByMinBytesRefCollector.
private GroupingCollector createGroupByMinBytesRefCollector(Functions functions) {
InputCollectExpression keyInput = new InputCollectExpression(0);
List<Input<?>> keyInputs = Arrays.<Input<?>>asList(keyInput);
CollectExpression[] collectExpressions = new CollectExpression[] { keyInput };
FunctionIdent minBytesRefFuncIdent = new FunctionIdent(MinimumAggregation.NAME, Arrays.asList(DataTypes.STRING));
FunctionInfo minBytesRefFuncInfo = new FunctionInfo(minBytesRefFuncIdent, DataTypes.INTEGER, FunctionInfo.Type.AGGREGATE);
AggregationFunction minAgg = (AggregationFunction) functions.get(minBytesRefFuncIdent);
Aggregation aggregation = Aggregation.finalAggregation(minBytesRefFuncInfo, Arrays.asList(new InputColumn(0)), Aggregation.Step.ITER);
Aggregator[] aggregators = new Aggregator[] { new Aggregator(RAM_ACCOUNTING_CONTEXT, aggregation, minAgg, new Input[] { keyInput }) };
return GroupingCollector.singleKey(collectExpressions, aggregators, RAM_ACCOUNTING_CONTEXT, keyInputs.get(0), DataTypes.STRING);
}
use of io.crate.operation.collect.InputCollectExpression in project crate by crate.
the class IndexWriterCountBatchIteratorTest method testIndexWriterIterator.
@Test
public void testIndexWriterIterator() throws Exception {
execute("create table bulk_import (id int primary key) with (number_of_replicas=0)");
ensureGreen();
Supplier<BatchIterator> sourceSupplier = () -> RowsBatchIterator.newInstance(RowGenerator.fromSingleColValues(() -> IntStream.range(0, 10).mapToObj(i -> new BytesRef("{\"id\": " + i + "}")).iterator()), 1);
Supplier<String> indexNameResolver = IndexNameResolver.forTable(new TableIdent(null, "bulk_import"));
Input<?> sourceInput = new InputCollectExpression(0);
List<CollectExpression<Row, ?>> collectExpressions = Collections.singletonList((InputCollectExpression) sourceInput);
RowShardResolver rowShardResolver = getRowShardResolver();
BulkShardProcessor bulkShardProcessor = getBulkShardProcessor();
Supplier<ShardUpsertRequest.Item> updateItemSupplier = () -> new ShardUpsertRequest.Item(rowShardResolver.id(), null, new Object[] { sourceInput.value() }, null);
List<Object[]> expectedResult = Collections.singletonList(new Object[] { 10L });
BatchIteratorTester tester = new BatchIteratorTester(() -> IndexWriterCountBatchIterator.newIndexInstance(sourceSupplier.get(), indexNameResolver, collectExpressions, rowShardResolver, bulkShardProcessor, updateItemSupplier));
tester.verifyResultAndEdgeCaseBehaviour(expectedResult);
}
Aggregations