This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch release-2.3
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/release-2.3 by this push:
new ec5e1b70821 [FLINK-40705][python] Reset Arrow map child writers
between batches
ec5e1b70821 is described below
commit ec5e1b708219ec5bffb726cbf246769ba9801d2a
Author: auroflow <[email protected]>
AuthorDate: Wed Sep 16 16:09:07 2026 +0800
[FLINK-40705][python] Reset Arrow map child writers between batches
Reset key and value writers together with map vectors so consecutive
batches reuse aligned offsets. Cover direct and nested maps, nulls, empty maps,
and partial batches through Arrow IPC round trips.
This closes #29247.
---
.../table/runtime/arrow/writers/MapWriter.java | 14 ++++
.../table/runtime/arrow/ArrowReaderWriterTest.java | 92 ++++++++++++++++++++++
2 files changed, 106 insertions(+)
diff --git
a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/writers/MapWriter.java
b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/writers/MapWriter.java
index c2c960c3328..1440e2bac5f 100644
---
a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/writers/MapWriter.java
+++
b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/writers/MapWriter.java
@@ -84,6 +84,20 @@ public abstract class MapWriter<T> extends
ArrowFieldWriter<T> {
}
}
+ @Override
+ public void finish() {
+ super.finish();
+ keyWriter.finish();
+ valueWriter.finish();
+ }
+
+ @Override
+ public void reset() {
+ super.reset();
+ keyWriter.reset();
+ valueWriter.reset();
+ }
+
//
------------------------------------------------------------------------------------------
/** {@link MapWriter} for {@link RowData} input. */
diff --git
a/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowReaderWriterTest.java
b/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowReaderWriterTest.java
index 6441246a9be..2e4c18a3a47 100644
---
a/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowReaderWriterTest.java
+++
b/flink-python/src/test/java/org/apache/flink/table/runtime/arrow/ArrowReaderWriterTest.java
@@ -21,6 +21,7 @@ package org.apache.flink.table.runtime.arrow;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.table.data.DecimalData;
import org.apache.flink.table.data.GenericArrayData;
+import org.apache.flink.table.data.GenericMapData;
import org.apache.flink.table.data.GenericRowData;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.data.StringData;
@@ -39,6 +40,7 @@ import org.apache.flink.table.types.logical.FloatType;
import org.apache.flink.table.types.logical.IntType;
import org.apache.flink.table.types.logical.LocalZonedTimestampType;
import org.apache.flink.table.types.logical.LogicalType;
+import org.apache.flink.table.types.logical.MapType;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.table.types.logical.SmallIntType;
import org.apache.flink.table.types.logical.TimeType;
@@ -53,15 +55,24 @@ import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowStreamReader;
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Arrays;
+import java.util.Collections;
+import java.util.LinkedHashMap;
import java.util.List;
+import java.util.Map;
import java.util.Objects;
+import static org.assertj.core.api.Assertions.assertThat;
+
/** Tests for {@link ArrowReader} and {@link ArrowWriter} of RowData. */
class ArrowReaderWriterTest extends ArrowReaderWriterTestBase<RowData> {
private static List<LogicalType> fieldTypes = new ArrayList<>();
@@ -168,6 +179,87 @@ class ArrowReaderWriterTest extends
ArrowReaderWriterTestBase<RowData> {
return Tuple2.of(arrowWriter, arrowStreamWriter);
}
+ @ParameterizedTest
+ @ValueSource(ints = {1, 2})
+ void testMapsAcrossBatches(int batchSize) throws IOException {
+ final MapType mapType = new MapType(new IntType(false), new IntType());
+ final MapType nestedMapType =
+ new MapType(new IntType(false), RowType.of(RowType.of(new
IntType())));
+ final RowType type =
+ RowType.of(
+ mapType, new ArrayType(mapType), new
ArrayType(RowType.of(nestedMapType)));
+ final List<RowData> rows = new ArrayList<>();
+ for (Map<Integer, Integer> values :
+ Arrays.asList(
+ Map.of(1, 11, 2, 22),
+ Collections.<Integer, Integer>emptyMap(),
+ null,
+ Collections.<Integer, Integer>singletonMap(3, null),
+ Map.of(4, 44))) {
+ final GenericMapData map = values == null ? null : new
GenericMapData(values);
+ final Map<Integer, RowData> nestedValues = new LinkedHashMap<>();
+ if (values != null) {
+ values.forEach(
+ (key, value) ->
+ nestedValues.put(
+ key,
+ value == null
+ ? null
+ :
GenericRowData.of(GenericRowData.of(value))));
+ }
+ rows.add(
+ GenericRowData.of(
+ map,
+ new GenericArrayData(new Object[] {map}),
+ new GenericArrayData(
+ new Object[] {
+ GenericRowData.of(
+ values == null
+ ? null
+ : new
GenericMapData(nestedValues))
+ })));
+ }
+
+ final ByteArrayOutputStream output = new ByteArrayOutputStream();
+ try (BufferAllocator batchAllocator =
+ ArrowUtils.getRootAllocator()
+ .newChildAllocator("map-batches", 0,
Long.MAX_VALUE);
+ VectorSchemaRoot root =
+
VectorSchemaRoot.create(ArrowUtils.toArrowSchema(type), batchAllocator);
+ ArrowStreamWriter streamWriter = new ArrowStreamWriter(root,
null, output)) {
+ final ArrowWriter<RowData> writer =
ArrowUtils.createRowDataArrowWriter(root, type);
+ streamWriter.start();
+ for (int start = 0; start < rows.size(); start += batchSize) {
+ for (RowData row : rows.subList(start, Math.min(start +
batchSize, rows.size()))) {
+ writer.write(row);
+ }
+ writer.finish();
+ streamWriter.writeBatch();
+ writer.reset();
+ }
+ streamWriter.end();
+
+ try (ArrowStreamReader streamReader =
+ new ArrowStreamReader(
+ new ByteArrayInputStream(output.toByteArray()),
batchAllocator)) {
+ final RowDataSerializer serializer = new
RowDataSerializer(type);
+ final ArrowReader reader =
+
ArrowUtils.createArrowReader(streamReader.getVectorSchemaRoot(), type);
+ int rowIndex = 0;
+ while (streamReader.loadNextBatch()) {
+ final int rowCount =
streamReader.getVectorSchemaRoot().getRowCount();
+ assertThat(rowCount).isEqualTo(Math.min(batchSize,
rows.size() - rowIndex));
+ for (int i = 0; i < rowCount; i++) {
+
assertThat(serializer.toBinaryRow(reader.read(i)).copy())
+ .as("row %s", rowIndex)
+
.isEqualTo(serializer.toBinaryRow(rows.get(rowIndex++)).copy());
+ }
+ }
+ assertThat(rowIndex).isEqualTo(rows.size());
+ }
+ }
+ }
+
@Override
public RowData[] getTestData() {
RowData row1 =
