Skip to content

Support single type union for ORC-vectorization reader #114

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
merged 3 commits into from
Jul 12, 2022
Merged

Support single type union for ORC-vectorization reader #114

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a9ca437
Support single type union for ORC-vectorization reader
Jun 15, 2022
8b8b87b
Support single type union for ORC-vectorization reader
Jun 15, 2022
8931f0f
Merge branch 'yiqdingin/single-type-union-non-vectorized' of https://…
Jun 15, 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
42 changes: 24 additions & 18 deletions spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkOrcReaders.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -432,36 +432,42 @@ public ColumnVector getChild(int ordinal) {
}

private static class UnionConverter implements Converter {
private final Types.StructType structType;
private final Type type;
private final List<Converter> optionConverters;

private UnionConverter(Type type, List<Converter> optionConverters) {
this.structType = type.asStructType();
this.type = type;
this.optionConverters = optionConverters;
}

@Override
public ColumnVector convert(org.apache.orc.storage.ql.exec.vector.ColumnVector vector, int batchSize,
long batchOffsetInFile) {
UnionColumnVector unionColumnVector = (UnionColumnVector) vector;
List<Types.NestedField> fields = structType.fields();
List<ColumnVector> fieldVectors = Lists.newArrayListWithExpectedSize(fields.size());

LongColumnVector longColumnVector = new LongColumnVector();
longColumnVector.vector = Arrays.stream(unionColumnVector.tags).asLongStream().toArray();
if (optionConverters.size() > 1) {
// the case of complex union with multiple types
List<Types.NestedField> fields = type.asStructType().fields();
List<ColumnVector> fieldVectors = Lists.newArrayListWithExpectedSize(fields.size());

LongColumnVector longColumnVector = new LongColumnVector();
longColumnVector.vector = Arrays.stream(unionColumnVector.tags).asLongStream().toArray();

fieldVectors.add(new PrimitiveOrcColumnVector(Types.IntegerType.get(), batchSize, longColumnVector,
OrcValueReaders.ints(), batchOffsetInFile));
for (int i = 0; i < fields.size() - 1; i += 1) {
fieldVectors.add(optionConverters.get(i).convert(unionColumnVector.fields[i], batchSize, batchOffsetInFile));
}

fieldVectors.add(new PrimitiveOrcColumnVector(Types.IntegerType.get(), batchSize, longColumnVector,
OrcValueReaders.ints(), batchOffsetInFile));
for (int i = 0; i < fields.size() - 1; i += 1) {
fieldVectors.add(optionConverters.get(i).convert(unionColumnVector.fields[i], batchSize, batchOffsetInFile));
return new BaseOrcColumnVector(type.asStructType(), batchSize, vector) {
@Override
public ColumnVector getChild(int ordinal) {
return fieldVectors.get(ordinal);
}
};
} else {
// the case of single type union
return optionConverters.get(0).convert(unionColumnVector.fields[0], batchSize, batchOffsetInFile);
}

return new BaseOrcColumnVector(structType, batchSize, vector) {
@Override
public ColumnVector getChild(int ordinal) {
return fieldVectors.get(ordinal);
}
};
}
}
}
6 changes: 0 additions & 6 deletions spark/src/test/java/org/apache/iceberg/spark/data/TestSparkOrcUnions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -293,7 +293,6 @@ public void testSingleTypeUnion() throws IOException {
}

// Test vectorized reader
/*
try (CloseableIterable<ColumnarBatch> reader = ORC.read(Files.localInput(orcFile))
.project(expectedSchema)
.createBatchedReaderFunc(readOrcSchema ->
Expand All @@ -304,7 +303,6 @@ public void testSingleTypeUnion() throws IOException {
assertEquals(expectedSchema, expectedFirstRow, actualRowsIt.next());
assertEquals(expectedSchema, expectedSecondRow, actualRowsIt.next());
}
*/
}

@Test
Expand Down Expand Up @@ -366,7 +364,6 @@ public void testSingleTypeUnionOfStruct() throws IOException {
}

// Test vectorized reader
/*
try (CloseableIterable<ColumnarBatch> reader = ORC.read(Files.localInput(orcFile))
.project(expectedSchema)
.createBatchedReaderFunc(readOrcSchema ->
Expand All @@ -377,7 +374,6 @@ public void testSingleTypeUnionOfStruct() throws IOException {
assertEquals(expectedSchema, expectedFirstRow, actualRowsIt.next());
assertEquals(expectedSchema, expectedSecondRow, actualRowsIt.next());
}
*/
}

@Test
Expand Down Expand Up @@ -441,7 +437,6 @@ public void testDeepNestedSingleTypeUnion() throws IOException {
}

// Test vectorized reader
/*
try (CloseableIterable<ColumnarBatch> reader = ORC.read(Files.localInput(orcFile))
.project(expectedSchema)
.createBatchedReaderFunc(readOrcSchema ->
Expand All @@ -452,7 +447,6 @@ public void testDeepNestedSingleTypeUnion() throws IOException {
assertEquals(expectedSchema, expectedFirstRow, actualRowsIt.next());
assertEquals(expectedSchema, expectedSecondRow, actualRowsIt.next());
}
*/
}

private Iterator<InternalRow> batchesToRows(Iterator<ColumnarBatch> batches) {
Expand Down