Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -185,4 +185,34 @@ public int numValues() {
return numRows;
}
}

public static class VariantVectorHolder extends VectorHolder {
private final VectorHolder metadataHolder;
private final VectorHolder valueHolder;
private final int numRows;

public VariantVectorHolder(
Types.NestedField icebergField,
int numRows,
VectorHolder metadataHolder,
VectorHolder valueHolder) {
super(icebergField);
this.numRows = numRows;
this.metadataHolder = metadataHolder;
this.valueHolder = valueHolder;
}

@Override
public int numValues() {
return numRows;
}

public VectorHolder metadataHolder() {
return metadataHolder;
}

public VectorHolder valueHolder() {
return valueHolder;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -965,4 +965,57 @@ public String toString() {
@Override
public void setBatchSize(int batchSize) {}
}

public static class VectorizedVariantReader extends VectorizedArrowReader {
private final VectorizedArrowReader metadataReader;
private final VectorizedArrowReader valueReader;

public VectorizedVariantReader(
Types.NestedField icebergField,
VectorizedArrowReader metadataReader,
VectorizedArrowReader valueReader) {
super(icebergField);
this.metadataReader = metadataReader;
this.valueReader = valueReader;
}

@Override
public VectorHolder read(VectorHolder reuse, int numValsToRead) {
VectorHolder reuseMetadata = null;
VectorHolder reuseValue = null;
if (reuse instanceof VectorHolder.VariantVectorHolder) {
VectorHolder.VariantVectorHolder variantReuse = (VectorHolder.VariantVectorHolder) reuse;
reuseMetadata = variantReuse.metadataHolder();
reuseValue = variantReuse.valueHolder();
}
VectorHolder metadataHolder = metadataReader.read(reuseMetadata, numValsToRead);
VectorHolder valueHolder = valueReader.read(reuseValue, numValsToRead);
return new VectorHolder.VariantVectorHolder(
icebergField(), numValsToRead, metadataHolder, valueHolder);
}

@Override
public void setRowGroupInfo(
PageReadStore source, Map<ColumnPath, ColumnChunkMetaData> metadata) {
metadataReader.setRowGroupInfo(source, metadata);
valueReader.setRowGroupInfo(source, metadata);
}

@Override
public void setBatchSize(int batchSize) {
metadataReader.setBatchSize(batchSize);
valueReader.setBatchSize(batchSize);
}

@Override
public void close() {
metadataReader.close();
valueReader.close();
}

@Override
public String toString() {
return "VectorizedVariantReader";
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import org.apache.iceberg.Schema;
import org.apache.iceberg.arrow.ArrowAllocation;
import org.apache.iceberg.arrow.vectorized.VectorizedArrowReader.ConstantVectorReader;
import org.apache.iceberg.parquet.ParquetVariantVisitor;
import org.apache.iceberg.parquet.TypeWithSchemaVisitor;
import org.apache.iceberg.parquet.VectorizedReader;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
Expand Down Expand Up @@ -154,14 +155,16 @@ public VectorizedReader<?> struct(
return null;
}

@Override
public ParquetVariantVisitor<VectorizedReader<?>> variantVisitor() {
return new VectorizedVariantVisitor(
currentPath(), parquetSchema, icebergSchema, rootAllocator, setArrowValidityVector);
}

@Override
public VectorizedReader<?> variant(
Types.VariantType iVariant, GroupType variant, VectorizedReader<?> result) {
if (iVariant != null) {
throw new UnsupportedOperationException(
"Vectorized reads are not supported yet for variant fields");
}
return null;
return result;
}

@Override
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,142 @@
/*
* 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.iceberg.arrow.vectorized;

import java.util.List;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.iceberg.Schema;
import org.apache.iceberg.parquet.ParquetVariantVisitor;
import org.apache.iceberg.parquet.VectorizedReader;
import org.apache.iceberg.types.Types;
import org.apache.parquet.column.ColumnDescriptor;
import org.apache.parquet.io.InvalidRecordException;
import org.apache.parquet.schema.GroupType;
import org.apache.parquet.schema.MessageType;
import org.apache.parquet.schema.PrimitiveType;

/** Creates Arrow readers for a variant column's metadata and value leaves. */
class VectorizedVariantVisitor extends ParquetVariantVisitor<VectorizedReader<?>> {
private final String[] variantGroupPath;
private final MessageType parquetSchema;
private final Schema icebergSchema;
private final BufferAllocator allocator;
private final boolean setArrowValidityVector;

VectorizedVariantVisitor(
String[] variantGroupPath,
MessageType parquetSchema,
Schema icebergSchema,
BufferAllocator allocator,
boolean setArrowValidityVector) {
this.variantGroupPath = variantGroupPath;
this.parquetSchema = parquetSchema;
this.icebergSchema = icebergSchema;
this.allocator = allocator;
this.setArrowValidityVector = setArrowValidityVector;
}

@Override
public VectorizedReader<?> metadata(PrimitiveType metadata) {
ColumnDescriptor desc = resolveDescriptor(metadata);
if (desc == null) {
return null;
}

Types.NestedField field =
Types.NestedField.required(-1, metadata.getName(), Types.BinaryType.get());
return new VectorizedArrowReader(desc, field, allocator, setArrowValidityVector);
}

@Override
public VectorizedReader<?> serialized(PrimitiveType value) {
ColumnDescriptor desc = resolveDescriptor(value);
if (desc == null) {
return null;
}

Types.NestedField field =
Types.NestedField.optional(-1, value.getName(), Types.BinaryType.get());
return new VectorizedArrowReader(desc, field, allocator, setArrowValidityVector);
}

@Override
public VectorizedReader<?> variant(
GroupType variant, VectorizedReader<?> metadataResult, VectorizedReader<?> valueResult) {
if (metadataResult instanceof VectorizedArrowReader
&& valueResult instanceof VectorizedArrowReader) {
Types.NestedField field = findVariantField(variant);
if (field != null) {
return new VectorizedArrowReader.VectorizedVariantReader(
field, (VectorizedArrowReader) metadataResult, (VectorizedArrowReader) valueResult);
}
}

return null;
}

@Override
public VectorizedReader<?> primitive(PrimitiveType primitive) {
throw new UnsupportedOperationException("Unsupported variant: shredded typed_value primitive");
}

@Override
public VectorizedReader<?> value(
GroupType value, VectorizedReader<?> valueResult, VectorizedReader<?> typedResult) {
if (typedResult != null) {
throw new UnsupportedOperationException(
"Unsupported variant: shredded typed_value primitive");
}
return valueResult;
Comment thread
nssalian marked this conversation as resolved.
}

@Override
public VectorizedReader<?> object(
GroupType object, VectorizedReader<?> valueResult, List<VectorizedReader<?>> fieldResults) {
throw new UnsupportedOperationException(
"Unsupported variant: shredded typed_value object with " + fieldResults.size() + " fields");
}

@Override
public VectorizedReader<?> array(
GroupType array, VectorizedReader<?> valueResult, VectorizedReader<?> elementResult) {
throw new UnsupportedOperationException("Unsupported variant: shredded typed_value array");
}

private ColumnDescriptor resolveDescriptor(PrimitiveType primitive) {
// Build full column path: variant group path + primitive name
// e.g., ["v1", "metadata"] or ["v2", "value"]
String[] path = new String[variantGroupPath.length + 1];
System.arraycopy(variantGroupPath, 0, path, 0, variantGroupPath.length);
path[variantGroupPath.length] = primitive.getName();

try {
return parquetSchema.getColumnDescription(path);
} catch (InvalidRecordException e) {
return null;
}
}

private Types.NestedField findVariantField(GroupType variant) {
if (variant.getId() != null) {
return icebergSchema.findField(variant.getId().intValue());
}

return null;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
package org.apache.iceberg.arrow.vectorized;

import static org.assertj.core.api.Assertions.assertThatNoException;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

import org.apache.iceberg.Schema;
import org.apache.iceberg.parquet.TypeWithSchemaVisitor;
Expand All @@ -38,7 +37,7 @@
public class TestVectorizedReaderBuilder {

@Test
public void testVariantNotSupportedInVectorizedReads() {
public void testVariantSupportedInVectorizedReads() {
Schema icebergSchema =
new Schema(
NestedField.required(1, "id", IntegerType.get()),
Expand All @@ -50,10 +49,9 @@ public void testVariantNotSupportedInVectorizedReads() {
new VectorizedReaderBuilder(
icebergSchema, parquetSchema, false, ImmutableMap.of(), readers -> null);

assertThatThrownBy(
() -> TypeWithSchemaVisitor.visit(icebergSchema.asStruct(), parquetSchema, builder))
.isInstanceOf(UnsupportedOperationException.class)
.hasMessageContaining("Vectorized reads are not supported yet for variant fields");
assertThatNoException()
.isThrownBy(
() -> TypeWithSchemaVisitor.visit(icebergSchema.asStruct(), parquetSchema, builder));
}

@Test
Expand Down
Loading
Loading