diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/ArrowScanRecords.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/ArrowScanRecords.java index 048eba4e1ae..80f83463116 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/ArrowScanRecords.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/ArrowScanRecords.java @@ -17,7 +17,7 @@ package org.apache.fluss.client.table.scanner.log; -import org.apache.fluss.annotation.Internal; +import org.apache.fluss.annotation.PublicEvolving; import org.apache.fluss.metadata.TableBucket; import org.apache.fluss.record.ArrowBatchData; import org.apache.fluss.utils.AbstractIterator; @@ -38,7 +38,7 @@ *

Each {@link ArrowBatchData} holds off-heap Arrow memory. Callers should use try-with-resources * on this container to ensure all batches are released if processing fails mid-iteration. */ -@Internal +@PublicEvolving public class ArrowScanRecords implements Iterable, AutoCloseable { public static final ArrowScanRecords EMPTY = new ArrowScanRecords(Collections.emptyMap()); diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogScanner.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogScanner.java index d255e340a94..ca17e6621a0 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogScanner.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogScanner.java @@ -56,6 +56,24 @@ public interface LogScanner extends AutoCloseable { */ ScanRecords poll(Duration timeout); + /** + * Polls Arrow record batches from the tablet server without materializing individual rows. + * + *

This method is supported for tables whose log format is {@code ARROW}. Each returned batch + * owns its Arrow memory and, for primary-key tables, carries the stored per-row changelog + * types. Callers must close the returned {@link ArrowScanRecords} after processing it. + * + * @param timeout the timeout to poll + * @return the Arrow batches returned by this poll + * @throws UnsupportedOperationException if the table does not use the {@code ARROW} log format + * or the scanner implementation does not support Arrow record batch polling + * @throws IllegalStateException if the scanner is not subscribed to any buckets + */ + default ArrowScanRecords pollRecordBatch(Duration timeout) { + throw new UnsupportedOperationException( + "This LogScanner implementation does not support Arrow record batch polling"); + } + /** * Subscribe to the given table bucket in given offset dynamically. If the table bucket is * already subscribed, the offset will be updated. diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogScannerImpl.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogScannerImpl.java index 4637f87f615..c2d32960614 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogScannerImpl.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/log/LogScannerImpl.java @@ -17,12 +17,10 @@ package org.apache.fluss.client.table.scanner.log; -import org.apache.fluss.annotation.Internal; import org.apache.fluss.annotation.PublicEvolving; import org.apache.fluss.client.metadata.MetadataUpdater; import org.apache.fluss.client.metrics.ScannerMetricGroup; import org.apache.fluss.client.table.scanner.RemoteFileDownloader; -import org.apache.fluss.client.table.scanner.Scan; import org.apache.fluss.config.Configuration; import org.apache.fluss.exception.WakeupException; import org.apache.fluss.metadata.LogFormat; @@ -59,8 +57,6 @@ public class LogScannerImpl extends AbstractLogScanner implements L private final long tableId; private final boolean isPartitionedTable; private final boolean isArrowLogFormat; - private final boolean isLogTable; - private final boolean hasProjection; private final TableInfo tableInfo; private final @Nullable Projection projection; @@ -114,8 +110,6 @@ private LogScannerImpl( this.tableId = tableInfo.getTableId(); this.isPartitionedTable = tableInfo.isPartitioned(); this.isArrowLogFormat = tableInfo.getTableConfig().getLogFormat() == LogFormat.ARROW; - this.isLogTable = !tableInfo.hasPrimaryKey(); - this.hasProjection = projectedFields != null; // add this table to metadata updater. metadataUpdater.checkAndUpdateTableMetadata(Collections.singleton(tablePath)); this.metadataUpdater = metadataUpdater; @@ -164,26 +158,12 @@ protected String notSubscribedMessage() { return "LogScanner is not subscribed any buckets."; } - /** - * Polls Arrow record batches for internal callers. - * - *

This method is intentionally kept off the public {@link Scan} API surface for now. - */ - @Internal + @Override public ArrowScanRecords pollRecordBatch(Duration timeout) { if (!isArrowLogFormat) { throw new UnsupportedOperationException( "Arrow record batch polling is only supported for tables whose log format is ARROW."); } - if (!isLogTable) { - throw new UnsupportedOperationException( - "Arrow record batch polling is only supported for log tables. CDC scanning is not supported."); - } - if (hasProjection) { - throw new UnsupportedOperationException( - "Arrow record batch polling does not support projection. Please create the scanner without projection."); - } - acquireAndEnsureOpen(); try { if (!logScannerStatus.prepareToPoll()) { diff --git a/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogScannerITCase.java b/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogScannerITCase.java index 7f6700acf3f..ac27b7d7c82 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogScannerITCase.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/log/LogScannerITCase.java @@ -43,6 +43,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; +import java.nio.ByteBuffer; import java.time.Duration; import java.util.ArrayList; import java.util.Collections; @@ -506,6 +507,65 @@ void testSubscribeOutOfRangeLog() throws Exception { } } + @Test + void testPollArrowBatchesWithPrimaryKeyChangelog() throws Exception { + TablePath tablePath = TablePath.of("test_db_1", "test_arrow_batches_with_changelog"); + TableDescriptor tableDescriptor = + TableDescriptor.builder() + .schema( + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .primaryKey("a") + .build()) + .distributedBy(1) + .logFormat(LogFormat.ARROW) + .build(); + createTable(tablePath, tableDescriptor, false); + + try (Table table = conn.getTable(tablePath)) { + UpsertWriter writer = table.newUpsert().createWriter(); + writer.upsert(row(1, "old")).get(); + writer.flush(); + writer.upsert(row(1, "new")).get(); + writer.flush(); + writer.upsert(row(2, "deleted")).get(); + writer.flush(); + writer.delete(row(2, "deleted")).get(); + writer.flush(); + + ChangeType[] expectedChangeTypes = { + ChangeType.INSERT, + ChangeType.UPDATE_BEFORE, + ChangeType.UPDATE_AFTER, + ChangeType.INSERT, + ChangeType.DELETE + }; + int[] expectedKeys = {1, 1, 1, 2, 2}; + String[] expectedValues = {"old", "old", "new", "deleted", "deleted"}; + + try (LogScanner scanner = table.newScan().createLogScanner()) { + scanner.subscribeFromBeginning(0); + pollAndVerifyChangelogArrowBatches( + scanner, expectedChangeTypes, expectedKeys, expectedValues, 0); + } + + try (LogScanner scanner = table.newScan().project(new int[] {1}).createLogScanner()) { + scanner.subscribeFromBeginning(0); + pollAndVerifyProjectedChangelogArrowBatches( + scanner, expectedChangeTypes, expectedValues); + } + + // Offset 2 starts in the middle of the update batch and verifies that slicing keeps + // the change-type vector aligned with the Arrow rows. + try (LogScanner scanner = table.newScan().createLogScanner()) { + scanner.subscribe(0, 2L); + pollAndVerifyChangelogArrowBatches( + scanner, expectedChangeTypes, expectedKeys, expectedValues, 2); + } + } + } + @Test void testPollArrowBatchesWithSchemaEvolution() throws Exception { TablePath tablePath = TablePath.of("test_db_1", "test_arrow_batches_with_schema_evolution"); @@ -548,7 +608,7 @@ void testPollArrowBatchesWithSchemaEvolution() throws Exception { int totalRecords = 6; // subscribe from beginning and verify all 6 records - try (LogScannerImpl scanner = (LogScannerImpl) table.newScan().createLogScanner()) { + try (LogScanner scanner = table.newScan().createLogScanner()) { scanner.subscribeFromBeginning(0); pollAndVerifyArrowBatches(scanner, totalRecords, 0); } @@ -556,7 +616,7 @@ void testPollArrowBatchesWithSchemaEvolution() throws Exception { // subscribe from the middle of the first batch (offset 1) // to ensure records before the subscribe offset are not returned int subscribeOffset = 1; - try (LogScannerImpl scanner2 = (LogScannerImpl) table.newScan().createLogScanner()) { + try (LogScanner scanner2 = table.newScan().createLogScanner()) { scanner2.subscribe(0, subscribeOffset); pollAndVerifyArrowBatches( scanner2, totalRecords - subscribeOffset, subscribeOffset); @@ -565,7 +625,7 @@ void testPollArrowBatchesWithSchemaEvolution() throws Exception { } private void pollAndVerifyArrowBatches( - LogScannerImpl scanner, int expectedRecords, int minExpectedOffset) { + LogScanner scanner, int expectedRecords, int minExpectedOffset) { int count = 0; long deadline = System.nanoTime() + Duration.ofSeconds(30).toNanos(); while (count < expectedRecords) { @@ -601,6 +661,78 @@ private void pollAndVerifyArrowBatches( assertThat(count).isEqualTo(expectedRecords); } + private void pollAndVerifyChangelogArrowBatches( + LogScanner scanner, + ChangeType[] expectedChangeTypes, + int[] expectedKeys, + String[] expectedValues, + int startingOffset) { + int count = startingOffset; + long deadline = System.nanoTime() + Duration.ofSeconds(30).toNanos(); + while (count < expectedChangeTypes.length) { + assertThat(System.nanoTime()) + .as( + "Timed out waiting for %s changelog records, got %s", + expectedChangeTypes.length - startingOffset, count - startingOffset) + .isLessThan(deadline); + try (ArrowScanRecords records = scanner.pollRecordBatch(Duration.ofSeconds(1))) { + for (ArrowBatchData batch : records) { + try (ArrowBatchData b = batch) { + assertThat(b.isAppendOnly()).isFalse(); + ByteBuffer changeTypes = b.getChangeTypes().get(); + assertThat(changeTypes.remaining()).isEqualTo(b.getRecordCount()); + IntVector keys = (IntVector) b.getVectorSchemaRoot().getVector(0); + VarCharVector values = (VarCharVector) b.getVectorSchemaRoot().getVector(1); + for (int rowId = 0; rowId < b.getRecordCount(); rowId++) { + long offset = b.getBaseLogOffset() + rowId; + assertThat(offset).isEqualTo(count); + assertThat(b.getChangeType(rowId)) + .isEqualTo(expectedChangeTypes[count]); + assertThat(changeTypes.get(rowId)) + .isEqualTo(expectedChangeTypes[count].toByteValue()); + assertThat(keys.get(rowId)).isEqualTo(expectedKeys[count]); + assertThat(values.getObject(rowId).toString()) + .isEqualTo(expectedValues[count]); + count++; + } + } + } + } + } + assertThat(count).isEqualTo(expectedChangeTypes.length); + } + + private void pollAndVerifyProjectedChangelogArrowBatches( + LogScanner scanner, ChangeType[] expectedChangeTypes, String[] expectedValues) { + int count = 0; + long deadline = System.nanoTime() + Duration.ofSeconds(30).toNanos(); + while (count < expectedChangeTypes.length) { + assertThat(System.nanoTime()) + .as( + "Timed out waiting for %s projected changelog records, got %s", + expectedChangeTypes.length, count) + .isLessThan(deadline); + try (ArrowScanRecords records = scanner.pollRecordBatch(Duration.ofSeconds(1))) { + for (ArrowBatchData batch : records) { + try (ArrowBatchData b = batch) { + assertThat(b.getVectorSchemaRoot().getFieldVectors()).hasSize(1); + VarCharVector values = (VarCharVector) b.getVectorSchemaRoot().getVector(0); + for (int rowId = 0; rowId < b.getRecordCount(); rowId++) { + long offset = b.getBaseLogOffset() + rowId; + assertThat(offset).isEqualTo(count); + assertThat(b.getChangeType(rowId)) + .isEqualTo(expectedChangeTypes[count]); + assertThat(values.getObject(rowId).toString()) + .isEqualTo(expectedValues[count]); + count++; + } + } + } + } + } + assertThat(count).isEqualTo(expectedChangeTypes.length); + } + private static ScanRecords pollUntilProgressOnly( LogScanner logScanner, TableBucket tableBucket) { long deadline = System.nanoTime() + Duration.ofSeconds(30).toNanos(); diff --git a/fluss-common/src/main/java/org/apache/fluss/record/ArrowBatchData.java b/fluss-common/src/main/java/org/apache/fluss/record/ArrowBatchData.java index c3e55fadf71..b8cb8da63b0 100644 --- a/fluss-common/src/main/java/org/apache/fluss/record/ArrowBatchData.java +++ b/fluss-common/src/main/java/org/apache/fluss/record/ArrowBatchData.java @@ -17,12 +17,18 @@ package org.apache.fluss.record; -import org.apache.fluss.annotation.Internal; +import org.apache.fluss.annotation.PublicEvolving; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; +import javax.annotation.Nullable; + +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Optional; + import static org.apache.fluss.utils.Preconditions.checkArgument; import static org.apache.fluss.utils.Preconditions.checkNotNull; @@ -33,21 +39,37 @@ * *

The caller must close this object after use in order to release the underlying Arrow memory. */ -@Internal +@PublicEvolving public class ArrowBatchData implements AutoCloseable { private final VectorSchemaRoot vectorSchemaRoot; private final long baseLogOffset; private final long timestamp; private final int schemaId; + @Nullable private final byte[] changeTypes; private boolean closed; public ArrowBatchData( VectorSchemaRoot vectorSchemaRoot, long baseLogOffset, long timestamp, int schemaId) { + this(vectorSchemaRoot, baseLogOffset, timestamp, schemaId, null); + } + + ArrowBatchData( + VectorSchemaRoot vectorSchemaRoot, + long baseLogOffset, + long timestamp, + int schemaId, + @Nullable byte[] changeTypes) { this.vectorSchemaRoot = checkNotNull(vectorSchemaRoot, "vectorSchemaRoot must not be null"); this.baseLogOffset = baseLogOffset; this.timestamp = timestamp; this.schemaId = schemaId; + checkArgument( + changeTypes == null || changeTypes.length == vectorSchemaRoot.getRowCount(), + "changeTypes length must match row count %s, but is %s", + vectorSchemaRoot.getRowCount(), + changeTypes == null ? 0 : changeTypes.length); + this.changeTypes = changeTypes; } /** Returns the Arrow vectors of this batch. */ @@ -75,6 +97,41 @@ public int getRecordCount() { return vectorSchemaRoot.getRowCount(); } + /** Returns whether every row in this batch is append-only. */ + public boolean isAppendOnly() { + return changeTypes == null; + } + + /** + * Returns the stored change type for the given row. + * + *

Append-only batches do not materialize a change-type vector and return {@link + * ChangeType#APPEND_ONLY} for every row. + */ + public ChangeType getChangeType(int rowId) { + checkArgument( + rowId >= 0 && rowId < getRecordCount(), + "rowId must be in [0, %s), but is %s", + getRecordCount(), + rowId); + return changeTypes == null + ? ChangeType.APPEND_ONLY + : ChangeType.fromByteValue(changeTypes[rowId]); + } + + /** + * Returns the stored per-row change-type vector as a read-only buffer. + * + *

The buffer contains one {@link ChangeType#toByteValue() encoded byte} for every row and is + * absent for append-only batches. The returned buffer remains valid for the lifetime of this + * object. + */ + public Optional getChangeTypes() { + return changeTypes == null + ? Optional.empty() + : Optional.of(ByteBuffer.wrap(changeTypes).asReadOnlyBuffer()); + } + /** Returns the total size in bytes of the underlying Arrow buffers. */ public long getSizeInBytes() { long size = 0; @@ -83,6 +140,9 @@ public long getSizeInBytes() { size += buf.readableBytes(); } } + if (changeTypes != null) { + size += changeTypes.length; + } return size; } @@ -105,9 +165,14 @@ skipRows < getRecordCount(), getRecordCount()); int remainingRows = getRecordCount() - skipRows; VectorSchemaRoot slicedRoot = vectorSchemaRoot.slice(skipRows, remainingRows); - // release original vector buffers; sliced vectors hold independent copies + byte[] slicedChangeTypes = + changeTypes == null + ? null + : Arrays.copyOfRange(changeTypes, skipRows, changeTypes.length); + // VectorSchemaRoot.slice transfers each slice into independently owned vectors. close(); - return new ArrowBatchData(slicedRoot, baseLogOffset + skipRows, timestamp, schemaId); + return new ArrowBatchData( + slicedRoot, baseLogOffset + skipRows, timestamp, schemaId, slicedChangeTypes); } /** @@ -128,8 +193,11 @@ public ArrowBatchData truncateAndTransferOwnership(int rowCount) { rowCount, getRecordCount()); VectorSchemaRoot slicedRoot = vectorSchemaRoot.slice(0, rowCount); + byte[] slicedChangeTypes = + changeTypes == null ? null : Arrays.copyOfRange(changeTypes, 0, rowCount); close(); - return new ArrowBatchData(slicedRoot, baseLogOffset, timestamp, schemaId); + return new ArrowBatchData( + slicedRoot, baseLogOffset, timestamp, schemaId, slicedChangeTypes); } @Override diff --git a/fluss-common/src/main/java/org/apache/fluss/record/ArrowRecordBatchContext.java b/fluss-common/src/main/java/org/apache/fluss/record/ArrowRecordBatchContext.java index 64cf19ba42d..cc0e65dc4fb 100644 --- a/fluss-common/src/main/java/org/apache/fluss/record/ArrowRecordBatchContext.java +++ b/fluss-common/src/main/java/org/apache/fluss/record/ArrowRecordBatchContext.java @@ -20,6 +20,8 @@ import org.apache.fluss.annotation.Internal; import org.apache.fluss.memory.MemorySegment; +import javax.annotation.Nullable; + /** Internal context for loading Arrow record batches with unshaded Arrow resources. */ @Internal interface ArrowRecordBatchContext extends LogRecordBatch.ReadContext { @@ -38,6 +40,7 @@ interface UnshadedArrowBatchAccess extends AutoCloseable { * Applies schema-evolution projection (if needed) internally and builds the final {@link * ArrowBatchData}, transferring ownership to the caller. */ - ArrowBatchData createArrowBatchData(long baseLogOffset, long timestamp, int schemaId); + ArrowBatchData createArrowBatchData( + long baseLogOffset, long timestamp, int schemaId, @Nullable byte[] changeTypes); } } 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..d250715f429 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 @@ -277,17 +277,24 @@ public ArrowBatchData loadArrowBatch(ReadContext context) { "Arrow batch loading requires context to implement ArrowRecordBatchContext, but is %s.", context.getClass().getName()); ArrowRecordBatchContext arrowRecordBatchContext = (ArrowRecordBatchContext) context; - checkArgument(isAppendOnly(), "Arrow batch loading only supports append-only batches."); ArrowRecordBatchContext.UnshadedArrowBatchAccess batchAccess = arrowRecordBatchContext.createUnshadedArrowBatchAccess(schemaId); try { int recordsDataOffset = recordsDataOffset(); - int arrowOffset = position + recordsDataOffset; - int arrowLength = sizeInBytes() - recordsDataOffset; + boolean appendOnly = isAppendOnly(); + int changeTypesLength = appendOnly ? 0 : getRecordCount(); + int arrowOffset = position + recordsDataOffset + changeTypesLength; + byte[] changeTypes = null; + if (!appendOnly) { + changeTypes = new byte[changeTypesLength]; + segment.get(position + recordsDataOffset, changeTypes); + } + int arrowLength = sizeInBytes() - recordsDataOffset - changeTypesLength; batchAccess.loadArrowBatch(segment, arrowOffset, arrowLength); ArrowBatchData arrowBatchData = - batchAccess.createArrowBatchData(baseLogOffset(), commitTimestamp(), schemaId); + batchAccess.createArrowBatchData( + baseLogOffset(), commitTimestamp(), schemaId, changeTypes); batchAccess = null; return arrowBatchData; } catch (Throwable t) { 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 cc5bcbdbae1..b8850c5d22b 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 @@ -443,7 +443,7 @@ public void loadArrowBatch(MemorySegment segment, int arrowOffset, int arrowLeng @Override public ArrowBatchData createArrowBatchData( - long baseLogOffset, long timestamp, int schemaId) { + long baseLogOffset, long timestamp, int schemaId, @Nullable byte[] changeTypes) { int[] schemaMapping = getSchemaEvolutionMapping(schemaId); if (schemaMapping != null) { outputRoot = @@ -453,7 +453,7 @@ public ArrowBatchData createArrowBatchData( readRoot = null; } ArrowBatchData arrowBatchData = - new ArrowBatchData(outputRoot, baseLogOffset, timestamp, schemaId); + new ArrowBatchData(outputRoot, baseLogOffset, timestamp, schemaId, changeTypes); outputRoot = null; readRoot = null; return arrowBatchData; diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringSplitReader.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringSplitReader.java index 60b911b6f74..e1c77f84967 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringSplitReader.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/tiering/source/TieringSplitReader.java @@ -23,7 +23,6 @@ import org.apache.fluss.client.table.scanner.ScanRecord; import org.apache.fluss.client.table.scanner.log.ArrowScanRecords; import org.apache.fluss.client.table.scanner.log.LogScanner; -import org.apache.fluss.client.table.scanner.log.LogScannerImpl; import org.apache.fluss.client.table.scanner.log.ScanRecords; import org.apache.fluss.flink.source.reader.BoundedSplitReader; import org.apache.fluss.flink.source.reader.RecordAndPos; @@ -209,7 +208,7 @@ public RecordsWithSplitIds> fetch() throws I } if (useRecordBatchPath()) { try (ArrowScanRecords arrowScanRecords = - ((LogScannerImpl) currentLogScanner).pollRecordBatch(pollTimeout)) { + currentLogScanner.pollRecordBatch(pollTimeout)) { return processLogRecords( arrowScanRecords.buckets(), arrowScanRecords::records,