Do not slice ListVector

trinodb · May 21, 2024 · d35d036 · d35d036
1 parent 59209f5
commit d35d036
Showing 1 changed file with 4 additions and 21 deletions.
diff --git a/...n/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowToPageConverter.java b/...n/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowToPageConverter.java
@@ -28,7 +28,6 @@
 import io.trino.spi.type.Type;
 import io.trino.spi.type.VarbinaryType;
 import io.trino.spi.type.VarcharType;
-import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.BigIntVector;
 import org.apache.arrow.vector.BitVector;
@@ -48,7 +47,6 @@
 import org.apache.arrow.vector.complex.StructVector;
 import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
 import org.apache.arrow.vector.types.pojo.Schema;
-import org.apache.arrow.vector.util.TransferPair;
 
 import java.math.BigDecimal;
 import java.util.List;
@@ -76,7 +74,6 @@
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 import static org.apache.arrow.compression.CommonsCompressionFactory.INSTANCE;
-import static org.apache.arrow.vector.complex.BaseRepeatedValueVector.OFFSET_WIDTH;
 import static org.apache.arrow.vector.types.Types.MinorType.DECIMAL256;
 
 public class BigQueryArrowToPageConverter
@@ -85,14 +82,12 @@ public class BigQueryArrowToPageConverter
     private final BigQueryTypeManager typeManager;
     private final VectorSchemaRoot root;
     private final VectorLoader loader;
-    private final BufferAllocator allocator;
     private final List<Type> columnTypes;
     private final List<String> columnNames;
 
     public BigQueryArrowToPageConverter(BigQueryTypeManager typeManager, BufferAllocator allocator, Schema schema, List<BigQueryColumnHandle> columns)
     {
         this.typeManager = requireNonNull(typeManager, "typeManager is null");
-        this.allocator = requireNonNull(allocator, "allocator is null");
         this.columnTypes = requireNonNull(columns, "columns is null").stream()
                 .map(BigQueryColumnHandle::trinoType)
                 .collect(toImmutableList());
@@ -165,7 +160,7 @@ else if (javaType == LongTimestampWithTimeZone.class) {
                 writeVectorValues(output, vector, index -> writeObjectTimestampWithTimezone(output, type, vector, index), offset, length);
             }
             else if (type instanceof ArrayType arrayType) {
-                writeVectorValues(output, vector, index -> writeArrayBlock(output, arrayType, vector, index), offset, length);
+                writeVectorValues(output, vector, _ -> writeArrayBlock(output, arrayType, vector), offset, length);
             }
             else if (type instanceof RowType rowType) {
                 writeVectorValues(output, vector, index -> writeRowBlock(output, rowType, vector, index), offset, length);
@@ -234,23 +229,11 @@ private void writeObjectTimestampWithTimezone(BlockBuilder output, Type type, Fi
         type.writeObject(output, fromEpochMillisAndFraction(floorDiv(epochMicros, MICROSECONDS_PER_MILLISECOND), picosOfMillis, UTC_KEY));
     }
 
-    private void writeArrayBlock(BlockBuilder output, ArrayType arrayType, FieldVector vector, int index)
+    private void writeArrayBlock(BlockBuilder output, ArrayType arrayType, FieldVector vector)
     {
         Type elementType = arrayType.getElementType();
-        ((ArrayBlockBuilder) output).buildEntry(elementBuilder -> {
-            ArrowBuf offsetBuffer = vector.getOffsetBuffer();
-
-            int start = offsetBuffer.getInt((long) index * OFFSET_WIDTH);
-            int end = offsetBuffer.getInt((long) (index + 1) * OFFSET_WIDTH);
-
-            FieldVector innerVector = ((ListVector) vector).getDataVector();
-
-            TransferPair transferPair = innerVector.getTransferPair(allocator);
-            transferPair.splitAndTransfer(start, end - start);
-            try (FieldVector sliced = (FieldVector) transferPair.getTo()) {
-                convertType(elementBuilder, elementType, sliced, 0, sliced.getValueCount());
-            }
-        });
+        FieldVector innerVector = ((ListVector) vector).getDataVector();
+        ((ArrayBlockBuilder) output).buildEntry(elementBuilder -> convertType(elementBuilder, elementType, innerVector, 0, innerVector.getValueCount()));
     }
 
     private void writeRowBlock(BlockBuilder output, RowType rowType, FieldVector vector, int index)