diff --git a/fluss-common/src/main/java/org/apache/fluss/record/ArrowReadPlan.java b/fluss-common/src/main/java/org/apache/fluss/record/ArrowReadPlan.java new file mode 100644 index 00000000000..c86598f06c3 --- /dev/null +++ b/fluss-common/src/main/java/org/apache/fluss/record/ArrowReadPlan.java @@ -0,0 +1,200 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.fluss.record; + +import org.apache.fluss.annotation.Internal; +import org.apache.fluss.compression.ArrowCompressionFactory; +import org.apache.fluss.memory.MemorySegment; +import org.apache.fluss.row.ProjectedRow; +import org.apache.fluss.row.arrow.ArrowReader; +import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocator; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.TypeLayout; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.types.pojo.Field; +import org.apache.fluss.types.RowType; +import org.apache.fluss.utils.ArrowUtils; + +import javax.annotation.Nullable; +import javax.annotation.concurrent.NotThreadSafe; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.BitSet; +import java.util.List; + +/** + * Reads Arrow record batches of one source schema. The loaded vectors, the buffers selected before + * decompression and the output row mapping are all derived from the same selected source fields, so + * they cannot disagree. + * + *

The plan owns a {@link VectorSchemaRoot} that is reused across batches: each {@link + * #load(MemorySegment, int, int)} replaces the buffers of the previous one, and {@link #clear()} + * releases them once a batch has been consumed. + */ +@Internal +@NotThreadSafe +public final class ArrowReadPlan implements AutoCloseable { + private final BufferAllocator allocator; + private final RowType loadedRowType; + private final VectorSchemaRoot root; + // null when every source field is loaded and no buffer needs to be skipped + @Nullable private final int[] nodeIndexes; + @Nullable private final int[] bufferIndexes; + private final int nodeCount; + private final int bufferCount; + // null when loaded rows can be returned as-is + @Nullable private final int[] outputMapping; + + /** + * Creates a plan reading batches of {@code sourceType} into {@code allocator}. + * + * @param sourceType the row type of the serialized batch + * @param outputToSource maps each output field to its source field (-1 if absent), or null if + * the output fields are the source fields + * @param selectedFields the output fields read by the caller; unselected ones are never loaded + */ + ArrowReadPlan( + RowType sourceType, + @Nullable int[] outputToSource, + int[] selectedFields, + BufferAllocator allocator) { + int sourceFieldCount = sourceType.getFieldCount(); + int[] mapping = outputToSource; + if (mapping == null) { + mapping = new int[sourceFieldCount]; + Arrays.setAll(mapping, i -> i); + } + BitSet sourceFields = new BitSet(sourceFieldCount); + for (int selectedField : selectedFields) { + if (mapping[selectedField] >= 0) { + sourceFields.set(mapping[selectedField]); + } + } + int[] selectedSourceFields = sourceFields.stream().toArray(); + boolean pruneBuffers = selectedSourceFields.length < sourceFieldCount; + + this.allocator = allocator; + this.loadedRowType = pruneBuffers ? sourceType.project(selectedSourceFields) : sourceType; + this.outputMapping = + outputToSource == null && !pruneBuffers + ? null + : toLoadedPositions(mapping, selectedSourceFields, sourceFieldCount); + + List flattenedFields = new ArrayList<>(); + BitSet selectedNodes = new BitSet(); + List fields = ArrowUtils.toArrowSchema(sourceType).getFields(); + for (int i = 0; i < fields.size(); i++) { + int start = flattenedFields.size(); + flattenField(fields.get(i), flattenedFields); + if (sourceFields.get(i)) { + selectedNodes.set(start, flattenedFields.size()); + } + } + BitSet selectedBuffers = new BitSet(); + int bufferIndex = 0; + for (int i = 0; i < flattenedFields.size(); i++) { + int count = TypeLayout.getTypeBufferCount(flattenedFields.get(i).getType()); + if (selectedNodes.get(i)) { + selectedBuffers.set(bufferIndex, bufferIndex + count); + } + bufferIndex += count; + } + this.nodeCount = flattenedFields.size(); + this.bufferCount = bufferIndex; + this.nodeIndexes = pruneBuffers ? selectedNodes.stream().toArray() : null; + this.bufferIndexes = pruneBuffers ? selectedBuffers.stream().toArray() : null; + this.root = VectorSchemaRoot.create(ArrowUtils.toArrowSchema(loadedRowType), allocator); + } + + /** Returns the row type of the loaded vectors: the selected source fields in source order. */ + public RowType getLoadedRowType() { + return loadedRowType; + } + + /** + * Deserializes the Arrow batch and loads only the selected fields; buffers of unselected fields + * are neither copied nor decompressed. Buffers of a previously loaded batch are replaced. + */ + public ArrowReader load(MemorySegment segment, int arrowOffset, int arrowLength) { + try (ArrowRecordBatch batch = + nodeIndexes == null + ? ArrowUtils.deserializeArrowRecordBatch( + segment, arrowOffset, arrowLength, allocator) + : ArrowUtils.deserializeProjectedArrowRecordBatch( + segment, + arrowOffset, + arrowLength, + allocator, + nodeCount, + bufferCount, + nodeIndexes, + bufferIndexes)) { + new FlussVectorLoader(root, ArrowCompressionFactory.INSTANCE).load(batch); + return ArrowUtils.createArrowReader(root, loadedRowType); + } catch (RuntimeException | Error e) { + // The root is shared across batches, so partially loaded buffers must not linger. + root.clear(); + throw e; + } + } + + /** + * Creates a row view exposing loaded rows at the output positions expected by the caller, or + * null if loaded rows already match them. Only selected output fields may be read; the others, + * like fields absent from the source schema, are reported as null. + */ + @Nullable + public ProjectedRow createOutputRow() { + return outputMapping == null ? null : ProjectedRow.from(outputMapping); + } + + /** Releases the buffers of the last loaded batch while keeping the plan reusable. */ + public void clear() { + root.clear(); + } + + @Override + public void close() { + root.close(); + } + + private static int[] toLoadedPositions( + int[] outputToSource, int[] selectedSourceFields, int sourceFieldCount) { + int[] sourceToLoaded = new int[sourceFieldCount]; + Arrays.fill(sourceToLoaded, -1); + for (int i = 0; i < selectedSourceFields.length; i++) { + sourceToLoaded[selectedSourceFields[i]] = i; + } + int[] outputToLoaded = new int[outputToSource.length]; + for (int i = 0; i < outputToSource.length; i++) { + int sourceField = outputToSource[i]; + outputToLoaded[i] = sourceField < 0 ? -1 : sourceToLoaded[sourceField]; + } + return outputToLoaded; + } + + private static void flattenField(Field field, List flattenedFields) { + flattenedFields.add(field); + for (Field child : field.getChildren()) { + flattenField(child, flattenedFields); + } + } +} diff --git a/fluss-common/src/main/java/org/apache/fluss/record/DefaultLogRecordBatch.java b/fluss-common/src/main/java/org/apache/fluss/record/DefaultLogRecordBatch.java index d4bfcb9adc0..72f23eff89b 100644 --- a/fluss-common/src/main/java/org/apache/fluss/record/DefaultLogRecordBatch.java +++ b/fluss-common/src/main/java/org/apache/fluss/record/DefaultLogRecordBatch.java @@ -25,11 +25,8 @@ import org.apache.fluss.row.ProjectedRow; import org.apache.fluss.row.arrow.ArrowReader; import org.apache.fluss.row.columnar.ColumnarRow; -import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocator; -import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorSchemaRoot; import org.apache.fluss.types.DataType; import org.apache.fluss.types.RowType; -import org.apache.fluss.utils.ArrowUtils; import org.apache.fluss.utils.CloseableIterator; import org.apache.fluss.utils.IOUtils; import org.apache.fluss.utils.MurmurHashUtils; @@ -244,21 +241,20 @@ public CloseableIterator records(ReadContext context) { int schemaId = schemaId(); long timestamp = commitTimestamp(); LogFormat logFormat = context.getLogFormat(); - RowType rowType = context.getRowType(schemaId); - ProjectedRow outputProjection = context.getOutputProjectedRow(schemaId); switch (logFormat) { case ARROW: - return columnRecordIterator( - rowType, - outputProjection, - context.getVectorSchemaRoot(schemaId), - context.getBufferAllocator(), - timestamp); + return columnRecordIterator(context.getArrowReadPlan(schemaId), timestamp); case INDEXED: - return rowRecordIterator(rowType, outputProjection, timestamp); + return rowRecordIterator( + context.getRowType(schemaId), + context.getOutputProjectedRow(schemaId), + timestamp); case COMPACTED: - return compactedRowRecordIterator(rowType, outputProjection, timestamp); + return compactedRowRecordIterator( + context.getRowType(schemaId), + context.getOutputProjectedRow(schemaId), + timestamp); default: throw new IllegalArgumentException("Unsupported log format: " + logFormat); } @@ -388,12 +384,8 @@ public void close() {} }; } - private CloseableIterator columnRecordIterator( - RowType rowType, - @Nullable ProjectedRow outputProjection, - VectorSchemaRoot root, - BufferAllocator allocator, - long timestamp) { + private CloseableIterator columnRecordIterator(ArrowReadPlan plan, long timestamp) { + ProjectedRow outputProjection = plan.createOutputRow(); boolean isAppendOnly = isAppendOnly(); int recordsDataOffset = recordsDataOffset(); if (isAppendOnly) { @@ -401,10 +393,8 @@ private CloseableIterator columnRecordIterator( // the start of the arrow data is the beginning of the batch records int arrowOffset = position + recordsDataOffset; int arrowLength = sizeInBytes() - recordsDataOffset; - ArrowReader reader = - ArrowUtils.createArrowReader( - segment, arrowOffset, arrowLength, root, allocator, rowType); - return new ArrowLogRecordIterator(root, reader, timestamp, outputProjection) { + ArrowReader reader = plan.load(segment, arrowOffset, arrowLength); + return new ArrowLogRecordIterator(plan, reader, timestamp, outputProjection) { @Override protected ChangeType getChangeType(int rowId) { return ChangeType.APPEND_ONLY; @@ -418,10 +408,8 @@ protected ChangeType getChangeType(int rowId) { new ChangeTypeVector(segment, changeTypeOffset, getRecordCount()); int arrowOffset = changeTypeOffset + changeTypeVector.sizeInBytes(); int arrowLength = sizeInBytes() - recordsDataOffset - changeTypeVector.sizeInBytes(); - ArrowReader reader = - ArrowUtils.createArrowReader( - segment, arrowOffset, arrowLength, root, allocator, rowType); - return new ArrowLogRecordIterator(root, reader, timestamp, outputProjection) { + ArrowReader reader = plan.load(segment, arrowOffset, arrowLength); + return new ArrowLogRecordIterator(plan, reader, timestamp, outputProjection) { @Override protected ChangeType getChangeType(int rowId) { return changeTypeVector.getChangeType(rowId); @@ -436,18 +424,18 @@ private boolean isAppendOnly() { /** The basic implementation for Arrow log record iterator. */ private abstract class ArrowLogRecordIterator extends LogRecordIterator { - private final VectorSchemaRoot root; + private final ArrowReadPlan plan; private final ArrowReader reader; private final long timestamp; private int rowId = 0; @Nullable private final ProjectedRow outputProjection; private ArrowLogRecordIterator( - VectorSchemaRoot root, + ArrowReadPlan plan, ArrowReader reader, long timestamp, @Nullable ProjectedRow outputProjection) { - this.root = root; + this.plan = plan; this.reader = reader; this.timestamp = timestamp; this.outputProjection = outputProjection; @@ -484,7 +472,7 @@ protected boolean ensureNoneRemaining() { public void close() { // Clear old buffers before the next batch load to avoid temporary // duplication of buffers (old + new) during loadFieldBuffers. - root.clear(); + plan.clear(); } } diff --git a/fluss-common/src/main/java/org/apache/fluss/record/LogRecordBatch.java b/fluss-common/src/main/java/org/apache/fluss/record/LogRecordBatch.java index 5d38bb2f692..0e55fe5b55b 100644 --- a/fluss-common/src/main/java/org/apache/fluss/record/LogRecordBatch.java +++ b/fluss-common/src/main/java/org/apache/fluss/record/LogRecordBatch.java @@ -20,8 +20,6 @@ import org.apache.fluss.annotation.PublicEvolving; import org.apache.fluss.metadata.LogFormat; import org.apache.fluss.row.ProjectedRow; -import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocator; -import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorSchemaRoot; import org.apache.fluss.types.RowType; import org.apache.fluss.utils.CloseableIterator; @@ -199,27 +197,22 @@ interface ReadContext { RowType getRowType(int schemaId); /** - * Gets the Arrow {@link VectorSchemaRoot} for the given schema id. The returned schema root - * is projected if the record batch is a projected {@link LogRecordBatch}. + * Gets the plan to read {@link LogFormat#ARROW} record batches of the given schema id. The + * plan decides which fields are loaded (projected if the record batch is a projected {@link + * LogRecordBatch}, or if the reader only selects some fields) and how loaded rows map to + * output rows. Client-side field selection does not change the serialized batch or its + * checksum. * - *

The schema root is used to read the Arrow records in the batch, if this is a {@link - * LogFormat#ARROW} record batch. - * - *

Lifecycle: DO NOT close the returned {@link VectorSchemaRoot} because it is - * shared across multiple batches. Each new batch load replaces the old buffers inside the - * same root. To avoid temporary buffer duplication (old and new buffers coexisting during - * the load), callers should call {@link VectorSchemaRoot#clear()} after finishing each - * batch to release old buffers immediately before the next batch is loaded. If you need to - * retain the data beyond the current batch, use {@link VectorSchemaRoot#slice(int)} to - * create an independent copy and close it after use. + *

Lifecycle: DO NOT close the returned plan because it is shared across multiple + * batches. Each new batch load replaces the old buffers inside the same plan. To avoid + * temporary buffer duplication (old and new buffers coexisting during the load), callers + * should call {@link ArrowReadPlan#clear()} after finishing each batch to release old + * buffers immediately before the next batch is loaded. * * @param schemaId The schema id of the record batch. - * @return The (maybe projected) schema root of the record batch. + * @return The read plan of the record batch. */ - VectorSchemaRoot getVectorSchemaRoot(int schemaId); - - /** Gets the buffer allocator. */ - BufferAllocator getBufferAllocator(); + ArrowReadPlan getArrowReadPlan(int schemaId); /** * If the read context defines an output projection (for example, log records may add new @@ -228,6 +221,9 @@ interface ReadContext { * The returned object is used by readers to transform or materialize rows according to the * output projection. Returns {@code null} if no output projection is configured. * + *

For {@link LogFormat#ARROW}, this is the output row of {@link #getArrowReadPlan(int)}, + * on which only the selected fields may be read. + * * @param schemaId the current row schema id * @return a {@link ProjectedRow} describing the output projection, or {@code null} if none */ diff --git a/fluss-common/src/main/java/org/apache/fluss/record/LogRecordReadContext.java b/fluss-common/src/main/java/org/apache/fluss/record/LogRecordReadContext.java index aa4a9b0a7c7..26e429eb953 100644 --- a/fluss-common/src/main/java/org/apache/fluss/record/LogRecordReadContext.java +++ b/fluss-common/src/main/java/org/apache/fluss/record/LogRecordReadContext.java @@ -20,7 +20,6 @@ import org.apache.fluss.annotation.VisibleForTesting; import org.apache.fluss.memory.MemorySegment; import org.apache.fluss.metadata.LogFormat; -import org.apache.fluss.metadata.Schema; import org.apache.fluss.metadata.SchemaGetter; import org.apache.fluss.metadata.TableInfo; import org.apache.fluss.row.InternalRow; @@ -29,12 +28,11 @@ import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.AllocationManager; import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocator; import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocatorUtil; -import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorSchemaRoot; import org.apache.fluss.types.DataType; import org.apache.fluss.types.RowType; -import org.apache.fluss.utils.ArrowUtils; import org.apache.fluss.utils.IOUtils; import org.apache.fluss.utils.Projection; +import org.apache.fluss.utils.SchemaUtil; import org.apache.fluss.utils.UnshadedArrowReadUtils; import javax.annotation.Nullable; @@ -86,7 +84,7 @@ public enum SchemaResolution { @Nullable private volatile AutoCloseable unshadedBufferAllocator; private final boolean projectionPushDowned; private final SchemaGetter schemaGetter; - private final ConcurrentHashMap vectorSchemaRootMap = + private final ConcurrentHashMap arrowReadPlans = new ConcurrentHashMap<>(); private final Object unshadedArrowResourceLock = new Object(); private final ConcurrentHashMap fieldGetterCache = @@ -324,29 +322,32 @@ public SchemaGetter getSchemaGetter() { */ @Nullable private int[] getSchemaEvolutionMapping(int schemaId) { - ProjectedRow projectedRow = getOutputProjectedRow(schemaId); - if (projectedRow == null) { + if (target == null || isProjectionPushDowned() || schemaId == target.schemaId) { return null; } - return projectedRow.getIndexMapping(); + return SchemaUtil.getIndexMapping( + schemaGetter.getSchema(schemaId), schemaGetter.getSchema(target.schemaId)); } @Override - public VectorSchemaRoot getVectorSchemaRoot(int schemaId) { + public ArrowReadPlan getArrowReadPlan(int schemaId) { if (logFormat != LogFormat.ARROW) { - throw new IllegalArgumentException( - "Only Arrow log format provides vector schema root."); + throw new IllegalArgumentException("Only Arrow log format provides Arrow read plan."); } + return arrowReadPlans.computeIfAbsent(schemaId, this::createArrowReadPlan); + } + private ArrowReadPlan createArrowReadPlan(int schemaId) { RowType rowType = getRowType(schemaId); - return vectorSchemaRootMap.computeIfAbsent( - schemaId, - (id) -> - VectorSchemaRoot.create( - ArrowUtils.toArrowSchema(rowType), bufferAllocator)); + int[] selectedFields = + target == null + ? IntStream.range(0, rowType.getFieldCount()).toArray() + : target.selectedFields; + return new ArrowReadPlan( + rowType, getSchemaEvolutionMapping(schemaId), selectedFields, bufferAllocator); } - @Override + /** Gets the buffer allocator of the {@link LogFormat#ARROW} read plans. */ public BufferAllocator getBufferAllocator() { if (logFormat != LogFormat.ARROW) { throw new IllegalArgumentException("Only Arrow log format provides buffer allocator."); @@ -367,20 +368,16 @@ public UnshadedArrowBatchAccess createUnshadedArrowBatchAccess(int schemaId) { @Nullable @Override public ProjectedRow getOutputProjectedRow(int schemaId) { - - if (target == null || isProjectionPushDowned() || schemaId == target.schemaId) { - return null; + if (logFormat == LogFormat.ARROW) { + return getArrowReadPlan(schemaId).createOutputRow(); } - - // TODO: should we cache the projection? - Schema originSchema = schemaGetter.getSchema(schemaId); - Schema expectedSchema = schemaGetter.getSchema(target.schemaId); - return ProjectedRow.from(originSchema, expectedSchema); + int[] mapping = getSchemaEvolutionMapping(schemaId); + return mapping == null ? null : ProjectedRow.from(mapping); } public void close() { - vectorSchemaRootMap.values().forEach(VectorSchemaRoot::close); - vectorSchemaRootMap.clear(); + arrowReadPlans.values().forEach(ArrowReadPlan::close); + arrowReadPlans.clear(); if (bufferAllocator != null) { bufferAllocator.close(); } diff --git a/fluss-common/src/main/java/org/apache/fluss/utils/ArrowUtils.java b/fluss-common/src/main/java/org/apache/fluss/utils/ArrowUtils.java index 4219d45ffd8..606f043a118 100644 --- a/fluss-common/src/main/java/org/apache/fluss/utils/ArrowUtils.java +++ b/fluss-common/src/main/java/org/apache/fluss/utils/ArrowUtils.java @@ -63,6 +63,9 @@ import org.apache.fluss.row.columnar.ColumnVector; import org.apache.fluss.row.columnar.VectorizedColumnBatch; import org.apache.fluss.shaded.arrow.com.google.flatbuffers.FlatBufferBuilder; +import org.apache.fluss.shaded.arrow.org.apache.arrow.flatbuf.BodyCompression; +import org.apache.fluss.shaded.arrow.org.apache.arrow.flatbuf.Buffer; +import org.apache.fluss.shaded.arrow.org.apache.arrow.flatbuf.FieldNode; import org.apache.fluss.shaded.arrow.org.apache.arrow.flatbuf.MessageHeader; import org.apache.fluss.shaded.arrow.org.apache.arrow.flatbuf.RecordBatch; import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.ArrowBuf; @@ -100,6 +103,7 @@ import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.ipc.message.FBSerializables; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.ipc.message.IpcOption; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.ipc.message.MessageMetadataResult; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.types.DateUnit; import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.types.FloatingPointPrecision; @@ -168,26 +172,98 @@ public static ArrowReader createArrowReader( VectorSchemaRoot schemaRoot, BufferAllocator allocator, RowType rowType) { + try (ArrowRecordBatch batch = + deserializeArrowRecordBatch(segment, arrowOffset, arrowLength, allocator)) { + new FlussVectorLoader(schemaRoot, ArrowCompressionFactory.INSTANCE).load(batch); + } + return createArrowReader(schemaRoot, rowType); + } + + /** Creates an {@link ArrowReader} over the vectors already loaded into the schema root. */ + public static ArrowReader createArrowReader(VectorSchemaRoot loadedRoot, RowType rowType) { + List fieldVectors = loadedRoot.getFieldVectors(); + ColumnVector[] columnVectors = new ColumnVector[fieldVectors.size()]; + for (int i = 0; i < columnVectors.length; i++) { + columnVectors[i] = createArrowColumnVector(fieldVectors.get(i), rowType.getTypeAt(i)); + } + return new ArrowReader(columnVectors, loadedRoot.getRowCount()); + } + + /** + * Deserializes an {@link ArrowRecordBatch} from the specified memory segment. The buffers are + * still compressed; the caller owns the returned batch and must close it. + */ + public static ArrowRecordBatch deserializeArrowRecordBatch( + MemorySegment segment, int arrowOffset, int arrowLength, BufferAllocator allocator) { ByteBuffer arrowBatchBuffer = segment.wrap(arrowOffset, arrowLength); try (ReadChannel channel = - new ReadChannel(new ByteBufferReadableChannel(arrowBatchBuffer)); - ArrowRecordBatch batch = deserializeRecordBatch(channel, allocator)) { - FlussVectorLoader vectorLoader = - new FlussVectorLoader(schemaRoot, ArrowCompressionFactory.INSTANCE); - vectorLoader.load(batch); - List columnVectors = new ArrayList<>(); - List fieldVectors = schemaRoot.getFieldVectors(); - for (int i = 0; i < fieldVectors.size(); i++) { - columnVectors.add( - createArrowColumnVector(fieldVectors.get(i), rowType.getTypeAt(i))); - } - return new ArrowReader( - columnVectors.toArray(new ColumnVector[0]), schemaRoot.getRowCount()); + new ReadChannel(new ByteBufferReadableChannel(arrowBatchBuffer))) { + return deserializeRecordBatch(channel, allocator); } catch (IOException e) { throw new RuntimeException("Failed to deserialize ArrowRecordBatch.", e); } } + /** + * Deserializes only the selected field nodes and buffers of an {@link ArrowRecordBatch}. + * Unselected buffers are never copied. The selected buffers are still compressed; the caller + * owns the returned batch and must close it. + * + * @param nodeCount the expected total number of field nodes in the source batch + * @param bufferCount the expected total number of buffers in the source batch + * @param nodeIndexes source field node indexes in projected schema order, including child nodes + * @param bufferIndexes source buffer indexes corresponding to the selected field nodes + */ + public static ArrowRecordBatch deserializeProjectedArrowRecordBatch( + MemorySegment segment, + int arrowOffset, + int arrowLength, + BufferAllocator allocator, + int nodeCount, + int bufferCount, + int[] nodeIndexes, + int[] bufferIndexes) { + ByteBuffer input = segment.wrap(arrowOffset, arrowLength); + try (ReadChannel channel = new ReadChannel(new ByteBufferReadableChannel(input))) { + // Read only the IPC metadata. MessageSerializer.deserializeRecordBatch would first + // allocate and copy the entire body, including all unselected columns. + MessageMetadataResult metadata = MessageSerializer.readMessage(channel); + checkArgument(metadata != null, "Missing Arrow record batch metadata."); + checkArgument( + metadata.headerType() == MessageHeader.RecordBatch, + "Expected Arrow record batch, but found message type %s.", + metadata.headerType()); + long bodyLength = metadata.getMessageBodyLength(); + checkArgument( + bodyLength >= 0 && bodyLength <= input.remaining(), + "Invalid Arrow body length %s: only %s bytes available.", + bodyLength, + input.remaining()); + RecordBatch batch = (RecordBatch) metadata.getMessage().header(new RecordBatch()); + checkArgument( + batch.nodesLength() == nodeCount, + "Expected %s Arrow field nodes, but found %s.", + nodeCount, + batch.nodesLength()); + checkArgument( + batch.buffersLength() == bufferCount, + "Expected %s Arrow buffers, but found %s.", + bufferCount, + batch.buffersLength()); + return copySelectedBuffers( + segment, + input.position(), + bodyLength, + batch, + allocator, + nodeIndexes, + bufferIndexes); + } catch (IOException e) { + throw new FlussRuntimeException( + "Failed to deserialize projected Arrow record batch.", e); + } + } + /** * Serialize metadata of a {@link ArrowRecordBatch} into write channel. This avoids to create an * instance of {@link ArrowRecordBatch}. @@ -265,6 +341,67 @@ public static long estimateArrowBodyLength(VectorSchemaRoot root) { // ------------------------------------------------------------------------------------------ + private static ArrowRecordBatch copySelectedBuffers( + MemorySegment segment, + int bodyOffset, + long bodyLength, + RecordBatch batch, + BufferAllocator allocator, + int[] nodeIndexes, + int[] bufferIndexes) { + checkArgument( + batch.length() >= 0 && batch.length() <= Integer.MAX_VALUE, + "Invalid Arrow record count: %s.", + batch.length()); + List nodes = new ArrayList<>(nodeIndexes.length); + for (int index : nodeIndexes) { + FieldNode node = batch.nodes(index); + nodes.add(new ArrowFieldNode(node.length(), node.nullCount())); + } + + int[] offsets = new int[bufferIndexes.length]; + int[] lengths = new int[bufferIndexes.length]; + long selectedBodyLength = 0; + for (int i = 0; i < bufferIndexes.length; i++) { + Buffer buffer = batch.buffers(bufferIndexes[i]); + long offset = buffer.offset(); + long length = buffer.length(); + checkArgument( + offset >= 0 && length >= 0 && offset <= bodyLength - length, + "Arrow buffer %s range [%s, %s) exceeds body length %s.", + bufferIndexes[i], + offset, + offset + length, + bodyLength); + offsets[i] = (int) offset; + lengths[i] = (int) length; + selectedBodyLength += DataSizeRoundingUtil.roundUpTo8Multiple(length); + } + checkArgument( + selectedBodyLength <= Integer.MAX_VALUE, + "Projected Arrow body is too large: %s bytes.", + selectedBodyLength); + BodyCompression compression = batch.compression(); + ArrowBodyCompression bodyCompression = + compression == null + ? NoCompressionCodec.DEFAULT_BODY_COMPRESSION + : new ArrowBodyCompression(compression.codec(), compression.method()); + + // Use one compact allocation and keep each selected buffer aligned. The record batch + // retains its slices, so closing this allocation releases only our initial reference. + try (ArrowBuf body = allocator.buffer(selectedBodyLength)) { + List buffers = new ArrayList<>(bufferIndexes.length); + long offset = 0; + for (int i = 0; i < bufferIndexes.length; i++) { + int length = lengths[i]; + body.setBytes(offset, segment.wrap(bodyOffset + offsets[i], length)); + buffers.add(body.slice(offset, length)); + offset += DataSizeRoundingUtil.roundUpTo8Multiple((long) length); + } + return new ArrowRecordBatch((int) batch.length(), nodes, buffers, bodyCompression); + } + } + private static List flattenFields(List fields) { List allFields = new ArrayList<>(); for (Field f : fields) { diff --git a/fluss-common/src/test/java/org/apache/fluss/record/ArrowReadPlanTest.java b/fluss-common/src/test/java/org/apache/fluss/record/ArrowReadPlanTest.java new file mode 100644 index 00000000000..f722f5237c1 --- /dev/null +++ b/fluss-common/src/test/java/org/apache/fluss/record/ArrowReadPlanTest.java @@ -0,0 +1,329 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.fluss.record; + +import org.apache.fluss.compression.ArrowCompressionFactory; +import org.apache.fluss.compression.ArrowCompressionType; +import org.apache.fluss.compression.ZstdArrowCompressionCodec; +import org.apache.fluss.memory.MemorySegment; +import org.apache.fluss.row.ProjectedRow; +import org.apache.fluss.row.arrow.ArrowReader; +import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.ArrowBuf; +import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocator; +import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.OutOfMemoryException; +import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.RootAllocator; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.IntVector; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VarCharVector; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.VectorUnloader; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.ipc.WriteChannel; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.ipc.message.ArrowBuffer; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.ipc.message.IpcOption; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.ipc.message.MessageSerializer; +import org.apache.fluss.shaded.arrow.org.apache.arrow.vector.types.MetadataVersion; +import org.apache.fluss.types.DataTypes; +import org.apache.fluss.types.RowType; +import org.apache.fluss.utils.ArrowUtils; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; + +import java.io.ByteArrayOutputStream; +import java.nio.channels.Channels; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import java.util.stream.Stream; + +import static org.apache.fluss.compression.ArrowCompressionFactory.toArrowCompressionCodecType; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link ArrowReadPlan}. */ +class ArrowReadPlanTest { + private static final RowType ROW_TYPE = + DataTypes.ROW( + DataTypes.FIELD("id", DataTypes.INT()), + DataTypes.FIELD("discarded", DataTypes.STRING()), + DataTypes.FIELD("value", DataTypes.STRING())); + + @ParameterizedTest + @MethodSource("projectionReadModes") + void testOnlySelectedBuffersAreCopied( + ArrowCompressionType compression, boolean offHeap, boolean legacyIpc) throws Exception { + byte[] bytes = serializedWideBatch(compression, legacyIpc); + // The discarded column alone exceeds the allocator limit, even when compressed. Reading + // the full body before pruning would fail, but the two selected columns easily fit. + assertThat(bytes.length).isGreaterThan(64 * 1024); + int arrowOffset = 17; + MemorySegment segment = + offHeap + ? MemorySegment.allocateOffHeapMemory(bytes.length + 40) + : MemorySegment.allocateHeapMemory(bytes.length + 40); + try (BufferAllocator allocator = new RootAllocator(16 * 1024); + ArrowReadPlan plan = + new ArrowReadPlan(ROW_TYPE, null, new int[] {2, 0}, allocator)) { + segment.put(arrowOffset, bytes); + ArrowReader reader = plan.load(segment, arrowOffset, bytes.length); + // The loaded vectors must own their data independently of the input segment. + segment.put(arrowOffset, new byte[bytes.length]); + assertThat(reader.getRowCount()).isEqualTo(64); + for (int i = 0; i < reader.getRowCount(); i++) { + assertThat(reader.read(i).getInt(0)).isEqualTo(i); + assertThat(reader.read(i).getString(1).toString()).isEqualTo("value-" + i); + } + plan.clear(); + assertThat(allocator.getAllocatedMemory()).isZero(); + } finally { + segment.free(); + } + } + + @ParameterizedTest + @EnumSource(ArrowCompressionType.class) + void testEmptyProjectionDoesNotAllocateBody(ArrowCompressionType compression) throws Exception { + byte[] bytes = serializedWideBatch(compression, false); + try (BufferAllocator allocator = new RootAllocator(0); + ArrowReadPlan plan = new ArrowReadPlan(ROW_TYPE, null, new int[0], allocator)) { + ArrowReader reader = plan.load(MemorySegment.wrap(bytes), 0, bytes.length); + assertThat(reader.getRowCount()).isEqualTo(64); + assertThat(plan.getLoadedRowType().getFieldCount()).isZero(); + assertThat(allocator.getAllocatedMemory()).isZero(); + } + } + + @Test + void testTruncatedBodyIsRejectedBeforeCopying() throws Exception { + byte[] bytes = serializedWideBatch(ArrowCompressionType.NONE, false); + try (BufferAllocator allocator = new RootAllocator(); + ArrowReadPlan plan = new ArrowReadPlan(ROW_TYPE, null, new int[] {0}, allocator)) { + // Keep the complete body in the backing segment. The supplied length must bound reads. + assertThatThrownBy(() -> plan.load(MemorySegment.wrap(bytes), 0, bytes.length - 1)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Invalid Arrow body length"); + assertThat(allocator.getAllocatedMemory()).isZero(); + ArrowReader reader = plan.load(MemorySegment.wrap(bytes), 0, bytes.length); + assertThat(reader.read(63).getInt(0)).isEqualTo(63); + } + } + + @Test + void testFailedCopyCanBeRetried() throws Exception { + byte[] bytes = serializedWideBatch(ArrowCompressionType.NONE, false); + try (BufferAllocator allocator = new RootAllocator(0); + ArrowReadPlan plan = new ArrowReadPlan(ROW_TYPE, null, new int[] {0}, allocator)) { + assertThatThrownBy(() -> plan.load(MemorySegment.wrap(bytes), 0, bytes.length)) + .isInstanceOf(OutOfMemoryException.class); + assertThat(allocator.getAllocatedMemory()).isZero(); + allocator.setLimit(16 * 1024); + ArrowReader reader = plan.load(MemorySegment.wrap(bytes), 0, bytes.length); + assertThat(reader.read(63).getInt(0)).isEqualTo(63); + plan.clear(); + assertThat(allocator.getAllocatedMemory()).isZero(); + } + } + + @Test + void testSelectedBufferCannotReadPastDeclaredBody() throws Exception { + byte[] bytes = serializedWideBatch(ArrowCompressionType.NONE, false); + byte[] invalidBytes; + try (BufferAllocator allocator = new RootAllocator(); + ArrowRecordBatch batch = + ArrowUtils.deserializeArrowRecordBatch( + MemorySegment.wrap(bytes), 0, bytes.length, allocator); + ByteArrayOutputStream output = new ByteArrayOutputStream()) { + List layout = new ArrayList<>(batch.getBuffersLayout()); + layout.set(1, new ArrowBuffer(batch.computeBodyLength(), 8)); + WriteChannel channel = new WriteChannel(Channels.newChannel(output)); + ArrowUtils.serializeArrowRecordBatchMetadata( + channel, + batch.getLength(), + batch.getNodes(), + layout, + batch.getBodyCompression(), + batch.computeBodyLength()); + MessageSerializer.writeBatchBuffers(channel, batch); + // The segment has enough trailing bytes, but the selected buffer must fit in the body + // declared by this message rather than borrowing bytes from the next message. + invalidBytes = Arrays.copyOf(output.toByteArray(), output.size() + 8); + } + try (BufferAllocator allocator = new RootAllocator(0); + ArrowReadPlan plan = new ArrowReadPlan(ROW_TYPE, null, new int[] {0}, allocator)) { + assertThatThrownBy( + () -> + plan.load( + MemorySegment.wrap(invalidBytes), + 0, + invalidBytes.length)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Arrow buffer 1 range") + .hasMessageContaining("exceeds body length"); + assertThat(allocator.getAllocatedMemory()).isZero(); + } + } + + @Test + void testUnselectedBuffersAreNotDecompressed() throws Exception { + // The discarded column's data buffer contains an invalid ZSTD frame. Reading the selected + // columns must not invoke the decompressor on that buffer. + byte[] bytes = serializedBatch(4); + try (BufferAllocator allocator = new RootAllocator(); + ArrowReadPlan plan = + new ArrowReadPlan(ROW_TYPE, null, new int[] {2, 0}, allocator)) { + assertThat(plan.getLoadedRowType().getFieldNames()).containsExactly("id", "value"); + ArrowReader reader = plan.load(MemorySegment.wrap(bytes), 0, bytes.length); + assertThat(reader.getRowCount()).isEqualTo(64); + for (int i = 0; i < reader.getRowCount(); i++) { + assertThat(reader.read(i).getInt(0)).isEqualTo(i); + assertThat(reader.read(i).getString(1).toString()).isEqualTo(value()); + } + ProjectedRow output = plan.createOutputRow(); + assertThat(output).isNotNull(); + output.replaceRow(reader.read(1)); + assertThat(output.getInt(0)).isEqualTo(1); + assertThat(output.isNullAt(1)).isTrue(); + assertThat(output.getString(2).toString()).isEqualTo(value()); + plan.clear(); + assertThat(allocator.getAllocatedMemory()).isZero(); + } + } + + @Test + void testAllSelectedFieldsLoadWithoutOutputRow() throws Exception { + byte[] bytes = serializedBatch(-1); + try (BufferAllocator allocator = new RootAllocator(); + ArrowReadPlan plan = + new ArrowReadPlan(ROW_TYPE, null, new int[] {2, 1, 0}, allocator)) { + assertThat(plan.getLoadedRowType()).isEqualTo(ROW_TYPE); + assertThat(plan.createOutputRow()).isNull(); + ArrowReader reader = plan.load(MemorySegment.wrap(bytes), 0, bytes.length); + assertThat(reader.read(63).getInt(0)).isEqualTo(63); + assertThat(reader.read(63).getString(1).toString()).isEqualTo(value()); + plan.clear(); + assertThat(allocator.getAllocatedMemory()).isZero(); + } + } + + @Test + void testFailedLoadReleasesSelectedBuffers() throws Exception { + // The id field fits, but the value field requires a 64 KiB decompression buffer. + byte[] bytes = serializedBatch(-1); + try (BufferAllocator allocator = new RootAllocator(32 * 1024); + ArrowReadPlan plan = + new ArrowReadPlan(ROW_TYPE, null, new int[] {0, 2}, allocator)) { + assertThatThrownBy(() -> plan.load(MemorySegment.wrap(bytes), 0, bytes.length)) + .isInstanceOf(IllegalArgumentException.class) + .hasRootCauseInstanceOf(OutOfMemoryException.class); + assertThat(allocator.getAllocatedMemory()).isZero(); + + allocator.setLimit(Long.MAX_VALUE); + ArrowReader reader = plan.load(MemorySegment.wrap(bytes), 0, bytes.length); + assertThat(reader.read(63).getInt(0)).isEqualTo(63); + plan.clear(); + assertThat(allocator.getAllocatedMemory()).isZero(); + } + } + + private static Stream projectionReadModes() { + List modes = new ArrayList<>(); + for (ArrowCompressionType compression : ArrowCompressionType.values()) { + for (boolean offHeap : new boolean[] {false, true}) { + for (boolean legacyIpc : new boolean[] {false, true}) { + modes.add(Arguments.of(compression, offHeap, legacyIpc)); + } + } + } + return modes.stream(); + } + + private static byte[] serializedWideBatch(ArrowCompressionType compression, boolean legacyIpc) + throws Exception { + try (BufferAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = + VectorSchemaRoot.create(ArrowUtils.toArrowSchema(ROW_TYPE), allocator)) { + Random random = new Random(42); + byte[] discarded = new byte[4096]; + for (int i = 0; i < 64; i++) { + random.nextBytes(discarded); + for (int j = 0; j < discarded.length; j++) { + discarded[j] = (byte) (' ' + (discarded[j] & 0x3f)); + } + ((IntVector) root.getVector(0)).setSafe(i, i); + ((VarCharVector) root.getVector(1)).setSafe(i, discarded); + ((VarCharVector) root.getVector(2)) + .setSafe(i, ("value-" + i).getBytes(StandardCharsets.UTF_8)); + } + root.setRowCount(64); + try (ArrowRecordBatch batch = + new VectorUnloader( + root, + true, + ArrowCompressionFactory.INSTANCE.createCodec( + toArrowCompressionCodecType(compression)), + true) + .getRecordBatch(); + ByteArrayOutputStream output = new ByteArrayOutputStream()) { + MessageSerializer.serialize( + new WriteChannel(Channels.newChannel(output)), + batch, + new IpcOption(legacyIpc, MetadataVersion.V5)); + return output.toByteArray(); + } + } + } + + private static byte[] serializedBatch(int corruptBuffer) throws Exception { + try (BufferAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = + VectorSchemaRoot.create(ArrowUtils.toArrowSchema(ROW_TYPE), allocator)) { + byte[] value = value().getBytes(StandardCharsets.UTF_8); + for (int i = 0; i < 64; i++) { + ((IntVector) root.getVector(0)).setSafe(i, i); + ((VarCharVector) root.getVector(1)).setSafe(i, value); + ((VarCharVector) root.getVector(2)).setSafe(i, value); + } + root.setRowCount(64); + try (ArrowRecordBatch batch = + new VectorUnloader(root, true, new ZstdArrowCompressionCodec(), true) + .getRecordBatch(); + ByteArrayOutputStream output = new ByteArrayOutputStream()) { + if (corruptBuffer >= 0) { + ArrowBuf buffer = batch.getBuffers().get(corruptBuffer); + assertThat(buffer.getLong(0)).isPositive(); + buffer.setByte(8, buffer.getByte(8) ^ 0xff); + } + MessageSerializer.serialize(new WriteChannel(Channels.newChannel(output)), batch); + return output.toByteArray(); + } + } + } + + private static String value() { + char[] chars = new char[1024]; + Arrays.fill(chars, 'a'); + return new String(chars); + } +} diff --git a/fluss-common/src/test/java/org/apache/fluss/record/LogRecordReadContextTest.java b/fluss-common/src/test/java/org/apache/fluss/record/LogRecordReadContextTest.java index 935eeb0d81e..4f1919fe6f1 100644 --- a/fluss-common/src/test/java/org/apache/fluss/record/LogRecordReadContextTest.java +++ b/fluss-common/src/test/java/org/apache/fluss/record/LogRecordReadContextTest.java @@ -17,11 +17,15 @@ package org.apache.fluss.record; +import org.apache.fluss.compression.ArrowCompressionInfo; +import org.apache.fluss.compression.ArrowCompressionType; import org.apache.fluss.memory.ManagedPagedOutputView; import org.apache.fluss.memory.TestingMemorySegmentPool; +import org.apache.fluss.memory.UnmanagedPagedOutputView; import org.apache.fluss.metadata.KvFormat; import org.apache.fluss.metadata.LogFormat; import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.SchemaInfo; import org.apache.fluss.metadata.TableDescriptor; import org.apache.fluss.metadata.TableInfo; import org.apache.fluss.metadata.TablePath; @@ -39,9 +43,13 @@ import org.apache.fluss.row.InternalRow.FieldGetter; import org.apache.fluss.row.TimestampLtz; import org.apache.fluss.row.TimestampNtz; +import org.apache.fluss.row.arrow.ArrowWriter; +import org.apache.fluss.row.arrow.ArrowWriterPool; import org.apache.fluss.row.compacted.CompactedRow; import org.apache.fluss.row.encode.RowEncoder; import org.apache.fluss.row.indexed.IndexedRow; +import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.BufferAllocator; +import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.RootAllocator; import org.apache.fluss.types.ArrayType; import org.apache.fluss.types.DataField; import org.apache.fluss.types.DataType; @@ -51,6 +59,7 @@ import org.apache.fluss.types.MapType; import org.apache.fluss.types.RowType; import org.apache.fluss.utils.CloseableIterator; +import org.apache.fluss.utils.Projection; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -100,6 +109,240 @@ class LogRecordReadContextTest { private static final RowType ROW_TYPE = allTypesRowType(); + @ParameterizedTest + @EnumSource(ArrowCompressionType.class) + void testArrowProjectionLoadsOnlySelectedColumns(ArrowCompressionType compression) + throws Exception { + RowType rowType = + DataTypes.ROW( + DataTypes.FIELD("id", DataTypes.INT()), + DataTypes.FIELD("unused", DataTypes.STRING()), + DataTypes.FIELD("value", DataTypes.STRING())); + Schema schema = Schema.newBuilder().fromRowType(rowType).build(); + GenericRow row = + GenericRow.of( + 42, BinaryString.fromString("unused"), BinaryString.fromString("selected")); + for (boolean appendOnly : new boolean[] {true, false}) { + LogRecordBatch batch = + buildArrowBatch( + rowType, + DEFAULT_SCHEMA_ID, + compression, + appendOnly, + row, + GenericRow.of(43, null, null)); + try (LogRecordReadContext context = + LogRecordReadContext.createReadContext( + tableInfo(LogFormat.ARROW, schema, DEFAULT_SCHEMA_ID), + true, + SchemaResolution.TARGET, + Projection.of(new int[] {2, 0}), + new TestingSchemaGetter(DEFAULT_SCHEMA_ID, schema), + null)) { + assertThat(context.isProjectionPushDowned()).isFalse(); + batch.ensureValid(); + try (CloseableIterator records = batch.records(context)) { + FieldGetter[] getters = context.getSelectedFieldGetters(DEFAULT_SCHEMA_ID); + LogRecord first = records.next(); + assertThat(getters[0].getFieldOrNull(first.getRow()).toString()) + .isEqualTo("selected"); + assertThat(getters[1].getFieldOrNull(first.getRow())).isEqualTo(42); + assertThat(first.getChangeType()) + .isEqualTo(appendOnly ? APPEND_ONLY : ChangeType.UPDATE_AFTER); + InternalRow second = records.next().getRow(); + assertThat(getters[0].getFieldOrNull(second)).isNull(); + assertThat(getters[1].getFieldOrNull(second)).isEqualTo(43); + assertThat(records.hasNext()).isFalse(); + assertThat( + context.getArrowReadPlan(DEFAULT_SCHEMA_ID) + .getLoadedRowType() + .getFieldNames()) + .containsExactly("id", "value"); + } + assertThat(context.getBufferAllocator().getAllocatedMemory()).isZero(); + batch.ensureValid(); + } + } + } + + @ParameterizedTest + @EnumSource(ArrowCompressionType.class) + void testArrowProjectionWithNestedColumns(ArrowCompressionType compression) throws Exception { + Schema schema = Schema.newBuilder().fromRowType(ROW_TYPE).build(); + GenericRow expected = (GenericRow) sampleValue(ROW_TYPE); + int count = ROW_TYPE.getFieldCount(); + int[] selected = {count - 1, count - 2, count - 3, 2, 1, 0}; + LogRecordBatch batch = + buildArrowBatch( + ROW_TYPE, + DEFAULT_SCHEMA_ID, + compression, + true, + expected, + new GenericRow(count)); + try (LogRecordReadContext context = + LogRecordReadContext.createReadContext( + tableInfo(LogFormat.ARROW, schema, DEFAULT_SCHEMA_ID), + true, + SchemaResolution.TARGET, + Projection.of(selected), + new TestingSchemaGetter(DEFAULT_SCHEMA_ID, schema), + null)) { + // Reusing the root must not retain buffers from the preceding batch. + for (int iteration = 0; iteration < 2; iteration++) { + try (CloseableIterator records = batch.records(context)) { + FieldGetter[] getters = context.getSelectedFieldGetters(DEFAULT_SCHEMA_ID); + InternalRow first = records.next().getRow(); + for (int i = 0; i < selected.length; i++) { + DataType type = ROW_TYPE.getTypeAt(selected[i]); + assertThat(toJava(getters[i].getFieldOrNull(first), type)) + .isEqualTo(toJava(expected.getField(selected[i]), type)); + } + InternalRow second = records.next().getRow(); + for (FieldGetter getter : getters) { + assertThat(getter.getFieldOrNull(second)).isNull(); + } + assertThat(records.hasNext()).isFalse(); + assertThat( + context.getArrowReadPlan(DEFAULT_SCHEMA_ID) + .getLoadedRowType() + .getFieldCount()) + .isEqualTo(selected.length); + } + assertThat(context.getBufferAllocator().getAllocatedMemory()).isZero(); + } + } + } + + @ParameterizedTest + @EnumSource(SchemaResolution.class) + void testArrowProjectionWithSchemaEvolution(SchemaResolution resolution) throws Exception { + Schema source = + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("unused", DataTypes.STRING()) + .column("value", DataTypes.STRING()) + .build(); + Schema target = + Schema.newBuilder() + .fromColumns( + Arrays.asList( + new Schema.Column("added", DataTypes.STRING(), null, 100), + new Schema.Column( + "renamed", + DataTypes.STRING(), + null, + source.getColumn("value").getColumnId()), + source.getColumn("id"))) + .build(); + TestingSchemaGetter schemas = new TestingSchemaGetter(0, source); + schemas.updateLatestSchemaInfo(new SchemaInfo(target, 1)); + LogRecordBatch oldBatch = + buildArrowBatch( + source.getRowType(), + 0, + ArrowCompressionType.ZSTD, + true, + GenericRow.of( + 42, + BinaryString.fromString("skip"), + BinaryString.fromString("old"))); + LogRecordBatch newBatch = + buildArrowBatch( + target.getRowType(), + 1, + ArrowCompressionType.LZ4_FRAME, + true, + GenericRow.of( + BinaryString.fromString("new"), + BinaryString.fromString("value"), + 43)); + try (LogRecordReadContext context = + LogRecordReadContext.createReadContext( + tableInfo(LogFormat.ARROW, target, 1), + true, + resolution, + Projection.of(new int[] {1, 0}), + schemas, + null)) { + assertThat(context.getArrowReadPlan(0)).isSameAs(context.getArrowReadPlan(0)); + assertThat(context.getArrowReadPlan(0).getLoadedRowType().getFieldNames()) + .containsExactly("value"); + assertThat(context.getArrowReadPlan(1).getLoadedRowType().getFieldNames()) + .containsExactly("added", "renamed"); + FieldGetter[] getters = context.getSelectedFieldGetters(0); + try (CloseableIterator records = oldBatch.records(context)) { + InternalRow row = records.next().getRow(); + assertThat(getters[0].getFieldOrNull(row).toString()).isEqualTo("old"); + assertThat(getters[1].getFieldOrNull(row)).isNull(); + } + try (CloseableIterator records = newBatch.records(context)) { + InternalRow row = records.next().getRow(); + assertThat(getters[0].getFieldOrNull(row).toString()).isEqualTo("value"); + assertThat(getters[1].getFieldOrNull(row).toString()).isEqualTo("new"); + } + assertThat(context.getBufferAllocator().getAllocatedMemory()).isZero(); + } + + // A requested column can be absent from every physical field of a historical batch. + try (LogRecordReadContext context = + LogRecordReadContext.createReadContext( + tableInfo(LogFormat.ARROW, target, 1), + true, + resolution, + Projection.of(new int[] {0}), + schemas, + null); + CloseableIterator records = oldBatch.records(context)) { + assertThat(context.getArrowReadPlan(0).getLoadedRowType().getFieldCount()).isZero(); + assertThat( + context.getSelectedFieldGetters(0)[0].getFieldOrNull( + records.next().getRow())) + .isNull(); + assertThat(records.hasNext()).isFalse(); + } + } + + @Test + void testArrowProjectionFastPaths() { + Schema schema = + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .build(); + TableInfo table = tableInfo(LogFormat.ARROW, schema, DEFAULT_SCHEMA_ID); + TestingSchemaGetter schemas = new TestingSchemaGetter(DEFAULT_SCHEMA_ID, schema); + try (LogRecordReadContext allFields = + LogRecordReadContext.createReadContext( + table, + true, + SchemaResolution.TARGET, + Projection.of(new int[] {1, 0}), + schemas, + null); + LogRecordReadContext dynamic = + LogRecordReadContext.createReadContext( + table, true, SchemaResolution.DYNAMIC, null, schemas, null); + LogRecordReadContext serverProjected = + LogRecordReadContext.createReadContext( + table, + false, + SchemaResolution.TARGET, + Projection.of(new int[] {1}), + schemas, + null)) { + for (LogRecordReadContext context : Arrays.asList(allFields, dynamic)) { + ArrowReadPlan plan = context.getArrowReadPlan(DEFAULT_SCHEMA_ID); + assertThat(plan.getLoadedRowType()).isEqualTo(schema.getRowType()); + assertThat(plan.createOutputRow()).isNull(); + } + assertThat(serverProjected.isProjectionPushDowned()).isTrue(); + ArrowReadPlan serverPlan = serverProjected.getArrowReadPlan(DEFAULT_SCHEMA_ID); + assertThat(serverPlan.getLoadedRowType().getFieldNames()).containsExactly("b"); + assertThat(serverPlan.createOutputRow()).isNull(); + } + } + @Test void testRowTypeCoversEveryTypeRoot() { Set roots = EnumSet.noneOf(DataTypeRoot.class); @@ -336,19 +579,59 @@ private static byte[] buildBatch(LogFormat logFormat, GenericRow row) throws Exc } private static TableInfo tableInfo(LogFormat logFormat) { + return tableInfo( + logFormat, Schema.newBuilder().fromRowType(ROW_TYPE).build(), DEFAULT_SCHEMA_ID); + } + + private static TableInfo tableInfo(LogFormat logFormat, Schema schema, int schemaId) { TableDescriptor descriptor = TableDescriptor.builder() - .schema(Schema.newBuilder().fromRowType(ROW_TYPE).build()) + .schema(schema) .distributedBy(1) .logFormat(logFormat) .build(); return TableInfo.of( TablePath.of("test_db", "test_table"), 1L, - DEFAULT_SCHEMA_ID, + schemaId, descriptor, DEFAULT_REMOTE_DATA_DIR, 0L, 0L); } + + private static LogRecordBatch buildArrowBatch( + RowType rowType, + int schemaId, + ArrowCompressionType compression, + boolean appendOnly, + GenericRow... rows) + throws Exception { + try (BufferAllocator allocator = new RootAllocator(); + ArrowWriterPool pool = new ArrowWriterPool(allocator)) { + ArrowWriter writer = + pool.getOrCreateWriter( + 1L, + schemaId, + 1024 * 1024, + rowType, + new ArrowCompressionInfo(compression, 3)); + try (MemoryLogRecordsArrowBuilder builder = + MemoryLogRecordsArrowBuilder.builder( + 0L, + DEFAULT_MAGIC, + schemaId, + writer, + new UnmanagedPagedOutputView(1024 * 1024), + appendOnly)) { + for (GenericRow row : rows) { + builder.append(appendOnly ? APPEND_ONLY : ChangeType.UPDATE_AFTER, row); + } + MemoryLogRecords records = MemoryLogRecords.pointToBytesView(builder.build()); + byte[] bytes = new byte[records.sizeInBytes()]; + records.getMemorySegment().get(records.getPosition(), bytes, 0, bytes.length); + return MemoryLogRecords.pointToBytes(bytes).batches().iterator().next(); + } + } + } }