Skip to content

Refactor DefLevelIterables to improve optimized parquet writer #13714

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

Merged
raunaqmorarka merged 2 commits into from
Sep 8, 2022
Merged

Refactor DefLevelIterables to improve optimized parquet writer #13714

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e671258
Rename DefLevelIterables to DefLevelWriterProviders
raunaqmorarka Aug 17, 2022
390448e
Refactor DefLevelWriterProviders to improve optimized parquet writer
raunaqmorarka Aug 17, 2022
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
7 changes: 6 additions & 1 deletion core/trino-spi/src/main/java/io/trino/spi/block/ColumnarArray.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ private static ColumnarArray toColumnarArray(RunLengthEncodedBlock rleBlock)

private ColumnarArray(Block nullCheckBlock, int offsetsOffset, int[] offsets, Block elementsBlock)
{
this.nullCheckBlock = nullCheckBlock;
this.nullCheckBlock = requireNonNull(nullCheckBlock, "nullCheckBlock is null");
Copy link
Copy Markdown
Member

@sopel39 sopel39 Aug 19, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: separate commit

this.offsetsOffset = offsetsOffset;
this.offsets = offsets;
this.elementsBlock = elementsBlock;
Expand All @@ -128,6 +128,11 @@ public int getPositionCount()
return nullCheckBlock.getPositionCount();
}

public boolean mayHaveNull()
{
return nullCheckBlock.mayHaveNull();
}

