Skip to content

Flink: Avoid per-row allocations in Parquet array and map writers #16791

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wombatu-kun wants to merge 1 commit into
base: main
Choose a base branch
from
+117 −105
Open

Flink: Avoid per-row allocations in Parquet array and map writers #16791

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 39 additions & 35 deletions flink/v1.20/flink/src/main/java/org/apache/iceberg/flink/data/FlinkParquetWriters.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -490,33 +490,36 @@ public void write(int repetitionLevel, byte[] bytes) {
}

private static class ArrayDataWriter<E> extends ParquetValueWriters.RepeatedWriter<ArrayData, E> {
private final LogicalType elementType;
private final ArrayData.ElementGetter elementGetter;
private final ElementIterator elementIterator;

private ArrayDataWriter(
int definitionLevel,
int repetitionLevel,
ParquetValueWriter<E> writer,
LogicalType elementType) {
super(definitionLevel, repetitionLevel, writer);
this.elementType = elementType;
this.elementGetter = ArrayData.createElementGetter(elementType);
this.elementIterator = new ElementIterator();
}

@Override
protected Iterator<E> elements(ArrayData list) {
return new ElementIterator<>(list);
// The parent writer fully consumes the iterator inside a single write() call, so a single
// reusable instance avoids allocating an iterator per row.
elementIterator.reset(list);
return elementIterator;
}

private class ElementIterator<E> implements Iterator<E> {
private final int size;
private final ArrayData list;
private final ArrayData.ElementGetter getter;
private class ElementIterator implements Iterator<E> {
private ArrayData list;
private int size;
private int index;

private ElementIterator(ArrayData list) {
this.list = list;
size = list.size();
getter = ArrayData.createElementGetter(elementType);
index = 0;
private void reset(ArrayData newList) {
this.list = newList;
this.size = newList.size();
this.index = 0;
}

@Override
Expand All @@ -531,7 +534,7 @@ public E next() {
throw new NoSuchElementException();
}

E element = (E) getter.getElementOrNull(list, index);
E element = (E) elementGetter.getElementOrNull(list, index);
index += 1;

return element;
Expand All @@ -541,8 +544,9 @@ public E next() {

private static class MapDataWriter<K, V>
extends ParquetValueWriters.RepeatedKeyValueWriter<MapData, K, V> {
private final LogicalType keyType;
private final LogicalType valueType;
private final ArrayData.ElementGetter keyGetter;
private final ArrayData.ElementGetter valueGetter;
private final EntryIterator entryIterator;

private MapDataWriter(
int definitionLevel,
Expand All @@ -552,32 +556,32 @@ private MapDataWriter(
LogicalType keyType,
LogicalType valueType) {
super(definitionLevel, repetitionLevel, keyWriter, valueWriter);
this.keyType = keyType;
this.valueType = valueType;
this.keyGetter = ArrayData.createElementGetter(keyType);
this.valueGetter = ArrayData.createElementGetter(valueType);
this.entryIterator = new EntryIterator();
}

@Override
protected Iterator<Map.Entry<K, V>> pairs(MapData map) {
return new EntryIterator<>(map);
}

private class EntryIterator<K, V> implements Iterator<Map.Entry<K, V>> {
private final int size;
private final ArrayData keys;
private final ArrayData values;
private final ParquetValueReaders.ReusableEntry<K, V> entry;
private final ArrayData.ElementGetter keyGetter;
private final ArrayData.ElementGetter valueGetter;
// The parent writer fully consumes the iterator inside a single write() call, so a single
// reusable instance (and its reusable entry) avoids allocating per row.
entryIterator.reset(map);
return entryIterator;
}

private class EntryIterator implements Iterator<Map.Entry<K, V>> {
private final ParquetValueReaders.ReusableEntry<K, V> entry =
new ParquetValueReaders.ReusableEntry<>();
private ArrayData keys;
private ArrayData values;
private int size;
private int index;

private EntryIterator(MapData map) {
size = map.size();
keys = map.keyArray();
values = map.valueArray();
entry = new ParquetValueReaders.ReusableEntry<>();
keyGetter = ArrayData.createElementGetter(keyType);
valueGetter = ArrayData.createElementGetter(valueType);
index = 0;
private void reset(MapData map) {
this.keys = map.keyArray();
this.values = map.valueArray();
this.size = map.size();
this.index = 0;
}

@Override
Expand Down
74 changes: 39 additions & 35 deletions flink/v2.0/flink/src/main/java/org/apache/iceberg/flink/data/FlinkParquetWriters.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -490,33 +490,36 @@ public void write(int repetitionLevel, byte[] bytes) {
}

private static class ArrayDataWriter<E> extends ParquetValueWriters.RepeatedWriter<ArrayData, E> {
private final LogicalType elementType;
private final ArrayData.ElementGetter elementGetter;
private final ElementIterator elementIterator;

private ArrayDataWriter(
int definitionLevel,
int repetitionLevel,
ParquetValueWriter<E> writer,
LogicalType elementType) {
super(definitionLevel, repetitionLevel, writer);
this.elementType = elementType;
this.elementGetter = ArrayData.createElementGetter(elementType);
this.elementIterator = new ElementIterator();
}

@Override
protected Iterator<E> elements(ArrayData list) {
return new ElementIterator<>(list);
// The parent writer fully consumes the iterator inside a single write() call, so a single
// reusable instance avoids allocating an iterator per row.
elementIterator.reset(list);
return elementIterator;
}

private class ElementIterator<E> implements Iterator<E> {
private final int size;
private final ArrayData list;
private final ArrayData.ElementGetter getter;
private class ElementIterator implements Iterator<E> {
private ArrayData list;
private int size;
private int index;

private ElementIterator(ArrayData list) {
this.list = list;
size = list.size();
getter = ArrayData.createElementGetter(elementType);
index = 0;
private void reset(ArrayData newList) {
this.list = newList;
this.size = newList.size();
this.index = 0;
}

@Override
Expand All @@ -531,7 +534,7 @@ public E next() {
throw new NoSuchElementException();
}

E element = (E) getter.getElementOrNull(list, index);
E element = (E) elementGetter.getElementOrNull(list, index);
index += 1;

return element;
Expand All @@ -541,8 +544,9 @@ public E next() {

private static class MapDataWriter<K, V>
extends ParquetValueWriters.RepeatedKeyValueWriter<MapData, K, V> {
private final LogicalType keyType;
private final LogicalType valueType;
private final ArrayData.ElementGetter keyGetter;
private final ArrayData.ElementGetter valueGetter;
private final EntryIterator entryIterator;

private MapDataWriter(
int definitionLevel,
Expand All @@ -552,32 +556,32 @@ private MapDataWriter(
LogicalType keyType,
LogicalType valueType) {
super(definitionLevel, repetitionLevel, keyWriter, valueWriter);
this.keyType = keyType;
this.valueType = valueType;
this.keyGetter = ArrayData.createElementGetter(keyType);
this.valueGetter = ArrayData.createElementGetter(valueType);
this.entryIterator = new EntryIterator();
}

@Override
protected Iterator<Map.Entry<K, V>> pairs(MapData map) {
return new EntryIterator<>(map);
}

private class EntryIterator<K, V> implements Iterator<Map.Entry<K, V>> {
private final int size;
private final ArrayData keys;
private final ArrayData values;
private final ParquetValueReaders.ReusableEntry<K, V> entry;
private final ArrayData.ElementGetter keyGetter;
private final ArrayData.ElementGetter valueGetter;
// The parent writer fully consumes the iterator inside a single write() call, so a single
// reusable instance (and its reusable entry) avoids allocating per row.
entryIterator.reset(map);
return entryIterator;
}

private class EntryIterator implements Iterator<Map.Entry<K, V>> {
private final ParquetValueReaders.ReusableEntry<K, V> entry =
new ParquetValueReaders.ReusableEntry<>();
private ArrayData keys;
private ArrayData values;
private int size;
private int index;

private EntryIterator(MapData map) {
size = map.size();
keys = map.keyArray();
values = map.valueArray();
entry = new ParquetValueReaders.ReusableEntry<>();
keyGetter = ArrayData.createElementGetter(keyType);
valueGetter = ArrayData.createElementGetter(valueType);
index = 0;
private void reset(MapData map) {
this.keys = map.keyArray();
this.values = map.valueArray();
this.size = map.size();
this.index = 0;
}

@Override
Expand Down
74 changes: 39 additions & 35 deletions flink/v2.1/flink/src/main/java/org/apache/iceberg/flink/data/FlinkParquetWriters.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -512,33 +512,36 @@ public void write(int repetitionLevel, byte[] bytes) {
}

private static class ArrayDataWriter<E> extends ParquetValueWriters.RepeatedWriter<ArrayData, E> {
private final LogicalType elementType;
private final ArrayData.ElementGetter elementGetter;
private final ElementIterator elementIterator;

private ArrayDataWriter(
int definitionLevel,
int repetitionLevel,
ParquetValueWriter<E> writer,
LogicalType elementType) {
super(definitionLevel, repetitionLevel, writer);
this.elementType = elementType;
this.elementGetter = ArrayData.createElementGetter(elementType);
this.elementIterator = new ElementIterator();
}

@Override
protected Iterator<E> elements(ArrayData list) {
return new ElementIterator<>(list);
// The parent writer fully consumes the iterator inside a single write() call, so a single
// reusable instance avoids allocating an iterator per row.
elementIterator.reset(list);
return elementIterator;
}

private class ElementIterator<E> implements Iterator<E> {
private final int size;
private final ArrayData list;
private final ArrayData.ElementGetter getter;
private class ElementIterator implements Iterator<E> {
private ArrayData list;
private int size;
private int index;

private ElementIterator(ArrayData list) {
this.list = list;
size = list.size();
getter = ArrayData.createElementGetter(elementType);
index = 0;
private void reset(ArrayData newList) {
this.list = newList;
this.size = newList.size();
this.index = 0;
}

@Override
Expand All @@ -553,7 +556,7 @@ public E next() {
throw new NoSuchElementException();
}

E element = (E) getter.getElementOrNull(list, index);
E element = (E) elementGetter.getElementOrNull(list, index);
index += 1;

return element;
Expand All @@ -563,8 +566,9 @@ public E next() {

private static class MapDataWriter<K, V>
extends ParquetValueWriters.RepeatedKeyValueWriter<MapData, K, V> {
private final LogicalType keyType;
private final LogicalType valueType;
private final ArrayData.ElementGetter keyGetter;
private final ArrayData.ElementGetter valueGetter;
private final EntryIterator entryIterator;

private MapDataWriter(
int definitionLevel,
Expand All @@ -574,32 +578,32 @@ private MapDataWriter(
LogicalType keyType,
LogicalType valueType) {
super(definitionLevel, repetitionLevel, keyWriter, valueWriter);
this.keyType = keyType;
this.valueType = valueType;
this.keyGetter = ArrayData.createElementGetter(keyType);
this.valueGetter = ArrayData.createElementGetter(valueType);
this.entryIterator = new EntryIterator();
}

@Override
protected Iterator<Map.Entry<K, V>> pairs(MapData map) {
return new EntryIterator<>(map);
}

private class EntryIterator<K, V> implements Iterator<Map.Entry<K, V>> {
private final int size;
private final ArrayData keys;
private final ArrayData values;
private final ParquetValueReaders.ReusableEntry<K, V> entry;
private final ArrayData.ElementGetter keyGetter;
private final ArrayData.ElementGetter valueGetter;
// The parent writer fully consumes the iterator inside a single write() call, so a single
// reusable instance (and its reusable entry) avoids allocating per row.
entryIterator.reset(map);
return entryIterator;
}

private class EntryIterator implements Iterator<Map.Entry<K, V>> {
private final ParquetValueReaders.ReusableEntry<K, V> entry =
new ParquetValueReaders.ReusableEntry<>();
private ArrayData keys;
private ArrayData values;
private int size;
private int index;

private EntryIterator(MapData map) {
size = map.size();
keys = map.keyArray();
values = map.valueArray();
entry = new ParquetValueReaders.ReusableEntry<>();
keyGetter = ArrayData.createElementGetter(keyType);
valueGetter = ArrayData.createElementGetter(valueType);
index = 0;
private void reset(MapData map) {
this.keys = map.keyArray();
this.values = map.valueArray();
this.size = map.size();
this.index = 0;
}

@Override
Expand Down