This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 6509987c60 [avro] Optimize array in avro (#8567)
6509987c60 is described below
commit 6509987c60321892d8c8acb00e9b5281fd61d78d
Author: yuzelin <[email protected]>
AuthorDate: Mon Jul 13 10:18:31 2026 +0800
[avro] Optimize array in avro (#8567)
---
.../benchmark/compact/ArrayCompactBenchmark.java | 142 ++++++++
.../apache/paimon/format/avro/AvroBytesArray.java | 380 +++++++++++++++++++++
.../paimon/format/avro/FieldReaderFactory.java | 17 +-
.../paimon/format/avro/FieldWriterFactory.java | 36 ++
.../paimon/format/avro/AvroBytesArrayTest.java | 335 ++++++++++++++++++
5 files changed, 907 insertions(+), 3 deletions(-)
diff --git
a/paimon-benchmark/paimon-micro-benchmarks/src/test/java/org/apache/paimon/benchmark/compact/ArrayCompactBenchmark.java
b/paimon-benchmark/paimon-micro-benchmarks/src/test/java/org/apache/paimon/benchmark/compact/ArrayCompactBenchmark.java
new file mode 100644
index 0000000000..83f0238517
--- /dev/null
+++
b/paimon-benchmark/paimon-micro-benchmarks/src/test/java/org/apache/paimon/benchmark/compact/ArrayCompactBenchmark.java
@@ -0,0 +1,142 @@
+/*
+ * 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.paimon.benchmark.compact;
+
+import org.apache.paimon.benchmark.Benchmark;
+import org.apache.paimon.catalog.FileSystemCatalog;
+import org.apache.paimon.catalog.Identifier;
+import org.apache.paimon.data.BinaryString;
+import org.apache.paimon.data.GenericArray;
+import org.apache.paimon.data.GenericRow;
+import org.apache.paimon.fs.Path;
+import org.apache.paimon.fs.local.LocalFileIO;
+import org.apache.paimon.schema.Schema;
+import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.sink.TableWriteImpl;
+import org.apache.paimon.types.DataType;
+import org.apache.paimon.types.DataTypes;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+/** Benchmark for compacting rows with array. */
+public class ArrayCompactBenchmark {
+
+ @TempDir java.nio.file.Path tempDir;
+
+ private final int rowCount = 4_000_000;
+ private final int numberArraySize = 100;
+ private final int stringArraySize = 20;
+
+ @Test
+ public void testArrayInt() throws Exception {
+ runBenchmark("array-int", DataTypes.INT(), this::intArray);
+ }
+
+ @Test
+ public void testArrayBigInt() throws Exception {
+ runBenchmark("array-bigint", DataTypes.BIGINT(), this::bigIntArray);
+ }
+
+ @Test
+ public void testArrayString() throws Exception {
+ runBenchmark("array-string", DataTypes.STRING(), this::stringArray);
+ }
+
+ private void runBenchmark(String name, DataType elementType, ArrayFactory
arrayFactory)
+ throws Exception {
+ Benchmark benchmark =
+ new Benchmark("avro-" + name + "-compact-benchmark", rowCount)
+ .setNumWarmupIters(1)
+ .setOutputPerIteration(true);
+
+ FileStoreTable table = createTable(elementType);
+ benchmark.addCase(
+ "write-and-compact",
+ 1,
+ () -> {
+ try {
+ writeAndCompact(table, arrayFactory);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
+ benchmark.run();
+ }
+
+ private FileStoreTable createTable(DataType elementType) throws Exception {
+ FileSystemCatalog catalog =
+ new FileSystemCatalog(LocalFileIO.create(), new
Path(tempDir.toString()));
+ catalog.createDatabase("default", false);
+ Identifier identifier = Identifier.create("default", "T");
+ catalog.createTable(
+ identifier,
+ Schema.newBuilder()
+ .column("k", DataTypes.STRING())
+ .column("v", DataTypes.ARRAY(elementType))
+ .primaryKey("k")
+ .option("bucket", "1")
+ .option("file.format", "avro")
+ .build(),
+ false);
+ return (FileStoreTable) catalog.getTable(identifier);
+ }
+
+ private void writeAndCompact(FileStoreTable table, ArrayFactory
arrayFactory) {
+ try (TableWriteImpl<?> write = table.newWrite("test")) {
+ for (int i = 0; i < rowCount; i++) {
+ write.write(
+ GenericRow.of(
+ BinaryString.fromString(String.valueOf(i)),
+ arrayFactory.create(i)));
+ }
+ write.prepareCommit(true, 1);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ private GenericArray intArray(int rowId) {
+ Integer[] array = new Integer[numberArraySize];
+ for (int i = 0; i < numberArraySize; i++) {
+ array[i] = rowId * numberArraySize + i;
+ }
+ return new GenericArray(array);
+ }
+
+ private GenericArray bigIntArray(int rowId) {
+ Long[] array = new Long[numberArraySize];
+ for (int i = 0; i < numberArraySize; i++) {
+ array[i] = (long) rowId * numberArraySize + i;
+ }
+ return new GenericArray(array);
+ }
+
+ private GenericArray stringArray(int rowId) {
+ BinaryString[] array = new BinaryString[stringArraySize];
+ for (int i = 0; i < stringArraySize; i++) {
+ array[i] = BinaryString.fromString(String.valueOf(rowId *
stringArraySize + i));
+ }
+ return new GenericArray(array);
+ }
+
+ private interface ArrayFactory {
+ GenericArray create(int rowId);
+ }
+}
diff --git
a/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroBytesArray.java
b/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroBytesArray.java
new file mode 100644
index 0000000000..5babf19f5c
--- /dev/null
+++
b/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroBytesArray.java
@@ -0,0 +1,380 @@
+/*
+ * 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.paimon.format.avro;
+
+import org.apache.paimon.data.BinaryString;
+import org.apache.paimon.data.Blob;
+import org.apache.paimon.data.Decimal;
+import org.apache.paimon.data.GenericArray;
+import org.apache.paimon.data.InternalArray;
+import org.apache.paimon.data.InternalMap;
+import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.data.InternalVector;
+import org.apache.paimon.data.Timestamp;
+import org.apache.paimon.data.variant.Variant;
+import org.apache.paimon.utils.IntArrayList;
+
+import org.apache.avro.Schema;
+import org.apache.avro.io.BinaryData;
+import org.apache.avro.io.BinaryDecoder;
+import org.apache.avro.io.BinaryEncoder;
+import org.apache.avro.io.DecoderFactory;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+
+/** An {@link InternalArray} whose elements are stored as Avro binary
payloads. */
+public class AvroBytesArray implements InternalArray {
+
+ // stores original element bytes from avro files
+ private final byte[] bytes;
+ // total bytes length
+ private final int lengthInBytes;
+ // offset of each element payload
+ private final IntArrayList off;
+ // length of each element payload
+ private final IntArrayList len;
+ // reader to decode element payloads lazily
+ private final FieldReader elementReader;
+
+ // decoded array, lazily initialized when elements are accessed
+ private GenericArray decodedArray;
+
+ private AvroBytesArray(
+ byte[] bytes,
+ int lengthInBytes,
+ IntArrayList off,
+ IntArrayList len,
+ FieldReader elementReader) {
+ this.bytes = bytes;
+ this.lengthInBytes = lengthInBytes;
+ this.off = off;
+ this.len = len;
+ this.elementReader = elementReader;
+ }
+
+ @Nullable
+ public static AvroBytesArray create(
+ BinaryDecoder decoder, FieldReader elementReader, Schema
elementType)
+ throws IOException {
+ RawFieldReader rawFieldReader = createRawFieldReader(elementType);
+ if (rawFieldReader == null) {
+ return null;
+ }
+
+ BytesBuilder builder = new BytesBuilder();
+ IntArrayList off = new IntArrayList(16);
+ IntArrayList len = new IntArrayList(16);
+
+ long chunkLength = decoder.readArrayStart();
+ while (chunkLength > 0) {
+ for (int i = 0; i < chunkLength; i++) {
+ int offset = builder.lengthInBytes;
+ rawFieldReader.read(decoder, builder);
+ off.add(offset);
+ len.add(builder.lengthInBytes - offset);
+ }
+ chunkLength = decoder.arrayNext();
+ }
+
+ return new AvroBytesArray(builder.bytes, builder.lengthInBytes, off,
len, elementReader);
+ }
+
+ @Nullable
+ private static RawFieldReader createRawFieldReader(Schema elementSchema) {
+ // handle nullable
+ if (elementSchema.getType() == Schema.Type.UNION) {
+ RawFieldReader nested =
createRawFieldReader(elementSchema.getTypes().get(1));
+ if (nested == null) {
+ return null;
+ }
+ return (decoder, builder) -> {
+ int index = decoder.readInt();
+ builder.addInt(index);
+ if (index == 1) {
+ nested.read(decoder, builder);
+ }
+ };
+ }
+
+ switch (elementSchema.getType()) {
+ case BOOLEAN:
+ return (decoder, builder) ->
builder.addBoolean(decoder.readBoolean());
+ case INT:
+ return (decoder, builder) -> builder.addInt(decoder.readInt());
+ case LONG:
+ return (decoder, builder) ->
builder.addLong(decoder.readLong());
+ case FLOAT:
+ return (decoder, builder) ->
builder.addFloat(decoder.readFloat());
+ case DOUBLE:
+ return (decoder, builder) ->
builder.addDouble(decoder.readDouble());
+ case STRING:
+ case BYTES:
+ return (decoder, builder) -> builder.addBytes(decoder);
+ default:
+ // unknown or unsupported
+ return null;
+ }
+ }
+
+ public void writeRawElements(BinaryEncoder encoder, int offset, int
length) throws IOException {
+ if (length == 0) {
+ return;
+ }
+
+ checkRange(offset, length);
+ int rawOffset = off.get(offset);
+ int last = offset + length - 1;
+ int rawLength = off.get(last) + len.get(last) - rawOffset;
+ encoder.writeFixed(bytes, rawOffset, rawLength);
+ }
+
+ private void checkRange(int offset, int length) {
+ if (offset < 0 || length < 0 || offset + length > size()) {
+ throw new IndexOutOfBoundsException(
+ String.format(
+ "Invalid range offset %s and length %s for array
size %s.",
+ offset, length, size()));
+ }
+ }
+
+ public byte[] bytes() {
+ return bytes;
+ }
+
+ public int lengthInBytes() {
+ return lengthInBytes;
+ }
+
+ @Override
+ public int size() {
+ return off.size();
+ }
+
+ @Override
+ public boolean isNullAt(int pos) {
+ return decodedArray().isNullAt(pos);
+ }
+
+ @Override
+ public boolean getBoolean(int pos) {
+ return decodedArray().getBoolean(pos);
+ }
+
+ @Override
+ public byte getByte(int pos) {
+ return decodedArray().getByte(pos);
+ }
+
+ @Override
+ public short getShort(int pos) {
+ return decodedArray().getShort(pos);
+ }
+
+ @Override
+ public int getInt(int pos) {
+ return decodedArray().getInt(pos);
+ }
+
+ @Override
+ public long getLong(int pos) {
+ return decodedArray().getLong(pos);
+ }
+
+ @Override
+ public float getFloat(int pos) {
+ return decodedArray().getFloat(pos);
+ }
+
+ @Override
+ public double getDouble(int pos) {
+ return decodedArray().getDouble(pos);
+ }
+
+ @Override
+ public BinaryString getString(int pos) {
+ return decodedArray().getString(pos);
+ }
+
+ @Override
+ public Decimal getDecimal(int pos, int precision, int scale) {
+ return decodedArray().getDecimal(pos, precision, scale);
+ }
+
+ @Override
+ public Timestamp getTimestamp(int pos, int precision) {
+ return decodedArray().getTimestamp(pos, precision);
+ }
+
+ @Override
+ public byte[] getBinary(int pos) {
+ return decodedArray().getBinary(pos);
+ }
+
+ @Override
+ public Variant getVariant(int pos) {
+ return decodedArray().getVariant(pos);
+ }
+
+ @Override
+ public Blob getBlob(int pos) {
+ return decodedArray().getBlob(pos);
+ }
+
+ @Override
+ public InternalArray getArray(int pos) {
+ return decodedArray().getArray(pos);
+ }
+
+ @Override
+ public InternalVector getVector(int pos) {
+ return decodedArray().getVector(pos);
+ }
+
+ @Override
+ public InternalMap getMap(int pos) {
+ return decodedArray().getMap(pos);
+ }
+
+ @Override
+ public InternalRow getRow(int pos, int numFields) {
+ return decodedArray().getRow(pos, numFields);
+ }
+
+ @Override
+ public boolean[] toBooleanArray() {
+ return decodedArray().toBooleanArray();
+ }
+
+ @Override
+ public byte[] toByteArray() {
+ return decodedArray().toByteArray();
+ }
+
+ @Override
+ public short[] toShortArray() {
+ return decodedArray().toShortArray();
+ }
+
+ @Override
+ public int[] toIntArray() {
+ return decodedArray().toIntArray();
+ }
+
+ @Override
+ public long[] toLongArray() {
+ return decodedArray().toLongArray();
+ }
+
+ @Override
+ public float[] toFloatArray() {
+ return decodedArray().toFloatArray();
+ }
+
+ @Override
+ public double[] toDoubleArray() {
+ return decodedArray().toDoubleArray();
+ }
+
+ private GenericArray decodedArray() {
+ if (decodedArray != null) {
+ return decodedArray;
+ }
+
+ try {
+ BinaryDecoder decoder =
+ DecoderFactory.get().binaryDecoder(bytes, 0,
lengthInBytes, null);
+ Object[] values = new Object[size()];
+ for (int i = 0; i < size(); i++) {
+ values[i] = elementReader.read(decoder, null);
+ }
+ decodedArray = new GenericArray(values);
+ return decodedArray;
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to decode Avro array bytes.",
e);
+ }
+ }
+
+ /** Read the raw bytes from decoder and add to the {@link BytesBuilder}. */
+ private interface RawFieldReader {
+
+ void read(BinaryDecoder decoder, BytesBuilder builder) throws
IOException;
+ }
+
+ /**
+ * A growable byte buffer used to copy Avro binary-encoded element
payloads from a decoder.
+ *
+ * <p>Please see the #encodeX method comments to know how many spaces
should be ensured.
+ */
+ private static class BytesBuilder {
+
+ private byte[] bytes = new byte[256];
+ private int lengthInBytes;
+
+ void addBoolean(boolean value) {
+ ensure(1);
+ lengthInBytes += BinaryData.encodeBoolean(value, bytes,
lengthInBytes);
+ }
+
+ void addInt(int value) {
+ ensure(5);
+ lengthInBytes += BinaryData.encodeInt(value, bytes, lengthInBytes);
+ }
+
+ void addLong(long value) {
+ ensure(10);
+ lengthInBytes += BinaryData.encodeLong(value, bytes,
lengthInBytes);
+ }
+
+ void addFloat(float value) {
+ ensure(4);
+ lengthInBytes += BinaryData.encodeFloat(value, bytes,
lengthInBytes);
+ }
+
+ void addDouble(double value) {
+ ensure(8);
+ lengthInBytes += BinaryData.encodeDouble(value, bytes,
lengthInBytes);
+ }
+
+ void addBytes(BinaryDecoder decoder) throws IOException {
+ // BinaryDecoder#readString,readBytes
+ int l = (int) decoder.readLong();
+ ensure(l + 10);
+ lengthInBytes += BinaryData.encodeLong(l, bytes, lengthInBytes);
+ decoder.readFixed(bytes, lengthInBytes, l);
+ lengthInBytes += l;
+ }
+
+ private void ensure(int need) {
+ if (lengthInBytes + need <= bytes.length) {
+ return;
+ }
+
+ int cap = bytes.length;
+ while (lengthInBytes + need > cap) {
+ cap *= 2;
+ }
+
+ byte[] newBytes = new byte[cap];
+ System.arraycopy(bytes, 0, newBytes, 0, lengthInBytes);
+ bytes = newBytes;
+ }
+ }
+}
diff --git
a/paimon-format/src/main/java/org/apache/paimon/format/avro/FieldReaderFactory.java
b/paimon-format/src/main/java/org/apache/paimon/format/avro/FieldReaderFactory.java
index 4801134507..41e71d4fde 100644
---
a/paimon-format/src/main/java/org/apache/paimon/format/avro/FieldReaderFactory.java
+++
b/paimon-format/src/main/java/org/apache/paimon/format/avro/FieldReaderFactory.java
@@ -164,7 +164,7 @@ public class FieldReaderFactory implements
AvroSchemaVisitor<FieldReader> {
@Override
public FieldReader visitArray(Schema schema, @Nullable DataType
elementType) {
FieldReader elementReader = visit(schema.getElementType(),
elementType);
- return new ArrayReader(elementReader);
+ return new ArrayReader(schema.getElementType(), elementReader);
}
@Override
@@ -429,15 +429,26 @@ public class FieldReaderFactory implements
AvroSchemaVisitor<FieldReader> {
private static class ArrayReader implements FieldReader {
+ private final Schema elementSchema;
private final FieldReader elementReader;
private final List<Object> reusedList = new ArrayList<>();
- private ArrayReader(FieldReader elementReader) {
+ private ArrayReader(@Nullable Schema elementSchema, FieldReader
elementReader) {
+ this.elementSchema = elementSchema;
this.elementReader = elementReader;
}
@Override
public Object read(Decoder decoder, Object reuse) throws IOException {
+ if (elementSchema != null && decoder instanceof BinaryDecoder) {
+ AvroBytesArray avroBytesArray =
+ AvroBytesArray.create(
+ (BinaryDecoder) decoder, elementReader,
elementSchema);
+ if (avroBytesArray != null) {
+ return avroBytesArray;
+ }
+ }
+
reusedList.clear();
long chunkLength = decoder.readArrayStart();
@@ -471,7 +482,7 @@ public class FieldReaderFactory implements
AvroSchemaVisitor<FieldReader> {
private final DataType elementType;
private ArrayVectorReader(FieldReader elementReader, DataType
elementType) {
- super(elementReader);
+ super(null, elementReader);
this.elementType = elementType;
}
diff --git
a/paimon-format/src/main/java/org/apache/paimon/format/avro/FieldWriterFactory.java
b/paimon-format/src/main/java/org/apache/paimon/format/avro/FieldWriterFactory.java
index 8ca985bd1a..9cefc48e09 100644
---
a/paimon-format/src/main/java/org/apache/paimon/format/avro/FieldWriterFactory.java
+++
b/paimon-format/src/main/java/org/apache/paimon/format/avro/FieldWriterFactory.java
@@ -186,6 +186,11 @@ public class FieldWriterFactory implements
AvroSchemaVisitor<FieldWriter> {
FieldWriter elementWriter = visit(schema.getElementType(),
elementType);
return (container, index, encoder) -> {
InternalArray array = container.getArray(index);
+ if (encoder instanceof BinaryEncoder) {
+ writeArrayData(array, elementWriter, (BinaryEncoder) encoder);
+ return;
+ }
+
encoder.writeArrayStart();
int numElements = array.size();
encoder.setItemCount(numElements);
@@ -197,6 +202,37 @@ public class FieldWriterFactory implements
AvroSchemaVisitor<FieldWriter> {
};
}
+ private void writeArrayData(
+ InternalArray array, FieldWriter elementWriter, BinaryEncoder
encoder)
+ throws IOException {
+ encoder.writeArrayStart();
+ encoder.setItemCount(array.size());
+ writeArrayItems(array, 0, array.size(), elementWriter, encoder);
+ encoder.writeArrayEnd();
+ }
+
+ private void writeArrayItems(
+ InternalArray array,
+ int offset,
+ int length,
+ FieldWriter elementWriter,
+ BinaryEncoder encoder)
+ throws IOException {
+ if (length == 0) {
+ return;
+ }
+
+ if (array instanceof AvroBytesArray) {
+ ((AvroBytesArray) array).writeRawElements(encoder, offset, length);
+ } else {
+ // fallback to element by element
+ for (int i = 0; i < length; i++) {
+ encoder.startItem();
+ elementWriter.write(array, offset + i, encoder);
+ }
+ }
+ }
+
@Override
public FieldWriter visitArrayVector(Schema schema, DataType elementType) {
FieldWriter elementWriter = visit(schema.getElementType(),
elementType);
diff --git
a/paimon-format/src/test/java/org/apache/paimon/format/avro/AvroBytesArrayTest.java
b/paimon-format/src/test/java/org/apache/paimon/format/avro/AvroBytesArrayTest.java
new file mode 100644
index 0000000000..12280c5a2f
--- /dev/null
+++
b/paimon-format/src/test/java/org/apache/paimon/format/avro/AvroBytesArrayTest.java
@@ -0,0 +1,335 @@
+/*
+ * 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.paimon.format.avro;
+
+import org.apache.paimon.data.BinaryString;
+import org.apache.paimon.data.Decimal;
+import org.apache.paimon.data.GenericArray;
+import org.apache.paimon.data.InternalArray;
+import org.apache.paimon.data.Timestamp;
+import org.apache.paimon.types.DataType;
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.DecimalType;
+import org.apache.paimon.types.LocalZonedTimestampType;
+import org.apache.paimon.types.TimestampType;
+
+import org.apache.avro.Schema;
+import org.apache.avro.io.BinaryDecoder;
+import org.apache.avro.io.BinaryEncoder;
+import org.apache.avro.io.DecoderFactory;
+import org.apache.avro.io.EncoderFactory;
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayOutputStream;
+import java.math.BigDecimal;
+import java.util.HashMap;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link AvroBytesArray}. */
+public class AvroBytesArrayTest {
+
+ @Test
+ public void testNotNullElementsForSupportedTypes() throws Exception {
+ checkResult(DataTypes.BOOLEAN(), objectArray(true, false, true),
false);
+ checkResult(DataTypes.TINYINT(), objectArray((byte) 1, (byte) 2,
(byte) 3), false);
+ checkResult(DataTypes.SMALLINT(), objectArray((short) 1, (short) 2,
(short) 3), false);
+ checkResult(DataTypes.INT(), objectArray(1, 2, 3), false);
+ checkResult(DataTypes.BIGINT(), objectArray(1L, 2L, 3L), false);
+ checkResult(DataTypes.FLOAT(), objectArray(1.25F, 2.5F, 3.75F), false);
+ checkResult(DataTypes.DOUBLE(), objectArray(1.25D, 2.5D, 3.75D),
false);
+ checkResult(
+ DataTypes.CHAR(10),
+ objectArray(
+ BinaryString.fromString("a"),
+ BinaryString.fromString("b"),
+ BinaryString.fromString("c")),
+ false);
+ checkResult(
+ DataTypes.STRING(),
+ objectArray(
+ BinaryString.fromString("d"),
+ BinaryString.fromString("e"),
+ BinaryString.fromString("f")),
+ false);
+ checkResult(
+ DataTypes.BINARY(10),
+ objectArray(new byte[] {1, 2}, new byte[] {3, 4}, new byte[]
{5, 6}),
+ false);
+ checkResult(
+ DataTypes.VARBINARY(10),
+ objectArray(new byte[] {1, 2}, new byte[] {3, 4}, new byte[]
{5, 6}),
+ false);
+ checkResult(DataTypes.DATE(), objectArray(1, 2, 3), false);
+ checkResult(DataTypes.TIME(), objectArray(1000, 2000, 3000), false);
+ checkResult(
+ DataTypes.TIMESTAMP(3),
+ objectArray(
+ Timestamp.fromEpochMillis(1000),
+ Timestamp.fromEpochMillis(2000),
+ Timestamp.fromEpochMillis(3000)),
+ false);
+ checkResult(
+ DataTypes.TIMESTAMP(6),
+ objectArray(
+ Timestamp.fromMicros(1001),
+ Timestamp.fromMicros(2002),
+ Timestamp.fromMicros(3003)),
+ false);
+ checkResult(
+ DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(3),
+ objectArray(
+ Timestamp.fromEpochMillis(1000),
+ Timestamp.fromEpochMillis(2000),
+ Timestamp.fromEpochMillis(3000)),
+ false);
+ checkResult(
+ DataTypes.DECIMAL(10, 2),
+ objectArray(decimal("1.23"), decimal("2.34"), decimal("3.45")),
+ false);
+ }
+
+ @Test
+ public void testNullableElementsForSupportedTypes() throws Exception {
+ checkResult(DataTypes.BOOLEAN(), objectArray(true, null, false), true);
+ checkResult(DataTypes.TINYINT(), objectArray((byte) 1, null, (byte)
2), true);
+ checkResult(DataTypes.SMALLINT(), objectArray((short) 1, null, (short)
2), true);
+ checkResult(DataTypes.INT(), objectArray(1, null, 2), true);
+ checkResult(DataTypes.BIGINT(), objectArray(1L, null, 2L), true);
+ checkResult(DataTypes.FLOAT(), objectArray(1.25F, null, 2.5F), true);
+ checkResult(DataTypes.DOUBLE(), objectArray(1.25D, null, 2.5D), true);
+ checkResult(
+ DataTypes.CHAR(10),
+ objectArray(BinaryString.fromString("a"), null,
BinaryString.fromString("b")),
+ true);
+ checkResult(
+ DataTypes.STRING(),
+ objectArray(BinaryString.fromString("c"), null,
BinaryString.fromString("d")),
+ true);
+ checkResult(
+ DataTypes.BINARY(10),
+ objectArray(new byte[] {1, 2}, null, new byte[] {3, 4}),
+ true);
+ checkResult(
+ DataTypes.VARBINARY(10),
+ objectArray(new byte[] {1, 2}, null, new byte[] {3, 4}),
+ true);
+ checkResult(DataTypes.DATE(), objectArray(1, null, 2), true);
+ checkResult(DataTypes.TIME(), objectArray(1000, null, 2000), true);
+ checkResult(
+ DataTypes.TIMESTAMP(3),
+ objectArray(Timestamp.fromEpochMillis(1000), null,
Timestamp.fromEpochMillis(2000)),
+ true);
+ checkResult(
+ DataTypes.TIMESTAMP(6),
+ objectArray(Timestamp.fromMicros(1001), null,
Timestamp.fromMicros(2002)),
+ true);
+ checkResult(
+ DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(3),
+ objectArray(Timestamp.fromEpochMillis(1000), null,
Timestamp.fromEpochMillis(2000)),
+ true);
+ checkResult(
+ DataTypes.DECIMAL(10, 2),
+ objectArray(decimal("1.23"), null, decimal("2.34")),
+ true);
+ }
+
+ private void checkResult(DataType elementType, GenericArray values,
boolean nullable)
+ throws Exception {
+ DataType arrayElementType = nullable ? elementType :
elementType.notNull();
+ Schema arraySchema =
+ AvroSchemaConverter.convertToSchema(
+ DataTypes.ARRAY(arrayElementType).notNull(), new
HashMap<>());
+ Schema elementSchema = arraySchema.getElementType();
+
+ byte[] bytes = writeRawArrayBytes(arrayElementType, values);
+ BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(bytes,
null);
+ FieldReader elementReader = new
FieldReaderFactory().visit(elementSchema, arrayElementType);
+ AvroBytesArray array = AvroBytesArray.create(decoder, elementReader,
elementSchema);
+ assertThat(array).isNotNull();
+ assertThat(array.size()).isEqualTo(values.size());
+
+ checkArrayValues(elementType, values, array);
+
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ BinaryEncoder encoder = EncoderFactory.get().binaryEncoder(baos, null);
+ encoder.writeArrayStart();
+ encoder.setItemCount(array.size());
+ array.writeRawElements(encoder, 0, array.size());
+ encoder.writeArrayEnd();
+ encoder.flush();
+ assertThat(baos.toByteArray()).isEqualTo(bytes);
+ }
+
+ private byte[] writeRawArrayBytes(DataType elementType, GenericArray
values) throws Exception {
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ BinaryEncoder encoder = EncoderFactory.get().binaryEncoder(baos, null);
+ encoder.writeArrayStart();
+ encoder.setItemCount(values.size());
+ for (int i = 0; i < values.size(); i++) {
+ encoder.startItem();
+ if (values.isNullAt(i)) {
+ encoder.writeIndex(0);
+ } else {
+ if (elementType.isNullable()) {
+ encoder.writeIndex(1);
+ }
+ writeElement(
+ elementType,
+
InternalArray.createElementGetter(elementType).getElementOrNull(values, i),
+ encoder);
+ }
+ }
+ encoder.writeArrayEnd();
+ encoder.flush();
+ return baos.toByteArray();
+ }
+
+ private void checkArrayValues(
+ DataType elementType, GenericArray expectedValues, InternalArray
actualArray) {
+ InternalArray.ElementGetter getter =
InternalArray.createElementGetter(elementType);
+ for (int i = 0; i < expectedValues.size(); i++) {
+ if (expectedValues.isNullAt(i)) {
+ assertThat(actualArray.isNullAt(i)).isTrue();
+ } else {
+ checkValue(elementType,
getter.getElementOrNull(expectedValues, i), actualArray, i);
+ }
+ }
+ }
+
+ private void checkValue(
+ DataType elementType, Object expected, InternalArray actualArray,
int pos) {
+ switch (elementType.getTypeRoot()) {
+ case CHAR:
+ case VARCHAR:
+
assertThat(actualArray.getString(pos).toString()).isEqualTo(expected.toString());
+ return;
+ case BOOLEAN:
+ assertThat(actualArray.getBoolean(pos)).isEqualTo(expected);
+ return;
+ case BINARY:
+ case VARBINARY:
+
assertThat(actualArray.getBinary(pos)).containsExactly((byte[]) expected);
+ return;
+ case DECIMAL:
+ DecimalType decimalType = (DecimalType) elementType;
+ assertThat(
+ actualArray.getDecimal(
+ pos, decimalType.getPrecision(),
decimalType.getScale()))
+ .isEqualTo(expected);
+ return;
+ case TINYINT:
+ assertThat(actualArray.getByte(pos)).isEqualTo(expected);
+ return;
+ case SMALLINT:
+ assertThat(actualArray.getShort(pos)).isEqualTo(expected);
+ return;
+ case INTEGER:
+ case DATE:
+ case TIME_WITHOUT_TIME_ZONE:
+ assertThat(actualArray.getInt(pos)).isEqualTo(expected);
+ return;
+ case BIGINT:
+ assertThat(actualArray.getLong(pos)).isEqualTo(expected);
+ return;
+ case FLOAT:
+ assertThat(actualArray.getFloat(pos)).isEqualTo(expected);
+ return;
+ case DOUBLE:
+ assertThat(actualArray.getDouble(pos)).isEqualTo(expected);
+ return;
+ case TIMESTAMP_WITHOUT_TIME_ZONE:
+ case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
+ assertThat(actualArray.getTimestamp(pos,
timestampPrecision(elementType)))
+ .isEqualTo(expected);
+ return;
+ default:
+ throw new UnsupportedOperationException(
+ "Unsupported AvroBytesArray test type: " +
elementType);
+ }
+ }
+
+ private void writeElement(DataType elementType, Object value,
BinaryEncoder encoder)
+ throws Exception {
+ switch (elementType.getTypeRoot()) {
+ case BOOLEAN:
+ encoder.writeBoolean((Boolean) value);
+ return;
+ case TINYINT:
+ case SMALLINT:
+ case INTEGER:
+ case DATE:
+ case TIME_WITHOUT_TIME_ZONE:
+ encoder.writeInt(((Number) value).intValue());
+ return;
+ case BIGINT:
+ encoder.writeLong((Long) value);
+ return;
+ case FLOAT:
+ encoder.writeFloat((Float) value);
+ return;
+ case DOUBLE:
+ encoder.writeDouble((Double) value);
+ return;
+ case VARCHAR:
+ case CHAR:
+ encoder.writeString(value.toString());
+ return;
+ case BINARY:
+ case VARBINARY:
+ encoder.writeBytes((byte[]) value);
+ return;
+ case DECIMAL:
+ encoder.writeBytes(((Decimal) value).toUnscaledBytes());
+ return;
+ case TIMESTAMP_WITHOUT_TIME_ZONE:
+ case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
+ Timestamp timestamp = (Timestamp) value;
+ if (timestampPrecision(elementType) <= 3) {
+ encoder.writeLong(timestamp.getMillisecond());
+ } else {
+ encoder.writeLong(timestamp.toMicros());
+ }
+ return;
+ default:
+ throw new UnsupportedOperationException(
+ "Unsupported E2E AvroBytesArray test type: " +
elementType);
+ }
+ }
+
+ private int timestampPrecision(DataType elementType) {
+ switch (elementType.getTypeRoot()) {
+ case TIMESTAMP_WITHOUT_TIME_ZONE:
+ return ((TimestampType) elementType).getPrecision();
+ case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
+ return ((LocalZonedTimestampType) elementType).getPrecision();
+ default:
+ throw new UnsupportedOperationException(
+ "Unsupported timestamp type: " + elementType);
+ }
+ }
+
+ private GenericArray objectArray(Object... values) {
+ return new GenericArray(values);
+ }
+
+ private Decimal decimal(String value) {
+ return Decimal.fromBigDecimal(new BigDecimal(value), 10, 2);
+ }
+}