public boolean isNull(int position)
{
return nullCheckBlock.isNull(position);
Expand Down
7 changes: 6 additions & 1 deletion core/trino-spi/src/main/java/io/trino/spi/block/ColumnarMap.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ private static ColumnarMap toColumnarMap(RunLengthEncodedBlock rleBlock)

private ColumnarMap(Block nullCheckBlock, int offsetsOffset, int[] offsets, Block keysBlock, Block valuesBlock)
{
this.nullCheckBlock = nullCheckBlock;
this.nullCheckBlock = requireNonNull(nullCheckBlock, "nullCheckBlock is null");
this.offsetsOffset = offsetsOffset;
this.offsets = offsets;
this.keysBlock = keysBlock;
Expand All @@ -131,6 +131,11 @@ public int getPositionCount()
return nullCheckBlock.getPositionCount();
}

public boolean mayHaveNull()
{
return nullCheckBlock.mayHaveNull();
}

public boolean isNull(int position)
{
return nullCheckBlock.isNull(position);
Expand Down
10 changes: 5 additions & 5 deletions lib/trino-parquet/src/main/java/io/trino/parquet/writer/ArrayColumnWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
package io.trino.parquet.writer;

import com.google.common.collect.ImmutableList;
import io.trino.parquet.writer.repdef.DefLevelIterable;
import io.trino.parquet.writer.repdef.DefLevelIterables;
import io.trino.parquet.writer.repdef.DefLevelWriterProvider;
import io.trino.parquet.writer.repdef.DefLevelWriterProviders;
import io.trino.parquet.writer.repdef.RepLevelIterable;
import io.trino.parquet.writer.repdef.RepLevelIterables;
import io.trino.spi.block.ColumnarArray;
Expand Down Expand Up @@ -49,9 +49,9 @@ public void writeBlock(ColumnChunk columnChunk)
ColumnarArray columnarArray = ColumnarArray.toColumnarArray(columnChunk.getBlock());
elementWriter.writeBlock(
new ColumnChunk(columnarArray.getElementsBlock(),
ImmutableList.<DefLevelIterable>builder()
.addAll(columnChunk.getDefLevelIterables())
.add(DefLevelIterables.of(columnarArray, maxDefinitionLevel))
ImmutableList.<DefLevelWriterProvider>builder()
.addAll(columnChunk.getDefLevelWriterProviders())
.add(DefLevelWriterProviders.of(columnarArray, maxDefinitionLevel))
.build(),
ImmutableList.<RepLevelIterable>builder()
.addAll(columnChunk.getRepLevelIterables())
Expand Down
12 changes: 6 additions & 6 deletions lib/trino-parquet/src/main/java/io/trino/parquet/writer/ColumnChunk.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
package io.trino.parquet.writer;

import com.google.common.collect.ImmutableList;
import io.trino.parquet.writer.repdef.DefLevelIterable;
import io.trino.parquet.writer.repdef.DefLevelWriterProvider;
import io.trino.parquet.writer.repdef.RepLevelIterable;
import io.trino.spi.block.Block;

Expand All @@ -25,24 +25,24 @@
public class ColumnChunk
{
private final Block block;
private final List<DefLevelIterable> defLevelIterables;
private final List<DefLevelWriterProvider> defLevelWriterProviders;
private final List<RepLevelIterable> repLevelIterables;

ColumnChunk(Block block)
{
this(block, ImmutableList.of(), ImmutableList.of());
}

ColumnChunk(Block block, List<DefLevelIterable> defLevelIterables, List<RepLevelIterable> repLevelIterables)
ColumnChunk(Block block, List<DefLevelWriterProvider> defLevelWriterProviders, List<RepLevelIterable> repLevelIterables)
{
this.block = requireNonNull(block, "block is null");
this.defLevelIterables = ImmutableList.copyOf(defLevelIterables);
this.defLevelWriterProviders = ImmutableList.copyOf(defLevelWriterProviders);
this.repLevelIterables = ImmutableList.copyOf(repLevelIterables);
}

List<DefLevelIterable> getDefLevelIterables()
List<DefLevelWriterProvider> getDefLevelWriterProviders()
{
return defLevelIterables;
return defLevelWriterProviders;
}

List<RepLevelIterable> getRepLevelIterables()
Expand Down
14 changes: 7 additions & 7 deletions lib/trino-parquet/src/main/java/io/trino/parquet/writer/MapColumnWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
package io.trino.parquet.writer;

import com.google.common.collect.ImmutableList;
import io.trino.parquet.writer.repdef.DefLevelIterable;
import io.trino.parquet.writer.repdef.DefLevelIterables;
import io.trino.parquet.writer.repdef.DefLevelWriterProvider;
import io.trino.parquet.writer.repdef.DefLevelWriterProviders;
import io.trino.parquet.writer.repdef.RepLevelIterable;
import io.trino.parquet.writer.repdef.RepLevelIterables;
import io.trino.spi.block.ColumnarMap;
Expand Down Expand Up @@ -50,16 +50,16 @@ public void writeBlock(ColumnChunk columnChunk)
{
ColumnarMap columnarMap = ColumnarMap.toColumnarMap(columnChunk.getBlock());

ImmutableList<DefLevelIterable> defLevelIterables = ImmutableList.<DefLevelIterable>builder()
.addAll(columnChunk.getDefLevelIterables())
.add(DefLevelIterables.of(columnarMap, maxDefinitionLevel)).build();
ImmutableList<DefLevelWriterProvider> defLevelWriterProviders = ImmutableList.<DefLevelWriterProvider>builder()
.addAll(columnChunk.getDefLevelWriterProviders())
.add(DefLevelWriterProviders.of(columnarMap, maxDefinitionLevel)).build();

ImmutableList<RepLevelIterable> repLevelIterables = ImmutableList.<RepLevelIterable>builder()
.addAll(columnChunk.getRepLevelIterables())
.add(RepLevelIterables.of(columnarMap, maxRepetitionLevel)).build();

keyWriter.writeBlock(new ColumnChunk(columnarMap.getKeysBlock(), defLevelIterables, repLevelIterables));
valueWriter.writeBlock(new ColumnChunk(columnarMap.getValuesBlock(), defLevelIterables, repLevelIterables));
keyWriter.writeBlock(new ColumnChunk(columnarMap.getKeysBlock(), defLevelWriterProviders, repLevelIterables));
valueWriter.writeBlock(new ColumnChunk(columnarMap.getValuesBlock(), defLevelWriterProviders, repLevelIterables));
}

@Override
Expand Down
48 changes: 13 additions & 35 deletions lib/trino-parquet/src/main/java/io/trino/parquet/writer/PrimitiveColumnWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,11 @@

import com.google.common.collect.ImmutableList;
import io.airlift.slice.Slices;
import io.trino.parquet.writer.repdef.DefLevelIterable;
import io.trino.parquet.writer.repdef.DefLevelIterables;
import io.trino.parquet.writer.repdef.DefLevelWriterProvider;
import io.trino.parquet.writer.repdef.DefLevelWriterProviders;
import io.trino.parquet.writer.repdef.RepLevelIterable;
import io.trino.parquet.writer.repdef.RepLevelIterables;
import io.trino.parquet.writer.valuewriter.PrimitiveValueWriter;
import io.trino.spi.block.Block;
import org.apache.parquet.bytes.BytesInput;
import org.apache.parquet.column.ColumnDescriptor;
import org.apache.parquet.column.Encoding;
Expand Down Expand Up @@ -50,6 +49,8 @@
import static com.google.common.collect.ImmutableList.toImmutableList;
import static io.trino.parquet.writer.ParquetCompressor.getCompressor;
import static io.trino.parquet.writer.ParquetDataOutput.createDataOutput;
import static io.trino.parquet.writer.repdef.DefLevelWriterProvider.DefinitionLevelWriter;
import static io.trino.parquet.writer.repdef.DefLevelWriterProvider.getRootDefinitionLevelWriter;
import static java.lang.Math.toIntExact;
import static java.util.Objects.requireNonNull;

Expand Down Expand Up @@ -116,38 +117,15 @@ public void writeBlock(ColumnChunk columnChunk)
// write values
primitiveValueWriter.write(columnChunk.getBlock());

if (columnChunk.getDefLevelIterables().isEmpty()) {
// write definition levels for flat data types
Block block = columnChunk.getBlock();
if (!block.mayHaveNull()) {
for (int position = 0; position < block.getPositionCount(); position++) {
definitionLevelWriter.writeInteger(maxDefinitionLevel);
}
}
else {
for (int position = 0; position < block.getPositionCount(); position++) {
byte isNull = (byte) (block.isNull(position) ? 1 : 0);
definitionLevelWriter.writeInteger(maxDefinitionLevel - isNull);
currentPageNullCounts += isNull;
}
}
valueCount += block.getPositionCount();
}
else {
// write definition levels for nested data types
Iterator<Integer> defIterator = DefLevelIterables.getIterator(ImmutableList.<DefLevelIterable>builder()
.addAll(columnChunk.getDefLevelIterables())
.add(DefLevelIterables.of(columnChunk.getBlock(), maxDefinitionLevel))
.build());
while (defIterator.hasNext()) {
int next = defIterator.next();
definitionLevelWriter.writeInteger(next);
if (next != maxDefinitionLevel) {
currentPageNullCounts++;
}
valueCount++;
}
}
List<DefLevelWriterProvider> defLevelWriterProviders = ImmutableList.<DefLevelWriterProvider>builder()
.addAll(columnChunk.getDefLevelWriterProviders())
.add(DefLevelWriterProviders.of(columnChunk.getBlock(), maxDefinitionLevel))
.build();
DefinitionLevelWriter rootDefinitionLevelWriter = getRootDefinitionLevelWriter(defLevelWriterProviders, definitionLevelWriter);

DefLevelWriterProvider.ValuesCount valuesCount = rootDefinitionLevelWriter.writeDefinitionLevels();
currentPageNullCounts += valuesCount.totalValuesCount() - valuesCount.maxDefinitionLevelValuesCount();
valueCount += valuesCount.totalValuesCount();

if (columnDescriptor.getMaxRepetitionLevel() > 0) {
// write repetition levels for nested types
Expand Down
12 changes: 6 additions & 6 deletions lib/trino-parquet/src/main/java/io/trino/parquet/writer/StructColumnWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
package io.trino.parquet.writer;

import com.google.common.collect.ImmutableList;
import io.trino.parquet.writer.repdef.DefLevelIterable;
import io.trino.parquet.writer.repdef.DefLevelIterables;
import io.trino.parquet.writer.repdef.DefLevelWriterProvider;
import io.trino.parquet.writer.repdef.DefLevelWriterProviders;
import io.trino.parquet.writer.repdef.RepLevelIterable;
import io.trino.parquet.writer.repdef.RepLevelIterables;
import io.trino.spi.block.Block;
Expand Down Expand Up @@ -50,9 +50,9 @@ public void writeBlock(ColumnChunk columnChunk)
ColumnarRow columnarRow = toColumnarRow(columnChunk.getBlock());
checkArgument(columnarRow.getFieldCount() == columnWriters.size(), "ColumnarRow field size %s is not equal to columnWriters size %s", columnarRow.getFieldCount(), columnWriters.size());

List<DefLevelIterable> defLevelIterables = ImmutableList.<DefLevelIterable>builder()
.addAll(columnChunk.getDefLevelIterables())
.add(DefLevelIterables.of(columnarRow, maxDefinitionLevel))
List<DefLevelWriterProvider> defLevelWriterProviders = ImmutableList.<DefLevelWriterProvider>builder()
.addAll(columnChunk.getDefLevelWriterProviders())
.add(DefLevelWriterProviders.of(columnarRow, maxDefinitionLevel))
.build();
List<RepLevelIterable> repLevelIterables = ImmutableList.<RepLevelIterable>builder()
.addAll(columnChunk.getRepLevelIterables())
Expand All @@ -62,7 +62,7 @@ public void writeBlock(ColumnChunk columnChunk)
for (int i = 0; i < columnWriters.size(); ++i) {
ColumnWriter columnWriter = columnWriters.get(i);
Block block = columnarRow.getField(i);
columnWriter.writeBlock(new ColumnChunk(block, defLevelIterables, repLevelIterables));
columnWriter.writeBlock(new ColumnChunk(block, defLevelWriterProviders, repLevelIterables));
}
}

Expand Down
29 changes: 0 additions & 29 deletions lib/trino-parquet/src/main/java/io/trino/parquet/writer/repdef/DefLevelIterable.java
View file
Open in desktop

This file was deleted.

Loading