diff --git a/core/trino-main/src/main/java/io/trino/connector/system/RuleStatsSystemTable.java b/core/trino-main/src/main/java/io/trino/connector/system/RuleStatsSystemTable.java index 73aef6f8bf5c..92c5a5e99e97 100644 --- a/core/trino-main/src/main/java/io/trino/connector/system/RuleStatsSystemTable.java +++ b/core/trino-main/src/main/java/io/trino/connector/system/RuleStatsSystemTable.java @@ -18,6 +18,7 @@ import io.trino.spi.Page; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; @@ -98,12 +99,13 @@ public ConnectorPageSource pageSource(ConnectorTransactionHandle transactionHand BIGINT.writeLong(blockBuilders.get("failures"), stats.getFailures()); DOUBLE.writeDouble(blockBuilders.get("average_time"), stats.getTime().getAvg()); - BlockBuilder mapWriter = blockBuilders.get("time_distribution_percentiles").beginBlockEntry(); - for (Map.Entry percentile : stats.getTime().getPercentiles().entrySet()) { - DOUBLE.writeDouble(mapWriter, percentile.getKey()); - DOUBLE.writeDouble(mapWriter, percentile.getValue()); - } - blockBuilders.get("time_distribution_percentiles").closeEntry(); + MapBlockBuilder blockBuilder = (MapBlockBuilder) blockBuilders.get("time_distribution_percentiles"); + blockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + stats.getTime().getPercentiles().forEach((key, value) -> { + DOUBLE.writeDouble(keyBuilder, key); + DOUBLE.writeDouble(valueBuilder, value); + }); + }); } Block[] blocks = ruleStatsTable.getColumns().stream() diff --git a/core/trino-main/src/main/java/io/trino/operator/HashAggregationOperator.java b/core/trino-main/src/main/java/io/trino/operator/HashAggregationOperator.java index f9ce939dd89a..6a2f969281a9 100644 --- a/core/trino-main/src/main/java/io/trino/operator/HashAggregationOperator.java +++ b/core/trino-main/src/main/java/io/trino/operator/HashAggregationOperator.java @@ -46,6 +46,7 @@ import static com.google.common.base.Preconditions.checkState; import static io.airlift.units.DataSize.Unit.MEGABYTE; import static io.trino.operator.aggregation.builder.InMemoryHashAggregationBuilder.toTypes; +import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.sql.planner.optimizations.HashGenerationOptimizer.INITIAL_HASH_VALUE; import static io.trino.type.TypeUtils.NULL_HASH_CODE; import static java.util.Objects.requireNonNull; @@ -572,7 +573,7 @@ private Page getGlobalAggregationOutput() while (channel < groupByTypes.size()) { if (channel == groupIdChannel.orElseThrow()) { - output.getBlockBuilder(channel).writeLong(groupId); + BIGINT.writeLong(output.getBlockBuilder(channel), groupId); } else { output.getBlockBuilder(channel).appendNull(); @@ -582,7 +583,7 @@ private Page getGlobalAggregationOutput() if (hashChannel.isPresent()) { long hashValue = calculateDefaultOutputHash(groupByTypes, groupIdChannel.orElseThrow(), groupId); - output.getBlockBuilder(channel).writeLong(hashValue); + BIGINT.writeLong(output.getBlockBuilder(channel), hashValue); channel++; } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/AccumulatorCompiler.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/AccumulatorCompiler.java index 130a361b861d..daac29c71fed 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/AccumulatorCompiler.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/AccumulatorCompiler.java @@ -31,6 +31,8 @@ import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.ColumnarRow; +import io.trino.spi.block.RowBlockBuilder; +import io.trino.spi.block.RowValueBuilder; import io.trino.spi.function.AccumulatorState; import io.trino.spi.function.AccumulatorStateFactory; import io.trino.spi.function.AccumulatorStateSerializer; @@ -73,6 +75,7 @@ import static io.trino.sql.gen.Bootstrap.BOOTSTRAP_METHOD; import static io.trino.sql.gen.BytecodeUtils.invoke; import static io.trino.sql.gen.BytecodeUtils.loadConstant; +import static io.trino.sql.gen.LambdaMetafactoryGenerator.generateMetafactory; import static io.trino.util.CompilerUtils.defineClass; import static io.trino.util.CompilerUtils.makeClassName; import static java.lang.String.format; @@ -776,18 +779,13 @@ private static void generateGroupedEvaluateIntermediate(ClassDefinition definiti .ret(); } else { - Variable rowBuilder = method.getScope().declareVariable(BlockBuilder.class, "rowBuilder"); - body.append(rowBuilder.set(out.invoke("beginBlockEntry", BlockBuilder.class))); - for (StateFieldAndDescriptor stateFieldAndDescriptor : stateFieldAndDescriptors) { - BytecodeExpression stateSerializer = thisVariable.getField(stateFieldAndDescriptor.getStateSerializerField()); BytecodeExpression state = thisVariable.getField(stateFieldAndDescriptor.getStateField()); - - body.append(state.invoke("setGroupId", void.class, groupId.cast(long.class))) - .append(stateSerializer.invoke("serialize", void.class, state.cast(AccumulatorState.class), rowBuilder)); + body.append(state.invoke("setGroupId", void.class, groupId.cast(long.class))); } - body.append(out.invoke("closeEntry", BlockBuilder.class).pop()) - .ret(); + + generateSerializeState(definition, stateFieldAndDescriptors, out, thisVariable, body); + body.ret(); } } @@ -818,17 +816,36 @@ private static void generateEvaluateIntermediate(ClassDefinition definition, Lis .ret(); } else { - Variable rowBuilder = method.getScope().declareVariable(BlockBuilder.class, "rowBuilder"); - body.append(rowBuilder.set(out.invoke("beginBlockEntry", BlockBuilder.class))); + generateSerializeState(definition, stateFieldAndDescriptors, out, thisVariable, body); + body.ret(); + } + } - for (StateFieldAndDescriptor stateFieldAndDescriptor : stateFieldAndDescriptors) { - BytecodeExpression stateSerializer = thisVariable.getField(stateFieldAndDescriptor.getStateSerializerField()); - BytecodeExpression state = thisVariable.getField(stateFieldAndDescriptor.getStateField()); - body.append(stateSerializer.invoke("serialize", void.class, state.cast(AccumulatorState.class), rowBuilder)); - } - body.append(out.invoke("closeEntry", BlockBuilder.class).pop()) - .ret(); + private static void generateSerializeState(ClassDefinition definition, List stateFieldAndDescriptors, Parameter out, Variable thisVariable, BytecodeBlock body) + { + MethodDefinition serializeState = generateSerializeStateMethod(definition, stateFieldAndDescriptors); + + BytecodeExpression rowEntryBuilder = generateMetafactory(RowValueBuilder.class, serializeState, ImmutableList.of(thisVariable)); + body.append(out.cast(RowBlockBuilder.class).invoke("buildEntry", void.class, rowEntryBuilder)); + } + + private static MethodDefinition generateSerializeStateMethod(ClassDefinition definition, List stateFieldAndDescriptors) + { + Parameter fieldBuilders = arg("fieldBuilders", type(List.class, BlockBuilder.class)); + MethodDefinition method = definition.declareMethod(a(PRIVATE), "serializeState", type(void.class), fieldBuilders); + + Variable thisVariable = method.getThis(); + BytecodeBlock body = method.getBody(); + + for (int i = 0; i < stateFieldAndDescriptors.size(); i++) { + StateFieldAndDescriptor stateFieldAndDescriptor = stateFieldAndDescriptors.get(i); + BytecodeExpression stateSerializer = thisVariable.getField(stateFieldAndDescriptor.getStateSerializerField()); + BytecodeExpression state = thisVariable.getField(stateFieldAndDescriptor.getStateField()); + BytecodeExpression fieldBuilder = fieldBuilders.invoke("get", Object.class, constantInt(i)).cast(BlockBuilder.class); + body.append(stateSerializer.invoke("serialize", void.class, state.cast(AccumulatorState.class), fieldBuilder)); } + body.ret(); + return method; } private static void generateGroupedEvaluateFinal( diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateDoublePercentileArrayAggregations.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateDoublePercentileArrayAggregations.java index 8d48116a94a7..b9384fef2f5d 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateDoublePercentileArrayAggregations.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateDoublePercentileArrayAggregations.java @@ -17,6 +17,7 @@ import com.google.common.primitives.Doubles; import io.airlift.stats.TDigest; import io.trino.operator.aggregation.state.TDigestAndPercentileArrayState; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.AggregationFunction; @@ -95,14 +96,12 @@ public static void output(@AggregationState TDigestAndPercentileArrayState state return; } - BlockBuilder blockBuilder = out.beginBlockEntry(); - List valuesAtPercentiles = valuesAtPercentiles(digest, percentiles); - for (double value : valuesAtPercentiles) { - DOUBLE.writeDouble(blockBuilder, value); - } - - out.closeEntry(); + ((ArrayBlockBuilder) out).buildEntry(elementBuilder -> { + for (double value : valuesAtPercentiles) { + DOUBLE.writeDouble(elementBuilder, value); + } + }); } public static List valuesAtPercentiles(TDigest digest, List percentiles) diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateLongPercentileArrayAggregations.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateLongPercentileArrayAggregations.java index 3a13d7867a5d..1aa649e1e870 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateLongPercentileArrayAggregations.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateLongPercentileArrayAggregations.java @@ -15,6 +15,7 @@ import io.airlift.stats.TDigest; import io.trino.operator.aggregation.state.TDigestAndPercentileArrayState; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.AggregationFunction; @@ -65,13 +66,11 @@ public static void output(@AggregationState TDigestAndPercentileArrayState state return; } - BlockBuilder blockBuilder = out.beginBlockEntry(); - - List valuesAtPercentiles = valuesAtPercentiles(digest, percentiles); - for (double value : valuesAtPercentiles) { - BIGINT.writeLong(blockBuilder, Math.round(value)); - } - - out.closeEntry(); + ((ArrayBlockBuilder) out).buildEntry(elementBuilder -> { + List valuesAtPercentiles = valuesAtPercentiles(digest, percentiles); + for (double value : valuesAtPercentiles) { + BIGINT.writeLong(elementBuilder, Math.round(value)); + } + }); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateRealPercentileArrayAggregations.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateRealPercentileArrayAggregations.java index d3577b86fc6a..6679e8feff0a 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateRealPercentileArrayAggregations.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/ApproximateRealPercentileArrayAggregations.java @@ -15,6 +15,7 @@ import io.airlift.stats.TDigest; import io.trino.operator.aggregation.state.TDigestAndPercentileArrayState; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.AggregationFunction; @@ -66,13 +67,11 @@ public static void output(@AggregationState TDigestAndPercentileArrayState state return; } - BlockBuilder blockBuilder = out.beginBlockEntry(); - List valuesAtPercentiles = valuesAtPercentiles(digest, percentiles); - for (double value : valuesAtPercentiles) { - REAL.writeLong(blockBuilder, floatToRawIntBits((float) value)); - } - - out.closeEntry(); + ((ArrayBlockBuilder) out).buildEntry(elementBuilder -> { + for (double value : valuesAtPercentiles) { + REAL.writeLong(elementBuilder, floatToRawIntBits((float) value)); + } + }); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/BigintApproximateMostFrequent.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/BigintApproximateMostFrequent.java index bdcf87242691..ecd50b8ac4e0 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/BigintApproximateMostFrequent.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/BigintApproximateMostFrequent.java @@ -14,6 +14,7 @@ package io.trino.operator.aggregation; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.function.AccumulatorState; import io.trino.spi.function.AccumulatorStateMetadata; import io.trino.spi.function.AggregationFunction; @@ -94,12 +95,10 @@ public static void output(@AggregationState State state, BlockBuilder out) out.appendNull(); } else { - BlockBuilder entryBuilder = out.beginBlockEntry(); - state.get().forEachBucket((key, value) -> { - BigintType.BIGINT.writeLong(entryBuilder, key); - BigintType.BIGINT.writeLong(entryBuilder, value); - }); - out.closeEntry(); + ((MapBlockBuilder) out).buildEntry((keyBuilder, valueBuilder) -> state.get().forEachBucket((key, value) -> { + BigintType.BIGINT.writeLong(keyBuilder, key); + BigintType.BIGINT.writeLong(valueBuilder, value); + })); } } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/DecimalSumAggregation.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/DecimalSumAggregation.java index 4681cd56c62e..6439dbc23483 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/DecimalSumAggregation.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/DecimalSumAggregation.java @@ -16,6 +16,7 @@ import io.trino.operator.aggregation.state.LongDecimalWithOverflowState; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.Int128ArrayBlockBuilder; import io.trino.spi.function.AggregationFunction; import io.trino.spi.function.AggregationState; import io.trino.spi.function.BlockIndex; @@ -129,9 +130,7 @@ public static void outputDecimal(@AggregationState LongDecimalWithOverflowState long rawLow = decimal[offset + 1]; Decimals.throwIfOverflows(rawHigh, rawLow); - out.writeLong(rawHigh); - out.writeLong(rawLow); - out.closeEntry(); + ((Int128ArrayBlockBuilder) out).writeInt128(rawHigh, rawLow); } else { out.appendNull(); diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/DoubleHistogramAggregation.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/DoubleHistogramAggregation.java index 71688ff75cb7..fd40d4d35c5f 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/DoubleHistogramAggregation.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/DoubleHistogramAggregation.java @@ -15,6 +15,7 @@ import io.trino.operator.aggregation.state.DoubleHistogramStateSerializer; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.function.AccumulatorState; import io.trino.spi.function.AccumulatorStateMetadata; import io.trino.spi.function.AggregationFunction; @@ -92,13 +93,12 @@ public static void output(@AggregationState State state, BlockBuilder out) } else { Map value = state.get().getBuckets(); - - BlockBuilder entryBuilder = out.beginBlockEntry(); - for (Map.Entry entry : value.entrySet()) { - DoubleType.DOUBLE.writeDouble(entryBuilder, entry.getKey()); - DoubleType.DOUBLE.writeDouble(entryBuilder, entry.getValue()); - } - out.closeEntry(); + ((MapBlockBuilder) out).buildEntry((keyBuilder, valueBuilder) -> { + for (Map.Entry entry : value.entrySet()) { + DoubleType.DOUBLE.writeDouble(keyBuilder, entry.getKey()); + DoubleType.DOUBLE.writeDouble(valueBuilder, entry.getValue()); + } + }); } } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/KeyValuePairs.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/KeyValuePairs.java index 21c50c9804f7..81f37c4f073d 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/KeyValuePairs.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/KeyValuePairs.java @@ -16,6 +16,7 @@ import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.type.Type; import io.trino.type.BlockTypeOperators.BlockPositionEqual; import io.trino.type.BlockTypeOperators.BlockPositionHashCode; @@ -126,12 +127,12 @@ private void deserialize(Block block) public void serialize(BlockBuilder out) { - BlockBuilder mapBlockBuilder = out.beginBlockEntry(); - for (int i = 0; i < keyBlockBuilder.getPositionCount(); i++) { - keyType.appendTo(keyBlockBuilder, i, mapBlockBuilder); - valueType.appendTo(valueBlockBuilder, i, mapBlockBuilder); - } - out.closeEntry(); + ((MapBlockBuilder) out).buildEntry((keyBuilder, valueBuilder) -> { + for (int i = 0; i < keyBlockBuilder.getPositionCount(); i++) { + keyType.appendTo(keyBlockBuilder, i, keyBuilder); + valueType.appendTo(valueBlockBuilder, i, valueBuilder); + } + }); } public long estimatedInMemorySize() diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/RealHistogramAggregation.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/RealHistogramAggregation.java index e4720a9b97f4..ce5820565715 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/RealHistogramAggregation.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/RealHistogramAggregation.java @@ -14,6 +14,7 @@ package io.trino.operator.aggregation; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.function.AggregationFunction; import io.trino.spi.function.AggregationState; import io.trino.spi.function.CombineFunction; @@ -59,12 +60,12 @@ public static void output(@AggregationState DoubleHistogramAggregation.State sta } else { Map value = state.get().getBuckets(); - BlockBuilder entryBuilder = out.beginBlockEntry(); - for (Map.Entry entry : value.entrySet()) { - REAL.writeLong(entryBuilder, floatToRawIntBits(entry.getKey().floatValue())); - REAL.writeLong(entryBuilder, floatToRawIntBits(entry.getValue().floatValue())); - } - out.closeEntry(); + ((MapBlockBuilder) out).buildEntry((keyBuilder, valueBuilder) -> { + for (Map.Entry entry : value.entrySet()) { + REAL.writeLong(keyBuilder, floatToRawIntBits(entry.getKey().floatValue())); + REAL.writeLong(valueBuilder, floatToRawIntBits(entry.getValue().floatValue())); + } + }); } } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/VarcharApproximateMostFrequent.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/VarcharApproximateMostFrequent.java index da94e4420cd2..b63a79d5ebff 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/VarcharApproximateMostFrequent.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/VarcharApproximateMostFrequent.java @@ -15,6 +15,7 @@ import io.airlift.slice.Slice; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.function.AccumulatorState; import io.trino.spi.function.AccumulatorStateMetadata; import io.trino.spi.function.AggregationFunction; @@ -97,12 +98,10 @@ public static void output(@AggregationState State state, BlockBuilder out) out.appendNull(); } else { - BlockBuilder entryBuilder = out.beginBlockEntry(); - state.get().forEachBucket((key, value) -> { - VarcharType.VARCHAR.writeSlice(entryBuilder, key); - BigintType.BIGINT.writeLong(entryBuilder, value); - }); - out.closeEntry(); + ((MapBlockBuilder) out).buildEntry((keyBuilder, valueBuilder) -> state.get().forEachBucket((key, value) -> { + VarcharType.VARCHAR.writeSlice(keyBuilder, key); + BigintType.BIGINT.writeLong(valueBuilder, value); + })); } } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/arrayagg/ArrayAggregationFunction.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/arrayagg/ArrayAggregationFunction.java index a4c4b295f292..ca12a2edd6e9 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/arrayagg/ArrayAggregationFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/arrayagg/ArrayAggregationFunction.java @@ -14,6 +14,7 @@ package io.trino.operator.aggregation.arrayagg; import io.trino.operator.aggregation.NullablePosition; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.AggregationFunction; @@ -62,9 +63,7 @@ public static void output( out.appendNull(); } else { - BlockBuilder entryBuilder = out.beginBlockEntry(); - state.forEach((block, position) -> elementType.appendTo(block, position, entryBuilder)); - out.closeEntry(); + ((ArrayBlockBuilder) out).buildEntry(elementBuilder -> state.forEach((block, position) -> elementType.appendTo(block, position, elementBuilder))); } } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/arrayagg/ArrayAggregationStateSerializer.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/arrayagg/ArrayAggregationStateSerializer.java index 5b9d93a0c7e2..88bec652f1e1 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/arrayagg/ArrayAggregationStateSerializer.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/arrayagg/ArrayAggregationStateSerializer.java @@ -13,6 +13,7 @@ */ package io.trino.operator.aggregation.arrayagg; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.AccumulatorStateSerializer; @@ -45,9 +46,7 @@ public void serialize(ArrayAggregationState state, BlockBuilder out) out.appendNull(); } else { - BlockBuilder entryBuilder = out.beginBlockEntry(); - state.forEach((block, position) -> elementType.appendTo(block, position, entryBuilder)); - out.closeEntry(); + ((ArrayBlockBuilder) out).buildEntry(elementBuilder -> state.forEach((block, position) -> elementType.appendTo(block, position, elementBuilder))); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/histogram/GroupedTypedHistogram.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/histogram/GroupedTypedHistogram.java index 53695bb1da7c..2693afb2af9f 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/histogram/GroupedTypedHistogram.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/histogram/GroupedTypedHistogram.java @@ -18,6 +18,7 @@ import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.type.Type; import io.trino.type.BlockTypeOperators.BlockPositionEqual; import io.trino.type.BlockTypeOperators.BlockPositionHashCode; @@ -171,15 +172,11 @@ public void serialize(BlockBuilder out) out.appendNull(); } else { - BlockBuilder blockBuilder = out.beginBlockEntry(); - - iterateGroupNodes(currentGroupId, nodePointer -> { + ((MapBlockBuilder) out).buildEntry((keyBuilder, valueBuilder) -> iterateGroupNodes(currentGroupId, nodePointer -> { checkArgument(nodePointer != NULL, "should never see null here as we exclude in iterateGroupNodesCall"); ValueNode valueNode = bucketNodeFactory.createValueNode(nodePointer); - valueNode.writeNodeAsBlock(values, blockBuilder); - }); - - out.closeEntry(); + valueNode.writeNodeAsBlock(values, keyBuilder, valueBuilder); + })); } } @@ -384,10 +381,10 @@ void add(long count) * * @param valuesBlock - values.build() is called externally */ - void writeNodeAsBlock(Block valuesBlock, BlockBuilder outputBlockBuilder) + void writeNodeAsBlock(Block valuesBlock, BlockBuilder keyBuilder, BlockBuilder valueBuilder) { - type.appendTo(valuesBlock, getValuePosition(), outputBlockBuilder); - BIGINT.writeLong(outputBlockBuilder, getCount()); + type.appendTo(valuesBlock, getValuePosition(), keyBuilder); + BIGINT.writeLong(valueBuilder, getCount()); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/histogram/SingleTypedHistogram.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/histogram/SingleTypedHistogram.java index c0e7b9f1304e..2ff8c75d3305 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/histogram/SingleTypedHistogram.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/histogram/SingleTypedHistogram.java @@ -18,6 +18,7 @@ import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.type.Type; import io.trino.type.BlockTypeOperators.BlockPositionEqual; import io.trino.type.BlockTypeOperators.BlockPositionHashCode; @@ -106,12 +107,12 @@ public void serialize(BlockBuilder out) } else { Block valuesBlock = values.build(); - BlockBuilder blockBuilder = out.beginBlockEntry(); - for (int i = 0; i < valuesBlock.getPositionCount(); i++) { - type.appendTo(valuesBlock, i, blockBuilder); - BIGINT.writeLong(blockBuilder, counts.get(i)); - } - out.closeEntry(); + ((MapBlockBuilder) out).buildEntry((keyBuilder, valueBuilder) -> { + for (int i = 0; i < valuesBlock.getPositionCount(); i++) { + type.appendTo(valuesBlock, i, keyBuilder); + BIGINT.writeLong(valueBuilder, counts.get(i)); + } + }); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/listagg/ListaggAggregationFunction.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/listagg/ListaggAggregationFunction.java index 8355dd9b5dda..152b33e06dde 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/listagg/ListaggAggregationFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/listagg/ListaggAggregationFunction.java @@ -19,6 +19,7 @@ import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.function.AggregationFunction; import io.trino.spi.function.AggregationState; import io.trino.spi.function.BlockIndex; @@ -88,56 +89,56 @@ public static void output(ListaggAggregationState state, BlockBuilder out) out.appendNull(); } else { - outputState(state, out, MAX_OUTPUT_LENGTH); + outputState(state, (VariableWidthBlockBuilder) out, MAX_OUTPUT_LENGTH); } } @VisibleForTesting - public static void outputState(ListaggAggregationState state, BlockBuilder out, int maxOutputLength) + public static void outputState(ListaggAggregationState state, VariableWidthBlockBuilder blockBuilder, int maxOutputLength) { Slice separator = state.getSeparator(); int separatorLength = separator.length(); OutputContext context = new OutputContext(); - state.forEach((block, position) -> { - int entryLength = block.getSliceLength(position); - int spaceRequired = entryLength + (context.emittedEntryCount > 0 ? separatorLength : 0); - - if (context.outputLength + spaceRequired > maxOutputLength) { - context.overflow = true; - return false; - } - - if (context.emittedEntryCount > 0) { - out.writeBytes(separator, 0, separatorLength); - context.outputLength += separatorLength; + blockBuilder.buildEntry(out -> { + state.forEach((block, position) -> { + int entryLength = block.getSliceLength(position); + int spaceRequired = entryLength + (context.emittedEntryCount > 0 ? separatorLength : 0); + + if (context.outputLength + spaceRequired > maxOutputLength) { + context.overflow = true; + return false; + } + + if (context.emittedEntryCount > 0) { + out.writeBytes(separator, 0, separatorLength); + context.outputLength += separatorLength; + } + + block.writeSliceTo(position, 0, entryLength, out); + context.outputLength += entryLength; + context.emittedEntryCount++; + + return true; + }); + + if (context.overflow) { + if (state.isOverflowError()) { + throw new TrinoException(EXCEEDED_FUNCTION_MEMORY_LIMIT, format("Concatenated string has the length in bytes larger than the maximum output length %d", maxOutputLength)); + } + + if (context.emittedEntryCount > 0) { + out.writeBytes(separator, 0, separatorLength); + } + out.writeBytes(state.getOverflowFiller(), 0, state.getOverflowFiller().length()); + + if (state.showOverflowEntryCount()) { + out.writeBytes(Slices.utf8Slice("("), 0, 1); + Slice count = Slices.utf8Slice(Integer.toString(state.getEntryCount() - context.emittedEntryCount)); + out.writeBytes(count, 0, count.length()); + out.writeBytes(Slices.utf8Slice(")"), 0, 1); + } } - - block.writeBytesTo(position, 0, entryLength, out); - context.outputLength += entryLength; - context.emittedEntryCount++; - - return true; }); - - if (context.overflow) { - if (state.isOverflowError()) { - throw new TrinoException(EXCEEDED_FUNCTION_MEMORY_LIMIT, format("Concatenated string has the length in bytes larger than the maximum output length %d", maxOutputLength)); - } - - if (context.emittedEntryCount > 0) { - out.writeBytes(separator, 0, separatorLength); - } - out.writeBytes(state.getOverflowFiller(), 0, state.getOverflowFiller().length()); - - if (state.showOverflowEntryCount()) { - out.writeBytes(Slices.utf8Slice("("), 0, 1); - Slice count = Slices.utf8Slice(Integer.toString(state.getEntryCount() - context.emittedEntryCount)); - out.writeBytes(count, 0, count.length()); - out.writeBytes(Slices.utf8Slice(")"), 0, 1); - } - } - - out.closeEntry(); } private static class OutputContext diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/listagg/ListaggAggregationStateSerializer.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/listagg/ListaggAggregationStateSerializer.java index 965fc2dcf05c..ea6f077d227e 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/listagg/ListaggAggregationStateSerializer.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/listagg/ListaggAggregationStateSerializer.java @@ -16,9 +16,11 @@ import com.google.common.collect.ImmutableList; import io.airlift.slice.Slice; import io.trino.spi.block.AbstractRowBlock; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.ColumnarRow; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.function.AccumulatorStateSerializer; import io.trino.spi.type.ArrayType; import io.trino.spi.type.RowType; @@ -54,20 +56,17 @@ public void serialize(ListaggAggregationState state, BlockBuilder out) out.appendNull(); } else { - BlockBuilder rowBlockBuilder = out.beginBlockEntry(); - VARCHAR.writeSlice(rowBlockBuilder, state.getSeparator()); - BOOLEAN.writeBoolean(rowBlockBuilder, state.isOverflowError()); - VARCHAR.writeSlice(rowBlockBuilder, state.getOverflowFiller()); - BOOLEAN.writeBoolean(rowBlockBuilder, state.showOverflowEntryCount()); + ((RowBlockBuilder) out).buildEntry(fieldBuilders -> { + VARCHAR.writeSlice(fieldBuilders.get(0), state.getSeparator()); + BOOLEAN.writeBoolean(fieldBuilders.get(1), state.isOverflowError()); + VARCHAR.writeSlice(fieldBuilders.get(2), state.getOverflowFiller()); + BOOLEAN.writeBoolean(fieldBuilders.get(3), state.showOverflowEntryCount()); - BlockBuilder stateElementsBlockBuilder = rowBlockBuilder.beginBlockEntry(); - state.forEach((block, position) -> { - VARCHAR.appendTo(block, position, stateElementsBlockBuilder); - return true; + ((ArrayBlockBuilder) fieldBuilders.get(4)).buildEntry(elementBuilder -> state.forEach((block, position) -> { + VARCHAR.appendTo(block, position, elementBuilder); + return true; + })); }); - rowBlockBuilder.closeEntry(); - - out.closeEntry(); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxbyn/MinMaxByNStateFactory.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxbyn/MinMaxByNStateFactory.java index 1b828b6fcf19..7439b55591d4 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxbyn/MinMaxByNStateFactory.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxbyn/MinMaxByNStateFactory.java @@ -14,6 +14,7 @@ package io.trino.operator.aggregation.minmaxbyn; import io.trino.array.ObjectBigArray; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.AccumulatorState; @@ -119,13 +120,9 @@ public final void popAll(BlockBuilder out) return; } - BlockBuilder arrayBlockBuilder = out.beginBlockEntry(); - size -= typedHeap.getEstimatedSize(); - typedHeap.popAllReverse(arrayBlockBuilder); + ((ArrayBlockBuilder) out).buildEntry(typedHeap::popAllReverse); size += typedHeap.getEstimatedSize(); - - out.closeEntry(); } @Override @@ -238,9 +235,7 @@ public final void popAll(BlockBuilder out) return; } - BlockBuilder arrayBlockBuilder = out.beginBlockEntry(); - typedHeap.popAllReverse(arrayBlockBuilder); - out.closeEntry(); + ((ArrayBlockBuilder) out).buildEntry(typedHeap::popAllReverse); } @Override diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxbyn/TypedKeyValueHeap.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxbyn/TypedKeyValueHeap.java index 3f2586295da7..df07183843a4 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxbyn/TypedKeyValueHeap.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxbyn/TypedKeyValueHeap.java @@ -15,8 +15,10 @@ import com.google.common.base.Throwables; import com.google.common.collect.ImmutableList; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.RowType; import io.trino.spi.type.Type; @@ -96,22 +98,21 @@ public boolean isEmpty() public void serialize(BlockBuilder out) { - BlockBuilder blockBuilder = out.beginBlockEntry(); - BIGINT.writeLong(blockBuilder, getCapacity()); - - BlockBuilder keyElements = blockBuilder.beginBlockEntry(); - for (int i = 0; i < positionCount; i++) { - keyType.appendTo(keyBlockBuilder, heapIndex[i], keyElements); - } - blockBuilder.closeEntry(); - - BlockBuilder valueElements = blockBuilder.beginBlockEntry(); - for (int i = 0; i < positionCount; i++) { - valueType.appendTo(valueBlockBuilder, heapIndex[i], valueElements); - } - blockBuilder.closeEntry(); - - out.closeEntry(); + ((RowBlockBuilder) out).buildEntry(fieldBuilders -> { + BIGINT.writeLong(fieldBuilders.get(0), getCapacity()); + + ((ArrayBlockBuilder) fieldBuilders.get(1)).buildEntry(elementBuilder -> { + for (int i = 0; i < positionCount; i++) { + keyType.appendTo(keyBlockBuilder, heapIndex[i], elementBuilder); + } + }); + + ((ArrayBlockBuilder) fieldBuilders.get(2)).buildEntry(elementBuilder -> { + for (int i = 0; i < positionCount; i++) { + valueType.appendTo(valueBlockBuilder, heapIndex[i], elementBuilder); + } + }); + }); } public static TypedKeyValueHeap deserialize(boolean min, MethodHandle compare, Type keyType, Type valueType, Block rowBlock) diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxn/MinMaxNStateFactory.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxn/MinMaxNStateFactory.java index 86a65446871b..e3ef60fd4312 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxn/MinMaxNStateFactory.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxn/MinMaxNStateFactory.java @@ -14,6 +14,7 @@ package io.trino.operator.aggregation.minmaxn; import io.trino.array.ObjectBigArray; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.AccumulatorState; @@ -122,11 +123,7 @@ public final void writeAll(BlockBuilder out) return; } - BlockBuilder arrayBlockBuilder = out.beginBlockEntry(); - - typedHeap.writeAll(arrayBlockBuilder); - - out.closeEntry(); + ((ArrayBlockBuilder) out).buildEntry(typedHeap::writeAll); } @Override @@ -238,9 +235,7 @@ public final void writeAll(BlockBuilder out) return; } - BlockBuilder arrayBlockBuilder = out.beginBlockEntry(); - typedHeap.writeAll(arrayBlockBuilder); - out.closeEntry(); + ((ArrayBlockBuilder) out).buildEntry(typedHeap::writeAll); } @Override diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxn/TypedHeap.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxn/TypedHeap.java index aa5dc34df168..352f765fc82a 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxn/TypedHeap.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/minmaxn/TypedHeap.java @@ -14,8 +14,10 @@ package io.trino.operator.aggregation.minmaxn; import com.google.common.base.Throwables; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.Type; import it.unimi.dsi.fastutil.ints.IntArrays; @@ -84,16 +86,15 @@ public boolean isEmpty() public void serialize(BlockBuilder out) { - BlockBuilder blockBuilder = out.beginBlockEntry(); - BIGINT.writeLong(blockBuilder, capacity); + ((RowBlockBuilder) out).buildEntry(fieldBuilders -> { + BIGINT.writeLong(fieldBuilders.get(0), capacity); - BlockBuilder elements = blockBuilder.beginBlockEntry(); - for (int i = 0; i < positionCount; i++) { - elementType.appendTo(heapBlockBuilder, heapIndex[i], elements); - } - blockBuilder.closeEntry(); - - out.closeEntry(); + ((ArrayBlockBuilder) fieldBuilders.get(1)).buildEntry(elementBuilder -> { + for (int i = 0; i < positionCount; i++) { + elementType.appendTo(heapBlockBuilder, heapIndex[i], elementBuilder); + } + }); + }); } public static TypedHeap deserialize(boolean min, MethodHandle compare, Type elementType, Block rowBlock) diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/multimapagg/MultimapAggregationFunction.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/multimapagg/MultimapAggregationFunction.java index 986d03094142..da1c9fa7e03d 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/multimapagg/MultimapAggregationFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/multimapagg/MultimapAggregationFunction.java @@ -18,6 +18,7 @@ import io.trino.operator.aggregation.TypedSet; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.function.AggregationFunction; import io.trino.spi.function.AggregationState; import io.trino.spi.function.BlockIndex; @@ -109,12 +110,12 @@ public static void output( // Write keys and value arrays into one Block Type valueArrayType = new ArrayType(valueType); - BlockBuilder multimapBlockBuilder = out.beginBlockEntry(); - for (int i = 0; i < distinctKeyBlockBuilder.getPositionCount(); i++) { - keyType.appendTo(distinctKeyBlockBuilder, i, multimapBlockBuilder); - valueArrayType.writeObject(multimapBlockBuilder, valueArrayBlockBuilders.get(i).build()); - } - out.closeEntry(); + ((MapBlockBuilder) out).buildEntry((keyBuilder, valueBuilder) -> { + for (int i = 0; i < distinctKeyBlockBuilder.getPositionCount(); i++) { + keyType.appendTo(distinctKeyBlockBuilder, i, keyBuilder); + valueArrayType.writeObject(valueBuilder, valueArrayBlockBuilders.get(i).build()); + } + }); } } } diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/multimapagg/MultimapAggregationStateSerializer.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/multimapagg/MultimapAggregationStateSerializer.java index e3acd8549b8b..d90066d141e2 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/multimapagg/MultimapAggregationStateSerializer.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/multimapagg/MultimapAggregationStateSerializer.java @@ -14,9 +14,11 @@ package io.trino.operator.aggregation.multimapagg; import com.google.common.collect.ImmutableList; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.ColumnarRow; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.function.AccumulatorStateSerializer; import io.trino.spi.function.TypeParameter; import io.trino.spi.type.ArrayType; @@ -55,14 +57,13 @@ public void serialize(MultimapAggregationState state, BlockBuilder out) out.appendNull(); return; } - BlockBuilder entryBuilder = out.beginBlockEntry(); - state.forEach((keyBlock, valueBlock, position) -> { - BlockBuilder rowBlockBuilder = entryBuilder.beginBlockEntry(); - valueType.appendTo(valueBlock, position, rowBlockBuilder); - keyType.appendTo(keyBlock, position, rowBlockBuilder); - entryBuilder.closeEntry(); + ((ArrayBlockBuilder) out).buildEntry(elementBuilder -> { + RowBlockBuilder rowBlockBuilder = (RowBlockBuilder) elementBuilder; + state.forEach((keyBlock, valueBlock, position) -> rowBlockBuilder.buildEntry(fieldBuilders -> { + valueType.appendTo(valueBlock, position, fieldBuilders.get(0)); + keyType.appendTo(keyBlock, position, fieldBuilders.get(1)); + })); }); - out.closeEntry(); } @Override diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/state/StateCompiler.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/state/StateCompiler.java index 4a3efbbec031..015305b0a78f 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/state/StateCompiler.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/state/StateCompiler.java @@ -40,6 +40,8 @@ import io.trino.array.SliceBigArray; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; +import io.trino.spi.block.RowValueBuilder; import io.trino.spi.function.AccumulatorState; import io.trino.spi.function.AccumulatorStateFactory; import io.trino.spi.function.AccumulatorStateMetadata; @@ -101,6 +103,7 @@ import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.TinyintType.TINYINT; import static io.trino.spi.type.VarbinaryType.VARBINARY; +import static io.trino.sql.gen.LambdaMetafactoryGenerator.generateMetafactory; import static io.trino.sql.gen.SqlTypeBytecodeExpression.constantType; import static io.trino.type.UnknownType.UNKNOWN; import static io.trino.util.CompilerUtils.defineClass; @@ -283,7 +286,7 @@ private static void generateSerialize(ClassDefinition definition, CallSiteBi Scope scope = method.getScope(); BytecodeBlock serializerBody = method.getBody(); - if (fields.size() == 0) { + if (fields.isEmpty()) { serializerBody.append(out.invoke("appendNull", BlockBuilder.class).pop()); } else if (fields.size() == 1) { @@ -302,29 +305,52 @@ else if (fields.size() == 1) { serializerBody.append(sqlType.writeValue(out, fieldValue.cast(getOnlyElement(fields).getSqlType().getJavaType()))); } } - else if (fields.size() > 1) { - Variable rowBuilder = scope.declareVariable(BlockBuilder.class, "rowBuilder"); - serializerBody.append(rowBuilder.set(out.invoke("beginBlockEntry", BlockBuilder.class))); - for (StateField field : fields) { - Method getter = getGetter(clazz, field); - SqlTypeBytecodeExpression sqlType = constantType(binder, field.getSqlType()); - Variable fieldValue = scope.createTempVariable(getter.getReturnType()); - serializerBody.append(fieldValue.set(state.cast(getter.getDeclaringClass()).invoke(getter))); - if (!field.isPrimitiveType()) { - serializerBody.append(new IfStatement().condition(equal(fieldValue, constantNull(getter.getReturnType()))) - .ifTrue(rowBuilder.invoke("appendNull", BlockBuilder.class).pop()) - .ifFalse(sqlType.writeValue(rowBuilder, fieldValue))); - } - else { - // For primitive type, we need to cast here because we serialize byte fields with TINYINT/INTEGER (whose java type is long). - serializerBody.append(sqlType.writeValue(rowBuilder, fieldValue.cast(field.getSqlType().getJavaType()))); - } - } - serializerBody.append(out.invoke("closeEntry", BlockBuilder.class).pop()); + else { + MethodDefinition serializeToRow = generateSerializeToRow(definition, binder, clazz, fields); + BytecodeExpression rowEntryBuilder = generateMetafactory(RowValueBuilder.class, serializeToRow, ImmutableList.of(state)); + serializerBody.append(out.cast(RowBlockBuilder.class).invoke("buildEntry", void.class, rowEntryBuilder)); } serializerBody.ret(); } + private static MethodDefinition generateSerializeToRow( + ClassDefinition definition, + CallSiteBinder binder, + Class clazz, + List fields) + { + Parameter state = arg("state", AccumulatorState.class); + Parameter fieldBuilders = arg("fieldBuilders", type(List.class, BlockBuilder.class)); + MethodDefinition method = definition.declareMethod(a(PRIVATE, STATIC), "serialize", type(void.class), state, fieldBuilders); + Scope scope = method.getScope(); + BytecodeBlock body = method.getBody(); + + Variable fieldBuilder = scope.createTempVariable(BlockBuilder.class); + for (int i = 0; i < fields.size(); i++) { + StateField field = fields.get(i); + Method getter = getGetter(clazz, field); + + SqlTypeBytecodeExpression sqlType = constantType(binder, field.getSqlType()); + + Variable fieldValue = scope.createTempVariable(getter.getReturnType()); + body.append(fieldValue.set(state.cast(getter.getDeclaringClass()).invoke(getter))); + + body.append(fieldBuilder.set(fieldBuilders.invoke("get", Object.class, constantInt(i)).cast(BlockBuilder.class))); + + if (!field.isPrimitiveType()) { + body.append(new IfStatement().condition(equal(fieldValue, constantNull(getter.getReturnType()))) + .ifTrue(fieldBuilder.invoke("appendNull", BlockBuilder.class).pop()) + .ifFalse(sqlType.writeValue(fieldBuilder, fieldValue))); + } + else { + // For primitive type, we need to cast here because we serialize byte fields with TINYINT/INTEGER (whose java type is long). + body.append(sqlType.writeValue(fieldBuilder, fieldValue.cast(field.getSqlType().getJavaType()))); + } + } + body.ret(); + return method; + } + private static Method getSetter(Class clazz, StateField field) { try { diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/state/TriStateBooleanStateSerializer.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/state/TriStateBooleanStateSerializer.java index bc9ca23e935f..fe2222fc4d39 100644 --- a/core/trino-main/src/main/java/io/trino/operator/aggregation/state/TriStateBooleanStateSerializer.java +++ b/core/trino-main/src/main/java/io/trino/operator/aggregation/state/TriStateBooleanStateSerializer.java @@ -39,7 +39,7 @@ public void serialize(TriStateBooleanState state, BlockBuilder out) out.appendNull(); } else { - out.writeByte(state.getValue() == TRUE_VALUE ? 1 : 0).closeEntry(); + BOOLEAN.writeBoolean(out, state.getValue() == TRUE_VALUE); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/ArrayJoin.java b/core/trino-main/src/main/java/io/trino/operator/scalar/ArrayJoin.java index 9abc95f268dd..d3d7aa8b7744 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/ArrayJoin.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/ArrayJoin.java @@ -20,7 +20,7 @@ import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.function.BoundSignature; import io.trino.spi.function.FunctionDependencies; @@ -213,38 +213,36 @@ public static Slice arrayJoin( pageBuilder.reset(); } int numElements = arrayBlock.getPositionCount(); - BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0); - - boolean needsDelimiter = false; - for (int i = 0; i < numElements; i++) { - Slice value = null; - if (!arrayBlock.isNull(i)) { - try { - value = (Slice) castFunction.invokeExact(session, arrayBlock, i); + VariableWidthBlockBuilder blockBuilder = (VariableWidthBlockBuilder) pageBuilder.getBlockBuilder(0); + blockBuilder.buildEntry(valueWriter -> { + boolean needsDelimiter = false; + for (int i = 0; i < numElements; i++) { + Slice value = null; + if (!arrayBlock.isNull(i)) { + try { + value = (Slice) castFunction.invokeExact(session, arrayBlock, i); + } + catch (Throwable throwable) { + // Restore pageBuilder into a consistent state + pageBuilder.declarePosition(); + throw new TrinoException(GENERIC_INTERNAL_ERROR, "Error casting array element to VARCHAR", throwable); + } } - catch (Throwable throwable) { - // Restore pageBuilder into a consistent state - blockBuilder.closeEntry(); - pageBuilder.declarePosition(); - throw new TrinoException(GENERIC_INTERNAL_ERROR, "Error casting array element to VARCHAR", throwable); - } - } - if (value == null) { - value = nullReplacement; if (value == null) { - continue; + value = nullReplacement; + if (value == null) { + continue; + } } - } - if (needsDelimiter) { - blockBuilder.writeBytes(delimiter, 0, delimiter.length()); + if (needsDelimiter) { + valueWriter.writeBytes(delimiter, 0, delimiter.length()); + } + valueWriter.writeBytes(value, 0, value.length()); + needsDelimiter = true; } - blockBuilder.writeBytes(value, 0, value.length()); - needsDelimiter = true; - } - - blockBuilder.closeEntry(); + }); pageBuilder.declarePosition(); return VARCHAR.getSlice(blockBuilder, blockBuilder.getPositionCount() - 1); } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/EmptyMapConstructor.java b/core/trino-main/src/main/java/io/trino/operator/scalar/EmptyMapConstructor.java index d40c6b5cabaa..c83d9948b08b 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/EmptyMapConstructor.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/EmptyMapConstructor.java @@ -14,7 +14,6 @@ package io.trino.operator.scalar; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.Description; import io.trino.spi.function.ScalarFunction; import io.trino.spi.function.SqlType; @@ -22,6 +21,8 @@ import io.trino.spi.type.MapType; import io.trino.spi.type.Type; +import static io.trino.spi.block.MapValueBuilder.buildMapValue; + @Description("Creates an empty map") @ScalarFunction("map") public final class EmptyMapConstructor @@ -30,10 +31,7 @@ public final class EmptyMapConstructor public EmptyMapConstructor(@TypeParameter("map(unknown,unknown)") Type mapType) { - BlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1); - mapBlockBuilder.beginBlockEntry(); - mapBlockBuilder.closeEntry(); - emptyMap = ((MapType) mapType).getObject(mapBlockBuilder.build(), 0); + emptyMap = buildMapValue(((MapType) mapType), 0, (keyBuilder, valueBuilder) -> {}); } @SqlType("map(unknown,unknown)") diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapConcatFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapConcatFunction.java index fc8744523208..23ec1b58e104 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/MapConcatFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/MapConcatFunction.java @@ -13,14 +13,12 @@ */ package io.trino.operator.scalar; -import com.google.common.collect.ImmutableList; import io.trino.annotation.UsedByGeneratedCode; import io.trino.metadata.SqlScalarFunction; import io.trino.operator.aggregation.TypedSet; -import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedMapValueBuilder; import io.trino.spi.function.BoundSignature; import io.trino.spi.function.FunctionMetadata; import io.trino.spi.function.Signature; @@ -111,7 +109,7 @@ protected SpecializedSqlScalarFunction specialize(BoundSignature boundSignature) @UsedByGeneratedCode public static Object createMapState(MapType mapType) { - return new PageBuilder(ImmutableList.of(mapType)); + return BufferedMapValueBuilder.createBuffered(mapType); } @UsedByGeneratedCode @@ -130,47 +128,41 @@ public static Block mapConcat(MapType mapType, BlockPositionIsDistinctFrom keysD if (lastMapIndex == firstMapIndex) { return maps[lastMapIndex]; } + int last = lastMapIndex; + int first = firstMapIndex; - PageBuilder pageBuilder = (PageBuilder) state; - if (pageBuilder.isFull()) { - pageBuilder.reset(); - } + BufferedMapValueBuilder mapValueBuilder = (BufferedMapValueBuilder) state; // TODO: we should move TypedSet into user state as well Type keyType = mapType.getKeyType(); Type valueType = mapType.getValueType(); TypedSet typedSet = createDistinctTypedSet(keyType, keysDistinctOperator, keyHashCode, entries / 2, FUNCTION_NAME); - BlockBuilder mapBlockBuilder = pageBuilder.getBlockBuilder(0); - BlockBuilder blockBuilder = mapBlockBuilder.beginBlockEntry(); - - // the last map - Block map = maps[lastMapIndex]; - for (int i = 0; i < map.getPositionCount(); i += 2) { - typedSet.add(map, i); - keyType.appendTo(map, i, blockBuilder); - valueType.appendTo(map, i + 1, blockBuilder); - } - // the map between the last and the first - for (int idx = lastMapIndex - 1; idx > firstMapIndex; idx--) { - map = maps[idx]; + return mapValueBuilder.build(entries / 2, (keyBuilder, valueBuilder) -> { + // the last map + Block map = maps[last]; for (int i = 0; i < map.getPositionCount(); i += 2) { - if (typedSet.add(map, i)) { - keyType.appendTo(map, i, blockBuilder); - valueType.appendTo(map, i + 1, blockBuilder); + typedSet.add(map, i); + keyType.appendTo(map, i, keyBuilder); + valueType.appendTo(map, i + 1, valueBuilder); + } + // the map between the last and the first + for (int idx = last - 1; idx > first; idx--) { + map = maps[idx]; + for (int i = 0; i < map.getPositionCount(); i += 2) { + if (typedSet.add(map, i)) { + keyType.appendTo(map, i, keyBuilder); + valueType.appendTo(map, i + 1, valueBuilder); + } } } - } - // the first map - map = maps[firstMapIndex]; - for (int i = 0; i < map.getPositionCount(); i += 2) { - if (!typedSet.contains(map, i)) { - keyType.appendTo(map, i, blockBuilder); - valueType.appendTo(map, i + 1, blockBuilder); + // the first map + map = maps[first]; + for (int i = 0; i < map.getPositionCount(); i += 2) { + if (!typedSet.contains(map, i)) { + keyType.appendTo(map, i, keyBuilder); + valueType.appendTo(map, i + 1, valueBuilder); + } } - } - - mapBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); - return mapType.getObject(mapBlockBuilder, mapBlockBuilder.getPositionCount() - 1); + }); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapConstructor.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapConstructor.java index 2f522a9da9d8..b43f480767e1 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/MapConstructor.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/MapConstructor.java @@ -16,12 +16,10 @@ import com.google.common.collect.ImmutableList; import io.trino.annotation.UsedByGeneratedCode; import io.trino.metadata.SqlScalarFunction; -import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedMapValueBuilder; import io.trino.spi.block.DuplicateMapKeyException; -import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.function.BoundSignature; import io.trino.spi.function.FunctionDependencies; @@ -119,59 +117,43 @@ public static Block createMap( Block valueBlock) { checkCondition(keyBlock.getPositionCount() == valueBlock.getPositionCount(), INVALID_FUNCTION_ARGUMENT, "Key and value arrays must be the same length"); - PageBuilder pageBuilder = state.getPageBuilder(); - if (pageBuilder.isFull()) { - pageBuilder.reset(); - } - - MapBlockBuilder mapBlockBuilder = (MapBlockBuilder) pageBuilder.getBlockBuilder(0); - mapBlockBuilder.strict(); - BlockBuilder blockBuilder = mapBlockBuilder.beginBlockEntry(); - for (int i = 0; i < keyBlock.getPositionCount(); i++) { - if (keyBlock.isNull(i)) { - // close block builder before throwing as we may be in a TRY() call - // so that subsequent calls do not find it in an inconsistent state - mapBlockBuilder.closeEntry(); - throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null"); - } - Object keyObject = readNativeValue(mapType.getKeyType(), keyBlock, i); - try { - if ((boolean) keyIndeterminate.invoke(keyObject)) { - throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be indeterminate: " + mapType.getKeyType().getObjectValue(session, keyBlock, i)); - } - } - catch (Throwable t) { - mapBlockBuilder.closeEntry(); - throw internalError(t); - } - mapType.getKeyType().appendTo(keyBlock, i, blockBuilder); - mapType.getValueType().appendTo(valueBlock, i, blockBuilder); - } try { - mapBlockBuilder.closeEntry(); + return state.getMapValueBuilder().build(keyBlock.getPositionCount(), (keyBuilder, valueBuilder) -> { + for (int i = 0; i < keyBlock.getPositionCount(); i++) { + if (keyBlock.isNull(i)) { + throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null"); + } + Object keyObject = readNativeValue(mapType.getKeyType(), keyBlock, i); + try { + if ((boolean) keyIndeterminate.invoke(keyObject)) { + throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be indeterminate: " + mapType.getKeyType().getObjectValue(session, keyBlock, i)); + } + } + catch (Throwable t) { + throw internalError(t); + } + mapType.getKeyType().appendTo(keyBlock, i, keyBuilder); + mapType.getValueType().appendTo(valueBlock, i, valueBuilder); + } + }); } catch (DuplicateMapKeyException e) { throw e.withDetailedMessage(mapType.getKeyType(), session); } - finally { - pageBuilder.declarePosition(); - } - - return mapType.getObject(mapBlockBuilder, mapBlockBuilder.getPositionCount() - 1); } public static final class State { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public State(MapType mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBufferedStrict(mapType); } - public PageBuilder getPageBuilder() + public BufferedMapValueBuilder getMapValueBuilder() { - return pageBuilder; + return mapValueBuilder; } } } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapEntriesFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapEntriesFunction.java index 46281427fd76..e54c34f9ce0f 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/MapEntriesFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/MapEntriesFunction.java @@ -13,10 +13,9 @@ */ package io.trino.operator.scalar; -import com.google.common.collect.ImmutableList; -import io.trino.spi.PageBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedArrayValueBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.function.Description; import io.trino.spi.function.ScalarFunction; import io.trino.spi.function.SqlType; @@ -31,13 +30,13 @@ @Description("Construct an array of entries from a given map") public class MapEntriesFunction { - private final PageBuilder pageBuilder; + private final BufferedArrayValueBuilder arrayValueBuilder; @TypeParameter("K") @TypeParameter("V") public MapEntriesFunction(@TypeParameter("array(row(K,V))") Type arrayType) { - pageBuilder = new PageBuilder(ImmutableList.of(arrayType)); + arrayValueBuilder = BufferedArrayValueBuilder.createBuffered((ArrayType) arrayType); } @TypeParameter("K") @@ -52,24 +51,16 @@ public Block mapFromEntries( Type keyType = rowType.getTypeParameters().get(0); Type valueType = rowType.getTypeParameters().get(1); - ArrayType arrayType = new ArrayType(rowType); - - if (pageBuilder.isFull()) { - pageBuilder.reset(); - } int entryCount = block.getPositionCount() / 2; - BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0); - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (int i = 0; i < entryCount; i++) { - BlockBuilder rowBuilder = entryBuilder.beginBlockEntry(); - keyType.appendTo(block, 2 * i, rowBuilder); - valueType.appendTo(block, 2 * i + 1, rowBuilder); - entryBuilder.closeEntry(); - } - - blockBuilder.closeEntry(); - pageBuilder.declarePosition(); - return arrayType.getObject(blockBuilder, blockBuilder.getPositionCount() - 1); + return arrayValueBuilder.build(entryCount, valueBuilder -> { + for (int i = 0; i < entryCount; i++) { + int position = 2 * i; + ((RowBlockBuilder) valueBuilder).buildEntry(fieldBuilders -> { + keyType.appendTo(block, position, fieldBuilders.get(0)); + valueType.appendTo(block, position + 1, fieldBuilders.get(1)); + }); + } + }); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapFilterFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapFilterFunction.java index 88a2095ae2fd..471bef6e7484 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/MapFilterFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/MapFilterFunction.java @@ -24,11 +24,13 @@ import io.airlift.bytecode.Variable; import io.airlift.bytecode.control.ForLoop; import io.airlift.bytecode.control.IfStatement; +import io.airlift.bytecode.expression.BytecodeExpression; import io.trino.annotation.UsedByGeneratedCode; import io.trino.metadata.SqlScalarFunction; -import io.trino.spi.PageBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedMapValueBuilder; +import io.trino.spi.block.MapValueBuilder; import io.trino.spi.function.BoundSignature; import io.trino.spi.function.FunctionMetadata; import io.trino.spi.function.Signature; @@ -53,9 +55,9 @@ import static io.airlift.bytecode.expression.BytecodeExpressions.and; import static io.airlift.bytecode.expression.BytecodeExpressions.constantInt; import static io.airlift.bytecode.expression.BytecodeExpressions.constantNull; +import static io.airlift.bytecode.expression.BytecodeExpressions.divide; import static io.airlift.bytecode.expression.BytecodeExpressions.lessThan; import static io.airlift.bytecode.expression.BytecodeExpressions.notEqual; -import static io.airlift.bytecode.expression.BytecodeExpressions.subtract; import static io.airlift.bytecode.instruction.VariableInstruction.incrementVariable; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.FUNCTION; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.NEVER_NULL; @@ -63,6 +65,7 @@ import static io.trino.spi.type.BooleanType.BOOLEAN; import static io.trino.spi.type.TypeSignature.functionType; import static io.trino.spi.type.TypeSignature.mapType; +import static io.trino.sql.gen.LambdaMetafactoryGenerator.generateMetafactory; import static io.trino.sql.gen.SqlTypeBytecodeExpression.constantType; import static io.trino.type.UnknownType.UNKNOWN; import static io.trino.util.CompilerUtils.defineClass; @@ -107,23 +110,20 @@ public SpecializedSqlScalarFunction specialize(BoundSignature boundSignature) @UsedByGeneratedCode public static Object createState(MapType mapType) { - return new PageBuilder(ImmutableList.of(mapType)); + return BufferedMapValueBuilder.createBuffered(mapType); } private static MethodHandle generateFilter(MapType mapType) { CallSiteBinder binder = new CallSiteBinder(); - Type keyType = mapType.getKeyType(); - Type valueType = mapType.getValueType(); - Class keyJavaType = Primitives.wrap(keyType.getJavaType()); - Class valueJavaType = Primitives.wrap(valueType.getJavaType()); - ClassDefinition definition = new ClassDefinition( a(PUBLIC, FINAL), makeClassName("MapFilter"), type(Object.class)); definition.declareDefaultConstructor(a(PRIVATE)); + MethodDefinition filterKeyValue = generateFilterInner(definition, binder, mapType); + Parameter state = arg("state", Object.class); Parameter block = arg("block", Block.class); Parameter function = arg("function", BinaryFunctionInterface.class); @@ -135,11 +135,40 @@ private static MethodHandle generateFilter(MapType mapType) BytecodeBlock body = method.getBody(); Scope scope = method.getScope(); + + Variable mapValueBuilder = scope.declareVariable(BufferedMapValueBuilder.class, "mapValueBuilder"); + body.append(mapValueBuilder.set(state.cast(BufferedMapValueBuilder.class))); + + BytecodeExpression mapEntryBuilder = generateMetafactory(MapValueBuilder.class, filterKeyValue, ImmutableList.of(block, function)); + BytecodeExpression entryCount = divide(block.invoke("getPositionCount", int.class), constantInt(2)); + body.append(mapValueBuilder.invoke("build", Block.class, entryCount, mapEntryBuilder).ret()); + + Class generatedClass = defineClass(definition, Object.class, binder.getBindings(), MapFilterFunction.class.getClassLoader()); + return methodHandle(generatedClass, "filter", Object.class, Block.class, BinaryFunctionInterface.class); + } + + private static MethodDefinition generateFilterInner(ClassDefinition definition, CallSiteBinder binder, MapType mapType) + { + Parameter block = arg("block", Block.class); + Parameter function = arg("function", BinaryFunctionInterface.class); + Parameter keyBuilder = arg("keyBuilder", BlockBuilder.class); + Parameter valueBuilder = arg("valueBuilder", BlockBuilder.class); + MethodDefinition method = definition.declareMethod( + a(PRIVATE, STATIC), + "filter", + type(void.class), + ImmutableList.of(block, function, keyBuilder, valueBuilder)); + + BytecodeBlock body = method.getBody(); + Scope scope = method.getScope(); + + Type keyType = mapType.getKeyType(); + Type valueType = mapType.getValueType(); + Class keyJavaType = Primitives.wrap(keyType.getJavaType()); + Class valueJavaType = Primitives.wrap(valueType.getJavaType()); + Variable positionCount = scope.declareVariable(int.class, "positionCount"); Variable position = scope.declareVariable(int.class, "position"); - Variable pageBuilder = scope.declareVariable(PageBuilder.class, "pageBuilder"); - Variable mapBlockBuilder = scope.declareVariable(BlockBuilder.class, "mapBlockBuilder"); - Variable singleMapBlockWriter = scope.declareVariable(BlockBuilder.class, "singleMapBlockWriter"); Variable keyElement = scope.declareVariable(keyJavaType, "keyElement"); Variable valueElement = scope.declareVariable(valueJavaType, "valueElement"); Variable keep = scope.declareVariable(Boolean.class, "keep"); @@ -147,14 +176,6 @@ private static MethodHandle generateFilter(MapType mapType) // invoke block.getPositionCount() body.append(positionCount.set(block.invoke("getPositionCount", int.class))); - // prepare the single map block builder - body.append(pageBuilder.set(state.cast(PageBuilder.class))); - body.append(new IfStatement() - .condition(pageBuilder.invoke("isFull", boolean.class)) - .ifTrue(pageBuilder.invoke("reset", void.class))); - body.append(mapBlockBuilder.set(pageBuilder.invoke("getBlockBuilder", BlockBuilder.class, constantInt(0)))); - body.append(singleMapBlockWriter.set(mapBlockBuilder.invoke("beginBlockEntry", BlockBuilder.class))); - SqlTypeBytecodeExpression keySqlType = constantType(binder, keyType); BytecodeNode loadKeyElement; if (!keyType.equals(UNKNOWN)) { @@ -188,22 +209,10 @@ private static MethodHandle generateFilter(MapType mapType) .append(new IfStatement("if (keep != null && keep) ...") .condition(and(notEqual(keep, constantNull(Boolean.class)), keep.cast(boolean.class))) .ifTrue(new BytecodeBlock() - .append(keySqlType.invoke("appendTo", void.class, block, position, singleMapBlockWriter)) - .append(valueSqlType.invoke("appendTo", void.class, block, add(position, constantInt(1)), singleMapBlockWriter)))))); - - body.append(mapBlockBuilder - .invoke("closeEntry", BlockBuilder.class) - .pop()); - body.append(pageBuilder.invoke("declarePosition", void.class)); - body.append(constantType(binder, mapType) - .invoke( - "getObject", - Object.class, - mapBlockBuilder.cast(Block.class), - subtract(mapBlockBuilder.invoke("getPositionCount", int.class), constantInt(1))) - .ret()); + .append(keySqlType.invoke("appendTo", void.class, block, position, keyBuilder)) + .append(valueSqlType.invoke("appendTo", void.class, block, add(position, constantInt(1)), valueBuilder)))))); + body.ret(); - Class generatedClass = defineClass(definition, Object.class, binder.getBindings(), MapFilterFunction.class.getClassLoader()); - return methodHandle(generatedClass, "filter", Object.class, Block.class, BinaryFunctionInterface.class); + return method; } } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapFromEntriesFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapFromEntriesFunction.java index f4131c3d5ac6..587690e39ee7 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/MapFromEntriesFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/MapFromEntriesFunction.java @@ -14,11 +14,10 @@ package io.trino.operator.scalar; import com.google.common.collect.ImmutableList; -import io.trino.operator.aggregation.TypedSet; -import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedMapValueBuilder; +import io.trino.spi.block.DuplicateMapKeyException; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.function.Convention; import io.trino.spi.function.Description; @@ -33,25 +32,23 @@ import io.trino.type.BlockTypeOperators.BlockPositionHashCode; import io.trino.type.BlockTypeOperators.BlockPositionIsDistinctFrom; -import static io.trino.operator.aggregation.TypedSet.createDistinctTypedSet; import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.BLOCK_POSITION; import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL; import static io.trino.spi.function.OperatorType.HASH_CODE; import static io.trino.spi.function.OperatorType.IS_DISTINCT_FROM; -import static java.lang.String.format; @ScalarFunction("map_from_entries") @Description("Construct a map from an array of entries") public final class MapFromEntriesFunction { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; @TypeParameter("K") @TypeParameter("V") public MapFromEntriesFunction(@TypeParameter("map(K,V)") Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBufferedDistinctStrict((MapType) mapType); } @TypeParameter("K") @@ -75,42 +72,27 @@ public Block mapFromEntries( Type valueType = mapType.getValueType(); RowType mapEntryType = RowType.anonymous(ImmutableList.of(keyType, valueType)); - if (pageBuilder.isFull()) { - pageBuilder.reset(); - } - int entryCount = mapEntries.getPositionCount(); - BlockBuilder mapBlockBuilder = pageBuilder.getBlockBuilder(0); - BlockBuilder resultBuilder = mapBlockBuilder.beginBlockEntry(); - TypedSet uniqueKeys = createDistinctTypedSet(keyType, keysDistinctOperator, keyHashCode, entryCount, "map_from_entries"); - - for (int i = 0; i < entryCount; i++) { - if (mapEntries.isNull(i)) { - mapBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); - throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map entry cannot be null"); - } - Block mapEntryBlock = mapEntryType.getObject(mapEntries, i); + try { + return mapValueBuilder.build(entryCount, (keyBuilder, valueBuilder) -> { + for (int i = 0; i < entryCount; i++) { + if (mapEntries.isNull(i)) { + throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map entry cannot be null"); + } + Block mapEntryBlock = mapEntryType.getObject(mapEntries, i); - if (mapEntryBlock.isNull(0)) { - mapBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); - throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null"); - } + if (mapEntryBlock.isNull(0)) { + throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null"); + } - if (!uniqueKeys.add(mapEntryBlock, 0)) { - mapBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); - throw new TrinoException(INVALID_FUNCTION_ARGUMENT, format("Duplicate keys (%s) are not allowed", keyType.getObjectValue(session, mapEntryBlock, 0))); - } - - keyType.appendTo(mapEntryBlock, 0, resultBuilder); - valueType.appendTo(mapEntryBlock, 1, resultBuilder); + keyType.appendTo(mapEntryBlock, 0, keyBuilder); + valueType.appendTo(mapEntryBlock, 1, valueBuilder); + } + }); + } + catch (DuplicateMapKeyException e) { + throw e.withDetailedMessage(mapType.getKeyType(), session); } - - mapBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); - return mapType.getObject(mapBlockBuilder, mapBlockBuilder.getPositionCount() - 1); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapToMapCast.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapToMapCast.java index ece2aae99f0b..c5924a79eb43 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/MapToMapCast.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/MapToMapCast.java @@ -40,9 +40,9 @@ import java.lang.invoke.MethodHandles; import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.base.Preconditions.checkState; import static io.trino.operator.aggregation.TypedSet.createDistinctTypedSet; import static io.trino.spi.StandardErrorCode.INVALID_CAST_ARGUMENT; +import static io.trino.spi.block.MapValueBuilder.buildMapValue; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.BLOCK_POSITION; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.NEVER_NULL; import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL; @@ -249,47 +249,45 @@ public static Block mapCast( ConnectorSession session, Block fromMap) { - checkState(targetType.getTypeParameters().size() == 2, "Expect two type parameters for targetType"); - Type toKeyType = targetType.getTypeParameters().get(0); + MapType mapType = (MapType) targetType; + Type toKeyType = mapType.getKeyType(); TypedSet resultKeys = createDistinctTypedSet(toKeyType, keyDistinctOperator, keyHashCode, fromMap.getPositionCount() / 2, "map-to-map cast"); // Cast the keys into a new block - BlockBuilder keyBlockBuilder = toKeyType.createBlockBuilder(null, fromMap.getPositionCount() / 2); + BlockBuilder tempKeyBlockBuilder = toKeyType.createBlockBuilder(null, fromMap.getPositionCount() / 2); for (int i = 0; i < fromMap.getPositionCount(); i += 2) { try { - keyProcessFunction.invokeExact(fromMap, i, session, keyBlockBuilder); + keyProcessFunction.invokeExact(fromMap, i, session, tempKeyBlockBuilder); } catch (Throwable t) { throw internalError(t); } } - Block keyBlock = keyBlockBuilder.build(); + Block keyBlock = tempKeyBlockBuilder.build(); - BlockBuilder mapBlockBuilder = targetType.createBlockBuilder(null, 1); - BlockBuilder blockBuilder = mapBlockBuilder.beginBlockEntry(); - for (int i = 0; i < fromMap.getPositionCount(); i += 2) { - if (resultKeys.add(keyBlock, i / 2)) { - toKeyType.appendTo(keyBlock, i / 2, blockBuilder); - if (fromMap.isNull(i + 1)) { - blockBuilder.appendNull(); - continue; - } + // TODO this should build the value block directly and then construct a single map from the two blocks + return buildMapValue(mapType, fromMap.getPositionCount() / 2, (keyBuilder, valueBuilder) -> { + for (int i = 0; i < fromMap.getPositionCount(); i += 2) { + if (resultKeys.add(keyBlock, i / 2)) { + toKeyType.appendTo(keyBlock, i / 2, keyBuilder); + if (fromMap.isNull(i + 1)) { + valueBuilder.appendNull(); + continue; + } - try { - valueProcessFunction.invokeExact(fromMap, i + 1, session, blockBuilder); + try { + valueProcessFunction.invokeExact(fromMap, i + 1, session, valueBuilder); + } + catch (Throwable t) { + throw internalError(t); + } } - catch (Throwable t) { - throw internalError(t); + else { + // if there are duplicated keys, fail it! + throw new TrinoException(INVALID_CAST_ARGUMENT, "duplicate keys"); } } - else { - // if there are duplicated keys, fail it! - throw new TrinoException(INVALID_CAST_ARGUMENT, "duplicate keys"); - } - } - - mapBlockBuilder.closeEntry(); - return (Block) targetType.getObject(mapBlockBuilder, mapBlockBuilder.getPositionCount() - 1); + }); } public static MethodHandle nativeValueWriter(Type type) diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapTransformKeysFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapTransformKeysFunction.java index 4e01c7d22e16..0c6d7546ff67 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/MapTransformKeysFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/MapTransformKeysFunction.java @@ -24,14 +24,17 @@ import io.airlift.bytecode.Variable; import io.airlift.bytecode.control.ForLoop; import io.airlift.bytecode.control.IfStatement; +import io.airlift.bytecode.control.TryCatch; +import io.airlift.bytecode.expression.BytecodeExpression; import io.trino.annotation.UsedByGeneratedCode; import io.trino.metadata.SqlScalarFunction; -import io.trino.operator.aggregation.TypedSet; import io.trino.spi.ErrorCodeSupplier; -import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedMapValueBuilder; +import io.trino.spi.block.DuplicateMapKeyException; +import io.trino.spi.block.MapValueBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.function.BoundSignature; import io.trino.spi.function.FunctionMetadata; @@ -43,8 +46,6 @@ import io.trino.sql.gen.SqlTypeBytecodeExpression; import io.trino.sql.gen.lambda.BinaryFunctionInterface; import io.trino.type.BlockTypeOperators; -import io.trino.type.BlockTypeOperators.BlockPositionEqual; -import io.trino.type.BlockTypeOperators.BlockPositionHashCode; import java.lang.invoke.MethodHandle; import java.util.Optional; @@ -63,11 +64,8 @@ import static io.airlift.bytecode.expression.BytecodeExpressions.divide; import static io.airlift.bytecode.expression.BytecodeExpressions.equal; import static io.airlift.bytecode.expression.BytecodeExpressions.getStatic; -import static io.airlift.bytecode.expression.BytecodeExpressions.invokeStatic; import static io.airlift.bytecode.expression.BytecodeExpressions.lessThan; -import static io.airlift.bytecode.expression.BytecodeExpressions.newArray; import static io.airlift.bytecode.expression.BytecodeExpressions.newInstance; -import static io.airlift.bytecode.expression.BytecodeExpressions.subtract; import static io.airlift.bytecode.instruction.VariableInstruction.incrementVariable; import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.FUNCTION; @@ -75,13 +73,12 @@ import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL; import static io.trino.spi.type.TypeSignature.functionType; import static io.trino.spi.type.TypeSignature.mapType; -import static io.trino.sql.gen.BytecodeUtils.loadConstant; +import static io.trino.sql.gen.LambdaMetafactoryGenerator.generateMetafactory; import static io.trino.sql.gen.SqlTypeBytecodeExpression.constantType; import static io.trino.type.UnknownType.UNKNOWN; import static io.trino.util.CompilerUtils.defineClass; import static io.trino.util.CompilerUtils.makeClassName; import static io.trino.util.Reflection.methodHandle; -import static java.util.Objects.requireNonNull; public final class MapTransformKeysFunction extends SqlScalarFunction @@ -89,8 +86,6 @@ public final class MapTransformKeysFunction public static final String NAME = "transform_keys"; private static final MethodHandle STATE_FACTORY = methodHandle(MapTransformKeysFunction.class, "createState", MapType.class); - private final BlockTypeOperators blockTypeOperators; - public MapTransformKeysFunction(BlockTypeOperators blockTypeOperators) { super(FunctionMetadata.scalarBuilder() @@ -106,7 +101,6 @@ public MapTransformKeysFunction(BlockTypeOperators blockTypeOperators) .nondeterministic() .description("Apply lambda to each entry of the map and transform the key") .build()); - this.blockTypeOperators = requireNonNull(blockTypeOperators, "blockTypeOperators is null"); } @Override @@ -123,29 +117,27 @@ protected SpecializedSqlScalarFunction specialize(BoundSignature boundSignature) FAIL_ON_NULL, ImmutableList.of(NEVER_NULL, FUNCTION), ImmutableList.of(BinaryFunctionInterface.class), - generateTransformKey(inputKeyType, outputKeyType, valueType, outputMapType), + generateTransformKey(inputKeyType, outputKeyType, valueType), Optional.of(STATE_FACTORY.bindTo(outputMapType))); } @UsedByGeneratedCode public static Object createState(MapType mapType) { - return new PageBuilder(ImmutableList.of(mapType)); + return BufferedMapValueBuilder.createBufferedDistinctStrict(mapType); } - private MethodHandle generateTransformKey(Type keyType, Type transformedKeyType, Type valueType, Type resultMapType) + private static MethodHandle generateTransformKey(Type keyType, Type transformedKeyType, Type valueType) { CallSiteBinder binder = new CallSiteBinder(); - Class keyJavaType = Primitives.wrap(keyType.getJavaType()); - Class transformedKeyJavaType = Primitives.wrap(transformedKeyType.getJavaType()); - Class valueJavaType = Primitives.wrap(valueType.getJavaType()); - ClassDefinition definition = new ClassDefinition( a(PUBLIC, FINAL), makeClassName("MapTransformKey"), type(Object.class)); definition.declareDefaultConstructor(a(PRIVATE)); + MethodDefinition transformMap = generateTransformKeyInner(definition, binder, keyType, transformedKeyType, valueType); + Parameter state = arg("state", Object.class); Parameter session = arg("session", ConnectorSession.class); Parameter block = arg("block", Block.class); @@ -158,12 +150,50 @@ private MethodHandle generateTransformKey(Type keyType, Type transformedKeyType, BytecodeBlock body = method.getBody(); Scope scope = method.getScope(); + + Variable mapValueBuilder = scope.declareVariable(BufferedMapValueBuilder.class, "mapValueBuilder"); + body.append(mapValueBuilder.set(state.cast(BufferedMapValueBuilder.class))); + + BytecodeExpression mapEntryBuilder = generateMetafactory(MapValueBuilder.class, transformMap, ImmutableList.of(session, block, function)); + BytecodeExpression entryCount = divide(block.invoke("getPositionCount", int.class), constantInt(2)); + + Variable duplicateKeyException = scope.declareVariable(DuplicateMapKeyException.class, "e"); + body.append(new TryCatch( + mapValueBuilder.invoke("build", Block.class, entryCount, mapEntryBuilder).ret(), + ImmutableList.of( + new TryCatch.CatchBlock( + new BytecodeBlock() + .putVariable(duplicateKeyException) + .append(duplicateKeyException.invoke("withDetailedMessage", DuplicateMapKeyException.class, constantType(binder, transformedKeyType), session)) + .throwObject(), + ImmutableList.of(type(DuplicateMapKeyException.class)))))); + + Class generatedClass = defineClass(definition, Object.class, binder.getBindings(), MapTransformKeysFunction.class.getClassLoader()); + return methodHandle(generatedClass, "transform", Object.class, ConnectorSession.class, Block.class, BinaryFunctionInterface.class); + } + + private static MethodDefinition generateTransformKeyInner(ClassDefinition definition, CallSiteBinder binder, Type keyType, Type transformedKeyType, Type valueType) + { + Parameter session = arg("session", ConnectorSession.class); + Parameter block = arg("block", Block.class); + Parameter function = arg("function", BinaryFunctionInterface.class); + Parameter keyBuilder = arg("keyBuilder", BlockBuilder.class); + Parameter valueBuilder = arg("valueBuilder", BlockBuilder.class); + MethodDefinition method = definition.declareMethod( + a(PRIVATE, STATIC), + "transform", + type(void.class), + ImmutableList.of(session, block, function, keyBuilder, valueBuilder)); + + BytecodeBlock body = method.getBody(); + Scope scope = method.getScope(); + + Class keyJavaType = Primitives.wrap(keyType.getJavaType()); + Class transformedKeyJavaType = Primitives.wrap(transformedKeyType.getJavaType()); + Class valueJavaType = Primitives.wrap(valueType.getJavaType()); + Variable positionCount = scope.declareVariable(int.class, "positionCount"); Variable position = scope.declareVariable(int.class, "position"); - Variable pageBuilder = scope.declareVariable(PageBuilder.class, "pageBuilder"); - Variable mapBlockBuilder = scope.declareVariable(BlockBuilder.class, "mapBlockBuilder"); - Variable blockBuilder = scope.declareVariable(BlockBuilder.class, "blockBuilder"); - Variable typedSet = scope.declareVariable(TypedSet.class, "typeSet"); Variable keyElement = scope.declareVariable(keyJavaType, "keyElement"); Variable transformedKeyElement = scope.declareVariable(transformedKeyJavaType, "transformedKeyElement"); Variable valueElement = scope.declareVariable(valueJavaType, "valueElement"); @@ -171,25 +201,6 @@ private MethodHandle generateTransformKey(Type keyType, Type transformedKeyType, // invoke block.getPositionCount() body.append(positionCount.set(block.invoke("getPositionCount", int.class))); - // prepare the single map block builder - body.append(pageBuilder.set(state.cast(PageBuilder.class))); - body.append(new IfStatement() - .condition(pageBuilder.invoke("isFull", boolean.class)) - .ifTrue(pageBuilder.invoke("reset", void.class))); - body.append(mapBlockBuilder.set(pageBuilder.invoke("getBlockBuilder", BlockBuilder.class, constantInt(0)))); - body.append(blockBuilder.set(mapBlockBuilder.invoke("beginBlockEntry", BlockBuilder.class))); - - // create typed set - body.append(typedSet.set(invokeStatic( - TypedSet.class, - "createEqualityTypedSet", - TypedSet.class, - constantType(binder, transformedKeyType), - loadConstant(binder, blockTypeOperators.getEqualOperator(transformedKeyType), BlockPositionEqual.class), - loadConstant(binder, blockTypeOperators.getHashCodeOperator(transformedKeyType), BlockPositionHashCode.class), - divide(positionCount, constantInt(2)), - constantString(NAME)))); - // throw null key exception block BytecodeNode throwNullKeyException = new BytecodeBlock() .append(newInstance( @@ -204,12 +215,11 @@ private MethodHandle generateTransformKey(Type keyType, Type transformedKeyType, loadKeyElement = new BytecodeBlock().append(keyElement.set(keySqlType.getValue(block, position).cast(keyJavaType))); } else { - // make sure invokeExact will not take uninitialized keys during compile time - // but if we reach this point during runtime, it is an exception + // make sure invokeExact will not take uninitialized keys during compile time but, + // if we reach this point during runtime, it is an exception // also close the block builder before throwing as we may be in a TRY() call // so that subsequent calls do not find it in an inconsistent state loadKeyElement = new BytecodeBlock() - .append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop()) .append(keyElement.set(constantNull(keyJavaType))) .append(throwNullKeyException); } @@ -227,9 +237,7 @@ private MethodHandle generateTransformKey(Type keyType, Type transformedKeyType, loadValueElement = new BytecodeBlock().append(valueElement.set(constantNull(valueJavaType))); } - SqlTypeBytecodeExpression transformedKeySqlType = constantType(binder, transformedKeyType); BytecodeNode writeKeyElement; - BytecodeNode throwDuplicatedKeyException; if (!transformedKeyType.equals(UNKNOWN)) { writeKeyElement = new BytecodeBlock() .append(transformedKeyElement.set(function.invoke("apply", Object.class, keyElement.cast(Object.class), valueElement.cast(Object.class)).cast(transformedKeyJavaType))) @@ -237,28 +245,13 @@ private MethodHandle generateTransformKey(Type keyType, Type transformedKeyType, .condition(equal(transformedKeyElement, constantNull(transformedKeyJavaType))) .ifTrue(throwNullKeyException) .ifFalse(new BytecodeBlock() - .append(constantType(binder, transformedKeyType).writeValue(blockBuilder, transformedKeyElement.cast(transformedKeyType.getJavaType()))) - .append(valueSqlType.invoke("appendTo", void.class, block, add(position, constantInt(1)), blockBuilder)))); - - // make sure getObjectValue takes a known key type - throwDuplicatedKeyException = new BytecodeBlock() - .append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop()) - .append(newInstance( - TrinoException.class, - getStatic(INVALID_FUNCTION_ARGUMENT.getDeclaringClass(), "INVALID_FUNCTION_ARGUMENT").cast(ErrorCodeSupplier.class), - invokeStatic( - String.class, - "format", - String.class, - constantString("Duplicate keys (%s) are not allowed"), - newArray(type(Object[].class), ImmutableList.of(transformedKeySqlType.invoke("getObjectValue", Object.class, session, blockBuilder.cast(Block.class), position)))))) - .throwObject(); + .append(constantType(binder, transformedKeyType).writeValue(keyBuilder, transformedKeyElement.cast(transformedKeyType.getJavaType()))) + .append(valueSqlType.invoke("appendTo", void.class, block, add(position, constantInt(1)), valueBuilder)))); } else { // key cannot be unknown // if we reach this point during runtime, it is an exception writeKeyElement = throwNullKeyException; - throwDuplicatedKeyException = throwNullKeyException; } body.append(new ForLoop() @@ -268,24 +261,8 @@ private MethodHandle generateTransformKey(Type keyType, Type transformedKeyType, .body(new BytecodeBlock() .append(loadKeyElement) .append(loadValueElement) - .append(writeKeyElement) - .append(new IfStatement() - .condition(typedSet.invoke("add", boolean.class, blockBuilder.cast(Block.class), position)) - .ifFalse(throwDuplicatedKeyException)))); - - body.append(mapBlockBuilder - .invoke("closeEntry", BlockBuilder.class) - .pop()); - body.append(pageBuilder.invoke("declarePosition", void.class)); - body.append(constantType(binder, resultMapType) - .invoke( - "getObject", - Object.class, - mapBlockBuilder.cast(Block.class), - subtract(mapBlockBuilder.invoke("getPositionCount", int.class), constantInt(1))) - .ret()); - - Class generatedClass = defineClass(definition, Object.class, binder.getBindings(), MapTransformKeysFunction.class.getClassLoader()); - return methodHandle(generatedClass, "transform", Object.class, ConnectorSession.class, Block.class, BinaryFunctionInterface.class); + .append(writeKeyElement))); + body.ret(); + return method; } } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapTransformValuesFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapTransformValuesFunction.java index 2b0cbfde1833..f2e74f2faec5 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/MapTransformValuesFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/MapTransformValuesFunction.java @@ -26,13 +26,15 @@ import io.airlift.bytecode.control.ForLoop; import io.airlift.bytecode.control.IfStatement; import io.airlift.bytecode.control.TryCatch; +import io.airlift.bytecode.expression.BytecodeExpression; import io.trino.annotation.UsedByGeneratedCode; import io.trino.metadata.SqlScalarFunction; import io.trino.spi.ErrorCodeSupplier; -import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedMapValueBuilder; +import io.trino.spi.block.MapValueBuilder; import io.trino.spi.function.BoundSignature; import io.trino.spi.function.FunctionMetadata; import io.trino.spi.function.Signature; @@ -57,12 +59,12 @@ import static io.airlift.bytecode.expression.BytecodeExpressions.constantInt; import static io.airlift.bytecode.expression.BytecodeExpressions.constantNull; import static io.airlift.bytecode.expression.BytecodeExpressions.constantString; +import static io.airlift.bytecode.expression.BytecodeExpressions.divide; import static io.airlift.bytecode.expression.BytecodeExpressions.equal; import static io.airlift.bytecode.expression.BytecodeExpressions.getStatic; import static io.airlift.bytecode.expression.BytecodeExpressions.invokeStatic; import static io.airlift.bytecode.expression.BytecodeExpressions.lessThan; import static io.airlift.bytecode.expression.BytecodeExpressions.newInstance; -import static io.airlift.bytecode.expression.BytecodeExpressions.subtract; import static io.airlift.bytecode.instruction.VariableInstruction.incrementVariable; import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.FUNCTION; @@ -70,6 +72,7 @@ import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL; import static io.trino.spi.type.TypeSignature.functionType; import static io.trino.spi.type.TypeSignature.mapType; +import static io.trino.sql.gen.LambdaMetafactoryGenerator.generateMetafactory; import static io.trino.sql.gen.SqlTypeBytecodeExpression.constantType; import static io.trino.type.UnknownType.UNKNOWN; import static io.trino.util.CompilerUtils.defineClass; @@ -113,29 +116,27 @@ protected SpecializedSqlScalarFunction specialize(BoundSignature boundSignature) FAIL_ON_NULL, ImmutableList.of(NEVER_NULL, FUNCTION), ImmutableList.of(BinaryFunctionInterface.class), - generateTransform(keyType, inputValueType, outputValueType, outputMapType), + generateTransform(keyType, inputValueType, outputValueType), Optional.of(STATE_FACTORY.bindTo(outputMapType))); } @UsedByGeneratedCode public static Object createState(MapType mapType) { - return new PageBuilder(ImmutableList.of(mapType)); + return BufferedMapValueBuilder.createBuffered(mapType); } - private static MethodHandle generateTransform(Type keyType, Type valueType, Type transformedValueType, Type resultMapType) + private static MethodHandle generateTransform(Type keyType, Type valueType, Type transformedValueType) { CallSiteBinder binder = new CallSiteBinder(); - Class keyJavaType = Primitives.wrap(keyType.getJavaType()); - Class valueJavaType = Primitives.wrap(valueType.getJavaType()); - Class transformedValueJavaType = Primitives.wrap(transformedValueType.getJavaType()); - ClassDefinition definition = new ClassDefinition( a(PUBLIC, FINAL), makeClassName("MapTransformValue"), type(Object.class)); definition.declareDefaultConstructor(a(PRIVATE)); + MethodDefinition transformMap = generateTransformInner(definition, binder, keyType, valueType, transformedValueType); + // define transform method Parameter state = arg("state", Object.class); Parameter block = arg("block", Block.class); @@ -148,11 +149,39 @@ private static MethodHandle generateTransform(Type keyType, Type valueType, Type BytecodeBlock body = method.getBody(); Scope scope = method.getScope(); + + Variable mapValueBuilder = scope.declareVariable(BufferedMapValueBuilder.class, "mapValueBuilder"); + body.append(mapValueBuilder.set(state.cast(BufferedMapValueBuilder.class))); + + BytecodeExpression mapEntryBuilder = generateMetafactory(MapValueBuilder.class, transformMap, ImmutableList.of(block, function)); + BytecodeExpression entryCount = divide(block.invoke("getPositionCount", int.class), constantInt(2)); + body.append(mapValueBuilder.invoke("build", Block.class, entryCount, mapEntryBuilder).ret()); + + Class generatedClass = defineClass(definition, Object.class, binder.getBindings(), MapTransformValuesFunction.class.getClassLoader()); + return methodHandle(generatedClass, "transform", Object.class, Block.class, BinaryFunctionInterface.class); + } + + private static MethodDefinition generateTransformInner(ClassDefinition definition, CallSiteBinder binder, Type keyType, Type valueType, Type transformedValueType) + { + Parameter block = arg("block", Block.class); + Parameter function = arg("function", BinaryFunctionInterface.class); + Parameter keyBuilder = arg("keyBuilder", BlockBuilder.class); + Parameter valueBuilder = arg("valueBuilder", BlockBuilder.class); + MethodDefinition method = definition.declareMethod( + a(PRIVATE, STATIC), + "transform", + type(void.class), + ImmutableList.of(block, function, keyBuilder, valueBuilder)); + + BytecodeBlock body = method.getBody(); + Scope scope = method.getScope(); + + Class keyJavaType = Primitives.wrap(keyType.getJavaType()); + Class valueJavaType = Primitives.wrap(valueType.getJavaType()); + Class transformedValueJavaType = Primitives.wrap(transformedValueType.getJavaType()); + Variable positionCount = scope.declareVariable(int.class, "positionCount"); Variable position = scope.declareVariable(int.class, "position"); - Variable pageBuilder = scope.declareVariable(PageBuilder.class, "pageBuilder"); - Variable mapBlockBuilder = scope.declareVariable(BlockBuilder.class, "mapBlockBuilder"); - Variable blockBuilder = scope.declareVariable(BlockBuilder.class, "blockBuilder"); Variable keyElement = scope.declareVariable(keyJavaType, "keyElement"); Variable valueElement = scope.declareVariable(valueJavaType, "valueElement"); Variable transformedValueElement = scope.declareVariable(transformedValueJavaType, "transformedValueElement"); @@ -160,14 +189,6 @@ private static MethodHandle generateTransform(Type keyType, Type valueType, Type // invoke block.getPositionCount() body.append(positionCount.set(block.invoke("getPositionCount", int.class))); - // prepare the single map block builder - body.append(pageBuilder.set(state.cast(PageBuilder.class))); - body.append(new IfStatement() - .condition(pageBuilder.invoke("isFull", boolean.class)) - .ifTrue(pageBuilder.invoke("reset", void.class))); - body.append(mapBlockBuilder.set(pageBuilder.invoke("getBlockBuilder", BlockBuilder.class, constantInt(0)))); - body.append(blockBuilder.set(mapBlockBuilder.invoke("beginBlockEntry", BlockBuilder.class))); - // throw null key exception block BytecodeNode throwNullKeyException = new BytecodeBlock() .append(newInstance( @@ -182,12 +203,11 @@ private static MethodHandle generateTransform(Type keyType, Type valueType, Type loadKeyElement = new BytecodeBlock().append(keyElement.set(keySqlType.getValue(block, position).cast(keyJavaType))); } else { - // make sure invokeExact will not take uninitialized keys during compile time + // make sure invokeExact will not take uninitialized keys during compile time, // but if we reach this point during runtime, it is an exception // also close the block builder before throwing as we may be in a TRY() call // so that subsequent calls do not find it in an inconsistent state loadKeyElement = new BytecodeBlock() - .append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop()) .append(keyElement.set(constantNull(keyJavaType))) .append(throwNullKeyException); } @@ -208,11 +228,11 @@ private static MethodHandle generateTransform(Type keyType, Type valueType, Type if (!transformedValueType.equals(UNKNOWN)) { writeTransformedValueElement = new IfStatement() .condition(equal(transformedValueElement, constantNull(transformedValueJavaType))) - .ifTrue(blockBuilder.invoke("appendNull", BlockBuilder.class).pop()) - .ifFalse(constantType(binder, transformedValueType).writeValue(blockBuilder, transformedValueElement.cast(transformedValueType.getJavaType()))); + .ifTrue(valueBuilder.invoke("appendNull", BlockBuilder.class).pop()) + .ifFalse(constantType(binder, transformedValueType).writeValue(valueBuilder, transformedValueElement.cast(transformedValueType.getJavaType()))); } else { - writeTransformedValueElement = new BytecodeBlock().append(blockBuilder.invoke("appendNull", BlockBuilder.class).pop()); + writeTransformedValueElement = new BytecodeBlock().append(valueBuilder.invoke("appendNull", BlockBuilder.class).pop()); } Variable transformationException = scope.declareVariable(Throwable.class, "transformationException"); @@ -231,29 +251,15 @@ private static MethodHandle generateTransform(Type keyType, Type valueType, Type ImmutableList.of( new TryCatch.CatchBlock( new BytecodeBlock() - .append(mapBlockBuilder.invoke("closeEntry", BlockBuilder.class).pop()) - .append(pageBuilder.invoke("declarePosition", void.class)) .putVariable(transformationException) .append(invokeStatic(Throwables.class, "throwIfUnchecked", void.class, transformationException)) .append(newInstance(RuntimeException.class, transformationException)) .throwObject(), ImmutableList.of(type(Throwable.class)))))) - .append(keySqlType.invoke("appendTo", void.class, block, position, blockBuilder)) + .append(keySqlType.invoke("appendTo", void.class, block, position, keyBuilder)) .append(writeTransformedValueElement))); - body.append(mapBlockBuilder - .invoke("closeEntry", BlockBuilder.class) - .pop()); - body.append(pageBuilder.invoke("declarePosition", void.class)); - body.append(constantType(binder, resultMapType) - .invoke( - "getObject", - Object.class, - mapBlockBuilder.cast(Block.class), - subtract(mapBlockBuilder.invoke("getPositionCount", int.class), constantInt(1))) - .ret()); - - Class generatedClass = defineClass(definition, Object.class, binder.getBindings(), MapTransformValuesFunction.class.getClassLoader()); - return methodHandle(generatedClass, "transform", Object.class, Block.class, BinaryFunctionInterface.class); + body.ret(); + return method; } } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapZipWithFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapZipWithFunction.java index 4c5b3165e52b..01e74aff1839 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/MapZipWithFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/MapZipWithFunction.java @@ -15,9 +15,8 @@ import com.google.common.collect.ImmutableList; import io.trino.metadata.SqlScalarFunction; -import io.trino.spi.PageBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedMapValueBuilder; import io.trino.spi.block.SingleMapBlock; import io.trino.spi.function.BoundSignature; import io.trino.spi.function.FunctionMetadata; @@ -30,7 +29,6 @@ import java.lang.invoke.MethodHandle; import java.util.Optional; -import static com.google.common.base.Throwables.throwIfUnchecked; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.FUNCTION; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.NEVER_NULL; import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL; @@ -85,7 +83,7 @@ protected SpecializedSqlScalarFunction specialize(BoundSignature boundSignature) public static Object createState(MapType mapType) { - return new PageBuilder(ImmutableList.of(mapType)); + return BufferedMapValueBuilder.createBuffered(mapType); } public static Block mapZipWith( @@ -102,70 +100,41 @@ public static Block mapZipWith( SingleMapBlock rightMapBlock = (SingleMapBlock) rightBlock; Type outputValueType = outputMapType.getValueType(); - PageBuilder pageBuilder = (PageBuilder) state; - if (pageBuilder.isFull()) { - pageBuilder.reset(); - } - BlockBuilder mapBlockBuilder = pageBuilder.getBlockBuilder(0); - BlockBuilder blockBuilder = mapBlockBuilder.beginBlockEntry(); - - // seekKey() can take non-trivial time when key is complicated value, such as a long VARCHAR or ROW. - boolean[] keyFound = new boolean[rightMapBlock.getPositionCount()]; - for (int leftKeyPosition = 0; leftKeyPosition < leftMapBlock.getPositionCount(); leftKeyPosition += 2) { - Object key = readNativeValue(keyType, leftMapBlock, leftKeyPosition); - Object leftValue = readNativeValue(leftValueType, leftMapBlock, leftKeyPosition + 1); - - int rightValuePosition = rightMapBlock.seekKey(key); - Object rightValue = null; - if (rightValuePosition != -1) { - rightValue = readNativeValue(rightValueType, rightMapBlock, rightValuePosition); - keyFound[rightValuePosition / 2] = true; - } + int maxOutputSize = (leftMapBlock.getPositionCount() + rightMapBlock.getPositionCount()) / 2; + BufferedMapValueBuilder mapValueBuilder = (BufferedMapValueBuilder) state; + return mapValueBuilder.build(maxOutputSize, (keyBuilder, valueBuilder) -> { + // seekKey() can take non-trivial time when key is complicated value, such as a long VARCHAR or ROW. + boolean[] keyFound = new boolean[rightMapBlock.getPositionCount()]; + for (int leftKeyPosition = 0; leftKeyPosition < leftMapBlock.getPositionCount(); leftKeyPosition += 2) { + Object key = readNativeValue(keyType, leftMapBlock, leftKeyPosition); + Object leftValue = readNativeValue(leftValueType, leftMapBlock, leftKeyPosition + 1); + + int rightValuePosition = rightMapBlock.seekKey(key); + Object rightValue = null; + if (rightValuePosition != -1) { + rightValue = readNativeValue(rightValueType, rightMapBlock, rightValuePosition); + keyFound[rightValuePosition / 2] = true; + } - Object outputValue; - try { - outputValue = function.apply(key, leftValue, rightValue); - } - catch (Throwable throwable) { - // Restore pageBuilder into a consistent state. - mapBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); + Object outputValue = function.apply(key, leftValue, rightValue); - throwIfUnchecked(throwable); - throw new RuntimeException(throwable); + keyType.appendTo(leftMapBlock, leftKeyPosition, keyBuilder); + writeNativeValue(outputValueType, valueBuilder, outputValue); } - keyType.appendTo(leftMapBlock, leftKeyPosition, blockBuilder); - writeNativeValue(outputValueType, blockBuilder, outputValue); - } + // iterate over keys that only exists in rightMapBlock + for (int rightKeyPosition = 0; rightKeyPosition < rightMapBlock.getPositionCount(); rightKeyPosition += 2) { + if (!keyFound[rightKeyPosition / 2]) { + Object key = readNativeValue(keyType, rightMapBlock, rightKeyPosition); + Object rightValue = readNativeValue(rightValueType, rightMapBlock, rightKeyPosition + 1); - // iterate over keys that only exists in rightMapBlock - for (int rightKeyPosition = 0; rightKeyPosition < rightMapBlock.getPositionCount(); rightKeyPosition += 2) { - if (!keyFound[rightKeyPosition / 2]) { - Object key = readNativeValue(keyType, rightMapBlock, rightKeyPosition); - Object rightValue = readNativeValue(rightValueType, rightMapBlock, rightKeyPosition + 1); + Object outputValue = function.apply(key, null, rightValue); - Object outputValue; - try { - outputValue = function.apply(key, null, rightValue); + keyType.appendTo(rightMapBlock, rightKeyPosition, keyBuilder); + writeNativeValue(outputValueType, valueBuilder, outputValue); } - catch (Throwable throwable) { - // Restore pageBuilder into a consistent state. - mapBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); - - throwIfUnchecked(throwable); - throw new RuntimeException(throwable); - } - - keyType.appendTo(rightMapBlock, rightKeyPosition, blockBuilder); - writeNativeValue(outputValueType, blockBuilder, outputValue); } - } - - mapBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); - return outputMapType.getObject(mapBlockBuilder, mapBlockBuilder.getPositionCount() - 1); + }); } @FunctionalInterface diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MultimapFromEntriesFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MultimapFromEntriesFunction.java index cda225e87a3a..d3a4e6c9a40a 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/MultimapFromEntriesFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/MultimapFromEntriesFunction.java @@ -15,10 +15,10 @@ import com.google.common.collect.ImmutableList; import io.trino.operator.aggregation.TypedSet; -import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedMapValueBuilder; import io.trino.spi.function.Convention; import io.trino.spi.function.Description; import io.trino.spi.function.OperatorDependency; @@ -50,14 +50,14 @@ public final class MultimapFromEntriesFunction private static final String NAME = "multimap_from_entries"; private static final int INITIAL_ENTRY_COUNT = 128; - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; private IntList[] entryIndicesList; @TypeParameter("K") @TypeParameter("V") public MultimapFromEntriesFunction(@TypeParameter("map(K,array(V))") Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); initializeEntryIndicesList(INITIAL_ENTRY_COUNT); } @@ -81,10 +81,6 @@ public Block multimapFromEntries( Type valueType = ((ArrayType) mapType.getValueType()).getElementType(); RowType mapEntryType = RowType.anonymous(ImmutableList.of(keyType, valueType)); - if (pageBuilder.isFull()) { - pageBuilder.reset(); - } - int entryCount = mapEntries.getPositionCount(); if (entryCount > entryIndicesList.length) { initializeEntryIndicesList(entryCount); @@ -111,21 +107,19 @@ public Block multimapFromEntries( } } - BlockBuilder multimapBlockBuilder = pageBuilder.getBlockBuilder(0); - BlockBuilder mapWriter = multimapBlockBuilder.beginBlockEntry(); - for (int i = 0; i < keySet.size(); i++) { - keyType.appendTo(mapEntryType.getObject(mapEntries, entryIndicesList[i].getInt(0)), 0, mapWriter); - BlockBuilder valuesArray = mapWriter.beginBlockEntry(); - for (int entryIndex : entryIndicesList[i]) { - valueType.appendTo(mapEntryType.getObject(mapEntries, entryIndex), 1, valuesArray); + Block resultMap = mapValueBuilder.build(keySet.size(), (keyBuilder, valueBuilder) -> { + for (int i = 0; i < keySet.size(); i++) { + IntList indexList = entryIndicesList[i]; + keyType.appendTo(mapEntryType.getObject(mapEntries, indexList.getInt(0)), 0, keyBuilder); + ((ArrayBlockBuilder) valueBuilder).buildEntry(elementBuilder -> { + for (int entryIndex : indexList) { + valueType.appendTo(mapEntryType.getObject(mapEntries, entryIndex), 1, elementBuilder); + } + }); } - mapWriter.closeEntry(); - } - - multimapBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); + }); clearEntryIndices(keySet.size()); - return mapType.getObject(multimapBlockBuilder, multimapBlockBuilder.getPositionCount() - 1); + return resultMap; } private void clearEntryIndices(int entryCount) diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/RowToRowCast.java b/core/trino-main/src/main/java/io/trino/operator/scalar/RowToRowCast.java index d96baba6fe64..aa6791ae79e9 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/RowToRowCast.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/RowToRowCast.java @@ -23,12 +23,14 @@ import io.airlift.bytecode.Parameter; import io.airlift.bytecode.Scope; import io.airlift.bytecode.Variable; +import io.airlift.bytecode.expression.BytecodeExpression; import io.trino.metadata.SqlScalarFunction; import io.trino.spi.StandardErrorCode; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.BlockBuilderStatus; +import io.trino.spi.block.RowBlock; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.function.BoundSignature; import io.trino.spi.function.FunctionDependencies; @@ -40,23 +42,25 @@ import io.trino.spi.function.TypeVariableConstraint; import io.trino.spi.type.Type; import io.trino.spi.type.TypeSignature; -import io.trino.sql.gen.CachedInstanceBinder; import io.trino.sql.gen.CallSiteBinder; import java.lang.invoke.MethodHandle; import java.util.List; import java.util.Objects; +import java.util.Optional; import static io.airlift.bytecode.Access.FINAL; +import static io.airlift.bytecode.Access.PRIVATE; import static io.airlift.bytecode.Access.PUBLIC; import static io.airlift.bytecode.Access.STATIC; import static io.airlift.bytecode.Access.a; import static io.airlift.bytecode.Parameter.arg; import static io.airlift.bytecode.ParameterizedType.type; -import static io.airlift.bytecode.expression.BytecodeExpressions.constantBoolean; import static io.airlift.bytecode.expression.BytecodeExpressions.constantInt; import static io.airlift.bytecode.expression.BytecodeExpressions.constantNull; import static io.airlift.bytecode.expression.BytecodeExpressions.invokeDynamic; +import static io.airlift.bytecode.expression.BytecodeExpressions.invokeStatic; +import static io.airlift.bytecode.expression.BytecodeExpressions.newArray; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.BLOCK_POSITION; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.NEVER_NULL; import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL; @@ -112,7 +116,7 @@ private RowToRowCast() .signature(Signature.builder() .operatorType(CAST) .typeVariableConstraint( - // this is technically a recursive constraint for cast, but TypeRegistry.canCast has explicit handling for row to row cast + // this is technically a recursive constraint for cast, but SignatureBinder has explicit handling for row-to-row cast TypeVariableConstraint.builder("F") .variadicBound("row") .castableTo(new TypeSignature("T")) @@ -171,6 +175,7 @@ private static Class generateRowCast(Type fromType, Type toType, FunctionDepe a(PUBLIC, FINAL), makeClassName(Joiner.on("$").join("RowCast", BaseEncoding.base16().encode(hashSuffix))), type(Object.class)); + definition.declareDefaultConstructor(a(PRIVATE)); Parameter session = arg("session", ConnectorSession.class); Parameter row = arg("row", Block.class); @@ -185,65 +190,53 @@ private static Class generateRowCast(Type fromType, Type toType, FunctionDepe Scope scope = method.getScope(); BytecodeBlock body = method.getBody(); - Variable wasNull = scope.declareVariable(boolean.class, "wasNull"); - Variable blockBuilder = scope.createTempVariable(BlockBuilder.class); - Variable singleRowBlockWriter = scope.createTempVariable(BlockBuilder.class); + Variable fieldBlocks = scope.createTempVariable(Block[].class); + body.append(fieldBlocks.set(newArray(type(Block[].class), toTypes.size()))); - body.append(wasNull.set(constantBoolean(false))); - - CachedInstanceBinder cachedInstanceBinder = new CachedInstanceBinder(definition, binder); - - // create the row block builder - body.append(blockBuilder.set( - constantType(binder, toType).invoke( - "createBlockBuilder", - BlockBuilder.class, - constantNull(BlockBuilderStatus.class), - constantInt(1)))); - body.append(singleRowBlockWriter.set(blockBuilder.invoke("beginBlockEntry", BlockBuilder.class))); - - // loop through to append member blocks + Variable fieldBuilder = scope.createTempVariable(BlockBuilder.class); for (int i = 0; i < toTypes.size(); i++) { Type fromElementType = fromTypes.get(i); Type toElementType = toTypes.get(i); - Type currentFromType = fromElementType; - if (currentFromType.equals(UNKNOWN)) { - body.append(singleRowBlockWriter.invoke("appendNull", BlockBuilder.class).pop()); - continue; - } + body.append(fieldBuilder.set(constantType(binder, toElementType).invoke( + "createBlockBuilder", + BlockBuilder.class, + constantNull(BlockBuilderStatus.class), + constantInt(1)))); - MethodHandle castMethod = getNullSafeCast(functionDependencies, fromElementType, toElementType); - MethodHandle writeMethod = getNullSafeWrite(toElementType); - MethodHandle castAndWrite = collectArguments(writeMethod, 1, castMethod); - body.append(invokeDynamic( - BOOTSTRAP_METHOD, - ImmutableList.of(binder.bind(castAndWrite).getBindingId()), - "castAndWriteField", - castAndWrite.type(), - singleRowBlockWriter, - scope.getVariable("session"), - row, - constantInt(i))); + if (fromElementType.equals(UNKNOWN)) { + body.append(fieldBuilder.invoke("appendNull", BlockBuilder.class).pop()); + } + else { + MethodHandle castMethod = getNullSafeCast(functionDependencies, fromElementType, toElementType); + MethodHandle writeMethod = getNullSafeWrite(toElementType); + MethodHandle castAndWrite = collectArguments(writeMethod, 1, castMethod); + body.append(invokeDynamic( + BOOTSTRAP_METHOD, + ImmutableList.of(binder.bind(castAndWrite).getBindingId()), + "castAndWriteField", + castAndWrite.type(), + fieldBuilder, + session, + row, + constantInt(i))); + } + body.append(fieldBlocks.setElement(i, fieldBuilder.invoke("build", Block.class))); } - // call blockBuilder.closeEntry() and return the single row block - body.append(blockBuilder.invoke("closeEntry", BlockBuilder.class).pop()); + BytecodeExpression rowBlock = invokeStatic( + RowBlock.class, + "fromFieldBlocks", + Block.class, + constantInt(1), + invokeStatic(Optional.class, "empty", Optional.class), + fieldBlocks); + body.append(constantType(binder, toType) - .invoke("getObject", Object.class, blockBuilder.cast(Block.class), constantInt(0)) + .invoke("getObject", Object.class, rowBlock, constantInt(0)) .cast(Block.class) .ret()); - // create constructor - MethodDefinition constructorDefinition = definition.declareConstructor(a(PUBLIC)); - BytecodeBlock constructorBody = constructorDefinition.getBody(); - Variable thisVariable = constructorDefinition.getThis(); - constructorBody.comment("super();") - .append(thisVariable) - .invokeConstructor(Object.class); - cachedInstanceBinder.generateInitializations(thisVariable, constructorBody); - constructorBody.ret(); - return defineClass(definition, Object.class, binder.getBindings(), RowToRowCast.class.getClassLoader()); } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/SplitToMapFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/SplitToMapFunction.java index ccb5b63adfb0..31dcf6dbf8d9 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/SplitToMapFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/SplitToMapFunction.java @@ -14,16 +14,15 @@ package io.trino.operator.scalar; -import com.google.common.collect.ImmutableList; import io.airlift.slice.Slice; -import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedMapValueBuilder; import io.trino.spi.function.Description; import io.trino.spi.function.ScalarFunction; import io.trino.spi.function.SqlType; import io.trino.spi.function.TypeParameter; +import io.trino.spi.type.MapType; import io.trino.spi.type.StandardTypes; import io.trino.spi.type.Type; @@ -39,11 +38,11 @@ @ScalarFunction("split_to_map") public class SplitToMapFunction { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public SplitToMapFunction(@TypeParameter("map(varchar,varchar)") Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType("map(varchar,varchar)") @@ -95,18 +94,11 @@ public Block splitToMap(@TypeParameter("map(varchar,varchar)") Type mapType, @Sq entryStart = entryEnd + entryDelimiter.length(); } - if (pageBuilder.isFull()) { - pageBuilder.reset(); - } - BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0); - BlockBuilder singleMapBlockBuilder = blockBuilder.beginBlockEntry(); - for (Map.Entry entry : map.entrySet()) { - VARCHAR.writeSlice(singleMapBlockBuilder, entry.getKey()); - VARCHAR.writeSlice(singleMapBlockBuilder, entry.getValue()); - } - blockBuilder.closeEntry(); - pageBuilder.declarePosition(); - - return (Block) mapType.getObject(blockBuilder, blockBuilder.getPositionCount() - 1); + return mapValueBuilder.build(map.size(), (keyBuilder, valueBuilder) -> { + map.forEach((key, value) -> { + VARCHAR.writeSlice(keyBuilder, key); + VARCHAR.writeSlice(valueBuilder, value); + }); + }); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/SplitToMultimapFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/SplitToMultimapFunction.java index 38b6c7a16cc3..6cd805612326 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/SplitToMultimapFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/SplitToMultimapFunction.java @@ -15,23 +15,20 @@ package io.trino.operator.scalar; import com.google.common.collect.ArrayListMultimap; -import com.google.common.collect.ImmutableList; import com.google.common.collect.Multimap; import io.airlift.slice.Slice; -import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedMapValueBuilder; import io.trino.spi.function.Description; import io.trino.spi.function.ScalarFunction; import io.trino.spi.function.SqlType; import io.trino.spi.function.TypeParameter; +import io.trino.spi.type.MapType; import io.trino.spi.type.StandardTypes; import io.trino.spi.type.Type; -import java.util.Collection; -import java.util.Map; - import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; import static io.trino.spi.type.VarcharType.VARCHAR; import static io.trino.util.Failures.checkCondition; @@ -40,11 +37,11 @@ @ScalarFunction("split_to_multimap") public class SplitToMultimapFunction { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public SplitToMultimapFunction(@TypeParameter("map(varchar,array(varchar))") Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType("map(varchar,array(varchar))") @@ -95,24 +92,15 @@ public Block splitToMultimap( entryStart = entryEnd + entryDelimiter.length(); } - if (pageBuilder.isFull()) { - pageBuilder.reset(); - } - - pageBuilder.declarePosition(); - BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0); - BlockBuilder singleMapBlockBuilder = blockBuilder.beginBlockEntry(); - for (Map.Entry> entry : multimap.asMap().entrySet()) { - VARCHAR.writeSlice(singleMapBlockBuilder, entry.getKey()); - Collection values = entry.getValue(); - BlockBuilder valueBlockBuilder = singleMapBlockBuilder.beginBlockEntry(); - for (Slice value : values) { - VARCHAR.writeSlice(valueBlockBuilder, value); - } - singleMapBlockBuilder.closeEntry(); - } - blockBuilder.closeEntry(); - - return (Block) mapType.getObject(blockBuilder, blockBuilder.getPositionCount() - 1); + return mapValueBuilder.build(multimap.size(), (keyBuilder, valueBuilder) -> { + multimap.asMap().forEach((key, values) -> { + VARCHAR.writeSlice(keyBuilder, key); + ((ArrayBlockBuilder) valueBuilder).buildEntry(elementBuilder -> { + for (Slice value : values) { + VARCHAR.writeSlice(elementBuilder, value); + } + }); + }); + }); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/ZipFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/ZipFunction.java index e792ffc75ac5..7093c4934f01 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/ZipFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/ZipFunction.java @@ -16,7 +16,7 @@ import io.trino.annotation.UsedByGeneratedCode; import io.trino.metadata.SqlScalarFunction; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.function.BoundSignature; import io.trino.spi.function.FunctionMetadata; import io.trino.spi.function.Signature; @@ -103,19 +103,24 @@ public static Block zip(List types, Block... arrays) biggestCardinality = Math.max(biggestCardinality, array.getPositionCount()); } RowType rowType = RowType.anonymous(types); - BlockBuilder outputBuilder = rowType.createBlockBuilder(null, biggestCardinality); + RowBlockBuilder outputBuilder = rowType.createBlockBuilder(null, biggestCardinality); for (int outputPosition = 0; outputPosition < biggestCardinality; outputPosition++) { - BlockBuilder rowBuilder = outputBuilder.beginBlockEntry(); + buildRow(types, outputBuilder, outputPosition, arrays); + } + return outputBuilder.build(); + } + + private static void buildRow(List types, RowBlockBuilder outputBuilder, int outputPosition, Block[] arrays) + { + outputBuilder.buildEntry(fieldBuilders -> { for (int fieldIndex = 0; fieldIndex < arrays.length; fieldIndex++) { if (arrays[fieldIndex].getPositionCount() <= outputPosition) { - rowBuilder.appendNull(); + fieldBuilders.get(fieldIndex).appendNull(); } else { - types.get(fieldIndex).appendTo(arrays[fieldIndex], outputPosition, rowBuilder); + types.get(fieldIndex).appendTo(arrays[fieldIndex], outputPosition, fieldBuilders.get(fieldIndex)); } } - outputBuilder.closeEntry(); - } - return outputBuilder.build(); + }); } } diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/ZipWithFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/ZipWithFunction.java index 9c686c7c68fa..fbda5e9021f8 100644 --- a/core/trino-main/src/main/java/io/trino/operator/scalar/ZipWithFunction.java +++ b/core/trino-main/src/main/java/io/trino/operator/scalar/ZipWithFunction.java @@ -15,9 +15,8 @@ import com.google.common.collect.ImmutableList; import io.trino.metadata.SqlScalarFunction; -import io.trino.spi.PageBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedArrayValueBuilder; import io.trino.spi.function.BoundSignature; import io.trino.spi.function.FunctionMetadata; import io.trino.spi.function.Signature; @@ -29,7 +28,6 @@ import java.lang.invoke.MethodHandle; import java.util.Optional; -import static com.google.common.base.Throwables.throwIfUnchecked; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.FUNCTION; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.NEVER_NULL; import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL; @@ -84,7 +82,7 @@ protected SpecializedSqlScalarFunction specialize(BoundSignature boundSignature) public static Object createState(ArrayType arrayType) { - return new PageBuilder(ImmutableList.of(arrayType)); + return BufferedArrayValueBuilder.createBuffered(arrayType); } public static Block zipWith( @@ -101,33 +99,14 @@ public static Block zipWith( int rightPositionCount = rightBlock.getPositionCount(); int outputPositionCount = max(leftPositionCount, rightPositionCount); - PageBuilder pageBuilder = (PageBuilder) state; - if (pageBuilder.isFull()) { - pageBuilder.reset(); - } - BlockBuilder arrayBlockBuilder = pageBuilder.getBlockBuilder(0); - BlockBuilder blockBuilder = arrayBlockBuilder.beginBlockEntry(); - - for (int position = 0; position < outputPositionCount; position++) { - Object left = position < leftPositionCount ? readNativeValue(leftElementType, leftBlock, position) : null; - Object right = position < rightPositionCount ? readNativeValue(rightElementType, rightBlock, position) : null; - Object output; - try { - output = function.apply(left, right); - } - catch (Throwable throwable) { - // Restore pageBuilder into a consistent state. - arrayBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); - - throwIfUnchecked(throwable); - throw new RuntimeException(throwable); + BufferedArrayValueBuilder arrayValueBuilder = (BufferedArrayValueBuilder) state; + return arrayValueBuilder.build(outputPositionCount, valueBuilder -> { + for (int position = 0; position < outputPositionCount; position++) { + Object left = position < leftPositionCount ? readNativeValue(leftElementType, leftBlock, position) : null; + Object right = position < rightPositionCount ? readNativeValue(rightElementType, rightBlock, position) : null; + Object output = function.apply(left, right); + writeNativeValue(outputElementType, valueBuilder, output); } - writeNativeValue(outputElementType, blockBuilder, output); - } - - arrayBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); - return outputArrayType.getObject(arrayBlockBuilder, arrayBlockBuilder.getPositionCount() - 1); + }); } } diff --git a/core/trino-main/src/main/java/io/trino/sql/gen/LambdaMetafactoryGenerator.java b/core/trino-main/src/main/java/io/trino/sql/gen/LambdaMetafactoryGenerator.java new file mode 100644 index 000000000000..cafd610f861a --- /dev/null +++ b/core/trino-main/src/main/java/io/trino/sql/gen/LambdaMetafactoryGenerator.java @@ -0,0 +1,128 @@ +/* + * Licensed 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.trino.sql.gen; + +import com.google.common.collect.ImmutableList; +import io.airlift.bytecode.Access; +import io.airlift.bytecode.MethodDefinition; +import io.airlift.bytecode.ParameterizedType; +import io.airlift.bytecode.expression.BytecodeExpression; +import org.objectweb.asm.Handle; +import org.objectweb.asm.Opcodes; +import org.objectweb.asm.Type; + +import java.lang.invoke.LambdaMetafactory; +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.MethodType; +import java.lang.reflect.Method; +import java.lang.reflect.Modifier; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static io.airlift.bytecode.ParameterizedType.type; +import static io.airlift.bytecode.expression.BytecodeExpressions.invokeDynamic; +import static org.objectweb.asm.Type.getMethodType; +import static org.objectweb.asm.Type.getType; + +public final class LambdaMetafactoryGenerator +{ + private static final Method METAFACTORY; + + static { + try { + METAFACTORY = LambdaMetafactory.class.getMethod("metafactory", MethodHandles.Lookup.class, String.class, MethodType.class, MethodType.class, MethodHandle.class, MethodType.class); + } + catch (NoSuchMethodException e) { + throw new AssertionError(e); + } + } + + private LambdaMetafactoryGenerator() {} + + public static BytecodeExpression generateMetafactory(Class interfaceType, MethodDefinition targetMethod, List additionalArguments) + { + Method interfaceMethod = getSingleAbstractMethod(interfaceType); + + // verify target method has signature of additionalArguments + interfaceMethod + List expectedTypes = new ArrayList<>(); + if (targetMethod.getAccess().contains(Access.STATIC)) { + additionalArguments.forEach(argument -> expectedTypes.add(argument.getType())); + } + else { + checkArgument(!additionalArguments.isEmpty() && additionalArguments.get(0).getType().equals(targetMethod.getDeclaringClass().getType()), + "Expected first additional argument to be 'this' for non-static method"); + additionalArguments + .subList(1, additionalArguments.size()) + .forEach(argument -> expectedTypes.add(argument.getType())); + } + Arrays.stream(interfaceMethod.getParameterTypes()).forEach(type -> expectedTypes.add(type(type))); + checkArgument(expectedTypes.equals(targetMethod.getParameterTypes()), + "Expected target method to have parameter types %s, but has %s", expectedTypes, targetMethod.getParameterTypes()); + + Type interfaceMethodType = toMethodType(interfaceMethod); + return invokeDynamic( + METAFACTORY, + ImmutableList.of( + interfaceMethodType, + new Handle( + targetMethod.getAccess().contains(Access.STATIC) ? Opcodes.H_INVOKESTATIC : Opcodes.H_INVOKEVIRTUAL, + targetMethod.getDeclaringClass().getName(), + targetMethod.getName(), + targetMethod.getMethodDescriptor(), + false), + interfaceMethodType), + "build", + type(interfaceType), + additionalArguments); + } + + private static Type toMethodType(Method interfaceMethod) + { + return getMethodType( + getType(interfaceMethod.getReturnType()), + Arrays.stream(interfaceMethod.getParameterTypes()).map(Type::getType).toArray(Type[]::new)); + } + + private static Method getSingleAbstractMethod(Class interfaceType) + { + List interfaceMethods = Arrays.stream(interfaceType.getMethods()) + .filter(m -> Modifier.isAbstract(m.getModifiers())) + .filter(m -> Modifier.isPublic(m.getModifiers())) + .filter(LambdaMetafactoryGenerator::notJavaObjectMethod) + .collect(toImmutableList()); + if (interfaceMethods.size() != 1) { + throw new IllegalArgumentException(interfaceType.getSimpleName() + " does not have a single abstract method"); + } + return interfaceMethods.get(0); + } + + private static boolean notJavaObjectMethod(Method method) + { + return !methodMatches(method, "toString", String.class) && + !methodMatches(method, "hashCode", int.class) && + !methodMatches(method, "equals", boolean.class, Object.class); + } + + private static boolean methodMatches(Method method, String name, Class returnType, Class... parameterTypes) + { + return method.getParameterCount() == parameterTypes.length && + method.getReturnType() == returnType && + name.equals(method.getName()) && + Arrays.equals(method.getParameterTypes(), parameterTypes); + } +} diff --git a/core/trino-main/src/main/java/io/trino/sql/gen/RowConstructorCodeGenerator.java b/core/trino-main/src/main/java/io/trino/sql/gen/RowConstructorCodeGenerator.java index 124a04a78bbd..e87639b02364 100644 --- a/core/trino-main/src/main/java/io/trino/sql/gen/RowConstructorCodeGenerator.java +++ b/core/trino-main/src/main/java/io/trino/sql/gen/RowConstructorCodeGenerator.java @@ -18,18 +18,24 @@ import io.airlift.bytecode.Scope; import io.airlift.bytecode.Variable; import io.airlift.bytecode.control.IfStatement; +import io.airlift.bytecode.expression.BytecodeExpression; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.BlockBuilderStatus; +import io.trino.spi.block.RowBlock; import io.trino.spi.type.Type; import io.trino.sql.relational.RowExpression; import io.trino.sql.relational.SpecialForm; import java.util.List; +import java.util.Optional; +import static io.airlift.bytecode.ParameterizedType.type; import static io.airlift.bytecode.expression.BytecodeExpressions.constantFalse; import static io.airlift.bytecode.expression.BytecodeExpressions.constantInt; import static io.airlift.bytecode.expression.BytecodeExpressions.constantNull; +import static io.airlift.bytecode.expression.BytecodeExpressions.invokeStatic; +import static io.airlift.bytecode.expression.BytecodeExpressions.newArray; import static io.trino.sql.gen.SqlTypeBytecodeExpression.constantType; import static java.util.Objects.requireNonNull; @@ -54,32 +60,41 @@ public BytecodeNode generateExpression(BytecodeGeneratorContext context) Scope scope = context.getScope(); List types = rowType.getTypeParameters(); - block.comment("Create new RowBlockBuilder; beginBlockEntry;"); - Variable blockBuilder = scope.createTempVariable(BlockBuilder.class); - Variable singleRowBlockWriter = scope.createTempVariable(BlockBuilder.class); - block.append(blockBuilder.set( - constantType(binder, rowType).invoke( - "createBlockBuilder", - BlockBuilder.class, - constantNull(BlockBuilderStatus.class), - constantInt(1)))); - block.append(singleRowBlockWriter.set(blockBuilder.invoke("beginBlockEntry", BlockBuilder.class))); + Variable fieldBlocks = scope.createTempVariable(Block[].class); + block.append(fieldBlocks.set(newArray(type(Block[].class), arguments.size()))); + Variable blockBuilder = scope.createTempVariable(BlockBuilder.class); for (int i = 0; i < arguments.size(); ++i) { Type fieldType = types.get(i); Variable field = scope.createTempVariable(fieldType.getJavaType()); + + block.append(blockBuilder.set(constantType(binder, fieldType).invoke( + "createBlockBuilder", + BlockBuilder.class, + constantNull(BlockBuilderStatus.class), + constantInt(1)))); + block.comment("Clean wasNull and Generate + " + i + "-th field of row"); block.append(context.wasNull().set(constantFalse())); block.append(context.generate(arguments.get(i))); block.putVariable(field); block.append(new IfStatement() .condition(context.wasNull()) - .ifTrue(singleRowBlockWriter.invoke("appendNull", BlockBuilder.class).pop()) - .ifFalse(constantType(binder, fieldType).writeValue(singleRowBlockWriter, field).pop())); + .ifTrue(blockBuilder.invoke("appendNull", BlockBuilder.class).pop()) + .ifFalse(constantType(binder, fieldType).writeValue(blockBuilder, field).pop())); + + block.append(fieldBlocks.setElement(i, blockBuilder.invoke("build", Block.class))); } - block.comment("closeEntry; slice the SingleRowBlock; wasNull = false;"); - block.append(blockBuilder.invoke("closeEntry", BlockBuilder.class).pop()); - block.append(constantType(binder, rowType).invoke("getObject", Object.class, blockBuilder.cast(Block.class), constantInt(0)) + + BytecodeExpression rowBlock = invokeStatic( + RowBlock.class, + "fromFieldBlocks", + Block.class, + constantInt(1), + invokeStatic(Optional.class, "empty", Optional.class), + fieldBlocks); + + block.append(constantType(binder, rowType).invoke("getObject", Object.class, rowBlock, constantInt(0)) .cast(Block.class)); block.append(context.wasNull().set(constantFalse())); return block; diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/ExpressionInterpreter.java b/core/trino-main/src/main/java/io/trino/sql/planner/ExpressionInterpreter.java index a92868fdd0e4..012dbe55f2a6 100644 --- a/core/trino-main/src/main/java/io/trino/sql/planner/ExpressionInterpreter.java +++ b/core/trino-main/src/main/java/io/trino/sql/planner/ExpressionInterpreter.java @@ -31,7 +31,6 @@ import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; -import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.block.SingleRowBlock; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.function.FunctionNullability; @@ -133,6 +132,7 @@ import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.StandardErrorCode.TYPE_MISMATCH; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.NEVER_NULL; import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL; import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.NULLABLE_RETURN; @@ -1444,13 +1444,11 @@ protected Object visitRow(Row node, Object context) if (hasUnresolvedValue(values)) { return new Row(toExpressions(values, parameterTypes)); } - BlockBuilder blockBuilder = new RowBlockBuilder(parameterTypes, null, 1); - BlockBuilder singleRowBlockWriter = blockBuilder.beginBlockEntry(); - for (int i = 0; i < cardinality; ++i) { - writeNativeValue(parameterTypes.get(i), singleRowBlockWriter, values.get(i)); - } - blockBuilder.closeEntry(); - return rowType.getObject(blockBuilder, 0); + return buildRowValue(rowType, fields -> { + for (int i = 0; i < cardinality; ++i) { + writeNativeValue(parameterTypes.get(i), fields.get(i), values.get(i)); + } + }); } @Override @@ -1483,17 +1481,15 @@ protected Object visitFormat(Format node, Object context) .resolveFunction(session, QualifiedName.of(FormatFunction.NAME), TypeSignatureProvider.fromTypes(VARCHAR, rowType)); // Construct a row with arguments [1..n] and invoke the underlying function - BlockBuilder rowBuilder = new RowBlockBuilder(argumentTypes, null, 1); - BlockBuilder singleRowBlockWriter = rowBuilder.beginBlockEntry(); - for (int i = 0; i < arguments.size(); ++i) { - writeNativeValue(argumentTypes.get(i), singleRowBlockWriter, processedArguments.get(i)); - } - rowBuilder.closeEntry(); - + Block row = buildRowValue(rowType, fields -> { + for (int i = 0; i < arguments.size(); ++i) { + writeNativeValue(argumentTypes.get(i), fields.get(i), processedArguments.get(i)); + } + }); return functionInvoker.invoke( function, connectorSession, - ImmutableList.of(format, rowType.getObject(rowBuilder, 0))); + ImmutableList.of(format, row)); } @Override diff --git a/core/trino-main/src/main/java/io/trino/testing/MaterializedResult.java b/core/trino-main/src/main/java/io/trino/testing/MaterializedResult.java index e72ad60689ea..f3fffd9a1026 100644 --- a/core/trino-main/src/main/java/io/trino/testing/MaterializedResult.java +++ b/core/trino-main/src/main/java/io/trino/testing/MaterializedResult.java @@ -22,8 +22,11 @@ import io.trino.client.Warning; import io.trino.spi.Page; import io.trino.spi.PageBuilder; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.ArrayType; @@ -385,31 +388,31 @@ else if (TIMESTAMP_TZ_MILLIS.equals(type)) { else if (type instanceof ArrayType) { List list = (List) value; Type elementType = ((ArrayType) type).getElementType(); - BlockBuilder arrayBlockBuilder = blockBuilder.beginBlockEntry(); - for (Object element : list) { - writeValue(elementType, arrayBlockBuilder, element); - } - blockBuilder.closeEntry(); + ((ArrayBlockBuilder) blockBuilder).buildEntry(elementBuilder -> { + for (Object element : list) { + writeValue(elementType, elementBuilder, element); + } + }); } else if (type instanceof MapType) { Map map = (Map) value; Type keyType = ((MapType) type).getKeyType(); Type valueType = ((MapType) type).getValueType(); - BlockBuilder mapBlockBuilder = blockBuilder.beginBlockEntry(); - for (Entry entry : map.entrySet()) { - writeValue(keyType, mapBlockBuilder, entry.getKey()); - writeValue(valueType, mapBlockBuilder, entry.getValue()); - } - blockBuilder.closeEntry(); + ((MapBlockBuilder) blockBuilder).buildEntry((keyBuilder, valueBuilder) -> { + for (Entry entry : map.entrySet()) { + writeValue(keyType, keyBuilder, entry.getKey()); + writeValue(valueType, valueBuilder, entry.getValue()); + } + }); } else if (type instanceof RowType) { List row = (List) value; List fieldTypes = type.getTypeParameters(); - BlockBuilder rowBlockBuilder = blockBuilder.beginBlockEntry(); - for (int field = 0; field < row.size(); field++) { - writeValue(fieldTypes.get(field), rowBlockBuilder, row.get(field)); - } - blockBuilder.closeEntry(); + ((RowBlockBuilder) blockBuilder).buildEntry(fieldBuilders -> { + for (int field = 0; field < row.size(); field++) { + writeValue(fieldTypes.get(field), fieldBuilders.get(field), row.get(field)); + } + }); } else { throw new IllegalArgumentException("Unsupported type " + type); diff --git a/core/trino-main/src/main/java/io/trino/type/CodePointsType.java b/core/trino-main/src/main/java/io/trino/type/CodePointsType.java index 768918af9597..a2560ff4901f 100644 --- a/core/trino-main/src/main/java/io/trino/type/CodePointsType.java +++ b/core/trino-main/src/main/java/io/trino/type/CodePointsType.java @@ -17,6 +17,7 @@ import io.airlift.slice.Slices; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.TypeSignature; @@ -61,6 +62,6 @@ public Object getObject(Block block, int position) public void writeObject(BlockBuilder blockBuilder, Object value) { Slice slice = Slices.wrappedIntArray((int[]) value); - blockBuilder.writeBytes(slice, 0, slice.length()).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(slice); } } diff --git a/core/trino-main/src/main/java/io/trino/type/IpAddressType.java b/core/trino-main/src/main/java/io/trino/type/IpAddressType.java index 440e47d09453..5c04f49acc9e 100644 --- a/core/trino-main/src/main/java/io/trino/type/IpAddressType.java +++ b/core/trino-main/src/main/java/io/trino/type/IpAddressType.java @@ -130,9 +130,9 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(block.getLong(position, 0)); - blockBuilder.writeLong(block.getLong(position, SIZE_OF_LONG)); - blockBuilder.closeEntry(); + ((Int128ArrayBlockBuilder) blockBuilder).writeInt128( + block.getLong(position, 0), + block.getLong(position, SIZE_OF_LONG)); } } @@ -148,9 +148,9 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int l if (length != INT128_BYTES) { throw new IllegalStateException("Expected entry size to be exactly " + INT128_BYTES + " but was " + length); } - blockBuilder.writeLong(value.getLong(offset)); - blockBuilder.writeLong(value.getLong(offset + SIZE_OF_LONG)); - blockBuilder.closeEntry(); + ((Int128ArrayBlockBuilder) blockBuilder).writeInt128( + value.getLong(offset), + value.getLong(offset + SIZE_OF_LONG)); } @Override diff --git a/core/trino-main/src/main/java/io/trino/type/JoniRegexpType.java b/core/trino-main/src/main/java/io/trino/type/JoniRegexpType.java index 881de477ffae..4f8b547c42f3 100644 --- a/core/trino-main/src/main/java/io/trino/type/JoniRegexpType.java +++ b/core/trino-main/src/main/java/io/trino/type/JoniRegexpType.java @@ -16,6 +16,7 @@ import io.airlift.slice.Slice; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.TypeSignature; @@ -59,6 +60,6 @@ public Object getObject(Block block, int position) public void writeObject(BlockBuilder blockBuilder, Object value) { Slice pattern = ((JoniRegexp) value).pattern(); - blockBuilder.writeBytes(pattern, 0, pattern.length()).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(pattern); } } diff --git a/core/trino-main/src/main/java/io/trino/type/Json2016Type.java b/core/trino-main/src/main/java/io/trino/type/Json2016Type.java index 3c6468939323..ce4578dad268 100644 --- a/core/trino-main/src/main/java/io/trino/type/Json2016Type.java +++ b/core/trino-main/src/main/java/io/trino/type/Json2016Type.java @@ -21,6 +21,7 @@ import io.trino.operator.scalar.json.JsonOutputConversionError; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.StandardTypes; @@ -87,6 +88,6 @@ public void writeObject(BlockBuilder blockBuilder, Object value) } } Slice bytes = utf8Slice(json); - blockBuilder.writeBytes(bytes, 0, bytes.length()).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(bytes); } } diff --git a/core/trino-main/src/main/java/io/trino/type/JsonPath2016Type.java b/core/trino-main/src/main/java/io/trino/type/JsonPath2016Type.java index 02e54b8757a4..4d754efb93c6 100644 --- a/core/trino-main/src/main/java/io/trino/type/JsonPath2016Type.java +++ b/core/trino-main/src/main/java/io/trino/type/JsonPath2016Type.java @@ -23,6 +23,7 @@ import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.BlockEncodingSerde; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.Type; @@ -71,7 +72,7 @@ public void writeObject(BlockBuilder blockBuilder, Object value) { String json = jsonPathCodec.toJson((IrJsonPath) value); Slice bytes = utf8Slice(json); - blockBuilder.writeBytes(bytes, 0, bytes.length()).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(bytes); } private static JsonCodec getCodec(TypeDeserializer typeDeserializer, BlockEncodingSerde blockEncodingSerde) diff --git a/core/trino-main/src/main/java/io/trino/type/JsonPathType.java b/core/trino-main/src/main/java/io/trino/type/JsonPathType.java index b7d7eb68225f..24a9ed298afa 100644 --- a/core/trino-main/src/main/java/io/trino/type/JsonPathType.java +++ b/core/trino-main/src/main/java/io/trino/type/JsonPathType.java @@ -18,6 +18,7 @@ import io.trino.operator.scalar.JsonPath; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.TypeSignature; @@ -59,6 +60,6 @@ public Object getObject(Block block, int position) public void writeObject(BlockBuilder blockBuilder, Object value) { Slice pattern = Slices.utf8Slice(((JsonPath) value).pattern()); - blockBuilder.writeBytes(pattern, 0, pattern.length()).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(pattern); } } diff --git a/core/trino-main/src/main/java/io/trino/type/JsonType.java b/core/trino-main/src/main/java/io/trino/type/JsonType.java index f50aed7f5e4c..3dedba3a2f8e 100644 --- a/core/trino-main/src/main/java/io/trino/type/JsonType.java +++ b/core/trino-main/src/main/java/io/trino/type/JsonType.java @@ -18,6 +18,7 @@ import io.airlift.slice.XxHash64; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.function.BlockIndex; import io.trino.spi.function.BlockPosition; @@ -77,8 +78,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -102,7 +102,7 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value) @Override public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length) { - blockBuilder.writeBytes(value, offset, length).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length); } @ScalarOperator(EQUAL) diff --git a/core/trino-main/src/main/java/io/trino/type/LikePatternType.java b/core/trino-main/src/main/java/io/trino/type/LikePatternType.java index 5edce0406160..2d690b41517f 100644 --- a/core/trino-main/src/main/java/io/trino/type/LikePatternType.java +++ b/core/trino-main/src/main/java/io/trino/type/LikePatternType.java @@ -17,6 +17,7 @@ import io.trino.likematcher.LikeMatcher; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.TypeSignature; @@ -78,18 +79,18 @@ public Object getObject(Block block, int position) public void writeObject(BlockBuilder blockBuilder, Object value) { LikeMatcher matcher = (LikeMatcher) value; - - Slice pattern = utf8Slice(matcher.getPattern()); - int length = pattern.length(); - blockBuilder.writeInt(length); - blockBuilder.writeBytes(pattern, 0, length); - if (matcher.getEscape().isEmpty()) { - blockBuilder.writeByte(0); - } - else { - blockBuilder.writeByte(1); - blockBuilder.writeInt(matcher.getEscape().get()); - } - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueWriter -> { + Slice pattern = utf8Slice(matcher.getPattern()); + int length = pattern.length(); + valueWriter.writeInt(length); + valueWriter.writeBytes(pattern, 0, length); + if (matcher.getEscape().isEmpty()) { + valueWriter.writeByte(0); + } + else { + valueWriter.writeByte(1); + valueWriter.writeInt(matcher.getEscape().get()); + } + }); } } diff --git a/core/trino-main/src/main/java/io/trino/type/Re2JRegexpType.java b/core/trino-main/src/main/java/io/trino/type/Re2JRegexpType.java index 639d5553eaa4..1a3b10b48d14 100644 --- a/core/trino-main/src/main/java/io/trino/type/Re2JRegexpType.java +++ b/core/trino-main/src/main/java/io/trino/type/Re2JRegexpType.java @@ -18,6 +18,7 @@ import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.TypeSignature; @@ -72,6 +73,6 @@ public Object getObject(Block block, int position) public void writeObject(BlockBuilder blockBuilder, Object value) { Slice pattern = Slices.utf8Slice(((Re2JRegexp) value).pattern()); - blockBuilder.writeBytes(pattern, 0, pattern.length()).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(pattern); } } diff --git a/core/trino-main/src/main/java/io/trino/type/TDigestType.java b/core/trino-main/src/main/java/io/trino/type/TDigestType.java index 5c6a95b8438c..066d21d235a8 100644 --- a/core/trino-main/src/main/java/io/trino/type/TDigestType.java +++ b/core/trino-main/src/main/java/io/trino/type/TDigestType.java @@ -17,6 +17,7 @@ import io.airlift.stats.TDigest; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.SqlVarbinary; @@ -40,8 +41,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -55,7 +55,7 @@ public Object getObject(Block block, int position) public void writeObject(BlockBuilder blockBuilder, Object value) { Slice serialized = ((TDigest) value).serialize(); - blockBuilder.writeBytes(serialized, 0, serialized.length()).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(serialized); } @Override diff --git a/core/trino-main/src/main/java/io/trino/type/setdigest/SetDigestFunctions.java b/core/trino-main/src/main/java/io/trino/type/setdigest/SetDigestFunctions.java index c7e8e23722e1..89e184aeb5af 100644 --- a/core/trino-main/src/main/java/io/trino/type/setdigest/SetDigestFunctions.java +++ b/core/trino-main/src/main/java/io/trino/type/setdigest/SetDigestFunctions.java @@ -16,15 +16,14 @@ import io.airlift.slice.Slice; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.ScalarFunction; import io.trino.spi.function.SqlType; import io.trino.spi.function.TypeParameter; +import io.trino.spi.type.MapType; import io.trino.spi.type.StandardTypes; import io.trino.spi.type.Type; -import java.util.Map; - +import static io.trino.spi.block.MapValueBuilder.buildMapValue; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.SmallintType.SMALLINT; import static io.trino.type.setdigest.SetDigest.exactIntersectionCardinality; @@ -81,16 +80,14 @@ public static double jaccardIndex(@SqlType(SetDigestType.NAME) Slice slice1, @Sq public static Block hashCounts(@TypeParameter("map(bigint,smallint)") Type mapType, @SqlType(SetDigestType.NAME) Slice slice) { SetDigest digest = SetDigest.newInstance(slice); - - // Maybe use static BlockBuilderStatus in order avoid `new`? - BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 1); - BlockBuilder singleMapBlockBuilder = blockBuilder.beginBlockEntry(); - for (Map.Entry entry : digest.getHashCounts().entrySet()) { - BIGINT.writeLong(singleMapBlockBuilder, entry.getKey()); - SMALLINT.writeLong(singleMapBlockBuilder, entry.getValue()); - } - blockBuilder.closeEntry(); - - return (Block) mapType.getObject(blockBuilder, 0); + return buildMapValue( + ((MapType) mapType), + digest.getHashCounts().size(), + (keyBuilder, valueBuilder) -> { + digest.getHashCounts().forEach((key, value) -> { + BIGINT.writeLong(keyBuilder, key); + SMALLINT.writeLong(valueBuilder, value); + }); + }); } } diff --git a/core/trino-main/src/main/java/io/trino/type/setdigest/SetDigestType.java b/core/trino-main/src/main/java/io/trino/type/setdigest/SetDigestType.java index a45b95b5f2b0..e921fef765e7 100644 --- a/core/trino-main/src/main/java/io/trino/type/setdigest/SetDigestType.java +++ b/core/trino-main/src/main/java/io/trino/type/setdigest/SetDigestType.java @@ -17,6 +17,7 @@ import io.airlift.slice.Slice; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.SqlVarbinary; @@ -59,8 +60,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -73,6 +73,6 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value) @Override public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length) { - blockBuilder.writeBytes(value, offset, length).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length); } } diff --git a/core/trino-main/src/main/java/io/trino/util/JsonUtil.java b/core/trino-main/src/main/java/io/trino/util/JsonUtil.java index 2449665235cb..a16a1328550c 100644 --- a/core/trino-main/src/main/java/io/trino/util/JsonUtil.java +++ b/core/trino-main/src/main/java/io/trino/util/JsonUtil.java @@ -24,11 +24,12 @@ import io.airlift.slice.SliceOutput; import io.airlift.slice.Slices; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.DuplicateMapKeyException; -import io.trino.spi.block.SingleMapBlockWriter; -import io.trino.spi.block.SingleRowBlockWriter; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.BigintType; import io.trino.spi.type.BooleanType; @@ -1174,11 +1175,11 @@ public void append(JsonParser parser, BlockBuilder blockBuilder) if (parser.getCurrentToken() != START_ARRAY) { throw new JsonCastException(format("Expected a json array, but got %s", parser.getText())); } - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - while (parser.nextToken() != END_ARRAY) { - elementAppender.append(parser, entryBuilder); - } - blockBuilder.closeEntry(); + ((ArrayBlockBuilder) blockBuilder).buildEntry(elementBuilder -> { + while (parser.nextToken() != END_ARRAY) { + elementAppender.append(parser, elementBuilder); + } + }); } } @@ -1206,20 +1207,26 @@ public void append(JsonParser parser, BlockBuilder blockBuilder) if (parser.getCurrentToken() != START_OBJECT) { throw new JsonCastException(format("Expected a json object, but got %s", parser.getText())); } - SingleMapBlockWriter entryBuilder = (SingleMapBlockWriter) blockBuilder.beginBlockEntry(); - entryBuilder.strict(); - while (parser.nextToken() != END_OBJECT) { - keyAppender.append(parser, entryBuilder); - parser.nextToken(); - valueAppender.append(parser, entryBuilder); - } + + MapBlockBuilder mapBlockBuilder = (MapBlockBuilder) blockBuilder; + mapBlockBuilder.strict(); try { - blockBuilder.closeEntry(); + mapBlockBuilder.buildEntry((keyBuilder, valueBuilder) -> appendMap(parser, keyBuilder, valueBuilder)); } catch (DuplicateMapKeyException e) { throw new JsonCastException("Duplicate keys are not allowed"); } } + + private void appendMap(JsonParser parser, BlockBuilder keyBuilder, BlockBuilder valueBuilder) + throws IOException + { + while (parser.nextToken() != END_OBJECT) { + keyAppender.append(parser, keyBuilder); + parser.nextToken(); + valueAppender.append(parser, valueBuilder); + } + } } private static class RowBlockBuilderAppender @@ -1247,12 +1254,7 @@ public void append(JsonParser parser, BlockBuilder blockBuilder) throw new JsonCastException(format("Expected a json array or object, but got %s", parser.getText())); } - parseJsonToSingleRowBlock( - parser, - (SingleRowBlockWriter) blockBuilder.beginBlockEntry(), - fieldAppenders, - fieldNameToIndex); - blockBuilder.closeEntry(); + ((RowBlockBuilder) blockBuilder).buildEntry(fieldBuilders -> parseJsonToSingleRowBlock(parser, fieldBuilders, fieldAppenders, fieldNameToIndex)); } } @@ -1272,9 +1274,9 @@ public static Optional> getFieldNameToIndex(List row // TODO: Once CAST function supports cachedInstanceFactory or directly write to BlockBuilder, // JsonToRowCast::toRow can use RowBlockBuilderAppender::append to parse JSON and append to the block builder. // Thus there will be single call to this method, so this method can be inlined. - public static void parseJsonToSingleRowBlock( + private static void parseJsonToSingleRowBlock( JsonParser parser, - SingleRowBlockWriter singleRowBlockWriter, + List fieldBuilders, BlockBuilderAppender[] fieldAppenders, Optional> fieldNameToIndex) throws IOException @@ -1282,7 +1284,7 @@ public static void parseJsonToSingleRowBlock( if (parser.getCurrentToken() == START_ARRAY) { for (int i = 0; i < fieldAppenders.length; i++) { parser.nextToken(); - fieldAppenders[i].append(parser, singleRowBlockWriter); + fieldAppenders[i].append(parser, fieldBuilders.get(i)); } if (parser.nextToken() != JsonToken.END_ARRAY) { throw new JsonCastException(format("Expected json array ending, but got %s", parser.getText())); @@ -1309,7 +1311,7 @@ public static void parseJsonToSingleRowBlock( } fieldWritten[fieldIndex] = true; numFieldsWritten++; - fieldAppenders[fieldIndex].append(parser, singleRowBlockWriter.getFieldBlockBuilder(fieldIndex)); + fieldAppenders[fieldIndex].append(parser, fieldBuilders.get(fieldIndex)); } else { parser.skipChildren(); @@ -1319,7 +1321,7 @@ public static void parseJsonToSingleRowBlock( if (numFieldsWritten != fieldAppenders.length) { for (int i = 0; i < fieldWritten.length; i++) { if (!fieldWritten[i]) { - singleRowBlockWriter.getFieldBlockBuilder(i).appendNull(); + fieldBuilders.get(i).appendNull(); } } } diff --git a/core/trino-main/src/test/java/io/trino/block/AbstractTestBlock.java b/core/trino-main/src/test/java/io/trino/block/AbstractTestBlock.java index 8c7fc3801d07..a56b02549d1f 100644 --- a/core/trino-main/src/test/java/io/trino/block/AbstractTestBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/AbstractTestBlock.java @@ -24,9 +24,10 @@ import io.trino.spi.block.BlockEncodingSerde; import io.trino.spi.block.DictionaryBlock; import io.trino.spi.block.DictionaryId; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.block.MapHashTables; -import io.trino.spi.block.SingleRowBlockWriter; import io.trino.spi.block.TestingBlockEncodingSerde; +import io.trino.spi.block.VariableWidthBlockBuilder; import java.lang.invoke.MethodHandle; import java.lang.reflect.Array; @@ -125,9 +126,6 @@ else if (type == BlockBuilder[].class || type == Block[].class) { else if (type == SliceOutput.class) { retainedSize += ((SliceOutput) field.get(block)).getRetainedSize(); } - else if (type == SingleRowBlockWriter.class) { - retainedSize += SingleRowBlockWriter.INSTANCE_SIZE; - } else if (type == int[].class) { retainedSize += sizeOf((int[]) field.get(block)); } @@ -149,12 +147,19 @@ else if (type == DictionaryId.class) { else if (type == MapHashTables.class) { retainedSize += ((MapHashTables) field.get(block)).getRetainedSizeInBytes(); } + else if (type.getEnclosingClass() == MapBlockBuilder.class) { + // ignore nested enum + } else if (type == MethodHandle.class) { // MethodHandles are only used in MapBlock/MapBlockBuilder, // and they are shared among blocks created by the same MapType. // So we don't account for the memory held onto by MethodHandle instances. // Otherwise, we will be counting it multiple times. } + else if (field.getName().equals("fieldBlockBuildersList")) { + // RowBlockBuilder fieldBlockBuildersList is a simple wrapper around the + // array already accounted for in the instance + } else { throw new IllegalArgumentException(format("Unknown type encountered: %s", type)); } @@ -344,9 +349,8 @@ protected void assertSlicePosition(Block block, int position, Slice expectedSlic assertTrue(block.equals(position, offset, expectedBlock, 0, offset, 3)); assertEquals(block.compareTo(position, offset, 3, expectedBlock, 0, offset, 3), 0); - BlockBuilder blockBuilder = VARBINARY.createBlockBuilder(null, 1); - block.writeBytesTo(position, offset, 3, blockBuilder); - blockBuilder.closeEntry(); + VariableWidthBlockBuilder blockBuilder = VARBINARY.createBlockBuilder(null, 1); + blockBuilder.writeEntry(block.getSlice(position, offset, 3)); Block segment = blockBuilder.build(); assertTrue(block.equals(position, offset, segment, 0, 0, 3)); diff --git a/core/trino-main/src/test/java/io/trino/block/BenchmarkMapCopy.java b/core/trino-main/src/test/java/io/trino/block/BenchmarkMapCopy.java index e3372a74c9d0..9fb639b32cfc 100644 --- a/core/trino-main/src/test/java/io/trino/block/BenchmarkMapCopy.java +++ b/core/trino-main/src/test/java/io/trino/block/BenchmarkMapCopy.java @@ -15,7 +15,7 @@ import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; -import io.trino.spi.block.BlockBuilderStatus; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.type.MapType; import org.openjdk.jmh.annotations.Benchmark; import org.openjdk.jmh.annotations.BenchmarkMode; @@ -72,8 +72,7 @@ public static class BenchmarkData private int mapSize; private Block dataBlock; - private BlockBuilder blockBuilder; - private BlockBuilderStatus status; + private MapBlockBuilder blockBuilder; @Setup public void setup() @@ -81,12 +80,12 @@ public void setup() MapType mapType = mapType(VARCHAR, BIGINT); blockBuilder = mapType.createBlockBuilder(null, POSITIONS); for (int position = 0; position < POSITIONS; position++) { - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (int i = 0; i < mapSize; i++) { - VARCHAR.writeString(entryBuilder, String.valueOf(ThreadLocalRandom.current().nextInt())); - BIGINT.writeLong(entryBuilder, ThreadLocalRandom.current().nextInt()); - } - blockBuilder.closeEntry(); + blockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + for (int i = 0; i < mapSize; i++) { + VARCHAR.writeString(keyBuilder, String.valueOf(ThreadLocalRandom.current().nextInt())); + BIGINT.writeLong(valueBuilder, ThreadLocalRandom.current().nextInt()); + } + }); } dataBlock = blockBuilder.build(); @@ -99,7 +98,7 @@ public Block getDataBlock() public BlockBuilder getBlockBuilder() { - return blockBuilder.newBlockBuilderLike(status); + return blockBuilder.newBlockBuilderLike(null); } } diff --git a/core/trino-main/src/test/java/io/trino/block/BenchmarkRowBlockBuilder.java b/core/trino-main/src/test/java/io/trino/block/BenchmarkRowBlockBuilder.java index 41e3d88af528..e95bc62247ad 100644 --- a/core/trino-main/src/test/java/io/trino/block/BenchmarkRowBlockBuilder.java +++ b/core/trino-main/src/test/java/io/trino/block/BenchmarkRowBlockBuilder.java @@ -14,7 +14,6 @@ package io.trino.block; import io.trino.spi.block.RowBlockBuilder; -import io.trino.spi.block.SingleRowBlockWriter; import io.trino.spi.type.RowType; import io.trino.spi.type.Type; import org.openjdk.jmh.annotations.Benchmark; @@ -53,12 +52,11 @@ public class BenchmarkRowBlockBuilder public void benchmarkBeginBlockEntry(BenchmarkData data, Blackhole blackhole) { for (int i = 0; i < data.rows; i++) { - SingleRowBlockWriter singleRowBlockWriter = data.getBlockBuilder().beginBlockEntry(); - for (int fieldIndex = 0; fieldIndex < data.getTypes().size(); fieldIndex++) { - singleRowBlockWriter.writeLong(data.getRandom().nextLong()).closeEntry(); - } - blackhole.consume(singleRowBlockWriter); - data.getBlockBuilder().closeEntry(); + data.getBlockBuilder().buildEntry(fieldBuilders -> { + for (int fieldIndex = 0; fieldIndex < data.getTypes().size(); fieldIndex++) { + BIGINT.writeLong(fieldBuilders.get(fieldIndex), data.getRandom().nextLong()); + } + }); } blackhole.consume(data.getBlockBuilder()); } diff --git a/core/trino-main/src/test/java/io/trino/block/BlockAssertions.java b/core/trino-main/src/test/java/io/trino/block/BlockAssertions.java index 98d1f1f5591d..22e7efb99509 100644 --- a/core/trino-main/src/test/java/io/trino/block/BlockAssertions.java +++ b/core/trino-main/src/test/java/io/trino/block/BlockAssertions.java @@ -603,48 +603,48 @@ public static Block createIntsBlock(Iterable values) public static Block createRowBlock(List fieldTypes, Object[]... rows) { - BlockBuilder rowBlockBuilder = new RowBlockBuilder(fieldTypes, null, 1); + RowBlockBuilder rowBlockBuilder = new RowBlockBuilder(fieldTypes, null, 1); for (Object[] row : rows) { if (row == null) { rowBlockBuilder.appendNull(); continue; } verify(row.length == fieldTypes.size()); - BlockBuilder singleRowBlockWriter = rowBlockBuilder.beginBlockEntry(); - for (int fieldIndex = 0; fieldIndex < fieldTypes.size(); fieldIndex++) { - Type fieldType = fieldTypes.get(fieldIndex); - Object fieldValue = row[fieldIndex]; - if (fieldValue == null) { - singleRowBlockWriter.appendNull(); - continue; + rowBlockBuilder.buildEntry(fieldBuilders -> { + for (int fieldIndex = 0; fieldIndex < fieldTypes.size(); fieldIndex++) { + Type fieldType = fieldTypes.get(fieldIndex); + Object fieldValue = row[fieldIndex]; + if (fieldValue == null) { + fieldBuilders.get(fieldIndex).appendNull(); + continue; + } + + if (fieldValue instanceof String) { + fieldType.writeSlice(fieldBuilders.get(fieldIndex), utf8Slice((String) fieldValue)); + } + else if (fieldValue instanceof Slice) { + fieldType.writeSlice(fieldBuilders.get(fieldIndex), (Slice) fieldValue); + } + else if (fieldValue instanceof Double) { + fieldType.writeDouble(fieldBuilders.get(fieldIndex), (Double) fieldValue); + } + else if (fieldValue instanceof Long) { + fieldType.writeLong(fieldBuilders.get(fieldIndex), (Long) fieldValue); + } + else if (fieldValue instanceof Boolean) { + fieldType.writeBoolean(fieldBuilders.get(fieldIndex), (Boolean) fieldValue); + } + else if (fieldValue instanceof Block) { + fieldType.writeObject(fieldBuilders.get(fieldIndex), fieldValue); + } + else if (fieldValue instanceof Integer) { + fieldType.writeLong(fieldBuilders.get(fieldIndex), (Integer) fieldValue); + } + else { + throw new IllegalArgumentException(); + } } - - if (fieldValue instanceof String) { - fieldType.writeSlice(singleRowBlockWriter, utf8Slice((String) fieldValue)); - } - else if (fieldValue instanceof Slice) { - fieldType.writeSlice(singleRowBlockWriter, (Slice) fieldValue); - } - else if (fieldValue instanceof Double) { - fieldType.writeDouble(singleRowBlockWriter, (Double) fieldValue); - } - else if (fieldValue instanceof Long) { - fieldType.writeLong(singleRowBlockWriter, (Long) fieldValue); - } - else if (fieldValue instanceof Boolean) { - fieldType.writeBoolean(singleRowBlockWriter, (Boolean) fieldValue); - } - else if (fieldValue instanceof Block) { - fieldType.writeObject(singleRowBlockWriter, fieldValue); - } - else if (fieldValue instanceof Integer) { - fieldType.writeLong(singleRowBlockWriter, (Integer) fieldValue); - } - else { - throw new IllegalArgumentException(); - } - } - rowBlockBuilder.closeEntry(); + }); } return rowBlockBuilder.build(); diff --git a/core/trino-main/src/test/java/io/trino/block/TestArrayBlock.java b/core/trino-main/src/test/java/io/trino/block/TestArrayBlock.java index 5ae7957d5580..476ffc4af855 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestArrayBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/TestArrayBlock.java @@ -21,6 +21,7 @@ import io.trino.spi.block.ByteArrayBlock; import org.testng.annotations.Test; +import java.util.Arrays; import java.util.Optional; import java.util.Random; @@ -130,7 +131,7 @@ public void testLazyBlockBuilderInitialization() for (int i = 0; i < ARRAY_SIZES.length; i++) { expectedValues[i] = rand.longs(ARRAY_SIZES[i]).toArray(); } - BlockBuilder emptyBlockBuilder = new ArrayBlockBuilder(BIGINT, null, 0, 0); + ArrayBlockBuilder emptyBlockBuilder = new ArrayBlockBuilder(BIGINT, null, 0, 0); BlockBuilder blockBuilder = new ArrayBlockBuilder(BIGINT, null, 100, 100); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); @@ -191,26 +192,22 @@ public void testCompactBlock() private static BlockBuilder createBlockBuilderWithValues(long[][][] expectedValues) { - BlockBuilder blockBuilder = new ArrayBlockBuilder(new ArrayBlockBuilder(BIGINT, null, 100, 100), null, 100); + ArrayBlockBuilder blockBuilder = new ArrayBlockBuilder(new ArrayBlockBuilder(BIGINT, null, 100, 100), null, 100); for (long[][] expectedValue : expectedValues) { if (expectedValue == null) { blockBuilder.appendNull(); } else { - BlockBuilder intermediateBlockBuilder = blockBuilder.beginBlockEntry(); - for (int j = 0; j < expectedValue.length; j++) { - if (expectedValue[j] == null) { - intermediateBlockBuilder.appendNull(); - } - else { - BlockBuilder innerMostBlockBuilder = intermediateBlockBuilder.beginBlockEntry(); - for (long v : expectedValue[j]) { - BIGINT.writeLong(innerMostBlockBuilder, v); + blockBuilder.buildEntry(elementBuilder -> { + for (long[] values : expectedValue) { + if (values == null) { + elementBuilder.appendNull(); + } + else { + ((ArrayBlockBuilder) elementBuilder).buildEntry(innerBuilder -> Arrays.stream(values).forEach(value -> BIGINT.writeLong(innerBuilder, value))); } - intermediateBlockBuilder.closeEntry(); } - } - blockBuilder.closeEntry(); + }); } } return blockBuilder; @@ -218,7 +215,7 @@ private static BlockBuilder createBlockBuilderWithValues(long[][][] expectedValu private static BlockBuilder createBlockBuilderWithValues(long[][] expectedValues) { - BlockBuilder blockBuilder = new ArrayBlockBuilder(BIGINT, null, 100, 100); + ArrayBlockBuilder blockBuilder = new ArrayBlockBuilder(BIGINT, null, 100, 100); return writeValues(expectedValues, blockBuilder); } @@ -229,11 +226,11 @@ private static BlockBuilder writeValues(long[][] expectedValues, BlockBuilder bl blockBuilder.appendNull(); } else { - BlockBuilder elementBlockBuilder = blockBuilder.beginBlockEntry(); - for (long v : expectedValue) { - BIGINT.writeLong(elementBlockBuilder, v); - } - blockBuilder.closeEntry(); + ((ArrayBlockBuilder) blockBuilder).buildEntry(elementBuilder -> { + for (long v : expectedValue) { + BIGINT.writeLong(elementBuilder, v); + } + }); } } return blockBuilder; @@ -241,17 +238,17 @@ private static BlockBuilder writeValues(long[][] expectedValues, BlockBuilder bl private static BlockBuilder createBlockBuilderWithValues(Slice[][] expectedValues) { - BlockBuilder blockBuilder = new ArrayBlockBuilder(VARCHAR, null, 100, 100); + ArrayBlockBuilder blockBuilder = new ArrayBlockBuilder(VARCHAR, null, 100, 100); for (Slice[] expectedValue : expectedValues) { if (expectedValue == null) { blockBuilder.appendNull(); } else { - BlockBuilder elementBlockBuilder = blockBuilder.beginBlockEntry(); - for (Slice v : expectedValue) { - VARCHAR.writeSlice(elementBlockBuilder, v); - } - blockBuilder.closeEntry(); + blockBuilder.buildEntry(elementBuilder -> { + for (Slice v : expectedValue) { + VARCHAR.writeSlice(elementBuilder, v); + } + }); } } return blockBuilder; diff --git a/core/trino-main/src/test/java/io/trino/block/TestBlockBuilder.java b/core/trino-main/src/test/java/io/trino/block/TestBlockBuilder.java index 75b08ccb4ec5..7817b093932b 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestBlockBuilder.java +++ b/core/trino-main/src/test/java/io/trino/block/TestBlockBuilder.java @@ -16,6 +16,7 @@ import com.google.common.collect.ImmutableList; import io.airlift.slice.Slices; import io.trino.spi.PageBuilder; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.type.ArrayType; @@ -55,20 +56,21 @@ public void testMultipleValuesWithNull() @Test public void testNewBlockBuilderLike() { - ArrayType longArrayType = new ArrayType(BIGINT); - ArrayType arrayType = new ArrayType(longArrayType); - List channels = ImmutableList.of(BIGINT, VARCHAR, arrayType); + List channels = ImmutableList.of(BIGINT, VARCHAR, new ArrayType(new ArrayType(BIGINT))); PageBuilder pageBuilder = new PageBuilder(channels); BlockBuilder bigintBlockBuilder = pageBuilder.getBlockBuilder(0); BlockBuilder varcharBlockBuilder = pageBuilder.getBlockBuilder(1); - BlockBuilder arrayBlockBuilder = pageBuilder.getBlockBuilder(2); + ArrayBlockBuilder arrayBlockBuilder = (ArrayBlockBuilder) pageBuilder.getBlockBuilder(2); for (int i = 0; i < 100; i++) { - BIGINT.writeLong(bigintBlockBuilder, i); - VARCHAR.writeSlice(varcharBlockBuilder, Slices.utf8Slice("test" + i)); - BlockBuilder blockBuilder = longArrayType.createBlockBuilder(null, 1); - longArrayType.writeObject(blockBuilder, BIGINT.createBlockBuilder(null, 2).writeLong(i).writeLong(i * 2).build()); - arrayType.writeObject(arrayBlockBuilder, blockBuilder); + int value = i; + BIGINT.writeLong(bigintBlockBuilder, value); + VARCHAR.writeSlice(varcharBlockBuilder, Slices.utf8Slice("test" + value)); + arrayBlockBuilder.buildEntry(elementBuilder -> { + ArrayBlockBuilder nestedArrayBuilder = (ArrayBlockBuilder) elementBuilder; + nestedArrayBuilder.buildEntry(valueBuilder -> BIGINT.writeLong(valueBuilder, value)); + nestedArrayBuilder.buildEntry(valueBuilder -> BIGINT.writeLong(valueBuilder, value * 2)); + }); pageBuilder.declarePosition(); } @@ -85,15 +87,20 @@ public void testNewBlockBuilderLike() @Test public void testGetPositions() { - BlockBuilder blockBuilder = BIGINT.createFixedSizeBlockBuilder(5).appendNull().writeLong(42L).appendNull().writeLong(43L).appendNull(); + BlockBuilder blockBuilder = BIGINT.createFixedSizeBlockBuilder(5); + blockBuilder.appendNull(); + BIGINT.writeLong(blockBuilder, 42L); + blockBuilder.appendNull(); + BIGINT.writeLong(blockBuilder, 43L); + blockBuilder.appendNull(); int[] positions = new int[] {0, 1, 1, 1, 4}; // test getPositions for block builder - assertBlockEquals(BIGINT, blockBuilder.getPositions(positions, 0, positions.length), BIGINT.createFixedSizeBlockBuilder(5).appendNull().writeLong(42).writeLong(42).writeLong(42).appendNull().build()); - assertBlockEquals(BIGINT, blockBuilder.getPositions(positions, 1, 4), BIGINT.createFixedSizeBlockBuilder(5).writeLong(42).writeLong(42).writeLong(42).appendNull().build()); - assertBlockEquals(BIGINT, blockBuilder.getPositions(positions, 2, 1), BIGINT.createFixedSizeBlockBuilder(5).writeLong(42).build()); - assertBlockEquals(BIGINT, blockBuilder.getPositions(positions, 0, 0), BIGINT.createFixedSizeBlockBuilder(5).build()); - assertBlockEquals(BIGINT, blockBuilder.getPositions(positions, 1, 0), BIGINT.createFixedSizeBlockBuilder(5).build()); + assertBlockEquals(BIGINT, blockBuilder.getPositions(positions, 0, positions.length), buildBigintBlock(null, 42, 42, 42, null)); + assertBlockEquals(BIGINT, blockBuilder.getPositions(positions, 1, 4), buildBigintBlock(42, 42, 42, null)); + assertBlockEquals(BIGINT, blockBuilder.getPositions(positions, 2, 1), buildBigintBlock(42)); + assertBlockEquals(BIGINT, blockBuilder.getPositions(positions, 0, 0), buildBigintBlock()); + assertBlockEquals(BIGINT, blockBuilder.getPositions(positions, 1, 0), buildBigintBlock()); // out of range assertInvalidPosition(blockBuilder, new int[] {-1}, 0, 1); @@ -104,11 +111,11 @@ public void testGetPositions() // test getPositions for block Block block = blockBuilder.build(); - assertBlockEquals(BIGINT, block.getPositions(positions, 0, positions.length), BIGINT.createFixedSizeBlockBuilder(5).appendNull().writeLong(42).writeLong(42).writeLong(42).appendNull().build()); - assertBlockEquals(BIGINT, block.getPositions(positions, 1, 4), BIGINT.createFixedSizeBlockBuilder(5).writeLong(42).writeLong(42).writeLong(42).appendNull().build()); - assertBlockEquals(BIGINT, block.getPositions(positions, 2, 1), BIGINT.createFixedSizeBlockBuilder(5).writeLong(42).build()); - assertBlockEquals(BIGINT, block.getPositions(positions, 0, 0), BIGINT.createFixedSizeBlockBuilder(5).build()); - assertBlockEquals(BIGINT, block.getPositions(positions, 1, 0), BIGINT.createFixedSizeBlockBuilder(5).build()); + assertBlockEquals(BIGINT, block.getPositions(positions, 0, positions.length), buildBigintBlock(null, 42, 42, 42, null)); + assertBlockEquals(BIGINT, block.getPositions(positions, 1, 4), buildBigintBlock(42, 42, 42, null)); + assertBlockEquals(BIGINT, block.getPositions(positions, 2, 1), buildBigintBlock(42)); + assertBlockEquals(BIGINT, block.getPositions(positions, 0, 0), buildBigintBlock()); + assertBlockEquals(BIGINT, block.getPositions(positions, 1, 0), buildBigintBlock()); // out of range assertInvalidPosition(block, new int[] {-1}, 0, 1); @@ -127,6 +134,20 @@ public void testGetPositions() assertTrue(isIdentical.get()); } + private static Block buildBigintBlock(Integer... values) + { + BlockBuilder blockBuilder = BIGINT.createFixedSizeBlockBuilder(5); + for (Integer value : values) { + if (value == null) { + blockBuilder.appendNull(); + } + else { + BIGINT.writeLong(blockBuilder, value); + } + } + return blockBuilder.build(); + } + private static void assertInvalidPosition(Block block, int[] positions, int offset, int length) { assertThatThrownBy(() -> block.getPositions(positions, offset, length).getLong(0, 0)) diff --git a/core/trino-main/src/test/java/io/trino/block/TestByteArrayBlock.java b/core/trino-main/src/test/java/io/trino/block/TestByteArrayBlock.java index 8a6d42c55043..dfe0338ed476 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestByteArrayBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/TestByteArrayBlock.java @@ -50,7 +50,7 @@ public void testLazyBlockBuilderInitialization() Slice[] expectedValues = createTestValue(100); BlockBuilder emptyBlockBuilder = new ByteArrayBlockBuilder(null, 0); - BlockBuilder blockBuilder = new ByteArrayBlockBuilder(null, expectedValues.length); + ByteArrayBlockBuilder blockBuilder = new ByteArrayBlockBuilder(null, expectedValues.length); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); assertEquals(blockBuilder.getRetainedSizeInBytes(), emptyBlockBuilder.getRetainedSizeInBytes()); @@ -58,7 +58,7 @@ public void testLazyBlockBuilderInitialization() assertTrue(blockBuilder.getSizeInBytes() > emptyBlockBuilder.getSizeInBytes()); assertTrue(blockBuilder.getRetainedSizeInBytes() > emptyBlockBuilder.getRetainedSizeInBytes()); - blockBuilder = blockBuilder.newBlockBuilderLike(null); + blockBuilder = (ByteArrayBlockBuilder) blockBuilder.newBlockBuilderLike(null); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); assertEquals(blockBuilder.getRetainedSizeInBytes(), emptyBlockBuilder.getRetainedSizeInBytes()); } @@ -95,14 +95,14 @@ private static BlockBuilder createBlockBuilderWithValues(Slice[] expectedValues) return blockBuilder; } - private static void writeValues(Slice[] expectedValues, BlockBuilder blockBuilder) + private static void writeValues(Slice[] expectedValues, ByteArrayBlockBuilder blockBuilder) { for (Slice expectedValue : expectedValues) { if (expectedValue == null) { blockBuilder.appendNull(); } else { - blockBuilder.writeByte(expectedValue.getByte(0)).closeEntry(); + blockBuilder.writeByte(expectedValue.getByte(0)); } } } diff --git a/core/trino-main/src/test/java/io/trino/block/TestColumnarMap.java b/core/trino-main/src/test/java/io/trino/block/TestColumnarMap.java index 2d316822f463..ed526b940105 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestColumnarMap.java +++ b/core/trino-main/src/test/java/io/trino/block/TestColumnarMap.java @@ -147,34 +147,33 @@ private static void assertColumnarMap(Block block, Slice[][][] expectedValues) public static BlockBuilder createBlockBuilderWithValues(Slice[][][] expectedValues) { - BlockBuilder blockBuilder = createMapBuilder(100); + MapBlockBuilder blockBuilder = createMapBuilder(100); for (Slice[][] expectedMap : expectedValues) { if (expectedMap == null) { blockBuilder.appendNull(); } else { - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - VARCHAR.createBlockBuilder(null, expectedMap.length); - for (Slice[] entry : expectedMap) { - Slice key = entry[0]; - assertNotNull(key); - VARCHAR.writeSlice(entryBuilder, key); - - Slice value = entry[1]; - if (value == null) { - entryBuilder.appendNull(); + blockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + for (Slice[] entry : expectedMap) { + Slice key = entry[0]; + assertNotNull(key); + VARCHAR.writeSlice(keyBuilder, key); + + Slice value = entry[1]; + if (value == null) { + valueBuilder.appendNull(); + } + else { + VARCHAR.writeSlice(valueBuilder, value); + } } - else { - VARCHAR.writeSlice(entryBuilder, value); - } - } - blockBuilder.closeEntry(); + }); } } return blockBuilder; } - private static BlockBuilder createMapBuilder(int expectedEntries) + private static MapBlockBuilder createMapBuilder(int expectedEntries) { MapType mapType = (MapType) TESTING_TYPE_MANAGER.getType(new TypeSignature(MAP, TypeSignatureParameter.typeParameter(VARCHAR.getTypeSignature()), TypeSignatureParameter.typeParameter(VARCHAR.getTypeSignature()))); return new MapBlockBuilder(mapType, null, expectedEntries); diff --git a/core/trino-main/src/test/java/io/trino/block/TestColumnarRow.java b/core/trino-main/src/test/java/io/trino/block/TestColumnarRow.java index 6c39000aef6e..60cb8cd9cbcc 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestColumnarRow.java +++ b/core/trino-main/src/test/java/io/trino/block/TestColumnarRow.java @@ -131,28 +131,29 @@ private static void assertColumnarRow(Block block, T[] expectedValues) public static BlockBuilder createBlockBuilderWithValues(Slice[][] expectedValues) { - BlockBuilder blockBuilder = createBlockBuilder(null, 100); + RowBlockBuilder blockBuilder = createBlockBuilder(null, 100); for (Slice[] expectedValue : expectedValues) { if (expectedValue == null) { blockBuilder.appendNull(); } else { - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (Slice v : expectedValue) { - if (v == null) { - entryBuilder.appendNull(); + blockBuilder.buildEntry(fieldBuilders -> { + for (int i = 0; i < expectedValue.length; i++) { + Slice v = expectedValue[i]; + if (v == null) { + fieldBuilders.get(i).appendNull(); + } + else { + VARCHAR.writeSlice(fieldBuilders.get(i), v); + } } - else { - VARCHAR.writeSlice(entryBuilder, v); - } - } - blockBuilder.closeEntry(); + }); } } return blockBuilder; } - private static BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) + private static RowBlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) { return new RowBlockBuilder(Collections.nCopies(FIELD_COUNT, VARCHAR), blockBuilderStatus, expectedEntries); } diff --git a/core/trino-main/src/test/java/io/trino/block/TestFixed12Block.java b/core/trino-main/src/test/java/io/trino/block/TestFixed12Block.java index 4be36b36282c..50202f59f88e 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestFixed12Block.java +++ b/core/trino-main/src/test/java/io/trino/block/TestFixed12Block.java @@ -18,7 +18,6 @@ import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.Fixed12Block; import io.trino.spi.block.Fixed12BlockBuilder; -import io.trino.spi.block.VariableWidthBlockBuilder; import org.testng.annotations.Test; import java.util.Optional; @@ -50,9 +49,9 @@ public void testCopyPositions() public void testLazyBlockBuilderInitialization() { Slice[] expectedValues = createTestValue(100); - BlockBuilder emptyBlockBuilder = new VariableWidthBlockBuilder(null, 0, 0); + Fixed12BlockBuilder emptyBlockBuilder = new Fixed12BlockBuilder(null, 0); - BlockBuilder blockBuilder = new VariableWidthBlockBuilder(null, expectedValues.length, 32 * expectedValues.length); + Fixed12BlockBuilder blockBuilder = new Fixed12BlockBuilder(null, expectedValues.length); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); assertEquals(blockBuilder.getRetainedSizeInBytes(), emptyBlockBuilder.getRetainedSizeInBytes()); @@ -60,7 +59,7 @@ public void testLazyBlockBuilderInitialization() assertTrue(blockBuilder.getSizeInBytes() > emptyBlockBuilder.getSizeInBytes()); assertTrue(blockBuilder.getRetainedSizeInBytes() > emptyBlockBuilder.getRetainedSizeInBytes()); - blockBuilder = blockBuilder.newBlockBuilderLike(null); + blockBuilder = (Fixed12BlockBuilder) blockBuilder.newBlockBuilderLike(null); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); assertEquals(blockBuilder.getRetainedSizeInBytes(), emptyBlockBuilder.getRetainedSizeInBytes()); } @@ -103,16 +102,16 @@ private static BlockBuilder createBlockBuilderWithValues(Slice[] expectedValues) return blockBuilder; } - private static void writeValues(Slice[] expectedValues, BlockBuilder blockBuilder) + private static void writeValues(Slice[] expectedValues, Fixed12BlockBuilder blockBuilder) { for (Slice expectedValue : expectedValues) { if (expectedValue == null) { blockBuilder.appendNull(); } else { - blockBuilder.writeLong(expectedValue.getLong(0)); - blockBuilder.writeInt(expectedValue.getInt(8)); - blockBuilder.closeEntry(); + blockBuilder.writeFixed12( + expectedValue.getLong(0), + expectedValue.getInt(8)); } } } diff --git a/core/trino-main/src/test/java/io/trino/block/TestInt128ArrayBlock.java b/core/trino-main/src/test/java/io/trino/block/TestInt128ArrayBlock.java index 7c7e44429d82..23ac0fceb7fe 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestInt128ArrayBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/TestInt128ArrayBlock.java @@ -17,7 +17,6 @@ import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.Int128ArrayBlock; import io.trino.spi.block.Int128ArrayBlockBuilder; -import io.trino.spi.block.VariableWidthBlockBuilder; import org.testng.annotations.Test; import java.util.Optional; @@ -49,9 +48,9 @@ public void testCopyPositions() public void testLazyBlockBuilderInitialization() { Slice[] expectedValues = createTestValue(100); - BlockBuilder emptyBlockBuilder = new VariableWidthBlockBuilder(null, 0, 0); + BlockBuilder emptyBlockBuilder = new Int128ArrayBlockBuilder(null, 0); - BlockBuilder blockBuilder = new VariableWidthBlockBuilder(null, expectedValues.length, 32 * expectedValues.length); + BlockBuilder blockBuilder = new Int128ArrayBlockBuilder(null, expectedValues.length); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); assertEquals(blockBuilder.getRetainedSizeInBytes(), emptyBlockBuilder.getRetainedSizeInBytes()); @@ -103,9 +102,9 @@ private static void writeValues(Slice[] expectedValues, BlockBuilder blockBuilde blockBuilder.appendNull(); } else { - blockBuilder.writeLong(expectedValue.getLong(0)); - blockBuilder.writeLong(expectedValue.getLong(8)); - blockBuilder.closeEntry(); + ((Int128ArrayBlockBuilder) blockBuilder).writeInt128( + expectedValue.getLong(0), + expectedValue.getLong(8)); } } } diff --git a/core/trino-main/src/test/java/io/trino/block/TestIntArrayBlock.java b/core/trino-main/src/test/java/io/trino/block/TestIntArrayBlock.java index c1c6aab457ba..983e114f1adf 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestIntArrayBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/TestIntArrayBlock.java @@ -50,7 +50,7 @@ public void testLazyBlockBuilderInitialization() Slice[] expectedValues = createTestValue(100); BlockBuilder emptyBlockBuilder = new IntArrayBlockBuilder(null, 0); - BlockBuilder blockBuilder = new IntArrayBlockBuilder(null, expectedValues.length); + IntArrayBlockBuilder blockBuilder = new IntArrayBlockBuilder(null, expectedValues.length); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); assertEquals(blockBuilder.getRetainedSizeInBytes(), emptyBlockBuilder.getRetainedSizeInBytes()); @@ -58,7 +58,7 @@ public void testLazyBlockBuilderInitialization() assertTrue(blockBuilder.getSizeInBytes() > emptyBlockBuilder.getSizeInBytes()); assertTrue(blockBuilder.getRetainedSizeInBytes() > emptyBlockBuilder.getRetainedSizeInBytes()); - blockBuilder = blockBuilder.newBlockBuilderLike(null); + blockBuilder = (IntArrayBlockBuilder) blockBuilder.newBlockBuilderLike(null); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); assertEquals(blockBuilder.getRetainedSizeInBytes(), emptyBlockBuilder.getRetainedSizeInBytes()); } @@ -95,14 +95,14 @@ private static BlockBuilder createBlockBuilderWithValues(Slice[] expectedValues) return blockBuilder; } - private static void writeValues(Slice[] expectedValues, BlockBuilder blockBuilder) + private static void writeValues(Slice[] expectedValues, IntArrayBlockBuilder blockBuilder) { for (Slice expectedValue : expectedValues) { if (expectedValue == null) { blockBuilder.appendNull(); } else { - blockBuilder.writeInt(expectedValue.getInt(0)).closeEntry(); + blockBuilder.writeInt(expectedValue.getInt(0)); } } } diff --git a/core/trino-main/src/test/java/io/trino/block/TestLongArrayBlock.java b/core/trino-main/src/test/java/io/trino/block/TestLongArrayBlock.java index c95ad5795ff6..8421932a9021 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestLongArrayBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/TestLongArrayBlock.java @@ -17,7 +17,6 @@ import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.LongArrayBlock; import io.trino.spi.block.LongArrayBlockBuilder; -import io.trino.spi.block.VariableWidthBlockBuilder; import org.testng.annotations.Test; import java.util.Optional; @@ -49,9 +48,9 @@ public void testCopyPositions() public void testLazyBlockBuilderInitialization() { Slice[] expectedValues = createTestValue(100); - BlockBuilder emptyBlockBuilder = new VariableWidthBlockBuilder(null, 0, 0); + BlockBuilder emptyBlockBuilder = new LongArrayBlockBuilder(null, 0); - BlockBuilder blockBuilder = new VariableWidthBlockBuilder(null, expectedValues.length, 32 * expectedValues.length); + LongArrayBlockBuilder blockBuilder = new LongArrayBlockBuilder(null, expectedValues.length); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); assertEquals(blockBuilder.getRetainedSizeInBytes(), emptyBlockBuilder.getRetainedSizeInBytes()); @@ -59,7 +58,7 @@ public void testLazyBlockBuilderInitialization() assertTrue(blockBuilder.getSizeInBytes() > emptyBlockBuilder.getSizeInBytes()); assertTrue(blockBuilder.getRetainedSizeInBytes() > emptyBlockBuilder.getRetainedSizeInBytes()); - blockBuilder = blockBuilder.newBlockBuilderLike(null); + blockBuilder = (LongArrayBlockBuilder) blockBuilder.newBlockBuilderLike(null); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); assertEquals(blockBuilder.getRetainedSizeInBytes(), emptyBlockBuilder.getRetainedSizeInBytes()); } @@ -96,14 +95,14 @@ private static BlockBuilder createBlockBuilderWithValues(Slice[] expectedValues) return blockBuilder; } - private static void writeValues(Slice[] expectedValues, BlockBuilder blockBuilder) + private static void writeValues(Slice[] expectedValues, LongArrayBlockBuilder blockBuilder) { for (Slice expectedValue : expectedValues) { if (expectedValue == null) { blockBuilder.appendNull(); } else { - blockBuilder.writeLong(expectedValue.getLong(0)).closeEntry(); + blockBuilder.writeLong(expectedValue.getLong(0)); } } } diff --git a/core/trino-main/src/test/java/io/trino/block/TestMapBlock.java b/core/trino-main/src/test/java/io/trino/block/TestMapBlock.java index 767dce557022..828e5eb7c957 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestMapBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/TestMapBlock.java @@ -239,7 +239,7 @@ private void testWith(Map[] expectedValues) private BlockBuilder createBlockBuilderWithValues(Map[] maps) { MapType mapType = mapType(VARCHAR, BIGINT); - BlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1); + MapBlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1); for (Map map : maps) { createBlockBuilderWithValues(map, mapBlockBuilder); } @@ -276,23 +276,23 @@ private MapBlock createBlockWithValuesFromKeyValueBlock(Map[] maps createLongsBlock(values)); } - private void createBlockBuilderWithValues(Map map, BlockBuilder mapBlockBuilder) + private void createBlockBuilderWithValues(Map map, MapBlockBuilder mapBlockBuilder) { if (map == null) { mapBlockBuilder.appendNull(); } else { - BlockBuilder elementBlockBuilder = mapBlockBuilder.beginBlockEntry(); - for (Map.Entry entry : map.entrySet()) { - VARCHAR.writeSlice(elementBlockBuilder, utf8Slice(entry.getKey())); - if (entry.getValue() == null) { - elementBlockBuilder.appendNull(); - } - else { - BIGINT.writeLong(elementBlockBuilder, entry.getValue()); + mapBlockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + for (Map.Entry entry : map.entrySet()) { + VARCHAR.writeSlice(keyBuilder, utf8Slice(entry.getKey())); + if (entry.getValue() == null) { + valueBuilder.appendNull(); + } + else { + BIGINT.writeLong(valueBuilder, entry.getValue()); + } } - } - mapBlockBuilder.closeEntry(); + }); } } @@ -353,61 +353,37 @@ private void assertValue(Block mapBlock, int position, Map map) public void testStrict() { MapType mapType = mapType(BIGINT, BIGINT); - MapBlockBuilder mapBlockBuilder = (MapBlockBuilder) mapType.createBlockBuilder(null, 1); + MapBlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1); mapBlockBuilder.strict(); // Add 100 maps with only one entry but the same key for (int i = 0; i < 100; i++) { - BlockBuilder entryBuilder = mapBlockBuilder.beginBlockEntry(); - BIGINT.writeLong(entryBuilder, 1); - BIGINT.writeLong(entryBuilder, -1); - mapBlockBuilder.closeEntry(); - } - - BlockBuilder entryBuilder = mapBlockBuilder.beginBlockEntry(); - // Add 50 keys so we get some chance to get hash conflict - // The purpose of this test is to make sure offset is calculated correctly in MapBlockBuilder.closeEntryStrict() - for (int i = 0; i < 50; i++) { - BIGINT.writeLong(entryBuilder, i); - BIGINT.writeLong(entryBuilder, -1); + mapBlockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + BIGINT.writeLong(keyBuilder, 1); + BIGINT.writeLong(valueBuilder, -1); + }); } - mapBlockBuilder.closeEntry(); - entryBuilder = mapBlockBuilder.beginBlockEntry(); - for (int i = 0; i < 2; i++) { - BIGINT.writeLong(entryBuilder, 99); - BIGINT.writeLong(entryBuilder, -1); - } - assertThatThrownBy(mapBlockBuilder::closeEntry) + mapBlockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + // Add 50 keys so we get some chance to get hash conflict + // The purpose of this test is to make sure offset is calculated correctly in MapBlockBuilder.closeEntryStrict() + for (int i = 0; i < 50; i++) { + BIGINT.writeLong(keyBuilder, i); + BIGINT.writeLong(valueBuilder, -1); + } + }); + + assertThatThrownBy( + () -> mapBlockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + for (int i = 0; i < 2; i++) { + BIGINT.writeLong(keyBuilder, 99); + BIGINT.writeLong(valueBuilder, -1); + } + })) .isInstanceOf(DuplicateMapKeyException.class) .hasMessage("Duplicate map keys are not allowed"); } - @Test - public void testCloseEntryStrict() - throws Exception - { - MapType mapType = mapType(BIGINT, BIGINT); - MapBlockBuilder mapBlockBuilder = (MapBlockBuilder) mapType.createBlockBuilder(null, 1); - - // Add 100 maps with only one entry but the same key - for (int i = 0; i < 100; i++) { - BlockBuilder entryBuilder = mapBlockBuilder.beginBlockEntry(); - BIGINT.writeLong(entryBuilder, 1); - BIGINT.writeLong(entryBuilder, -1); - mapBlockBuilder.closeEntry(); - } - - BlockBuilder entryBuilder = mapBlockBuilder.beginBlockEntry(); - // Add 50 keys so we get some chance to get hash conflict - // The purpose of this test is to make sure offset is calculated correctly in MapBlockBuilder.closeEntryStrict() - for (int i = 0; i < 50; i++) { - BIGINT.writeLong(entryBuilder, i); - BIGINT.writeLong(entryBuilder, -1); - } - mapBlockBuilder.closeEntryStrict(); - } - @Test public void testEstimatedDataSizeForStats() { diff --git a/core/trino-main/src/test/java/io/trino/block/TestRowBlock.java b/core/trino-main/src/test/java/io/trino/block/TestRowBlock.java index bdc07dd4d62d..a8504538660d 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestRowBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/TestRowBlock.java @@ -135,30 +135,31 @@ private void testWith(List fieldTypes, List[] expectedValues) private BlockBuilder createBlockBuilderWithValues(List fieldTypes, List[] rows) { - BlockBuilder rowBlockBuilder = new RowBlockBuilder(fieldTypes, null, 1); + RowBlockBuilder rowBlockBuilder = new RowBlockBuilder(fieldTypes, null, 1); for (List row : rows) { if (row == null) { rowBlockBuilder.appendNull(); } else { - BlockBuilder singleRowBlockWriter = rowBlockBuilder.beginBlockEntry(); - for (Object fieldValue : row) { - if (fieldValue == null) { - singleRowBlockWriter.appendNull(); - } - else { - if (fieldValue instanceof Long) { - BIGINT.writeLong(singleRowBlockWriter, ((Long) fieldValue).longValue()); - } - else if (fieldValue instanceof String) { - VARCHAR.writeSlice(singleRowBlockWriter, utf8Slice((String) fieldValue)); + rowBlockBuilder.buildEntry(fieldBuilders -> { + for (int i = 0; i < row.size(); i++) { + Object fieldValue = row.get(i); + if (fieldValue == null) { + fieldBuilders.get(i).appendNull(); } else { - throw new IllegalArgumentException(); + if (fieldValue instanceof Long) { + BIGINT.writeLong(fieldBuilders.get(i), ((Long) fieldValue).longValue()); + } + else if (fieldValue instanceof String) { + VARCHAR.writeSlice(fieldBuilders.get(i), utf8Slice((String) fieldValue)); + } + else { + throw new IllegalArgumentException(); + } } } - } - rowBlockBuilder.closeEntry(); + }); } } diff --git a/core/trino-main/src/test/java/io/trino/block/TestRunLengthEncodedBlock.java b/core/trino-main/src/test/java/io/trino/block/TestRunLengthEncodedBlock.java index 9268a96b547d..da184d934c84 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestRunLengthEncodedBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/TestRunLengthEncodedBlock.java @@ -52,8 +52,8 @@ private void assertRleBlock(int positionCount) private static Block createSingleValueBlock(Slice expectedValue) { - BlockBuilder blockBuilder = new VariableWidthBlockBuilder(null, 1, expectedValue.length()); - blockBuilder.writeBytes(expectedValue, 0, expectedValue.length()).closeEntry(); + VariableWidthBlockBuilder blockBuilder = new VariableWidthBlockBuilder(null, 1, expectedValue.length()); + blockBuilder.writeEntry(expectedValue); return blockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/block/TestShortArrayBlock.java b/core/trino-main/src/test/java/io/trino/block/TestShortArrayBlock.java index 246efd161c80..578fb2f1cf67 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestShortArrayBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/TestShortArrayBlock.java @@ -48,9 +48,9 @@ public void testCopyPositions() public void testLazyBlockBuilderInitialization() { Slice[] expectedValues = createTestValue(100); - BlockBuilder emptyBlockBuilder = new ShortArrayBlockBuilder(null, 0); + ShortArrayBlockBuilder emptyBlockBuilder = new ShortArrayBlockBuilder(null, 0); - BlockBuilder blockBuilder = new ShortArrayBlockBuilder(null, expectedValues.length); + ShortArrayBlockBuilder blockBuilder = new ShortArrayBlockBuilder(null, expectedValues.length); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); assertEquals(blockBuilder.getRetainedSizeInBytes(), emptyBlockBuilder.getRetainedSizeInBytes()); @@ -58,7 +58,7 @@ public void testLazyBlockBuilderInitialization() assertTrue(blockBuilder.getSizeInBytes() > emptyBlockBuilder.getSizeInBytes()); assertTrue(blockBuilder.getRetainedSizeInBytes() > emptyBlockBuilder.getRetainedSizeInBytes()); - blockBuilder = blockBuilder.newBlockBuilderLike(null); + blockBuilder = (ShortArrayBlockBuilder) blockBuilder.newBlockBuilderLike(null); assertEquals(blockBuilder.getSizeInBytes(), emptyBlockBuilder.getSizeInBytes()); assertEquals(blockBuilder.getRetainedSizeInBytes(), emptyBlockBuilder.getRetainedSizeInBytes()); } @@ -95,14 +95,14 @@ private static BlockBuilder createBlockBuilderWithValues(Slice[] expectedValues) return blockBuilder; } - private static void writeValues(Slice[] expectedValues, BlockBuilder blockBuilder) + private static void writeValues(Slice[] expectedValues, ShortArrayBlockBuilder blockBuilder) { for (Slice expectedValue : expectedValues) { if (expectedValue == null) { blockBuilder.appendNull(); } else { - blockBuilder.writeShort(expectedValue.getShort(0)).closeEntry(); + blockBuilder.writeShort(expectedValue.getShort(0)); } } } diff --git a/core/trino-main/src/test/java/io/trino/block/TestSingleRowBlockWriter.java b/core/trino-main/src/test/java/io/trino/block/TestSingleRowBlockWriter.java deleted file mode 100644 index ed1185edf079..000000000000 --- a/core/trino-main/src/test/java/io/trino/block/TestSingleRowBlockWriter.java +++ /dev/null @@ -1,66 +0,0 @@ -/* - * Licensed 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.trino.block; - -import com.google.common.collect.ImmutableList; -import io.trino.spi.block.RowBlockBuilder; -import io.trino.spi.block.SingleRowBlockWriter; -import io.trino.spi.type.RowType; -import io.trino.spi.type.Type; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Test; - -import java.util.List; - -import static io.trino.spi.type.BigintType.BIGINT; -import static io.trino.spi.type.BooleanType.BOOLEAN; -import static org.testng.Assert.assertEquals; - -public class TestSingleRowBlockWriter -{ - private RowBlockBuilder rowBlockBuilder; - - @BeforeClass - public void setup() - { - List types = ImmutableList.of(BIGINT, BOOLEAN); - rowBlockBuilder = (RowBlockBuilder) RowType.anonymous(types).createBlockBuilder(null, 8); - } - - @Test - public void testGetSizeInBytes() - { - SingleRowBlockWriter singleRowBlockWriter = rowBlockBuilder.beginBlockEntry(); - // Test whether new singleRowBlockWriter has size equal to 0 - assertEquals(0, singleRowBlockWriter.getSizeInBytes()); - - singleRowBlockWriter.writeLong(10).closeEntry(); - assertEquals(9, singleRowBlockWriter.getSizeInBytes()); - - singleRowBlockWriter.writeByte(10).closeEntry(); - assertEquals(11, singleRowBlockWriter.getSizeInBytes()); - rowBlockBuilder.closeEntry(); - - // Test whether previous entry does not mix to the next entry (for size). Does reset works on size? - singleRowBlockWriter = rowBlockBuilder.beginBlockEntry(); - assertEquals(0, singleRowBlockWriter.getSizeInBytes()); - - singleRowBlockWriter.writeLong(10).closeEntry(); - assertEquals(9, singleRowBlockWriter.getSizeInBytes()); - - singleRowBlockWriter.writeByte(10).closeEntry(); - assertEquals(11, singleRowBlockWriter.getSizeInBytes()); - rowBlockBuilder.closeEntry(); - } -} diff --git a/core/trino-main/src/test/java/io/trino/block/TestVariableWidthBlock.java b/core/trino-main/src/test/java/io/trino/block/TestVariableWidthBlock.java index 3c0b2cb0b3ad..35d4ab2b716a 100644 --- a/core/trino-main/src/test/java/io/trino/block/TestVariableWidthBlock.java +++ b/core/trino-main/src/test/java/io/trino/block/TestVariableWidthBlock.java @@ -154,7 +154,7 @@ private static BlockBuilder writeValues(Slice[] expectedValues, BlockBuilder blo blockBuilder.appendNull(); } else { - blockBuilder.writeBytes(expectedValue, 0, expectedValue.length()).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(expectedValue); } } return blockBuilder; diff --git a/core/trino-main/src/test/java/io/trino/connector/TestingTableFunctions.java b/core/trino-main/src/test/java/io/trino/connector/TestingTableFunctions.java index 24204b767eee..f00edbfaf7df 100644 --- a/core/trino-main/src/test/java/io/trino/connector/TestingTableFunctions.java +++ b/core/trino-main/src/test/java/io/trino/connector/TestingTableFunctions.java @@ -682,7 +682,7 @@ public TableFunctionProcessorState process(List> input) BlockBuilder builder = BIGINT.createBlockBuilder(null, page.getPositionCount()); for (long index = processedPositions; index < processedPositions + page.getPositionCount(); index++) { // TODO check for long overflow - builder.writeLong(index); + BIGINT.writeLong(builder, index); } processedPositions = processedPositions + page.getPositionCount(); return usedInputAndProduced(new Page(builder.build())); @@ -789,7 +789,7 @@ public TableFunctionProcessorState process(List> input) BlockBuilder builder = BIGINT.createBlockBuilder(null, page.getPositionCount()); for (long index = processedPositions; index < processedPositions + page.getPositionCount(); index++) { // TODO check for long overflow - builder.writeLong(index); + BIGINT.writeLong(builder, index); } processedPositions = processedPositions + page.getPositionCount(); indexes = builder.build(); @@ -1079,7 +1079,7 @@ public TableFunctionProcessorState process(List> input) // pass-through index for input_1 BlockBuilder input1PassThroughBuilder = BIGINT.createBlockBuilder(null, 1); if (input1Present) { - input1PassThroughBuilder.writeLong(input1EndIndex - 1); + BIGINT.writeLong(input1PassThroughBuilder, input1EndIndex - 1); } else { input1PassThroughBuilder.appendNull(); @@ -1088,7 +1088,7 @@ public TableFunctionProcessorState process(List> input) // pass-through index for input_2 BlockBuilder input2PassThroughBuilder = BIGINT.createBlockBuilder(null, 1); if (input2Present) { - input2PassThroughBuilder.writeLong(input2EndIndex - 1); + BIGINT.writeLong(input2PassThroughBuilder, input2EndIndex - 1); } else { input2PassThroughBuilder.appendNull(); diff --git a/core/trino-main/src/test/java/io/trino/execution/TestPageSplitterUtil.java b/core/trino-main/src/test/java/io/trino/execution/TestPageSplitterUtil.java index 873eeec05600..f2074110f70e 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestPageSplitterUtil.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestPageSplitterUtil.java @@ -17,8 +17,8 @@ import io.airlift.slice.Slice; import io.trino.spi.Page; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.RunLengthEncodedBlock; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.type.Type; import io.trino.testing.MaterializedResult; import org.testng.annotations.Test; @@ -82,8 +82,8 @@ public void testSplitPageNonDecreasingPageSize() List types = ImmutableList.of(VARCHAR); Slice expectedValue = wrappedBuffer("test".getBytes(UTF_8)); - BlockBuilder blockBuilder = VARCHAR.createBlockBuilder(null, 1, expectedValue.length()); - blockBuilder.writeBytes(expectedValue, 0, expectedValue.length()).closeEntry(); + VariableWidthBlockBuilder blockBuilder = VARCHAR.createBlockBuilder(null, 1, expectedValue.length()); + blockBuilder.writeEntry(expectedValue); Block rleBlock = RunLengthEncodedBlock.create(blockBuilder.build(), positionCount); Page initialPage = new Page(rleBlock); List pages = splitPage(initialPage, maxPageSizeInBytes); diff --git a/core/trino-main/src/test/java/io/trino/execution/buffer/BenchmarkBlockSerde.java b/core/trino-main/src/test/java/io/trino/execution/buffer/BenchmarkBlockSerde.java index 4c3e38f8cab8..a22c69452842 100644 --- a/core/trino-main/src/test/java/io/trino/execution/buffer/BenchmarkBlockSerde.java +++ b/core/trino-main/src/test/java/io/trino/execution/buffer/BenchmarkBlockSerde.java @@ -22,6 +22,7 @@ import io.trino.spi.Page; import io.trino.spi.PageBuilder; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.block.TestingBlockEncodingSerde; import io.trino.spi.type.DecimalType; import io.trino.spi.type.Int128; @@ -252,26 +253,29 @@ else if (TIMESTAMP_PICOS.equals(type)) { TIMESTAMP_PICOS.writeObject(blockBuilder, value); } else if (INTEGER.equals(type)) { - blockBuilder.writeInt((int) value); + INTEGER.writeInt(blockBuilder, (int) value); } else if (SMALLINT.equals(type)) { - blockBuilder.writeShort((short) value); + SMALLINT.writeShort(blockBuilder, (short) value); } else if (TINYINT.equals(type)) { - blockBuilder.writeByte((byte) value); + TINYINT.writeByte(blockBuilder, (byte) value); } else if (type instanceof RowType) { - BlockBuilder row = blockBuilder.beginBlockEntry(); List values = (List) value; if (values.size() != type.getTypeParameters().size()) { throw new IllegalArgumentException("Size of types and values must have the same size"); } - List> pairs = new ArrayList<>(); - for (int i = 0; i < type.getTypeParameters().size(); i++) { - pairs.add(new SimpleEntry<>(type.getTypeParameters().get(i), ((List) value).get(i))); - } - pairs.forEach(p -> writeValue(p.getKey(), p.getValue(), row)); - blockBuilder.closeEntry(); + ((RowBlockBuilder) blockBuilder).buildEntry(fieldBuilders -> { + List> pairs = new ArrayList<>(); + for (int i = 0; i < type.getTypeParameters().size(); i++) { + pairs.add(new SimpleEntry<>(type.getTypeParameters().get(i), ((List) value).get(i))); + } + for (int i = 0; i < pairs.size(); i++) { + SimpleEntry p = pairs.get(i); + writeValue(p.getKey(), p.getValue(), fieldBuilders.get(i)); + } + }); } else { throw new IllegalArgumentException("Unsupported type " + type); diff --git a/core/trino-main/src/test/java/io/trino/execution/buffer/TestPagesSerde.java b/core/trino-main/src/test/java/io/trino/execution/buffer/TestPagesSerde.java index 744ccee92c78..25d24b34d7b8 100644 --- a/core/trino-main/src/test/java/io/trino/execution/buffer/TestPagesSerde.java +++ b/core/trino-main/src/test/java/io/trino/execution/buffer/TestPagesSerde.java @@ -283,11 +283,12 @@ private void testDeserializationWithRollover(boolean encryptionEnabled, boolean private static Page createTestPage(int numberOfEntries) { VariableWidthBlockBuilder blockBuilder = new VariableWidthBlockBuilder(null, 1, 1000); - blockBuilder.writeInt(numberOfEntries); - for (int i = 0; i < numberOfEntries; i++) { - blockBuilder.writeLong(i); - } - blockBuilder.closeEntry(); + blockBuilder.buildEntry(value -> { + value.writeInt(numberOfEntries); + for (int i = 0; i < numberOfEntries; i++) { + value.writeLong(i); + } + }); return new Page(blockBuilder.build()); } @@ -299,12 +300,13 @@ public Block readBlock(SliceInput input) { int numberOfEntries = input.readInt(); VariableWidthBlockBuilder blockBuilder = new VariableWidthBlockBuilder(null, 1, 1000); - blockBuilder.writeInt(numberOfEntries); - for (int i = 0; i < numberOfEntries; ++i) { - // read 8 bytes at a time - blockBuilder.writeLong(input.readLong()); - } - blockBuilder.closeEntry(); + blockBuilder.buildEntry(value -> { + value.writeInt(numberOfEntries); + for (int i = 0; i < numberOfEntries; ++i) { + // read 8 bytes at a time + value.writeLong(input.readLong()); + } + }); return blockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/execution/buffer/TestSpoolingExchangeOutputBuffer.java b/core/trino-main/src/test/java/io/trino/execution/buffer/TestSpoolingExchangeOutputBuffer.java index 9073e4d31641..b6fa9e36aa6c 100644 --- a/core/trino-main/src/test/java/io/trino/execution/buffer/TestSpoolingExchangeOutputBuffer.java +++ b/core/trino-main/src/test/java/io/trino/execution/buffer/TestSpoolingExchangeOutputBuffer.java @@ -25,8 +25,8 @@ import io.trino.spi.Page; import io.trino.spi.PageBuilder; import io.trino.spi.QueryId; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.TestingBlockEncodingSerde; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.exchange.ExchangeSink; import io.trino.spi.exchange.ExchangeSinkInstanceHandle; import org.testng.annotations.Test; @@ -309,9 +309,8 @@ private static Slice createPage(String value) PageBuilder pageBuilder = new PageBuilder(ImmutableList.of(VARCHAR)); pageBuilder.declarePosition(); Slice valueSlice = utf8Slice(value); - BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0); - blockBuilder.writeBytes(valueSlice, 0, valueSlice.length()); - blockBuilder.closeEntry(); + VariableWidthBlockBuilder blockBuilder = (VariableWidthBlockBuilder) pageBuilder.getBlockBuilder(0); + blockBuilder.writeEntry(valueSlice); Page page = pageBuilder.build(); PageSerializer serializer = new PagesSerdeFactory(new TestingBlockEncodingSerde(), false).createSerializer(Optional.empty()); return serializer.serialize(page); diff --git a/core/trino-main/src/test/java/io/trino/operator/BenchmarkGroupByHashOnSimulatedData.java b/core/trino-main/src/test/java/io/trino/operator/BenchmarkGroupByHashOnSimulatedData.java index 9fccff0b2647..efb5f85bea06 100644 --- a/core/trino-main/src/test/java/io/trino/operator/BenchmarkGroupByHashOnSimulatedData.java +++ b/core/trino-main/src/test/java/io/trino/operator/BenchmarkGroupByHashOnSimulatedData.java @@ -20,6 +20,7 @@ import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.DictionaryBlock; +import io.trino.spi.block.LongArrayBlockBuilder; import io.trino.spi.type.BigintType; import io.trino.spi.type.CharType; import io.trino.spi.type.DoubleType; @@ -163,19 +164,19 @@ public enum ColumnType BIGINT(BigintType.BIGINT, (blockBuilder, positionCount, seed) -> { Random r = new Random(seed); for (int i = 0; i < positionCount; i++) { - blockBuilder.writeLong((r.nextLong() >>> 1)); // Only positives + BigintType.BIGINT.writeLong(blockBuilder, r.nextLong() >>> 1); // Only positives } }), INT(IntegerType.INTEGER, (blockBuilder, positionCount, seed) -> { Random r = new Random(seed); for (int i = 0; i < positionCount; i++) { - blockBuilder.writeInt(r.nextInt()); + IntegerType.INTEGER.writeInt(blockBuilder, r.nextInt()); } }), DOUBLE(DoubleType.DOUBLE, (blockBuilder, positionCount, seed) -> { Random r = new Random(seed); for (int i = 0; i < positionCount; i++) { - blockBuilder.writeLong((r.nextLong() >>> 1)); // Only positives + ((LongArrayBlockBuilder) blockBuilder).writeLong(r.nextLong() >>> 1); // Only positives } }), VARCHAR_25(VarcharType.VARCHAR, (blockBuilder, positionCount, seed) -> { diff --git a/core/trino-main/src/test/java/io/trino/operator/BenchmarkTopNOperator.java b/core/trino-main/src/test/java/io/trino/operator/BenchmarkTopNOperator.java index 2303a05e41f3..9ca091ea9d68 100644 --- a/core/trino-main/src/test/java/io/trino/operator/BenchmarkTopNOperator.java +++ b/core/trino-main/src/test/java/io/trino/operator/BenchmarkTopNOperator.java @@ -50,7 +50,6 @@ import static io.trino.spi.connector.SortOrder.DESC_NULLS_LAST; import static io.trino.spi.type.DateType.DATE; import static io.trino.spi.type.DoubleType.DOUBLE; -import static io.trino.spi.type.VarcharType.VARCHAR; import static java.util.concurrent.Executors.newCachedThreadPool; import static java.util.concurrent.Executors.newScheduledThreadPool; import static org.testng.Assert.assertEquals; @@ -88,7 +87,7 @@ public void setup() executor = newCachedThreadPool(daemonThreadsNamed(getClass().getSimpleName() + "-%s")); scheduledExecutor = newScheduledThreadPool(2, daemonThreadsNamed(getClass().getSimpleName() + "-scheduledExecutor-%s")); - List types = ImmutableList.of(DOUBLE, DOUBLE, VARCHAR, DOUBLE); + List types = ImmutableList.of(DOUBLE, DOUBLE, DATE, DOUBLE); pages = createInputPages(Integer.valueOf(positionsPerPage), types); operatorFactory = TopNOperator.createOperatorFactory( 0, diff --git a/core/trino-main/src/test/java/io/trino/operator/BenchmarkWindowOperator.java b/core/trino-main/src/test/java/io/trino/operator/BenchmarkWindowOperator.java index ee08482f6be7..7f970f4ed0a5 100644 --- a/core/trino-main/src/test/java/io/trino/operator/BenchmarkWindowOperator.java +++ b/core/trino-main/src/test/java/io/trino/operator/BenchmarkWindowOperator.java @@ -193,8 +193,8 @@ private RowPagesBuilder buildPages(int currentPartitionIdentifier, List ty currentGroupIdentifier = groupIdentifier++; } - firstColumnBlockBuilder.writeLong(currentGroupIdentifier); - secondColumnBlockBuilder.writeLong(currentPartitionIdentifier); + BIGINT.writeLong(firstColumnBlockBuilder, currentGroupIdentifier); + BIGINT.writeLong(secondColumnBlockBuilder, currentPartitionIdentifier); ++currentNumberOfRowsInPartition; } diff --git a/core/trino-main/src/test/java/io/trino/operator/OperatorAssertion.java b/core/trino-main/src/test/java/io/trino/operator/OperatorAssertion.java index ac664e55cf50..153c80daec52 100644 --- a/core/trino-main/src/test/java/io/trino/operator/OperatorAssertion.java +++ b/core/trino-main/src/test/java/io/trino/operator/OperatorAssertion.java @@ -20,8 +20,6 @@ import io.trino.Session; import io.trino.spi.Page; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; -import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.RowType; import io.trino.spi.type.Type; import io.trino.testing.MaterializedResult; @@ -44,6 +42,7 @@ import static io.airlift.concurrent.MoreFutures.tryGetFutureValue; import static io.airlift.testing.Assertions.assertEqualsIgnoreOrder; import static io.trino.operator.PageAssertions.assertPageEquals; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.util.StructuralTestUtil.appendToBlockBuilder; import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; @@ -199,14 +198,11 @@ public static Block toRow(List parameterTypes, Object... values) { checkArgument(parameterTypes.size() == values.length, "parameterTypes.size(" + parameterTypes.size() + ") does not equal to values.length(" + values.length + ")"); - RowType rowType = RowType.anonymous(parameterTypes); - BlockBuilder blockBuilder = new RowBlockBuilder(parameterTypes, null, 1); - BlockBuilder singleRowBlockWriter = blockBuilder.beginBlockEntry(); - for (int i = 0; i < values.length; i++) { - appendToBlockBuilder(parameterTypes.get(i), values[i], singleRowBlockWriter); - } - blockBuilder.closeEntry(); - return rowType.getObject(blockBuilder, 0); + return buildRowValue(RowType.anonymous(parameterTypes), fields -> { + for (int i = 0; i < values.length; i++) { + appendToBlockBuilder(parameterTypes.get(i), values[i], fields.get(i)); + } + }); } public static void assertOperatorEquals(OperatorFactory operatorFactory, List types, DriverContext driverContext, List input, List expected) diff --git a/core/trino-main/src/test/java/io/trino/operator/TestDynamicFilterSourceOperator.java b/core/trino-main/src/test/java/io/trino/operator/TestDynamicFilterSourceOperator.java index 49d90a3dfcf2..af1ec02e8b45 100644 --- a/core/trino-main/src/test/java/io/trino/operator/TestDynamicFilterSourceOperator.java +++ b/core/trino-main/src/test/java/io/trino/operator/TestDynamicFilterSourceOperator.java @@ -271,12 +271,11 @@ public void testCollectOnlyLastColumn() @Test public void testCollectWithNulls() { - Block blockWithNulls = INTEGER - .createFixedSizeBlockBuilder(0) - .writeInt(3) - .appendNull() - .writeInt(4) - .build(); + BlockBuilder blockBuilder = INTEGER.createFixedSizeBlockBuilder(3); + INTEGER.writeInt(blockBuilder, 3); + blockBuilder.appendNull(); + INTEGER.writeInt(blockBuilder, 4); + Block blockWithNulls = blockBuilder.build(); OperatorFactory operatorFactory = createOperatorFactory(channel(0, INTEGER)); verifyPassthrough(createOperator(operatorFactory), diff --git a/core/trino-main/src/test/java/io/trino/operator/TestTableFinishOperator.java b/core/trino-main/src/test/java/io/trino/operator/TestTableFinishOperator.java index 251d393430b1..0089a5e678c9 100644 --- a/core/trino-main/src/test/java/io/trino/operator/TestTableFinishOperator.java +++ b/core/trino-main/src/test/java/io/trino/operator/TestTableFinishOperator.java @@ -24,7 +24,6 @@ import io.trino.operator.TableFinishOperator.TableFinishOperatorFactory; import io.trino.operator.TableFinishOperator.TableFinisher; import io.trino.operator.aggregation.TestingAggregationFunction; -import io.trino.spi.block.Block; import io.trino.spi.block.LongArrayBlockBuilder; import io.trino.spi.connector.ConnectorOutputMetadata; import io.trino.spi.statistics.ColumnStatisticMetadata; @@ -144,11 +143,10 @@ public void testStatisticsAggregation() assertEquals(tableFinisher.getFragments(), ImmutableList.of(Slices.wrappedBuffer(new byte[] {1}), Slices.wrappedBuffer(new byte[] {2}))); assertEquals(tableFinisher.getComputedStatistics().size(), 1); assertEquals(getOnlyElement(tableFinisher.getComputedStatistics()).getColumnStatistics().size(), 1); - Block expectedStatisticsBlock = new LongArrayBlockBuilder(null, 1) - .writeLong(7) - .closeEntry() - .build(); - assertBlockEquals(BIGINT, getOnlyElement(tableFinisher.getComputedStatistics()).getColumnStatistics().get(statisticMetadata), expectedStatisticsBlock); + + LongArrayBlockBuilder expectedStatistics = new LongArrayBlockBuilder(null, 1); + BIGINT.writeLong(expectedStatistics, 7); + assertBlockEquals(BIGINT, getOnlyElement(tableFinisher.getComputedStatistics()).getColumnStatistics().get(statisticMetadata), expectedStatistics.build()); assertEquals(driverContext.getMemoryUsage(), 0, "memoryUsage"); } diff --git a/core/trino-main/src/test/java/io/trino/operator/aggregation/TestApproximatePercentileAggregation.java b/core/trino-main/src/test/java/io/trino/operator/aggregation/TestApproximatePercentileAggregation.java index ba207ee9892f..032368bcda7f 100644 --- a/core/trino-main/src/test/java/io/trino/operator/aggregation/TestApproximatePercentileAggregation.java +++ b/core/trino-main/src/test/java/io/trino/operator/aggregation/TestApproximatePercentileAggregation.java @@ -15,6 +15,7 @@ import com.google.common.collect.ImmutableList; import io.trino.metadata.TestingFunctionResolution; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.RunLengthEncodedBlock; @@ -631,15 +632,12 @@ private static Block createRleBlock(double percentile, int positionCount) private static Block createRleBlock(Iterable percentiles, int positionCount) { - BlockBuilder rleBlockBuilder = new ArrayType(DOUBLE).createBlockBuilder(null, 1); - BlockBuilder arrayBlockBuilder = rleBlockBuilder.beginBlockEntry(); - - for (double percentile : percentiles) { - DOUBLE.writeDouble(arrayBlockBuilder, percentile); - } - - rleBlockBuilder.closeEntry(); - - return RunLengthEncodedBlock.create(rleBlockBuilder.build(), positionCount); + ArrayBlockBuilder arrayBuilder = new ArrayType(DOUBLE).createBlockBuilder(null, 1); + arrayBuilder.buildEntry(elementBuilder -> { + for (double percentile : percentiles) { + DOUBLE.writeDouble(elementBuilder, percentile); + } + }); + return RunLengthEncodedBlock.create(arrayBuilder.build(), positionCount); } } diff --git a/core/trino-main/src/test/java/io/trino/operator/aggregation/listagg/TestListaggAggregationFunction.java b/core/trino-main/src/test/java/io/trino/operator/aggregation/listagg/TestListaggAggregationFunction.java index 1aabfd0ed015..9e867e7803c1 100644 --- a/core/trino-main/src/test/java/io/trino/operator/aggregation/listagg/TestListaggAggregationFunction.java +++ b/core/trino-main/src/test/java/io/trino/operator/aggregation/listagg/TestListaggAggregationFunction.java @@ -13,17 +13,18 @@ */ package io.trino.operator.aggregation.listagg; +import io.airlift.slice.DynamicSliceOutput; import io.airlift.slice.Slice; import io.trino.block.BlockAssertions; import io.trino.metadata.TestingFunctionResolution; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.sql.analyzer.TypeSignatureProvider; import io.trino.sql.tree.QualifiedName; import org.testng.annotations.Test; +import java.nio.charset.StandardCharsets; import java.util.List; import static io.airlift.slice.Slices.utf8Slice; @@ -68,13 +69,12 @@ public void testInputEmptyState() assertEquals(state.getOverflowFiller(), overflowFiller); assertTrue(state.showOverflowEntryCount()); - BlockBuilder out = new VariableWidthBlockBuilder(null, 16, 128); + DynamicSliceOutput out = new DynamicSliceOutput(state.getEntryCount() * 128); state.forEach((block, position) -> { - block.writeBytesTo(position, 0, block.getSliceLength(position), out); + out.writeBytes(block.getSlice(position, 0, block.getSliceLength(position))); return true; }); - out.closeEntry(); - String result = (String) BlockAssertions.getOnlyValue(VARCHAR, out); + String result = out.toString(StandardCharsets.UTF_8); assertEquals(result, s); } @@ -111,7 +111,7 @@ public void testOutputStateWithOverflowError() SingleListaggAggregationState state = createListaggAggregationState("", true, "...", false, "overflowvalue1", "overflowvalue2"); - BlockBuilder out = new VariableWidthBlockBuilder(null, 16, 128); + VariableWidthBlockBuilder out = new VariableWidthBlockBuilder(null, 16, 128); assertThatThrownBy(() -> ListaggAggregationFunction.outputState(state, out, 20)) .isInstanceOf(TrinoException.class) .matches(throwable -> ((TrinoException) throwable).getErrorCode() == EXCEEDED_FUNCTION_MEMORY_LIMIT.toErrorCode()); @@ -244,7 +244,7 @@ public void testExecute() private static String getOutputStateOnlyValue(SingleListaggAggregationState state, int maxOutputLengthInBytes) { - BlockBuilder out = new VariableWidthBlockBuilder(null, 32, 256); + VariableWidthBlockBuilder out = new VariableWidthBlockBuilder(null, 32, 256); ListaggAggregationFunction.outputState(state, out, maxOutputLengthInBytes); return (String) BlockAssertions.getOnlyValue(VARCHAR, out); } diff --git a/core/trino-main/src/test/java/io/trino/operator/aggregation/minmaxn/TestArrayMaxNAggregation.java b/core/trino-main/src/test/java/io/trino/operator/aggregation/minmaxn/TestArrayMaxNAggregation.java index 5fea0777d28b..3d65d123cddf 100644 --- a/core/trino-main/src/test/java/io/trino/operator/aggregation/minmaxn/TestArrayMaxNAggregation.java +++ b/core/trino-main/src/test/java/io/trino/operator/aggregation/minmaxn/TestArrayMaxNAggregation.java @@ -16,8 +16,8 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; import io.trino.operator.aggregation.AbstractTestAggregationFunction; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.Type; import org.testng.annotations.Test; @@ -37,15 +37,13 @@ public class TestArrayMaxNAggregation public static Block createLongArraysBlock(Long[] values) { ArrayType arrayType = new ArrayType(BIGINT); - BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, values.length); + ArrayBlockBuilder blockBuilder = arrayType.createBlockBuilder(null, values.length); for (Long value : values) { if (value == null) { blockBuilder.appendNull(); } else { - BlockBuilder elementBlockBuilder = blockBuilder.beginBlockEntry(); - BIGINT.writeLong(elementBlockBuilder, value); - blockBuilder.closeEntry(); + blockBuilder.buildEntry(elementBuilder -> BIGINT.writeLong(elementBuilder, value)); } } return blockBuilder.build(); diff --git a/core/trino-main/src/test/java/io/trino/operator/aggregation/state/TestStateCompiler.java b/core/trino-main/src/test/java/io/trino/operator/aggregation/state/TestStateCompiler.java index ec20e6a236c7..11364de7e7d3 100644 --- a/core/trino-main/src/test/java/io/trino/operator/aggregation/state/TestStateCompiler.java +++ b/core/trino-main/src/test/java/io/trino/operator/aggregation/state/TestStateCompiler.java @@ -26,6 +26,7 @@ import io.trino.array.SliceBigArray; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.function.AccumulatorState; import io.trino.spi.function.AccumulatorStateFactory; import io.trino.spi.function.AccumulatorStateSerializer; @@ -317,11 +318,11 @@ public void testComplexStateEstimatedSize() Block array = createLongsBlock(45); retainedSize += array.getRetainedSizeInBytes(); groupedState.setBlock(array); - BlockBuilder mapBlockBuilder = mapType(BIGINT, VARCHAR).createBlockBuilder(null, 1); - BlockBuilder singleMapBlockWriter = mapBlockBuilder.beginBlockEntry(); - BIGINT.writeLong(singleMapBlockWriter, 123L); - VARCHAR.writeSlice(singleMapBlockWriter, utf8Slice("testBlock")); - mapBlockBuilder.closeEntry(); + MapBlockBuilder mapBlockBuilder = mapType(BIGINT, VARCHAR).createBlockBuilder(null, 1); + mapBlockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + BIGINT.writeLong(keyBuilder, 123L); + VARCHAR.writeSlice(valueBuilder, utf8Slice("testBlock")); + }); Block map = mapBlockBuilder.build(); retainedSize += map.getRetainedSizeInBytes(); groupedState.setAnotherBlock(map); @@ -346,11 +347,11 @@ public void testComplexStateEstimatedSize() Block array = createLongsBlock(45); retainedSize += array.getRetainedSizeInBytes(); groupedState.setBlock(array); - BlockBuilder mapBlockBuilder = mapType(BIGINT, VARCHAR).createBlockBuilder(null, 1); - BlockBuilder singleMapBlockWriter = mapBlockBuilder.beginBlockEntry(); - BIGINT.writeLong(singleMapBlockWriter, 123L); - VARCHAR.writeSlice(singleMapBlockWriter, utf8Slice("testBlock")); - mapBlockBuilder.closeEntry(); + MapBlockBuilder mapBlockBuilder = mapType(BIGINT, VARCHAR).createBlockBuilder(null, 1); + mapBlockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + BIGINT.writeLong(keyBuilder, 123L); + VARCHAR.writeSlice(valueBuilder, utf8Slice("testBlock")); + }); Block map = mapBlockBuilder.build(); retainedSize += map.getRetainedSizeInBytes(); groupedState.setAnotherBlock(map); diff --git a/core/trino-main/src/test/java/io/trino/operator/project/TestDictionaryAwarePageProjection.java b/core/trino-main/src/test/java/io/trino/operator/project/TestDictionaryAwarePageProjection.java index 945915644caf..293fa9fcb078 100644 --- a/core/trino-main/src/test/java/io/trino/operator/project/TestDictionaryAwarePageProjection.java +++ b/core/trino-main/src/test/java/io/trino/operator/project/TestDictionaryAwarePageProjection.java @@ -435,7 +435,7 @@ public boolean process() int offset = selectedPositions.getOffset(); int[] positions = selectedPositions.getPositions(); for (int index = nextIndexOrPosition + offset; index < offset + selectedPositions.size(); index++) { - blockBuilder.writeLong(verifyPositive(block.getLong(positions[index], 0))); + BIGINT.writeLong(blockBuilder, verifyPositive(block.getLong(positions[index], 0))); if (yieldSignal.isSet()) { nextIndexOrPosition = index + 1 - offset; return false; @@ -445,7 +445,7 @@ public boolean process() else { int offset = selectedPositions.getOffset(); for (int position = nextIndexOrPosition + offset; position < offset + selectedPositions.size(); position++) { - blockBuilder.writeLong(verifyPositive(block.getLong(position, 0))); + BIGINT.writeLong(blockBuilder, verifyPositive(block.getLong(position, 0))); if (yieldSignal.isSet()) { nextIndexOrPosition = position + 1 - offset; return false; diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayDistinct.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayDistinct.java index 46ef750afeed..638ecc7cc8da 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayDistinct.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayDistinct.java @@ -21,6 +21,7 @@ import io.trino.operator.aggregation.TypedSet; import io.trino.operator.project.PageProcessor; import io.trino.spi.Page; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.ScalarFunction; @@ -129,21 +130,21 @@ public void setup() private static Block createChannel(int positionCount, int arraySize, ArrayType arrayType) { - BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); + ArrayBlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); for (int position = 0; position < positionCount; position++) { - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (int i = 0; i < arraySize; i++) { - if (arrayType.getElementType().getJavaType() == long.class) { - arrayType.getElementType().writeLong(entryBuilder, ThreadLocalRandom.current().nextLong()); + blockBuilder.buildEntry(elementBuilder -> { + for (int i = 0; i < arraySize; i++) { + if (arrayType.getElementType().getJavaType() == long.class) { + arrayType.getElementType().writeLong(elementBuilder, ThreadLocalRandom.current().nextLong()); + } + else if (arrayType.getElementType().equals(VARCHAR)) { + arrayType.getElementType().writeSlice(elementBuilder, Slices.utf8Slice("test_string")); + } + else { + throw new UnsupportedOperationException(); + } } - else if (arrayType.getElementType().equals(VARCHAR)) { - arrayType.getElementType().writeSlice(entryBuilder, Slices.utf8Slice("test_string")); - } - else { - throw new UnsupportedOperationException(); - } - } - blockBuilder.closeEntry(); + }); } return blockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayEqualOperator.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayEqualOperator.java index 7f3fc622af95..8b2e229cf0fa 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayEqualOperator.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayEqualOperator.java @@ -13,8 +13,8 @@ */ package io.trino.operator.scalar; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.type.ArrayType; import org.openjdk.jmh.annotations.Benchmark; import org.openjdk.jmh.annotations.BenchmarkMode; @@ -86,18 +86,18 @@ public void setup() private static Block[] createChannels(int positionCount, int arraySize, ArrayType arrayType) { ThreadLocalRandom random = ThreadLocalRandom.current(); - BlockBuilder leftBlockBuilder = arrayType.createBlockBuilder(null, positionCount); - BlockBuilder rightBlockBuilder = arrayType.createBlockBuilder(null, positionCount); + ArrayBlockBuilder leftBlockBuilder = arrayType.createBlockBuilder(null, positionCount); + ArrayBlockBuilder rightBlockBuilder = arrayType.createBlockBuilder(null, positionCount); for (int position = 0; position < positionCount; position++) { - BlockBuilder leftEntryBuilder = leftBlockBuilder.beginBlockEntry(); - BlockBuilder rightEntryBuilder = rightBlockBuilder.beginBlockEntry(); - for (int i = 0; i < arraySize - 1; i++) { - addElement(arrayType.getElementType(), random, leftEntryBuilder, rightEntryBuilder, true); - } - // last element has a 50% chance of being equal - addElement(arrayType.getElementType(), random, leftEntryBuilder, rightEntryBuilder, random.nextBoolean()); - leftBlockBuilder.closeEntry(); - rightBlockBuilder.closeEntry(); + leftBlockBuilder.buildEntry(leftElementBuilder -> { + rightBlockBuilder.buildEntry(rightElementBuilder -> { + for (int i = 0; i < arraySize - 1; i++) { + addElement(arrayType.getElementType(), random, leftElementBuilder, rightElementBuilder, true); + } + // last element has a 50% chance of being equal + addElement(arrayType.getElementType(), random, leftElementBuilder, rightElementBuilder, random.nextBoolean()); + }); + }); } return new Block[] {leftBlockBuilder.build(), rightBlockBuilder.build()}; } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayFilter.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayFilter.java index f1ad4b805550..9b159995fb39 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayFilter.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayFilter.java @@ -22,6 +22,7 @@ import io.trino.operator.DriverYieldSignal; import io.trino.operator.project.PageProcessor; import io.trino.spi.Page; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.BoundSignature; @@ -166,18 +167,18 @@ public void setup() private static Block createChannel(int positionCount, int arraySize, ArrayType arrayType) { - BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); + ArrayBlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); for (int position = 0; position < positionCount; position++) { - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (int i = 0; i < arraySize; i++) { - if (arrayType.getElementType().getJavaType() == long.class) { - arrayType.getElementType().writeLong(entryBuilder, ThreadLocalRandom.current().nextLong()); + blockBuilder.buildEntry(elementBuilder -> { + for (int i = 0; i < arraySize; i++) { + if (arrayType.getElementType().getJavaType() == long.class) { + arrayType.getElementType().writeLong(elementBuilder, ThreadLocalRandom.current().nextLong()); + } + else { + throw new UnsupportedOperationException(); + } } - else { - throw new UnsupportedOperationException(); - } - } - blockBuilder.closeEntry(); + }); } return blockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayHashCodeOperator.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayHashCodeOperator.java index 9d3d32100b70..387994a38d26 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayHashCodeOperator.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayHashCodeOperator.java @@ -13,8 +13,8 @@ */ package io.trino.operator.scalar; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.type.ArrayType; import org.openjdk.jmh.annotations.Benchmark; import org.openjdk.jmh.annotations.BenchmarkMode; @@ -81,13 +81,13 @@ public void setup() private static Block createChannel(int positionCount, int arraySize, ArrayType arrayType) { - BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); + ArrayBlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); for (int position = 0; position < positionCount; position++) { - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (int i = 0; i < arraySize; i++) { - addElement(arrayType.getElementType(), ThreadLocalRandom.current(), entryBuilder); - } - blockBuilder.closeEntry(); + blockBuilder.buildEntry(elementBuilder -> { + for (int i = 0; i < arraySize; i++) { + addElement(arrayType.getElementType(), ThreadLocalRandom.current(), elementBuilder); + } + }); } return blockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayIntersect.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayIntersect.java index f690efd433e6..278d01ae0603 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayIntersect.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayIntersect.java @@ -19,8 +19,8 @@ import io.trino.operator.DriverYieldSignal; import io.trino.operator.project.PageProcessor; import io.trino.spi.Page; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.Type; import io.trino.sql.gen.ExpressionCompiler; @@ -129,28 +129,28 @@ public void setup() private static Block createChannel(int positionCount, int arraySize, Type elementType) { ArrayType arrayType = new ArrayType(elementType); - BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); + ArrayBlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); for (int position = 0; position < positionCount; position++) { - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (int i = 0; i < arraySize; i++) { - if (elementType.getJavaType() == long.class) { - elementType.writeLong(entryBuilder, ThreadLocalRandom.current().nextLong() % arraySize); + blockBuilder.buildEntry(elementBuilder -> { + for (int i = 0; i < arraySize; i++) { + if (elementType.getJavaType() == long.class) { + elementType.writeLong(elementBuilder, ThreadLocalRandom.current().nextLong() % arraySize); + } + else if (elementType.getJavaType() == double.class) { + elementType.writeDouble(elementBuilder, ThreadLocalRandom.current().nextDouble() % arraySize); + } + else if (elementType.getJavaType() == boolean.class) { + elementType.writeBoolean(elementBuilder, ThreadLocalRandom.current().nextBoolean()); + } + else if (elementType.equals(VARCHAR)) { + // make sure the size of a varchar is rather small; otherwise the aggregated slice may overflow + elementType.writeSlice(elementBuilder, Slices.utf8Slice(Long.toString(ThreadLocalRandom.current().nextLong() % arraySize))); + } + else { + throw new UnsupportedOperationException(); + } } - else if (elementType.getJavaType() == double.class) { - elementType.writeDouble(entryBuilder, ThreadLocalRandom.current().nextDouble() % arraySize); - } - else if (elementType.getJavaType() == boolean.class) { - elementType.writeBoolean(entryBuilder, ThreadLocalRandom.current().nextBoolean()); - } - else if (elementType.equals(VARCHAR)) { - // make sure the size of a varchar is rather small; otherwise the aggregated slice may overflow - elementType.writeSlice(entryBuilder, Slices.utf8Slice(Long.toString(ThreadLocalRandom.current().nextLong() % arraySize))); - } - else { - throw new UnsupportedOperationException(); - } - } - blockBuilder.closeEntry(); + }); } return blockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayJoin.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayJoin.java index 64280c0dc6bc..d1e6906e6d48 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayJoin.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayJoin.java @@ -20,8 +20,8 @@ import io.trino.operator.DriverYieldSignal; import io.trino.operator.project.PageProcessor; import io.trino.spi.Page; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.sql.relational.CallExpression; import io.trino.sql.relational.RowExpression; @@ -100,13 +100,13 @@ private static Block createChannel(int positionCount, int arraySize) { ArrayType arrayType = new ArrayType(BIGINT); - BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); + ArrayBlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); for (int position = 0; position < positionCount; position++) { - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (int i = 0; i < arraySize; i++) { - arrayType.getElementType().writeLong(entryBuilder, ThreadLocalRandom.current().nextLong()); - } - blockBuilder.closeEntry(); + blockBuilder.buildEntry(elementBuilder -> { + for (int i = 0; i < arraySize; i++) { + arrayType.getElementType().writeLong(elementBuilder, ThreadLocalRandom.current().nextLong()); + } + }); } return blockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArraySort.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArraySort.java index daa89ec75281..fae8dd242c30 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArraySort.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArraySort.java @@ -21,6 +21,7 @@ import io.trino.operator.DriverYieldSignal; import io.trino.operator.project.PageProcessor; import io.trino.spi.Page; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.ScalarFunction; @@ -123,21 +124,21 @@ public void setup() private static Block createChannel(int positionCount, int arraySize, ArrayType arrayType) { - BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); + ArrayBlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); for (int position = 0; position < positionCount; position++) { - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (int i = 0; i < arraySize; i++) { - if (arrayType.getElementType().getJavaType() == long.class) { - arrayType.getElementType().writeLong(entryBuilder, ThreadLocalRandom.current().nextLong()); + blockBuilder.buildEntry(elementBuilder -> { + for (int i = 0; i < arraySize; i++) { + if (arrayType.getElementType().getJavaType() == long.class) { + arrayType.getElementType().writeLong(elementBuilder, ThreadLocalRandom.current().nextLong()); + } + else if (arrayType.getElementType().equals(VARCHAR)) { + arrayType.getElementType().writeSlice(elementBuilder, Slices.utf8Slice("test_string")); + } + else { + throw new UnsupportedOperationException(); + } } - else if (arrayType.getElementType().equals(VARCHAR)) { - arrayType.getElementType().writeSlice(entryBuilder, Slices.utf8Slice("test_string")); - } - else { - throw new UnsupportedOperationException(); - } - } - blockBuilder.closeEntry(); + }); } return blockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayTransform.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayTransform.java index 89d71455c1ba..298e692ec523 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayTransform.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkArrayTransform.java @@ -20,8 +20,8 @@ import io.trino.operator.project.PageProcessor; import io.trino.spi.Page; import io.trino.spi.PageBuilder; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.Type; import io.trino.sql.gen.ExpressionCompiler; @@ -130,18 +130,18 @@ public void setup() private static Block createChannel(int positionCount, int arraySize, ArrayType arrayType) { - BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); + ArrayBlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount); for (int position = 0; position < positionCount; position++) { - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (int i = 0; i < arraySize; i++) { - if (arrayType.getElementType().getJavaType() == long.class) { - arrayType.getElementType().writeLong(entryBuilder, ThreadLocalRandom.current().nextLong()); + blockBuilder.buildEntry(elementBuilder -> { + for (int i = 0; i < arraySize; i++) { + if (arrayType.getElementType().getJavaType() == long.class) { + arrayType.getElementType().writeLong(elementBuilder, ThreadLocalRandom.current().nextLong()); + } + else { + throw new UnsupportedOperationException(); + } } - else { - throw new UnsupportedOperationException(); - } - } - blockBuilder.closeEntry(); + }); } return blockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkRowEqualOperator.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkRowEqualOperator.java index f22756ae3dad..a977844a7fd9 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkRowEqualOperator.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkRowEqualOperator.java @@ -15,7 +15,7 @@ import com.google.common.collect.ImmutableList; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.RowType; import io.trino.spi.type.RowType.Field; import org.openjdk.jmh.annotations.Benchmark; @@ -91,20 +91,20 @@ public void setup() private static Block[] createChannels(int positionCount, RowType rowType) { ThreadLocalRandom random = ThreadLocalRandom.current(); - BlockBuilder leftBlockBuilder = rowType.createBlockBuilder(null, positionCount); - BlockBuilder rightBlockBuilder = rowType.createBlockBuilder(null, positionCount); + RowBlockBuilder leftBlockBuilder = rowType.createBlockBuilder(null, positionCount); + RowBlockBuilder rightBlockBuilder = rowType.createBlockBuilder(null, positionCount); for (int position = 0; position < positionCount; position++) { - BlockBuilder leftEntryBuilder = leftBlockBuilder.beginBlockEntry(); - BlockBuilder rightEntryBuilder = rightBlockBuilder.beginBlockEntry(); - - List fields = rowType.getFields(); - for (int i = 0; i < fields.size() - 1; i++) { - addElement(fields.get(i).getType(), random, leftEntryBuilder, rightEntryBuilder, true); - } - // last field has a 50% chance of being equal - addElement(fields.get(fields.size() - 1).getType(), random, leftEntryBuilder, rightEntryBuilder, random.nextBoolean()); - leftBlockBuilder.closeEntry(); - rightBlockBuilder.closeEntry(); + leftBlockBuilder.buildEntry(leftFieldBuilders -> { + rightBlockBuilder.buildEntry(rightFieldBuilders -> { + List fields = rowType.getFields(); + for (int i = 0; i < fields.size() - 1; i++) { + addElement(fields.get(i).getType(), random, leftFieldBuilders.get(i), rightFieldBuilders.get(i), true); + } + // last field has a 50% chance of being equal + int lastIndex = fields.size() - 1; + addElement(fields.get(lastIndex).getType(), random, leftFieldBuilders.get(lastIndex), rightFieldBuilders.get(lastIndex), random.nextBoolean()); + }); + }); } return new Block[] {leftBlockBuilder.build(), rightBlockBuilder.build()}; } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkRowHashCodeOperator.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkRowHashCodeOperator.java index 926b546b3f19..23bce932eb4d 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkRowHashCodeOperator.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkRowHashCodeOperator.java @@ -15,7 +15,7 @@ import com.google.common.collect.ImmutableList; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.RowType; import io.trino.spi.type.RowType.Field; import org.openjdk.jmh.annotations.Benchmark; @@ -86,15 +86,15 @@ public void setup() private static Block createChannel(int positionCount, RowType rowType) { ThreadLocalRandom random = ThreadLocalRandom.current(); - BlockBuilder blockBuilder = rowType.createBlockBuilder(null, positionCount); + RowBlockBuilder blockBuilder = rowType.createBlockBuilder(null, positionCount); for (int position = 0; position < positionCount; position++) { - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - - List fields = rowType.getFields(); - for (Field field : fields) { - addElement(field.getType(), random, entryBuilder); - } - blockBuilder.closeEntry(); + blockBuilder.buildEntry(fieldBuilders -> { + List fields = rowType.getFields(); + for (int i = 0; i < fields.size(); i++) { + Field field = fields.get(i); + addElement(field.getType(), random, fieldBuilders.get(i)); + } + }); } return blockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkTransformKey.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkTransformKey.java index 34a8fbd670b1..4a9608c6ddd7 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkTransformKey.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkTransformKey.java @@ -22,7 +22,7 @@ import io.trino.operator.project.PageProcessor; import io.trino.spi.Page; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.type.MapType; import io.trino.spi.type.Type; import io.trino.sql.gen.ExpressionCompiler; @@ -139,24 +139,24 @@ public void setup() private static Block createChannel(int positionCount, MapType mapType, Type elementType) { - BlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1); - BlockBuilder singleMapBlockWriter = mapBlockBuilder.beginBlockEntry(); - Object value; - for (int position = 0; position < positionCount; position++) { - if (elementType.equals(BIGINT)) { - value = ThreadLocalRandom.current().nextLong(); + MapBlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1); + mapBlockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + Object value; + for (int position = 0; position < positionCount; position++) { + if (elementType.equals(BIGINT)) { + value = ThreadLocalRandom.current().nextLong(); + } + else if (elementType.equals(DOUBLE)) { + value = ThreadLocalRandom.current().nextDouble(); + } + else { + throw new UnsupportedOperationException(); + } + // Use position as the key to avoid collision + writeNativeValue(elementType, keyBuilder, position); + writeNativeValue(elementType, valueBuilder, value); } - else if (elementType.equals(DOUBLE)) { - value = ThreadLocalRandom.current().nextDouble(); - } - else { - throw new UnsupportedOperationException(); - } - // Use position as the key to avoid collision - writeNativeValue(elementType, singleMapBlockWriter, position); - writeNativeValue(elementType, singleMapBlockWriter, value); - } - mapBlockBuilder.closeEntry(); + }); return mapBlockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkTransformValue.java b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkTransformValue.java index ee6663f3bea6..3a148243ee70 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkTransformValue.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/BenchmarkTransformValue.java @@ -23,7 +23,7 @@ import io.trino.operator.project.PageProcessor; import io.trino.spi.Page; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.type.MapType; import io.trino.spi.type.Type; import io.trino.sql.gen.ExpressionCompiler; @@ -145,31 +145,31 @@ public void setup() private static Block createChannel(int positionCount, MapType mapType, Type elementType) { - BlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1); - BlockBuilder singleMapBlockWriter = mapBlockBuilder.beginBlockEntry(); - Object key; - Object value; - for (int position = 0; position < positionCount; position++) { - if (elementType.equals(BIGINT)) { - key = position; - value = ThreadLocalRandom.current().nextLong(); + MapBlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1); + mapBlockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + Object key; + Object value; + for (int position = 0; position < positionCount; position++) { + if (elementType.equals(BIGINT)) { + key = position; + value = ThreadLocalRandom.current().nextLong(); + } + else if (elementType.equals(DOUBLE)) { + key = position; + value = ThreadLocalRandom.current().nextDouble(); + } + else if (elementType.equals(VARCHAR)) { + key = Slices.utf8Slice(Integer.toString(position)); + value = Slices.utf8Slice(Double.toString(ThreadLocalRandom.current().nextDouble())); + } + else { + throw new UnsupportedOperationException(); + } + // Use position as the key to avoid collision + writeNativeValue(elementType, keyBuilder, key); + writeNativeValue(elementType, valueBuilder, value); } - else if (elementType.equals(DOUBLE)) { - key = position; - value = ThreadLocalRandom.current().nextDouble(); - } - else if (elementType.equals(VARCHAR)) { - key = Slices.utf8Slice(Integer.toString(position)); - value = Slices.utf8Slice(Double.toString(ThreadLocalRandom.current().nextDouble())); - } - else { - throw new UnsupportedOperationException(); - } - // Use position as the key to avoid collision - writeNativeValue(elementType, singleMapBlockWriter, key); - writeNativeValue(elementType, singleMapBlockWriter, value); - } - mapBlockBuilder.closeEntry(); + }); return mapBlockBuilder.build(); } diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/TestMapTransformKeysFunction.java b/core/trino-main/src/test/java/io/trino/operator/scalar/TestMapTransformKeysFunction.java index dbdb2a472edd..afad49f4a35d 100644 --- a/core/trino-main/src/test/java/io/trino/operator/scalar/TestMapTransformKeysFunction.java +++ b/core/trino-main/src/test/java/io/trino/operator/scalar/TestMapTransformKeysFunction.java @@ -134,42 +134,42 @@ public void testDuplicateKeys() assertTrinoExceptionThrownBy(() -> assertions.expression("transform_keys(a, (k, v) -> k % 3)") .binding("a", "map(ARRAY[1, 2, 3, 4], ARRAY['a', 'b', 'c', 'd'])") .evaluate()) - .hasMessage("Duplicate keys (1) are not allowed"); + .hasMessage("Duplicate map keys (1) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.expression("transform_keys(a, (k, v) -> k % 2 = 0)") .binding("a", "map(ARRAY[1, 2, 3], ARRAY['a', 'b', 'c'])") .evaluate()) - .hasMessage("Duplicate keys (false) are not allowed"); + .hasMessage("Duplicate map keys (false) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.expression("transform_keys(a, (k, v) -> k - floor(k))") .binding("a", "map(ARRAY[1.5E0, 2.5E0, 3.5E0], ARRAY['a', 'b', 'c'])") .evaluate()) - .hasMessage("Duplicate keys (0.5) are not allowed"); + .hasMessage("Duplicate map keys (0.5) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.expression("transform_keys(a, (k, v) -> v)") .binding("a", "map(ARRAY[1, 2, 3, 4], ARRAY['a', 'b', 'c', 'b'])") .evaluate()) - .hasMessage("Duplicate keys (b) are not allowed"); + .hasMessage("Duplicate map keys (b) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.expression("transform_keys(a, (k, v) -> substr(k, 1, 3))") .binding("a", "map(ARRAY['abc1', 'cba2', 'abc3'], ARRAY[1, 2, 3])") .evaluate()) - .hasMessage("Duplicate keys (abc) are not allowed"); + .hasMessage("Duplicate map keys (abc) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.expression("transform_keys(a, (k, v) -> array_sort(k || v))") .binding("a", "map(ARRAY[ARRAY[1], ARRAY[2]], ARRAY[2, 1])") .evaluate()) - .hasMessage("Duplicate keys ([1, 2]) are not allowed"); + .hasMessage("Duplicate map keys ([1, 2]) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.expression("transform_keys(a, (k, v) -> DATE '2001-08-22')") .binding("a", "map(ARRAY[1, 2], ARRAY[null, null])") .evaluate()) - .hasMessage("Duplicate keys (2001-08-22) are not allowed"); + .hasMessage("Duplicate map keys (2001-08-22) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.expression("transform_keys(a, (k, v) -> TIMESTAMP '2001-08-22 03:04:05.321')") .binding("a", "map(ARRAY[1, 2], ARRAY[null, null])") .evaluate()) - .hasMessage("Duplicate keys (2001-08-22 03:04:05.321) are not allowed"); + .hasMessage("Duplicate map keys (2001-08-22 03:04:05.321) are not allowed"); } @Test diff --git a/core/trino-main/src/test/java/io/trino/operator/unnest/TestingUnnesterUtil.java b/core/trino-main/src/test/java/io/trino/operator/unnest/TestingUnnesterUtil.java index 7915eb495b0e..dc2a270bca22 100644 --- a/core/trino-main/src/test/java/io/trino/operator/unnest/TestingUnnesterUtil.java +++ b/core/trino-main/src/test/java/io/trino/operator/unnest/TestingUnnesterUtil.java @@ -22,6 +22,7 @@ import io.trino.spi.block.ColumnarArray; import io.trino.spi.block.ColumnarMap; import io.trino.spi.block.ColumnarRow; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.MapType; import io.trino.spi.type.RowType; @@ -85,22 +86,23 @@ public static Block createArrayBlockOfRowBlocks(Slice[][][] elements, RowType ro } else { Slice[][] expectedValues = elements[i]; - BlockBuilder elementBlockBuilder = rowType.createBlockBuilder(null, elements[i].length); + RowBlockBuilder elementBlockBuilder = rowType.createBlockBuilder(null, elements[i].length); for (Slice[] expectedValue : expectedValues) { if (expectedValue == null) { elementBlockBuilder.appendNull(); } else { - BlockBuilder entryBuilder = elementBlockBuilder.beginBlockEntry(); - for (Slice v : expectedValue) { - if (v == null) { - entryBuilder.appendNull(); + elementBlockBuilder.buildEntry(fieldBuilders -> { + for (int fieldId = 0; fieldId < expectedValue.length; fieldId++) { + Slice v = expectedValue[fieldId]; + if (v == null) { + fieldBuilders.get(fieldId).appendNull(); + } + else { + VARCHAR.writeSlice(fieldBuilders.get(fieldId), v); + } } - else { - VARCHAR.writeSlice(entryBuilder, v); - } - } - elementBlockBuilder.closeEntry(); + }); } } arrayType.writeObject(arrayBlockBuilder, elementBlockBuilder.build()); diff --git a/core/trino-main/src/test/java/io/trino/spiller/TestBinaryFileSpiller.java b/core/trino-main/src/test/java/io/trino/spiller/TestBinaryFileSpiller.java index b9dce3aca0c1..ea55cf0d6f33 100644 --- a/core/trino-main/src/test/java/io/trino/spiller/TestBinaryFileSpiller.java +++ b/core/trino-main/src/test/java/io/trino/spiller/TestBinaryFileSpiller.java @@ -14,6 +14,7 @@ package io.trino.spiller; import com.google.common.collect.ImmutableList; +import io.airlift.slice.Slices; import io.trino.FeaturesConfig; import io.trino.RowPagesBuilder; import io.trino.execution.buffer.PageSerializer; @@ -45,7 +46,6 @@ import static io.trino.spi.type.DoubleType.DOUBLE; import static io.trino.spi.type.VarbinaryType.VARBINARY; import static io.trino.spi.type.VarcharType.VARCHAR; -import static java.lang.Double.doubleToLongBits; import static org.testng.Assert.assertEquals; @Test(singleThreaded = true) @@ -110,9 +110,9 @@ public void testFileVarbinarySpiller() BlockBuilder col2 = DOUBLE.createBlockBuilder(null, 1); BlockBuilder col3 = VARBINARY.createBlockBuilder(null, 1); - col1.writeLong(42).closeEntry(); - col2.writeLong(doubleToLongBits(43.0)).closeEntry(); - col3.writeLong(doubleToLongBits(43.0)).writeLong(1).closeEntry(); + BIGINT.writeLong(col1, 42); + DOUBLE.writeDouble(col2, 43.0); + VARBINARY.writeSlice(col3, Slices.allocate(16).getOutput().appendDouble(43.0).appendLong(1).slice()); Page page = new Page(col1.build(), col2.build(), col3.build()); diff --git a/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpiller.java b/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpiller.java index b7ff11f37ffe..7aafebd8c5d6 100644 --- a/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpiller.java +++ b/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpiller.java @@ -17,6 +17,7 @@ import com.google.common.collect.Iterators; import com.google.common.util.concurrent.ListeningExecutorService; import io.airlift.slice.Slice; +import io.airlift.slice.Slices; import io.trino.execution.buffer.PagesSerdeUtil; import io.trino.memory.context.LocalMemoryContext; import io.trino.operator.PageAssertions; @@ -45,7 +46,6 @@ import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.DoubleType.DOUBLE; import static io.trino.spi.type.VarbinaryType.VARBINARY; -import static java.lang.Double.doubleToLongBits; import static java.nio.file.Files.newInputStream; import static java.util.concurrent.Executors.newCachedThreadPool; import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy; @@ -168,9 +168,9 @@ private Page buildPage() BlockBuilder col2 = DOUBLE.createBlockBuilder(null, 1); BlockBuilder col3 = VARBINARY.createBlockBuilder(null, 1); - col1.writeLong(42).closeEntry(); - col2.writeLong(doubleToLongBits(43.0)).closeEntry(); - col3.writeLong(doubleToLongBits(43.0)).writeLong(1).closeEntry(); + BIGINT.writeLong(col1, 42); + DOUBLE.writeDouble(col2, 43.0); + VARBINARY.writeSlice(col3, Slices.allocate(16).getOutput().appendDouble(43.0).appendLong(1).slice()); return new Page(col1.build(), col2.build(), col3.build()); } diff --git a/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpillerFactory.java b/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpillerFactory.java index f45c9557e101..952cc0a65612 100644 --- a/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpillerFactory.java +++ b/core/trino-main/src/test/java/io/trino/spiller/TestFileSingleStreamSpillerFactory.java @@ -140,7 +140,7 @@ public void testDistributesSpillOverPathsBadDisk() private Page buildPage() { BlockBuilder col1 = BIGINT.createBlockBuilder(null, 1); - col1.writeLong(42).closeEntry(); + BIGINT.writeLong(col1, 42); return new Page(col1.build()); } diff --git a/core/trino-main/src/test/java/io/trino/type/TestArrayOperators.java b/core/trino-main/src/test/java/io/trino/type/TestArrayOperators.java index c45416b72562..445acd7fe591 100644 --- a/core/trino-main/src/test/java/io/trino/type/TestArrayOperators.java +++ b/core/trino-main/src/test/java/io/trino/type/TestArrayOperators.java @@ -20,8 +20,8 @@ import io.airlift.slice.DynamicSliceOutput; import io.airlift.slice.Slice; import io.trino.metadata.InternalFunctionBundle; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.function.LiteralParameters; import io.trino.spi.function.ScalarFunction; import io.trino.spi.function.SqlType; @@ -119,9 +119,12 @@ public void testStackRepresentation() DynamicSliceOutput actualSliceOutput = new DynamicSliceOutput(100); writeBlock(((LocalQueryRunner) assertions.getQueryRunner()).getPlannerContext().getBlockEncodingSerde(), actualSliceOutput, actualBlock); - BlockBuilder expectedBlockBuilder = arrayType.createBlockBuilder(null, 3); - arrayType.writeObject(expectedBlockBuilder, BIGINT.createBlockBuilder(null, 2).writeLong(1).writeLong(2).build()); - arrayType.writeObject(expectedBlockBuilder, BIGINT.createBlockBuilder(null, 1).writeLong(3).build()); + ArrayBlockBuilder expectedBlockBuilder = arrayType.createBlockBuilder(null, 3); + expectedBlockBuilder.buildEntry(elementBuilder -> { + BIGINT.writeLong(elementBuilder, 1); + BIGINT.writeLong(elementBuilder, 2); + }); + expectedBlockBuilder.buildEntry(elementBuilder -> BIGINT.writeLong(elementBuilder, 3)); Block expectedBlock = expectedBlockBuilder.build(); DynamicSliceOutput expectedSliceOutput = new DynamicSliceOutput(100); writeBlock(((LocalQueryRunner) assertions.getQueryRunner()).getPlannerContext().getBlockEncodingSerde(), expectedSliceOutput, expectedBlock); diff --git a/core/trino-main/src/test/java/io/trino/type/TestCharType.java b/core/trino-main/src/test/java/io/trino/type/TestCharType.java index c33aad758a39..882aad182a42 100644 --- a/core/trino-main/src/test/java/io/trino/type/TestCharType.java +++ b/core/trino-main/src/test/java/io/trino/type/TestCharType.java @@ -18,6 +18,7 @@ import io.airlift.slice.Slices; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.type.CharType; import io.trino.spi.type.Type; import org.testng.annotations.Test; @@ -72,10 +73,9 @@ public void testGetObjectValue() CharType charType = createCharType(3); for (int codePoint : ImmutableList.of(0, 1, 10, 17, (int) ' ', 127, 1011, 11_000, 65_891, MIN_SUPPLEMENTARY_CODE_POINT, MAX_CODE_POINT)) { - BlockBuilder blockBuilder = charType.createBlockBuilder(null, 1); + VariableWidthBlockBuilder blockBuilder = charType.createBlockBuilder(null, 1); Slice slice = (codePoint != ' ') ? codePointToUtf8(codePoint) : EMPTY_SLICE; - blockBuilder.writeBytes(slice, 0, slice.length()); - blockBuilder.closeEntry(); + blockBuilder.writeEntry(slice); Block block = blockBuilder.build(); int codePointLengthInUtf16 = isSupplementaryCodePoint(codePoint) ? 2 : 1; diff --git a/core/trino-main/src/test/java/io/trino/type/TestDoubleType.java b/core/trino-main/src/test/java/io/trino/type/TestDoubleType.java index 7ab8a0dd4329..f2e602b4e8e6 100644 --- a/core/trino-main/src/test/java/io/trino/type/TestDoubleType.java +++ b/core/trino-main/src/test/java/io/trino/type/TestDoubleType.java @@ -59,7 +59,8 @@ protected Object getGreaterValue(Object value) @Test public void testNaNHash() { - BlockBuilder blockBuilder = new LongArrayBlockBuilder(null, 4); + LongArrayBlockBuilder blockBuilder = (LongArrayBlockBuilder) DOUBLE.createBlockBuilder(null, 5); + DOUBLE.writeDouble(blockBuilder, Double.NaN); blockBuilder.writeLong(doubleToLongBits(Double.NaN)); blockBuilder.writeLong(doubleToRawLongBits(Double.NaN)); // the following two are the long values of a double NaN @@ -70,10 +71,12 @@ public void testNaNHash() assertEquals(hashCodeOperator.hashCode(blockBuilder, 0), hashCodeOperator.hashCode(blockBuilder, 1)); assertEquals(hashCodeOperator.hashCode(blockBuilder, 0), hashCodeOperator.hashCode(blockBuilder, 2)); assertEquals(hashCodeOperator.hashCode(blockBuilder, 0), hashCodeOperator.hashCode(blockBuilder, 3)); + assertEquals(hashCodeOperator.hashCode(blockBuilder, 0), hashCodeOperator.hashCode(blockBuilder, 4)); BlockPositionXxHash64 xxHash64Operator = blockTypeOperators.getXxHash64Operator(DOUBLE); assertEquals(xxHash64Operator.xxHash64(blockBuilder, 0), xxHash64Operator.xxHash64(blockBuilder, 1)); assertEquals(xxHash64Operator.xxHash64(blockBuilder, 0), xxHash64Operator.xxHash64(blockBuilder, 2)); assertEquals(xxHash64Operator.xxHash64(blockBuilder, 0), xxHash64Operator.xxHash64(blockBuilder, 3)); + assertEquals(xxHash64Operator.xxHash64(blockBuilder, 0), xxHash64Operator.xxHash64(blockBuilder, 4)); } } diff --git a/core/trino-main/src/test/java/io/trino/type/TestMapOperators.java b/core/trino-main/src/test/java/io/trino/type/TestMapOperators.java index f2c4d3aecdf5..7828fff69436 100644 --- a/core/trino-main/src/test/java/io/trino/type/TestMapOperators.java +++ b/core/trino-main/src/test/java/io/trino/type/TestMapOperators.java @@ -1662,22 +1662,22 @@ public void testMapFromEntries() // invalid invocation assertTrinoExceptionThrownBy(() -> assertions.function("map_from_entries", "ARRAY[('a', 1), ('a', 2)]").evaluate()) - .hasMessage("Duplicate keys (a) are not allowed"); + .hasMessage("Duplicate map keys (a) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.function("map_from_entries", "ARRAY[(1, 1), (1, 2)]").evaluate()) - .hasMessage("Duplicate keys (1) are not allowed"); + .hasMessage("Duplicate map keys (1) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.function("map_from_entries", "ARRAY[(1.0, 1), (1.0, 2)]").evaluate()) - .hasMessage("Duplicate keys (1.0) are not allowed"); + .hasMessage("Duplicate map keys (1.0) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.function("map_from_entries", "ARRAY[(ARRAY[1, 2], 1), (ARRAY[1, 2], 2)]").evaluate()) - .hasMessage("Duplicate keys ([1, 2]) are not allowed"); + .hasMessage("Duplicate map keys ([1, 2]) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.function("map_from_entries", "ARRAY[(MAP(ARRAY[1], ARRAY[2]), 1), (MAP(ARRAY[1], ARRAY[2]), 2)]").evaluate()) - .hasMessage("Duplicate keys ({1=2}) are not allowed"); + .hasMessage("Duplicate map keys ({1=2}) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.function("map_from_entries", "ARRAY[(NaN(), 1), (NaN(), 2)]").evaluate()) - .hasMessage("Duplicate keys (NaN) are not allowed"); + .hasMessage("Duplicate map keys (NaN) are not allowed"); assertTrinoExceptionThrownBy(() -> assertions.function("map_from_entries", "ARRAY[(null, 1), (null, 2)]").evaluate()) .hasMessage("map key cannot be null"); diff --git a/core/trino-main/src/test/java/io/trino/type/TestRealType.java b/core/trino-main/src/test/java/io/trino/type/TestRealType.java index ae63fbcf3cc6..8ecb3dfa6062 100644 --- a/core/trino-main/src/test/java/io/trino/type/TestRealType.java +++ b/core/trino-main/src/test/java/io/trino/type/TestRealType.java @@ -61,16 +61,18 @@ protected Object getGreaterValue(Object value) @Test public void testNaNHash() { - BlockBuilder blockBuilder = new IntArrayBlockBuilder(null, 4); - blockBuilder.writeInt(floatToIntBits(Float.NaN)); - blockBuilder.writeInt(floatToRawIntBits(Float.NaN)); + BlockBuilder blockBuilder = new IntArrayBlockBuilder(null, 5); + REAL.writeFloat(blockBuilder, Float.NaN); + REAL.writeInt(blockBuilder, floatToIntBits(Float.NaN)); + REAL.writeInt(blockBuilder, floatToRawIntBits(Float.NaN)); // the following two are the integer values of a float NaN - blockBuilder.writeInt(-0x400000); - blockBuilder.writeInt(0x7fc00000); + REAL.writeInt(blockBuilder, -0x400000); + REAL.writeInt(blockBuilder, 0x7fc00000); BlockPositionHashCode hashCodeOperator = blockTypeOperators.getHashCodeOperator(REAL); assertEquals(hashCodeOperator.hashCode(blockBuilder, 0), hashCodeOperator.hashCode(blockBuilder, 1)); assertEquals(hashCodeOperator.hashCode(blockBuilder, 0), hashCodeOperator.hashCode(blockBuilder, 2)); assertEquals(hashCodeOperator.hashCode(blockBuilder, 0), hashCodeOperator.hashCode(blockBuilder, 3)); + assertEquals(hashCodeOperator.hashCode(blockBuilder, 0), hashCodeOperator.hashCode(blockBuilder, 4)); } } diff --git a/core/trino-main/src/test/java/io/trino/type/TestSimpleRowType.java b/core/trino-main/src/test/java/io/trino/type/TestSimpleRowType.java index ef718bfc9ce3..81bd8f62af3b 100644 --- a/core/trino-main/src/test/java/io/trino/type/TestSimpleRowType.java +++ b/core/trino-main/src/test/java/io/trino/type/TestSimpleRowType.java @@ -16,13 +16,12 @@ import com.google.common.collect.ImmutableList; import io.trino.spi.block.Block; import io.trino.spi.block.RowBlockBuilder; -import io.trino.spi.block.SingleRowBlockWriter; import io.trino.spi.type.RowType; -import io.trino.spi.type.Type; import java.util.List; import static io.airlift.slice.Slices.utf8Slice; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.RowType.field; import static io.trino.spi.type.VarcharType.VARCHAR; @@ -30,7 +29,7 @@ public class TestSimpleRowType extends AbstractTestType { - private static final Type TYPE = RowType.from(ImmutableList.of( + private static final RowType TYPE = RowType.from(ImmutableList.of( field("a", BIGINT), field("b", VARCHAR))); @@ -41,24 +40,22 @@ public TestSimpleRowType() private static Block createTestBlock() { - RowBlockBuilder blockBuilder = (RowBlockBuilder) TYPE.createBlockBuilder(null, 3); + RowBlockBuilder blockBuilder = TYPE.createBlockBuilder(null, 3); - SingleRowBlockWriter singleRowBlockWriter; + blockBuilder.buildEntry(fieldBuilders -> { + BIGINT.writeLong(fieldBuilders.get(0), 1); + VARCHAR.writeSlice(fieldBuilders.get(1), utf8Slice("cat")); + }); - singleRowBlockWriter = blockBuilder.beginBlockEntry(); - BIGINT.writeLong(singleRowBlockWriter, 1); - VARCHAR.writeSlice(singleRowBlockWriter, utf8Slice("cat")); - blockBuilder.closeEntry(); + blockBuilder.buildEntry(fieldBuilders -> { + BIGINT.writeLong(fieldBuilders.get(0), 2); + VARCHAR.writeSlice(fieldBuilders.get(1), utf8Slice("cats")); + }); - singleRowBlockWriter = blockBuilder.beginBlockEntry(); - BIGINT.writeLong(singleRowBlockWriter, 2); - VARCHAR.writeSlice(singleRowBlockWriter, utf8Slice("cats")); - blockBuilder.closeEntry(); - - singleRowBlockWriter = blockBuilder.beginBlockEntry(); - BIGINT.writeLong(singleRowBlockWriter, 3); - VARCHAR.writeSlice(singleRowBlockWriter, utf8Slice("dog")); - blockBuilder.closeEntry(); + blockBuilder.buildEntry(fieldBuilders -> { + BIGINT.writeLong(fieldBuilders.get(0), 3); + VARCHAR.writeSlice(fieldBuilders.get(1), utf8Slice("dog")); + }); return blockBuilder.build(); } @@ -66,15 +63,10 @@ private static Block createTestBlock() @Override protected Object getGreaterValue(Object value) { - RowBlockBuilder blockBuilder = (RowBlockBuilder) TYPE.createBlockBuilder(null, 1); - SingleRowBlockWriter singleRowBlockWriter; - - Block block = (Block) value; - singleRowBlockWriter = blockBuilder.beginBlockEntry(); - BIGINT.writeLong(singleRowBlockWriter, block.getSingleValueBlock(0).getLong(0, 0) + 1); - VARCHAR.writeSlice(singleRowBlockWriter, block.getSingleValueBlock(1).getSlice(0, 0, 1)); - blockBuilder.closeEntry(); - - return TYPE.getObject(blockBuilder.build(), 0); + return buildRowValue(TYPE, fieldBuilders -> { + Block block = (Block) value; + BIGINT.writeLong(fieldBuilders.get(0), block.getSingleValueBlock(0).getLong(0, 0) + 1); + VARCHAR.writeSlice(fieldBuilders.get(1), block.getSingleValueBlock(1).getSlice(0, 0, 1)); + }); } } diff --git a/core/trino-main/src/test/java/io/trino/util/StructuralTestUtil.java b/core/trino-main/src/test/java/io/trino/util/StructuralTestUtil.java index d0f46c5d530d..ca12ddaee824 100644 --- a/core/trino-main/src/test/java/io/trino/util/StructuralTestUtil.java +++ b/core/trino-main/src/test/java/io/trino/util/StructuralTestUtil.java @@ -16,8 +16,11 @@ import com.google.common.collect.ImmutableList; import io.airlift.slice.Slice; import io.airlift.slice.Slices; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.Decimals; import io.trino.spi.type.Int128; @@ -31,6 +34,7 @@ import java.math.BigDecimal; import java.util.Map; +import static io.trino.spi.block.MapValueBuilder.buildMapValue; import static io.trino.spi.type.RealType.REAL; import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; import static java.lang.Float.floatToRawIntBits; @@ -48,17 +52,17 @@ public static Block arrayBlockOf(Type elementType, Object... values) return blockBuilder.build(); } - public static Block mapBlockOf(Type keyType, Type valueType, Map value) + public static Block mapBlockOf(Type keyType, Type valueType, Map map) { - MapType mapType = mapType(keyType, valueType); - BlockBuilder mapArrayBuilder = mapType.createBlockBuilder(null, 1); - BlockBuilder singleMapWriter = mapArrayBuilder.beginBlockEntry(); - for (Map.Entry entry : value.entrySet()) { - appendToBlockBuilder(keyType, entry.getKey(), singleMapWriter); - appendToBlockBuilder(valueType, entry.getValue(), singleMapWriter); - } - mapArrayBuilder.closeEntry(); - return mapType.getObject(mapArrayBuilder, 0); + return buildMapValue( + mapType(keyType, valueType), + map.size(), + (keyBuilder, valueBuilder) -> { + map.forEach((key, value) -> { + appendToBlockBuilder(keyType, key, keyBuilder); + appendToBlockBuilder(valueType, value, valueBuilder); + }); + }); } public static MapType mapType(Type keyType, Type valueType) @@ -75,28 +79,28 @@ public static void appendToBlockBuilder(Type type, Object element, BlockBuilder blockBuilder.appendNull(); } else if (type instanceof ArrayType && element instanceof Iterable) { - BlockBuilder subBlockBuilder = blockBuilder.beginBlockEntry(); - for (Object subElement : (Iterable) element) { - appendToBlockBuilder(type.getTypeParameters().get(0), subElement, subBlockBuilder); - } - blockBuilder.closeEntry(); + ((ArrayBlockBuilder) blockBuilder).buildEntry(elementBuilder -> { + for (Object subElement : (Iterable) element) { + appendToBlockBuilder(type.getTypeParameters().get(0), subElement, elementBuilder); + } + }); } else if (type instanceof RowType && element instanceof Iterable) { - BlockBuilder subBlockBuilder = blockBuilder.beginBlockEntry(); - int field = 0; - for (Object subElement : (Iterable) element) { - appendToBlockBuilder(type.getTypeParameters().get(field), subElement, subBlockBuilder); - field++; - } - blockBuilder.closeEntry(); + ((RowBlockBuilder) blockBuilder).buildEntry(fieldBuilders -> { + int field = 0; + for (Object subElement : (Iterable) element) { + appendToBlockBuilder(type.getTypeParameters().get(field), subElement, fieldBuilders.get(field)); + field++; + } + }); } - else if (type instanceof MapType && element instanceof Map) { - BlockBuilder subBlockBuilder = blockBuilder.beginBlockEntry(); - for (Map.Entry entry : ((Map) element).entrySet()) { - appendToBlockBuilder(type.getTypeParameters().get(0), entry.getKey(), subBlockBuilder); - appendToBlockBuilder(type.getTypeParameters().get(1), entry.getValue(), subBlockBuilder); - } - blockBuilder.closeEntry(); + else if (type instanceof MapType mapType && element instanceof Map) { + ((MapBlockBuilder) blockBuilder).buildEntry((keyBuilder, valueBuilder) -> { + for (Map.Entry entry : ((Map) element).entrySet()) { + appendToBlockBuilder(mapType.getKeyType(), entry.getKey(), keyBuilder); + appendToBlockBuilder(mapType.getValueType(), entry.getValue(), valueBuilder); + } + }); } else if (javaType == boolean.class) { type.writeBoolean(blockBuilder, (Boolean) element); diff --git a/core/trino-spi/pom.xml b/core/trino-spi/pom.xml index 1e4f8213f4c6..019ea450582a 100644 --- a/core/trino-spi/pom.xml +++ b/core/trino-spi/pom.xml @@ -234,25 +234,222 @@ - true - java.method.returnTypeTypeParametersChanged - method java.util.Set<io.trino.spi.ptf.ConnectorTableFunction> io.trino.spi.connector.Connector::getTableFunctions() - method java.util.Set<io.trino.spi.function.table.ConnectorTableFunction> io.trino.spi.connector.Connector::getTableFunctions() - Relocate table function infrastructure + java.method.removed + method io.trino.spi.block.SingleArrayBlockWriter io.trino.spi.block.ArrayBlockBuilder::beginBlockEntry() - true - java.method.parameterTypeChanged - parameter java.util.Optional<io.trino.spi.connector.TableFunctionApplicationResult<io.trino.spi.connector.ConnectorTableHandle>> io.trino.spi.connector.ConnectorMetadata::applyTableFunction(io.trino.spi.connector.ConnectorSession, ===io.trino.spi.ptf.ConnectorTableFunctionHandle===) - parameter java.util.Optional<io.trino.spi.connector.TableFunctionApplicationResult<io.trino.spi.connector.ConnectorTableHandle>> io.trino.spi.connector.ConnectorMetadata::applyTableFunction(io.trino.spi.connector.ConnectorSession, ===io.trino.spi.function.table.ConnectorTableFunctionHandle===) - 1 - Relocate table function infrastructure + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.BlockBuilder::beginBlockEntry() + + + java.method.removed + method io.trino.spi.block.SingleMapBlockWriter io.trino.spi.block.MapBlockBuilder::beginBlockEntry() + + + java.method.removed + method io.trino.spi.block.SingleRowBlockWriter io.trino.spi.block.RowBlockBuilder::beginBlockEntry() - true java.class.removed - interface io.trino.spi.ptf.TableFunctionDataProcessor - Relocate table function infrastructure + class io.trino.spi.block.SingleArrayBlockWriter + + + java.class.removed + class io.trino.spi.block.SingleMapBlockWriter + + + java.class.removed + class io.trino.spi.block.SingleRowBlockWriter + + + java.method.returnTypeChangedCovariantly + method io.trino.spi.block.BlockBuilder io.trino.spi.type.ArrayType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + method io.trino.spi.block.ArrayBlockBuilder io.trino.spi.type.ArrayType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + + + java.method.returnTypeChangedCovariantly + method io.trino.spi.block.BlockBuilder io.trino.spi.type.ArrayType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int, int) + method io.trino.spi.block.ArrayBlockBuilder io.trino.spi.type.ArrayType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int, int) + + + java.method.returnTypeChangedCovariantly + method io.trino.spi.block.BlockBuilder io.trino.spi.type.MapType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + method io.trino.spi.block.MapBlockBuilder io.trino.spi.type.MapType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + + + java.method.returnTypeChangedCovariantly + method io.trino.spi.block.BlockBuilder io.trino.spi.type.MapType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int, int) + method io.trino.spi.block.MapBlockBuilder io.trino.spi.type.MapType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int, int) + + + java.method.returnTypeChangedCovariantly + method io.trino.spi.block.BlockBuilder io.trino.spi.type.RowType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + method io.trino.spi.block.RowBlockBuilder io.trino.spi.type.RowType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + + + java.method.returnTypeChangedCovariantly + method io.trino.spi.block.BlockBuilder io.trino.spi.type.RowType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int, int) + method io.trino.spi.block.RowBlockBuilder io.trino.spi.type.RowType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int, int) + + + java.method.removed + method void io.trino.spi.block.AbstractSingleArrayBlock::writeBytesTo(int, int, int, io.trino.spi.block.BlockBuilder) + + + java.method.removed + method void io.trino.spi.block.AbstractSingleMapBlock::writeBytesTo(int, int, int, io.trino.spi.block.BlockBuilder) + + + java.method.removed + method void io.trino.spi.block.AbstractSingleRowBlock::writeBytesTo(int, int, int, io.trino.spi.block.BlockBuilder) + + + java.method.removed + method void io.trino.spi.block.AbstractVariableWidthBlock::writeBytesTo(int, int, int, io.trino.spi.block.BlockBuilder) + + + java.method.removed + method void io.trino.spi.block.Block::writeBytesTo(int, int, int, io.trino.spi.block.BlockBuilder) + + + java.method.removed + method void io.trino.spi.block.DictionaryBlock::writeBytesTo(int, int, int, io.trino.spi.block.BlockBuilder) + + + java.method.removed + method void io.trino.spi.block.LazyBlock::writeBytesTo(int, int, int, io.trino.spi.block.BlockBuilder) + + + java.method.removed + method void io.trino.spi.block.RunLengthEncodedBlock::writeBytesTo(int, int, int, io.trino.spi.block.BlockBuilder) + + + java.method.returnTypeChangedCovariantly + method io.trino.spi.block.BlockBuilder io.trino.spi.type.AbstractVariableWidthType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + method io.trino.spi.block.VariableWidthBlockBuilder io.trino.spi.type.AbstractVariableWidthType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + + + java.method.returnTypeChangedCovariantly + method io.trino.spi.block.BlockBuilder io.trino.spi.type.AbstractVariableWidthType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int, int) + method io.trino.spi.block.VariableWidthBlockBuilder io.trino.spi.type.AbstractVariableWidthType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int, int) + + + java.method.returnTypeChangedCovariantly + method io.trino.spi.block.BlockBuilder io.trino.spi.type.CharType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + method io.trino.spi.block.VariableWidthBlockBuilder io.trino.spi.type.CharType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + + + java.method.returnTypeChangedCovariantly + method io.trino.spi.block.BlockBuilder io.trino.spi.type.VarcharType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + method io.trino.spi.block.VariableWidthBlockBuilder io.trino.spi.type.VarcharType::createBlockBuilder(io.trino.spi.block.BlockBuilderStatus, int) + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.ArrayBlockBuilder::closeEntry() + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.BlockBuilder::closeEntry() + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.ByteArrayBlockBuilder::closeEntry() + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.Fixed12BlockBuilder::closeEntry() + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.Int128ArrayBlockBuilder::closeEntry() + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.IntArrayBlockBuilder::closeEntry() + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.LongArrayBlockBuilder::closeEntry() + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.MapBlockBuilder::closeEntry() + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.RowBlockBuilder::closeEntry() + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.ShortArrayBlockBuilder::closeEntry() + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.VariableWidthBlockBuilder::closeEntry() + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.BlockBuilder::writeBytes(io.airlift.slice.Slice, int, int) + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.VariableWidthBlockBuilder::writeBytes(io.airlift.slice.Slice, int, int) + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.BlockBuilder::writeByte(int) + + + java.method.parameterTypeChanged + parameter io.trino.spi.block.BlockBuilder io.trino.spi.block.ByteArrayBlockBuilder::writeByte(===int===) + parameter io.trino.spi.block.BlockBuilder io.trino.spi.block.ByteArrayBlockBuilder::writeByte(===byte===) + 0 + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.VariableWidthBlockBuilder::writeByte(int) + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.BlockBuilder::writeShort(int) + + + java.method.parameterTypeChanged + parameter io.trino.spi.block.BlockBuilder io.trino.spi.block.ShortArrayBlockBuilder::writeShort(===int===) + parameter io.trino.spi.block.BlockBuilder io.trino.spi.block.ShortArrayBlockBuilder::writeShort(===short===) + 0 + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.VariableWidthBlockBuilder::writeShort(int) + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.BlockBuilder::writeInt(int) + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.Fixed12BlockBuilder::writeInt(int) + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.VariableWidthBlockBuilder::writeInt(int) + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.BlockBuilder::writeLong(long) + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.Fixed12BlockBuilder::writeLong(long) + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.Int128ArrayBlockBuilder::writeLong(long) + + + java.method.removed + method io.trino.spi.block.BlockBuilder io.trino.spi.block.VariableWidthBlockBuilder::writeLong(long) diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleArrayBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleArrayBlock.java index 13fe12bb783b..9c8d71902f41 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleArrayBlock.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleArrayBlock.java @@ -14,6 +14,7 @@ package io.trino.spi.block; import io.airlift.slice.Slice; +import io.airlift.slice.SliceOutput; import java.util.List; @@ -80,6 +81,13 @@ public Slice getSlice(int position, int offset, int length) return getBlock().getSlice(position + start, offset, length); } + @Override + public void writeSliceTo(int position, int offset, int length, SliceOutput output) + { + checkReadablePosition(this, position); + getBlock().writeSliceTo(position + start, offset, length, output); + } + @Override public T getObject(int position, Class clazz) { @@ -101,13 +109,6 @@ public int bytesCompare(int position, int offset, int length, Slice otherSlice, return getBlock().bytesCompare(position + start, offset, length, otherSlice, otherOffset, otherLength); } - @Override - public void writeBytesTo(int position, int offset, int length, BlockBuilder blockBuilder) - { - checkReadablePosition(this, position); - getBlock().writeBytesTo(position + start, offset, length, blockBuilder); - } - @Override public boolean equals(int position, int offset, Block otherBlock, int otherPosition, int otherOffset, int length) { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleMapBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleMapBlock.java index 72d5b2fb7d95..b63df6bb750e 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleMapBlock.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleMapBlock.java @@ -15,6 +15,7 @@ package io.trino.spi.block; import io.airlift.slice.Slice; +import io.airlift.slice.SliceOutput; import java.util.List; @@ -104,6 +105,18 @@ public Slice getSlice(int position, int offset, int length) return getRawValueBlock().getSlice(position / 2, offset, length); } + @Override + public void writeSliceTo(int position, int offset, int length, SliceOutput output) + { + position = getAbsolutePosition(position); + if (position % 2 == 0) { + getRawKeyBlock().writeSliceTo(position / 2, offset, length, output); + } + else { + getRawValueBlock().writeSliceTo(position / 2, offset, length, output); + } + } + @Override public int getSliceLength(int position) { @@ -144,18 +157,6 @@ public int bytesCompare(int position, int offset, int length, Slice otherSlice, return getRawValueBlock().bytesCompare(position / 2, offset, length, otherSlice, otherOffset, otherLength); } - @Override - public void writeBytesTo(int position, int offset, int length, BlockBuilder blockBuilder) - { - position = getAbsolutePosition(position); - if (position % 2 == 0) { - getRawKeyBlock().writeBytesTo(position / 2, offset, length, blockBuilder); - } - else { - getRawValueBlock().writeBytesTo(position / 2, offset, length, blockBuilder); - } - } - @Override public boolean equals(int position, int offset, Block otherBlock, int otherPosition, int otherOffset, int length) { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleRowBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleRowBlock.java index 22d736dd54da..63b7cad095e8 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleRowBlock.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/AbstractSingleRowBlock.java @@ -15,6 +15,7 @@ package io.trino.spi.block; import io.airlift.slice.Slice; +import io.airlift.slice.SliceOutput; import java.util.List; @@ -82,6 +83,13 @@ public Slice getSlice(int position, int offset, int length) return getRawFieldBlock(position).getSlice(getRowIndex(), offset, length); } + @Override + public void writeSliceTo(int position, int offset, int length, SliceOutput output) + { + checkFieldIndex(position); + getRawFieldBlock(position).writeSliceTo(getRowIndex(), offset, length, output); + } + @Override public int getSliceLength(int position) { @@ -110,13 +118,6 @@ public int bytesCompare(int position, int offset, int length, Slice otherSlice, return getRawFieldBlock(position).bytesCompare(getRowIndex(), offset, length, otherSlice, otherOffset, otherLength); } - @Override - public void writeBytesTo(int position, int offset, int length, BlockBuilder blockBuilder) - { - checkFieldIndex(position); - getRawFieldBlock(position).writeBytesTo(getRowIndex(), offset, length, blockBuilder); - } - @Override public boolean equals(int position, int offset, Block otherBlock, int otherPosition, int otherOffset, int length) { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/AbstractVariableWidthBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/AbstractVariableWidthBlock.java index 597c4833c365..5c4da27502e0 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/AbstractVariableWidthBlock.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/AbstractVariableWidthBlock.java @@ -14,6 +14,7 @@ package io.trino.spi.block; import io.airlift.slice.Slice; +import io.airlift.slice.SliceOutput; import io.airlift.slice.Slices; import io.airlift.slice.XxHash64; @@ -70,6 +71,13 @@ public Slice getSlice(int position, int offset, int length) return getRawSlice(position).slice(getPositionOffset(position) + offset, length); } + @Override + public void writeSliceTo(int position, int offset, int length, SliceOutput output) + { + checkReadablePosition(this, position); + output.writeBytes(getRawSlice(position), getPositionOffset(position) + offset, length); + } + @Override public boolean equals(int position, int offset, Block otherBlock, int otherPosition, int otherOffset, int length) { @@ -113,13 +121,6 @@ public int bytesCompare(int position, int offset, int length, Slice otherSlice, return getRawSlice(position).compareTo(getPositionOffset(position) + offset, length, otherSlice, otherOffset, otherLength); } - @Override - public void writeBytesTo(int position, int offset, int length, BlockBuilder blockBuilder) - { - checkReadablePosition(this, position); - blockBuilder.writeBytes(getRawSlice(position), getPositionOffset(position) + offset, length); - } - @Override public Block getSingleValueBlock(int position) { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/ArrayBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/ArrayBlockBuilder.java index 2b53dc8f952e..d660e6035465 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/ArrayBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/ArrayBlockBuilder.java @@ -148,26 +148,17 @@ public boolean mayHaveNull() return hasNullValue; } - @Override - public SingleArrayBlockWriter beginBlockEntry() + public void buildEntry(ArrayValueBuilder builder) + throws E { if (currentEntryOpened) { throw new IllegalStateException("Expected current entry to be closed but was opened"); } - currentEntryOpened = true; - return new SingleArrayBlockWriter(values, values.getPositionCount()); - } - - @Override - public BlockBuilder closeEntry() - { - if (!currentEntryOpened) { - throw new IllegalStateException("Expected entry to be opened but was closed"); - } + currentEntryOpened = true; + builder.build(values); entryAdded(false); currentEntryOpened = false; - return this; } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/ArrayValueBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/ArrayValueBuilder.java new file mode 100644 index 000000000000..89ce9c000254 --- /dev/null +++ b/core/trino-spi/src/main/java/io/trino/spi/block/ArrayValueBuilder.java @@ -0,0 +1,29 @@ +/* + * Licensed 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.trino.spi.block; + +import io.trino.spi.type.ArrayType; + +public interface ArrayValueBuilder +{ + static Block buildArrayValue(ArrayType arrayType, int entryCount, ArrayValueBuilder builder) + throws E + { + return new BufferedArrayValueBuilder(arrayType, 1) + .build(entryCount, builder); + } + + void build(BlockBuilder elementBuilder) + throws E; +} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/Block.java b/core/trino-spi/src/main/java/io/trino/spi/block/Block.java index 057ba169447e..3281bb7e6727 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/Block.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/Block.java @@ -14,6 +14,7 @@ package io.trino.spi.block; import io.airlift.slice.Slice; +import io.airlift.slice.SliceOutput; import java.util.Collections; import java.util.List; @@ -73,6 +74,14 @@ default Slice getSlice(int position, int offset, int length) throw new UnsupportedOperationException(getClass().getName()); } + /** + * Writes a slice at {@code offset} in the value at {@code position} into the {@code output} slice output. + */ + default void writeSliceTo(int position, int offset, int length, SliceOutput output) + { + throw new UnsupportedOperationException(getClass().getName()); + } + /** * Gets an object in the value at {@code position}. */ @@ -101,16 +110,6 @@ default int bytesCompare(int position, int offset, int length, Slice otherSlice, throw new UnsupportedOperationException(getClass().getName()); } - /** - * Appends the byte sequences at {@code offset} in the value at {@code position} - * to {@code blockBuilder}. - * This method must be implemented if @{code getSlice} is implemented. - */ - default void writeBytesTo(int position, int offset, int length, BlockBuilder blockBuilder) - { - throw new UnsupportedOperationException(getClass().getName()); - } - /** * Is the byte sequences at {@code offset} in the value at {@code position} equal * to the byte sequence at {@code otherOffset} in the value at {@code otherPosition} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/BlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/BlockBuilder.java index 47ac89760802..47aaf3e2736e 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/BlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/BlockBuilder.java @@ -13,62 +13,11 @@ */ package io.trino.spi.block; -import io.airlift.slice.Slice; - import static io.trino.spi.block.BlockUtil.calculateBlockResetSize; public interface BlockBuilder extends Block { - /** - * Write a byte to the current entry; - */ - default BlockBuilder writeByte(int value) - { - throw new UnsupportedOperationException(getClass().getName()); - } - - /** - * Write a short to the current entry; - */ - default BlockBuilder writeShort(int value) - { - throw new UnsupportedOperationException(getClass().getName()); - } - - /** - * Write a int to the current entry; - */ - default BlockBuilder writeInt(int value) - { - throw new UnsupportedOperationException(getClass().getName()); - } - - /** - * Write a long to the current entry; - */ - default BlockBuilder writeLong(long value) - { - throw new UnsupportedOperationException(getClass().getName()); - } - - /** - * Write a byte sequences to the current entry; - */ - default BlockBuilder writeBytes(Slice source, int sourceIndex, int length) - { - throw new UnsupportedOperationException(getClass().getName()); - } - - /** - * Return a writer to the current entry. The caller can operate on the returned caller to incrementally build the object. This is generally more efficient than - * building the object elsewhere and call writeObject afterwards because a large chunk of memory could potentially be unnecessarily copied in this process. - */ - default BlockBuilder beginBlockEntry() - { - throw new UnsupportedOperationException(getClass().getName()); - } - /** * Create a new block from the current materialized block by keeping the same elements * only with respect to {@code visiblePositions}. @@ -79,11 +28,6 @@ default Block getPositions(int[] visiblePositions, int offset, int length) return build().getPositions(visiblePositions, offset, length); } - /** - * Close the current entry. - */ - BlockBuilder closeEntry(); - /** * Appends a null value to the block. */ diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/BufferedArrayValueBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/BufferedArrayValueBuilder.java new file mode 100644 index 000000000000..a744880e2711 --- /dev/null +++ b/core/trino-spi/src/main/java/io/trino/spi/block/BufferedArrayValueBuilder.java @@ -0,0 +1,61 @@ +/* + * Licensed 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.trino.spi.block; + +import io.trino.spi.type.ArrayType; + +import static io.airlift.slice.SizeOf.instanceSize; + +public class BufferedArrayValueBuilder +{ + private static final int INSTANCE_SIZE = instanceSize(BufferedArrayValueBuilder.class); + + private int bufferSize; + private BlockBuilder valueBuilder; + + public static BufferedArrayValueBuilder createBuffered(ArrayType arrayType) + { + return new BufferedArrayValueBuilder(arrayType, 1024); + } + + BufferedArrayValueBuilder(ArrayType arrayType, int bufferSize) + { + this.bufferSize = bufferSize; + this.valueBuilder = arrayType.getElementType().createBlockBuilder(null, bufferSize); + } + + public long getRetainedSizeInBytes() + { + return INSTANCE_SIZE + valueBuilder.getRetainedSizeInBytes(); + } + + public Block build(int entryCount, ArrayValueBuilder builder) + throws E + { + // grow or reset builders if necessary + if (valueBuilder.getPositionCount() + entryCount > bufferSize) { + if (bufferSize < entryCount) { + bufferSize = entryCount; + } + valueBuilder = valueBuilder.newBlockBuilderLike(bufferSize, null); + } + + int startSize = valueBuilder.getPositionCount(); + + builder.build(valueBuilder); + + int endSize = valueBuilder.getPositionCount(); + return valueBuilder.build().getRegion(startSize, endSize - startSize); + } +} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/BufferedMapValueBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/BufferedMapValueBuilder.java new file mode 100644 index 000000000000..4474fc6fd5b6 --- /dev/null +++ b/core/trino-spi/src/main/java/io/trino/spi/block/BufferedMapValueBuilder.java @@ -0,0 +1,148 @@ +/* + * Licensed 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.trino.spi.block; + +import io.trino.spi.type.MapType; + +import java.util.Arrays; +import java.util.Optional; + +import static io.airlift.slice.SizeOf.instanceSize; +import static io.trino.spi.block.MapHashTables.HASH_MULTIPLIER; +import static java.lang.Math.max; +import static java.util.Objects.requireNonNull; + +public class BufferedMapValueBuilder +{ + private static final int INSTANCE_SIZE = instanceSize(BufferedMapValueBuilder.class); + + enum HashBuildMode + { + DUPLICATE_NOT_CHECKED, STRICT_EQUALS, STRICT_NOT_DISTINCT_FROM + } + + private final MapType mapType; + private final HashBuildMode hashBuildMode; + private BlockBuilder keyBlockBuilder; + private BlockBuilder valueBlockBuilder; + private int bufferSize; + + public static BufferedMapValueBuilder createBuffered(MapType mapType) + { + return new BufferedMapValueBuilder(mapType, HashBuildMode.DUPLICATE_NOT_CHECKED, 1024); + } + + public static BufferedMapValueBuilder createBufferedStrict(MapType mapType) + { + return new BufferedMapValueBuilder(mapType, HashBuildMode.STRICT_EQUALS, 1024); + } + + public static BufferedMapValueBuilder createBufferedDistinctStrict(MapType mapType) + { + return new BufferedMapValueBuilder(mapType, HashBuildMode.STRICT_NOT_DISTINCT_FROM, 1024); + } + + BufferedMapValueBuilder(MapType mapType, HashBuildMode hashBuildMode, int bufferSize) + { + this.mapType = requireNonNull(mapType, "mapType is null"); + this.hashBuildMode = hashBuildMode; + this.keyBlockBuilder = mapType.getKeyType().createBlockBuilder(null, bufferSize); + this.valueBlockBuilder = mapType.getValueType().createBlockBuilder(null, bufferSize); + this.bufferSize = bufferSize; + } + + public long getRetainedSizeInBytes() + { + return INSTANCE_SIZE + keyBlockBuilder.getRetainedSizeInBytes() + valueBlockBuilder.getRetainedSizeInBytes(); + } + + public Block build(int entryCount, MapValueBuilder builder) + throws E + { + if (keyBlockBuilder.getPositionCount() != valueBlockBuilder.getPositionCount()) { + // we could fix this by appending nulls to the shorter builder, but this is a sign the buffer is being used in a multithreaded environment which is not supported + throw new IllegalStateException("Key and value builders were corrupted by a previous call to buildValue"); + } + + // grow or reset builders if necessary + if (keyBlockBuilder.getPositionCount() + entryCount > bufferSize) { + if (bufferSize < entryCount) { + bufferSize = entryCount; + } + keyBlockBuilder = keyBlockBuilder.newBlockBuilderLike(bufferSize, null); + valueBlockBuilder = valueBlockBuilder.newBlockBuilderLike(bufferSize, null); + } + + int startSize = keyBlockBuilder.getPositionCount(); + + // build the map + try { + builder.build(keyBlockBuilder, valueBlockBuilder); + } + catch (Exception e) { + equalizeBlockBuilders(); + throw e; + } + + // check that key and value builders have the same size + if (equalizeBlockBuilders()) { + throw new IllegalStateException("Expected key and value builders to have the same size"); + } + int endSize = keyBlockBuilder.getPositionCount(); + + // build the map block + Block keyBlock = keyBlockBuilder.build().getRegion(startSize, endSize - startSize); + Block valueBlock = valueBlockBuilder.build().getRegion(startSize, endSize - startSize); + + // build the hash table + int[] table = new int[keyBlock.getPositionCount() * HASH_MULTIPLIER]; + Arrays.fill(table, -1); + MapHashTables hashTables = new MapHashTables(mapType, Optional.of(table)); + switch (hashBuildMode) { + case DUPLICATE_NOT_CHECKED -> hashTables.buildHashTable(keyBlock, 0, keyBlock.getPositionCount()); + case STRICT_EQUALS -> hashTables.buildHashTableStrict(keyBlock, 0, keyBlock.getPositionCount()); + case STRICT_NOT_DISTINCT_FROM -> hashTables.buildDistinctHashTableStrict(keyBlock, 0, keyBlock.getPositionCount()); + } + + MapBlock mapBlock = MapBlock.createMapBlockInternal( + mapType, + 0, + 1, + Optional.empty(), + new int[] {0, keyBlock.getPositionCount()}, + keyBlock, + valueBlock, + hashTables); + + return new SingleMapBlock(0, keyBlock.getPositionCount() * 2, mapBlock); + } + + private boolean equalizeBlockBuilders() + { + int keyBlockSize = keyBlockBuilder.getPositionCount(); + if (keyBlockSize == valueBlockBuilder.getPositionCount()) { + return false; + } + + // append nulls to even out the blocks + int expectedSize = max(keyBlockSize, valueBlockBuilder.getPositionCount()); + while (keyBlockBuilder.getPositionCount() < expectedSize) { + keyBlockBuilder.appendNull(); + } + while (valueBlockBuilder.getPositionCount() < expectedSize) { + valueBlockBuilder.appendNull(); + } + return true; + } +} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/BufferedRowValueBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/BufferedRowValueBuilder.java new file mode 100644 index 000000000000..3bdc882c995d --- /dev/null +++ b/core/trino-spi/src/main/java/io/trino/spi/block/BufferedRowValueBuilder.java @@ -0,0 +1,101 @@ +/* + * Licensed 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.trino.spi.block; + +import io.trino.spi.type.RowType; + +import java.util.List; + +import static io.airlift.slice.SizeOf.instanceSize; + +public class BufferedRowValueBuilder +{ + private static final int INSTANCE_SIZE = instanceSize(BufferedRowValueBuilder.class); + + private final int bufferSize; + private List fieldBuilders; + + public static BufferedRowValueBuilder createBuffered(RowType rowType) + { + return new BufferedRowValueBuilder(rowType, 1024); + } + + BufferedRowValueBuilder(RowType rowType, int bufferSize) + { + this.bufferSize = bufferSize; + this.fieldBuilders = rowType.getTypeParameters().stream() + .map(fieldType -> fieldType.createBlockBuilder(null, bufferSize)) + .toList(); + } + + public long getRetainedSizeInBytes() + { + return INSTANCE_SIZE + fieldBuilders.stream().mapToLong(BlockBuilder::getRetainedSizeInBytes).sum(); + } + + public Block build(RowValueBuilder builder) + throws E + { + int expectedSize = fieldBuilders.get(0).getPositionCount(); + if (!fieldBuilders.stream().allMatch(field -> field.getPositionCount() == expectedSize)) { + // we could fix this by appending nulls to the shorter builders, but this is a sign the buffer is being used in a multithreaded environment which is not supported + throw new IllegalStateException("Field builders were corrupted by a previous call to buildValue"); + } + + // grow or reset builders if necessary + if (fieldBuilders.get(0).getPositionCount() + 1 > bufferSize) { + fieldBuilders = fieldBuilders.stream() + .map(field -> field.newBlockBuilderLike(bufferSize, null)) + .toList(); + } + + int startSize = fieldBuilders.get(0).getPositionCount(); + + try { + builder.build(fieldBuilders); + } + catch (Exception e) { + equalizeBlockBuilders(); + throw e; + } + + // check that field builders have the same size + if (equalizeBlockBuilders()) { + throw new IllegalStateException("Expected field builders to have the same size"); + } + int endSize = fieldBuilders.get(0).getPositionCount(); + if (endSize != startSize + 1) { + throw new IllegalStateException("Expected exactly one entry added to each field builder"); + } + + List blocks = fieldBuilders.stream() + .map(field -> field.build().getRegion(startSize, 1)) + .toList(); + return new SingleRowBlock(0, blocks.toArray(new Block[0])); + } + + private boolean equalizeBlockBuilders() + { + // append nulls to even out the blocks + boolean nullsAppended = false; + int newBlockSize = fieldBuilders.stream().mapToInt(BlockBuilder::getPositionCount).max().orElseThrow(); + for (BlockBuilder fieldBuilder : fieldBuilders) { + while (fieldBuilder.getPositionCount() < newBlockSize) { + fieldBuilder.appendNull(); + nullsAppended = true; + } + } + return nullsAppended; + } +} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/ByteArrayBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/ByteArrayBlockBuilder.java index 6223b87de036..acbd7cd8beff 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/ByteArrayBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/ByteArrayBlockBuilder.java @@ -58,14 +58,13 @@ public ByteArrayBlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, in updateDataSize(); } - @Override - public BlockBuilder writeByte(int value) + public BlockBuilder writeByte(byte value) { if (values.length <= positionCount) { growCapacity(); } - values[positionCount] = (byte) value; + values[positionCount] = value; hasNonNullValue = true; positionCount++; @@ -75,12 +74,6 @@ public BlockBuilder writeByte(int value) return this; } - @Override - public BlockBuilder closeEntry() - { - return this; - } - @Override public BlockBuilder appendNull() { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/DictionaryBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/DictionaryBlock.java index 38a8972fbfb1..f48fa8b1f16b 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/DictionaryBlock.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/DictionaryBlock.java @@ -14,6 +14,7 @@ package io.trino.spi.block; import io.airlift.slice.Slice; +import io.airlift.slice.SliceOutput; import io.airlift.slice.Slices; import java.util.ArrayList; @@ -176,6 +177,12 @@ public Slice getSlice(int position, int offset, int length) return dictionary.getSlice(getId(position), offset, length); } + @Override + public void writeSliceTo(int position, int offset, int length, SliceOutput output) + { + dictionary.writeSliceTo(getId(position), offset, length, output); + } + @Override public T getObject(int position, Class clazz) { @@ -194,12 +201,6 @@ public int bytesCompare(int position, int offset, int length, Slice otherSlice, return dictionary.bytesCompare(getId(position), offset, length, otherSlice, otherOffset, otherLength); } - @Override - public void writeBytesTo(int position, int offset, int length, BlockBuilder blockBuilder) - { - dictionary.writeBytesTo(getId(position), offset, length, blockBuilder); - } - @Override public boolean equals(int position, int offset, Block otherBlock, int otherPosition, int otherOffset, int length) { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/Fixed12Block.java b/core/trino-spi/src/main/java/io/trino/spi/block/Fixed12Block.java index 9466814fdbf4..f6076e6a470b 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/Fixed12Block.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/Fixed12Block.java @@ -255,14 +255,10 @@ public String toString() */ public static void encodeFixed12(long first, int second, int[] values, int position) { - encodeFirst(first, values, position); - values[position * 3 + 2] = second; - } - - static void encodeFirst(long first, int[] values, int position) - { - values[(position * 3)] = (int) first; - values[(position * 3) + 1] = (int) (first >>> 32); + int entryPosition = position * 3; + values[entryPosition] = (int) first; + values[entryPosition + 1] = (int) (first >>> 32); + values[entryPosition + 2] = second; } /** diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/Fixed12BlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/Fixed12BlockBuilder.java index 3d3021f75e9d..38b9c5f20bfc 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/Fixed12BlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/Fixed12BlockBuilder.java @@ -22,7 +22,6 @@ import java.util.OptionalInt; import java.util.function.ObjLongConsumer; -import static io.airlift.slice.SizeOf.SIZE_OF_INT; import static io.airlift.slice.SizeOf.instanceSize; import static io.airlift.slice.SizeOf.sizeOf; import static io.trino.spi.block.BlockUtil.checkArrayRange; @@ -31,7 +30,7 @@ import static io.trino.spi.block.BlockUtil.compactArray; import static io.trino.spi.block.Fixed12Block.FIXED12_BYTES; import static io.trino.spi.block.Fixed12Block.decodeFixed12First; -import static io.trino.spi.block.Fixed12Block.encodeFirst; +import static io.trino.spi.block.Fixed12Block.encodeFixed12; import static java.lang.Math.max; public class Fixed12BlockBuilder @@ -55,8 +54,6 @@ public class Fixed12BlockBuilder private long retainedSizeInBytes; - private int entryPositionCount; - public Fixed12BlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, int expectedEntries) { this.blockBuilderStatus = blockBuilderStatus; @@ -65,60 +62,24 @@ public Fixed12BlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, int updateDataSize(); } - @Override - public BlockBuilder writeLong(long value) + public void writeFixed12(long first, int second) { - if (entryPositionCount != 0) { - throw new IllegalArgumentException("long can only be written at the beginning of the entry"); - } - if (valueIsNull.length <= positionCount) { growCapacity(); } - encodeFirst(value, values, positionCount); - entryPositionCount += 2; + encodeFixed12(first, second, values, positionCount); hasNonNullValue = true; - return this; - } - - @Override - public BlockBuilder writeInt(int value) - { - if (valueIsNull.length <= positionCount) { - growCapacity(); - } - - values[(positionCount * 3) + entryPositionCount] = value; - entryPositionCount++; - - hasNonNullValue = true; - return this; - } - - @Override - public BlockBuilder closeEntry() - { - if (entryPositionCount != 3) { - throw new IllegalStateException("Expected entry size to be exactly " + FIXED12_BYTES + " bytes but was " + (entryPositionCount * SIZE_OF_INT)); - } - positionCount++; - entryPositionCount = 0; if (blockBuilderStatus != null) { blockBuilderStatus.addBytes(Byte.BYTES + FIXED12_BYTES); } - return this; } @Override public BlockBuilder appendNull() { - if (entryPositionCount != 0) { - throw new IllegalStateException("Current entry must be closed before a null can be written"); - } - if (valueIsNull.length <= positionCount) { growCapacity(); } diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/Int128ArrayBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/Int128ArrayBlockBuilder.java index e40c87b6747f..692e48610ce2 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/Int128ArrayBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/Int128ArrayBlockBuilder.java @@ -22,7 +22,6 @@ import java.util.OptionalInt; import java.util.function.ObjLongConsumer; -import static io.airlift.slice.SizeOf.SIZE_OF_LONG; import static io.airlift.slice.SizeOf.instanceSize; import static io.airlift.slice.SizeOf.sizeOf; import static io.trino.spi.block.BlockUtil.checkArrayRange; @@ -53,8 +52,6 @@ public class Int128ArrayBlockBuilder private long retainedSizeInBytes; - private int entryPositionCount; - public Int128ArrayBlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, int expectedEntries) { this.blockBuilderStatus = blockBuilderStatus; @@ -63,33 +60,21 @@ public Int128ArrayBlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, updateDataSize(); } - @Override - public BlockBuilder writeLong(long value) + public void writeInt128(long high, long low) { if (valueIsNull.length <= positionCount) { growCapacity(); } - values[(positionCount * 2) + entryPositionCount] = value; - entryPositionCount++; + int valueIndex = positionCount * 2; + values[valueIndex] = high; + values[valueIndex + 1] = low; hasNonNullValue = true; - return this; - } - - @Override - public BlockBuilder closeEntry() - { - if (entryPositionCount != 2) { - throw new IllegalStateException("Expected entry size to be exactly " + INT128_BYTES + " bytes but was " + (entryPositionCount * SIZE_OF_LONG)); - } - positionCount++; - entryPositionCount = 0; if (blockBuilderStatus != null) { blockBuilderStatus.addBytes(Int128ArrayBlock.SIZE_IN_BYTES_PER_POSITION); } - return this; } @Override @@ -98,9 +83,6 @@ public BlockBuilder appendNull() if (valueIsNull.length <= positionCount) { growCapacity(); } - if (entryPositionCount != 0) { - throw new IllegalStateException("Current entry must be closed before a null can be written"); - } valueIsNull[positionCount] = true; diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/IntArrayBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/IntArrayBlockBuilder.java index 41ff815c7d94..3d5f1125d49a 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/IntArrayBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/IntArrayBlockBuilder.java @@ -58,7 +58,6 @@ public IntArrayBlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, int updateDataSize(); } - @Override public BlockBuilder writeInt(int value) { if (values.length <= positionCount) { @@ -75,12 +74,6 @@ public BlockBuilder writeInt(int value) return this; } - @Override - public BlockBuilder closeEntry() - { - return this; - } - @Override public BlockBuilder appendNull() { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/LazyBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/LazyBlock.java index 1b428cd55656..3bc23a18a8cc 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/LazyBlock.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/LazyBlock.java @@ -14,6 +14,7 @@ package io.trino.spi.block; import io.airlift.slice.Slice; +import io.airlift.slice.SliceOutput; import javax.annotation.Nullable; import javax.annotation.concurrent.NotThreadSafe; @@ -88,6 +89,12 @@ public Slice getSlice(int position, int offset, int length) return getBlock().getSlice(position, offset, length); } + @Override + public void writeSliceTo(int position, int offset, int length, SliceOutput output) + { + getBlock().writeSliceTo(position, offset, length, output); + } + @Override public T getObject(int position, Class clazz) { @@ -112,12 +119,6 @@ public int bytesCompare(int position, int offset, int length, Slice otherSlice, otherLength); } - @Override - public void writeBytesTo(int position, int offset, int length, BlockBuilder blockBuilder) - { - getBlock().writeBytesTo(position, offset, length, blockBuilder); - } - @Override public boolean equals(int position, int offset, Block otherBlock, int otherPosition, int otherOffset, int length) { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/LongArrayBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/LongArrayBlockBuilder.java index 9ca2e64a225f..9a36b0a5ddfd 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/LongArrayBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/LongArrayBlockBuilder.java @@ -59,7 +59,6 @@ public LongArrayBlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, in updateDataSize(); } - @Override public BlockBuilder writeLong(long value) { if (values.length <= positionCount) { @@ -76,12 +75,6 @@ public BlockBuilder writeLong(long value) return this; } - @Override - public BlockBuilder closeEntry() - { - return this; - } - @Override public BlockBuilder appendNull() { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/MapBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/MapBlockBuilder.java index fe96206b0cfa..c64276b4e1ab 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/MapBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/MapBlockBuilder.java @@ -36,6 +36,11 @@ public class MapBlockBuilder { private static final int INSTANCE_SIZE = instanceSize(MapBlockBuilder.class); + private enum HashBuildMode + { + DUPLICATE_NOT_CHECKED, STRICT_EQUALS, STRICT_NOT_DISTINCT_FROM + } + @Nullable private final BlockBuilderStatus blockBuilderStatus; @@ -48,7 +53,7 @@ public class MapBlockBuilder private final MapHashTables hashTables; private boolean currentEntryOpened; - private boolean strict; + private HashBuildMode hashBuildMode = HashBuildMode.DUPLICATE_NOT_CHECKED; public MapBlockBuilder(MapType mapType, BlockBuilderStatus blockBuilderStatus, int expectedEntries) { @@ -86,7 +91,13 @@ private MapBlockBuilder( public MapBlockBuilder strict() { - this.strict = true; + this.hashBuildMode = HashBuildMode.STRICT_EQUALS; + return this; + } + + public MapBlockBuilder strictNotDistinctFrom() + { + this.hashBuildMode = HashBuildMode.STRICT_NOT_DISTINCT_FROM; return this; } @@ -173,23 +184,15 @@ public void retainedBytesForEachPart(ObjLongConsumer consumer) consumer.accept(this, INSTANCE_SIZE); } - @Override - public SingleMapBlockWriter beginBlockEntry() + public void buildEntry(MapValueBuilder builder) + throws E { if (currentEntryOpened) { throw new IllegalStateException("Expected current entry to be closed but was opened"); } - currentEntryOpened = true; - return new SingleMapBlockWriter(keyBlockBuilder.getPositionCount() * 2, keyBlockBuilder, valueBlockBuilder, this::strict); - } - - @Override - public BlockBuilder closeEntry() - { - if (!currentEntryOpened) { - throw new IllegalStateException("Expected entry to be opened but was closed"); - } + currentEntryOpened = true; + builder.build(keyBlockBuilder, valueBlockBuilder); entryAdded(false); currentEntryOpened = false; @@ -197,40 +200,11 @@ public BlockBuilder closeEntry() int previousAggregatedEntryCount = offsets[positionCount - 1]; int aggregatedEntryCount = offsets[positionCount]; int entryCount = aggregatedEntryCount - previousAggregatedEntryCount; - if (strict) { - hashTables.buildHashTableStrict(keyBlockBuilder, previousAggregatedEntryCount, entryCount); + switch (hashBuildMode) { + case DUPLICATE_NOT_CHECKED -> hashTables.buildHashTable(keyBlockBuilder, previousAggregatedEntryCount, entryCount); + case STRICT_EQUALS -> hashTables.buildHashTableStrict(keyBlockBuilder, previousAggregatedEntryCount, entryCount); + case STRICT_NOT_DISTINCT_FROM -> hashTables.buildDistinctHashTableStrict(keyBlockBuilder, previousAggregatedEntryCount, entryCount); } - else { - hashTables.buildHashTable(keyBlockBuilder, previousAggregatedEntryCount, entryCount); - } - return this; - } - - /** - * This method will check duplicate keys and close entry. - *

- * When duplicate keys are discovered, the block is guaranteed to be in - * a consistent state before {@link DuplicateMapKeyException} is thrown. - * In other words, one can continue to use this BlockBuilder. - * - * @deprecated use strict method instead - */ - @Deprecated - public void closeEntryStrict() - throws DuplicateMapKeyException - { - if (!currentEntryOpened) { - throw new IllegalStateException("Expected entry to be opened but was closed"); - } - - entryAdded(false); - currentEntryOpened = false; - - ensureHashTableSize(); - int previousAggregatedEntryCount = offsets[positionCount - 1]; - int aggregatedEntryCount = offsets[positionCount]; - int entryCount = aggregatedEntryCount - previousAggregatedEntryCount; - hashTables.buildHashTableStrict(keyBlockBuilder, previousAggregatedEntryCount, entryCount); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/MapHashTables.java b/core/trino-spi/src/main/java/io/trino/spi/block/MapHashTables.java index 35eed1eee353..3f1bd8778356 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/MapHashTables.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/MapHashTables.java @@ -202,6 +202,57 @@ synchronized void buildHashTableStrict(Block keyBlock, int keyOffset, int keyCou this.hashTables = hashTables; } + /** + * This method checks whether {@code keyBlock} has duplicates based on type NOT DISTINCT FROM. + */ + synchronized void buildDistinctHashTableStrict(Block keyBlock, int keyOffset, int keyCount) + throws DuplicateMapKeyException + { + int[] hashTables = this.hashTables; + if (hashTables == null) { + throw new IllegalStateException("hashTables not set"); + } + + int hashTableOffset = keyOffset * HASH_MULTIPLIER; + int hashTableSize = keyCount * HASH_MULTIPLIER; + + for (int i = 0; i < keyCount; i++) { + int hash = getHashPosition(keyBlock, keyOffset + i, hashTableSize); + while (true) { + if (hashTables[hashTableOffset + hash] == -1) { + hashTables[hashTableOffset + hash] = i; + break; + } + + if (keyBlock.isNull(keyOffset + i)) { + throw new TrinoException(NOT_SUPPORTED, "map key cannot be null or contain nulls"); + } + + boolean isDuplicateKey; + try { + // assuming maps with indeterminate keys are not supported + isDuplicateKey = (boolean) mapType.getKeyBlockNotDistinctFrom().invokeExact(keyBlock, keyOffset + i, keyBlock, keyOffset + hashTables[hashTableOffset + hash]); + } + catch (RuntimeException e) { + throw e; + } + catch (Throwable throwable) { + throw new RuntimeException(throwable); + } + + if (isDuplicateKey) { + throw new DuplicateMapKeyException(keyBlock, keyOffset + i); + } + + hash++; + if (hash == hashTableSize) { + hash = 0; + } + } + } + this.hashTables = hashTables; + } + private int getHashPosition(Block keyBlock, int position, int hashTableSize) { if (keyBlock.isNull(position)) { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/MapValueBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/MapValueBuilder.java new file mode 100644 index 000000000000..2c2642c43306 --- /dev/null +++ b/core/trino-spi/src/main/java/io/trino/spi/block/MapValueBuilder.java @@ -0,0 +1,30 @@ +/* + * Licensed 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.trino.spi.block; + +import io.trino.spi.block.BufferedMapValueBuilder.HashBuildMode; +import io.trino.spi.type.MapType; + +public interface MapValueBuilder +{ + static Block buildMapValue(MapType mapType, int entryCount, MapValueBuilder builder) + throws E + { + return new BufferedMapValueBuilder(mapType, HashBuildMode.DUPLICATE_NOT_CHECKED, entryCount) + .build(entryCount, builder); + } + + void build(BlockBuilder keyBuilder, BlockBuilder valueBuilder) + throws E; +} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java index 437ae97ec28c..7cf82a1541b8 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/RowBlockBuilder.java @@ -43,7 +43,7 @@ public class RowBlockBuilder private int[] fieldBlockOffsets; private boolean[] rowIsNull; private final BlockBuilder[] fieldBlockBuilders; - private final SingleRowBlockWriter singleRowBlockWriter; + private final List fieldBlockBuildersList; private boolean currentEntryOpened; private boolean hasNullRow; @@ -67,7 +67,7 @@ private RowBlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, BlockBu this.fieldBlockOffsets = requireNonNull(fieldBlockOffsets, "fieldBlockOffsets is null"); this.rowIsNull = requireNonNull(rowIsNull, "rowIsNull is null"); this.fieldBlockBuilders = requireNonNull(fieldBlockBuilders, "fieldBlockBuilders is null"); - this.singleRowBlockWriter = new SingleRowBlockWriter(fieldBlockBuilders); + this.fieldBlockBuildersList = List.of(fieldBlockBuilders); } private static BlockBuilder[] createFieldBlockBuilders(List fieldTypes, BlockBuilderStatus blockBuilderStatus, int expectedEntries) @@ -138,7 +138,6 @@ public long getRetainedSizeInBytes() if (blockBuilderStatus != null) { size += BlockBuilderStatus.INSTANCE_SIZE; } - size += SingleRowBlockWriter.INSTANCE_SIZE; return size; } @@ -153,28 +152,17 @@ public void retainedBytesForEachPart(ObjLongConsumer consumer) consumer.accept(this, INSTANCE_SIZE); } - @Override - public SingleRowBlockWriter beginBlockEntry() + public void buildEntry(RowValueBuilder builder) + throws E { if (currentEntryOpened) { throw new IllegalStateException("Expected current entry to be closed but was opened"); } - currentEntryOpened = true; - singleRowBlockWriter.setRowIndex(fieldBlockBuilders[0].getPositionCount()); - return singleRowBlockWriter; - } - - @Override - public BlockBuilder closeEntry() - { - if (!currentEntryOpened) { - throw new IllegalStateException("Expected entry to be opened but was closed"); - } + currentEntryOpened = true; + builder.build(fieldBlockBuildersList); entryAdded(false); currentEntryOpened = false; - singleRowBlockWriter.reset(); - return this; } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/RowValueBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/RowValueBuilder.java new file mode 100644 index 000000000000..c4f02476004f --- /dev/null +++ b/core/trino-spi/src/main/java/io/trino/spi/block/RowValueBuilder.java @@ -0,0 +1,31 @@ +/* + * Licensed 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.trino.spi.block; + +import io.trino.spi.type.RowType; + +import java.util.List; + +public interface RowValueBuilder +{ + static Block buildRowValue(RowType rowType, RowValueBuilder builder) + throws E + { + return new BufferedRowValueBuilder(rowType, 1) + .build(builder); + } + + void build(List fieldBuilders) + throws E; +} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/RunLengthEncodedBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/RunLengthEncodedBlock.java index eb23b72364a6..5da98c60a087 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/RunLengthEncodedBlock.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/RunLengthEncodedBlock.java @@ -14,6 +14,7 @@ package io.trino.spi.block; import io.airlift.slice.Slice; +import io.airlift.slice.SliceOutput; import io.trino.spi.predicate.Utils; import io.trino.spi.type.Type; @@ -247,6 +248,13 @@ public Slice getSlice(int position, int offset, int length) return value.getSlice(0, offset, length); } + @Override + public void writeSliceTo(int position, int offset, int length, SliceOutput output) + { + checkReadablePosition(this, position); + value.writeSliceTo(0, offset, length, output); + } + @Override public T getObject(int position, Class clazz) { @@ -268,13 +276,6 @@ public int bytesCompare(int position, int offset, int length, Slice otherSlice, return value.bytesCompare(0, offset, length, otherSlice, otherOffset, otherLength); } - @Override - public void writeBytesTo(int position, int offset, int length, BlockBuilder blockBuilder) - { - checkReadablePosition(this, position); - value.writeBytesTo(0, offset, length, blockBuilder); - } - @Override public boolean equals(int position, int offset, Block otherBlock, int otherPosition, int otherOffset, int length) { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/ShortArrayBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/ShortArrayBlockBuilder.java index a6dc1f4a4240..f99193068382 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/ShortArrayBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/ShortArrayBlockBuilder.java @@ -58,14 +58,13 @@ public ShortArrayBlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, i updateDataSize(); } - @Override - public BlockBuilder writeShort(int value) + public BlockBuilder writeShort(short value) { if (values.length <= positionCount) { growCapacity(); } - values[positionCount] = (short) value; + values[positionCount] = value; hasNonNullValue = true; positionCount++; @@ -75,12 +74,6 @@ public BlockBuilder writeShort(int value) return this; } - @Override - public BlockBuilder closeEntry() - { - return this; - } - @Override public BlockBuilder appendNull() { diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/SingleArrayBlockWriter.java b/core/trino-spi/src/main/java/io/trino/spi/block/SingleArrayBlockWriter.java deleted file mode 100644 index 65219571b509..000000000000 --- a/core/trino-spi/src/main/java/io/trino/spi/block/SingleArrayBlockWriter.java +++ /dev/null @@ -1,157 +0,0 @@ -/* - * Licensed 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.trino.spi.block; - -import io.airlift.slice.Slice; - -import java.util.OptionalInt; -import java.util.function.ObjLongConsumer; - -import static io.airlift.slice.SizeOf.instanceSize; -import static java.lang.String.format; - -public class SingleArrayBlockWriter - extends AbstractSingleArrayBlock - implements BlockBuilder -{ - private static final int INSTANCE_SIZE = instanceSize(SingleArrayBlockWriter.class); - - private final BlockBuilder blockBuilder; - private final long initialBlockBuilderSize; - private int positionsWritten; - - public SingleArrayBlockWriter(BlockBuilder blockBuilder, int start) - { - super(start); - this.blockBuilder = blockBuilder; - this.initialBlockBuilderSize = blockBuilder.getSizeInBytes(); - } - - @Override - protected Block getBlock() - { - return blockBuilder; - } - - @Override - public OptionalInt fixedSizeInBytesPerPosition() - { - return OptionalInt.empty(); - } - - @Override - public long getSizeInBytes() - { - return blockBuilder.getSizeInBytes() - initialBlockBuilderSize; - } - - @Override - public long getRetainedSizeInBytes() - { - return INSTANCE_SIZE + blockBuilder.getRetainedSizeInBytes(); - } - - @Override - public void retainedBytesForEachPart(ObjLongConsumer consumer) - { - consumer.accept(blockBuilder, blockBuilder.getRetainedSizeInBytes()); - consumer.accept(this, INSTANCE_SIZE); - } - - @Override - public BlockBuilder writeByte(int value) - { - blockBuilder.writeByte(value); - return this; - } - - @Override - public BlockBuilder writeShort(int value) - { - blockBuilder.writeShort(value); - return this; - } - - @Override - public BlockBuilder writeInt(int value) - { - blockBuilder.writeInt(value); - return this; - } - - @Override - public BlockBuilder writeLong(long value) - { - blockBuilder.writeLong(value); - return this; - } - - @Override - public BlockBuilder writeBytes(Slice source, int sourceIndex, int length) - { - blockBuilder.writeBytes(source, sourceIndex, length); - return this; - } - - @Override - public BlockBuilder beginBlockEntry() - { - return blockBuilder.beginBlockEntry(); - } - - @Override - public BlockBuilder appendNull() - { - blockBuilder.appendNull(); - entryAdded(); - return this; - } - - @Override - public BlockBuilder closeEntry() - { - blockBuilder.closeEntry(); - entryAdded(); - return this; - } - - private void entryAdded() - { - positionsWritten++; - } - - @Override - public int getPositionCount() - { - return positionsWritten; - } - - @Override - public Block build() - { - throw new UnsupportedOperationException(); - } - - @Override - public BlockBuilder newBlockBuilderLike(int expectedEntries, BlockBuilderStatus blockBuilderStatus) - { - throw new UnsupportedOperationException(); - } - - @Override - public String toString() - { - return format("SingleArrayBlockWriter{positionCount=%d}", getPositionCount()); - } -} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/SingleMapBlockWriter.java b/core/trino-spi/src/main/java/io/trino/spi/block/SingleMapBlockWriter.java deleted file mode 100644 index f90cea18374d..000000000000 --- a/core/trino-spi/src/main/java/io/trino/spi/block/SingleMapBlockWriter.java +++ /dev/null @@ -1,232 +0,0 @@ -/* - * Licensed 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.trino.spi.block; - -import io.airlift.slice.Slice; - -import java.util.OptionalInt; -import java.util.function.ObjLongConsumer; - -import static io.airlift.slice.SizeOf.instanceSize; -import static java.lang.String.format; - -public class SingleMapBlockWriter - extends AbstractSingleMapBlock - implements BlockBuilder -{ - private static final int INSTANCE_SIZE = instanceSize(SingleMapBlockWriter.class); - - private final int offset; - private final BlockBuilder keyBlockBuilder; - private final BlockBuilder valueBlockBuilder; - private final Runnable setStrict; - private final long initialBlockBuilderSize; - private int positionsWritten; - - private boolean writeToValueNext; - - SingleMapBlockWriter(int start, BlockBuilder keyBlockBuilder, BlockBuilder valueBlockBuilder, Runnable setStrict) - { - this.offset = start; - this.keyBlockBuilder = keyBlockBuilder; - this.valueBlockBuilder = valueBlockBuilder; - this.setStrict = setStrict; - this.initialBlockBuilderSize = keyBlockBuilder.getSizeInBytes() + valueBlockBuilder.getSizeInBytes(); - } - - public SingleMapBlockWriter strict() - { - setStrict.run(); - return this; - } - - @Override - int getOffset() - { - return offset; - } - - @Override - Block getRawKeyBlock() - { - return keyBlockBuilder; - } - - @Override - Block getRawValueBlock() - { - return valueBlockBuilder; - } - - @Override - public OptionalInt fixedSizeInBytesPerPosition() - { - return OptionalInt.empty(); - } - - @Override - public long getSizeInBytes() - { - return keyBlockBuilder.getSizeInBytes() + valueBlockBuilder.getSizeInBytes() - initialBlockBuilderSize; - } - - @Override - public long getRetainedSizeInBytes() - { - return INSTANCE_SIZE + keyBlockBuilder.getRetainedSizeInBytes() + valueBlockBuilder.getRetainedSizeInBytes(); - } - - @Override - public void retainedBytesForEachPart(ObjLongConsumer consumer) - { - consumer.accept(keyBlockBuilder, keyBlockBuilder.getRetainedSizeInBytes()); - consumer.accept(valueBlockBuilder, valueBlockBuilder.getRetainedSizeInBytes()); - consumer.accept(this, INSTANCE_SIZE); - } - - @Override - public BlockBuilder writeByte(int value) - { - if (writeToValueNext) { - valueBlockBuilder.writeByte(value); - } - else { - keyBlockBuilder.writeByte(value); - } - return this; - } - - @Override - public BlockBuilder writeShort(int value) - { - if (writeToValueNext) { - valueBlockBuilder.writeShort(value); - } - else { - keyBlockBuilder.writeShort(value); - } - return this; - } - - @Override - public BlockBuilder writeInt(int value) - { - if (writeToValueNext) { - valueBlockBuilder.writeInt(value); - } - else { - keyBlockBuilder.writeInt(value); - } - return this; - } - - @Override - public BlockBuilder writeLong(long value) - { - if (writeToValueNext) { - valueBlockBuilder.writeLong(value); - } - else { - keyBlockBuilder.writeLong(value); - } - return this; - } - - @Override - public BlockBuilder writeBytes(Slice source, int sourceIndex, int length) - { - if (writeToValueNext) { - valueBlockBuilder.writeBytes(source, sourceIndex, length); - } - else { - keyBlockBuilder.writeBytes(source, sourceIndex, length); - } - return this; - } - - @Override - public BlockBuilder beginBlockEntry() - { - BlockBuilder result; - if (writeToValueNext) { - result = valueBlockBuilder.beginBlockEntry(); - } - else { - result = keyBlockBuilder.beginBlockEntry(); - } - return result; - } - - @Override - public BlockBuilder appendNull() - { - if (writeToValueNext) { - valueBlockBuilder.appendNull(); - } - else { - keyBlockBuilder.appendNull(); - } - entryAdded(); - return this; - } - - @Override - public BlockBuilder closeEntry() - { - if (writeToValueNext) { - valueBlockBuilder.closeEntry(); - } - else { - keyBlockBuilder.closeEntry(); - } - entryAdded(); - return this; - } - - private void entryAdded() - { - writeToValueNext = !writeToValueNext; - positionsWritten++; - } - - @Override - public int getPositionCount() - { - return positionsWritten; - } - - @Override - public String getEncodingName() - { - throw new UnsupportedOperationException(); - } - - @Override - public Block build() - { - throw new UnsupportedOperationException(); - } - - @Override - public BlockBuilder newBlockBuilderLike(int expectedEntries, BlockBuilderStatus blockBuilderStatus) - { - throw new UnsupportedOperationException(); - } - - @Override - public String toString() - { - return format("SingleMapBlockWriter{positionCount=%d}", getPositionCount()); - } -} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/SingleRowBlockWriter.java b/core/trino-spi/src/main/java/io/trino/spi/block/SingleRowBlockWriter.java deleted file mode 100644 index 4cd96950712a..000000000000 --- a/core/trino-spi/src/main/java/io/trino/spi/block/SingleRowBlockWriter.java +++ /dev/null @@ -1,249 +0,0 @@ -/* - * Licensed 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.trino.spi.block; - -import io.airlift.slice.Slice; - -import java.util.OptionalInt; -import java.util.function.ObjLongConsumer; - -import static io.airlift.slice.SizeOf.instanceSize; -import static java.lang.String.format; - -public class SingleRowBlockWriter - extends AbstractSingleRowBlock - implements BlockBuilder -{ - public static final int INSTANCE_SIZE = instanceSize(SingleRowBlockWriter.class); - - private final BlockBuilder[] fieldBlockBuilders; - - private int currentFieldIndexToWrite; - private int rowIndex = -1; - private boolean fieldBlockBuilderReturned; - - SingleRowBlockWriter(BlockBuilder[] fieldBlockBuilders) - { - this.fieldBlockBuilders = fieldBlockBuilders; - } - - /** - * Obtains the field {@code BlockBuilder}. - *

- * This method is used to perform random write to {@code SingleRowBlockWriter}. - * Each field {@code BlockBuilder} must be written EXACTLY once. - *

- * Field {@code BlockBuilder} can only be obtained before any sequential write has done. - * Once obtained, sequential write is no longer allowed. - */ - public BlockBuilder getFieldBlockBuilder(int fieldIndex) - { - if (currentFieldIndexToWrite != 0) { - throw new IllegalStateException("field block builder can only be obtained before any sequential write has done"); - } - fieldBlockBuilderReturned = true; - return fieldBlockBuilders[fieldIndex]; - } - - @Override - Block[] getRawFieldBlocks() - { - return fieldBlockBuilders; - } - - @Override - protected Block getRawFieldBlock(int fieldIndex) - { - return fieldBlockBuilders[fieldIndex]; - } - - @Override - protected int getRowIndex() - { - return rowIndex; - } - - @Override - public OptionalInt fixedSizeInBytesPerPosition() - { - return OptionalInt.empty(); - } - - @Override - public long getSizeInBytes() - { - long currentBlockBuilderSize = 0; - /* - * We need to use subtraction in order to compute size because getRegionSizeInBytes(0, 1) - * returns non-zero result even if field block builder has no position appended - */ - for (BlockBuilder fieldBlockBuilder : fieldBlockBuilders) { - currentBlockBuilderSize += fieldBlockBuilder.getSizeInBytes() - fieldBlockBuilder.getRegionSizeInBytes(0, rowIndex); - } - return currentBlockBuilderSize; - } - - @Override - public long getRetainedSizeInBytes() - { - long size = INSTANCE_SIZE; - for (BlockBuilder fieldBlockBuilder : fieldBlockBuilders) { - size += fieldBlockBuilder.getRetainedSizeInBytes(); - } - return size; - } - - @Override - public void retainedBytesForEachPart(ObjLongConsumer consumer) - { - for (BlockBuilder fieldBlockBuilder : fieldBlockBuilders) { - consumer.accept(fieldBlockBuilder, fieldBlockBuilder.getRetainedSizeInBytes()); - } - consumer.accept(this, INSTANCE_SIZE); - } - - @Override - public BlockBuilder writeByte(int value) - { - checkFieldIndexToWrite(); - fieldBlockBuilders[currentFieldIndexToWrite].writeByte(value); - return this; - } - - @Override - public BlockBuilder writeShort(int value) - { - checkFieldIndexToWrite(); - fieldBlockBuilders[currentFieldIndexToWrite].writeShort(value); - return this; - } - - @Override - public BlockBuilder writeInt(int value) - { - checkFieldIndexToWrite(); - fieldBlockBuilders[currentFieldIndexToWrite].writeInt(value); - return this; - } - - @Override - public BlockBuilder writeLong(long value) - { - checkFieldIndexToWrite(); - fieldBlockBuilders[currentFieldIndexToWrite].writeLong(value); - return this; - } - - @Override - public BlockBuilder writeBytes(Slice source, int sourceIndex, int length) - { - checkFieldIndexToWrite(); - fieldBlockBuilders[currentFieldIndexToWrite].writeBytes(source, sourceIndex, length); - return this; - } - - @Override - public BlockBuilder beginBlockEntry() - { - checkFieldIndexToWrite(); - return fieldBlockBuilders[currentFieldIndexToWrite].beginBlockEntry(); - } - - @Override - public BlockBuilder appendNull() - { - checkFieldIndexToWrite(); - fieldBlockBuilders[currentFieldIndexToWrite].appendNull(); - entryAdded(); - return this; - } - - @Override - public BlockBuilder closeEntry() - { - checkFieldIndexToWrite(); - fieldBlockBuilders[currentFieldIndexToWrite].closeEntry(); - entryAdded(); - return this; - } - - private void entryAdded() - { - currentFieldIndexToWrite++; - } - - @Override - public int getPositionCount() - { - if (fieldBlockBuilderReturned) { - throw new IllegalStateException("field block builder has been returned"); - } - return currentFieldIndexToWrite; - } - - @Override - public String getEncodingName() - { - throw new UnsupportedOperationException(); - } - - @Override - public Block build() - { - throw new UnsupportedOperationException(); - } - - @Override - public BlockBuilder newBlockBuilderLike(int expectedEntries, BlockBuilderStatus blockBuilderStatus) - { - throw new UnsupportedOperationException(); - } - - @Override - public String toString() - { - if (!fieldBlockBuilderReturned) { - return format("SingleRowBlockWriter{numFields=%d, fieldBlockBuilderReturned=false, positionCount=%d}", fieldBlockBuilders.length, getPositionCount()); - } - return format("SingleRowBlockWriter{numFields=%d, fieldBlockBuilderReturned=true}", fieldBlockBuilders.length); - } - - void setRowIndex(int rowIndex) - { - if (this.rowIndex != -1) { - throw new IllegalStateException("SingleRowBlockWriter should be reset before usage"); - } - this.rowIndex = rowIndex; - } - - void reset() - { - if (this.rowIndex == -1) { - throw new IllegalStateException("SingleRowBlockWriter is already reset"); - } - this.rowIndex = -1; - this.currentFieldIndexToWrite = 0; - this.fieldBlockBuilderReturned = false; - } - - private void checkFieldIndexToWrite() - { - if (fieldBlockBuilderReturned) { - throw new IllegalStateException("cannot do sequential write after getFieldBlockBuilder is called"); - } - if (currentFieldIndexToWrite >= fieldBlockBuilders.length) { - throw new IllegalStateException("currentFieldIndexToWrite is not valid"); - } - } -} diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/VariableWidthBlockBuilder.java b/core/trino-spi/src/main/java/io/trino/spi/block/VariableWidthBlockBuilder.java index eff098b1ef43..7d53dd3027ee 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/block/VariableWidthBlockBuilder.java +++ b/core/trino-spi/src/main/java/io/trino/spi/block/VariableWidthBlockBuilder.java @@ -26,8 +26,6 @@ import static io.airlift.slice.SizeOf.SIZE_OF_BYTE; import static io.airlift.slice.SizeOf.SIZE_OF_INT; -import static io.airlift.slice.SizeOf.SIZE_OF_LONG; -import static io.airlift.slice.SizeOf.SIZE_OF_SHORT; import static io.airlift.slice.SizeOf.instanceSize; import static io.airlift.slice.SizeOf.sizeOf; import static io.airlift.slice.Slices.EMPTY_SLICE; @@ -64,7 +62,6 @@ public class VariableWidthBlockBuilder private int[] offsets = new int[1]; private int positions; - private int currentEntrySize; private long arraysRetainedSizeInBytes; @@ -191,76 +188,44 @@ public Block copyPositions(int[] positions, int offset, int length) return new VariableWidthBlock(0, length, newSlice.slice(), newOffsets, newValueIsNull); } - @Override - public BlockBuilder writeByte(int value) + public void writeEntry(Slice source) { - if (!initialized) { - initializeCapacity(); - } - sliceOutput.writeByte(value); - currentEntrySize += SIZE_OF_BYTE; - return this; + writeEntry(source, 0, source.length()); } - @Override - public BlockBuilder writeShort(int value) + public VariableWidthBlockBuilder writeEntry(Slice source, int sourceIndex, int length) { if (!initialized) { initializeCapacity(); } - sliceOutput.writeShort(value); - currentEntrySize += SIZE_OF_SHORT; - return this; - } - @Override - public BlockBuilder writeInt(int value) - { - if (!initialized) { - initializeCapacity(); - } - sliceOutput.writeInt(value); - currentEntrySize += SIZE_OF_INT; + sliceOutput.writeBytes(source, sourceIndex, length); + entryAdded(length, false); return this; } - @Override - public BlockBuilder writeLong(long value) + public void buildEntry(VariableWidthEntryBuilder builder) + throws E { if (!initialized) { initializeCapacity(); } - sliceOutput.writeLong(value); - currentEntrySize += SIZE_OF_LONG; - return this; - } - @Override - public BlockBuilder writeBytes(Slice source, int sourceIndex, int length) - { - if (!initialized) { - initializeCapacity(); - } - sliceOutput.writeBytes(source, sourceIndex, length); - currentEntrySize += length; - return this; + int start = sliceOutput.size(); + builder.build(sliceOutput); + int length = sliceOutput.size() - start; + entryAdded(length, false); } - @Override - public BlockBuilder closeEntry() + public interface VariableWidthEntryBuilder { - entryAdded(currentEntrySize, false); - currentEntrySize = 0; - return this; + void build(SliceOutput output) + throws E; } @Override public BlockBuilder appendNull() { - if (currentEntrySize > 0) { - throw new IllegalStateException("Current entry must be closed before a null can be written"); - } - hasNullValue = true; entryAdded(0, true); return this; @@ -295,7 +260,7 @@ private void growCapacity() private void initializeCapacity() { - if (positions != 0 || currentEntrySize != 0) { + if (positions != 0) { throw new IllegalStateException(getClass().getSimpleName() + " was used before initialization"); } initialized = true; @@ -357,9 +322,6 @@ public Block copyRegion(int positionOffset, int length) @Override public Block build() { - if (currentEntrySize > 0) { - throw new IllegalStateException("Current entry must be closed before the block can be built"); - } if (!hasNonNullValue) { return RunLengthEncodedBlock.create(NULL_VALUE_BLOCK, positions); } diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/AbstractIntType.java b/core/trino-spi/src/main/java/io/trino/spi/type/AbstractIntType.java index ebab8fe88256..34c53dad9cc7 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/AbstractIntType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/AbstractIntType.java @@ -90,7 +90,12 @@ public final Slice getSlice(Block block, int position) public void writeLong(BlockBuilder blockBuilder, long value) { checkValueValid(value); - blockBuilder.writeInt((int) value).closeEntry(); + writeInt(blockBuilder, (int) value); + } + + public BlockBuilder writeInt(BlockBuilder blockBuilder, int value) + { + return ((IntArrayBlockBuilder) blockBuilder).writeInt(value); } protected void checkValueValid(long value) @@ -110,7 +115,7 @@ public final void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeInt(block.getInt(position, 0)).closeEntry(); + writeInt(blockBuilder, block.getInt(position, 0)); } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/AbstractLongType.java b/core/trino-spi/src/main/java/io/trino/spi/type/AbstractLongType.java index 7ae2d0859e99..f335b3fe0e8e 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/AbstractLongType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/AbstractLongType.java @@ -82,7 +82,7 @@ public final Slice getSlice(Block block, int position) @Override public final void writeLong(BlockBuilder blockBuilder, long value) { - blockBuilder.writeLong(value).closeEntry(); + ((LongArrayBlockBuilder) blockBuilder).writeLong(value); } @Override @@ -92,7 +92,7 @@ public final void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(block.getLong(position, 0)).closeEntry(); + writeLong(blockBuilder, getLong(block, position)); } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/AbstractVariableWidthType.java b/core/trino-spi/src/main/java/io/trino/spi/type/AbstractVariableWidthType.java index 45f4a71c4d75..5402862f4cf3 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/AbstractVariableWidthType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/AbstractVariableWidthType.java @@ -13,7 +13,6 @@ */ package io.trino.spi.type; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.BlockBuilderStatus; import io.trino.spi.block.PageBuilderStatus; import io.trino.spi.block.VariableWidthBlockBuilder; @@ -32,7 +31,7 @@ protected AbstractVariableWidthType(TypeSignature signature, Class javaType) } @Override - public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry) + public VariableWidthBlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry) { int maxBlockSizeInBytes; if (blockBuilderStatus == null) { @@ -51,7 +50,7 @@ public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, in } @Override - public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) + public VariableWidthBlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) { return createBlockBuilder(blockBuilderStatus, expectedEntries, EXPECTED_BYTES_PER_ENTRY); } diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/ArrayType.java b/core/trino-spi/src/main/java/io/trino/spi/type/ArrayType.java index ac29e3095cd0..63e07721e6dd 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/ArrayType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/ArrayType.java @@ -13,7 +13,6 @@ */ package io.trino.spi.type; -import io.airlift.slice.Slice; import io.trino.spi.block.AbstractArrayBlock; import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; @@ -221,24 +220,6 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) } } - @Override - public Slice getSlice(Block block, int position) - { - return block.getSlice(position, 0, block.getSliceLength(position)); - } - - @Override - public void writeSlice(BlockBuilder blockBuilder, Slice value) - { - writeSlice(blockBuilder, value, 0, value.length()); - } - - @Override - public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length) - { - blockBuilder.writeBytes(value, offset, length).closeEntry(); - } - @Override public Block getObject(Block block, int position) { @@ -249,22 +230,21 @@ public Block getObject(Block block, int position) public void writeObject(BlockBuilder blockBuilder, Object value) { Block arrayBlock = (Block) value; - - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (int i = 0; i < arrayBlock.getPositionCount(); i++) { - elementType.appendTo(arrayBlock, i, entryBuilder); - } - blockBuilder.closeEntry(); + ((ArrayBlockBuilder) blockBuilder).buildEntry(elementBuilder -> { + for (int i = 0; i < arrayBlock.getPositionCount(); i++) { + elementType.appendTo(arrayBlock, i, elementBuilder); + } + }); } @Override - public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry) + public ArrayBlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry) { return new ArrayBlockBuilder(elementType, blockBuilderStatus, expectedEntries, expectedBytesPerEntry); } @Override - public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) + public ArrayBlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) { return createBlockBuilder(blockBuilderStatus, expectedEntries, 100); } diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/BooleanType.java b/core/trino-spi/src/main/java/io/trino/spi/type/BooleanType.java index 226248866dec..258eec78a640 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/BooleanType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/BooleanType.java @@ -134,7 +134,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeByte(block.getByte(position, 0)).closeEntry(); + ((ByteArrayBlockBuilder) blockBuilder).writeByte(block.getByte(position, 0)); } } @@ -147,7 +147,7 @@ public boolean getBoolean(Block block, int position) @Override public void writeBoolean(BlockBuilder blockBuilder, boolean value) { - blockBuilder.writeByte(value ? 1 : 0).closeEntry(); + ((ByteArrayBlockBuilder) blockBuilder).writeByte((byte) (value ? 1 : 0)); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/CharType.java b/core/trino-spi/src/main/java/io/trino/spi/type/CharType.java index f2b77164305a..2f52d7862f3d 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/CharType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/CharType.java @@ -20,6 +20,7 @@ import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.BlockBuilderStatus; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.function.BlockIndex; import io.trino.spi.function.BlockPosition; @@ -173,7 +174,7 @@ public Object getObjectValue(ConnectorSession session, Block block, int position } @Override - public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) + public VariableWidthBlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) { // If bound on length of char is smaller than EXPECTED_BYTES_PER_ENTRY, use that as expectedBytesPerEntry // The data can take up to 4 bytes per character due to UTF-8 encoding, but we assume it is ASCII and only needs one byte. @@ -187,8 +188,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -215,7 +215,7 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int l if (length > 0 && value.getByte(offset + length - 1) == ' ') { throw new IllegalArgumentException("Slice representing Char should not have trailing spaces"); } - blockBuilder.writeBytes(value, offset, length).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/Decimals.java b/core/trino-spi/src/main/java/io/trino/spi/type/Decimals.java index 6db9e935281f..9ed0664a1e0c 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/Decimals.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/Decimals.java @@ -16,6 +16,7 @@ import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.LongArrayBlockBuilder; import java.math.BigDecimal; import java.math.BigInteger; @@ -253,7 +254,7 @@ public static BigDecimal rescale(BigDecimal value, DecimalType type) public static void writeShortDecimal(BlockBuilder blockBuilder, long value) { - blockBuilder.writeLong(value).closeEntry(); + ((LongArrayBlockBuilder) blockBuilder).writeLong(value); } public static long rescale(long value, int fromScale, int toScale) diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/DoubleType.java b/core/trino-spi/src/main/java/io/trino/spi/type/DoubleType.java index 5ecbd37bf60c..3b98ea0fead3 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/DoubleType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/DoubleType.java @@ -91,7 +91,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(block.getLong(position, 0)).closeEntry(); + ((LongArrayBlockBuilder) blockBuilder).writeLong(block.getLong(position, 0)); } } @@ -104,7 +104,7 @@ public double getDouble(Block block, int position) @Override public void writeDouble(BlockBuilder blockBuilder, double value) { - blockBuilder.writeLong(doubleToLongBits(value)).closeEntry(); + ((LongArrayBlockBuilder) blockBuilder).writeLong(doubleToLongBits(value)); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/HyperLogLogType.java b/core/trino-spi/src/main/java/io/trino/spi/type/HyperLogLogType.java index 26a9ae187b75..4c1bcf0ed33f 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/HyperLogLogType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/HyperLogLogType.java @@ -17,6 +17,7 @@ import io.airlift.slice.Slice; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; // Layout is :, where @@ -40,8 +41,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -60,7 +60,7 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value) @Override public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length) { - blockBuilder.writeBytes(value, offset, length).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/LongDecimalType.java b/core/trino-spi/src/main/java/io/trino/spi/type/LongDecimalType.java index 22e6028e7e81..47eac8a7c1cb 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/LongDecimalType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/LongDecimalType.java @@ -102,9 +102,9 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(block.getLong(position, 0)); - blockBuilder.writeLong(block.getLong(position, SIZE_OF_LONG)); - blockBuilder.closeEntry(); + ((Int128ArrayBlockBuilder) blockBuilder).writeInt128( + block.getLong(position, 0), + block.getLong(position, SIZE_OF_LONG)); } } @@ -112,9 +112,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) public void writeObject(BlockBuilder blockBuilder, Object value) { Int128 decimal = (Int128) value; - blockBuilder.writeLong(decimal.getHigh()); - blockBuilder.writeLong(decimal.getLow()); - blockBuilder.closeEntry(); + ((Int128ArrayBlockBuilder) blockBuilder).writeInt128(decimal.getHigh(), decimal.getLow()); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/LongTimeWithTimeZoneType.java b/core/trino-spi/src/main/java/io/trino/spi/type/LongTimeWithTimeZoneType.java index 5ba9eddbd410..b02cfed4546c 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/LongTimeWithTimeZoneType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/LongTimeWithTimeZoneType.java @@ -96,9 +96,9 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(getPicos(block, position)); - blockBuilder.writeInt(getOffsetMinutes(block, position)); - blockBuilder.closeEntry(); + ((Fixed12BlockBuilder) blockBuilder).writeFixed12( + getPicos(block, position), + getOffsetMinutes(block, position)); } } @@ -112,9 +112,9 @@ public Object getObject(Block block, int position) public void writeObject(BlockBuilder blockBuilder, Object value) { LongTimeWithTimeZone timestamp = (LongTimeWithTimeZone) value; - blockBuilder.writeLong(timestamp.getPicoseconds()); - blockBuilder.writeInt(timestamp.getOffsetMinutes()); - blockBuilder.closeEntry(); + ((Fixed12BlockBuilder) blockBuilder).writeFixed12( + timestamp.getPicoseconds(), + timestamp.getOffsetMinutes()); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/LongTimestampType.java b/core/trino-spi/src/main/java/io/trino/spi/type/LongTimestampType.java index 8e4fd631b32b..f33f5028ceaa 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/LongTimestampType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/LongTimestampType.java @@ -109,9 +109,9 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(getEpochMicros(block, position)); - blockBuilder.writeInt(getFraction(block, position)); - blockBuilder.closeEntry(); + ((Fixed12BlockBuilder) blockBuilder).writeFixed12( + getEpochMicros(block, position), + getFraction(block, position)); } } @@ -130,9 +130,7 @@ public void writeObject(BlockBuilder blockBuilder, Object value) public void write(BlockBuilder blockBuilder, long epochMicros, int fraction) { - blockBuilder.writeLong(epochMicros); - blockBuilder.writeInt(fraction); - blockBuilder.closeEntry(); + ((Fixed12BlockBuilder) blockBuilder).writeFixed12(epochMicros, fraction); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/LongTimestampWithTimeZoneType.java b/core/trino-spi/src/main/java/io/trino/spi/type/LongTimestampWithTimeZoneType.java index 56f2b497b39a..101786dd3a46 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/LongTimestampWithTimeZoneType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/LongTimestampWithTimeZoneType.java @@ -107,9 +107,9 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(getPackedEpochMillis(block, position)); - blockBuilder.writeInt(getPicosOfMilli(block, position)); - blockBuilder.closeEntry(); + ((Fixed12BlockBuilder) blockBuilder).writeFixed12( + getPackedEpochMillis(block, position), + getPicosOfMilli(block, position)); } } @@ -127,9 +127,9 @@ public void writeObject(BlockBuilder blockBuilder, Object value) { LongTimestampWithTimeZone timestamp = (LongTimestampWithTimeZone) value; - blockBuilder.writeLong(packDateTimeWithZone(timestamp.getEpochMillis(), timestamp.getTimeZoneKey())); - blockBuilder.writeInt(timestamp.getPicosOfMilli()); - blockBuilder.closeEntry(); + ((Fixed12BlockBuilder) blockBuilder).writeFixed12( + packDateTimeWithZone(timestamp.getEpochMillis(), timestamp.getTimeZoneKey()), + timestamp.getPicosOfMilli()); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/MapType.java b/core/trino-spi/src/main/java/io/trino/spi/type/MapType.java index 99585bf89388..6025a03dbd12 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/MapType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/MapType.java @@ -42,12 +42,14 @@ import static io.trino.spi.type.TypeOperatorDeclaration.NO_TYPE_OPERATOR_DECLARATION; import static io.trino.spi.type.TypeUtils.NULL_HASH_CODE; import static java.lang.String.format; +import static java.lang.invoke.MethodHandles.filterReturnValue; import static java.lang.invoke.MethodType.methodType; import static java.util.Arrays.asList; public class MapType extends AbstractType { + private static final MethodHandle NOT; private static final InvocationConvention EQUAL_CONVENTION = simpleConvention(NULLABLE_RETURN, NEVER_NULL, NEVER_NULL); private static final InvocationConvention HASH_CODE_CONVENTION = simpleConvention(FAIL_ON_NULL, NEVER_NULL); private static final InvocationConvention DISTINCT_FROM_CONVENTION = simpleConvention(FAIL_ON_NULL, BOXED_NULLABLE, BOXED_NULLABLE); @@ -63,6 +65,7 @@ public class MapType static { try { Lookup lookup = MethodHandles.lookup(); + NOT = lookup.findStatic(MapType.class, "not", methodType(boolean.class, boolean.class)); EQUAL = lookup.findStatic(MapType.class, "equalOperator", methodType(Boolean.class, MethodHandle.class, MethodHandle.class, Block.class, Block.class)); HASH_CODE = lookup.findStatic(MapType.class, "hashOperator", methodType(long.class, MethodHandle.class, MethodHandle.class, Block.class)); DISTINCT_FROM = lookup.findStatic(MapType.class, "distinctFromOperator", methodType(boolean.class, MethodHandle.class, MethodHandle.class, Block.class, Block.class)); @@ -81,6 +84,8 @@ public class MapType private final Type valueType; private static final int EXPECTED_BYTES_PER_ENTRY = 32; + private final MethodHandle keyBlockNativeNotDistinctFrom; + private final MethodHandle keyBlockNotDistinctFrom; private final MethodHandle keyNativeHashCode; private final MethodHandle keyBlockHashCode; private final MethodHandle keyBlockNativeEqual; @@ -107,6 +112,11 @@ public MapType(Type keyType, Type valueType, TypeOperators typeOperators) keyBlockNativeEqual = typeOperators.getEqualOperator(keyType, simpleConvention(NULLABLE_RETURN, BLOCK_POSITION, NEVER_NULL)) .asType(methodType(Boolean.class, Block.class, int.class, keyType.getJavaType().isPrimitive() ? keyType.getJavaType() : Object.class)); keyBlockEqual = typeOperators.getEqualOperator(keyType, simpleConvention(NULLABLE_RETURN, BLOCK_POSITION, BLOCK_POSITION)); + + keyBlockNativeNotDistinctFrom = filterReturnValue(typeOperators.getDistinctFromOperator(keyType, simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, NEVER_NULL)), NOT) + .asType(methodType(boolean.class, Block.class, int.class, keyType.getJavaType().isPrimitive() ? keyType.getJavaType() : Object.class)); + keyBlockNotDistinctFrom = filterReturnValue(typeOperators.getDistinctFromOperator(keyType, simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, BLOCK_POSITION)), NOT); + keyNativeHashCode = typeOperators.getHashCodeOperator(keyType, HASH_CODE_CONVENTION) .asType(methodType(long.class, keyType.getJavaType().isPrimitive() ? keyType.getJavaType() : Object.class)); keyBlockHashCode = typeOperators.getHashCodeOperator(keyType, simpleConvention(FAIL_ON_NULL, BLOCK_POSITION)); @@ -182,13 +192,13 @@ private static OperatorMethodHandle getIndeterminateOperatorInvoker(TypeOperator } @Override - public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry) + public MapBlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry) { return new MapBlockBuilder(this, blockBuilderStatus, expectedEntries); } @Override - public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) + public MapBlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) { return createBlockBuilder(blockBuilderStatus, expectedEntries, EXPECTED_BYTES_PER_ENTRY); } @@ -252,14 +262,12 @@ public void writeObject(BlockBuilder blockBuilder, Object value) throw new IllegalArgumentException("Maps must be represented with SingleMapBlock"); } - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - - for (int i = 0; i < singleMapBlock.getPositionCount(); i += 2) { - keyType.appendTo(singleMapBlock, i, entryBuilder); - valueType.appendTo(singleMapBlock, i + 1, entryBuilder); - } - - blockBuilder.closeEntry(); + ((MapBlockBuilder) blockBuilder).buildEntry((keyBuilder, valueBuilder) -> { + for (int i = 0; i < singleMapBlock.getPositionCount(); i += 2) { + keyType.appendTo(singleMapBlock, i, keyBuilder); + valueType.appendTo(singleMapBlock, i + 1, valueBuilder); + } + }); } @Override @@ -316,6 +324,22 @@ public MethodHandle getKeyBlockEqual() return keyBlockEqual; } + /** + * Internal use by this package and io.trino.spi.block only. + */ + public MethodHandle getKeyBlockNativeNotDistinctFrom() + { + return keyBlockNativeNotDistinctFrom; + } + + /** + * Internal use by this package and io.trino.spi.block only. + */ + public MethodHandle getKeyBlockNotDistinctFrom() + { + return keyBlockNotDistinctFrom; + } + private static long hashOperator(MethodHandle keyOperator, MethodHandle valueOperator, Block block) throws Throwable { @@ -430,4 +454,9 @@ private static boolean indeterminate(MethodHandle valueIndeterminateFunction, Bl } return false; } + + private static boolean not(boolean value) + { + return !value; + } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/QuantileDigestType.java b/core/trino-spi/src/main/java/io/trino/spi/type/QuantileDigestType.java index ac6a0d19dd7b..404ecf2e4f65 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/QuantileDigestType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/QuantileDigestType.java @@ -17,6 +17,7 @@ import io.airlift.slice.Slice; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import java.util.List; @@ -42,8 +43,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -62,7 +62,7 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value) @Override public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length) { - blockBuilder.writeBytes(value, offset, length).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/RealType.java b/core/trino-spi/src/main/java/io/trino/spi/type/RealType.java index ac07c3f51190..f301dc792f6d 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/RealType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/RealType.java @@ -81,7 +81,12 @@ public void writeLong(BlockBuilder blockBuilder, long value) catch (ArithmeticException e) { throw new TrinoException(GENERIC_INTERNAL_ERROR, format("Value (%sb) is not a valid single-precision float", Long.toBinaryString(value))); } - blockBuilder.writeInt(floatValue).closeEntry(); + writeInt(blockBuilder, floatValue); + } + + public void writeFloat(BlockBuilder blockBuilder, float value) + { + writeInt(blockBuilder, floatToIntBits(value)); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/RowType.java b/core/trino-spi/src/main/java/io/trino/spi/type/RowType.java index 2cbc920cb32d..1c3d5f1b62a9 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/RowType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/RowType.java @@ -178,13 +178,13 @@ private static TypeSignature makeSignature(List fields) } @Override - public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry) + public RowBlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry) { return new RowBlockBuilder(getTypeParameters(), blockBuilderStatus, expectedEntries); } @Override - public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) + public RowBlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) { return new RowBlockBuilder(getTypeParameters(), blockBuilderStatus, expectedEntries); } @@ -249,13 +249,11 @@ public Block getObject(Block block, int position) public void writeObject(BlockBuilder blockBuilder, Object value) { Block rowBlock = (Block) value; - - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (int i = 0; i < rowBlock.getPositionCount(); i++) { - fields.get(i).getType().appendTo(rowBlock, i, entryBuilder); - } - - blockBuilder.closeEntry(); + ((RowBlockBuilder) blockBuilder).buildEntry(fieldBuilders -> { + for (int i = 0; i < rowBlock.getPositionCount(); i++) { + fields.get(i).getType().appendTo(rowBlock, i, fieldBuilders.get(i)); + } + }); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/ShortDecimalType.java b/core/trino-spi/src/main/java/io/trino/spi/type/ShortDecimalType.java index 769293a6b9dd..74b0e9dd1e95 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/ShortDecimalType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/ShortDecimalType.java @@ -122,7 +122,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(block.getLong(position, 0)).closeEntry(); + writeLong(blockBuilder, block.getLong(position, 0)); } } @@ -135,7 +135,7 @@ public long getLong(Block block, int position) @Override public void writeLong(BlockBuilder blockBuilder, long value) { - blockBuilder.writeLong(value).closeEntry(); + ((LongArrayBlockBuilder) blockBuilder).writeLong(value); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimeWithTimeZoneType.java b/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimeWithTimeZoneType.java index 2f1b8b944fdd..d67909eb80c7 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimeWithTimeZoneType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimeWithTimeZoneType.java @@ -81,7 +81,7 @@ public final Slice getSlice(Block block, int position) @Override public final void writeLong(BlockBuilder blockBuilder, long value) { - blockBuilder.writeLong(value).closeEntry(); + ((LongArrayBlockBuilder) blockBuilder).writeLong(value); } @Override @@ -91,7 +91,7 @@ public final void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(block.getLong(position, 0)).closeEntry(); + writeLong(blockBuilder, block.getLong(position, 0)); } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimestampType.java b/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimestampType.java index 6703b15a55ab..f225dd73e97e 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimestampType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimestampType.java @@ -86,7 +86,7 @@ public final long getLong(Block block, int position) @Override public final void writeLong(BlockBuilder blockBuilder, long value) { - blockBuilder.writeLong(value).closeEntry(); + ((LongArrayBlockBuilder) blockBuilder).writeLong(value); } @Override @@ -96,7 +96,7 @@ public final void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(getLong(block, position)).closeEntry(); + writeLong(blockBuilder, getLong(block, position)); } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimestampWithTimeZoneType.java b/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimestampWithTimeZoneType.java index 2a0684e64e35..c66f785a10c9 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimestampWithTimeZoneType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/ShortTimestampWithTimeZoneType.java @@ -81,7 +81,7 @@ public final Slice getSlice(Block block, int position) @Override public final void writeLong(BlockBuilder blockBuilder, long value) { - blockBuilder.writeLong(value).closeEntry(); + ((LongArrayBlockBuilder) blockBuilder).writeLong(value); } @Override @@ -91,7 +91,7 @@ public final void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(block.getLong(position, 0)).closeEntry(); + writeLong(blockBuilder, block.getLong(position, 0)); } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/SmallintType.java b/core/trino-spi/src/main/java/io/trino/spi/type/SmallintType.java index a0b7c0df7a7a..a82f7ccc5de2 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/SmallintType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/SmallintType.java @@ -153,7 +153,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeShort(block.getShort(position, 0)).closeEntry(); + ((ShortArrayBlockBuilder) blockBuilder).writeShort(block.getShort(position, 0)); } } @@ -172,7 +172,12 @@ public short getShort(Block block, int position) public void writeLong(BlockBuilder blockBuilder, long value) { checkValueValid(value); - blockBuilder.writeShort((int) value).closeEntry(); + writeShort(blockBuilder, (short) value); + } + + public void writeShort(BlockBuilder blockBuilder, short value) + { + ((ShortArrayBlockBuilder) blockBuilder).writeShort(value); } private void checkValueValid(long value) diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/TimestampTypes.java b/core/trino-spi/src/main/java/io/trino/spi/type/TimestampTypes.java index c5243dcbb3ac..5b260a5f14af 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/TimestampTypes.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/TimestampTypes.java @@ -14,6 +14,7 @@ package io.trino.spi.type; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.Fixed12BlockBuilder; public final class TimestampTypes { @@ -26,8 +27,6 @@ public static void writeLongTimestamp(BlockBuilder blockBuilder, LongTimestamp t public static void writeLongTimestamp(BlockBuilder blockBuilder, long epochMicros, int fraction) { - blockBuilder.writeLong(epochMicros); - blockBuilder.writeInt(fraction); - blockBuilder.closeEntry(); + ((Fixed12BlockBuilder) blockBuilder).writeFixed12(epochMicros, fraction); } } diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/TinyintType.java b/core/trino-spi/src/main/java/io/trino/spi/type/TinyintType.java index 95f761102062..5bdbf2dac224 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/TinyintType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/TinyintType.java @@ -153,7 +153,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeByte(block.getByte(position, 0)).closeEntry(); + writeByte(blockBuilder, block.getByte(position, 0)); } } @@ -172,7 +172,12 @@ public byte getByte(Block block, int position) public void writeLong(BlockBuilder blockBuilder, long value) { checkValueValid(value); - blockBuilder.writeByte((int) value).closeEntry(); + writeByte(blockBuilder, (byte) value); + } + + public void writeByte(BlockBuilder blockBuilder, byte value) + { + ((ByteArrayBlockBuilder) blockBuilder).writeByte(value); } private void checkValueValid(long value) diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/UuidType.java b/core/trino-spi/src/main/java/io/trino/spi/type/UuidType.java index 2cadfc33e2e4..a7a200e2ec10 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/UuidType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/UuidType.java @@ -124,9 +124,9 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(block.getLong(position, 0)); - blockBuilder.writeLong(block.getLong(position, SIZE_OF_LONG)); - blockBuilder.closeEntry(); + ((Int128ArrayBlockBuilder) blockBuilder).writeInt128( + block.getLong(position, 0), + block.getLong(position, SIZE_OF_LONG)); } } @@ -142,9 +142,9 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int l if (length != INT128_BYTES) { throw new IllegalStateException("Expected entry size to be exactly " + INT128_BYTES + " but was " + length); } - blockBuilder.writeLong(value.getLong(offset)); - blockBuilder.writeLong(value.getLong(offset + SIZE_OF_LONG)); - blockBuilder.closeEntry(); + ((Int128ArrayBlockBuilder) blockBuilder).writeInt128( + value.getLong(offset), + value.getLong(offset + SIZE_OF_LONG)); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/VarbinaryType.java b/core/trino-spi/src/main/java/io/trino/spi/type/VarbinaryType.java index 52454d3d5f06..4475f2868747 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/VarbinaryType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/VarbinaryType.java @@ -17,6 +17,7 @@ import io.airlift.slice.XxHash64; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.function.BlockIndex; import io.trino.spi.function.BlockPosition; @@ -85,8 +86,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -105,7 +105,7 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value) @Override public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length) { - blockBuilder.writeBytes(value, offset, length).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length); } @Override diff --git a/core/trino-spi/src/main/java/io/trino/spi/type/VarcharType.java b/core/trino-spi/src/main/java/io/trino/spi/type/VarcharType.java index ed3e0a24abfc..451198b4aa07 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/type/VarcharType.java +++ b/core/trino-spi/src/main/java/io/trino/spi/type/VarcharType.java @@ -20,6 +20,7 @@ import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.BlockBuilderStatus; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.function.BlockIndex; import io.trino.spi.function.BlockPosition; @@ -144,7 +145,7 @@ public Object getObjectValue(ConnectorSession session, Block block, int position } @Override - public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) + public VariableWidthBlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries) { return createBlockBuilder( blockBuilderStatus, @@ -192,8 +193,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -217,7 +217,7 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value) @Override public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length) { - blockBuilder.writeBytes(value, offset, length).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length); } @Override diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/BenchmarkCopyPositions.java b/core/trino-spi/src/test/java/io/trino/spi/block/BenchmarkCopyPositions.java index 282ea0b1b20c..a07fe8ccf2df 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/BenchmarkCopyPositions.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/BenchmarkCopyPositions.java @@ -135,13 +135,13 @@ private static boolean randomNullChance(Random random) private static BlockBuilder createBlockBuilderWithValues(Slice[] generatedValues) { - BlockBuilder blockBuilder = new VariableWidthBlockBuilder(null, generatedValues.length, 32 * generatedValues.length); + VariableWidthBlockBuilder blockBuilder = new VariableWidthBlockBuilder(null, generatedValues.length, 32 * generatedValues.length); for (Slice value : generatedValues) { if (value == null) { blockBuilder.appendNull(); } else { - blockBuilder.writeBytes(value, 0, value.length()).closeEntry(); + blockBuilder.writeEntry(value); } } return blockBuilder; diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestArrayBlockBuilder.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestArrayBlockBuilder.java index 6761efd6c4bf..8a8437e69f8c 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestArrayBlockBuilder.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestArrayBlockBuilder.java @@ -36,14 +36,14 @@ public void testArrayBlockIsFull() private void testIsFull(PageBuilderStatus pageBuilderStatus) { - BlockBuilder blockBuilder = new ArrayBlockBuilder(BIGINT, pageBuilderStatus.createBlockBuilderStatus(), EXPECTED_ENTRY_COUNT); + ArrayBlockBuilder blockBuilder = new ArrayBlockBuilder(BIGINT, pageBuilderStatus.createBlockBuilderStatus(), EXPECTED_ENTRY_COUNT); assertTrue(pageBuilderStatus.isEmpty()); while (!pageBuilderStatus.isFull()) { - BlockBuilder elementBuilder = blockBuilder.beginBlockEntry(); - BIGINT.writeLong(elementBuilder, 12); - elementBuilder.appendNull(); - BIGINT.writeLong(elementBuilder, 34); - blockBuilder.closeEntry(); + blockBuilder.buildEntry(elementBuilder -> { + BIGINT.writeLong(elementBuilder, 12); + elementBuilder.appendNull(); + BIGINT.writeLong(elementBuilder, 34); + }); } assertEquals(blockBuilder.getPositionCount(), EXPECTED_ENTRY_COUNT); assertEquals(pageBuilderStatus.isFull(), true); @@ -54,12 +54,11 @@ private void testIsFull(PageBuilderStatus pageBuilderStatus) public void testRetainedSizeInBytes() { int expectedEntries = 1000; - BlockBuilder arrayBlockBuilder = new ArrayBlockBuilder(BIGINT, null, expectedEntries); + ArrayBlockBuilder arrayBlockBuilder = new ArrayBlockBuilder(BIGINT, null, expectedEntries); long initialRetainedSize = arrayBlockBuilder.getRetainedSizeInBytes(); for (int i = 0; i < expectedEntries; i++) { - BlockBuilder arrayElementBuilder = arrayBlockBuilder.beginBlockEntry(); - BIGINT.writeLong(arrayElementBuilder, i); - arrayBlockBuilder.closeEntry(); + int value = i; + arrayBlockBuilder.buildEntry(elementBuilder -> BIGINT.writeLong(elementBuilder, value)); } assertTrue(arrayBlockBuilder.getRetainedSizeInBytes() >= (expectedEntries * BYTES + instanceSize(LongArrayBlockBuilder.class) + initialRetainedSize)); } @@ -67,12 +66,13 @@ public void testRetainedSizeInBytes() @Test public void testConcurrentWriting() { - BlockBuilder blockBuilder = new ArrayBlockBuilder(BIGINT, null, EXPECTED_ENTRY_COUNT); - BlockBuilder elementBlockWriter = blockBuilder.beginBlockEntry(); - elementBlockWriter.writeLong(45).closeEntry(); - assertThatThrownBy(blockBuilder::beginBlockEntry) - .isInstanceOf(IllegalStateException.class) - .hasMessage("Expected current entry to be closed but was opened"); + ArrayBlockBuilder blockBuilder = new ArrayBlockBuilder(BIGINT, null, EXPECTED_ENTRY_COUNT); + blockBuilder.buildEntry(elementBuilder -> { + BIGINT.writeLong(elementBuilder, 45); + assertThatThrownBy(() -> blockBuilder.buildEntry(ignore -> {})) + .isInstanceOf(IllegalStateException.class) + .hasMessage("Expected current entry to be closed but was opened"); + }); } @Test diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestBlockRetainedSizeBreakdown.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestBlockRetainedSizeBreakdown.java index d44d78c50cf9..ab54425f6fae 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestBlockRetainedSizeBreakdown.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestBlockRetainedSizeBreakdown.java @@ -38,11 +38,10 @@ public class TestBlockRetainedSizeBreakdown @Test public void testArrayBlock() { - BlockBuilder arrayBlockBuilder = new ArrayBlockBuilder(BIGINT, null, EXPECTED_ENTRIES); + ArrayBlockBuilder arrayBlockBuilder = new ArrayBlockBuilder(BIGINT, null, EXPECTED_ENTRIES); for (int i = 0; i < EXPECTED_ENTRIES; i++) { - BlockBuilder arrayElementBuilder = arrayBlockBuilder.beginBlockEntry(); - writeNativeValue(BIGINT, arrayElementBuilder, castIntegerToObject(i, BIGINT)); - arrayBlockBuilder.closeEntry(); + int value = i; + arrayBlockBuilder.buildEntry(elementBuilder -> writeNativeValue(BIGINT, elementBuilder, castIntegerToObject(value, BIGINT))); } checkRetainedSize(arrayBlockBuilder.build(), false); } @@ -50,9 +49,9 @@ public void testArrayBlock() @Test public void testByteArrayBlock() { - BlockBuilder blockBuilder = new ByteArrayBlockBuilder(null, EXPECTED_ENTRIES); + ByteArrayBlockBuilder blockBuilder = new ByteArrayBlockBuilder(null, EXPECTED_ENTRIES); for (int i = 0; i < EXPECTED_ENTRIES; i++) { - blockBuilder.writeByte(i); + blockBuilder.writeByte((byte) i); } checkRetainedSize(blockBuilder.build(), false); } @@ -95,9 +94,9 @@ public void testRunLengthEncodedBlock() @Test public void testShortArrayBlock() { - BlockBuilder blockBuilder = new ShortArrayBlockBuilder(null, EXPECTED_ENTRIES); + ShortArrayBlockBuilder blockBuilder = new ShortArrayBlockBuilder(null, EXPECTED_ENTRIES); for (int i = 0; i < EXPECTED_ENTRIES; i++) { - blockBuilder.writeShort(i); + blockBuilder.writeShort((short) i); } checkRetainedSize(blockBuilder.build(), false); } diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestByteArrayBlockEncoding.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestByteArrayBlockEncoding.java index c5a0e25d793b..f0488c1ea3fe 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestByteArrayBlockEncoding.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestByteArrayBlockEncoding.java @@ -31,7 +31,7 @@ protected Type getType() @Override protected void write(BlockBuilder blockBuilder, Byte value) { - blockBuilder.writeByte(value); + TINYINT.writeByte(blockBuilder, value); } @Override diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestIntegerArrayBlockEncoding.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestIntegerArrayBlockEncoding.java index 96afa7fbd06d..617e23632271 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestIntegerArrayBlockEncoding.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestIntegerArrayBlockEncoding.java @@ -31,7 +31,7 @@ protected Type getType() @Override protected void write(BlockBuilder blockBuilder, Integer value) { - blockBuilder.writeInt(value); + ((IntArrayBlockBuilder) blockBuilder).writeInt(value); } @Override diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlockEncoding.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlockEncoding.java index 1d662813c4d6..2bb1ddc76507 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlockEncoding.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestRowBlockEncoding.java @@ -37,10 +37,10 @@ protected Type getType() @Override protected void write(BlockBuilder blockBuilder, Object[] value) { - BlockBuilder row = blockBuilder.beginBlockEntry(); - BIGINT.writeLong(row, (long) value[0]); - VARCHAR.writeSlice(row, utf8Slice((String) value[1])); - blockBuilder.closeEntry(); + ((RowBlockBuilder) blockBuilder).buildEntry(fieldBuilders -> { + BIGINT.writeLong(fieldBuilders.get(0), (long) value[0]); + VARCHAR.writeSlice(fieldBuilders.get(1), utf8Slice((String) value[1])); + }); } @Override diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestShortArrayBlockEncoding.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestShortArrayBlockEncoding.java index 66d5e1626519..1251eb7fdae8 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestShortArrayBlockEncoding.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestShortArrayBlockEncoding.java @@ -31,7 +31,7 @@ protected Type getType() @Override protected void write(BlockBuilder blockBuilder, Short value) { - blockBuilder.writeShort(value); + ((ShortArrayBlockBuilder) blockBuilder).writeShort(value); } @Override diff --git a/core/trino-spi/src/test/java/io/trino/spi/block/TestVariableWidthBlockBuilder.java b/core/trino-spi/src/test/java/io/trino/spi/block/TestVariableWidthBlockBuilder.java index 8bad38440cdb..50341168ee73 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/block/TestVariableWidthBlockBuilder.java +++ b/core/trino-spi/src/test/java/io/trino/spi/block/TestVariableWidthBlockBuilder.java @@ -45,14 +45,13 @@ public void testNewBlockBuilderLike() { int entries = 12345; double resetSkew = 1.25; - BlockBuilder blockBuilder = new VariableWidthBlockBuilder(null, entries, entries); + VariableWidthBlockBuilder blockBuilder = new VariableWidthBlockBuilder(null, entries, entries); for (int i = 0; i < entries; i++) { - blockBuilder.writeByte(i); - blockBuilder.closeEntry(); + blockBuilder.writeEntry(Slices.wrappedBuffer((byte) i)); } - blockBuilder = blockBuilder.newBlockBuilderLike(null); + blockBuilder = (VariableWidthBlockBuilder) blockBuilder.newBlockBuilderLike(null); // force to initialize capacity - blockBuilder.writeByte(1); + blockBuilder.writeEntry(Slices.wrappedBuffer((byte) 1)); long actualArrayBytes = sizeOf(new int[(int) ceil(resetSkew * (entries + 1))]) + sizeOf(new boolean[(int) ceil(resetSkew * entries)]); long actualSliceBytes = SLICE_INSTANCE_SIZE + sizeOf(new byte[(int) ceil(resetSkew * entries)]); diff --git a/core/trino-spi/src/test/java/io/trino/spi/function/TestScalarFunctionAdapter.java b/core/trino-spi/src/test/java/io/trino/spi/function/TestScalarFunctionAdapter.java index 21f410ce5b55..ac33a8b90773 100644 --- a/core/trino-spi/src/test/java/io/trino/spi/function/TestScalarFunctionAdapter.java +++ b/core/trino-spi/src/test/java/io/trino/spi/function/TestScalarFunctionAdapter.java @@ -307,9 +307,9 @@ private static Object getTestValue(Type argumentType) if (argumentType.equals(ARRAY_TYPE)) { BlockBuilder blockBuilder = BIGINT.createBlockBuilder(null, 4); blockBuilder.appendNull(); - blockBuilder.writeLong(99); + BIGINT.writeLong(blockBuilder, 99); blockBuilder.appendNull(); - blockBuilder.writeLong(100); + BIGINT.writeLong(blockBuilder, 100); return blockBuilder.build(); } if (argumentType.equals(CHAR_TYPE)) { diff --git a/lib/trino-array/src/test/java/io/trino/array/TestBlockBigArray.java b/lib/trino-array/src/test/java/io/trino/array/TestBlockBigArray.java index c52bf3fd9063..8e047bfe9fad 100644 --- a/lib/trino-array/src/test/java/io/trino/array/TestBlockBigArray.java +++ b/lib/trino-array/src/test/java/io/trino/array/TestBlockBigArray.java @@ -14,7 +14,6 @@ package io.trino.array; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.IntArrayBlockBuilder; import org.testng.annotations.Test; @@ -27,7 +26,7 @@ public class TestBlockBigArray public void testRetainedSizeWithOverlappingBlocks() { int entries = 123; - BlockBuilder blockBuilder = new IntArrayBlockBuilder(null, entries); + IntArrayBlockBuilder blockBuilder = new IntArrayBlockBuilder(null, entries); for (int i = 0; i < entries; i++) { blockBuilder.writeInt(i); } diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/avro/AvroPageDataReader.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/avro/AvroPageDataReader.java index 4c4c08f7e052..f21299ef643b 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/avro/AvroPageDataReader.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/avro/AvroPageDataReader.java @@ -17,8 +17,10 @@ import io.airlift.slice.Slices; import io.trino.spi.Page; import io.trino.spi.PageBuilder; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.BlockBuilder; -import io.trino.spi.block.SingleRowBlockWriter; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.RowType; import io.trino.spi.type.Type; import org.apache.avro.Resolver; @@ -468,17 +470,17 @@ public ArrayBlockBuildingDecoder(Resolver.Container containerAction, AvroTypeMan protected void decodeIntoBlock(Decoder decoder, BlockBuilder builder) throws IOException { - BlockBuilder elementBuilder = builder.beginBlockEntry(); - long elementsInBlock = decoder.readArrayStart(); - if (elementsInBlock > 0) { - do { - for (int i = 0; i < elementsInBlock; i++) { - elementBlockBuildingDecoder.decodeIntoBlock(decoder, elementBuilder); + ((ArrayBlockBuilder) builder).buildEntry(elementBuilder -> { + long elementsInBlock = decoder.readArrayStart(); + if (elementsInBlock > 0) { + do { + for (int i = 0; i < elementsInBlock; i++) { + elementBlockBuildingDecoder.decodeIntoBlock(decoder, elementBuilder); + } } + while ((elementsInBlock = decoder.arrayNext()) > 0); } - while ((elementsInBlock = decoder.arrayNext()) > 0); - } - builder.closeEntry(); + }); } } @@ -500,19 +502,19 @@ public MapBlockBuildingDecoder(Resolver.Container containerAction, AvroTypeManag protected void decodeIntoBlock(Decoder decoder, BlockBuilder builder) throws IOException { - BlockBuilder entryBuilder = builder.beginBlockEntry(); - long entriesInBlock = decoder.readMapStart(); - // TODO need to filter out all but last value for key? - if (entriesInBlock > 0) { - do { - for (int i = 0; i < entriesInBlock; i++) { - keyBlockBuildingDecoder.decodeIntoBlock(decoder, entryBuilder); - valueBlockBuildingDecoder.decodeIntoBlock(decoder, entryBuilder); + ((MapBlockBuilder) builder).buildEntry((keyBuilder, valueBuilder) -> { + long entriesInBlock = decoder.readMapStart(); + // TODO need to filter out all but last value for key? + if (entriesInBlock > 0) { + do { + for (int i = 0; i < entriesInBlock; i++) { + keyBlockBuildingDecoder.decodeIntoBlock(decoder, keyBuilder); + valueBlockBuildingDecoder.decodeIntoBlock(decoder, valueBuilder); + } } + while ((entriesInBlock = decoder.mapNext()) > 0); } - while ((entriesInBlock = decoder.mapNext()) > 0); - } - builder.closeEntry(); + }); } } @@ -576,9 +578,7 @@ private RowBlockBuildingDecoder(Resolver.Action action, AvroTypeManager typeMana protected void decodeIntoBlock(Decoder decoder, BlockBuilder builder) throws IOException { - SingleRowBlockWriter currentBuilder = (SingleRowBlockWriter) builder.beginBlockEntry(); - decodeIntoBlockProvided(decoder, currentBuilder::getFieldBlockBuilder); - builder.closeEntry(); + ((RowBlockBuilder) builder).buildEntry(fieldBuilders -> decodeIntoBlockProvided(decoder, fieldBuilders::get)); } protected void decodeIntoPageBuilder(Decoder decoder, PageBuilder builder) @@ -880,19 +880,19 @@ protected void decodeIntoBlock(Decoder decoder, BlockBuilder builder) protected static void makeSingleRowWithTagAndAllFieldsNullButOne(int outputChannel, int totalChannels, BlockBuildingDecoder blockBuildingDecoder, Decoder decoder, BlockBuilder builder) throws IOException { - SingleRowBlockWriter currentBuilder = (SingleRowBlockWriter) builder.beginBlockEntry(); - //add tag with channel - UNION_FIELD_TAG_TYPE.writeLong(currentBuilder.getFieldBlockBuilder(0), outputChannel); - //add in null fields except one - for (int channel = 1; channel <= totalChannels; channel++) { - if (channel == outputChannel + 1) { - blockBuildingDecoder.decodeIntoBlock(decoder, currentBuilder.getFieldBlockBuilder(channel)); - } - else { - currentBuilder.getFieldBlockBuilder(channel).appendNull(); + ((RowBlockBuilder) builder).buildEntry(fieldBuilders -> { + //add tag with channel + UNION_FIELD_TAG_TYPE.writeLong(fieldBuilders.get(0), outputChannel); + //add in null fields except one + for (int channel = 1; channel <= totalChannels; channel++) { + if (channel == outputChannel + 1) { + blockBuildingDecoder.decodeIntoBlock(decoder, fieldBuilders.get(channel)); + } + else { + fieldBuilders.get(channel).appendNull(); + } } - } - builder.closeEntry(); + }); } protected static int[] getIndexToChannel(List schemas) diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/ListEncoding.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/ListEncoding.java index d6db5e0fab55..4aff509f9d29 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/ListEncoding.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/ListEncoding.java @@ -16,6 +16,7 @@ import io.airlift.slice.Slice; import io.airlift.slice.SliceOutput; import io.trino.hive.formats.ReadWriteUtils; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.type.Type; @@ -63,7 +64,11 @@ public void encodeValue(Block block, int position, SliceOutput output) @Override public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int length) { - // entries in list + ((ArrayBlockBuilder) builder).buildEntry(elementBuilder -> decodeArrayInto(elementBuilder, slice, offset)); + } + + private void decodeArrayInto(BlockBuilder elementBuilder, Slice slice, int offset) + { int entries = toIntExact(ReadWriteUtils.readVInt(slice, offset)); offset += ReadWriteUtils.decodeVIntSize(slice.getByte(offset)); @@ -73,24 +78,22 @@ public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int l // read elements starting after null bytes int elementOffset = nullByteEnd; - BlockBuilder arrayBuilder = builder.beginBlockEntry(); for (int i = 0; i < entries; i++) { if ((slice.getByte(nullByteCur) & (1 << (i % 8))) != 0) { int valueOffset = elementEncoding.getValueOffset(slice, elementOffset); int valueLength = elementEncoding.getValueLength(slice, elementOffset); - elementEncoding.decodeValueInto(arrayBuilder, slice, elementOffset + valueOffset, valueLength); + elementEncoding.decodeValueInto(elementBuilder, slice, elementOffset + valueOffset, valueLength); elementOffset = elementOffset + valueOffset + valueLength; } else { - arrayBuilder.appendNull(); + elementBuilder.appendNull(); } // move onto the next null byte if (7 == (i % 8)) { nullByteCur++; } } - builder.closeEntry(); } } diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/MapEncoding.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/MapEncoding.java index 77558d475f0d..587666385427 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/MapEncoding.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/MapEncoding.java @@ -20,6 +20,7 @@ import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.type.Type; import static java.lang.Math.toIntExact; @@ -82,6 +83,11 @@ public void encodeValue(Block block, int position, SliceOutput output) @Override public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int length) + { + ((MapBlockBuilder) builder).buildEntry((keyBuilder, valueBuilder) -> decodeValueInto(keyBuilder, valueBuilder, slice, offset)); + } + + private void decodeValueInto(BlockBuilder keyBuilder, BlockBuilder valueBuilder, Slice slice, int offset) { // entries in list int entries = toIntExact(ReadWriteUtils.readVInt(slice, offset)); @@ -93,7 +99,6 @@ public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int l // read elements starting after null bytes int elementOffset = nullByteEnd; - BlockBuilder mapBuilder = builder.beginBlockEntry(); for (int i = 0; i < entries; i++) { // read key boolean nullKey; @@ -101,7 +106,7 @@ public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int l int keyOffset = keyReader.getValueOffset(slice, elementOffset); int keyLength = keyReader.getValueLength(slice, elementOffset); - keyReader.decodeValueInto(mapBuilder, slice, elementOffset + keyOffset, keyLength); + keyReader.decodeValueInto(keyBuilder, slice, elementOffset + keyOffset, keyLength); nullKey = false; elementOffset = elementOffset + keyOffset + keyLength; @@ -119,7 +124,7 @@ public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int l // ignore entries with a null key if (!nullKey) { - valueReader.decodeValueInto(mapBuilder, slice, elementOffset + valueOffset, valueLength); + valueReader.decodeValueInto(valueBuilder, slice, elementOffset + valueOffset, valueLength); } elementOffset = elementOffset + valueOffset + valueLength; @@ -127,7 +132,7 @@ public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int l else { // ignore entries with a null key if (!nullKey) { - mapBuilder.appendNull(); + valueBuilder.appendNull(); } } @@ -136,6 +141,5 @@ public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int l nullByteCur++; } } - builder.closeEntry(); } } diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/StructEncoding.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/StructEncoding.java index 3c79434ea594..4690def5507b 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/StructEncoding.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/binary/StructEncoding.java @@ -18,6 +18,7 @@ import io.airlift.slice.SliceOutput; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.Type; import java.util.List; @@ -61,40 +62,38 @@ public void encodeValue(Block block, int position, SliceOutput output) @Override public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int length) { - int fieldId = 0; - int nullByte = 0; - int elementOffset = offset; - BlockBuilder rowBuilder = builder.beginBlockEntry(); - while (fieldId < structFields.size() && elementOffset < offset + length) { - BinaryColumnEncoding field = structFields.get(fieldId); + ((RowBlockBuilder) builder).buildEntry(fieldBuilders -> { + int fieldId = 0; + int nullByte = 0; + int elementOffset = offset; + while (fieldId < structFields.size() && elementOffset < offset + length) { + BinaryColumnEncoding field = structFields.get(fieldId); - // null byte prefixes every 8 fields - if ((fieldId % 8) == 0) { - nullByte = slice.getByte(elementOffset); - elementOffset++; - } + // null byte prefixes every 8 fields + if ((fieldId % 8) == 0) { + nullByte = slice.getByte(elementOffset); + elementOffset++; + } - // read field - if ((nullByte & (1 << (fieldId % 8))) != 0) { - int valueOffset = field.getValueOffset(slice, elementOffset); - int valueLength = field.getValueLength(slice, elementOffset); + // read field + if ((nullByte & (1 << (fieldId % 8))) != 0) { + int valueOffset = field.getValueOffset(slice, elementOffset); + int valueLength = field.getValueLength(slice, elementOffset); - field.decodeValueInto(rowBuilder, slice, elementOffset + valueOffset, valueLength); + field.decodeValueInto(fieldBuilders.get(fieldId), slice, elementOffset + valueOffset, valueLength); - elementOffset = elementOffset + valueOffset + valueLength; + elementOffset = elementOffset + valueOffset + valueLength; + } + else { + fieldBuilders.get(fieldId).appendNull(); + } + fieldId++; } - else { - rowBuilder.appendNull(); + // Sometimes a struct does not have all fields written, so we fill with nulls + while (fieldId < structFields.size()) { + fieldBuilders.get(fieldId).appendNull(); + fieldId++; } - fieldId++; - } - - // Sometimes a struct does not have all fields written, so we fill with nulls - while (fieldId < structFields.size()) { - rowBuilder.appendNull(); - fieldId++; - } - - builder.closeEntry(); + }); } } diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/ListEncoding.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/ListEncoding.java index 0541474082a3..2074e042a99b 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/ListEncoding.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/ListEncoding.java @@ -16,6 +16,7 @@ import io.airlift.slice.Slice; import io.airlift.slice.SliceOutput; import io.trino.hive.formats.FileCorruptionException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.type.Type; @@ -55,26 +56,31 @@ public void encodeValueInto(Block block, int position, SliceOutput output) public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int length) throws FileCorruptionException { - int end = offset + length; + ((ArrayBlockBuilder) builder).buildEntry(elementBuilder -> decodeArrayInto(elementBuilder, slice, offset, length)); + } + + private void decodeArrayInto(BlockBuilder elementBuilder, Slice slice, int offset, int length) + throws FileCorruptionException + { + if (length <= 0) { + return; + } - BlockBuilder arrayBlockBuilder = builder.beginBlockEntry(); - if (length > 0) { - int elementOffset = offset; - while (offset < end) { - byte currentByte = slice.getByte(offset); - if (currentByte == separator) { - decodeElementValueInto(arrayBlockBuilder, slice, elementOffset, offset - elementOffset); - elementOffset = offset + 1; - } - else if (isEscapeByte(currentByte) && offset + 1 < length) { - // ignore the char after escape_char - offset++; - } + int end = offset + length; + int elementOffset = offset; + while (offset < end) { + byte currentByte = slice.getByte(offset); + if (currentByte == separator) { + decodeElementValueInto(elementBuilder, slice, elementOffset, offset - elementOffset); + elementOffset = offset + 1; + } + else if (isEscapeByte(currentByte) && offset + 1 < length) { + // ignore the char after escape_char offset++; } - decodeElementValueInto(arrayBlockBuilder, slice, elementOffset, offset - elementOffset); + offset++; } - builder.closeEntry(); + decodeElementValueInto(elementBuilder, slice, elementOffset, offset - elementOffset); } private void decodeElementValueInto(BlockBuilder blockBuilder, Slice slice, int offset, int length) diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/MapEncoding.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/MapEncoding.java index e45025916b09..cdb0de0f7705 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/MapEncoding.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/MapEncoding.java @@ -21,6 +21,7 @@ import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.type.MapType; import io.trino.spi.type.Type; @@ -94,9 +95,10 @@ public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int l boolean[] distinctKeys = distinctMapKeys.selectDistinctKeys(keyBlock); // add the distinct entries to the map - BlockBuilder mapBuilder = builder.beginBlockEntry(); - processEntries(slice, offset, length, new DistinctEntryDecoder(distinctKeys, keyBlock, mapBuilder)); - builder.closeEntry(); + ((MapBlockBuilder) builder).buildEntry((keyBuilder, valueBuilder) -> { + DistinctEntryDecoder entryDecoder = new DistinctEntryDecoder(distinctKeys, keyBlock, keyBuilder, valueBuilder); + processEntries(slice, offset, length, entryDecoder); + }); } private void processEntries(Slice slice, int offset, int length, EntryDecoder entryDecoder) @@ -182,14 +184,16 @@ private class DistinctEntryDecoder { private final boolean[] distinctKeys; private final Block keyBlock; - private final BlockBuilder mapBuilder; + private final BlockBuilder keyBuilder; + private final BlockBuilder valueBuilder; private int entryPosition; - public DistinctEntryDecoder(boolean[] distinctKeys, Block keyBlock, BlockBuilder mapBuilder) + public DistinctEntryDecoder(boolean[] distinctKeys, Block keyBlock, BlockBuilder keyBuilder, BlockBuilder valueBuilder) { this.distinctKeys = distinctKeys; this.keyBlock = keyBlock; - this.mapBuilder = mapBuilder; + this.keyBuilder = keyBuilder; + this.valueBuilder = valueBuilder; } @Override @@ -197,13 +201,13 @@ public void decodeKeyValue(int depth, Slice slice, int keyOffset, int keyLength, throws FileCorruptionException { if (distinctKeys[entryPosition]) { - mapType.getKeyType().appendTo(keyBlock, entryPosition, mapBuilder); + mapType.getKeyType().appendTo(keyBlock, entryPosition, keyBuilder); if (hasValue && !isNullSequence(slice, valueOffset, valueLength)) { - valueEncoding.decodeValueInto(mapBuilder, slice, valueOffset, valueLength); + valueEncoding.decodeValueInto(valueBuilder, slice, valueOffset, valueLength); } else { - mapBuilder.appendNull(); + valueBuilder.appendNull(); } } entryPosition++; diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/StructEncoding.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/StructEncoding.java index 1a039d7cf1c7..15116590b06b 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/StructEncoding.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/encodings/text/StructEncoding.java @@ -18,6 +18,7 @@ import io.trino.hive.formats.FileCorruptionException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.Type; import java.util.List; @@ -67,42 +68,40 @@ public void decodeValueInto(BlockBuilder builder, Slice slice, int offset, int l throws FileCorruptionException { int end = offset + length; - - BlockBuilder structBuilder = builder.beginBlockEntry(); - int elementOffset = offset; - int fieldIndex = 0; - while (offset < end) { - byte currentByte = slice.getByte(offset); - if (currentByte == separator) { - decodeElementValueInto(fieldIndex, structBuilder, slice, elementOffset, offset - elementOffset); - elementOffset = offset + 1; - fieldIndex++; - if (lastColumnTakesRest && fieldIndex == structFields.size() - 1) { - // no need to process the remaining bytes as they are all assigned to the last column - break; + ((RowBlockBuilder) builder).buildEntry(fieldBuilders -> { + int currentOffset = offset; + int elementOffset = currentOffset; + int fieldIndex = 0; + while (currentOffset < end) { + byte currentByte = slice.getByte(currentOffset); + if (currentByte == separator) { + decodeElementValueInto(fieldIndex, fieldBuilders.get(fieldIndex), slice, elementOffset, currentOffset - elementOffset); + elementOffset = currentOffset + 1; + fieldIndex++; + if (lastColumnTakesRest && fieldIndex == structFields.size() - 1) { + // no need to process the remaining bytes as they are all assigned to the last column + break; + } + if (fieldIndex == structFields.size()) { + // this was the last field, so there is no more data to process + return; + } } - if (fieldIndex == structFields.size()) { - // this was the last field, so there is no more data to process - builder.closeEntry(); - return; + else if (isEscapeByte(currentByte)) { + // ignore the char after escape_char + currentOffset++; } + currentOffset++; } - else if (isEscapeByte(currentByte)) { - // ignore the char after escape_char - offset++; - } - offset++; - } - decodeElementValueInto(fieldIndex, structBuilder, slice, elementOffset, end - elementOffset); - fieldIndex++; - - // missing fields are null - while (fieldIndex < structFields.size()) { - structBuilder.appendNull(); + decodeElementValueInto(fieldIndex, fieldBuilders.get(fieldIndex), slice, elementOffset, end - elementOffset); fieldIndex++; - } - builder.closeEntry(); + // missing fields are null + while (fieldIndex < structFields.size()) { + fieldBuilders.get(fieldIndex).appendNull(); + fieldIndex++; + } + }); } private void decodeElementValueInto(int fieldIndex, BlockBuilder builder, Slice slice, int offset, int length) diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/json/JsonDeserializer.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/json/JsonDeserializer.java index e408082b4edc..506cdf67af29 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/json/JsonDeserializer.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/json/JsonDeserializer.java @@ -27,9 +27,11 @@ import io.trino.hive.formats.line.LineDeserializer; import io.trino.plugin.base.type.DecodedTimestamp; import io.trino.spi.PageBuilder; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; -import io.trino.spi.block.SingleRowBlockWriter; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.CharType; import io.trino.spi.type.DecimalType; @@ -528,15 +530,14 @@ public ArrayDecoder(ArrayType arrayType, Decoder elementDecoder) void decodeValue(LineBuffer lineBuffer, JsonParser parser, BlockBuilder builder) throws IOException { - BlockBuilder elementBuilder = builder.beginBlockEntry(); - - if (parser.currentToken() != START_ARRAY) { - throw invalidJson("start of array expected"); - } - while (nextTokenRequired(parser) != JsonToken.END_ARRAY) { - elementDecoder.decode(lineBuffer, parser, elementBuilder); - } - builder.closeEntry(); + ((ArrayBlockBuilder) builder).buildEntry(elementBuilder -> { + if (parser.currentToken() != START_ARRAY) { + throw invalidJson("start of array expected"); + } + while (nextTokenRequired(parser) != JsonToken.END_ARRAY) { + elementDecoder.decode(lineBuffer, parser, elementBuilder); + } + }); } } @@ -570,23 +571,23 @@ void decodeValue(LineBuffer lineBuffer, JsonParser parser, BlockBuilder builder) Block keyBlock = readKeys(createParserAt(parser.currentTokenLocation(), lineBuffer)); boolean[] distinctKeys = distinctMapKeys.selectDistinctKeys(keyBlock); - BlockBuilder entryBuilder = builder.beginBlockEntry(); - if (parser.currentToken() != START_OBJECT) { - throw invalidJson("start of object expected"); - } - int keyIndex = 0; - while (nextObjectField(parser)) { - if (distinctKeys[keyIndex]) { - keyType.appendTo(keyBlock, keyIndex, entryBuilder); - parser.nextToken(); - valueDecoder.decode(lineBuffer, parser, entryBuilder); + ((MapBlockBuilder) builder).buildEntry((keyBuilder, valueBuilder) -> { + if (parser.currentToken() != START_OBJECT) { + throw invalidJson("start of object expected"); } - else { - skipNextValue(parser); + int keyIndex = 0; + while (nextObjectField(parser)) { + if (distinctKeys[keyIndex]) { + keyType.appendTo(keyBlock, keyIndex, keyBuilder); + parser.nextToken(); + valueDecoder.decode(lineBuffer, parser, valueBuilder); + } + else { + skipNextValue(parser); + } + keyIndex++; } - keyIndex++; - } - builder.closeEntry(); + }); } private Block readKeys(JsonParser fieldNameParser) @@ -687,9 +688,7 @@ public void decode(LineBuffer lineBuffer, JsonParser parser, PageBuilder builder void decodeValue(LineBuffer lineBuffer, JsonParser parser, BlockBuilder builder) throws IOException { - SingleRowBlockWriter currentBuilder = (SingleRowBlockWriter) builder.beginBlockEntry(); - decodeValue(lineBuffer, parser, currentBuilder::getFieldBlockBuilder); - builder.closeEntry(); + ((RowBlockBuilder) builder).buildEntry(fieldBuilders -> decodeValue(lineBuffer, parser, fieldBuilders::get)); } private void decodeValue(LineBuffer lineBuffer, JsonParser parser, IntFunction fieldBuilders) diff --git a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/openxjson/OpenXJsonDeserializer.java b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/openxjson/OpenXJsonDeserializer.java index 80711b97955a..74698cb1956c 100644 --- a/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/openxjson/OpenXJsonDeserializer.java +++ b/lib/trino-hive-formats/src/main/java/io/trino/hive/formats/line/openxjson/OpenXJsonDeserializer.java @@ -24,9 +24,11 @@ import io.trino.plugin.base.type.DecodedTimestamp; import io.trino.plugin.base.type.TrinoTimestampEncoder; import io.trino.spi.PageBuilder; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; -import io.trino.spi.block.SingleRowBlockWriter; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.CharType; import io.trino.spi.type.DecimalType; @@ -658,18 +660,17 @@ void decodeValue(Object jsonValue, BlockBuilder builder) return; } - BlockBuilder elementBuilder = builder.beginBlockEntry(); - - if (jsonValue instanceof List jsonArray) { - for (Object element : jsonArray) { - elementDecoder.decode(element, elementBuilder); + ((ArrayBlockBuilder) builder).buildEntry(elementBuilder -> { + if (jsonValue instanceof List jsonArray) { + for (Object element : jsonArray) { + elementDecoder.decode(element, elementBuilder); + } } - } - else { - // all other values are coerced to a single element array - elementDecoder.decode(jsonValue, elementBuilder); - } - builder.closeEntry(); + else { + // all other values are coerced to a single element array + elementDecoder.decode(jsonValue, elementBuilder); + } + }); } } @@ -681,7 +682,7 @@ private static class MapDecoder private final Type keyType; private final DistinctMapKeys distinctMapKeys; - private BlockBuilder keyBlockBuilder; + private BlockBuilder tempKeyBlockBuilder; public MapDecoder(MapType mapType, Decoder keyDecoder, Decoder valueDecoder) { @@ -690,7 +691,7 @@ public MapDecoder(MapType mapType, Decoder keyDecoder, Decoder valueDecoder) this.valueDecoder = valueDecoder; this.distinctMapKeys = new DistinctMapKeys(mapType, true); - this.keyBlockBuilder = mapType.getKeyType().createBlockBuilder(null, 128); + this.tempKeyBlockBuilder = mapType.getKeyType().createBlockBuilder(null, 128); } @Override @@ -713,16 +714,16 @@ void decodeValue(Object jsonValue, BlockBuilder builder) Block keyBlock = readKeys(fieldNames); boolean[] distinctKeys = distinctMapKeys.selectDistinctKeys(keyBlock); - BlockBuilder entryBuilder = builder.beginBlockEntry(); - int keyIndex = 0; - for (Object fieldName : fieldNames) { - if (distinctKeys[keyIndex]) { - keyType.appendTo(keyBlock, keyIndex, entryBuilder); - valueDecoder.decode(jsonObject.get(fieldName), entryBuilder); + ((MapBlockBuilder) builder).buildEntry((keyBuilder, valueBuilder) -> { + int keyIndex = 0; + for (Object fieldName : fieldNames) { + if (distinctKeys[keyIndex]) { + keyType.appendTo(keyBlock, keyIndex, keyBuilder); + valueDecoder.decode(jsonObject.get(fieldName), valueBuilder); + } + keyIndex++; } - keyIndex++; - } - builder.closeEntry(); + }); } private Block readKeys(Collection fieldNames) @@ -731,11 +732,11 @@ private Block readKeys(Collection fieldNames) // field names are always processed as a quoted JSON string even though they may // have not been quoted in the original JSON text JsonString jsonValue = new JsonString(fieldName.toString(), true); - keyDecoder.decode(jsonValue, keyBlockBuilder); + keyDecoder.decode(jsonValue, tempKeyBlockBuilder); } - Block keyBlock = keyBlockBuilder.build(); - keyBlockBuilder = keyType.createBlockBuilder(null, keyBlock.getPositionCount()); + Block keyBlock = tempKeyBlockBuilder.build(); + tempKeyBlockBuilder = keyType.createBlockBuilder(null, keyBlock.getPositionCount()); return keyBlock; } } @@ -759,7 +760,6 @@ public RowDecoder(RowType rowType, OpenXJsonOptions options, List field } public void decode(Object jsonValue, PageBuilder builder) - throws IOException { builder.declarePosition(); decodeValue(jsonValue, builder::getBlockBuilder); @@ -768,9 +768,7 @@ public void decode(Object jsonValue, PageBuilder builder) @Override void decodeValue(Object jsonValue, BlockBuilder builder) { - SingleRowBlockWriter currentBuilder = (SingleRowBlockWriter) builder.beginBlockEntry(); - decodeValue(jsonValue, currentBuilder::getFieldBlockBuilder); - builder.closeEntry(); + ((RowBlockBuilder) builder).buildEntry(fieldBuilders -> decodeValue(jsonValue, fieldBuilders::get)); } private void decodeValue(Object jsonValue, IntFunction fieldBuilders) diff --git a/lib/trino-hive-formats/src/test/java/io/trino/hive/formats/FormatTestUtils.java b/lib/trino-hive-formats/src/test/java/io/trino/hive/formats/FormatTestUtils.java index 4d7063b1a023..ba40c71991fc 100644 --- a/lib/trino-hive-formats/src/test/java/io/trino/hive/formats/FormatTestUtils.java +++ b/lib/trino-hive-formats/src/test/java/io/trino/hive/formats/FormatTestUtils.java @@ -22,7 +22,10 @@ import io.trino.plugin.base.type.DecodedTimestamp; import io.trino.spi.Page; import io.trino.spi.PageBuilder; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.CharType; import io.trino.spi.type.DecimalType; @@ -538,32 +541,30 @@ else if (type instanceof TimestampType timestampType) { else if (type instanceof ArrayType) { List array = (List) value; Type elementType = type.getTypeParameters().get(0); - BlockBuilder arrayBlockBuilder = blockBuilder.beginBlockEntry(); - for (Object elementValue : array) { - writeTrinoValue(elementType, arrayBlockBuilder, elementValue); - } - blockBuilder.closeEntry(); + ((ArrayBlockBuilder) blockBuilder).buildEntry(elementBuilder -> { + for (Object elementValue : array) { + writeTrinoValue(elementType, elementBuilder, elementValue); + } + }); } else if (type instanceof MapType) { Map map = (Map) value; Type keyType = type.getTypeParameters().get(0); Type valueType = type.getTypeParameters().get(1); - BlockBuilder mapBlockBuilder = blockBuilder.beginBlockEntry(); - map.forEach((entryKey, entryValue) -> { - writeTrinoValue(keyType, mapBlockBuilder, entryKey); - writeTrinoValue(valueType, mapBlockBuilder, entryValue); - }); - blockBuilder.closeEntry(); + ((MapBlockBuilder) blockBuilder).buildEntry((keyBuilder, valueBuilder) -> map.forEach((entryKey, entryValue) -> { + writeTrinoValue(keyType, keyBuilder, entryKey); + writeTrinoValue(valueType, valueBuilder, entryValue); + })); } else if (type instanceof RowType) { List array = (List) value; List fieldTypes = type.getTypeParameters(); - BlockBuilder rowBlockBuilder = blockBuilder.beginBlockEntry(); - for (int fieldId = 0; fieldId < fieldTypes.size(); fieldId++) { - Type fieldType = fieldTypes.get(fieldId); - writeTrinoValue(fieldType, rowBlockBuilder, array.get(fieldId)); - } - blockBuilder.closeEntry(); + ((RowBlockBuilder) blockBuilder).buildEntry(fieldBuilders -> { + for (int fieldId = 0; fieldId < fieldTypes.size(); fieldId++) { + Type fieldType = fieldTypes.get(fieldId); + writeTrinoValue(fieldType, fieldBuilders.get(fieldId), array.get(fieldId)); + } + }); } else { throw new IllegalArgumentException("Unsupported type: " + type); diff --git a/lib/trino-orc/src/main/java/io/trino/orc/writer/DictionaryBuilder.java b/lib/trino-orc/src/main/java/io/trino/orc/writer/DictionaryBuilder.java index 1a2b6e1780b4..13d1bb6e7bc8 100644 --- a/lib/trino-orc/src/main/java/io/trino/orc/writer/DictionaryBuilder.java +++ b/lib/trino-orc/src/main/java/io/trino/orc/writer/DictionaryBuilder.java @@ -15,7 +15,6 @@ import io.trino.array.IntBigArray; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.VariableWidthBlockBuilder; import static com.google.common.base.Preconditions.checkArgument; @@ -39,7 +38,7 @@ public class DictionaryBuilder private static final int EXPECTED_BYTES_PER_ENTRY = 32; private final IntBigArray blockPositionByHash = new IntBigArray(); - private BlockBuilder elementBlock; + private VariableWidthBlockBuilder elementBlock; private int maxFill; private int hashMask; @@ -90,7 +89,7 @@ public void clear() { containsNullElement = false; blockPositionByHash.fill(EMPTY_SLOT); - elementBlock = elementBlock.newBlockBuilderLike(null); + elementBlock = (VariableWidthBlockBuilder) elementBlock.newBlockBuilderLike(null); // first position is always null elementBlock.appendNull(); } @@ -158,8 +157,7 @@ private long getHashPositionOfElement(Block block, int position) private int addNewElement(long hashPosition, Block block, int position) { checkArgument(!block.isNull(position), "position is null"); - block.writeBytesTo(position, 0, block.getSliceLength(position), elementBlock); - elementBlock.closeEntry(); + elementBlock.buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); int newElementPositionInBlock = elementBlock.getPositionCount() - 1; blockPositionByHash.set(hashPosition, newElementPositionInBlock); diff --git a/lib/trino-orc/src/test/java/io/trino/orc/OrcTester.java b/lib/trino-orc/src/test/java/io/trino/orc/OrcTester.java index 6230a777b1f1..14e48231ef34 100644 --- a/lib/trino-orc/src/test/java/io/trino/orc/OrcTester.java +++ b/lib/trino-orc/src/test/java/io/trino/orc/OrcTester.java @@ -28,8 +28,11 @@ import io.trino.orc.metadata.CompressionKind; import io.trino.orc.metadata.OrcType; import io.trino.spi.Page; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.CharType; import io.trino.spi.type.DecimalType; @@ -760,32 +763,32 @@ else if (TIMESTAMP_TZ_MICROS.equals(type) || TIMESTAMP_TZ_NANOS.equals(type)) { if (type instanceof ArrayType) { List array = (List) value; Type elementType = type.getTypeParameters().get(0); - BlockBuilder arrayBlockBuilder = blockBuilder.beginBlockEntry(); - for (Object elementValue : array) { - writeValue(elementType, arrayBlockBuilder, elementValue); - } - blockBuilder.closeEntry(); + ((ArrayBlockBuilder) blockBuilder).buildEntry(elementBuilder -> { + for (Object elementValue : array) { + writeValue(elementType, elementBuilder, elementValue); + } + }); } - else if (type instanceof MapType) { + else if (type instanceof MapType mapType) { Map map = (Map) value; - Type keyType = type.getTypeParameters().get(0); - Type valueType = type.getTypeParameters().get(1); - BlockBuilder mapBlockBuilder = blockBuilder.beginBlockEntry(); - for (Entry entry : map.entrySet()) { - writeValue(keyType, mapBlockBuilder, entry.getKey()); - writeValue(valueType, mapBlockBuilder, entry.getValue()); - } - blockBuilder.closeEntry(); + Type keyType = mapType.getKeyType(); + Type valueType = mapType.getValueType(); + ((MapBlockBuilder) blockBuilder).buildEntry((keyBuilder, valueBuilder) -> { + map.forEach((key, value1) -> { + writeValue(keyType, keyBuilder, key); + writeValue(valueType, valueBuilder, value1); + }); + }); } else if (type instanceof RowType) { List array = (List) value; List fieldTypes = type.getTypeParameters(); - BlockBuilder rowBlockBuilder = blockBuilder.beginBlockEntry(); - for (int fieldId = 0; fieldId < fieldTypes.size(); fieldId++) { - Type fieldType = fieldTypes.get(fieldId); - writeValue(fieldType, rowBlockBuilder, array.get(fieldId)); - } - blockBuilder.closeEntry(); + ((RowBlockBuilder) blockBuilder).buildEntry(fieldBuilders -> { + for (int fieldId = 0; fieldId < fieldTypes.size(); fieldId++) { + Type fieldType = fieldTypes.get(fieldId); + writeValue(fieldType, fieldBuilders.get(fieldId), array.get(fieldId)); + } + }); } else { throw new IllegalArgumentException("Unsupported type " + type); diff --git a/lib/trino-orc/src/test/java/io/trino/orc/TestOrcWriter.java b/lib/trino-orc/src/test/java/io/trino/orc/TestOrcWriter.java index c50bd89040df..19097430462f 100644 --- a/lib/trino-orc/src/test/java/io/trino/orc/TestOrcWriter.java +++ b/lib/trino-orc/src/test/java/io/trino/orc/TestOrcWriter.java @@ -31,7 +31,7 @@ import io.trino.orc.stream.OrcInputStream; import io.trino.spi.Page; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.type.Type; import org.testng.annotations.Test; @@ -88,17 +88,16 @@ public void testWriteOutputStreamsInOrder() String[] data = new String[] {"a", "bbbbb", "ccc", "dd", "eeee"}; Block[] blocks = new Block[data.length]; int entries = 65536; - BlockBuilder blockBuilder = VARCHAR.createBlockBuilder(null, entries); + VariableWidthBlockBuilder blockBuilder = VARCHAR.createBlockBuilder(null, entries); for (int i = 0; i < data.length; i++) { byte[] bytes = data[i].getBytes(UTF_8); for (int j = 0; j < entries; j++) { // force to write different data bytes[0] = (byte) ((bytes[0] + 1) % 128); - blockBuilder.writeBytes(Slices.wrappedBuffer(bytes, 0, bytes.length), 0, bytes.length); - blockBuilder.closeEntry(); + blockBuilder.writeEntry(Slices.wrappedBuffer(bytes, 0, bytes.length)); } blocks[i] = blockBuilder.build(); - blockBuilder = blockBuilder.newBlockBuilderLike(null); + blockBuilder = (VariableWidthBlockBuilder) blockBuilder.newBlockBuilderLike(null); } writer.write(new Page(blocks)); diff --git a/lib/trino-orc/src/test/java/io/trino/orc/TestSliceDictionaryColumnWriter.java b/lib/trino-orc/src/test/java/io/trino/orc/TestSliceDictionaryColumnWriter.java index 14b902e4ed1b..d483409e2a54 100644 --- a/lib/trino-orc/src/test/java/io/trino/orc/TestSliceDictionaryColumnWriter.java +++ b/lib/trino-orc/src/test/java/io/trino/orc/TestSliceDictionaryColumnWriter.java @@ -102,8 +102,7 @@ public void testBloomFiltersAfterConvertToDirect() testValues.add(value); base = (byte) (base + i); if (i % 9 == 0) { - blockBuilder.writeBytes(value, 0, value.length()); - blockBuilder.closeEntry(); + blockBuilder.writeEntry(value); } } diff --git a/lib/trino-orc/src/test/java/io/trino/orc/TestStructColumnReader.java b/lib/trino-orc/src/test/java/io/trino/orc/TestStructColumnReader.java index d146219d7dbd..bfc13d33649e 100644 --- a/lib/trino-orc/src/test/java/io/trino/orc/TestStructColumnReader.java +++ b/lib/trino-orc/src/test/java/io/trino/orc/TestStructColumnReader.java @@ -22,8 +22,8 @@ import io.trino.orc.metadata.OrcType; import io.trino.spi.Page; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.RowBlock; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.type.NamedTypeSignature; import io.trino.spi.type.RowFieldName; import io.trino.spi.type.StandardTypes; @@ -57,8 +57,6 @@ @Test(singleThreaded = true) public class TestStructColumnReader { - private static final Type TEST_DATA_TYPE = VARCHAR; - private static final String STRUCT_COL_NAME = "struct_col"; private TempFile tempFile; @@ -239,15 +237,14 @@ private void write(TempFile tempFile, Type writerType, List data) boolean[] rowIsNull = new boolean[entries]; Arrays.fill(rowIsNull, false); - BlockBuilder blockBuilder = TEST_DATA_TYPE.createBlockBuilder(null, entries); + VariableWidthBlockBuilder blockBuilder = VARCHAR.createBlockBuilder(null, entries); for (int i = 0; i < data.size(); i++) { byte[] bytes = data.get(i).getBytes(UTF_8); for (int j = 0; j < entries; j++) { - blockBuilder.writeBytes(Slices.wrappedBuffer(bytes), 0, bytes.length); - blockBuilder.closeEntry(); + blockBuilder.writeEntry(Slices.wrappedBuffer(bytes)); } fieldBlocks[i] = blockBuilder.build(); - blockBuilder = blockBuilder.newBlockBuilderLike(null); + blockBuilder = (VariableWidthBlockBuilder) blockBuilder.newBlockBuilderLike(null); } Block rowBlock = RowBlock.fromFieldBlocks(rowIsNull.length, Optional.of(rowIsNull), fieldBlocks); writer.write(new Page(rowBlock)); @@ -279,7 +276,7 @@ private Type getType(List fieldNames) { ImmutableList.Builder typeSignatureParameters = ImmutableList.builder(); for (String fieldName : fieldNames) { - typeSignatureParameters.add(TypeSignatureParameter.namedTypeParameter(new NamedTypeSignature(Optional.of(new RowFieldName(fieldName)), TEST_DATA_TYPE.getTypeSignature()))); + typeSignatureParameters.add(TypeSignatureParameter.namedTypeParameter(new NamedTypeSignature(Optional.of(new RowFieldName(fieldName)), VARCHAR.getTypeSignature()))); } return TESTING_TYPE_MANAGER.getParameterizedType(StandardTypes.ROW, typeSignatureParameters.build()); } @@ -289,7 +286,7 @@ private Type getTypeNullName(int numFields) ImmutableList.Builder typeSignatureParameters = ImmutableList.builder(); for (int i = 0; i < numFields; i++) { - typeSignatureParameters.add(TypeSignatureParameter.namedTypeParameter(new NamedTypeSignature(Optional.empty(), TEST_DATA_TYPE.getTypeSignature()))); + typeSignatureParameters.add(TypeSignatureParameter.namedTypeParameter(new NamedTypeSignature(Optional.empty(), VARCHAR.getTypeSignature()))); } return TESTING_TYPE_MANAGER.getParameterizedType(StandardTypes.ROW, typeSignatureParameters.build()); } diff --git a/lib/trino-record-decoder/src/main/java/io/trino/decoder/avro/AvroColumnDecoder.java b/lib/trino-record-decoder/src/main/java/io/trino/decoder/avro/AvroColumnDecoder.java index 1c33a6acac7f..89bfc06d9dec 100644 --- a/lib/trino-record-decoder/src/main/java/io/trino/decoder/avro/AvroColumnDecoder.java +++ b/lib/trino-record-decoder/src/main/java/io/trino/decoder/avro/AvroColumnDecoder.java @@ -22,6 +22,8 @@ import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.BigintType; import io.trino.spi.type.BooleanType; @@ -51,6 +53,7 @@ import static io.airlift.slice.Slices.utf8Slice; import static io.trino.decoder.DecoderErrorCode.DECODER_CONVERSION_NOT_SUPPORTED; import static io.trino.spi.StandardErrorCode.GENERIC_USER_ERROR; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.type.Varchars.truncateToLength; import static java.lang.Float.floatToIntBits; import static java.lang.String.format; @@ -241,8 +244,8 @@ private static Block serializeObject(BlockBuilder builder, Object value, Type ty if (type instanceof ArrayType) { return serializeList(builder, value, type, columnName); } - if (type instanceof MapType) { - return serializeMap(builder, value, type, columnName); + if (type instanceof MapType mapType) { + return serializeMap(builder, value, mapType, columnName); } if (type instanceof RowType) { return serializeRow(builder, value, type, columnName); @@ -310,7 +313,7 @@ private static void serializePrimitive(BlockBuilder blockBuilder, Object value, throw new TrinoException(DECODER_CONVERSION_NOT_SUPPORTED, format("cannot decode object of '%s' as '%s' for column '%s'", value.getClass(), type, columnName)); } - private static Block serializeMap(BlockBuilder parentBlockBuilder, Object value, Type type, String columnName) + private static Block serializeMap(BlockBuilder parentBlockBuilder, Object value, MapType type, String columnName) { if (value == null) { checkState(parentBlockBuilder != null, "parentBlockBuilder is null"); @@ -319,26 +322,25 @@ private static Block serializeMap(BlockBuilder parentBlockBuilder, Object value, } Map map = (Map) value; - List typeParameters = type.getTypeParameters(); - Type keyType = typeParameters.get(0); - Type valueType = typeParameters.get(1); + Type keyType = type.getKeyType(); + Type valueType = type.getValueType(); - BlockBuilder blockBuilder; + MapBlockBuilder blockBuilder; if (parentBlockBuilder != null) { - blockBuilder = parentBlockBuilder; + blockBuilder = (MapBlockBuilder) parentBlockBuilder; } else { blockBuilder = type.createBlockBuilder(null, 1); } - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (Map.Entry entry : map.entrySet()) { - if (entry.getKey() != null) { - keyType.writeSlice(entryBuilder, truncateToLength(utf8Slice(entry.getKey().toString()), keyType)); - serializeObject(entryBuilder, entry.getValue(), valueType, columnName); + blockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + for (Map.Entry entry : map.entrySet()) { + if (entry.getKey() != null) { + keyType.writeSlice(keyBuilder, truncateToLength(utf8Slice(entry.getKey().toString()), keyType)); + serializeObject(valueBuilder, entry.getValue(), valueType, columnName); + } } - } - blockBuilder.closeEntry(); + }); if (parentBlockBuilder == null) { return blockBuilder.getObject(0, Block.class); @@ -346,32 +348,30 @@ private static Block serializeMap(BlockBuilder parentBlockBuilder, Object value, return null; } - private static Block serializeRow(BlockBuilder parentBlockBuilder, Object value, Type type, String columnName) + private static Block serializeRow(BlockBuilder blockBuilder, Object value, Type type, String columnName) { if (value == null) { - checkState(parentBlockBuilder != null, "parent block builder is null"); - parentBlockBuilder.appendNull(); + checkState(blockBuilder != null, "block builder is null"); + blockBuilder.appendNull(); return null; } - BlockBuilder blockBuilder; - if (parentBlockBuilder != null) { - blockBuilder = parentBlockBuilder; - } - else { - blockBuilder = type.createBlockBuilder(null, 1); + RowType rowType = (RowType) type; + if (blockBuilder == null) { + return buildRowValue(rowType, fieldBuilders -> buildRow(rowType, columnName, (GenericRecord) value, fieldBuilders)); } - BlockBuilder singleRowBuilder = blockBuilder.beginBlockEntry(); - GenericRecord record = (GenericRecord) value; - List fields = ((RowType) type).getFields(); - for (Field field : fields) { + + ((RowBlockBuilder) blockBuilder).buildEntry(fieldBuilders -> buildRow(rowType, columnName, (GenericRecord) value, fieldBuilders)); + return null; + } + + private static void buildRow(RowType type, String columnName, GenericRecord record, List fieldBuilders) + { + List fields = type.getFields(); + for (int i = 0; i < fields.size(); i++) { + Field field = fields.get(i); checkState(field.getName().isPresent(), "field name not found"); - serializeObject(singleRowBuilder, record.get(field.getName().get()), field.getType(), columnName); + serializeObject(fieldBuilders.get(i), record.get(field.getName().get()), field.getType(), columnName); } - blockBuilder.closeEntry(); - if (parentBlockBuilder == null) { - return blockBuilder.getObject(0, Block.class); - } - return null; } } diff --git a/lib/trino-record-decoder/src/main/java/io/trino/decoder/protobuf/ProtobufValueProvider.java b/lib/trino-record-decoder/src/main/java/io/trino/decoder/protobuf/ProtobufValueProvider.java index da5daddd0391..bafffabf27f1 100644 --- a/lib/trino-record-decoder/src/main/java/io/trino/decoder/protobuf/ProtobufValueProvider.java +++ b/lib/trino-record-decoder/src/main/java/io/trino/decoder/protobuf/ProtobufValueProvider.java @@ -23,6 +23,8 @@ import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.BigintType; import io.trino.spi.type.BooleanType; @@ -31,6 +33,7 @@ import io.trino.spi.type.MapType; import io.trino.spi.type.RealType; import io.trino.spi.type.RowType; +import io.trino.spi.type.RowType.Field; import io.trino.spi.type.SmallintType; import io.trino.spi.type.TimestampType; import io.trino.spi.type.TinyintType; @@ -50,6 +53,7 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static io.airlift.slice.Slices.utf8Slice; import static io.trino.decoder.DecoderErrorCode.DECODER_CONVERSION_NOT_SUPPORTED; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.type.StandardTypes.JSON; import static io.trino.spi.type.TimestampType.MAX_SHORT_PRECISION; import static io.trino.spi.type.Timestamps.MICROSECONDS_PER_SECOND; @@ -157,8 +161,8 @@ private Block serializeObject(BlockBuilder builder, Object value, Type type, Str if (type instanceof ArrayType) { return serializeList(builder, value, type, columnName); } - if (type instanceof MapType) { - return serializeMap(builder, value, type, columnName); + if (type instanceof MapType mapType) { + return serializeMap(builder, value, mapType, columnName); } if (type instanceof RowType) { return serializeRow(builder, value, type, columnName); @@ -239,7 +243,7 @@ private void serializePrimitive(BlockBuilder blockBuilder, @Nullable Object valu } @Nullable - private Block serializeMap(BlockBuilder parentBlockBuilder, @Nullable Object value, Type type, String columnName) + private Block serializeMap(BlockBuilder parentBlockBuilder, @Nullable Object value, MapType type, String columnName) { if (value == null) { checkState(parentBlockBuilder != null, "parentBlockBuilder is null"); @@ -250,26 +254,25 @@ private Block serializeMap(BlockBuilder parentBlockBuilder, @Nullable Object val Collection dynamicMessages = ((Collection) value).stream() .map(DynamicMessage.class::cast) .collect(toImmutableList()); - List typeParameters = type.getTypeParameters(); - Type keyType = typeParameters.get(0); - Type valueType = typeParameters.get(1); + Type keyType = type.getKeyType(); + Type valueType = type.getValueType(); - BlockBuilder blockBuilder; + MapBlockBuilder blockBuilder; if (parentBlockBuilder != null) { - blockBuilder = parentBlockBuilder; + blockBuilder = (MapBlockBuilder) parentBlockBuilder; } else { blockBuilder = type.createBlockBuilder(null, 1); } - BlockBuilder entryBuilder = blockBuilder.beginBlockEntry(); - for (DynamicMessage dynamicMessage : dynamicMessages) { - if (dynamicMessage.getField(dynamicMessage.getDescriptorForType().findFieldByNumber(1)) != null) { - serializeObject(entryBuilder, dynamicMessage.getField(getFieldDescriptor(dynamicMessage, 1)), keyType, columnName); - serializeObject(entryBuilder, dynamicMessage.getField(getFieldDescriptor(dynamicMessage, 2)), valueType, columnName); + blockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + for (DynamicMessage dynamicMessage : dynamicMessages) { + if (dynamicMessage.getField(dynamicMessage.getDescriptorForType().findFieldByNumber(1)) != null) { + serializeObject(keyBuilder, dynamicMessage.getField(getFieldDescriptor(dynamicMessage, 1)), keyType, columnName); + serializeObject(valueBuilder, dynamicMessage.getField(getFieldDescriptor(dynamicMessage, 2)), valueType, columnName); + } } - } - blockBuilder.closeEntry(); + }); if (parentBlockBuilder == null) { return blockBuilder.getObject(0, Block.class); @@ -278,39 +281,36 @@ private Block serializeMap(BlockBuilder parentBlockBuilder, @Nullable Object val } @Nullable - private Block serializeRow(BlockBuilder parentBlockBuilder, @Nullable Object value, Type type, String columnName) + private Block serializeRow(BlockBuilder blockBuilder, @Nullable Object value, Type type, String columnName) { if (value == null) { - checkState(parentBlockBuilder != null, "parent block builder is null"); - parentBlockBuilder.appendNull(); + checkState(blockBuilder != null, "parent block builder is null"); + blockBuilder.appendNull(); return null; } - BlockBuilder blockBuilder; - if (parentBlockBuilder != null) { - blockBuilder = parentBlockBuilder; - } - else { - blockBuilder = type.createBlockBuilder(null, 1); + RowType rowType = (RowType) type; + if (blockBuilder == null) { + return buildRowValue(rowType, fieldBuilders -> buildRow(rowType, columnName, (DynamicMessage) value, fieldBuilders)); } - BlockBuilder singleRowBuilder = blockBuilder.beginBlockEntry(); - DynamicMessage record = (DynamicMessage) value; - List fields = ((RowType) type).getFields(); - for (RowType.Field field : fields) { + ((RowBlockBuilder) blockBuilder).buildEntry((fieldBuilders) -> buildRow(rowType, columnName, (DynamicMessage) value, fieldBuilders)); + return null; + } + + private void buildRow(RowType rowType, String columnName, DynamicMessage record, List fieldBuilders) + { + List fields = rowType.getFields(); + for (int i = 0; i < fields.size(); i++) { + Field field = fields.get(i); checkState(field.getName().isPresent(), "field name not found"); FieldDescriptor fieldDescriptor = getFieldDescriptor(record, field.getName().get()); checkState(fieldDescriptor != null, format("Unknown Field %s", field.getName().get())); serializeObject( - singleRowBuilder, + fieldBuilders.get(i), record.getField(fieldDescriptor), field.getType(), columnName); } - blockBuilder.closeEntry(); - if (parentBlockBuilder == null) { - return blockBuilder.getObject(0, Block.class); - } - return null; } @Nullable diff --git a/plugin/trino-accumulo/src/main/java/io/trino/plugin/accumulo/serializers/AccumuloRowSerializer.java b/plugin/trino-accumulo/src/main/java/io/trino/plugin/accumulo/serializers/AccumuloRowSerializer.java index 889b2ff96ce6..5f49e54ee38f 100644 --- a/plugin/trino-accumulo/src/main/java/io/trino/plugin/accumulo/serializers/AccumuloRowSerializer.java +++ b/plugin/trino-accumulo/src/main/java/io/trino/plugin/accumulo/serializers/AccumuloRowSerializer.java @@ -16,8 +16,11 @@ import com.google.common.collect.ImmutableList; import io.airlift.slice.Slice; import io.trino.plugin.accumulo.Types; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.type.MapType; import io.trino.spi.type.Type; import io.trino.spi.type.TypeUtils; import io.trino.spi.type.VarcharType; @@ -32,6 +35,8 @@ import java.util.Map; import java.util.Map.Entry; +import static io.trino.spi.block.MapValueBuilder.buildMapValue; + /** * Interface for deserializing the data in Accumulo into a Trino row. *

@@ -541,25 +546,22 @@ static Block getBlockFromArray(Type elementType, List array) /** * Encodes the given map into a Block. * - * @param mapType Trino type of the map + * @param type Trino type of the map * @param map Map of key/value pairs to encode * @return Trino Block */ - static Block getBlockFromMap(Type mapType, Map map) + static Block getBlockFromMap(Type type, Map map) { - Type keyType = mapType.getTypeParameters().get(0); - Type valueType = mapType.getTypeParameters().get(1); - - BlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1); - BlockBuilder builder = mapBlockBuilder.beginBlockEntry(); - - for (Entry entry : map.entrySet()) { - writeObject(builder, keyType, entry.getKey()); - writeObject(builder, valueType, entry.getValue()); - } - - mapBlockBuilder.closeEntry(); - return (Block) mapType.getObject(mapBlockBuilder, 0); + MapType mapType = (MapType) type; + Type keyType = mapType.getKeyType(); + Type valueType = mapType.getValueType(); + + return buildMapValue(mapType, map.size(), (keyBuilder, valueBuilder) -> { + map.forEach((key, value) -> { + writeObject(keyBuilder, keyType, key); + writeObject(valueBuilder, valueType, value); + }); + }); } /** @@ -574,20 +576,20 @@ static Block getBlockFromMap(Type mapType, Map map) static void writeObject(BlockBuilder builder, Type type, Object obj) { if (Types.isArrayType(type)) { - BlockBuilder arrayBldr = builder.beginBlockEntry(); - Type elementType = Types.getElementType(type); - for (Object item : (List) obj) { - writeObject(arrayBldr, elementType, item); - } - builder.closeEntry(); + ((ArrayBlockBuilder) builder).buildEntry(elementBuilder -> { + Type elementType = Types.getElementType(type); + for (Object item : (List) obj) { + writeObject(elementBuilder, elementType, item); + } + }); } - else if (Types.isMapType(type)) { - BlockBuilder mapBlockBuilder = builder.beginBlockEntry(); - for (Entry entry : ((Map) obj).entrySet()) { - writeObject(mapBlockBuilder, Types.getKeyType(type), entry.getKey()); - writeObject(mapBlockBuilder, Types.getValueType(type), entry.getValue()); - } - builder.closeEntry(); + else if (type instanceof MapType mapType) { + ((MapBlockBuilder) builder).buildEntry((keyBuilder, valueBuilder) -> { + for (Entry entry : ((Map) obj).entrySet()) { + writeObject(keyBuilder, mapType.getKeyType(), entry.getKey()); + writeObject(valueBuilder, mapType.getValueType(), entry.getValue()); + } + }); } else { TypeUtils.writeNativeValue(type, builder, obj); diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowToPageConverter.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowToPageConverter.java index 46279fd985bc..c80724efdb75 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowToPageConverter.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowToPageConverter.java @@ -17,8 +17,10 @@ import io.airlift.slice.Slice; import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.DecimalType; import io.trino.spi.type.Decimals; @@ -250,40 +252,40 @@ private void writeBlock(BlockBuilder output, Type type, FieldVector vector, int private void writeArrayBlock(BlockBuilder output, Type type, FieldVector vector, int index) { - BlockBuilder block = output.beginBlockEntry(); - Type elementType = getOnlyElement(type.getTypeParameters()); + ((ArrayBlockBuilder) output).buildEntry(elementBuilder -> { + Type elementType = getOnlyElement(type.getTypeParameters()); - ArrowBuf offsetBuffer = vector.getOffsetBuffer(); + ArrowBuf offsetBuffer = vector.getOffsetBuffer(); - int start = offsetBuffer.getInt((long) index * OFFSET_WIDTH); - int end = offsetBuffer.getInt((long) (index + 1) * OFFSET_WIDTH); + int start = offsetBuffer.getInt((long) index * OFFSET_WIDTH); + int end = offsetBuffer.getInt((long) (index + 1) * OFFSET_WIDTH); - FieldVector innerVector = ((ListVector) vector).getDataVector(); + FieldVector innerVector = ((ListVector) vector).getDataVector(); - TransferPair transferPair = innerVector.getTransferPair(allocator); - transferPair.splitAndTransfer(start, end - start); - try (FieldVector sliced = (FieldVector) transferPair.getTo()) { - convertType(block, elementType, sliced, 0, sliced.getValueCount()); - } - output.closeEntry(); + TransferPair transferPair = innerVector.getTransferPair(allocator); + transferPair.splitAndTransfer(start, end - start); + try (FieldVector sliced = (FieldVector) transferPair.getTo()) { + convertType(elementBuilder, elementType, sliced, 0, sliced.getValueCount()); + } + }); } private void writeRowBlock(BlockBuilder output, Type type, FieldVector vector, int index) { - BlockBuilder builder = output.beginBlockEntry(); - ImmutableList.Builder fieldNamesBuilder = ImmutableList.builder(); - for (int i = 0; i < type.getTypeSignature().getParameters().size(); i++) { - TypeSignatureParameter parameter = type.getTypeSignature().getParameters().get(i); - fieldNamesBuilder.add(parameter.getNamedTypeSignature().getName().orElse("field" + i)); - } - List fieldNames = fieldNamesBuilder.build(); - checkState(fieldNames.size() == type.getTypeParameters().size(), "fieldNames size differs from type %s type parameters size", type); + ((RowBlockBuilder) output).buildEntry(fieldBuilders -> { + ImmutableList.Builder fieldNamesBuilder = ImmutableList.builder(); + for (int i = 0; i < type.getTypeSignature().getParameters().size(); i++) { + TypeSignatureParameter parameter = type.getTypeSignature().getParameters().get(i); + fieldNamesBuilder.add(parameter.getNamedTypeSignature().getName().orElse("field" + i)); + } + List fieldNames = fieldNamesBuilder.build(); + checkState(fieldNames.size() == type.getTypeParameters().size(), "fieldNames size differs from type %s type parameters size", type); - for (int i = 0; i < type.getTypeParameters().size(); i++) { - FieldVector innerVector = ((StructVector) vector).getChild(fieldNames.get(i)); - convertType(builder, type.getTypeParameters().get(i), innerVector, index, 1); - } - output.closeEntry(); + for (int i = 0; i < type.getTypeParameters().size(); i++) { + FieldVector innerVector = ((StructVector) vector).getChild(fieldNames.get(i)); + convertType(fieldBuilders.get(i), type.getTypeParameters().get(i), innerVector, index, 1); + } + }); } @Override diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryQueryPageSource.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryQueryPageSource.java index 125166465f7c..dc679596f500 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryQueryPageSource.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryQueryPageSource.java @@ -23,8 +23,10 @@ import io.trino.spi.Page; import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.ArrayType; @@ -234,24 +236,20 @@ else if (type instanceof VarbinaryType) { private void writeBlock(BlockBuilder output, Type type, FieldValue value) { if (type instanceof ArrayType) { - BlockBuilder builder = output.beginBlockEntry(); - - for (FieldValue element : value.getRepeatedValue()) { - appendTo(type.getTypeParameters().get(0), element, builder); - } - - output.closeEntry(); + ((ArrayBlockBuilder) output).buildEntry(elementBuilder -> { + for (FieldValue element : value.getRepeatedValue()) { + appendTo(type.getTypeParameters().get(0), element, elementBuilder); + } + }); return; } if (type instanceof RowType) { FieldValueList record = value.getRecordValue(); - BlockBuilder builder = output.beginBlockEntry(); - - for (int index = 0; index < type.getTypeParameters().size(); index++) { - appendTo(type.getTypeParameters().get(index), record.get(index), builder); - } - output.closeEntry(); - return; + ((RowBlockBuilder) output).buildEntry(fieldBuilders -> { + for (int index = 0; index < type.getTypeParameters().size(); index++) { + appendTo(type.getTypeParameters().get(index), record.get(index), fieldBuilders.get(index)); + } + }); } throw new TrinoException(GENERIC_INTERNAL_ERROR, "Unhandled type for Block: " + type.getTypeSignature()); } diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryStorageAvroPageSource.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryStorageAvroPageSource.java index c6d1e47b100c..8df8d34aff6d 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryStorageAvroPageSource.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryStorageAvroPageSource.java @@ -21,8 +21,10 @@ import io.trino.spi.Page; import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.type.ArrayType; import io.trino.spi.type.DecimalType; @@ -250,29 +252,25 @@ private static void writeObject(BlockBuilder output, Type type, Object value) private void writeBlock(BlockBuilder output, Type type, Object value) { if (type instanceof ArrayType && value instanceof List) { - BlockBuilder builder = output.beginBlockEntry(); - - for (Object element : (List) value) { - appendTo(type.getTypeParameters().get(0), element, builder); - } - - output.closeEntry(); + ((ArrayBlockBuilder) output).buildEntry(elementBuilder -> { + for (Object element : (List) value) { + appendTo(type.getTypeParameters().get(0), element, elementBuilder); + } + }); return; } if (type instanceof RowType && value instanceof GenericRecord record) { - BlockBuilder builder = output.beginBlockEntry(); - - List fieldNames = new ArrayList<>(); - for (int i = 0; i < type.getTypeSignature().getParameters().size(); i++) { - TypeSignatureParameter parameter = type.getTypeSignature().getParameters().get(i); - fieldNames.add(parameter.getNamedTypeSignature().getName().orElse("field" + i)); - } - checkState(fieldNames.size() == type.getTypeParameters().size(), "fieldName doesn't match with type size : %s", type); - for (int index = 0; index < type.getTypeParameters().size(); index++) { - appendTo(type.getTypeParameters().get(index), record.get(fieldNames.get(index)), builder); - } - output.closeEntry(); - return; + ((RowBlockBuilder) output).buildEntry(fieldBuilders -> { + List fieldNames = new ArrayList<>(); + for (int i = 0; i < type.getTypeSignature().getParameters().size(); i++) { + TypeSignatureParameter parameter = type.getTypeSignature().getParameters().get(i); + fieldNames.add(parameter.getNamedTypeSignature().getName().orElse("field" + i)); + } + checkState(fieldNames.size() == type.getTypeParameters().size(), "fieldName doesn't match with type size : %s", type); + for (int index = 0; index < type.getTypeParameters().size(); index++) { + appendTo(type.getTypeParameters().get(index), record.get(fieldNames.get(index)), fieldBuilders.get(index)); + } + }); } throw new TrinoException(GENERIC_INTERNAL_ERROR, "Unhandled type for Block: " + type.getTypeSignature()); } diff --git a/plugin/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraTypeManager.java b/plugin/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraTypeManager.java index b0244367300e..2c92dd6eb95e 100644 --- a/plugin/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraTypeManager.java +++ b/plugin/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraTypeManager.java @@ -35,8 +35,7 @@ import io.airlift.slice.Slice; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; -import io.trino.spi.block.RowBlockBuilder; -import io.trino.spi.block.SingleRowBlockWriter; +import io.trino.spi.block.BlockBuilder; import io.trino.spi.predicate.NullableValue; import io.trino.spi.type.BigintType; import io.trino.spi.type.BooleanType; @@ -89,6 +88,7 @@ import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; import static io.trino.spi.StandardErrorCode.INVALID_CAST_ARGUMENT; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.type.DateTimeEncoding.packDateTimeWithZone; import static io.trino.spi.type.DateTimeEncoding.unpackMillisUtc; import static io.trino.spi.type.Timestamps.PICOSECONDS_PER_NANOSECOND; @@ -344,37 +344,37 @@ private Block buildTupleValue(CassandraType type, GettableByIndex row, int posit { verify(type.getKind() == TUPLE, "Not a TUPLE type"); TupleValue tupleValue = row.getTupleValue(position); - RowBlockBuilder blockBuilder = (RowBlockBuilder) type.getTrinoType().createBlockBuilder(null, 1); - SingleRowBlockWriter singleRowBlockWriter = blockBuilder.beginBlockEntry(); - int tuplePosition = 0; - for (CassandraType argumentType : type.getArgumentTypes()) { - int finalTuplePosition = tuplePosition; - NullableValue value = getColumnValue(argumentType, tupleValue, tuplePosition, () -> tupleValue.getType().getComponentTypes().get(finalTuplePosition)); - writeNativeValue(argumentType.getTrinoType(), singleRowBlockWriter, value.getValue()); - tuplePosition++; - } - // can I just return singleRowBlockWriter here? It extends AbstractSingleRowBlock and tests pass. - blockBuilder.closeEntry(); - return (Block) type.getTrinoType().getObject(blockBuilder, 0); + return buildRowValue((RowType) type.getTrinoType(), fieldBuilders -> { + int tuplePosition = 0; + List argumentTypes = type.getArgumentTypes(); + for (int i = 0; i < argumentTypes.size(); i++) { + CassandraType argumentType = argumentTypes.get(i); + BlockBuilder fieldBuilder = fieldBuilders.get(i); + int finalTuplePosition = tuplePosition; + NullableValue value = getColumnValue(argumentType, tupleValue, tuplePosition, () -> tupleValue.getType().getComponentTypes().get(finalTuplePosition)); + writeNativeValue(argumentType.getTrinoType(), fieldBuilder, value.getValue()); + tuplePosition++; + } + }); } private Block buildUserTypeValue(CassandraType type, GettableByIndex row, int position) { verify(type.getKind() == UDT, "Not a user defined type: %s", type.getKind()); UdtValue udtValue = row.getUdtValue(position); - RowBlockBuilder blockBuilder = (RowBlockBuilder) type.getTrinoType().createBlockBuilder(null, 1); - SingleRowBlockWriter singleRowBlockWriter = blockBuilder.beginBlockEntry(); - int tuplePosition = 0; - List udtTypeFieldTypes = udtValue.getType().getFieldTypes(); - for (CassandraType argumentType : type.getArgumentTypes()) { - int finalTuplePosition = tuplePosition; - NullableValue value = getColumnValue(argumentType, udtValue, tuplePosition, () -> udtTypeFieldTypes.get(finalTuplePosition)); - writeNativeValue(argumentType.getTrinoType(), singleRowBlockWriter, value.getValue()); - tuplePosition++; - } - - blockBuilder.closeEntry(); - return (Block) type.getTrinoType().getObject(blockBuilder, 0); + return buildRowValue((RowType) type.getTrinoType(), fieldBuilders -> { + int tuplePosition = 0; + List udtTypeFieldTypes = udtValue.getType().getFieldTypes(); + List argumentTypes = type.getArgumentTypes(); + for (int i = 0; i < argumentTypes.size(); i++) { + CassandraType argumentType = argumentTypes.get(i); + BlockBuilder fieldBuilder = fieldBuilders.get(i); + int finalTuplePosition = tuplePosition; + NullableValue value = getColumnValue(argumentType, udtValue, tuplePosition, () -> udtTypeFieldTypes.get(finalTuplePosition)); + writeNativeValue(argumentType.getTrinoType(), fieldBuilder, value.getValue()); + tuplePosition++; + } + }); } // TODO unify with toCqlLiteral diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/DeltaLakeParquetStatisticsUtils.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/DeltaLakeParquetStatisticsUtils.java index adbba2203695..7d0f8d3fff2b 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/DeltaLakeParquetStatisticsUtils.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/DeltaLakeParquetStatisticsUtils.java @@ -18,9 +18,7 @@ import io.airlift.slice.Slice; import io.trino.plugin.base.type.DecodedTimestamp; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.block.ColumnarRow; -import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.DateType; import io.trino.spi.type.DecimalType; @@ -60,6 +58,7 @@ import static io.airlift.slice.Slices.utf8Slice; import static io.trino.parquet.ParquetTimestampUtils.decodeInt96Timestamp; import static io.trino.spi.block.ColumnarRow.toColumnarRow; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.BooleanType.BOOLEAN; import static io.trino.spi.type.DateTimeEncoding.unpackMillisUtc; @@ -146,18 +145,15 @@ public static Object jsonValueToTrinoValue(Type type, @Nullable Object jsonValue if (type instanceof RowType rowType) { Map values = (Map) jsonValue; List fieldTypes = rowType.getTypeParameters(); - BlockBuilder blockBuilder = new RowBlockBuilder(fieldTypes, null, 1); - BlockBuilder singleRowBlockWriter = blockBuilder.beginBlockEntry(); - for (int i = 0; i < values.size(); ++i) { - Type fieldType = fieldTypes.get(i); - String fieldName = rowType.getFields().get(i).getName().orElseThrow(() -> new IllegalArgumentException("Field name must exist")); - Object fieldValue = jsonValueToTrinoValue(fieldType, values.remove(fieldName)); - writeNativeValue(fieldType, singleRowBlockWriter, fieldValue); - } - checkState(values.isEmpty(), "All fields must be converted into Trino value: %s", values); - - blockBuilder.closeEntry(); - return blockBuilder.build(); + return buildRowValue(rowType, fields -> { + for (int i = 0; i < values.size(); ++i) { + Type fieldType = fieldTypes.get(i); + String fieldName = rowType.getFields().get(i).getName().orElseThrow(() -> new IllegalArgumentException("Field name must exist")); + Object fieldValue = jsonValueToTrinoValue(fieldType, values.remove(fieldName)); + writeNativeValue(fieldType, fields.get(i), fieldValue); + } + checkState(values.isEmpty(), "All fields must be converted into Trino value: %s", values); + }); } throw new UnsupportedOperationException("Unsupported type: " + type); diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointWriter.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointWriter.java index 5d11a9ba4e99..570c6dab0e0f 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointWriter.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/checkpoint/CheckpointWriter.java @@ -29,12 +29,16 @@ import io.trino.plugin.deltalake.transactionlog.statistics.DeltaLakeJsonFileStatistics; import io.trino.plugin.deltalake.transactionlog.statistics.DeltaLakeParquetFileStatistics; import io.trino.spi.PageBuilder; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.DateTimeEncoding; import io.trino.spi.type.MapType; import io.trino.spi.type.RowType; +import io.trino.spi.type.RowType.Field; import io.trino.spi.type.TimestampType; import io.trino.spi.type.Type; import io.trino.spi.type.TypeManager; @@ -153,23 +157,22 @@ public void write(CheckpointEntries entries, TrinoOutputFile outputFile) private void writeMetadataEntry(PageBuilder pageBuilder, RowType entryType, MetadataEntry metadataEntry) { pageBuilder.declarePosition(); - BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(METADATA_BLOCK_CHANNEL); - BlockBuilder entryBlockBuilder = blockBuilder.beginBlockEntry(); - writeString(entryBlockBuilder, entryType, 0, "id", metadataEntry.getId()); - writeString(entryBlockBuilder, entryType, 1, "name", metadataEntry.getName()); - writeString(entryBlockBuilder, entryType, 2, "description", metadataEntry.getDescription()); - - RowType formatType = getInternalRowType(entryType, 3, "format"); - BlockBuilder formatBlockBuilder = entryBlockBuilder.beginBlockEntry(); - writeString(formatBlockBuilder, formatType, 0, "provider", metadataEntry.getFormat().getProvider()); - writeStringMap(formatBlockBuilder, formatType, 1, "options", metadataEntry.getFormat().getOptions()); - entryBlockBuilder.closeEntry(); - - writeString(entryBlockBuilder, entryType, 4, "schemaString", metadataEntry.getSchemaString()); - writeStringList(entryBlockBuilder, entryType, 5, "partitionColumns", metadataEntry.getOriginalPartitionColumns()); - writeStringMap(entryBlockBuilder, entryType, 6, "configuration", metadataEntry.getConfiguration()); - writeLong(entryBlockBuilder, entryType, 7, "createdTime", metadataEntry.getCreatedTime()); - blockBuilder.closeEntry(); + ((RowBlockBuilder) pageBuilder.getBlockBuilder(METADATA_BLOCK_CHANNEL)).buildEntry(fieldBuilders -> { + writeString(fieldBuilders.get(0), entryType, 0, "id", metadataEntry.getId()); + writeString(fieldBuilders.get(1), entryType, 1, "name", metadataEntry.getName()); + writeString(fieldBuilders.get(2), entryType, 2, "description", metadataEntry.getDescription()); + + RowType formatType = getInternalRowType(entryType, 3, "format"); + ((RowBlockBuilder) fieldBuilders.get(3)).buildEntry(formatBlockBuilders -> { + writeString(formatBlockBuilders.get(0), formatType, 0, "provider", metadataEntry.getFormat().getProvider()); + writeStringMap(formatBlockBuilders.get(1), formatType, 1, "options", metadataEntry.getFormat().getOptions()); + }); + + writeString(fieldBuilders.get(4), entryType, 4, "schemaString", metadataEntry.getSchemaString()); + writeStringList(fieldBuilders.get(5), entryType, 5, "partitionColumns", metadataEntry.getOriginalPartitionColumns()); + writeStringMap(fieldBuilders.get(6), entryType, 6, "configuration", metadataEntry.getConfiguration()); + writeLong(fieldBuilders.get(7), entryType, 7, "createdTime", metadataEntry.getCreatedTime()); + }); // null for others appendNullOtherBlocks(pageBuilder, METADATA_BLOCK_CHANNEL); @@ -178,18 +181,23 @@ private void writeMetadataEntry(PageBuilder pageBuilder, RowType entryType, Meta private void writeProtocolEntry(PageBuilder pageBuilder, RowType entryType, ProtocolEntry protocolEntry) { pageBuilder.declarePosition(); - BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(PROTOCOL_BLOCK_CHANNEL); - BlockBuilder entryBlockBuilder = blockBuilder.beginBlockEntry(); - int fieldId = 0; - writeLong(entryBlockBuilder, entryType, fieldId++, "minReaderVersion", (long) protocolEntry.getMinReaderVersion()); - writeLong(entryBlockBuilder, entryType, fieldId++, "minWriterVersion", (long) protocolEntry.getMinWriterVersion()); - if (protocolEntry.getReaderFeatures().isPresent()) { - writeStringList(entryBlockBuilder, entryType, fieldId++, "readerFeatures", protocolEntry.getReaderFeatures().get().stream().collect(toImmutableList())); - } - if (protocolEntry.getWriterFeatures().isPresent()) { - writeStringList(entryBlockBuilder, entryType, fieldId++, "writerFeatures", protocolEntry.getWriterFeatures().get().stream().collect(toImmutableList())); - } - blockBuilder.closeEntry(); + ((RowBlockBuilder) pageBuilder.getBlockBuilder(PROTOCOL_BLOCK_CHANNEL)).buildEntry(fieldBuilders -> { + int fieldId = 0; + writeLong(fieldBuilders.get(fieldId), entryType, fieldId, "minReaderVersion", (long) protocolEntry.getMinReaderVersion()); + fieldId++; + + writeLong(fieldBuilders.get(fieldId), entryType, fieldId, "minWriterVersion", (long) protocolEntry.getMinWriterVersion()); + fieldId++; + + if (protocolEntry.getReaderFeatures().isPresent()) { + writeStringList(fieldBuilders.get(fieldId), entryType, fieldId, "readerFeatures", protocolEntry.getReaderFeatures().get().stream().collect(toImmutableList())); + fieldId++; + } + + if (protocolEntry.getWriterFeatures().isPresent()) { + writeStringList(fieldBuilders.get(fieldId), entryType, fieldId, "writerFeatures", protocolEntry.getWriterFeatures().get().stream().collect(toImmutableList())); + } + }); // null for others appendNullOtherBlocks(pageBuilder, PROTOCOL_BLOCK_CHANNEL); @@ -198,12 +206,11 @@ private void writeProtocolEntry(PageBuilder pageBuilder, RowType entryType, Prot private void writeTransactionEntry(PageBuilder pageBuilder, RowType entryType, TransactionEntry transactionEntry) { pageBuilder.declarePosition(); - BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(TXN_BLOCK_CHANNEL); - BlockBuilder entryBlockBuilder = blockBuilder.beginBlockEntry(); - writeString(entryBlockBuilder, entryType, 0, "appId", transactionEntry.getAppId()); - writeLong(entryBlockBuilder, entryType, 1, "version", transactionEntry.getVersion()); - writeLong(entryBlockBuilder, entryType, 2, "lastUpdated", transactionEntry.getLastUpdated()); - blockBuilder.closeEntry(); + ((RowBlockBuilder) pageBuilder.getBlockBuilder(TXN_BLOCK_CHANNEL)).buildEntry(fieldBuilders -> { + writeString(fieldBuilders.get(0), entryType, 0, "appId", transactionEntry.getAppId()); + writeLong(fieldBuilders.get(1), entryType, 1, "version", transactionEntry.getVersion()); + writeLong(fieldBuilders.get(2), entryType, 2, "lastUpdated", transactionEntry.getLastUpdated()); + }); // null for others appendNullOtherBlocks(pageBuilder, TXN_BLOCK_CHANNEL); @@ -212,22 +219,36 @@ private void writeTransactionEntry(PageBuilder pageBuilder, RowType entryType, T private void writeAddFileEntry(PageBuilder pageBuilder, RowType entryType, AddFileEntry addFileEntry, MetadataEntry metadataEntry, boolean writeStatsAsJson, boolean writeStatsAsStruct) { pageBuilder.declarePosition(); - BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(ADD_BLOCK_CHANNEL); - BlockBuilder entryBlockBuilder = blockBuilder.beginBlockEntry(); - int fieldId = 0; - writeString(entryBlockBuilder, entryType, fieldId++, "path", addFileEntry.getPath()); - writeStringMap(entryBlockBuilder, entryType, fieldId++, "partitionValues", addFileEntry.getPartitionValues()); - writeLong(entryBlockBuilder, entryType, fieldId++, "size", addFileEntry.getSize()); - writeLong(entryBlockBuilder, entryType, fieldId++, "modificationTime", addFileEntry.getModificationTime()); - writeBoolean(entryBlockBuilder, entryType, fieldId++, "dataChange", addFileEntry.isDataChange()); - if (writeStatsAsJson) { - writeJsonStats(entryBlockBuilder, entryType, addFileEntry, metadataEntry, fieldId++); - } - if (writeStatsAsStruct) { - writeParsedStats(entryBlockBuilder, entryType, addFileEntry, fieldId++); - } - writeStringMap(entryBlockBuilder, entryType, fieldId++, "tags", addFileEntry.getTags()); - blockBuilder.closeEntry(); + RowBlockBuilder blockBuilder = (RowBlockBuilder) pageBuilder.getBlockBuilder(ADD_BLOCK_CHANNEL); + blockBuilder.buildEntry(fieldBuilders -> { + int fieldId = 0; + writeString(fieldBuilders.get(fieldId), entryType, fieldId, "path", addFileEntry.getPath()); + fieldId++; + + writeStringMap(fieldBuilders.get(fieldId), entryType, fieldId, "partitionValues", addFileEntry.getPartitionValues()); + fieldId++; + + writeLong(fieldBuilders.get(fieldId), entryType, fieldId, "size", addFileEntry.getSize()); + fieldId++; + + writeLong(fieldBuilders.get(fieldId), entryType, fieldId, "modificationTime", addFileEntry.getModificationTime()); + fieldId++; + + writeBoolean(fieldBuilders.get(fieldId), entryType, fieldId, "dataChange", addFileEntry.isDataChange()); + fieldId++; + + if (writeStatsAsJson) { + writeJsonStats(fieldBuilders.get(fieldId), entryType, addFileEntry, metadataEntry, fieldId); + fieldId++; + } + + if (writeStatsAsStruct) { + writeParsedStats(fieldBuilders.get(fieldId), entryType, addFileEntry, fieldId); + fieldId++; + } + + writeStringMap(fieldBuilders.get(fieldId), entryType, fieldId, "tags", addFileEntry.getTags()); + }); // null for others appendNullOtherBlocks(pageBuilder, ADD_BLOCK_CHANNEL); @@ -278,26 +299,30 @@ private void writeParsedStats(BlockBuilder entryBlockBuilder, RowType entryType, return; } DeltaLakeFileStatistics stats = addFileEntry.getStats().get(); - BlockBuilder statsBlockBuilder = entryBlockBuilder.beginBlockEntry(); - - if (stats instanceof DeltaLakeParquetFileStatistics) { - writeLong(statsBlockBuilder, statsType, 0, "numRecords", stats.getNumRecords().orElse(null)); - writeMinMaxMapAsFields(statsBlockBuilder, statsType, 1, "minValues", stats.getMinValues(), false); - writeMinMaxMapAsFields(statsBlockBuilder, statsType, 2, "maxValues", stats.getMaxValues(), false); - writeNullCountAsFields(statsBlockBuilder, statsType, 3, "nullCount", stats.getNullCount()); - } - else { - int internalFieldId = 0; - writeLong(statsBlockBuilder, statsType, internalFieldId++, "numRecords", stats.getNumRecords().orElse(null)); - if (statsType.getFields().stream().anyMatch(field -> field.getName().orElseThrow().equals("minValues"))) { - writeMinMaxMapAsFields(statsBlockBuilder, statsType, internalFieldId++, "minValues", stats.getMinValues(), true); + ((RowBlockBuilder) entryBlockBuilder).buildEntry(fieldBuilders -> { + if (stats instanceof DeltaLakeParquetFileStatistics) { + writeLong(fieldBuilders.get(0), statsType, 0, "numRecords", stats.getNumRecords().orElse(null)); + writeMinMaxMapAsFields(fieldBuilders.get(1), statsType, 1, "minValues", stats.getMinValues(), false); + writeMinMaxMapAsFields(fieldBuilders.get(2), statsType, 2, "maxValues", stats.getMaxValues(), false); + writeNullCountAsFields(fieldBuilders.get(3), statsType, 3, "nullCount", stats.getNullCount()); } - if (statsType.getFields().stream().anyMatch(field -> field.getName().orElseThrow().equals("maxValues"))) { - writeMinMaxMapAsFields(statsBlockBuilder, statsType, internalFieldId++, "maxValues", stats.getMaxValues(), true); + else { + int internalFieldId = 0; + + writeLong(fieldBuilders.get(internalFieldId), statsType, internalFieldId, "numRecords", stats.getNumRecords().orElse(null)); + internalFieldId++; + + if (statsType.getFields().stream().anyMatch(field -> field.getName().orElseThrow().equals("minValues"))) { + writeMinMaxMapAsFields(fieldBuilders.get(internalFieldId), statsType, internalFieldId, "minValues", stats.getMinValues(), true); + internalFieldId++; + } + if (statsType.getFields().stream().anyMatch(field -> field.getName().orElseThrow().equals("maxValues"))) { + writeMinMaxMapAsFields(fieldBuilders.get(internalFieldId), statsType, internalFieldId, "maxValues", stats.getMaxValues(), true); + internalFieldId++; + } + writeNullCountAsFields(fieldBuilders.get(internalFieldId), statsType, internalFieldId, "nullCount", stats.getNullCount()); } - writeNullCountAsFields(statsBlockBuilder, statsType, internalFieldId++, "nullCount", stats.getNullCount()); - } - entryBlockBuilder.closeEntry(); + }); } private void writeMinMaxMapAsFields(BlockBuilder blockBuilder, RowType type, int fieldId, String fieldName, Optional> values, boolean isJson) @@ -314,17 +339,21 @@ private void writeNullCountAsFields(BlockBuilder blockBuilder, RowType type, int private void writeObjectMapAsFields(BlockBuilder blockBuilder, RowType type, int fieldId, String fieldName, Optional> values) { - RowType.Field valuesField = validateAndGetField(type, fieldId, fieldName); - RowType valuesFieldType = (RowType) valuesField.getType(); if (values.isEmpty()) { blockBuilder.appendNull(); + return; } - else { - BlockBuilder fieldBlockBuilder = blockBuilder.beginBlockEntry(); - for (RowType.Field valueField : valuesFieldType.getFields()) { + + Field valuesField = validateAndGetField(type, fieldId, fieldName); + List fields = ((RowType) valuesField.getType()).getFields(); + ((RowBlockBuilder) blockBuilder).buildEntry(fieldBuilders -> { + for (int i = 0; i < fields.size(); i++) { + Field field = fields.get(i); + BlockBuilder fieldBlockBuilder = fieldBuilders.get(i); + // anonymous row fields are not expected here - Object value = values.get().get(valueField.getName().orElseThrow()); - if (valueField.getType() instanceof RowType) { + Object value = values.get().get(field.getName().orElseThrow()); + if (field.getType() instanceof RowType) { Block rowBlock = (Block) value; // Statistics were not collected if (rowBlock == null) { @@ -336,15 +365,14 @@ private void writeObjectMapAsFields(BlockBuilder blockBuilder, RowType type, int fieldBlockBuilder.appendNull(); } else { - valueField.getType().appendTo(rowBlock, 0, fieldBlockBuilder); + field.getType().appendTo(rowBlock, 0, fieldBlockBuilder); } } else { - writeNativeValue(valueField.getType(), fieldBlockBuilder, value); + writeNativeValue(field.getType(), fieldBlockBuilder, value); } } - blockBuilder.closeEntry(); - } + }); } private Optional> preprocessMinMaxValues(RowType valuesType, Optional> valuesOptional, boolean isJson) @@ -394,12 +422,11 @@ private Optional> preprocessNullCount(Optional { + writeString(fieldBuilders.get(0), entryType, 0, "path", removeFileEntry.getPath()); + writeLong(fieldBuilders.get(1), entryType, 1, "deletionTimestamp", removeFileEntry.getDeletionTimestamp()); + writeBoolean(fieldBuilders.get(2), entryType, 2, "dataChange", removeFileEntry.isDataChange()); + }); // null for others appendNullOtherBlocks(pageBuilder, REMOVE_BLOCK_CHANNEL); @@ -450,17 +477,17 @@ private void writeStringMap(BlockBuilder blockBuilder, RowType type, int fieldId return; } MapType mapType = (MapType) field.getType(); - BlockBuilder mapBuilder = blockBuilder.beginBlockEntry(); - for (Map.Entry entry : values.entrySet()) { - mapType.getKeyType().writeSlice(mapBuilder, utf8Slice(entry.getKey())); - if (entry.getValue() == null) { - mapBuilder.appendNull(); - } - else { - mapType.getValueType().writeSlice(mapBuilder, utf8Slice(entry.getValue())); + ((MapBlockBuilder) blockBuilder).buildEntry((keyBlockBuilder, valueBlockBuilder) -> { + for (Map.Entry entry : values.entrySet()) { + mapType.getKeyType().writeSlice(keyBlockBuilder, utf8Slice(entry.getKey())); + if (entry.getValue() == null) { + valueBlockBuilder.appendNull(); + } + else { + mapType.getValueType().writeSlice(valueBlockBuilder, utf8Slice(entry.getValue())); + } } - } - blockBuilder.closeEntry(); + }); } private void writeStringList(BlockBuilder blockBuilder, RowType type, int fieldId, String fieldName, @Nullable List values) @@ -472,16 +499,16 @@ private void writeStringList(BlockBuilder blockBuilder, RowType type, int fieldI return; } ArrayType arrayType = (ArrayType) field.getType(); - BlockBuilder mapBuilder = blockBuilder.beginBlockEntry(); - for (String value : values) { - if (value == null) { - mapBuilder.appendNull(); - } - else { - arrayType.getElementType().writeSlice(mapBuilder, utf8Slice(value)); + ((ArrayBlockBuilder) blockBuilder).buildEntry(elementBuilder -> { + for (String value : values) { + if (value == null) { + elementBuilder.appendNull(); + } + else { + arrayType.getElementType().writeSlice(elementBuilder, utf8Slice(value)); + } } - } - blockBuilder.closeEntry(); + }); } private RowType getInternalRowType(RowType type, int fieldId, String fieldName) diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/util/PageListBuilder.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/util/PageListBuilder.java index 1b520da08414..523d72040879 100644 --- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/util/PageListBuilder.java +++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/util/PageListBuilder.java @@ -17,6 +17,7 @@ import io.trino.spi.Page; import io.trino.spi.PageBuilder; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.ConnectorTableMetadata; import io.trino.spi.type.TimeZoneKey; @@ -109,23 +110,21 @@ public void appendVarchar(String value) public void appendVarcharVarcharMap(Map values) { - BlockBuilder column = nextColumn(); - BlockBuilder map = column.beginBlockEntry(); - values.forEach((key, value) -> { + MapBlockBuilder column = (MapBlockBuilder) nextColumn(); + column.buildEntry((keyBuilder, valueBuilder) -> values.forEach((key, value) -> { if (key == null) { - map.appendNull(); + keyBuilder.appendNull(); } else { - VARCHAR.writeString(map, key); + VARCHAR.writeString(keyBuilder, key); } if (value == null) { - map.appendNull(); + valueBuilder.appendNull(); } else { - VARCHAR.writeString(map, value); + VARCHAR.writeString(valueBuilder, value); } - }); - column.closeEntry(); + })); } public BlockBuilder nextColumn() diff --git a/plugin/trino-elasticsearch/src/main/java/io/trino/plugin/elasticsearch/decoders/ArrayDecoder.java b/plugin/trino-elasticsearch/src/main/java/io/trino/plugin/elasticsearch/decoders/ArrayDecoder.java index acb5453153a8..2c6bd89b500c 100644 --- a/plugin/trino-elasticsearch/src/main/java/io/trino/plugin/elasticsearch/decoders/ArrayDecoder.java +++ b/plugin/trino-elasticsearch/src/main/java/io/trino/plugin/elasticsearch/decoders/ArrayDecoder.java @@ -16,6 +16,7 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; import io.trino.plugin.elasticsearch.DecoderDescriptor; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.BlockBuilder; import org.elasticsearch.search.SearchHit; @@ -40,15 +41,11 @@ public void decode(SearchHit hit, Supplier getter, BlockBuilder output) if (data == null) { output.appendNull(); } - else if (data instanceof List) { - BlockBuilder array = output.beginBlockEntry(); - ((List) data).forEach(element -> elementDecoder.decode(hit, () -> element, array)); - output.closeEntry(); + else if (data instanceof List list) { + ((ArrayBlockBuilder) output).buildEntry(elementBuilder -> list.forEach(element -> elementDecoder.decode(hit, () -> element, elementBuilder))); } else { - BlockBuilder array = output.beginBlockEntry(); - elementDecoder.decode(hit, () -> data, array); - output.closeEntry(); + ((ArrayBlockBuilder) output).buildEntry(elementBuilder -> elementDecoder.decode(hit, () -> data, elementBuilder)); } } diff --git a/plugin/trino-elasticsearch/src/main/java/io/trino/plugin/elasticsearch/decoders/RowDecoder.java b/plugin/trino-elasticsearch/src/main/java/io/trino/plugin/elasticsearch/decoders/RowDecoder.java index 49616e47ca07..abcaf7b3e75b 100644 --- a/plugin/trino-elasticsearch/src/main/java/io/trino/plugin/elasticsearch/decoders/RowDecoder.java +++ b/plugin/trino-elasticsearch/src/main/java/io/trino/plugin/elasticsearch/decoders/RowDecoder.java @@ -18,6 +18,7 @@ import io.trino.plugin.elasticsearch.DecoderDescriptor; import io.trino.spi.TrinoException; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import org.elasticsearch.search.SearchHit; import java.util.List; @@ -53,12 +54,12 @@ public void decode(SearchHit hit, Supplier getter, BlockBuilder output) output.appendNull(); } else if (data instanceof Map) { - BlockBuilder row = output.beginBlockEntry(); - for (int i = 0; i < decoders.size(); i++) { - String field = fieldNames.get(i); - decoders.get(i).decode(hit, () -> getField((Map) data, field), row); - } - output.closeEntry(); + ((RowBlockBuilder) output).buildEntry(fieldBuilders -> { + for (int i = 0; i < decoders.size(); i++) { + String field = fieldNames.get(i); + decoders.get(i).decode(hit, () -> getField((Map) data, field), fieldBuilders.get(i)); + } + }); } else { throw new TrinoException(TYPE_MISMATCH, format("Expected object for field '%s' of type ROW: %s [%s]", path, data, data.getClass().getSimpleName())); diff --git a/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/BingTileFunctions.java b/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/BingTileFunctions.java index 6dbc00ebe57a..fce7ab43d459 100644 --- a/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/BingTileFunctions.java +++ b/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/BingTileFunctions.java @@ -18,10 +18,10 @@ import com.esri.core.geometry.ogc.OGCGeometry; import com.google.common.collect.ImmutableList; import io.airlift.slice.Slice; -import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedRowValueBuilder; import io.trino.spi.function.Description; import io.trino.spi.function.ScalarFunction; import io.trino.spi.function.SqlType; @@ -108,28 +108,21 @@ public static final class BingTileCoordinatesFunction { private static final RowType BING_TILE_COORDINATES_ROW_TYPE = RowType.anonymous(ImmutableList.of(INTEGER, INTEGER)); - private final PageBuilder pageBuilder; + private final BufferedRowValueBuilder rowValueBuilder; public BingTileCoordinatesFunction() { - pageBuilder = new PageBuilder(ImmutableList.of(BING_TILE_COORDINATES_ROW_TYPE)); + rowValueBuilder = BufferedRowValueBuilder.createBuffered(BING_TILE_COORDINATES_ROW_TYPE); } @SqlType("row(x integer,y integer)") public Block bingTileCoordinates(@SqlType(BingTileType.NAME) long input) { - if (pageBuilder.isFull()) { - pageBuilder.reset(); - } - BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0); BingTile tile = BingTile.decode(input); - BlockBuilder tileBlockBuilder = blockBuilder.beginBlockEntry(); - INTEGER.writeLong(tileBlockBuilder, tile.getX()); - INTEGER.writeLong(tileBlockBuilder, tile.getY()); - blockBuilder.closeEntry(); - pageBuilder.declarePosition(); - - return BING_TILE_COORDINATES_ROW_TYPE.getObject(blockBuilder, blockBuilder.getPositionCount() - 1); + return rowValueBuilder.build(fields -> { + INTEGER.writeLong(fields.get(0), tile.getX()); + INTEGER.writeLong(fields.get(1), tile.getY()); + }); } } diff --git a/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/GeoFunctions.java b/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/GeoFunctions.java index 6e56501b2fd5..88eff53d9413 100644 --- a/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/GeoFunctions.java +++ b/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/GeoFunctions.java @@ -48,7 +48,6 @@ import io.trino.geospatial.serde.GeometrySerde; import io.trino.geospatial.serde.GeometrySerializationType; import io.trino.geospatial.serde.JtsGeometrySerde; -import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; @@ -112,6 +111,7 @@ import static io.trino.plugin.geospatial.GeometryType.GEOMETRY_TYPE_NAME; import static io.trino.plugin.geospatial.SphericalGeographyType.SPHERICAL_GEOGRAPHY_TYPE_NAME; import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.type.StandardTypes.BIGINT; import static io.trino.spi.type.StandardTypes.BOOLEAN; import static io.trino.spi.type.StandardTypes.DOUBLE; @@ -1201,18 +1201,13 @@ public static Block geometryNearestPoints(@SqlType(GEOMETRY_TYPE_NAME) Slice lef } RowType rowType = RowType.anonymous(ImmutableList.of(GEOMETRY, GEOMETRY)); - PageBuilder pageBuilder = new PageBuilder(ImmutableList.of(rowType)); GeometryFactory geometryFactory = leftGeometry.getFactory(); Coordinate[] nearestCoordinates = DistanceOp.nearestPoints(leftGeometry, rightGeometry); - BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0); - BlockBuilder entryBlockBuilder = blockBuilder.beginBlockEntry(); - GEOMETRY.writeSlice(entryBlockBuilder, JtsGeometrySerde.serialize(geometryFactory.createPoint(nearestCoordinates[0]))); - GEOMETRY.writeSlice(entryBlockBuilder, JtsGeometrySerde.serialize(geometryFactory.createPoint(nearestCoordinates[1]))); - blockBuilder.closeEntry(); - pageBuilder.declarePosition(); - - return rowType.getObject(blockBuilder, blockBuilder.getPositionCount() - 1); + return buildRowValue(rowType, fieldBuilders -> { + GEOMETRY.writeSlice(fieldBuilders.get(0), serialize(geometryFactory.createPoint(nearestCoordinates[0]))); + GEOMETRY.writeSlice(fieldBuilders.get(1), serialize(geometryFactory.createPoint(nearestCoordinates[1]))); + }); } @SqlNullable @@ -1487,7 +1482,7 @@ private static Block spatialPartitions(KdbTree kdbTree, Rectangle envelope) for (Map.Entry partition : partitions.entrySet()) { if (envelope.getXMin() < partition.getValue().getXMax() && envelope.getYMin() < partition.getValue().getYMax()) { BlockBuilder blockBuilder = IntegerType.INTEGER.createFixedSizeBlockBuilder(1); - blockBuilder.writeInt(partition.getKey()); + IntegerType.INTEGER.writeInt(blockBuilder, partition.getKey()); return blockBuilder.build(); } } @@ -1496,7 +1491,7 @@ private static Block spatialPartitions(KdbTree kdbTree, Rectangle envelope) BlockBuilder blockBuilder = IntegerType.INTEGER.createFixedSizeBlockBuilder(partitions.size()); for (int id : partitions.keySet()) { - blockBuilder.writeInt(id); + IntegerType.INTEGER.writeInt(blockBuilder, id); } return blockBuilder.build(); diff --git a/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/GeometryType.java b/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/GeometryType.java index b9f4129ec658..79ac70f7ea9a 100644 --- a/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/GeometryType.java +++ b/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/GeometryType.java @@ -16,6 +16,7 @@ import io.airlift.slice.Slice; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.TypeSignature; @@ -45,8 +46,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -73,7 +73,7 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int l blockBuilder.appendNull(); return; } - blockBuilder.writeBytes(value, offset, length).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length); } @Override diff --git a/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/KdbTreeType.java b/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/KdbTreeType.java index e548371ca3bd..ea61aaba10e3 100644 --- a/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/KdbTreeType.java +++ b/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/KdbTreeType.java @@ -18,6 +18,7 @@ import io.trino.geospatial.KdbTreeUtils; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.TypeSignature; @@ -51,8 +52,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -61,7 +61,7 @@ public void writeObject(BlockBuilder blockBuilder, Object value) { String json = KdbTreeUtils.toJson(((KdbTree) value)); Slice bytes = utf8Slice(json); - blockBuilder.writeBytes(bytes, 0, bytes.length()).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(bytes); } @Override diff --git a/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/SphericalGeographyType.java b/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/SphericalGeographyType.java index 1dc207ca48ec..b63769cd6647 100644 --- a/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/SphericalGeographyType.java +++ b/plugin/trino-geospatial/src/main/java/io/trino/plugin/geospatial/SphericalGeographyType.java @@ -16,6 +16,7 @@ import io.airlift.slice.Slice; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.TypeSignature; @@ -40,8 +41,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -60,7 +60,7 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value) @Override public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length) { - blockBuilder.writeBytes(value, offset, length).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length); } @Override diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSink.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSink.java index a56df88554be..26e84921435f 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSink.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HivePageSink.java @@ -440,7 +440,7 @@ private Block buildBucketBlock(Page page) Page bucketColumnsPage = extractColumns(page, bucketColumns); for (int position = 0; position < page.getPositionCount(); position++) { int bucket = bucketFunction.getBucket(bucketColumnsPage, position); - bucketColumnBuilder.writeInt(bucket); + INTEGER.writeInt(bucketColumnBuilder, bucket); } return bucketColumnBuilder.build(); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/SerDeUtils.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/SerDeUtils.java index 1000cebca302..9f0f5744a633 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/SerDeUtils.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/SerDeUtils.java @@ -16,10 +16,16 @@ import com.google.common.annotations.VisibleForTesting; import io.airlift.slice.Slices; import io.trino.plugin.base.type.DecodedTimestamp; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; +import io.trino.spi.type.ArrayType; import io.trino.spi.type.CharType; import io.trino.spi.type.DecimalType; +import io.trino.spi.type.MapType; +import io.trino.spi.type.RowType; import io.trino.spi.type.TimestampType; import io.trino.spi.type.Type; import org.apache.hadoop.hive.common.type.HiveChar; @@ -53,9 +59,13 @@ import java.util.List; import java.util.Map; +import java.util.Map.Entry; import static com.google.common.base.Preconditions.checkArgument; import static io.trino.plugin.base.type.TrinoTimestampEncoderFactory.createTimestampEncoder; +import static io.trino.spi.block.ArrayValueBuilder.buildArrayValue; +import static io.trino.spi.block.MapValueBuilder.buildMapValue; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.type.Chars.truncateToLengthAndTrimSpaces; import static io.trino.spi.type.Timestamps.NANOSECONDS_PER_SECOND; import static io.trino.spi.type.Timestamps.round; @@ -90,7 +100,7 @@ public static Block serializeObject(Type type, BlockBuilder builder, Object obje case LIST: return serializeList(type, builder, object, (ListObjectInspector) inspector); case MAP: - return serializeMap(type, builder, object, (MapObjectInspector) inspector, filterNullMapKeys); + return serializeMap((MapType) type, (MapBlockBuilder) builder, object, (MapObjectInspector) inspector, filterNullMapKeys); case STRUCT: return serializeStruct(type, builder, object, (StructObjectInspector) inspector); case UNION: @@ -179,31 +189,23 @@ private static Block serializeList(Type type, BlockBuilder builder, Object objec return null; } - List typeParameters = type.getTypeParameters(); - checkArgument(typeParameters.size() == 1, "list must have exactly 1 type parameter"); - Type elementType = typeParameters.get(0); + ArrayType arrayType = (ArrayType) type; ObjectInspector elementInspector = inspector.getListElementObjectInspector(); - BlockBuilder currentBuilder; - if (builder != null) { - currentBuilder = builder.beginBlockEntry(); - } - else { - currentBuilder = elementType.createBlockBuilder(null, list.size()); + if (builder == null) { + return buildArrayValue(arrayType, list.size(), valuesBuilder -> buildList(list, arrayType.getElementType(), elementInspector, valuesBuilder)); } + ((ArrayBlockBuilder) builder).buildEntry(elementBuilder -> buildList(list, arrayType.getElementType(), elementInspector, elementBuilder)); + return null; + } + private static void buildList(List list, Type elementType, ObjectInspector elementInspector, BlockBuilder valueBuilder) + { for (Object element : list) { - serializeObject(elementType, currentBuilder, element, elementInspector); - } - - if (builder != null) { - builder.closeEntry(); - return null; + serializeObject(elementType, valueBuilder, element, elementInspector); } - Block resultBlock = currentBuilder.build(); - return resultBlock; } - private static Block serializeMap(Type type, BlockBuilder builder, Object object, MapObjectInspector inspector, boolean filterNullMapKeys) + private static Block serializeMap(MapType mapType, MapBlockBuilder builder, Object object, MapObjectInspector inspector, boolean filterNullMapKeys) { Map map = inspector.getMap(object); if (map == null) { @@ -211,34 +213,31 @@ private static Block serializeMap(Type type, BlockBuilder builder, Object object return null; } - List typeParameters = type.getTypeParameters(); - checkArgument(typeParameters.size() == 2, "map must have exactly 2 type parameter"); - Type keyType = typeParameters.get(0); - Type valueType = typeParameters.get(1); - ObjectInspector keyInspector = inspector.getMapKeyObjectInspector(); - ObjectInspector valueInspector = inspector.getMapValueObjectInspector(); - BlockBuilder currentBuilder; - - boolean builderSynthesized = false; if (builder == null) { - builderSynthesized = true; - builder = type.createBlockBuilder(null, 1); + return buildMapValue( + mapType, + map.size(), + (keyBuilder, valueBuilder) -> buildMap(mapType, keyBuilder, valueBuilder, map, inspector, filterNullMapKeys)); } - currentBuilder = builder.beginBlockEntry(); - for (Map.Entry entry : map.entrySet()) { + builder.buildEntry((keyBuilder, valueBuilder) -> buildMap(mapType, keyBuilder, valueBuilder, map, inspector, filterNullMapKeys)); + return null; + } + + private static void buildMap(MapType mapType, BlockBuilder keyBuilder, BlockBuilder valueBuilder, Map map, MapObjectInspector inspector, boolean filterNullMapKeys) + { + Type keyType = mapType.getKeyType(); + Type valueType = mapType.getValueType(); + ObjectInspector keyInspector = inspector.getMapKeyObjectInspector(); + ObjectInspector valueInspector = inspector.getMapValueObjectInspector(); + + for (Entry entry : map.entrySet()) { // Hive skips map entries with null keys if (!filterNullMapKeys || entry.getKey() != null) { - serializeObject(keyType, currentBuilder, entry.getKey(), keyInspector); - serializeObject(valueType, currentBuilder, entry.getValue(), valueInspector); + serializeObject(keyType, keyBuilder, entry.getKey(), keyInspector); + serializeObject(valueType, valueBuilder, entry.getValue(), valueInspector); } } - - builder.closeEntry(); - if (builderSynthesized) { - return (Block) type.getObject(builder, 0); - } - return null; } private static Block serializeStruct(Type type, BlockBuilder builder, Object object, StructObjectInspector inspector) @@ -248,28 +247,24 @@ private static Block serializeStruct(Type type, BlockBuilder builder, Object obj return null; } - List typeParameters = type.getTypeParameters(); - List allStructFieldRefs = inspector.getAllStructFieldRefs(); - checkArgument(typeParameters.size() == allStructFieldRefs.size()); - BlockBuilder currentBuilder; - - boolean builderSynthesized = false; + RowType rowType = (RowType) type; if (builder == null) { - builderSynthesized = true; - builder = type.createBlockBuilder(null, 1); + return buildRowValue(rowType, fieldBuilders -> buildStruct(rowType, object, inspector, fieldBuilders)); } - currentBuilder = builder.beginBlockEntry(); + ((RowBlockBuilder) builder).buildEntry(fieldBuilders -> buildStruct(rowType, object, inspector, fieldBuilders)); + return null; + } + + private static void buildStruct(RowType type, Object object, StructObjectInspector inspector, List fieldBuilders) + { + List typeParameters = type.getTypeParameters(); + List allStructFieldRefs = inspector.getAllStructFieldRefs(); + checkArgument(typeParameters.size() == allStructFieldRefs.size()); for (int i = 0; i < typeParameters.size(); i++) { StructField field = allStructFieldRefs.get(i); - serializeObject(typeParameters.get(i), currentBuilder, inspector.getStructFieldData(object, field), field.getFieldObjectInspector()); + serializeObject(typeParameters.get(i), fieldBuilders.get(i), inspector.getStructFieldData(object, field), field.getFieldObjectInspector()); } - - builder.closeEntry(); - if (builderSynthesized) { - return (Block) type.getObject(builder, 0); - } - return null; } // Use row blocks to represent union objects when reading @@ -280,32 +275,28 @@ private static Block serializeUnion(Type type, BlockBuilder builder, Object obje return null; } - boolean builderSynthesized = false; + RowType rowType = (RowType) type; if (builder == null) { - builderSynthesized = true; - builder = type.createBlockBuilder(null, 1); + return buildRowValue(rowType, fieldBuilders -> buildUnion(rowType, object, inspector, fieldBuilders)); } + ((RowBlockBuilder) builder).buildEntry(fieldBuilders -> buildUnion(rowType, object, inspector, fieldBuilders)); + return null; + } - BlockBuilder currentBuilder = builder.beginBlockEntry(); - + private static void buildUnion(RowType rowType, Object object, UnionObjectInspector inspector, List fieldBuilders) + { byte tag = inspector.getTag(object); - TINYINT.writeLong(currentBuilder, tag); + TINYINT.writeLong(fieldBuilders.get(0), tag); - List typeParameters = type.getTypeParameters(); + List typeParameters = rowType.getTypeParameters(); for (int i = 1; i < typeParameters.size(); i++) { if (i == tag + 1) { - serializeObject(typeParameters.get(i), currentBuilder, inspector.getField(object), inspector.getObjectInspectors().get(tag)); + serializeObject(typeParameters.get(i), fieldBuilders.get(i), inspector.getField(object), inspector.getObjectInspectors().get(tag)); } else { - currentBuilder.appendNull(); + fieldBuilders.get(i).appendNull(); } } - - builder.closeEntry(); - if (builderSynthesized) { - return (Block) type.getObject(builder, 0); - } - return null; } @SuppressWarnings("deprecation") diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveTestUtils.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveTestUtils.java index 5a419fd6d119..7589b912994d 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveTestUtils.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveTestUtils.java @@ -58,7 +58,6 @@ import io.trino.plugin.hive.s3select.TrinoS3ClientFactory; import io.trino.spi.PageSorter; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.ArrayType; @@ -87,6 +86,7 @@ import java.lang.invoke.MethodHandle; import java.math.BigDecimal; import java.nio.ByteBuffer; +import java.util.Collection; import java.util.List; import java.util.Map; import java.util.Optional; @@ -95,6 +95,9 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static io.airlift.units.DataSize.Unit.MEGABYTE; +import static io.trino.spi.block.ArrayValueBuilder.buildArrayValue; +import static io.trino.spi.block.MapValueBuilder.buildMapValue; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.NULL_FLAG; import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL; import static io.trino.spi.function.InvocationConvention.simpleConvention; @@ -311,35 +314,34 @@ public static Object toNativeContainerValue(Type type, Object hiveValue) return null; } - if (type instanceof ArrayType) { - BlockBuilder blockBuilder = type.createBlockBuilder(null, 1); - BlockBuilder subBlockBuilder = blockBuilder.beginBlockEntry(); - for (Object subElement : (Iterable) hiveValue) { - appendToBlockBuilder(type.getTypeParameters().get(0), subElement, subBlockBuilder); - } - blockBuilder.closeEntry(); - return type.getObject(blockBuilder, 0); + if (type instanceof ArrayType arrayType) { + Collection hiveArray = (Collection) hiveValue; + return buildArrayValue(arrayType, hiveArray.size(), valueBuilder -> { + for (Object subElement : hiveArray) { + appendToBlockBuilder(type.getTypeParameters().get(0), subElement, valueBuilder); + } + }); } - if (type instanceof RowType) { - BlockBuilder blockBuilder = type.createBlockBuilder(null, 1); - BlockBuilder subBlockBuilder = blockBuilder.beginBlockEntry(); - int field = 0; - for (Object subElement : (Iterable) hiveValue) { - appendToBlockBuilder(type.getTypeParameters().get(field), subElement, subBlockBuilder); - field++; - } - blockBuilder.closeEntry(); - return type.getObject(blockBuilder, 0); + if (type instanceof RowType rowType) { + return buildRowValue(rowType, fields -> { + int fieldIndex = 0; + for (Object subElement : (Iterable) hiveValue) { + appendToBlockBuilder(type.getTypeParameters().get(fieldIndex), subElement, fields.get(fieldIndex)); + fieldIndex++; + } + }); } - if (type instanceof MapType) { - BlockBuilder blockBuilder = type.createBlockBuilder(null, 1); - BlockBuilder subBlockBuilder = blockBuilder.beginBlockEntry(); - for (Map.Entry entry : ((Map) hiveValue).entrySet()) { - appendToBlockBuilder(type.getTypeParameters().get(0), entry.getKey(), subBlockBuilder); - appendToBlockBuilder(type.getTypeParameters().get(1), entry.getValue(), subBlockBuilder); - } - blockBuilder.closeEntry(); - return type.getObject(blockBuilder, 0); + if (type instanceof MapType mapType) { + Map hiveMap = (Map) hiveValue; + return buildMapValue( + mapType, + hiveMap.size(), + (keyBuilder, valueBuilder) -> { + hiveMap.forEach((key, value) -> { + appendToBlockBuilder(mapType.getKeyType(), key, keyBuilder); + appendToBlockBuilder(mapType.getValueType(), value, valueBuilder); + }); + }); } if (type instanceof BooleanType) { return hiveValue; diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHivePartitionedBucketFunction.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHivePartitionedBucketFunction.java index 425021644b53..05140c7550df 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHivePartitionedBucketFunction.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHivePartitionedBucketFunction.java @@ -138,8 +138,8 @@ public void testConsecutiveBucketsWithinPartition(BucketingVersion hiveBucketing BlockBuilder bucketColumn = BIGINT.createFixedSizeBlockBuilder(10); BlockBuilder partitionColumn = BIGINT.createFixedSizeBlockBuilder(10); for (int i = 0; i < 100; ++i) { - bucketColumn.writeLong(i); - partitionColumn.writeLong(42); + BIGINT.writeLong(bucketColumn, i); + BIGINT.writeLong(partitionColumn, 42); } Page page = new Page(bucketColumn, partitionColumn); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestReaderProjectionsAdapter.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestReaderProjectionsAdapter.java index b39e347310c1..819ff910d103 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestReaderProjectionsAdapter.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestReaderProjectionsAdapter.java @@ -211,7 +211,7 @@ private static Block createLongArrayBlock(List data) builder.appendNull(); } else { - builder.writeLong(value); + BIGINT.writeLong(builder, value); } } return builder.build(); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/BenchmarkHiveFileFormat.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/BenchmarkHiveFileFormat.java index aa47bde110cb..4670320d0e55 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/BenchmarkHiveFileFormat.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/benchmark/BenchmarkHiveFileFormat.java @@ -21,10 +21,13 @@ import io.trino.plugin.hive.parquet.ParquetReaderConfig; import io.trino.spi.Page; import io.trino.spi.PageBuilder; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.ArrayType; +import io.trino.spi.type.MapType; import io.trino.spi.type.Type; import io.trino.tpch.OrderColumn; import it.unimi.dsi.fastutil.ints.IntArrays; @@ -55,7 +58,6 @@ import static io.trino.jmh.Benchmarks.benchmark; import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT; import static io.trino.plugin.hive.HiveTestUtils.getHiveSession; -import static io.trino.plugin.hive.HiveTestUtils.mapType; import static io.trino.plugin.hive.benchmark.BenchmarkFileFormat.TRINO_OPTIMIZED_PARQUET; import static io.trino.plugin.hive.benchmark.BenchmarkFileFormatsUtils.MIN_DATA_SIZE; import static io.trino.plugin.hive.benchmark.BenchmarkFileFormatsUtils.createTempDir; @@ -64,9 +66,11 @@ import static io.trino.plugin.hive.benchmark.BenchmarkFileFormatsUtils.printResults; import static io.trino.spi.type.DoubleType.DOUBLE; import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.VarcharType.VARCHAR; import static io.trino.spi.type.VarcharType.createUnboundedVarcharType; import static io.trino.tpch.TpchTable.LINE_ITEM; import static io.trino.tpch.TpchTable.ORDERS; +import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; @State(Scope.Thread) @OutputTimeUnit(TimeUnit.SECONDS) @@ -263,7 +267,7 @@ public TestData createTestData(FileFormat format) @Override public TestData createTestData(FileFormat format) { - Type type = mapType(createUnboundedVarcharType(), DOUBLE); + MapType type = new MapType(VARCHAR, DOUBLE, TESTING_TYPE_MANAGER.getTypeOperators()); Random random = new Random(1234); PageBuilder pageBuilder = new PageBuilder(ImmutableList.of(type)); @@ -275,15 +279,15 @@ public TestData createTestData(FileFormat format) while (dataSize < MIN_DATA_SIZE) { pageBuilder.declarePosition(); - BlockBuilder builder = pageBuilder.getBlockBuilder(0); - BlockBuilder mapBuilder = builder.beginBlockEntry(); - int entries = nextRandomBetween(random, MIN_ENTRIES, MAX_ENTRIES); - IntArrays.shuffle(keys, random); - for (int entryId = 0; entryId < entries; entryId++) { - createUnboundedVarcharType().writeSlice(mapBuilder, Slices.utf8Slice("key" + keys[entryId])); - DOUBLE.writeDouble(mapBuilder, random.nextDouble()); - } - builder.closeEntry(); + MapBlockBuilder builder = (MapBlockBuilder) pageBuilder.getBlockBuilder(0); + builder.buildEntry((keyBuilder, valueBuilder) -> { + int entries = nextRandomBetween(random, MIN_ENTRIES, MAX_ENTRIES); + IntArrays.shuffle(keys, random); + for (int entryId = 0; entryId < entries; entryId++) { + VARCHAR.writeSlice(keyBuilder, Slices.utf8Slice("key" + keys[entryId])); + DOUBLE.writeDouble(valueBuilder, random.nextDouble()); + } + }); if (pageBuilder.isFull()) { Page page = pageBuilder.build(); @@ -302,7 +306,7 @@ public TestData createTestData(FileFormat format) @Override public TestData createTestData(FileFormat format) { - Type type = mapType(createUnboundedVarcharType(), DOUBLE); + MapType type = new MapType(VARCHAR, DOUBLE, TESTING_TYPE_MANAGER.getTypeOperators()); Random random = new Random(1234); PageBuilder pageBuilder = new PageBuilder(ImmutableList.of(type)); @@ -311,14 +315,14 @@ public TestData createTestData(FileFormat format) while (dataSize < MIN_DATA_SIZE) { pageBuilder.declarePosition(); - BlockBuilder builder = pageBuilder.getBlockBuilder(0); - BlockBuilder mapBuilder = builder.beginBlockEntry(); - int entries = nextRandomBetween(random, MIN_ENTRIES, MAX_ENTRIES); - for (int entryId = 0; entryId < entries; entryId++) { - createUnboundedVarcharType().writeSlice(mapBuilder, Slices.utf8Slice("key" + random.nextInt(10_000_000))); - DOUBLE.writeDouble(mapBuilder, random.nextDouble()); - } - builder.closeEntry(); + MapBlockBuilder builder = (MapBlockBuilder) pageBuilder.getBlockBuilder(0); + builder.buildEntry((keyBuilder, valueBuilder) -> { + int entries = nextRandomBetween(random, MIN_ENTRIES, MAX_ENTRIES); + for (int entryId = 0; entryId < entries; entryId++) { + VARCHAR.writeSlice(keyBuilder, Slices.utf8Slice("key" + random.nextInt(10_000_000))); + DOUBLE.writeDouble(valueBuilder, random.nextDouble()); + } + }); if (pageBuilder.isFull()) { Page page = pageBuilder.build(); @@ -337,7 +341,7 @@ public TestData createTestData(FileFormat format) @Override public TestData createTestData(FileFormat format) { - Type type = mapType(INTEGER, DOUBLE); + MapType type = new MapType(INTEGER, DOUBLE, TESTING_TYPE_MANAGER.getTypeOperators()); Random random = new Random(1234); PageBuilder pageBuilder = new PageBuilder(ImmutableList.of(type)); @@ -349,15 +353,15 @@ public TestData createTestData(FileFormat format) while (dataSize < MIN_DATA_SIZE) { pageBuilder.declarePosition(); - BlockBuilder builder = pageBuilder.getBlockBuilder(0); - BlockBuilder mapBuilder = builder.beginBlockEntry(); - int entries = nextRandomBetween(random, MIN_ENTRIES, MAX_ENTRIES); - IntArrays.shuffle(keys, random); - for (int entryId = 0; entryId < entries; entryId++) { - INTEGER.writeLong(mapBuilder, keys[entryId]); - DOUBLE.writeDouble(mapBuilder, random.nextDouble()); - } - builder.closeEntry(); + MapBlockBuilder builder = (MapBlockBuilder) pageBuilder.getBlockBuilder(0); + builder.buildEntry((keyBuilder, valueBuilder) -> { + int entries = nextRandomBetween(random, MIN_ENTRIES, MAX_ENTRIES); + IntArrays.shuffle(keys, random); + for (int entryId = 0; entryId < entries; entryId++) { + INTEGER.writeLong(keyBuilder, keys[entryId]); + DOUBLE.writeDouble(valueBuilder, random.nextDouble()); + } + }); if (pageBuilder.isFull()) { Page page = pageBuilder.build(); @@ -376,7 +380,7 @@ public TestData createTestData(FileFormat format) @Override public TestData createTestData(FileFormat format) { - Type type = mapType(INTEGER, DOUBLE); + MapType type = new MapType(INTEGER, DOUBLE, TESTING_TYPE_MANAGER.getTypeOperators()); Random random = new Random(1234); PageBuilder pageBuilder = new PageBuilder(ImmutableList.of(type)); @@ -385,14 +389,14 @@ public TestData createTestData(FileFormat format) while (dataSize < MIN_DATA_SIZE) { pageBuilder.declarePosition(); - BlockBuilder builder = pageBuilder.getBlockBuilder(0); - BlockBuilder mapBuilder = builder.beginBlockEntry(); - int entries = nextRandomBetween(random, MIN_ENTRIES, MAX_ENTRIES); - for (int entryId = 0; entryId < entries; entryId++) { - INTEGER.writeLong(mapBuilder, random.nextInt(10_000_000)); - DOUBLE.writeDouble(mapBuilder, random.nextDouble()); - } - builder.closeEntry(); + MapBlockBuilder builder = (MapBlockBuilder) pageBuilder.getBlockBuilder(0); + builder.buildEntry((keyBuilder, valueBuilder) -> { + int entries = nextRandomBetween(random, MIN_ENTRIES, MAX_ENTRIES); + for (int entryId = 0; entryId < entries; entryId++) { + INTEGER.writeLong(keyBuilder, random.nextInt(10_000_000)); + DOUBLE.writeDouble(valueBuilder, random.nextDouble()); + } + }); if (pageBuilder.isFull()) { Page page = pageBuilder.build(); @@ -421,12 +425,12 @@ public TestData createTestData(FileFormat format) pageBuilder.declarePosition(); BlockBuilder builder = pageBuilder.getBlockBuilder(0); - BlockBuilder mapBuilder = builder.beginBlockEntry(); - int entries = nextRandomBetween(random, MIN_ENTRIES, MAX_ENTRIES); - for (int entryId = 0; entryId < entries; entryId++) { - createUnboundedVarcharType().writeSlice(mapBuilder, Slices.utf8Slice("key" + random.nextInt(10_000_000))); - } - builder.closeEntry(); + ((ArrayBlockBuilder) builder).buildEntry(elementBuilder -> { + int entries = nextRandomBetween(random, MIN_ENTRIES, MAX_ENTRIES); + for (int entryId = 0; entryId < entries; entryId++) { + createUnboundedVarcharType().writeSlice(elementBuilder, Slices.utf8Slice("key" + random.nextInt(10_000_000))); + } + }); if (pageBuilder.isFull()) { Page page = pageBuilder.build(); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java index 89da15287220..e699a85c0396 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java @@ -42,6 +42,7 @@ import io.trino.plugin.hive.metastore.glue.converter.GlueInputConverter; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; +import io.trino.spi.block.BlockBuilder; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.ConnectorMetadata; import io.trino.spi.connector.ConnectorOutputTableHandle; @@ -1416,7 +1417,9 @@ public void testGlueObjectsWithoutStorageDescriptor() private Block singleValueBlock(long value) { - return BigintType.BIGINT.createBlockBuilder(null, 1).writeLong(value).build(); + BlockBuilder blockBuilder = BIGINT.createBlockBuilder(null, 1); + BIGINT.writeLong(blockBuilder, value); + return blockBuilder.build(); } private void doGetPartitionsFilterTest( diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcDeletedRows.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcDeletedRows.java index 7e138bcff0fd..669d6c7d2e56 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcDeletedRows.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/orc/TestOrcDeletedRows.java @@ -50,12 +50,14 @@ public class TestOrcDeletedRows public void setUp() { partitionDirectory = Location.of(getResource("fullacid_delete_delta_test").toString()); - rowIdBlock = BIGINT.createFixedSizeBlockBuilder(1) - .writeLong(0) - .build(); - bucketBlock = INTEGER.createFixedSizeBlockBuilder(1) - .writeInt(536870912) - .build(); + + BlockBuilder rowIdBlockBuilder = BIGINT.createFixedSizeBlockBuilder(1); + BIGINT.writeLong(rowIdBlockBuilder, 0); + rowIdBlock = rowIdBlockBuilder.build(); + + BlockBuilder bucketBlockBuilder = INTEGER.createFixedSizeBlockBuilder(1); + INTEGER.writeInt(bucketBlockBuilder, 536870912); + bucketBlock = bucketBlockBuilder.build(); } @Test @@ -175,7 +177,7 @@ private Page createTestPage(int originalTransactionStart, int originalTransactio int size = originalTransactionEnd - originalTransactionStart; BlockBuilder originalTransaction = BIGINT.createFixedSizeBlockBuilder(size); for (long i = originalTransactionStart; i < originalTransactionEnd; i++) { - originalTransaction.writeLong(i); + BIGINT.writeLong(originalTransaction, i); } return new Page( diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/ParquetTester.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/ParquetTester.java index e58929d25cc1..5042ab9eb738 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/ParquetTester.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/parquet/ParquetTester.java @@ -43,8 +43,11 @@ import io.trino.spi.Page; import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.RecordCursor; @@ -872,32 +875,32 @@ else if (TIMESTAMP_NANOS.equals(type)) { if (type instanceof ArrayType) { List array = (List) value; Type elementType = type.getTypeParameters().get(0); - BlockBuilder arrayBlockBuilder = blockBuilder.beginBlockEntry(); - for (Object elementValue : array) { - writeValue(elementType, arrayBlockBuilder, elementValue); - } - blockBuilder.closeEntry(); + ((ArrayBlockBuilder) blockBuilder).buildEntry(elementBuilder -> { + for (Object elementValue : array) { + writeValue(elementType, elementBuilder, elementValue); + } + }); } else if (type instanceof MapType) { Map map = (Map) value; Type keyType = type.getTypeParameters().get(0); Type valueType = type.getTypeParameters().get(1); - BlockBuilder mapBlockBuilder = blockBuilder.beginBlockEntry(); - for (Map.Entry entry : map.entrySet()) { - writeValue(keyType, mapBlockBuilder, entry.getKey()); - writeValue(valueType, mapBlockBuilder, entry.getValue()); - } - blockBuilder.closeEntry(); + ((MapBlockBuilder) blockBuilder).buildEntry((keyBuilder, valueBuilder) -> { + for (Map.Entry entry : map.entrySet()) { + writeValue(keyType, keyBuilder, entry.getKey()); + writeValue(valueType, valueBuilder, entry.getValue()); + } + }); } else if (type instanceof RowType) { List array = (List) value; List fieldTypes = type.getTypeParameters(); - BlockBuilder rowBlockBuilder = blockBuilder.beginBlockEntry(); - for (int fieldId = 0; fieldId < fieldTypes.size(); fieldId++) { - Type fieldType = fieldTypes.get(fieldId); - writeValue(fieldType, rowBlockBuilder, array.get(fieldId)); - } - blockBuilder.closeEntry(); + ((RowBlockBuilder) blockBuilder).buildEntry(fieldBuilders -> { + for (int fieldId = 0; fieldId < fieldTypes.size(); fieldId++) { + Type fieldType = fieldTypes.get(fieldId); + writeValue(fieldType, fieldBuilders.get(fieldId), array.get(fieldId)); + } + }); } else { throw new IllegalArgumentException("Unsupported type " + type); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestSerDeUtils.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestSerDeUtils.java index d31c937e7703..4f700175f41e 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestSerDeUtils.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestSerDeUtils.java @@ -41,6 +41,7 @@ import java.util.List; import java.util.Map; import java.util.TreeMap; +import java.util.function.Consumer; import static io.airlift.slice.Slices.utf8Slice; import static io.trino.plugin.hive.HiveTestUtils.mapType; @@ -61,9 +62,6 @@ import static io.trino.testing.StructuralTestUtil.arrayBlockOf; import static io.trino.testing.StructuralTestUtil.mapBlockOf; import static io.trino.testing.StructuralTestUtil.rowBlockOf; -import static io.trino.type.DateTimes.MICROSECONDS_PER_MILLISECOND; -import static java.lang.Double.doubleToLongBits; -import static java.lang.Float.floatToRawIntBits; import static java.lang.Math.toIntExact; import static java.nio.charset.StandardCharsets.UTF_8; import static org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory.ObjectInspectorOptions; @@ -126,64 +124,71 @@ private static synchronized ObjectInspector getInspector(Type type) public void testPrimitiveSlice() { // boolean - Block expectedBoolean = VARBINARY.createBlockBuilder(null, 1).writeByte(1).closeEntry().build(); + Block expectedBoolean = createSingleValue(BOOLEAN, blockBuilder -> BOOLEAN.writeBoolean(blockBuilder, true)); Block actualBoolean = toBinaryBlock(BOOLEAN, true, getInspector(Boolean.class)); assertBlockEquals(actualBoolean, expectedBoolean); // byte - Block expectedByte = VARBINARY.createBlockBuilder(null, 1).writeByte(5).closeEntry().build(); + Block expectedByte = createSingleValue(TINYINT, blockBuilder -> TINYINT.writeLong(blockBuilder, 5)); Block actualByte = toBinaryBlock(TINYINT, (byte) 5, getInspector(Byte.class)); assertBlockEquals(actualByte, expectedByte); // short - Block expectedShort = VARBINARY.createBlockBuilder(null, 1).writeShort(2).closeEntry().build(); + Block expectedShort = createSingleValue(SMALLINT, blockBuilder -> SMALLINT.writeLong(blockBuilder, 2)); Block actualShort = toBinaryBlock(SMALLINT, (short) 2, getInspector(Short.class)); assertBlockEquals(actualShort, expectedShort); // int - Block expectedInt = VARBINARY.createBlockBuilder(null, 1).writeInt(1).closeEntry().build(); + Block expectedInt = createSingleValue(INTEGER, blockBuilder -> INTEGER.writeLong(blockBuilder, 1)); Block actualInt = toBinaryBlock(INTEGER, 1, getInspector(Integer.class)); assertBlockEquals(actualInt, expectedInt); // long - Block expectedLong = VARBINARY.createBlockBuilder(null, 1).writeLong(10).closeEntry().build(); + Block expectedLong = createSingleValue(BIGINT, blockBuilder -> BIGINT.writeLong(blockBuilder, 10)); Block actualLong = toBinaryBlock(BIGINT, 10L, getInspector(Long.class)); assertBlockEquals(actualLong, expectedLong); // float - Block expectedFloat = VARBINARY.createBlockBuilder(null, 1).writeInt(floatToRawIntBits(20.0f)).closeEntry().build(); + Block expectedFloat = createSingleValue(REAL, blockBuilder -> REAL.writeLong(blockBuilder, Float.floatToIntBits(20.0f))); Block actualFloat = toBinaryBlock(REAL, 20.0f, getInspector(Float.class)); assertBlockEquals(actualFloat, expectedFloat); // double - Block expectedDouble = VARBINARY.createBlockBuilder(null, 1).writeLong(doubleToLongBits(30.12)).closeEntry().build(); + Block expectedDouble = createSingleValue(DOUBLE, blockBuilder -> DOUBLE.writeDouble(blockBuilder, 30.12d)); Block actualDouble = toBinaryBlock(DOUBLE, 30.12d, getInspector(Double.class)); assertBlockEquals(actualDouble, expectedDouble); // string - Block expectedString = VARBINARY.createBlockBuilder(null, 1).writeBytes(utf8Slice("abdd"), 0, 4).closeEntry().build(); - Block actualString = toBinaryBlock(createUnboundedVarcharType(), "abdd", getInspector(String.class)); + Block expectedString = createSingleValue(VARCHAR, blockBuilder -> VARCHAR.writeString(blockBuilder, "value")); + Block actualString = toBinaryBlock(VARCHAR, "value", getInspector(String.class)); assertBlockEquals(actualString, expectedString); // date int date = toIntExact(LocalDate.of(2008, 10, 28).toEpochDay()); - Block expectedDate = VARBINARY.createBlockBuilder(null, 1).writeInt(date).closeEntry().build(); + Block expectedDate = createSingleValue(DATE, blockBuilder -> DATE.writeLong(blockBuilder, date)); Block actualDate = toBinaryBlock(DATE, Date.ofEpochDay(date), getInspector(Date.class)); assertBlockEquals(actualDate, expectedDate); // timestamp DateTime dateTime = new DateTime(2008, 10, 28, 16, 7, 15, 123); - Block expectedTimestamp = VARBINARY.createBlockBuilder(null, 1).writeLong(dateTime.getMillis() * MICROSECONDS_PER_MILLISECOND).closeEntry().build(); + Block expectedTimestamp = createSingleValue(TIMESTAMP_MILLIS, blockBuilder -> TIMESTAMP_MILLIS.writeLong(blockBuilder, dateTime.getMillis() * 1000)); Block actualTimestamp = toBinaryBlock(TIMESTAMP_MILLIS, Timestamp.ofEpochMilli(dateTime.getMillis()), getInspector(Timestamp.class)); assertBlockEquals(actualTimestamp, expectedTimestamp); // binary byte[] byteArray = {81, 82, 84, 85}; - Block expectedBinary = VARBINARY.createBlockBuilder(null, 1).writeBytes(Slices.wrappedBuffer(byteArray), 0, 4).closeEntry().build(); + Block expectedBinary = createSingleValue(VARBINARY, blockBuilder -> VARBINARY.writeSlice(blockBuilder, Slices.wrappedBuffer(byteArray))); Block actualBinary = toBinaryBlock(VARBINARY, byteArray, getInspector(byte[].class)); assertBlockEquals(actualBinary, expectedBinary); } + private static Block createSingleValue(io.trino.spi.type.Type type, Consumer outputConsumer) + { + BlockBuilder blockBuilder = type.createBlockBuilder(null, 1); + outputConsumer.accept(blockBuilder); + return blockBuilder.build(); + } + @Test public void testListBlock() { @@ -334,7 +339,7 @@ private static Block toBinaryBlock(io.trino.spi.type.Type type, Object object, O private static Block getPrimitiveBlock(io.trino.spi.type.Type type, Object object, ObjectInspector inspector) { - BlockBuilder builder = VARBINARY.createBlockBuilder(null, 1); + BlockBuilder builder = type.createBlockBuilder(null, 1); serializeObject(type, builder, object, inspector); return builder.build(); } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestStatistics.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestStatistics.java index 6c8527c90bc8..02af9c7e2511 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestStatistics.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/util/TestStatistics.java @@ -24,9 +24,9 @@ import io.trino.plugin.hive.metastore.HiveColumnStatistics; import io.trino.plugin.hive.metastore.IntegerStatistics; import io.trino.spi.block.Block; +import io.trino.spi.block.BlockBuilder; import io.trino.spi.statistics.ComputedStatistics; import io.trino.spi.statistics.TableStatisticType; -import io.trino.spi.type.BigintType; import io.trino.spi.type.Type; import org.testng.annotations.Test; @@ -53,6 +53,7 @@ import static io.trino.plugin.hive.util.Statistics.merge; import static io.trino.plugin.hive.util.Statistics.reduce; import static io.trino.spi.predicate.Utils.nativeValueToBlock; +import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.DoubleType.DOUBLE; import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.RealType.REAL; @@ -318,7 +319,11 @@ public void testMergeHiveColumnStatisticsMap() public void testFromComputedStatistics() { Function singleIntegerValueBlock = value -> - BigintType.BIGINT.createBlockBuilder(null, 1).writeLong(value).build(); + { + BlockBuilder blockBuilder = BIGINT.createBlockBuilder(null, 1); + BIGINT.writeLong(blockBuilder, value); + return blockBuilder.build(); + }; ComputedStatistics statistics = ComputedStatistics.builder(ImmutableList.of(), ImmutableList.of()) .addTableStatistic(TableStatisticType.ROW_COUNT, singleIntegerValueBlock.apply(5)) diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/FilesTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/FilesTable.java index c361972707e5..7c4f317e4ce0 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/FilesTable.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/FilesTable.java @@ -29,6 +29,7 @@ import io.trino.spi.connector.SystemTable; import io.trino.spi.predicate.TupleDomain; import io.trino.spi.type.ArrayType; +import io.trino.spi.type.MapType; import io.trino.spi.type.TypeManager; import org.apache.iceberg.DataFile; import org.apache.iceberg.FileScanTask; @@ -56,6 +57,7 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static io.trino.spi.block.MapValueBuilder.buildMapValue; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.IntegerType.INTEGER; import static io.trino.spi.type.TypeSignature.mapType; @@ -135,16 +137,16 @@ private static class PlanFilesIterable private final Map idToTypeMapping; private final List types; private boolean closed; - private final io.trino.spi.type.Type integerToBigintMapType; - private final io.trino.spi.type.Type integerToVarcharMapType; + private final MapType integerToBigintMapType; + private final MapType integerToVarcharMapType; public PlanFilesIterable(CloseableIterable planFiles, Map idToTypeMapping, List types, TypeManager typeManager) { this.planFiles = requireNonNull(planFiles, "planFiles is null"); this.idToTypeMapping = ImmutableMap.copyOf(requireNonNull(idToTypeMapping, "idToTypeMapping is null")); this.types = ImmutableList.copyOf(requireNonNull(types, "types is null")); - this.integerToBigintMapType = typeManager.getType(mapType(INTEGER.getTypeSignature(), BIGINT.getTypeSignature())); - this.integerToVarcharMapType = typeManager.getType(mapType(INTEGER.getTypeSignature(), VARCHAR.getTypeSignature())); + this.integerToBigintMapType = new MapType(INTEGER, BIGINT, typeManager.getTypeOperators()); + this.integerToVarcharMapType = new MapType(INTEGER, VARCHAR, typeManager.getTypeOperators()); addCloseable(planFiles); } @@ -239,26 +241,24 @@ private Object getIntegerVarcharMapBlock(Map value) private Object toIntegerBigintMapBlock(Map values) { - BlockBuilder blockBuilder = integerToBigintMapType.createBlockBuilder(null, 1); - BlockBuilder singleMapBlockBuilder = blockBuilder.beginBlockEntry(); - values.forEach((key, value) -> { - INTEGER.writeLong(singleMapBlockBuilder, key); - BIGINT.writeLong(singleMapBlockBuilder, value); - }); - blockBuilder.closeEntry(); - return integerToBigintMapType.getObject(blockBuilder, 0); + return buildMapValue( + integerToBigintMapType, + values.size(), + (keyBuilder, valueBuilder) -> values.forEach((key, value) -> { + INTEGER.writeLong(keyBuilder, key); + BIGINT.writeLong(valueBuilder, value); + })); } private Object toIntegerVarcharMapBlock(Map values) { - BlockBuilder blockBuilder = integerToVarcharMapType.createBlockBuilder(null, 1); - BlockBuilder singleMapBlockBuilder = blockBuilder.beginBlockEntry(); - values.forEach((key, value) -> { - INTEGER.writeLong(singleMapBlockBuilder, key); - VARCHAR.writeString(singleMapBlockBuilder, value); - }); - blockBuilder.closeEntry(); - return integerToVarcharMapType.getObject(blockBuilder, 0); + return buildMapValue( + integerToVarcharMapType, + values.size(), + (keyBuilder, valueBuilder) -> values.forEach((key, value) -> { + INTEGER.writeLong(keyBuilder, key); + VARCHAR.writeString(valueBuilder, value); + })); } @Nullable diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroDataConversion.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroDataConversion.java index 5133e3f93877..d2ab5e7eb263 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroDataConversion.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroDataConversion.java @@ -17,8 +17,11 @@ import io.airlift.slice.Slices; import io.trino.spi.Page; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.DecimalType; import io.trino.spi.type.Decimals; @@ -310,11 +313,11 @@ public static void serializeToTrinoBlock(Type type, org.apache.iceberg.types.Typ Collection array = (Collection) object; Type elementType = ((ArrayType) type).getElementType(); org.apache.iceberg.types.Type elementIcebergType = icebergType.asListType().elementType(); - BlockBuilder currentBuilder = builder.beginBlockEntry(); - for (Object element : array) { - serializeToTrinoBlock(elementType, elementIcebergType, currentBuilder, element); - } - builder.closeEntry(); + ((ArrayBlockBuilder) builder).buildEntry(elementBuilder -> { + for (Object element : array) { + serializeToTrinoBlock(elementType, elementIcebergType, elementBuilder, element); + } + }); return; } if (type instanceof MapType) { @@ -323,23 +326,23 @@ public static void serializeToTrinoBlock(Type type, org.apache.iceberg.types.Typ Type valueType = ((MapType) type).getValueType(); org.apache.iceberg.types.Type keyIcebergType = icebergType.asMapType().keyType(); org.apache.iceberg.types.Type valueIcebergType = icebergType.asMapType().valueType(); - BlockBuilder currentBuilder = builder.beginBlockEntry(); - for (Map.Entry entry : map.entrySet()) { - serializeToTrinoBlock(keyType, keyIcebergType, currentBuilder, entry.getKey()); - serializeToTrinoBlock(valueType, valueIcebergType, currentBuilder, entry.getValue()); - } - builder.closeEntry(); + ((MapBlockBuilder) builder).buildEntry((keyBuilder, valueBuilder) -> { + for (Map.Entry entry : map.entrySet()) { + serializeToTrinoBlock(keyType, keyIcebergType, keyBuilder, entry.getKey()); + serializeToTrinoBlock(valueType, valueIcebergType, valueBuilder, entry.getValue()); + } + }); return; } if (type instanceof RowType) { Record record = (Record) object; List typeParameters = type.getTypeParameters(); List icebergFields = icebergType.asStructType().fields(); - BlockBuilder currentBuilder = builder.beginBlockEntry(); - for (int i = 0; i < typeParameters.size(); i++) { - serializeToTrinoBlock(typeParameters.get(i), icebergFields.get(i).type(), currentBuilder, record.get(i)); - } - builder.closeEntry(); + ((RowBlockBuilder) builder).buildEntry(fieldBuilders -> { + for (int i = 0; i < typeParameters.size(); i++) { + serializeToTrinoBlock(typeParameters.get(i), icebergFields.get(i).type(), fieldBuilders.get(i), record.get(i)); + } + }); return; } throw new TrinoException(NOT_SUPPORTED, "unsupported type: " + type); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroPageSource.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroPageSource.java index 3add0c0ceef6..65fbad86c794 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroPageSource.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergAvroPageSource.java @@ -38,6 +38,7 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static io.trino.plugin.iceberg.IcebergAvroDataConversion.serializeToTrinoBlock; +import static io.trino.spi.type.BigintType.BIGINT; import static java.util.Objects.requireNonNull; public class IcebergAvroPageSource @@ -131,7 +132,7 @@ public Page getNextPage() Record record = recordIterator.next(); for (int channel = 0; channel < columnTypes.size(); channel++) { if (isIndexColumn(channel)) { - pageBuilder.getBlockBuilder(channel).writeLong(rowId); + BIGINT.writeLong(pageBuilder.getBlockBuilder(channel), rowId); } else { String name = columnNames.get(channel); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ManifestsTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ManifestsTable.java index 58beeb6a005b..7b8da9a4baf6 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ManifestsTable.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/ManifestsTable.java @@ -17,7 +17,9 @@ import io.trino.plugin.iceberg.util.PageListBuilder; import io.trino.spi.Page; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.connector.ConnectorSession; @@ -135,27 +137,27 @@ private static List buildPages(ConnectorTableMetadata tableMetadata, Table private static void writePartitionSummaries(BlockBuilder arrayBlockBuilder, List summaries, PartitionSpec partitionSpec) { - BlockBuilder singleArrayWriter = arrayBlockBuilder.beginBlockEntry(); - for (int i = 0; i < summaries.size(); i++) { - PartitionFieldSummary summary = summaries.get(i); - PartitionField field = partitionSpec.fields().get(i); - Type nestedType = partitionSpec.partitionType().fields().get(i).type(); - - BlockBuilder rowBuilder = singleArrayWriter.beginBlockEntry(); - BOOLEAN.writeBoolean(rowBuilder, summary.containsNull()); - Boolean containsNan = summary.containsNaN(); - if (containsNan == null) { - rowBuilder.appendNull(); + ((ArrayBlockBuilder) arrayBlockBuilder).buildEntry(elementBuilder -> { + for (int i = 0; i < summaries.size(); i++) { + PartitionFieldSummary summary = summaries.get(i); + PartitionField field = partitionSpec.fields().get(i); + Type nestedType = partitionSpec.partitionType().fields().get(i).type(); + + ((RowBlockBuilder) elementBuilder).buildEntry(fieldBuilders -> { + BOOLEAN.writeBoolean(fieldBuilders.get(0), summary.containsNull()); + Boolean containsNan = summary.containsNaN(); + if (containsNan == null) { + fieldBuilders.get(1).appendNull(); + } + else { + BOOLEAN.writeBoolean(fieldBuilders.get(1), containsNan); + } + VARCHAR.writeString(fieldBuilders.get(2), field.transform().toHumanString( + nestedType, Conversions.fromByteBuffer(nestedType, summary.lowerBound()))); + VARCHAR.writeString(fieldBuilders.get(3), field.transform().toHumanString( + nestedType, Conversions.fromByteBuffer(nestedType, summary.upperBound()))); + }); } - else { - BOOLEAN.writeBoolean(rowBuilder, containsNan); - } - VARCHAR.writeString(rowBuilder, field.transform().toHumanString( - nestedType, Conversions.fromByteBuffer(nestedType, summary.lowerBound()))); - VARCHAR.writeString(rowBuilder, field.transform().toHumanString( - nestedType, Conversions.fromByteBuffer(nestedType, summary.upperBound()))); - singleArrayWriter.closeEntry(); - } - arrayBlockBuilder.closeEntry(); + }); } } diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTable.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTable.java index 226d942add0b..22108bd9348f 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTable.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/PartitionTable.java @@ -17,7 +17,6 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.ConnectorTableMetadata; @@ -61,6 +60,7 @@ import static io.trino.plugin.iceberg.IcebergUtil.getIdentityPartitions; import static io.trino.plugin.iceberg.IcebergUtil.primitiveFieldTypes; import static io.trino.plugin.iceberg.TypeConverter.toTrinoType; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.TypeUtils.writeNativeValue; import static java.util.Objects.requireNonNull; @@ -257,24 +257,22 @@ private RecordCursor buildRecordCursor(Map { - BlockBuilder partitionRowBlockBuilder = partitionColumnType.rowType.createBlockBuilder(null, 1); - BlockBuilder partitionBlockBuilder = partitionRowBlockBuilder.beginBlockEntry(); - List partitionColumnTypes = partitionColumnType.rowType.getFields().stream() - .map(RowType.Field::getType) - .collect(toImmutableList()); - for (int i = 0; i < partitionColumnTypes.size(); i++) { - io.trino.spi.type.Type trinoType = partitionColumnType.rowType.getFields().get(i).getType(); - Object value = null; - Integer fieldId = partitionColumnType.fieldIds.get(i); - if (partitionStruct.fieldIdToIndex.containsKey(fieldId)) { - value = convertIcebergValueToTrino( - partitionTypes.get(i), - partitionStruct.structLikeWrapper.get().get(partitionStruct.fieldIdToIndex.get(fieldId), partitionColumnClass.get(i))); + row.add(buildRowValue(partitionColumnType.rowType, fields -> { + List partitionColumnTypes = partitionColumnType.rowType.getFields().stream() + .map(RowType.Field::getType) + .collect(toImmutableList()); + for (int i = 0; i < partitionColumnTypes.size(); i++) { + io.trino.spi.type.Type trinoType = partitionColumnType.rowType.getFields().get(i).getType(); + Object value = null; + Integer fieldId = partitionColumnType.fieldIds.get(i); + if (partitionStruct.fieldIdToIndex.containsKey(fieldId)) { + value = convertIcebergValueToTrino( + partitionTypes.get(i), + partitionStruct.structLikeWrapper.get().get(partitionStruct.fieldIdToIndex.get(fieldId), partitionColumnClass.get(i))); + } + writeNativeValue(trinoType, fields.get(i), value); } - writeNativeValue(trinoType, partitionBlockBuilder, value); - } - partitionRowBlockBuilder.closeEntry(); - row.add(partitionColumnType.rowType.getObject(partitionRowBlockBuilder, 0)); + })); }); // add the top level metrics. @@ -284,25 +282,26 @@ private RecordCursor buildRecordCursor(Map { - BlockBuilder dataRowBlockBuilder = dataColumnType.createBlockBuilder(null, 1); - BlockBuilder dataBlockBuilder = dataRowBlockBuilder.beginBlockEntry(); - - for (int i = 0; i < columnMetricTypes.size(); i++) { - Integer fieldId = nonPartitionPrimitiveColumns.get(i).fieldId(); - Object min = icebergStatistics.getMinValues().get(fieldId); - Object max = icebergStatistics.getMaxValues().get(fieldId); - Long nullCount = icebergStatistics.getNullCounts().get(fieldId); - Long nanCount = icebergStatistics.getNanCounts().get(fieldId); - if (min == null && max == null && nullCount == null) { - row.add(null); - return; - } - - RowType columnMetricType = columnMetricTypes.get(i); - columnMetricType.writeObject(dataBlockBuilder, getColumnMetricBlock(columnMetricType, min, max, nullCount, nanCount)); + try { + row.add(buildRowValue(dataColumnType, fields -> { + for (int i = 0; i < columnMetricTypes.size(); i++) { + Integer fieldId = nonPartitionPrimitiveColumns.get(i).fieldId(); + Object min = icebergStatistics.getMinValues().get(fieldId); + Object max = icebergStatistics.getMaxValues().get(fieldId); + Long nullCount = icebergStatistics.getNullCounts().get(fieldId); + Long nanCount = icebergStatistics.getNanCounts().get(fieldId); + if (min == null && max == null && nullCount == null) { + throw new MissingColumnMetricsException(); + } + + RowType columnMetricType = columnMetricTypes.get(i); + columnMetricType.writeObject(fields.get(i), getColumnMetricBlock(columnMetricType, min, max, nullCount, nanCount)); + } + })); + } + catch (MissingColumnMetricsException ignored) { + row.add(null); } - dataRowBlockBuilder.closeEntry(); - row.add(dataColumnType.getObject(dataRowBlockBuilder, 0)); }); records.add(row); @@ -311,6 +310,10 @@ private RecordCursor buildRecordCursor(Map partitionTypes() { ImmutableList.Builder partitionTypeBuilder = ImmutableList.builder(); @@ -324,16 +327,13 @@ private List partitionTypes() private static Block getColumnMetricBlock(RowType columnMetricType, Object min, Object max, Long nullCount, Long nanCount) { - BlockBuilder rowBlockBuilder = columnMetricType.createBlockBuilder(null, 1); - BlockBuilder builder = rowBlockBuilder.beginBlockEntry(); - List fields = columnMetricType.getFields(); - writeNativeValue(fields.get(0).getType(), builder, min); - writeNativeValue(fields.get(1).getType(), builder, max); - writeNativeValue(fields.get(2).getType(), builder, nullCount); - writeNativeValue(fields.get(3).getType(), builder, nanCount); - - rowBlockBuilder.closeEntry(); - return columnMetricType.getObject(rowBlockBuilder, 0); + return buildRowValue(columnMetricType, fieldBuilders -> { + List fields = columnMetricType.getFields(); + writeNativeValue(fields.get(0).getType(), fieldBuilders.get(0), min); + writeNativeValue(fields.get(1).getType(), fieldBuilders.get(1), max); + writeNativeValue(fields.get(2).getType(), fieldBuilders.get(2), nullCount); + writeNativeValue(fields.get(3).getType(), fieldBuilders.get(3), nanCount); + }); } @VisibleForTesting diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/PageListBuilder.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/PageListBuilder.java index 90ec89715c8a..515186ca8d47 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/PageListBuilder.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/util/PageListBuilder.java @@ -17,7 +17,9 @@ import io.airlift.slice.Slice; import io.trino.spi.Page; import io.trino.spi.PageBuilder; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.ConnectorTableMetadata; import io.trino.spi.type.TimeZoneKey; @@ -130,65 +132,59 @@ public void appendVarbinary(Slice value) public void appendIntegerArray(Iterable values) { - BlockBuilder column = nextColumn(); - BlockBuilder array = column.beginBlockEntry(); - for (Integer value : values) { - INTEGER.writeLong(array, value); - } - column.closeEntry(); + ArrayBlockBuilder column = (ArrayBlockBuilder) nextColumn(); + column.buildEntry(elementBuilder -> { + for (Integer value : values) { + INTEGER.writeLong(elementBuilder, value); + } + }); } public void appendBigintArray(Iterable values) { - BlockBuilder column = nextColumn(); - BlockBuilder array = column.beginBlockEntry(); - for (Long value : values) { - BIGINT.writeLong(array, value); - } - column.closeEntry(); + ArrayBlockBuilder column = (ArrayBlockBuilder) nextColumn(); + column.buildEntry(elementBuilder -> { + for (Long value : values) { + BIGINT.writeLong(elementBuilder, value); + } + }); } public void appendVarcharArray(Iterable values) { - BlockBuilder column = nextColumn(); - BlockBuilder array = column.beginBlockEntry(); - for (String value : values) { - VARCHAR.writeString(array, value); - } - column.closeEntry(); + ArrayBlockBuilder column = (ArrayBlockBuilder) nextColumn(); + column.buildEntry(elementBuilder -> { + for (String value : values) { + VARCHAR.writeString(elementBuilder, value); + } + }); } public void appendVarcharVarcharMap(Map values) { - BlockBuilder column = nextColumn(); - BlockBuilder map = column.beginBlockEntry(); - values.forEach((key, value) -> { - VARCHAR.writeString(map, key); - VARCHAR.writeString(map, value); - }); - column.closeEntry(); + MapBlockBuilder column = (MapBlockBuilder) nextColumn(); + column.buildEntry((keyBuilder, valueBuilder) -> values.forEach((key, value) -> { + VARCHAR.writeString(keyBuilder, key); + VARCHAR.writeString(valueBuilder, value); + })); } public void appendIntegerBigintMap(Map values) { - BlockBuilder column = nextColumn(); - BlockBuilder map = column.beginBlockEntry(); - values.forEach((key, value) -> { - INTEGER.writeLong(map, key); - BIGINT.writeLong(map, value); - }); - column.closeEntry(); + MapBlockBuilder column = (MapBlockBuilder) nextColumn(); + column.buildEntry((keyBuilder, valueBuilder) -> values.forEach((key, value) -> { + INTEGER.writeLong(keyBuilder, key); + BIGINT.writeLong(valueBuilder, value); + })); } public void appendIntegerVarcharMap(Map values) { - BlockBuilder column = nextColumn(); - BlockBuilder map = column.beginBlockEntry(); - values.forEach((key, value) -> { - INTEGER.writeLong(map, key); - VARCHAR.writeString(map, value); - }); - column.closeEntry(); + MapBlockBuilder column = (MapBlockBuilder) nextColumn(); + column.buildEntry((keyBuilder, valueBuilder) -> values.forEach((key, value) -> { + INTEGER.writeLong(keyBuilder, key); + VARCHAR.writeString(valueBuilder, value); + })); } public BlockBuilder nextColumn() diff --git a/plugin/trino-kafka/src/main/java/io/trino/plugin/kafka/KafkaRecordSet.java b/plugin/trino-kafka/src/main/java/io/trino/plugin/kafka/KafkaRecordSet.java index 2ad607695906..e24d256e15bb 100644 --- a/plugin/trino-kafka/src/main/java/io/trino/plugin/kafka/KafkaRecordSet.java +++ b/plugin/trino-kafka/src/main/java/io/trino/plugin/kafka/KafkaRecordSet.java @@ -20,8 +20,8 @@ import io.trino.decoder.DecoderColumnHandle; import io.trino.decoder.FieldValueProvider; import io.trino.decoder.RowDecoder; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.RecordCursor; @@ -45,6 +45,7 @@ import static io.trino.decoder.FieldValueProviders.booleanValueProvider; import static io.trino.decoder.FieldValueProviders.bytesValueProvider; import static io.trino.decoder.FieldValueProviders.longValueProvider; +import static io.trino.spi.block.MapValueBuilder.buildMapValue; import static io.trino.spi.type.Timestamps.MICROSECONDS_PER_MILLISECOND; import static io.trino.spi.type.TypeUtils.writeNativeValue; import static java.lang.Math.max; @@ -268,25 +269,25 @@ public static FieldValueProvider headerMapValueProvider(MapType varcharMapType, Type valueArrayType = varcharMapType.getTypeParameters().get(1); Type valueType = valueArrayType.getTypeParameters().get(0); - BlockBuilder mapBlockBuilder = varcharMapType.createBlockBuilder(null, 1); - BlockBuilder builder = mapBlockBuilder.beginBlockEntry(); - // Group by keys and collect values as array. Multimap headerMap = ArrayListMultimap.create(); for (Header header : headers) { headerMap.put(header.key(), header.value()); } - for (String headerKey : headerMap.keySet()) { - writeNativeValue(keyType, builder, headerKey); - BlockBuilder arrayBuilder = builder.beginBlockEntry(); - for (byte[] value : headerMap.get(headerKey)) { - writeNativeValue(valueType, arrayBuilder, value); - } - builder.closeEntry(); - } - - mapBlockBuilder.closeEntry(); + Block map = buildMapValue( + varcharMapType, + headerMap.size(), + (keyBuilder, valueBuilder) -> { + for (String headerKey : headerMap.keySet()) { + writeNativeValue(keyType, keyBuilder, headerKey); + ((ArrayBlockBuilder) valueBuilder).buildEntry(elementBuilder -> { + for (byte[] value : headerMap.get(headerKey)) { + writeNativeValue(valueType, elementBuilder, value); + } + }); + } + }); return new FieldValueProvider() { @@ -299,7 +300,7 @@ public boolean isNull() @Override public Block getBlock() { - return varcharMapType.getObject(mapBlockBuilder, 0); + return map; } }; } diff --git a/plugin/trino-kafka/src/test/java/io/trino/plugin/kafka/encoder/raw/TestRawEncoderMapping.java b/plugin/trino-kafka/src/test/java/io/trino/plugin/kafka/encoder/raw/TestRawEncoderMapping.java index c8a4c8f11136..b297fbafdc11 100644 --- a/plugin/trino-kafka/src/test/java/io/trino/plugin/kafka/encoder/raw/TestRawEncoderMapping.java +++ b/plugin/trino-kafka/src/test/java/io/trino/plugin/kafka/encoder/raw/TestRawEncoderMapping.java @@ -59,10 +59,13 @@ public void testMapping() buf.put("abcdef".getBytes(StandardCharsets.UTF_8)); // 36-42 buf.put("abcdef".getBytes(StandardCharsets.UTF_8)); // 42-48 - Block longArrayBlock = new LongArrayBlockBuilder(null, 1).writeLong(123456789).closeEntry().build(); + LongArrayBlockBuilder longArrayBlockBuilder = new LongArrayBlockBuilder(null, 1); + BIGINT.writeLong(longArrayBlockBuilder, 123456789); + Block longArrayBlock = longArrayBlockBuilder.build(); + Block varArrayBlock = new VariableWidthBlockBuilder(null, 1, 6) - .writeBytes(Slices.wrappedBuffer("abcdef".getBytes(StandardCharsets.UTF_8)), 0, 6) - .closeEntry().build(); + .writeEntry(Slices.wrappedBuffer("abcdef".getBytes(StandardCharsets.UTF_8)), 0, 6) + .build(); rowEncoder.appendColumnValue(longArrayBlock, 0); rowEncoder.appendColumnValue(varArrayBlock, 0); diff --git a/plugin/trino-kafka/src/test/java/io/trino/plugin/kafka/protobuf/TestProtobufEncoder.java b/plugin/trino-kafka/src/test/java/io/trino/plugin/kafka/protobuf/TestProtobufEncoder.java index 4ff338361d66..5e1e97fe0c9d 100644 --- a/plugin/trino-kafka/src/test/java/io/trino/plugin/kafka/protobuf/TestProtobufEncoder.java +++ b/plugin/trino-kafka/src/test/java/io/trino/plugin/kafka/protobuf/TestProtobufEncoder.java @@ -23,8 +23,11 @@ import io.trino.plugin.kafka.encoder.EncoderColumnHandle; import io.trino.plugin.kafka.encoder.RowEncoder; import io.trino.plugin.kafka.encoder.protobuf.ProtobufRowEncoderFactory; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.type.ArrayType; import io.trino.spi.type.MapType; import io.trino.spi.type.RowType; @@ -157,35 +160,32 @@ public void testStructuralDataTypes(String stringData, Integer integerData, Long RowEncoder rowEncoder = createRowEncoder("structural_datatypes.proto", columnHandles.subList(0, 3)); - BlockBuilder arrayBlockBuilder = columnHandles.get(0).getType() + ArrayBlockBuilder arrayBlockBuilder = (ArrayBlockBuilder) columnHandles.get(0).getType() .createBlockBuilder(null, 1); - BlockBuilder singleArrayBlockWriter = arrayBlockBuilder.beginBlockEntry(); - writeNativeValue(createVarcharType(5), singleArrayBlockWriter, utf8Slice(stringData)); - arrayBlockBuilder.closeEntry(); + arrayBlockBuilder.buildEntry(elementBuilder -> writeNativeValue(createVarcharType(5), elementBuilder, utf8Slice(stringData))); rowEncoder.appendColumnValue(arrayBlockBuilder.build(), 0); - BlockBuilder mapBlockBuilder = columnHandles.get(1).getType() + MapBlockBuilder mapBlockBuilder = (MapBlockBuilder) columnHandles.get(1).getType() .createBlockBuilder(null, 1); - BlockBuilder singleMapBlockWriter = mapBlockBuilder.beginBlockEntry(); - writeNativeValue(VARCHAR, singleMapBlockWriter, utf8Slice("Key")); - writeNativeValue(VARCHAR, singleMapBlockWriter, utf8Slice("Value")); - mapBlockBuilder.closeEntry(); + mapBlockBuilder.buildEntry((keyBuilder, valueBuilder) -> { + writeNativeValue(VARCHAR, keyBuilder, utf8Slice("Key")); + writeNativeValue(VARCHAR, valueBuilder, utf8Slice("Value")); + }); rowEncoder.appendColumnValue(mapBlockBuilder.build(), 0); - BlockBuilder rowBlockBuilder = columnHandles.get(2).getType() + RowBlockBuilder rowBlockBuilder = (RowBlockBuilder) columnHandles.get(2).getType() .createBlockBuilder(null, 1); - BlockBuilder singleRowBlockWriter = rowBlockBuilder.beginBlockEntry(); - writeNativeValue(VARCHAR, singleRowBlockWriter, Slices.utf8Slice(stringData)); - writeNativeValue(INTEGER, singleRowBlockWriter, integerData.longValue()); - writeNativeValue(BIGINT, singleRowBlockWriter, longData); - writeNativeValue(DOUBLE, singleRowBlockWriter, doubleData); - writeNativeValue(REAL, singleRowBlockWriter, (long) floatToIntBits(floatData)); - writeNativeValue(BOOLEAN, singleRowBlockWriter, booleanData); - writeNativeValue(VARCHAR, singleRowBlockWriter, enumData); - writeNativeValue(createTimestampType(6), singleRowBlockWriter, sqlTimestamp.getEpochMicros()); - writeNativeValue(VARBINARY, singleRowBlockWriter, bytesData); - - rowBlockBuilder.closeEntry(); + rowBlockBuilder.buildEntry(fieldBuilders -> { + writeNativeValue(VARCHAR, fieldBuilders.get(0), utf8Slice(stringData)); + writeNativeValue(INTEGER, fieldBuilders.get(1), integerData.longValue()); + writeNativeValue(BIGINT, fieldBuilders.get(2), longData); + writeNativeValue(DOUBLE, fieldBuilders.get(3), doubleData); + writeNativeValue(REAL, fieldBuilders.get(4), (long) floatToIntBits(floatData)); + writeNativeValue(BOOLEAN, fieldBuilders.get(5), booleanData); + writeNativeValue(VARCHAR, fieldBuilders.get(6), enumData); + writeNativeValue(createTimestampType(6), fieldBuilders.get(7), sqlTimestamp.getEpochMicros()); + writeNativeValue(VARBINARY, fieldBuilders.get(8), bytesData); + }); rowEncoder.appendColumnValue(rowBlockBuilder.build(), 0); assertEquals(messageBuilder.build().toByteArray(), rowEncoder.toByteArray()); @@ -246,19 +246,18 @@ public void testNestedStructuralDataTypes(String stringData, Integer integerData RowEncoder rowEncoder = createRowEncoder("structural_datatypes.proto", columnHandles); - BlockBuilder rowBlockBuilder = rowType - .createBlockBuilder(null, 1); - BlockBuilder singleRowBlockWriter = rowBlockBuilder.beginBlockEntry(); - writeNativeValue(VARCHAR, singleRowBlockWriter, Slices.utf8Slice(stringData)); - writeNativeValue(INTEGER, singleRowBlockWriter, integerData.longValue()); - writeNativeValue(BIGINT, singleRowBlockWriter, longData); - writeNativeValue(DOUBLE, singleRowBlockWriter, doubleData); - writeNativeValue(REAL, singleRowBlockWriter, (long) floatToIntBits(floatData)); - writeNativeValue(BOOLEAN, singleRowBlockWriter, booleanData); - writeNativeValue(VARCHAR, singleRowBlockWriter, enumData); - writeNativeValue(createTimestampType(6), singleRowBlockWriter, sqlTimestamp.getEpochMicros()); - writeNativeValue(VARBINARY, singleRowBlockWriter, bytesData); - rowBlockBuilder.closeEntry(); + RowBlockBuilder rowBlockBuilder = rowType.createBlockBuilder(null, 1); + rowBlockBuilder.buildEntry(fieldBuilders -> { + writeNativeValue(VARCHAR, fieldBuilders.get(0), Slices.utf8Slice(stringData)); + writeNativeValue(INTEGER, fieldBuilders.get(1), integerData.longValue()); + writeNativeValue(BIGINT, fieldBuilders.get(2), longData); + writeNativeValue(DOUBLE, fieldBuilders.get(3), doubleData); + writeNativeValue(REAL, fieldBuilders.get(4), (long) floatToIntBits(floatData)); + writeNativeValue(BOOLEAN, fieldBuilders.get(5), booleanData); + writeNativeValue(VARCHAR, fieldBuilders.get(6), enumData); + writeNativeValue(createTimestampType(6), fieldBuilders.get(7), sqlTimestamp.getEpochMicros()); + writeNativeValue(VARBINARY, fieldBuilders.get(8), bytesData); + }); RowType nestedRowType = (RowType) columnHandles.get(0).getType(); diff --git a/plugin/trino-ml/src/main/java/io/trino/plugin/ml/MLFeaturesFunctions.java b/plugin/trino-ml/src/main/java/io/trino/plugin/ml/MLFeaturesFunctions.java index f4ebb38d0e8b..f518a5b0112c 100644 --- a/plugin/trino-ml/src/main/java/io/trino/plugin/ml/MLFeaturesFunctions.java +++ b/plugin/trino-ml/src/main/java/io/trino/plugin/ml/MLFeaturesFunctions.java @@ -14,14 +14,14 @@ package io.trino.plugin.ml; import com.google.common.collect.ImmutableList; -import io.trino.spi.PageBuilder; import io.trino.spi.block.Block; -import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.BufferedMapValueBuilder; import io.trino.spi.function.ScalarFunction; import io.trino.spi.function.SqlType; import io.trino.spi.function.TypeParameter; import io.trino.spi.type.BigintType; import io.trino.spi.type.DoubleType; +import io.trino.spi.type.MapType; import io.trino.spi.type.StandardTypes; import io.trino.spi.type.Type; @@ -38,189 +38,180 @@ private MLFeaturesFunctions() {} @ScalarFunction("features") public static class Features1 { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public Features1(@TypeParameter(MAP_BIGINT_DOUBLE) Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType(MAP_BIGINT_DOUBLE) public Block features(@SqlType(StandardTypes.DOUBLE) double f1) { - return featuresHelper(pageBuilder, f1); + return featuresHelper(mapValueBuilder, f1); } } @ScalarFunction("features") public static class Features2 { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public Features2(@TypeParameter(MAP_BIGINT_DOUBLE) Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType(MAP_BIGINT_DOUBLE) public Block features(@SqlType(StandardTypes.DOUBLE) double f1, @SqlType(StandardTypes.DOUBLE) double f2) { - return featuresHelper(pageBuilder, f1, f2); + return featuresHelper(mapValueBuilder, f1, f2); } } @ScalarFunction("features") public static class Features3 { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public Features3(@TypeParameter(MAP_BIGINT_DOUBLE) Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType(MAP_BIGINT_DOUBLE) public Block features(@SqlType(StandardTypes.DOUBLE) double f1, @SqlType(StandardTypes.DOUBLE) double f2, @SqlType(StandardTypes.DOUBLE) double f3) { - return featuresHelper(pageBuilder, f1, f2, f3); + return featuresHelper(mapValueBuilder, f1, f2, f3); } } @ScalarFunction("features") public static class Features4 { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public Features4(@TypeParameter(MAP_BIGINT_DOUBLE) Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType(MAP_BIGINT_DOUBLE) public Block features(@SqlType(StandardTypes.DOUBLE) double f1, @SqlType(StandardTypes.DOUBLE) double f2, @SqlType(StandardTypes.DOUBLE) double f3, @SqlType(StandardTypes.DOUBLE) double f4) { - return featuresHelper(pageBuilder, f1, f2, f3, f4); + return featuresHelper(mapValueBuilder, f1, f2, f3, f4); } } @ScalarFunction("features") public static class Features5 { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public Features5(@TypeParameter(MAP_BIGINT_DOUBLE) Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType(MAP_BIGINT_DOUBLE) public Block features(@SqlType(StandardTypes.DOUBLE) double f1, @SqlType(StandardTypes.DOUBLE) double f2, @SqlType(StandardTypes.DOUBLE) double f3, @SqlType(StandardTypes.DOUBLE) double f4, @SqlType(StandardTypes.DOUBLE) double f5) { - return featuresHelper(pageBuilder, f1, f2, f3, f4, f5); + return featuresHelper(mapValueBuilder, f1, f2, f3, f4, f5); } } @ScalarFunction("features") public static class Features6 { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public Features6(@TypeParameter(MAP_BIGINT_DOUBLE) Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType(MAP_BIGINT_DOUBLE) public Block features(@SqlType(StandardTypes.DOUBLE) double f1, @SqlType(StandardTypes.DOUBLE) double f2, @SqlType(StandardTypes.DOUBLE) double f3, @SqlType(StandardTypes.DOUBLE) double f4, @SqlType(StandardTypes.DOUBLE) double f5, @SqlType(StandardTypes.DOUBLE) double f6) { - return featuresHelper(pageBuilder, f1, f2, f3, f4, f5, f6); + return featuresHelper(mapValueBuilder, f1, f2, f3, f4, f5, f6); } } @ScalarFunction("features") public static class Features7 { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public Features7(@TypeParameter(MAP_BIGINT_DOUBLE) Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType(MAP_BIGINT_DOUBLE) public Block features(@SqlType(StandardTypes.DOUBLE) double f1, @SqlType(StandardTypes.DOUBLE) double f2, @SqlType(StandardTypes.DOUBLE) double f3, @SqlType(StandardTypes.DOUBLE) double f4, @SqlType(StandardTypes.DOUBLE) double f5, @SqlType(StandardTypes.DOUBLE) double f6, @SqlType(StandardTypes.DOUBLE) double f7) { - return featuresHelper(pageBuilder, f1, f2, f3, f4, f5, f6, f7); + return featuresHelper(mapValueBuilder, f1, f2, f3, f4, f5, f6, f7); } } @ScalarFunction("features") public static class Features8 { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public Features8(@TypeParameter(MAP_BIGINT_DOUBLE) Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType(MAP_BIGINT_DOUBLE) public Block features(@SqlType(StandardTypes.DOUBLE) double f1, @SqlType(StandardTypes.DOUBLE) double f2, @SqlType(StandardTypes.DOUBLE) double f3, @SqlType(StandardTypes.DOUBLE) double f4, @SqlType(StandardTypes.DOUBLE) double f5, @SqlType(StandardTypes.DOUBLE) double f6, @SqlType(StandardTypes.DOUBLE) double f7, @SqlType(StandardTypes.DOUBLE) double f8) { - return featuresHelper(pageBuilder, f1, f2, f3, f4, f5, f6, f7, f8); + return featuresHelper(mapValueBuilder, f1, f2, f3, f4, f5, f6, f7, f8); } } @ScalarFunction("features") public static class Features9 { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public Features9(@TypeParameter(MAP_BIGINT_DOUBLE) Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType(MAP_BIGINT_DOUBLE) public Block features(@SqlType(StandardTypes.DOUBLE) double f1, @SqlType(StandardTypes.DOUBLE) double f2, @SqlType(StandardTypes.DOUBLE) double f3, @SqlType(StandardTypes.DOUBLE) double f4, @SqlType(StandardTypes.DOUBLE) double f5, @SqlType(StandardTypes.DOUBLE) double f6, @SqlType(StandardTypes.DOUBLE) double f7, @SqlType(StandardTypes.DOUBLE) double f8, @SqlType(StandardTypes.DOUBLE) double f9) { - return featuresHelper(pageBuilder, f1, f2, f3, f4, f5, f6, f7, f8, f9); + return featuresHelper(mapValueBuilder, f1, f2, f3, f4, f5, f6, f7, f8, f9); } } @ScalarFunction("features") public static class Features10 { - private final PageBuilder pageBuilder; + private final BufferedMapValueBuilder mapValueBuilder; public Features10(@TypeParameter(MAP_BIGINT_DOUBLE) Type mapType) { - pageBuilder = new PageBuilder(ImmutableList.of(mapType)); + mapValueBuilder = BufferedMapValueBuilder.createBuffered((MapType) mapType); } @SqlType(MAP_BIGINT_DOUBLE) public Block features(@SqlType(StandardTypes.DOUBLE) double f1, @SqlType(StandardTypes.DOUBLE) double f2, @SqlType(StandardTypes.DOUBLE) double f3, @SqlType(StandardTypes.DOUBLE) double f4, @SqlType(StandardTypes.DOUBLE) double f5, @SqlType(StandardTypes.DOUBLE) double f6, @SqlType(StandardTypes.DOUBLE) double f7, @SqlType(StandardTypes.DOUBLE) double f8, @SqlType(StandardTypes.DOUBLE) double f9, @SqlType(StandardTypes.DOUBLE) double f10) { - return featuresHelper(pageBuilder, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10); + return featuresHelper(mapValueBuilder, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10); } } - private static Block featuresHelper(PageBuilder pageBuilder, double... features) + private static Block featuresHelper(BufferedMapValueBuilder mapValueBuilder, double... features) { - if (pageBuilder.isFull()) { - pageBuilder.reset(); - } - - BlockBuilder mapBlockBuilder = pageBuilder.getBlockBuilder(0); - BlockBuilder blockBuilder = mapBlockBuilder.beginBlockEntry(); - - for (int i = 0; i < features.length; i++) { - BigintType.BIGINT.writeLong(blockBuilder, i); - DoubleType.DOUBLE.writeDouble(blockBuilder, features[i]); - } - - mapBlockBuilder.closeEntry(); - pageBuilder.declarePosition(); - return mapBlockBuilder.getObject(mapBlockBuilder.getPositionCount() - 1, Block.class); + return mapValueBuilder.build(features.length, (keyBuilder, valueBuilder) -> { + for (int i = 0; i < features.length; i++) { + BigintType.BIGINT.writeLong(keyBuilder, i); + DoubleType.DOUBLE.writeDouble(valueBuilder, features[i]); + } + }); } } diff --git a/plugin/trino-ml/src/main/java/io/trino/plugin/ml/type/ModelType.java b/plugin/trino-ml/src/main/java/io/trino/plugin/ml/type/ModelType.java index c84168005ce0..de39c6bacc73 100644 --- a/plugin/trino-ml/src/main/java/io/trino/plugin/ml/type/ModelType.java +++ b/plugin/trino-ml/src/main/java/io/trino/plugin/ml/type/ModelType.java @@ -16,6 +16,7 @@ import io.airlift.slice.Slice; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.type.AbstractVariableWidthType; import io.trino.spi.type.TypeSignature; @@ -48,8 +49,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -68,7 +68,7 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value) @Override public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length) { - blockBuilder.writeBytes(value, offset, length).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length); } @Override diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoPageSource.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoPageSource.java index 5f4a14d1e754..73252eee1f96 100644 --- a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoPageSource.java +++ b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoPageSource.java @@ -21,13 +21,18 @@ import io.trino.spi.Page; import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; +import io.trino.spi.block.RowBlockBuilder; import io.trino.spi.connector.ConnectorPageSource; import io.trino.spi.type.CharType; import io.trino.spi.type.DecimalType; import io.trino.spi.type.Decimals; import io.trino.spi.type.Int128; +import io.trino.spi.type.MapType; +import io.trino.spi.type.RowType; import io.trino.spi.type.Type; import io.trino.spi.type.TypeSignatureParameter; import io.trino.spi.type.VarbinaryType; @@ -56,8 +61,6 @@ import static io.trino.plugin.mongodb.ObjectIdType.OBJECT_ID; import static io.trino.plugin.mongodb.TypeUtils.isArrayType; import static io.trino.plugin.mongodb.TypeUtils.isJsonType; -import static io.trino.plugin.mongodb.TypeUtils.isMapType; -import static io.trino.plugin.mongodb.TypeUtils.isRowType; import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.Chars.truncateToLengthAndTrimSpaces; @@ -275,85 +278,77 @@ else if (isJsonType(type)) { private void writeBlock(BlockBuilder output, Type type, Object value) { if (isArrayType(type)) { - if (value instanceof List) { - BlockBuilder builder = output.beginBlockEntry(); - - ((List) value).forEach(element -> - appendTo(type.getTypeParameters().get(0), element, builder)); - - output.closeEntry(); + if (value instanceof List list) { + ((ArrayBlockBuilder) output).buildEntry(elementBuilder -> list.forEach(element -> appendTo(type.getTypeParameters().get(0), element, elementBuilder))); return; } } - else if (isMapType(type)) { + else if (type instanceof MapType mapType) { if (value instanceof List) { - BlockBuilder builder = output.beginBlockEntry(); - for (Object element : (List) value) { - if (!(element instanceof Map document)) { - continue; - } + ((MapBlockBuilder) output).buildEntry((keyBuilder, valueBuilder) -> { + for (Object element : (List) value) { + if (!(element instanceof Map document)) { + continue; + } - if (document.containsKey("key") && document.containsKey("value")) { - appendTo(type.getTypeParameters().get(0), document.get("key"), builder); - appendTo(type.getTypeParameters().get(1), document.get("value"), builder); + if (document.containsKey("key") && document.containsKey("value")) { + appendTo(mapType.getKeyType(), document.get("key"), keyBuilder); + appendTo(mapType.getValueType(), document.get("value"), valueBuilder); + } } - } - - output.closeEntry(); + }); return; } if (value instanceof Map document) { - BlockBuilder builder = output.beginBlockEntry(); - for (Map.Entry entry : document.entrySet()) { - appendTo(type.getTypeParameters().get(0), entry.getKey(), builder); - appendTo(type.getTypeParameters().get(1), entry.getValue(), builder); - } - output.closeEntry(); + ((MapBlockBuilder) output).buildEntry((keyBuilder, valueBuilder) -> { + for (Map.Entry entry : document.entrySet()) { + appendTo(mapType.getKeyType(), entry.getKey(), keyBuilder); + appendTo(mapType.getValueType(), entry.getValue(), valueBuilder); + } + }); return; } } - else if (isRowType(type)) { + else if (type instanceof RowType rowType) { if (value instanceof Map mapValue) { - BlockBuilder builder = output.beginBlockEntry(); - - for (int i = 0; i < type.getTypeSignature().getParameters().size(); i++) { - TypeSignatureParameter parameter = type.getTypeSignature().getParameters().get(i); - String fieldName = parameter.getNamedTypeSignature().getName().orElse("field" + i); - appendTo(type.getTypeParameters().get(i), mapValue.get(fieldName), builder); - } - output.closeEntry(); + ((RowBlockBuilder) output).buildEntry(fieldBuilders -> { + for (int i = 0; i < type.getTypeSignature().getParameters().size(); i++) { + TypeSignatureParameter parameter = type.getTypeSignature().getParameters().get(i); + String fieldName = parameter.getNamedTypeSignature().getName().orElse("field" + i); + appendTo(type.getTypeParameters().get(i), mapValue.get(fieldName), fieldBuilders.get(i)); + } + }); return; } if (value instanceof DBRef dbRefValue) { - BlockBuilder builder = output.beginBlockEntry(); - checkState(type.getTypeParameters().size() == 3, "DBRef should have 3 fields : %s", type); - for (int i = 0; i < type.getTypeSignature().getParameters().size(); i++) { - TypeSignatureParameter parameter = type.getTypeSignature().getParameters().get(i); - Type fieldType = type.getTypeParameters().get(i); - String fieldName = parameter.getNamedTypeSignature().getName().orElseThrow(); - switch (fieldName) { - case DATABASE_NAME -> appendTo(fieldType, dbRefValue.getDatabaseName(), builder); - case COLLECTION_NAME -> appendTo(fieldType, dbRefValue.getCollectionName(), builder); - case ID -> appendTo(fieldType, dbRefValue.getId(), builder); - default -> throw new TrinoException(GENERIC_INTERNAL_ERROR, "Unexpected field name for DBRef: " + fieldName); + ((RowBlockBuilder) output).buildEntry(fieldBuilders -> { + for (int i = 0; i < type.getTypeSignature().getParameters().size(); i++) { + TypeSignatureParameter parameter = type.getTypeSignature().getParameters().get(i); + Type fieldType = type.getTypeParameters().get(i); + String fieldName = parameter.getNamedTypeSignature().getName().orElseThrow(); + BlockBuilder builder = fieldBuilders.get(i); + switch (fieldName) { + case DATABASE_NAME -> appendTo(fieldType, dbRefValue.getDatabaseName(), builder); + case COLLECTION_NAME -> appendTo(fieldType, dbRefValue.getCollectionName(), builder); + case ID -> appendTo(fieldType, dbRefValue.getId(), builder); + default -> throw new TrinoException(GENERIC_INTERNAL_ERROR, "Unexpected field name for DBRef: " + fieldName); + } } - } - - output.closeEntry(); + }); return; } if (value instanceof List listValue) { - BlockBuilder builder = output.beginBlockEntry(); - for (int index = 0; index < type.getTypeParameters().size(); index++) { - if (index < listValue.size()) { - appendTo(type.getTypeParameters().get(index), listValue.get(index), builder); + ((RowBlockBuilder) output).buildEntry(fieldBuilders -> { + for (int index = 0; index < type.getTypeParameters().size(); index++) { + if (index < listValue.size()) { + appendTo(type.getTypeParameters().get(index), listValue.get(index), fieldBuilders.get(index)); + } + else { + fieldBuilders.get(index).appendNull(); + } } - else { - builder.appendNull(); - } - } - output.closeEntry(); + }); return; } } diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/ObjectIdType.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/ObjectIdType.java index eba68e379d6c..3b01c75d29d1 100644 --- a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/ObjectIdType.java +++ b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/ObjectIdType.java @@ -21,6 +21,7 @@ import io.airlift.slice.XxHash64; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import io.trino.spi.connector.ConnectorSession; import io.trino.spi.function.BlockIndex; import io.trino.spi.function.BlockPosition; @@ -89,8 +90,7 @@ public void appendTo(Block block, int position, BlockBuilder blockBuilder) blockBuilder.appendNull(); } else { - block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder); - blockBuilder.closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).buildEntry(valueBuilder -> block.writeSliceTo(position, 0, block.getSliceLength(position), valueBuilder)); } } @@ -109,7 +109,7 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value) @Override public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length) { - blockBuilder.writeBytes(value, offset, length).closeEntry(); + ((VariableWidthBlockBuilder) blockBuilder).writeEntry(value, offset, length); } @ScalarOperator(EQUAL) diff --git a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotSegmentPageSource.java b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotSegmentPageSource.java index 23b85ff664db..bfc36fcfa20a 100755 --- a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotSegmentPageSource.java +++ b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotSegmentPageSource.java @@ -42,6 +42,9 @@ import static io.trino.plugin.pinot.PinotErrorCode.PINOT_DECODE_ERROR; import static io.trino.plugin.pinot.PinotErrorCode.PINOT_UNSUPPORTED_COLUMN_TYPE; import static io.trino.plugin.pinot.decoders.VarbinaryDecoder.toBytes; +import static io.trino.spi.type.BigintType.BIGINT; +import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.RealType.REAL; import static java.lang.Float.floatToIntBits; import static java.lang.String.format; import static java.util.Objects.requireNonNull; @@ -298,21 +301,21 @@ private Block getArrayBlock(int rowIndex, int columnIndex) int[] intArray = currentDataTable.getDataTable().getIntArray(rowIndex, columnIndex); blockBuilder = elementType.createBlockBuilder(null, intArray.length); for (int element : intArray) { - blockBuilder.writeInt(element); + INTEGER.writeInt(blockBuilder, element); } break; case LONG_ARRAY: long[] longArray = currentDataTable.getDataTable().getLongArray(rowIndex, columnIndex); blockBuilder = elementType.createBlockBuilder(null, longArray.length); for (long element : longArray) { - blockBuilder.writeLong(element); + BIGINT.writeLong(blockBuilder, element); } break; case FLOAT_ARRAY: float[] floatArray = currentDataTable.getDataTable().getFloatArray(rowIndex, columnIndex); blockBuilder = elementType.createBlockBuilder(null, floatArray.length); for (float element : floatArray) { - blockBuilder.writeInt(floatToIntBits(element)); + REAL.writeFloat(blockBuilder, element); } break; case DOUBLE_ARRAY: diff --git a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/IntegerDecoder.java b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/IntegerDecoder.java index d503b19a99e1..4597df429bcb 100644 --- a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/IntegerDecoder.java +++ b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/IntegerDecoder.java @@ -14,6 +14,7 @@ package io.trino.plugin.pinot.decoders; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.IntArrayBlockBuilder; import java.util.function.Supplier; @@ -28,7 +29,7 @@ public void decode(Supplier getter, BlockBuilder output) output.appendNull(); } else { - output.writeInt(((Number) value).intValue()); + ((IntArrayBlockBuilder) output).writeInt(((Number) value).intValue()); } } } diff --git a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/JsonDecoder.java b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/JsonDecoder.java index dcd0280be60c..baa69b5153d1 100644 --- a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/JsonDecoder.java +++ b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/JsonDecoder.java @@ -16,6 +16,7 @@ import io.airlift.slice.Slice; import io.trino.spi.TrinoException; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import java.util.function.Supplier; @@ -36,7 +37,7 @@ public void decode(Supplier getter, BlockBuilder output) } else if (value instanceof String) { Slice slice = jsonParse(utf8Slice((String) value)); - output.writeBytes(slice, 0, slice.length()).closeEntry(); + ((VariableWidthBlockBuilder) output).writeEntry(slice); } else { throw new TrinoException(TYPE_MISMATCH, format("Expected a json value of type STRING: %s [%s]", value, value.getClass().getSimpleName())); diff --git a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/RealDecoder.java b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/RealDecoder.java index ea6dd1a09cd2..2f292a9f26e3 100644 --- a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/RealDecoder.java +++ b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/RealDecoder.java @@ -17,7 +17,7 @@ import java.util.function.Supplier; -import static java.lang.Float.floatToIntBits; +import static io.trino.spi.type.RealType.REAL; public class RealDecoder implements Decoder @@ -31,10 +31,10 @@ public void decode(Supplier getter, BlockBuilder output) } else if (value instanceof String) { // Pinot returns NEGATIVE_INFINITY, POSITIVE_INFINITY as a String - output.writeInt(floatToIntBits(Float.valueOf((String) value))); + REAL.writeFloat(output, Float.parseFloat((String) value)); } else { - output.writeInt((floatToIntBits(((Number) value).floatValue()))); + REAL.writeFloat(output, ((Number) value).floatValue()); } } } diff --git a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/VarbinaryDecoder.java b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/VarbinaryDecoder.java index 472a479958f0..b09b6365b603 100644 --- a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/VarbinaryDecoder.java +++ b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/VarbinaryDecoder.java @@ -17,6 +17,7 @@ import io.airlift.slice.Slices; import io.trino.spi.TrinoException; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import org.apache.commons.codec.DecoderException; import org.apache.commons.codec.binary.Hex; @@ -37,7 +38,7 @@ public void decode(Supplier getter, BlockBuilder output) } else if (value instanceof String) { Slice slice = Slices.wrappedBuffer(toBytes((String) value)); - output.writeBytes(slice, 0, slice.length()).closeEntry(); + ((VariableWidthBlockBuilder) output).writeEntry(slice); } else { throw new TrinoException(TYPE_MISMATCH, format("Expected a string value of type VARBINARY: %s [%s]", value, value.getClass().getSimpleName())); diff --git a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/VarcharDecoder.java b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/VarcharDecoder.java index cf58f2519c24..87eb47c60a9a 100644 --- a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/VarcharDecoder.java +++ b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/decoders/VarcharDecoder.java @@ -16,6 +16,7 @@ import io.airlift.slice.Slice; import io.airlift.slice.Slices; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.VariableWidthBlockBuilder; import java.util.function.Supplier; @@ -31,7 +32,7 @@ public void decode(Supplier getter, BlockBuilder output) } else { Slice slice = Slices.utf8Slice(value.toString()); - output.writeBytes(slice, 0, slice.length()).closeEntry(); + ((VariableWidthBlockBuilder) output).writeEntry(slice); } } } diff --git a/plugin/trino-prometheus/src/main/java/io/trino/plugin/prometheus/PrometheusRecordCursor.java b/plugin/trino-prometheus/src/main/java/io/trino/plugin/prometheus/PrometheusRecordCursor.java index 34df2d0658de..f71fb46cd922 100644 --- a/plugin/trino-prometheus/src/main/java/io/trino/plugin/prometheus/PrometheusRecordCursor.java +++ b/plugin/trino-prometheus/src/main/java/io/trino/plugin/prometheus/PrometheusRecordCursor.java @@ -19,8 +19,10 @@ import io.airlift.slice.Slice; import io.airlift.slice.Slices; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; +import io.trino.spi.block.MapBlockBuilder; import io.trino.spi.connector.RecordCursor; import io.trino.spi.type.ArrayType; import io.trino.spi.type.MapType; @@ -42,6 +44,7 @@ import static com.google.common.base.Preconditions.checkState; import static io.trino.plugin.prometheus.PrometheusClient.TIMESTAMP_COLUMN_TYPE; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; +import static io.trino.spi.block.MapValueBuilder.buildMapValue; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.BooleanType.BOOLEAN; import static io.trino.spi.type.DateTimeEncoding.packDateTimeWithZone; @@ -191,25 +194,21 @@ private List prometheusResultsInStandardizedForm(List .collect(Collectors.toList()); } - static Block getBlockFromMap(Type mapType, Map map) + static Block getBlockFromMap(Type type, Map map) { // on functions like COUNT() the Type won't be a MapType - if (!(mapType instanceof MapType)) { + if (!(type instanceof MapType mapType)) { return null; } - Type keyType = mapType.getTypeParameters().get(0); - Type valueType = mapType.getTypeParameters().get(1); - - BlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1); - BlockBuilder builder = mapBlockBuilder.beginBlockEntry(); - - for (Map.Entry entry : map.entrySet()) { - writeObject(builder, keyType, entry.getKey()); - writeObject(builder, valueType, entry.getValue()); - } + Type keyType = mapType.getKeyType(); + Type valueType = mapType.getValueType(); - mapBlockBuilder.closeEntry(); - return (Block) mapType.getObject(mapBlockBuilder, 0); + return buildMapValue(mapType, map.size(), (keyBuilder, valueBuilder) -> { + map.forEach((key, value) -> { + writeObject(keyBuilder, keyType, key); + writeObject(valueBuilder, valueType, value); + }); + }); } static Map getMapFromBlock(Type type, Block block) @@ -228,19 +227,19 @@ private static void writeObject(BlockBuilder builder, Type type, Object obj) { if (type instanceof ArrayType arrayType) { Type elementType = arrayType.getElementType(); - BlockBuilder arrayBuilder = builder.beginBlockEntry(); - for (Object item : (List) obj) { - writeObject(arrayBuilder, elementType, item); - } - builder.closeEntry(); + ((ArrayBlockBuilder) builder).buildEntry(elementBuilder -> { + for (Object item : (List) obj) { + writeObject(elementBuilder, elementType, item); + } + }); } else if (type instanceof MapType mapType) { - BlockBuilder mapBlockBuilder = builder.beginBlockEntry(); - for (Map.Entry entry : ((Map) obj).entrySet()) { - writeObject(mapBlockBuilder, mapType.getKeyType(), entry.getKey()); - writeObject(mapBlockBuilder, mapType.getValueType(), entry.getValue()); - } - builder.closeEntry(); + ((MapBlockBuilder) builder).buildEntry((keyBuilder, valueBuilder) -> { + for (Map.Entry entry : ((Map) obj).entrySet()) { + writeObject(keyBuilder, mapType.getKeyType(), entry.getKey()); + writeObject(valueBuilder, mapType.getValueType(), entry.getValue()); + } + }); } else { if (BOOLEAN.equals(type) diff --git a/plugin/trino-raptor-legacy/src/main/java/io/trino/plugin/raptor/legacy/systemtables/TableMetadataSystemTable.java b/plugin/trino-raptor-legacy/src/main/java/io/trino/plugin/raptor/legacy/systemtables/TableMetadataSystemTable.java index b33f754f5923..28192417bca6 100644 --- a/plugin/trino-raptor-legacy/src/main/java/io/trino/plugin/raptor/legacy/systemtables/TableMetadataSystemTable.java +++ b/plugin/trino-raptor-legacy/src/main/java/io/trino/plugin/raptor/legacy/systemtables/TableMetadataSystemTable.java @@ -24,6 +24,7 @@ import io.trino.plugin.raptor.legacy.metadata.TableMetadataRow; import io.trino.spi.Page; import io.trino.spi.TrinoException; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.BlockBuilder; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.connector.ConnectorPageSource; @@ -198,11 +199,11 @@ private static void writeArray(BlockBuilder blockBuilder, Collection val blockBuilder.appendNull(); } else { - BlockBuilder array = blockBuilder.beginBlockEntry(); - for (String value : values) { - VARCHAR.writeSlice(array, utf8Slice(value)); - } - blockBuilder.closeEntry(); + ((ArrayBlockBuilder) blockBuilder).buildEntry(elementBuilder -> { + for (String value : values) { + VARCHAR.writeSlice(elementBuilder, utf8Slice(value)); + } + }); } } diff --git a/plugin/trino-thrift-api/src/test/java/io/trino/plugin/thrift/api/TestReadWrite.java b/plugin/trino-thrift-api/src/test/java/io/trino/plugin/thrift/api/TestReadWrite.java index 0665f1d27929..56c6a97cfac5 100644 --- a/plugin/trino-thrift-api/src/test/java/io/trino/plugin/thrift/api/TestReadWrite.java +++ b/plugin/trino-thrift-api/src/test/java/io/trino/plugin/thrift/api/TestReadWrite.java @@ -18,6 +18,7 @@ import io.airlift.stats.cardinality.HyperLogLog; import io.trino.operator.index.PageRecordSet; import io.trino.spi.Page; +import io.trino.spi.block.ArrayBlockBuilder; import io.trino.spi.block.Block; import io.trino.spi.block.BlockBuilder; import io.trino.spi.type.ArrayType; @@ -242,16 +243,16 @@ private static Slice nextHyperLogLog(Random random) private static void generateBigintArray(Random random, BlockBuilder parentBuilder) { int numberOfElements = random.nextInt(MAX_ARRAY_GENERATED_LENGTH); - BlockBuilder builder = parentBuilder.beginBlockEntry(); - for (int i = 0; i < numberOfElements; i++) { - if (random.nextDouble() < NULL_FRACTION) { - builder.appendNull(); - } - else { - builder.writeLong(random.nextLong()); + ((ArrayBlockBuilder) parentBuilder).buildEntry(elementBuilder -> { + for (int i = 0; i < numberOfElements; i++) { + if (random.nextDouble() < NULL_FRACTION) { + elementBuilder.appendNull(); + } + else { + BIGINT.writeLong(elementBuilder, random.nextLong()); + } } - } - parentBuilder.closeEntry(); + }); } private abstract static class ColumnDefinition diff --git a/plugin/trino-thrift-api/src/test/java/io/trino/plugin/thrift/api/datatypes/TestTrinoThriftBigint.java b/plugin/trino-thrift-api/src/test/java/io/trino/plugin/thrift/api/datatypes/TestTrinoThriftBigint.java index 8002605004cd..e41e35a5071b 100644 --- a/plugin/trino-thrift-api/src/test/java/io/trino/plugin/thrift/api/datatypes/TestTrinoThriftBigint.java +++ b/plugin/trino-thrift-api/src/test/java/io/trino/plugin/thrift/api/datatypes/TestTrinoThriftBigint.java @@ -174,7 +174,7 @@ private static Block longBlock(Integer... values) blockBuilder.appendNull(); } else { - blockBuilder.writeLong(value).closeEntry(); + BIGINT.writeLong(blockBuilder, value); } } return blockBuilder.build(); diff --git a/testing/trino-testing/src/main/java/io/trino/testing/StructuralTestUtil.java b/testing/trino-testing/src/main/java/io/trino/testing/StructuralTestUtil.java index 142ab680efe4..e3a77e23a14e 100644 --- a/testing/trino-testing/src/main/java/io/trino/testing/StructuralTestUtil.java +++ b/testing/trino-testing/src/main/java/io/trino/testing/StructuralTestUtil.java @@ -31,6 +31,8 @@ import java.util.List; import static com.google.common.base.Preconditions.checkArgument; +import static io.trino.spi.block.MapValueBuilder.buildMapValue; +import static io.trino.spi.block.RowValueBuilder.buildRowValue; import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER; import static io.trino.util.StructuralTestUtil.appendToBlockBuilder; @@ -90,41 +92,38 @@ public static Block arrayBlockOf(Type elementType, Object... values) public static Block mapBlockOf(Type keyType, Type valueType, Object key, Object value) { - MapType mapType = mapType(keyType, valueType); - BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 10); - BlockBuilder singleMapBlockWriter = blockBuilder.beginBlockEntry(); - appendToBlockBuilder(keyType, key, singleMapBlockWriter); - appendToBlockBuilder(valueType, value, singleMapBlockWriter); - blockBuilder.closeEntry(); - return mapType.getObject(blockBuilder, 0); + return buildMapValue( + mapType(keyType, valueType), + 1, + (keyBuilder, valueBuilder) -> { + appendToBlockBuilder(keyType, key, keyBuilder); + appendToBlockBuilder(valueType, value, valueBuilder); + }); } public static Block mapBlockOf(Type keyType, Type valueType, Object[] keys, Object[] values) { checkArgument(keys.length == values.length, "keys/values must have the same length"); - MapType mapType = mapType(keyType, valueType); - BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 10); - BlockBuilder singleMapBlockWriter = blockBuilder.beginBlockEntry(); - for (int i = 0; i < keys.length; i++) { - Object key = keys[i]; - Object value = values[i]; - appendToBlockBuilder(keyType, key, singleMapBlockWriter); - appendToBlockBuilder(valueType, value, singleMapBlockWriter); - } - blockBuilder.closeEntry(); - return mapType.getObject(blockBuilder, 0); + return buildMapValue( + mapType(keyType, valueType), + keys.length, + (keyBuilder, valueBuilder) -> { + for (int i = 0; i < keys.length; i++) { + Object key = keys[i]; + Object value = values[i]; + appendToBlockBuilder(keyType, key, keyBuilder); + appendToBlockBuilder(valueType, value, valueBuilder); + } + }); } public static Block rowBlockOf(List parameterTypes, Object... values) { - RowType rowType = RowType.anonymous(parameterTypes); - BlockBuilder blockBuilder = rowType.createBlockBuilder(null, 1); - BlockBuilder singleRowBlockWriter = blockBuilder.beginBlockEntry(); - for (int i = 0; i < values.length; i++) { - appendToBlockBuilder(parameterTypes.get(i), values[i], singleRowBlockWriter); - } - blockBuilder.closeEntry(); - return rowType.getObject(blockBuilder, 0); + return buildRowValue(RowType.anonymous(parameterTypes), fields -> { + for (int i = 0; i < values.length; i++) { + appendToBlockBuilder(parameterTypes.get(i), values[i], fields.get(i)); + } + }); } public static Block decimalArrayBlockOf(DecimalType type, BigDecimal decimal)