You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@paimon.apache.org by lz...@apache.org on 2024/08/19 09:37:15 UTC
(paimon) branch master updated: [arrow] Add reader from arrow to paimon (#3996)
This is an automated email from the ASF dual-hosted git repository.
lzljs3620320 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 cf8ef5eed [arrow] Add reader from arrow to paimon (#3996)
cf8ef5eed is described below
commit cf8ef5eed2583546777dd789be80576597c29145
Author: YeJunHao <41...@users.noreply.github.com>
AuthorDate: Mon Aug 19 17:37:09 2024 +0800
[arrow] Add reader from arrow to paimon (#3996)
---
.../java/org/apache/paimon/arrow/ArrowUtils.java | 2 +-
.../converter/Arrow2PaimonVectorConverter.java | 559 +++++++++++++++++++++
.../paimon/arrow/reader/ArrowBatchReader.java | 73 +++
.../paimon/arrow/vector/ArrowFormatWriter.java | 4 +-
.../paimon/arrow/vector/ArrowFormatWriterTest.java | 25 +-
5 files changed, 652 insertions(+), 11 deletions(-)
diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java
index 235ab9531..8a69fb63a 100644
--- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java
+++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/ArrowUtils.java
@@ -136,7 +136,7 @@ public class ArrowUtils {
return ArrowCStruct.of(array, schema);
}
- public void serializeToIpc(VectorSchemaRoot vsr, OutputStream out) {
+ public static void serializeToIpc(VectorSchemaRoot vsr, OutputStream out) {
try (ArrowStreamWriter writer = new ArrowStreamWriter(vsr, null, out)) {
writer.start();
writer.writeBatch();
diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/converter/Arrow2PaimonVectorConverter.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/converter/Arrow2PaimonVectorConverter.java
new file mode 100644
index 000000000..ff73f0694
--- /dev/null
+++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/converter/Arrow2PaimonVectorConverter.java
@@ -0,0 +1,559 @@
+/*
+ * 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.arrow.converter;
+
+import org.apache.paimon.data.Decimal;
+import org.apache.paimon.data.InternalArray;
+import org.apache.paimon.data.InternalMap;
+import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.data.Timestamp;
+import org.apache.paimon.data.columnar.ArrayColumnVector;
+import org.apache.paimon.data.columnar.BooleanColumnVector;
+import org.apache.paimon.data.columnar.ByteColumnVector;
+import org.apache.paimon.data.columnar.BytesColumnVector;
+import org.apache.paimon.data.columnar.ColumnVector;
+import org.apache.paimon.data.columnar.ColumnarArray;
+import org.apache.paimon.data.columnar.ColumnarMap;
+import org.apache.paimon.data.columnar.ColumnarRow;
+import org.apache.paimon.data.columnar.DecimalColumnVector;
+import org.apache.paimon.data.columnar.DoubleColumnVector;
+import org.apache.paimon.data.columnar.FloatColumnVector;
+import org.apache.paimon.data.columnar.IntColumnVector;
+import org.apache.paimon.data.columnar.LongColumnVector;
+import org.apache.paimon.data.columnar.MapColumnVector;
+import org.apache.paimon.data.columnar.RowColumnVector;
+import org.apache.paimon.data.columnar.ShortColumnVector;
+import org.apache.paimon.data.columnar.TimestampColumnVector;
+import org.apache.paimon.data.columnar.VectorizedColumnBatch;
+import org.apache.paimon.types.ArrayType;
+import org.apache.paimon.types.BigIntType;
+import org.apache.paimon.types.BinaryType;
+import org.apache.paimon.types.BooleanType;
+import org.apache.paimon.types.CharType;
+import org.apache.paimon.types.DataType;
+import org.apache.paimon.types.DataTypeVisitor;
+import org.apache.paimon.types.DateType;
+import org.apache.paimon.types.DecimalType;
+import org.apache.paimon.types.DoubleType;
+import org.apache.paimon.types.FloatType;
+import org.apache.paimon.types.IntType;
+import org.apache.paimon.types.LocalZonedTimestampType;
+import org.apache.paimon.types.MapType;
+import org.apache.paimon.types.MultisetType;
+import org.apache.paimon.types.RowType;
+import org.apache.paimon.types.SmallIntType;
+import org.apache.paimon.types.TimeType;
+import org.apache.paimon.types.TimestampType;
+import org.apache.paimon.types.TinyIntType;
+import org.apache.paimon.types.VarBinaryType;
+import org.apache.paimon.types.VarCharType;
+
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.BitVector;
+import org.apache.arrow.vector.DateDayVector;
+import org.apache.arrow.vector.DecimalVector;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.Float4Vector;
+import org.apache.arrow.vector.Float8Vector;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.SmallIntVector;
+import org.apache.arrow.vector.TimeMilliVector;
+import org.apache.arrow.vector.TinyIntVector;
+import org.apache.arrow.vector.VarBinaryVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.complex.MapVector;
+import org.apache.arrow.vector.complex.StructVector;
+
+import java.time.LocalDateTime;
+import java.util.ArrayList;
+import java.util.List;
+
+/** Convert a {@link FieldVector} to {@link ColumnVector}. */
+public interface Arrow2PaimonVectorConverter {
+
+ static Arrow2PaimonVectorConverter construct(DataType type) {
+ return type.accept(Arrow2PaimonVectorConvertorVisitor.INSTANCE);
+ }
+
+ ColumnVector convertVector(FieldVector vector);
+
+ /** Visitor to create convertor from arrow to paimon. */
+ class Arrow2PaimonVectorConvertorVisitor
+ implements DataTypeVisitor<Arrow2PaimonVectorConverter> {
+
+ private static final Arrow2PaimonVectorConvertorVisitor INSTANCE =
+ new Arrow2PaimonVectorConvertorVisitor();
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(CharType charType) {
+ return vector ->
+ new BytesColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public Bytes getBytes(int index) {
+ byte[] bytes = ((VarCharVector) vector).get(index);
+ return new Bytes(bytes, 0, bytes.length) {
+ @Override
+ public byte[] getBytes() {
+ return bytes;
+ }
+ };
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(VarCharType varCharType) {
+ return vector ->
+ new BytesColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public Bytes getBytes(int index) {
+ byte[] bytes = ((VarCharVector) vector).get(index);
+ return new Bytes(bytes, 0, bytes.length) {
+ @Override
+ public byte[] getBytes() {
+ return bytes;
+ }
+ };
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(BooleanType booleanType) {
+ return vector ->
+ new BooleanColumnVector() {
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public boolean getBoolean(int index) {
+ return ((BitVector) vector).getObject(index);
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(BinaryType binaryType) {
+ return vector ->
+ new BytesColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public Bytes getBytes(int index) {
+ byte[] bytes = ((VarBinaryVector) vector).getObject(index);
+ return new Bytes(bytes, 0, bytes.length) {
+ @Override
+ public byte[] getBytes() {
+ return bytes;
+ }
+ };
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(VarBinaryType varBinaryType) {
+ return vector ->
+ new BytesColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public Bytes getBytes(int index) {
+ byte[] bytes = ((VarBinaryVector) vector).getObject(index);
+ return new Bytes(bytes, 0, bytes.length) {
+ @Override
+ public byte[] getBytes() {
+ return bytes;
+ }
+ };
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(DecimalType decimalType) {
+ return vector ->
+ new DecimalColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public Decimal getDecimal(int index, int precision, int scale) {
+ return Decimal.fromBigDecimal(
+ ((DecimalVector) vector).getObject(index), precision, scale);
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(TinyIntType tinyIntType) {
+ return vector ->
+ new ByteColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public byte getByte(int index) {
+ return ((TinyIntVector) vector).getObject(index);
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(SmallIntType smallIntType) {
+ return vector ->
+ new ShortColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public short getShort(int index) {
+ return ((SmallIntVector) vector).getObject(index);
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(IntType intType) {
+ return vector ->
+ new IntColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public int getInt(int index) {
+ return ((IntVector) vector).getObject(index);
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(BigIntType bigIntType) {
+ return vector ->
+ new LongColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public long getLong(int index) {
+ return ((BigIntVector) vector).getObject(index);
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(FloatType floatType) {
+ return vector ->
+ new FloatColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public float getFloat(int index) {
+ return ((Float4Vector) vector).getObject(index);
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(DoubleType doubleType) {
+ return vector ->
+ new DoubleColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public double getDouble(int index) {
+ return ((Float8Vector) vector).getObject(index);
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(DateType dateType) {
+ return vector ->
+ new IntColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public int getInt(int index) {
+ return ((DateDayVector) vector).getObject(index);
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(TimeType timeType) {
+ return vector ->
+ new IntColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public int getInt(int index) {
+ return ((TimeMilliVector) vector).get(index);
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(TimestampType timestampType) {
+ return vector ->
+ new TimestampColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public Timestamp getTimestamp(int i, int precision) {
+ return Timestamp.fromLocalDateTime(
+ (LocalDateTime) (vector.getObject(i)));
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(LocalZonedTimestampType localZonedTimestampType) {
+ return vector ->
+ new TimestampColumnVector() {
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public Timestamp getTimestamp(int i, int precision) {
+ long value = (long) vector.getObject(i);
+ if (precision == 0) {
+ return Timestamp.fromEpochMillis(value * 1000);
+ } else if (precision >= 1 && precision <= 3) {
+ return Timestamp.fromEpochMillis(value);
+ } else if (precision >= 4 && precision <= 6) {
+ return Timestamp.fromMicros(value);
+ } else {
+ return Timestamp.fromEpochMillis(
+ value / 1_000_000, (int) value % 1_000_000);
+ }
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(ArrayType arrayType) {
+ final Arrow2PaimonVectorConverter arrowVectorConvertor =
+ arrayType.getElementType().accept(this);
+
+ return vector ->
+ new ArrayColumnVector() {
+
+ private boolean inited = false;
+ private ColumnVector columnVector;
+
+ private void init() {
+ if (!inited) {
+ FieldVector child = ((ListVector) vector).getDataVector();
+ this.columnVector = arrowVectorConvertor.convertVector(child);
+ inited = true;
+ }
+ }
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public InternalArray getArray(int index) {
+ init();
+ ListVector listVector = (ListVector) vector;
+ int start = listVector.getElementStartIndex(index);
+ int end = listVector.getElementEndIndex(index);
+ return new ColumnarArray(columnVector, start, end - start);
+ }
+
+ @Override
+ public ColumnVector getColumnVector() {
+ init();
+ return columnVector;
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(MultisetType multisetType) {
+ throw new UnsupportedOperationException("Doesn't support MultisetType.");
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(MapType mapType) {
+ final Arrow2PaimonVectorConverter keyConvertor = mapType.getKeyType().accept(this);
+ final Arrow2PaimonVectorConverter valueConverter = mapType.getValueType().accept(this);
+
+ return vector ->
+ new MapColumnVector() {
+
+ private boolean inited = false;
+ private MapVector mapVector;
+ private ColumnVector keyColumnVector;
+ private ColumnVector valueColumnVector;
+
+ private void init() {
+ if (!inited) {
+ this.mapVector = (MapVector) vector;
+ StructVector listVector = (StructVector) mapVector.getDataVector();
+
+ FieldVector keyVector = listVector.getChildrenFromFields().get(0);
+ FieldVector valueVector = listVector.getChildrenFromFields().get(1);
+
+ this.keyColumnVector = keyConvertor.convertVector(keyVector);
+ this.valueColumnVector = valueConverter.convertVector(valueVector);
+ inited = true;
+ }
+ }
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public InternalMap getMap(int index) {
+ init();
+
+ int start = mapVector.getElementStartIndex(index);
+ int end = mapVector.getElementEndIndex(index);
+
+ return new ColumnarMap(
+ keyColumnVector, valueColumnVector, start, end - start);
+ }
+
+ @Override
+ public ColumnVector getKeyColumnVector() {
+ init();
+ return keyColumnVector;
+ }
+
+ @Override
+ public ColumnVector getValueColumnVector() {
+ init();
+ return valueColumnVector;
+ }
+ };
+ }
+
+ @Override
+ public Arrow2PaimonVectorConverter visit(RowType rowType) {
+ final List<Arrow2PaimonVectorConverter> convertors = new ArrayList<>();
+ for (int i = 0; i < rowType.getFields().size(); i++) {
+ convertors.add(rowType.getTypeAt(i).accept(this));
+ }
+
+ return vector ->
+ new RowColumnVector() {
+
+ private boolean inited = false;
+ private VectorizedColumnBatch vectorizedColumnBatch;
+
+ private void init() {
+ if (!inited) {
+ List<FieldVector> children =
+ ((StructVector) vector).getChildrenFromFields();
+ ColumnVector[] vectors = new ColumnVector[children.size()];
+ for (int i = 0; i < children.size(); i++) {
+ vectors[i] = convertors.get(i).convertVector(children.get(i));
+ }
+ this.vectorizedColumnBatch = new VectorizedColumnBatch(vectors);
+ inited = true;
+ }
+ }
+
+ @Override
+ public boolean isNullAt(int index) {
+ return vector.isNull(index);
+ }
+
+ @Override
+ public InternalRow getRow(int index) {
+ init();
+ return new ColumnarRow(vectorizedColumnBatch, index);
+ }
+
+ @Override
+ public VectorizedColumnBatch getBatch() {
+ init();
+ return vectorizedColumnBatch;
+ }
+ };
+ }
+ }
+}
diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java
new file mode 100644
index 000000000..ec76cd95d
--- /dev/null
+++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/reader/ArrowBatchReader.java
@@ -0,0 +1,73 @@
+/*
+ * 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.arrow.reader;
+
+import org.apache.paimon.arrow.converter.Arrow2PaimonVectorConverter;
+import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.data.columnar.ColumnVector;
+import org.apache.paimon.data.columnar.ColumnarRow;
+import org.apache.paimon.data.columnar.VectorizedColumnBatch;
+import org.apache.paimon.types.RowType;
+
+import org.apache.arrow.vector.VectorSchemaRoot;
+
+import java.util.Iterator;
+
+/** Reader from a {@link VectorSchemaRoot} to paimon rows. */
+public class ArrowBatchReader {
+
+ private final VectorizedColumnBatch batch;
+ private final Arrow2PaimonVectorConverter[] convertors;
+
+ public ArrowBatchReader(RowType rowType) {
+ ColumnVector[] columnVectors = new ColumnVector[rowType.getFieldCount()];
+ this.convertors = new Arrow2PaimonVectorConverter[rowType.getFieldCount()];
+ this.batch = new VectorizedColumnBatch(columnVectors);
+
+ for (int i = 0; i < columnVectors.length; i++) {
+ this.convertors[i] = Arrow2PaimonVectorConverter.construct(rowType.getTypeAt(i));
+ }
+ }
+
+ public Iterable<InternalRow> readBatch(VectorSchemaRoot vsr) {
+ for (int i = 0; i < batch.columns.length; i++) {
+ batch.columns[i] = convertors[i].convertVector(vsr.getVector(i));
+ }
+ int rowCount = vsr.getRowCount();
+
+ batch.setNumRows(vsr.getRowCount());
+ ColumnarRow columnarRow = new ColumnarRow(batch);
+ return () ->
+ new Iterator<InternalRow>() {
+ private int position = 0;
+
+ @Override
+ public boolean hasNext() {
+ return position < rowCount;
+ }
+
+ @Override
+ public InternalRow next() {
+ columnarRow.setRowId(position);
+ position++;
+ return columnarRow;
+ }
+ };
+ }
+}
diff --git a/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java b/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java
index b148bfabe..e5d7289c3 100644
--- a/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java
+++ b/paimon-arrow/src/main/java/org/apache/paimon/arrow/vector/ArrowFormatWriter.java
@@ -18,7 +18,6 @@
package org.apache.paimon.arrow.vector;
-import org.apache.paimon.annotation.VisibleForTesting;
import org.apache.paimon.arrow.ArrowUtils;
import org.apache.paimon.arrow.writer.ArrowFieldWriter;
import org.apache.paimon.arrow.writer.ArrowFieldWriterFactoryVisitor;
@@ -103,8 +102,7 @@ public class ArrowFormatWriter implements AutoCloseable {
allocator.close();
}
- @VisibleForTesting
- VectorSchemaRoot getVectorSchemaRoot() {
+ public VectorSchemaRoot getVectorSchemaRoot() {
return vectorSchemaRoot;
}
}
diff --git a/paimon-arrow/src/test/java/org/apache/paimon/arrow/vector/ArrowFormatWriterTest.java b/paimon-arrow/src/test/java/org/apache/paimon/arrow/vector/ArrowFormatWriterTest.java
index 1a256f91c..8d7b5b7e6 100644
--- a/paimon-arrow/src/test/java/org/apache/paimon/arrow/vector/ArrowFormatWriterTest.java
+++ b/paimon-arrow/src/test/java/org/apache/paimon/arrow/vector/ArrowFormatWriterTest.java
@@ -18,6 +18,7 @@
package org.apache.paimon.arrow.vector;
+import org.apache.paimon.arrow.reader.ArrowBatchReader;
import org.apache.paimon.data.BinaryString;
import org.apache.paimon.data.Decimal;
import org.apache.paimon.data.GenericRow;
@@ -34,6 +35,7 @@ import org.junit.jupiter.api.Test;
import java.math.BigDecimal;
import java.util.ArrayList;
+import java.util.Iterator;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ThreadLocalRandom;
@@ -82,6 +84,11 @@ public class ArrowFormatWriterTest {
public void testWrite() {
try (ArrowFormatWriter writer = new ArrowFormatWriter(PRIMITIVE_TYPE, 4096)) {
List<InternalRow> list = new ArrayList<>();
+ List<InternalRow.FieldGetter> fieldGetters = new ArrayList<>();
+
+ for (int i = 0; i < PRIMITIVE_TYPE.getFieldCount(); i++) {
+ fieldGetters.add(InternalRow.createFieldGetter(PRIMITIVE_TYPE.getTypeAt(i), i));
+ }
for (int i = 0; i < 1000; i++) {
list.add(GenericRow.of(randomRowValues(null)));
}
@@ -91,14 +98,18 @@ public class ArrowFormatWriterTest {
writer.flush();
VectorSchemaRoot vectorSchemaRoot = writer.getVectorSchemaRoot();
- for (int i = 0; i < vectorSchemaRoot.getRowCount(); i++) {
- Assertions.assertThat(list.get(i).getString(0).toString())
- .isEqualTo(vectorSchemaRoot.getVector(0).getObject(i).toString());
- }
+ ArrowBatchReader arrowBatchReader = new ArrowBatchReader(PRIMITIVE_TYPE);
+ Iterable<InternalRow> rows = arrowBatchReader.readBatch(vectorSchemaRoot);
+
+ Iterator<InternalRow> iterator = rows.iterator();
+ for (int i = 0; i < 1000; i++) {
+ InternalRow actual = iterator.next();
+ InternalRow expectec = list.get(i);
- for (int i = 0; i < vectorSchemaRoot.getRowCount(); i++) {
- Assertions.assertThat(list.get(i).getInt(14))
- .isEqualTo(vectorSchemaRoot.getVector(14).getObject(i));
+ for (InternalRow.FieldGetter fieldGetter : fieldGetters) {
+ Assertions.assertThat(fieldGetter.getFieldOrNull(actual))
+ .isEqualTo(fieldGetter.getFieldOrNull(expectec));
+ }
}
}
}