You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hudi.apache.org by da...@apache.org on 2025/02/10 02:35:22 UTC
(hudi) branch master updated: [HUDI-8190] Implement efficient streaming reads for HoodieDataBlocks (#11924)
This is an automated email from the ASF dual-hosted git repository.
danny0405 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hudi.git
The following commit(s) were added to refs/heads/master by this push:
new b40d22a67d1 [HUDI-8190] Implement efficient streaming reads for HoodieDataBlocks (#11924)
b40d22a67d1 is described below
commit b40d22a67d1aa25d882b22b39ba858129f509427
Author: usberkeley <15...@users.noreply.github.com>
AuthorDate: Mon Feb 10 10:35:15 2025 +0800
[HUDI-8190] Implement efficient streaming reads for HoodieDataBlocks (#11924)
---
.../table/log/block/HoodieAvroDataBlock.java | 175 ++++++++++
.../common/table/log/block/HoodieDataBlock.java | 81 ++++-
.../table/log/block/HoodieHFileDataBlock.java | 18 +
.../common/table/log/block/HoodieLogBlock.java | 4 +
.../table/log/block/HoodieParquetDataBlock.java | 18 +
.../table/log/block/TestHoodieAvroDataBlock.java | 373 +++++++++++++++++++++
6 files changed, 666 insertions(+), 3 deletions(-)
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieAvroDataBlock.java b/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieAvroDataBlock.java
index 488ffb1ce43..0a440919ab6 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieAvroDataBlock.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieAvroDataBlock.java
@@ -27,6 +27,7 @@ import org.apache.hudi.common.model.HoodieRecord.HoodieRecordType;
import org.apache.hudi.common.util.Option;
import org.apache.hudi.common.util.collection.ClosableIterator;
import org.apache.hudi.common.util.collection.CloseableMappingIterator;
+import org.apache.hudi.exception.HoodieException;
import org.apache.hudi.exception.HoodieIOException;
import org.apache.hudi.internal.schema.InternalSchema;
import org.apache.hudi.io.SeekableDataInputStream;
@@ -52,6 +53,7 @@ import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
+import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
@@ -146,6 +148,26 @@ public class HoodieAvroDataBlock extends HoodieDataBlock {
return new CloseableMappingIterator<>(iterator, data -> (HoodieRecord<T>) new HoodieAvroIndexedRecord(data));
}
+ /**
+ * Streaming deserialization of records.
+ *
+ * @param inputStream The input stream from which to read the records.
+ * @param contentLocation The location within the input stream where the content starts.
+ * @param bufferSize The size of the buffer to use for reading the records.
+ * @return A ClosableIterator over HoodieRecord<T>.
+ * @throws IOException If there is an error reading or deserializing the records.
+ */
+ @Override
+ protected <T> ClosableIterator<HoodieRecord<T>> deserializeRecords(
+ SeekableDataInputStream inputStream,
+ HoodieLogBlockContentLocation contentLocation,
+ HoodieRecordType type,
+ int bufferSize
+ ) throws IOException {
+ StreamingRecordIterator iterator = StreamingRecordIterator.getInstance(this, inputStream, contentLocation, bufferSize);
+ return new CloseableMappingIterator<>(iterator, data -> (HoodieRecord<T>) new HoodieAvroIndexedRecord(data));
+ }
+
@Override
protected <T> ClosableIterator<T> deserializeRecords(HoodieReaderContext<T> readerContext, byte[] content) throws IOException {
checkState(this.readerSchema != null, "Reader's schema has to be non-null");
@@ -221,6 +243,159 @@ public class HoodieAvroDataBlock extends HoodieDataBlock {
}
}
+ /**
+ * {@code StreamingRecordIterator} is an iterator for reading records from a Hoodie log block in streaming manner.
+ * It decodes the given input stream into Avro records with optional schema promotion.
+ *
+ * <p>This iterator ensures that the buffer has enough data for each record and handles buffer setup,
+ * including compaction and resizing when necessary.
+ */
+ private static class StreamingRecordIterator implements ClosableIterator<IndexedRecord> {
+ private static final int RECORD_LENGTH_BYTES = 4;
+ // The minimum buffer size in bytes
+ private static final int MIN_BUFFER_SIZE = RECORD_LENGTH_BYTES;
+ private final SeekableDataInputStream inputStream;
+ private final GenericDatumReader<IndexedRecord> reader;
+ private final ThreadLocal<BinaryDecoder> decoderCache = new ThreadLocal<>();
+ private Option<Schema> promotedSchema = Option.empty();
+ private int totalRecords = 0;
+ private int readRecords = 0;
+ private ByteBuffer buffer;
+
+ private StreamingRecordIterator(Schema readerSchema, Schema writerSchema, SeekableDataInputStream inputStream,
+ HoodieLogBlockContentLocation contentLocation, int bufferSize) throws IOException {
+ // Negative values should not be used because they are generally considered to indicate the operation of closing stream reading,
+ // in order to avoid confusing users into thinking that stream reading can be closed.
+ checkArgument(bufferSize > 0, "Buffer size must be greater than zero");
+ bufferSize = Math.max(bufferSize, MIN_BUFFER_SIZE);
+
+ this.inputStream = inputStream;
+
+ // Seek to the start of the block
+ this.inputStream.seek(contentLocation.getContentPositionInLogFile());
+
+ // Read version for this data block
+ int version = this.inputStream.readInt();
+ if (new HoodieAvroDataBlockVersion(version).hasRecordCount()) {
+ this.totalRecords = this.inputStream.readInt();
+ }
+
+ if (recordNeedsRewriteForExtendedAvroTypePromotion(writerSchema, readerSchema)) {
+ this.reader = new GenericDatumReader<>(writerSchema, writerSchema);
+ this.promotedSchema = Option.of(readerSchema);
+ } else {
+ this.reader = new GenericDatumReader<>(writerSchema, readerSchema);
+ }
+
+ this.buffer = ByteBuffer.allocate(Math.min(bufferSize, Math.toIntExact(contentLocation.getBlockSize())));
+ // The buffer defaults to read mode
+ this.buffer.flip();
+ }
+
+ public static StreamingRecordIterator getInstance(HoodieAvroDataBlock dataBlock, SeekableDataInputStream inputStream,
+ HoodieLogBlockContentLocation contentLocation, int bufferSize) throws IOException {
+ return new StreamingRecordIterator(dataBlock.readerSchema, dataBlock.getSchemaFromHeader(), inputStream, contentLocation, bufferSize);
+ }
+
+ @Override
+ public void close() {
+ this.decoderCache.remove();
+ this.buffer = null;
+ try {
+ this.inputStream.close();
+ } catch (IOException ex) {
+ throw new HoodieIOException("Failed to close input stream", ex);
+ }
+ }
+
+ @Override
+ public boolean hasNext() {
+ return readRecords < totalRecords;
+ }
+
+ @Override
+ public IndexedRecord next() {
+ try {
+ ensureBufferHasData(RECORD_LENGTH_BYTES);
+
+ // Read the record length
+ int recordLength = buffer.getInt();
+
+ // Ensure buffer is large enough and has enough data
+ ensureBufferHasData(recordLength);
+
+ // Decode the record
+ BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(buffer.array(), buffer.position(), recordLength, this.decoderCache.get());
+ this.decoderCache.set(decoder);
+ IndexedRecord record = this.reader.read(null, decoder);
+ buffer.position(buffer.position() + recordLength);
+ this.readRecords++;
+ if (this.promotedSchema.isPresent()) {
+ return HoodieAvroUtils.rewriteRecordWithNewSchema(record, this.promotedSchema.get());
+ }
+ return record;
+ } catch (IOException e) {
+ throw new HoodieIOException("Unable to convert bytes to record", e);
+ }
+ }
+
+ /**
+ * Ensures that the buffer contains at least the specified amount of data.
+ *
+ * <p>This method checks if the buffer has the required amount of data. If not, it attempts to fill the buffer
+ * by reading more data from the input stream. If the buffer's capacity is insufficient, it allocates a larger buffer.
+ * If the end of the input stream is reached before the required amount of data is available, an exception is thrown.
+ *
+ * @param dataLength the amount of data (in bytes) that must be available in the buffer.
+ * @throws IOException if an I/O error occurs while reading from the input stream.
+ * @throws HoodieException if the end of the input stream is reached before the required amount of data is available.
+ */
+ private void ensureBufferHasData(int dataLength) throws IOException {
+ // Check if the current buffer has enough space to read the required data length
+ if (buffer.capacity() - buffer.position() < dataLength) {
+ buffer.compact();
+ // Reset the buffer to read mode
+ buffer.flip();
+ }
+
+ // Check again if the buffer still doesn't have enough space after compaction
+ if (buffer.capacity() - buffer.position() < dataLength) {
+ ByteBuffer newBuffer = ByteBuffer.allocate(buffer.position() + dataLength);
+ newBuffer.put(buffer);
+ // Reset the new buffer to read mode
+ newBuffer.flip();
+ buffer = newBuffer;
+ }
+
+ while (buffer.remaining() < dataLength) {
+ boolean hasMoreData = fillBuffer();
+ if (!hasMoreData && buffer.remaining() < dataLength) {
+ throw new HoodieException("Unable to read enough data from the input stream to fill the buffer");
+ }
+ }
+ }
+
+ /**
+ * Attempts to fill the buffer with more data from the input stream.
+ *
+ * <p>This method reads data from the input stream into the buffer, starting at the current limit
+ * and reading up to the capacity of the buffer. If the end of the input stream is reached,
+ * it returns false. Otherwise, it updates the buffer's limit to reflect the new data and returns true.
+ *
+ * @return true if data was successfully read into the buffer; false if the end of the input stream was reached.
+ * @throws IOException if an I/O error occurs while reading from the input stream.
+ */
+ private boolean fillBuffer() throws IOException {
+ int bytesRead = inputStream.read(buffer.array(), buffer.limit(), buffer.capacity() - buffer.limit());
+ if (bytesRead == -1) {
+ return false;
+ }
+
+ buffer.limit(buffer.limit() + bytesRead);
+ return true;
+ }
+ }
+
//----------------------------------------------------------------------------------------
// DEPRECATED METHODS
//
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieDataBlock.java b/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieDataBlock.java
index 6f48cb27a74..c2a707c8303 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieDataBlock.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieDataBlock.java
@@ -142,16 +142,38 @@ public abstract class HoodieDataBlock extends HoodieLogBlock {
}
/**
- * Returns all the records iterator contained w/in this block.
+ * Returns an iterator over all the records contained within this block.
+ * This method uses a default buffer size of 0, which means it will read
+ * the entire Data Block content at once.
+ *
+ * @param type The type of HoodieRecord.
+ * @param <T> The type parameter for HoodieRecord.
+ * @return A ClosableIterator over HoodieRecord<T>.
*/
public final <T> ClosableIterator<HoodieRecord<T>> getRecordIterator(HoodieRecordType type) {
+ return getRecordIterator(type, 0);
+ }
+
+ /**
+ * Returns an iterator over all the records contained within this block.
+ *
+ * @param type The type of HoodieRecord.
+ * @param bufferSize The size of the buffer for streaming read.
+ * A bufferSize less than or equal to 0 means that streaming read is disabled and
+ * the entire block content will be read at once.
+ * A bufferSize greater than 0 enables streaming read with the specified buffer size.
+ * @param <T> The type parameter for HoodieRecord.
+ * @return A ClosableIterator over HoodieRecord<T>.
+ * @throws HoodieIOException If there is an error reading records from the block payload.
+ */
+ public final <T> ClosableIterator<HoodieRecord<T>> getRecordIterator(HoodieRecordType type, int bufferSize) {
if (records.isPresent()) {
// TODO need convert record type
return list2Iterator(unsafeCast(records.get()));
}
try {
// in case records are absent, read content lazily and then convert to IndexedRecords
- return readRecordsFromBlockPayload(type);
+ return readRecordsFromBlockPayload(type, bufferSize);
} catch (IOException io) {
throw new HoodieIOException("Unable to convert content bytes to records", io);
}
@@ -170,6 +192,23 @@ public abstract class HoodieDataBlock extends HoodieLogBlock {
* @throws IOException in case of failures encountered when reading/parsing records
*/
public final <T> ClosableIterator<HoodieRecord<T>> getRecordIterator(List<String> keys, boolean fullKey, HoodieRecordType type) throws IOException {
+ return getRecordIterator(keys, fullKey, type, 0);
+ }
+
+ /**
+ * Batch get of keys of interest. Implementation can choose to either do full scan and return matched entries or
+ * do a seek based parsing and return matched entries.
+ *
+ * @param keys keys of interest.
+ * @param bufferSize The size of the buffer for streaming read.
+ * A bufferSize less than or equal to 0 means that streaming read is disabled and
+ * the entire block content will be read at once.
+ * A bufferSize greater than 0 enables streaming read with the specified buffer size.
+ * @param <T> The type parameter for HoodieRecord.
+ * @return List of IndexedRecords for the keys of interest.
+ * @throws IOException in case of failures encountered when reading/parsing records
+ */
+ public final <T> ClosableIterator<HoodieRecord<T>> getRecordIterator(List<String> keys, boolean fullKey, HoodieRecordType type, int bufferSize) throws IOException {
boolean fullScan = keys.isEmpty();
if (enablePointLookups && !fullScan) {
return lookupRecords(keys, fullKey);
@@ -177,7 +216,7 @@ public abstract class HoodieDataBlock extends HoodieLogBlock {
// Otherwise, we fetch all the records and filter out all the records, but the
// ones requested
- ClosableIterator<HoodieRecord<T>> allRecords = getRecordIterator(type);
+ ClosableIterator<HoodieRecord<T>> allRecords = getRecordIterator(type, bufferSize);
if (fullScan) {
return allRecords;
}
@@ -246,6 +285,26 @@ public abstract class HoodieDataBlock extends HoodieLogBlock {
}
}
+ /**
+ * Reads records from the block payload using a specified buffer size.
+ * This method attempts to read serialized records from the block payload, leveraging a buffer size for streaming reads.
+ * If the buffer size is less than or equal to 0, it reads the entire block content at once.
+ *
+ * @param bufferSize The size of the buffer for streaming read.
+ * A bufferSize less than or equal to 0 means that streaming read is disabled and
+ * the entire block content will be read at once.
+ * A bufferSize greater than 0 enables streaming read with the specified buffer size.
+ * @return A ClosableIterator over HoodieRecord<T>.
+ * @throws IOException If there is an error reading or deserializing the records.
+ */
+ protected <T> ClosableIterator<HoodieRecord<T>> readRecordsFromBlockPayload(HoodieRecordType type, int bufferSize) throws IOException {
+ if (getContent().isPresent() || bufferSize <= 0) {
+ return readRecordsFromBlockPayload(type);
+ }
+
+ return deserializeRecords(getInputStreamSupplier().get(), getBlockContentLocation().get(), type, bufferSize);
+ }
+
protected <T> ClosableIterator<T> readRecordsFromBlockPayload(HoodieReaderContext<T> readerContext) throws IOException {
if (readBlockLazily && !getContent().isPresent()) {
// read log block contents from disk
@@ -270,6 +329,22 @@ public abstract class HoodieDataBlock extends HoodieLogBlock {
protected abstract <T> ClosableIterator<HoodieRecord<T>> deserializeRecords(byte[] content, HoodieRecordType type) throws IOException;
+ /**
+ * Streaming deserialization of records.
+ *
+ * @param inputStream The input stream from which to read the records.
+ * @param contentLocation The location within the input stream where the content starts.
+ * @param bufferSize The size of the buffer to use for reading the records.
+ * @return A ClosableIterator over HoodieRecord<T>.
+ * @throws IOException If there is an error reading or deserializing the records.
+ */
+ protected abstract <T> ClosableIterator<HoodieRecord<T>> deserializeRecords(
+ SeekableDataInputStream inputStream,
+ HoodieLogBlockContentLocation contentLocation,
+ HoodieRecordType type,
+ int bufferSize
+ ) throws IOException;
+
/**
* Deserializes the content bytes of the data block to the records in engine-specific representation.
*
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieHFileDataBlock.java b/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieHFileDataBlock.java
index 4d1267701ba..ea474142463 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieHFileDataBlock.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieHFileDataBlock.java
@@ -125,6 +125,24 @@ public class HoodieHFileDataBlock extends HoodieDataBlock {
}
}
+ /**
+ * Streaming deserialization of records.
+ *
+ * @param inputStream The input stream from which to read the records.
+ * @param contentLocation The location within the input stream where the content starts.
+ * @param bufferSize The size of the buffer to use for reading the records.
+ * @return A ClosableIterator over HoodieRecord<T>.
+ * @throws IOException If there is an error reading or deserializing the records.
+ */
+ protected <T> ClosableIterator<HoodieRecord<T>> deserializeRecords(
+ SeekableDataInputStream inputStream,
+ HoodieLogBlockContentLocation contentLocation,
+ HoodieRecordType type,
+ int bufferSize
+ ) throws IOException {
+ throw new UnsupportedOperationException("Streaming deserialization is not supported for HoodieHFileDataBlock");
+ }
+
@Override
protected <T> ClosableIterator<T> deserializeRecords(HoodieReaderContext<T> readerContext, byte[] content) throws IOException {
checkState(readerSchema != null, "Reader's schema has to be non-null");
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieLogBlock.java b/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieLogBlock.java
index 04c7827b4e4..78c39ff5faf 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieLogBlock.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieLogBlock.java
@@ -345,6 +345,10 @@ public abstract class HoodieLogBlock {
return Option.of(content);
}
+ protected Supplier<SeekableDataInputStream> getInputStreamSupplier() {
+ return inputStreamSupplier;
+ }
+
/**
* Adds the record positions if the base file instant time of the positions exists
* in the log header and the record positions are all valid.
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieParquetDataBlock.java b/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieParquetDataBlock.java
index 8d9b1221c54..35c0a0f8536 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieParquetDataBlock.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/table/log/block/HoodieParquetDataBlock.java
@@ -159,6 +159,24 @@ public class HoodieParquetDataBlock extends HoodieDataBlock {
throw new UnsupportedOperationException("Should not be invoked");
}
+ /**
+ * Streaming deserialization of records.
+ *
+ * @param inputStream The input stream from which to read the records.
+ * @param contentLocation The location within the input stream where the content starts.
+ * @param bufferSize The size of the buffer to use for reading the records.
+ * @return A ClosableIterator over HoodieRecord<T>.
+ * @throws IOException If there is an error reading or deserializing the records.
+ */
+ protected <T> ClosableIterator<HoodieRecord<T>> deserializeRecords(
+ SeekableDataInputStream inputStream,
+ HoodieLogBlockContentLocation contentLocation,
+ HoodieRecordType type,
+ int bufferSize
+ ) throws IOException {
+ throw new UnsupportedOperationException("Should not be invoked");
+ }
+
@Override
protected <T> ClosableIterator<T> deserializeRecords(HoodieReaderContext<T> readerContext, byte[] content) throws IOException {
throw new UnsupportedOperationException("Should not be invoked");
diff --git a/hudi-common/src/test/java/org/apache/hudi/common/table/log/block/TestHoodieAvroDataBlock.java b/hudi-common/src/test/java/org/apache/hudi/common/table/log/block/TestHoodieAvroDataBlock.java
new file mode 100644
index 00000000000..98d918c1cb5
--- /dev/null
+++ b/hudi-common/src/test/java/org/apache/hudi/common/table/log/block/TestHoodieAvroDataBlock.java
@@ -0,0 +1,373 @@
+/*
+ * 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.hudi.common.table.log.block;
+
+import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.hudi.common.model.HoodieAvroIndexedRecord;
+import org.apache.hudi.common.model.HoodieKey;
+import org.apache.hudi.common.model.HoodieRecord;
+import org.apache.hudi.common.table.log.block.HoodieLogBlock.HoodieLogBlockContentLocation;
+import org.apache.hudi.common.table.log.block.HoodieLogBlock.HeaderMetadataType;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.collection.ClosableIterator;
+import org.apache.hudi.common.util.io.ByteBufferBackedInputStream;
+import org.apache.hudi.exception.HoodieIOException;
+import org.apache.hudi.io.ByteArraySeekableDataInputStream;
+import org.apache.hudi.io.SeekableDataInputStream;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.io.EOFException;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class TestHoodieAvroDataBlock {
+ // Record key field of the test data
+ private static final String RECORD_KEY_FIELD = HoodieRecord.RECORD_KEY_METADATA_FIELD;
+
+ private static final String SCHEMA_STRING = "{\n"
+ + " \"type\": \"record\",\n"
+ + " \"name\": \"RandomRecord\",\n"
+ + " \"fields\": [\n"
+ + " {\"name\": \"" + RECORD_KEY_FIELD + "\", \"type\": \"string\"},\n"
+ + " {\"name\": \"value\", \"type\": \"double\"}\n"
+ + " ]\n"
+ + "}";
+
+ private static final Schema SCHEMA = new Schema.Parser().parse(SCHEMA_STRING);
+
+ // Just to avoid errors when the HoodieAvroDataBlock constructor wraps HoodieLogBlockContentLocation as an Option and encounters a null value
+ private static final HoodieLogBlockContentLocation NULL_BLOCK_CONTENT_LOCATION = new HoodieLogBlockContentLocation(null, null, 0, 0, 0);
+
+ /**
+ * Tests the getRecordIterator method of HoodieAvroDataBlock when records are present.
+ *
+ * @param useKeyFilter Whether to use a key filter.
+ * @param keyFilterFullKeyMatch Whether the key filter should match the full key.
+ * @throws IOException If an I/O error occurs during the test.
+ */
+ @ParameterizedTest
+ @CsvSource({
+ "true, true",
+ "true, false",
+ "false, true",
+ "false, false"
+ })
+ public void testGetRecordIteratorWithRecordsPresent(boolean useKeyFilter, boolean keyFilterFullKeyMatch) throws IOException {
+ List<HoodieRecord> records = generateRandomHoodieRecords(SCHEMA, 1000);
+ HoodieAvroDataBlock block = createHoodieAvroDataBlock(SCHEMA, records);
+
+ // Expect records for using key filter
+ List<HoodieRecord> recordsForFilter = selectRandomRecords(records, keyFilterFullKeyMatch);
+ List<String> keysForFilter = recordsForFilter.stream().map(HoodieRecord::getRecordKey).collect(Collectors.toList());
+
+ // Get record iterator
+ try (ClosableIterator<HoodieRecord<Object>> recordIterator = useKeyFilter
+ ? block.getRecordIterator(keysForFilter, keyFilterFullKeyMatch, HoodieRecord.HoodieRecordType.AVRO)
+ : block.getRecordIterator(HoodieRecord.HoodieRecordType.AVRO)) {
+ List<HoodieRecord> retrievedRecords = new ArrayList<>();
+ recordIterator.forEachRemaining(retrievedRecords::add);
+
+ verifyRecords(useKeyFilter ? recordsForFilter : records, retrievedRecords);
+ }
+ }
+
+ /**
+ * Tests the getRecordIterator method of HoodieAvroDataBlock using content directly.
+ *
+ * @param useKeyFilter Whether to use a key filter.
+ * @param keyFilterFullKeyMatch Whether the key filter should match the full key.
+ * @throws IOException If an I/O error occurs during the test.
+ */
+ @ParameterizedTest
+ @CsvSource({
+ "true, true",
+ "true, false",
+ "false, true",
+ "false, false"
+ })
+ public void testGetRecordIteratorWithContent(boolean useKeyFilter, boolean keyFilterFullKeyMatch) throws IOException {
+ List<HoodieRecord> records = generateRandomHoodieRecords(SCHEMA, 1000);
+ byte[] content = createHoodieAvroDataBlockContent(SCHEMA, records);
+ Option<byte[]> contentOpt = Option.of(content);
+ Map<HoodieLogBlock.HeaderMetadataType, String> header = new HashMap<>();
+ header.put(HoodieLogBlock.HeaderMetadataType.SCHEMA, SCHEMA.toString());
+
+ // Expect records for using key filter
+ List<HoodieRecord> recordsForFilter = selectRandomRecords(records, keyFilterFullKeyMatch);
+ List<String> keysForFilter = recordsForFilter.stream().map(HoodieRecord::getRecordKey).collect(Collectors.toList());
+
+ // With log block content
+ HoodieAvroDataBlock dataBlock = new HoodieAvroDataBlock(() -> null, contentOpt, false,
+ NULL_BLOCK_CONTENT_LOCATION, Option.of(SCHEMA), header, new HashMap<>(), RECORD_KEY_FIELD);
+
+ // Get record iterator
+ try (ClosableIterator<HoodieRecord<Object>> recordIterator = useKeyFilter
+ ? dataBlock.getRecordIterator(keysForFilter, keyFilterFullKeyMatch, HoodieRecord.HoodieRecordType.AVRO)
+ : dataBlock.getRecordIterator(HoodieRecord.HoodieRecordType.AVRO)) {
+ List<HoodieRecord> retrievedRecords = new ArrayList<>();
+ recordIterator.forEachRemaining(retrievedRecords::add);
+
+ verifyRecords(useKeyFilter ? recordsForFilter : records, retrievedRecords);
+ }
+ }
+
+ /**
+ * Tests the getRecordIterator method of HoodieAvroDataBlock with input stream (read block lazily).
+ *
+ * @param useStreamingRead Whether to use streaming read.
+ * @param useKeyFilter Whether to use a key filter.
+ * @param keyFilterFullKeyMatch Whether the key filter should match the full key.
+ * @throws IOException If an I/O error occurs during the test.
+ */
+ @ParameterizedTest
+ @CsvSource({
+ "true, true, true",
+ "true, true, false",
+ "true, false, true",
+ "true, false, false",
+ "false, true, true",
+ "false, true, false",
+ "false, false, true",
+ "false, false, false"
+ })
+ public void testGetRecordIteratorWithInputStream(boolean useStreamingRead, boolean useKeyFilter, boolean keyFilterFullKeyMatch) throws IOException {
+ List<HoodieRecord> records = generateRandomHoodieRecords(SCHEMA, 1000);
+ byte[] blockContent = createHoodieAvroDataBlockContent(SCHEMA, records);
+ SeekableDataInputStream inputStream = createSeekableDataInputStream(blockContent);
+ Map<HeaderMetadataType, String> header = new HashMap<>();
+ header.put(HeaderMetadataType.SCHEMA, SCHEMA.toString());
+ HoodieLogBlockContentLocation logBlockContentLocation = new HoodieLogBlockContentLocation(null, null, 0, blockContent.length, blockContent.length);
+ int bufferSize = useStreamingRead ? 100 : 0; // bytes
+
+ // Expect records for using key filter
+ List<HoodieRecord> recordsForFilter = selectRandomRecords(records, keyFilterFullKeyMatch);
+ List<String> keysForFilter = recordsForFilter.stream().map(HoodieRecord::getRecordKey).collect(Collectors.toList());
+
+ // With input stream
+ HoodieAvroDataBlock dataBlock = new HoodieAvroDataBlock(() -> inputStream, Option.empty(), true,
+ logBlockContentLocation, Option.of(SCHEMA), header, new HashMap<>(), RECORD_KEY_FIELD);
+
+ // Get record iterator
+ try (ClosableIterator<HoodieRecord<Object>> recordIterator = useKeyFilter
+ ? dataBlock.getRecordIterator(keysForFilter, keyFilterFullKeyMatch, HoodieRecord.HoodieRecordType.AVRO, bufferSize)
+ : dataBlock.getRecordIterator(HoodieRecord.HoodieRecordType.AVRO, bufferSize)) {
+ List<HoodieRecord> retrievedRecords = new ArrayList<>();
+ recordIterator.forEachRemaining(retrievedRecords::add);
+
+ verifyRecords(useKeyFilter ? recordsForFilter : records, retrievedRecords);
+ }
+ }
+
+ /**
+ * Tests the getRecordIterator method of HoodieAvroDataBlock with streaming read.
+ *
+ * @param bufferSize The size of the buffer to be used for the streaming read.
+ * @throws IOException If an I/O error occurs during the test.
+ */
+ @ParameterizedTest
+ @ValueSource(ints = {Integer.MIN_VALUE, -2, -1, 0, 1, 2, 100, 200, 300, 1024, Integer.MAX_VALUE})
+ public void testGetRecordIteratorWithStreamingRead(int bufferSize) throws IOException {
+ List<HoodieRecord> records = generateRandomHoodieRecords(SCHEMA, 1000);
+ byte[] blockContent = createHoodieAvroDataBlockContent(SCHEMA, records);
+ SeekableDataInputStream inputStream = createSeekableDataInputStream(blockContent);
+ Map<HeaderMetadataType, String> header = new HashMap<>();
+ header.put(HeaderMetadataType.SCHEMA, SCHEMA.toString());
+ HoodieLogBlockContentLocation logBlockContentLocation = new HoodieLogBlockContentLocation(null, null, 0, blockContent.length, blockContent.length);
+
+ // Create an instance of HoodieAvroDataBlock with streaming read
+ HoodieAvroDataBlock dataBlock = new HoodieAvroDataBlock(() -> inputStream, Option.empty(), true,
+ logBlockContentLocation, Option.of(SCHEMA), header, new HashMap<>(), RECORD_KEY_FIELD);
+
+ // Get record iterator
+ try (ClosableIterator<HoodieRecord<Object>> recordIterator = dataBlock.getRecordIterator(HoodieRecord.HoodieRecordType.AVRO, bufferSize)) {
+ List<HoodieRecord> retrievedRecords = new ArrayList<>();
+ recordIterator.forEachRemaining(retrievedRecords::add);
+
+ verifyRecords(records, retrievedRecords);
+ }
+ }
+
+ /**
+ * Tests the getRecordIterator method of HoodieAvroDataBlock with empty content.
+ *
+ */
+ @Test
+ public void testGetRecordIteratorWithEmptyContent() {
+ byte[] content = new byte[0]; // Simulate empty content bytes
+ Option<byte[]> contentOption = Option.of(content);
+ Map<HeaderMetadataType, String> header = new HashMap<>();
+ header.put(HeaderMetadataType.SCHEMA, SCHEMA.toString());
+
+ // With content
+ HoodieAvroDataBlock avroDataBlock = new HoodieAvroDataBlock(() -> null, contentOption, false,
+ NULL_BLOCK_CONTENT_LOCATION, Option.of(SCHEMA), header, new HashMap<>(), RECORD_KEY_FIELD);
+
+ // Mock the behavior of deserializeRecords to throw IOException
+ assertThrows(EOFException.class, () -> avroDataBlock.deserializeRecords(content, HoodieRecord.HoodieRecordType.AVRO));
+
+ // Call getRecordIterator and verify the behavior
+ assertThrows(HoodieIOException.class, () -> avroDataBlock.getRecordIterator(HoodieRecord.HoodieRecordType.AVRO));
+ }
+
+ /**
+ * Tests the getRecordIterator method of HoodieAvroDataBlock with empty content and input stream.
+ *
+ */
+ @ParameterizedTest
+ @ValueSource(ints = {Integer.MIN_VALUE, -2, -1, 0, 1, 2, 100, 200, 300, 1024, Integer.MAX_VALUE})
+ public void testGetRecordIteratorWithEmptyContentAndInputStream(int bufferSize) throws IOException {
+ Map<HeaderMetadataType, String> header = new HashMap<>();
+ header.put(HeaderMetadataType.SCHEMA, SCHEMA.toString());
+
+ // empty input stream supplier and empty content
+ HoodieAvroDataBlock avroDataBlock = new HoodieAvroDataBlock(() -> null, Option.empty(), true,
+ NULL_BLOCK_CONTENT_LOCATION, Option.of(SCHEMA), header, new HashMap<>(), RECORD_KEY_FIELD);
+
+ assertThrows(NullPointerException.class, () -> avroDataBlock.getRecordIterator(HoodieRecord.HoodieRecordType.AVRO, bufferSize));
+ }
+
+ /**
+ * Generates a list of random HoodieRecord objects based on the provided Avro schema.
+ *
+ * @param schema The Avro schema to use for generating the records.
+ * @param recordCount The number of records to generate.
+ * @return A list of HoodieRecord objects.
+ */
+ private static List<HoodieRecord> generateRandomHoodieRecords(Schema schema, int recordCount) {
+ List<HoodieRecord> records = new ArrayList<>();
+ Random random = new Random();
+ for (int i = 0; i < recordCount; i++) {
+ GenericRecord record = new GenericData.Record(schema);
+ String recordKey = "key_" + i;
+ record.put("_hoodie_record_key", "key_" + i);
+ record.put("value", random.nextDouble());
+ records.add(new HoodieAvroIndexedRecord(new HoodieKey(recordKey, ""), record));
+ }
+
+ return records;
+ }
+
+ /**
+ * Creates a HoodieAvroDataBlock object from the provided schema and list of HoodieRecord objects.
+ *
+ * @param schema The Avro schema used to define the structure of the records.
+ * @param records The list of HoodieRecord objects to be included in the data block.
+ * @return A HoodieAvroDataBlock object containing the provided records and schema.
+ */
+ private static HoodieAvroDataBlock createHoodieAvroDataBlock(Schema schema, List<HoodieRecord> records) {
+ Map<HeaderMetadataType, String> header = new HashMap<>();
+ header.put(HeaderMetadataType.SCHEMA, schema.toString());
+
+ return new HoodieAvroDataBlock(records, header, RECORD_KEY_FIELD);
+ }
+
+ /**
+ * Creates the content bytes of a HoodieAvroDataBlock from the provided schema and list of HoodieRecord objects.
+ *
+ * @param schema The Avro schema used to define the structure of the records.
+ * @param records The list of HoodieRecord objects to be included in the data block.
+ * @return A byte array representing the content of the HoodieAvroDataBlock.
+ * @throws IOException If an I/O error occurs while generating the content bytes.
+ */
+ private static byte[] createHoodieAvroDataBlockContent(Schema schema, List<HoodieRecord> records) throws IOException {
+ Map<HeaderMetadataType, String> header = new HashMap<>();
+ header.put(HeaderMetadataType.SCHEMA, schema.toString());
+
+ return new HoodieAvroDataBlock(records, header, RECORD_KEY_FIELD).getContentBytes(null);
+ }
+
+ /**
+ * Creates a SeekableDataInputStream from the provided byte array content.
+ *
+ * @param content The byte array containing the data to be read.
+ * @return A SeekableDataInputStream that allows seeking within the provided byte array content.
+ * @throws IOException If an I/O error occurs while creating the SeekableDataInputStream.
+ */
+ private static SeekableDataInputStream createSeekableDataInputStream(byte[] content) throws IOException {
+ return new ByteArraySeekableDataInputStream(new ByteBufferBackedInputStream(content));
+ }
+
+ /**
+ * Selects a random subset of HoodieRecord objects from the provided list.
+ *
+ * @param records The list of HoodieRecord objects to sample from.
+ * @param fullKey A boolean flag indicating whether to match records by full key or by prefix.
+ * If true, records are matched by full key; if false, records are matched by key prefix.
+ * @return A list of randomly selected HoodieRecord objects.
+ */
+ private static List<HoodieRecord> selectRandomRecords(List<HoodieRecord> records, boolean fullKey) {
+ // number of sampled records
+ int count = new Random().nextInt(records.size());
+ // set of record keys of the sampled records
+ Set<String> keys = new Random()
+ .ints(count, 0, records.size())
+ .mapToObj(records::get)
+ .map(r -> r.getRecordKey(SCHEMA, RECORD_KEY_FIELD))
+ .collect(Collectors.toSet());
+
+ // simulate KeyFilter matching logic
+ return records.stream()
+ .filter(r -> fullKey
+ ? keys.contains(r.getRecordKey(SCHEMA, RECORD_KEY_FIELD))
+ : keys.stream().anyMatch(r.getRecordKey(SCHEMA, RECORD_KEY_FIELD)::startsWith))
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Verifies that the actual records match the expected records.
+ * This method converts the HoodieRecord objects to GenericRecord objects, sorts them,
+ * and then compares the sorted lists for equality.
+ *
+ * @param expectedRecords The list of expected HoodieRecord objects.
+ * @param actualRecords The list of actual HoodieRecord objects to verify.
+ * @throws AssertionError If the record counts do not match or if any record data mismatches.
+ */
+ private static void verifyRecords(List<HoodieRecord> expectedRecords, List<HoodieRecord> actualRecords) {
+ List<GenericRecord> expectedGenericRecords = convertAndSortRecords(expectedRecords);
+ List<GenericRecord> actualGenericRecords = convertAndSortRecords(actualRecords);
+
+ assertEquals(expectedGenericRecords.size(), actualGenericRecords.size(), "Record count mismatch");
+ for (int i = 0; i < expectedGenericRecords.size(); i++) {
+ assertEquals(expectedGenericRecords, actualGenericRecords, "Record data mismatch");
+ }
+ }
+
+ private static List<GenericRecord> convertAndSortRecords(List<HoodieRecord> hoodieRecords) {
+ return hoodieRecords.stream()
+ .map(HoodieRecord::getData)
+ .map(r -> (GenericRecord) r)
+ .sorted(Comparator.comparing(r -> r.get(RECORD_KEY_FIELD).toString()))
+ .collect(Collectors.toList());
+ }
+}