diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java new file mode 100644 index 000000000000..99f6d07429aa --- /dev/null +++ b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java @@ -0,0 +1,208 @@ +/* + * 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.spark.data; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assumptions.assumeThat; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collection; +import java.util.List; +import java.util.Set; +import org.apache.arrow.vector.ValueVector; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.DataGenerator; +import org.apache.iceberg.data.DataGenerators; +import org.apache.iceberg.data.ReadFormatModelTests; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableSet; +import org.apache.iceberg.spark.SparkSchemaUtil; +import org.apache.iceberg.spark.SparkUtil; +import org.apache.iceberg.spark.data.vectorized.IcebergArrowColumnVector; +import org.apache.iceberg.types.Type; +import org.apache.spark.sql.vectorized.ColumnarBatch; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.FieldSource; + +class TestSparkVectorizedFormatModel extends ReadFormatModelTests { + + private static final FileFormat[] FILE_FORMATS = {FileFormat.ORC, FileFormat.PARQUET}; + + // Spark 3.5 has no TIME, TIMESTAMP_NANO or VARIANT type + private static final Set SPARK_UNSUPPORTED_TYPE_IDS = + Set.of(Type.TypeID.TIME, Type.TypeID.TIMESTAMP_NANO, Type.TypeID.VARIANT); + + // The Arrow-based Parquet reader cannot read nested or geospatial columns; + // the vectorized ORC reader supports nested types. + private static final Set ARROW_UNSUPPORTED_TYPE_IDS = + Set.of( + Type.TypeID.STRUCT, + Type.TypeID.LIST, + Type.TypeID.MAP, + Type.TypeID.GEOMETRY, + Type.TypeID.GEOGRAPHY); + + private static final Set PARQUET_UNSUPPORTED_TYPE_IDS = + ImmutableSet.builder() + .addAll(SPARK_UNSUPPORTED_TYPE_IDS) + .addAll(ARROW_UNSUPPORTED_TYPE_IDS) + .build(); + + private static final List FORMAT_AND_GENERATOR = + Arrays.stream(FILE_FORMATS) + .flatMap( + format -> + Arrays.stream(DataGenerators.ALL) + .filter(generator -> supportsVectorized(format, generator)) + .map(generator -> Arguments.of(format, generator))) + .toList(); + + private static boolean supportsVectorized(FileFormat format, DataGenerator generator) { + return supportsGenerator(format, generator) + && (format != FileFormat.PARQUET + || !containsType(generator.schema(), ARROW_UNSUPPORTED_TYPE_IDS)); + } + + @Override + protected Collection unsupportedTypeIds(FileFormat fileFormat) { + return fileFormat == FileFormat.PARQUET + ? PARQUET_UNSUPPORTED_TYPE_IDS + : SPARK_UNSUPPORTED_TYPE_IDS; + } + + @Override + protected boolean supportsBatchReads() { + return true; + } + + @Override + protected int countRows(List batches) { + return batches.stream().mapToInt(ColumnarBatch::numRows).sum(); + } + + @Override + protected void assertContainersReused(ColumnarBatch first, ColumnarBatch second) { + assertThat(second.numCols()).isEqualTo(first.numCols()); + + for (int i = 0; i < first.numCols(); i++) { + assertThat(arrowVector(second, i)).isSameAs(arrowVector(first, i)); + } + } + + @Override + @ParameterizedTest + @FieldSource("FILE_FORMATS") + protected void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws IOException { + assumeThat(fileFormat) + .as("Vectorized Parquet reader only populates _deleted when a delete filter is set") + .isNotEqualTo(FileFormat.PARQUET); + super.testReadMetadataColumnIsDeleted(fileFormat); + } + + @Override + @ParameterizedTest + @FieldSource("FILE_FORMATS") + protected void testReadMetadataColumnRowLineage(FileFormat fileFormat) throws IOException { + assumeRowLineageSupported(fileFormat); + super.testReadMetadataColumnRowLineage(fileFormat); + } + + @Override + @ParameterizedTest + @FieldSource("FILE_FORMATS") + protected void testReadMetadataColumnRowLineageExistingValues(FileFormat fileFormat) + throws IOException { + assumeRowLineageSupported(fileFormat); + super.testReadMetadataColumnRowLineageExistingValues(fileFormat); + } + + // TODO: remove when https://github.com/apache/iceberg/issues/18543 is fixed + private static void assumeRowLineageSupported(FileFormat fileFormat) { + assumeThat(fileFormat) + .as("Vectorized ORC reader does not apply row lineage") + .isNotEqualTo(FileFormat.ORC); + } + + private ValueVector arrowVector(ColumnarBatch batch, int col) { + assertThat(batch.column(col)).isInstanceOf(IcebergArrowColumnVector.class); + + IcebergArrowColumnVector vector = (IcebergArrowColumnVector) batch.column(col); + return vector.vectorAccessor().getVector(); + } + + @Override + protected void assertContainersNotReused(ColumnarBatch first, ColumnarBatch second) { + assertThat(second.numCols()).isEqualTo(first.numCols()); + + for (int i = 0; i < first.numCols(); i++) { + assertThat(arrowVector(second, i)).isNotSameAs(arrowVector(first, i)); + } + } + + @Override + protected Object engineSchema(Schema schema) { + return SparkSchemaUtil.convert(schema); + } + + @Override + protected Object convertConstantToEngine(Type type, Object value) { + return SparkUtil.internalToSpark(type, value); + } + + @Override + protected Class engineType() { + return ColumnarBatch.class; + } + + @Override + protected ColumnarBatch convertToEngine(Record record, Schema schema) { + throw new UnsupportedOperationException("Vectorized reads are read-only"); + } + + @Override + protected void assertEquals( + Schema schema, List expected, List actual) { + throw new UnsupportedOperationException("Use assertRecordsEqual for batch reads"); + } + + @Override + protected void assertRecordsEqual( + Schema schema, List expected, CloseableIterable actual) { + int rowOffset = 0; + + for (ColumnarBatch batch : actual) { + assertThat(rowOffset + batch.numRows()).isLessThanOrEqualTo(expected.size()); + + for (int i = 0; i < batch.numRows(); i++) { + TestHelpers.assertEquals( + schema, + InternalRowConverter.convert(schema, expected.get(rowOffset + i)), + batch.getRow(i)); + } + + rowOffset += batch.numRows(); + } + + assertThat(rowOffset).isEqualTo(expected.size()); + } +} diff --git a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java new file mode 100644 index 000000000000..88b4d7e21441 --- /dev/null +++ b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java @@ -0,0 +1,207 @@ +/* + * 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.spark.data; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assumptions.assumeThat; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collection; +import java.util.List; +import java.util.Set; +import org.apache.arrow.vector.ValueVector; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.DataGenerator; +import org.apache.iceberg.data.DataGenerators; +import org.apache.iceberg.data.ReadFormatModelTests; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableSet; +import org.apache.iceberg.spark.SparkSchemaUtil; +import org.apache.iceberg.spark.SparkUtil; +import org.apache.iceberg.spark.data.vectorized.IcebergArrowColumnVector; +import org.apache.iceberg.types.Type; +import org.apache.spark.sql.vectorized.ColumnarBatch; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.FieldSource; + +class TestSparkVectorizedFormatModel extends ReadFormatModelTests { + + private static final FileFormat[] FILE_FORMATS = {FileFormat.ORC, FileFormat.PARQUET}; + + private static final Set SPARK_UNSUPPORTED_TYPE_IDS = + Set.of(Type.TypeID.TIME, Type.TypeID.TIMESTAMP_NANO); + + // The Arrow-based Parquet reader cannot read nested or geospatial columns; + // the vectorized ORC reader supports nested types. + private static final Set ARROW_UNSUPPORTED_TYPE_IDS = + Set.of( + Type.TypeID.STRUCT, + Type.TypeID.LIST, + Type.TypeID.MAP, + Type.TypeID.GEOMETRY, + Type.TypeID.GEOGRAPHY); + + private static final Set PARQUET_UNSUPPORTED_TYPE_IDS = + ImmutableSet.builder() + .addAll(SPARK_UNSUPPORTED_TYPE_IDS) + .addAll(ARROW_UNSUPPORTED_TYPE_IDS) + .build(); + + private static final List FORMAT_AND_GENERATOR = + Arrays.stream(FILE_FORMATS) + .flatMap( + format -> + Arrays.stream(DataGenerators.ALL) + .filter(generator -> supportsVectorized(format, generator)) + .map(generator -> Arguments.of(format, generator))) + .toList(); + + private static boolean supportsVectorized(FileFormat format, DataGenerator generator) { + return supportsGenerator(format, generator) + && (format != FileFormat.PARQUET + || !containsType(generator.schema(), ARROW_UNSUPPORTED_TYPE_IDS)); + } + + @Override + protected Collection unsupportedTypeIds(FileFormat fileFormat) { + return fileFormat == FileFormat.PARQUET + ? PARQUET_UNSUPPORTED_TYPE_IDS + : SPARK_UNSUPPORTED_TYPE_IDS; + } + + @Override + protected boolean supportsBatchReads() { + return true; + } + + @Override + protected int countRows(List batches) { + return batches.stream().mapToInt(ColumnarBatch::numRows).sum(); + } + + @Override + protected void assertContainersReused(ColumnarBatch first, ColumnarBatch second) { + assertThat(second.numCols()).isEqualTo(first.numCols()); + + for (int i = 0; i < first.numCols(); i++) { + assertThat(arrowVector(second, i)).isSameAs(arrowVector(first, i)); + } + } + + @Override + @ParameterizedTest + @FieldSource("FILE_FORMATS") + protected void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws IOException { + assumeThat(fileFormat) + .as("Vectorized Parquet reader only populates _deleted when a delete filter is set") + .isNotEqualTo(FileFormat.PARQUET); + super.testReadMetadataColumnIsDeleted(fileFormat); + } + + @Override + @ParameterizedTest + @FieldSource("FILE_FORMATS") + protected void testReadMetadataColumnRowLineage(FileFormat fileFormat) throws IOException { + assumeRowLineageSupported(fileFormat); + super.testReadMetadataColumnRowLineage(fileFormat); + } + + @Override + @ParameterizedTest + @FieldSource("FILE_FORMATS") + protected void testReadMetadataColumnRowLineageExistingValues(FileFormat fileFormat) + throws IOException { + assumeRowLineageSupported(fileFormat); + super.testReadMetadataColumnRowLineageExistingValues(fileFormat); + } + + // TODO: remove when https://github.com/apache/iceberg/issues/18543 is fixed + private static void assumeRowLineageSupported(FileFormat fileFormat) { + assumeThat(fileFormat) + .as("Vectorized ORC reader does not apply row lineage") + .isNotEqualTo(FileFormat.ORC); + } + + private ValueVector arrowVector(ColumnarBatch batch, int col) { + assertThat(batch.column(col)).isInstanceOf(IcebergArrowColumnVector.class); + + IcebergArrowColumnVector vector = (IcebergArrowColumnVector) batch.column(col); + return vector.vectorAccessor().getVector(); + } + + @Override + protected void assertContainersNotReused(ColumnarBatch first, ColumnarBatch second) { + assertThat(second.numCols()).isEqualTo(first.numCols()); + + for (int i = 0; i < first.numCols(); i++) { + assertThat(arrowVector(second, i)).isNotSameAs(arrowVector(first, i)); + } + } + + @Override + protected Object engineSchema(Schema schema) { + return SparkSchemaUtil.convert(schema); + } + + @Override + protected Object convertConstantToEngine(Type type, Object value) { + return SparkUtil.internalToSpark(type, value); + } + + @Override + protected Class engineType() { + return ColumnarBatch.class; + } + + @Override + protected ColumnarBatch convertToEngine(Record record, Schema schema) { + throw new UnsupportedOperationException("Vectorized reads are read-only"); + } + + @Override + protected void assertEquals( + Schema schema, List expected, List actual) { + throw new UnsupportedOperationException("Use assertRecordsEqual for batch reads"); + } + + @Override + protected void assertRecordsEqual( + Schema schema, List expected, CloseableIterable actual) { + int rowOffset = 0; + + for (ColumnarBatch batch : actual) { + assertThat(rowOffset + batch.numRows()).isLessThanOrEqualTo(expected.size()); + + for (int i = 0; i < batch.numRows(); i++) { + TestHelpers.assertEquals( + schema, + InternalRowConverter.convert(schema, expected.get(rowOffset + i)), + batch.getRow(i)); + } + + rowOffset += batch.numRows(); + } + + assertThat(rowOffset).isEqualTo(expected.size()); + } +} diff --git a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java new file mode 100644 index 000000000000..88b4d7e21441 --- /dev/null +++ b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java @@ -0,0 +1,207 @@ +/* + * 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.spark.data; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assumptions.assumeThat; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collection; +import java.util.List; +import java.util.Set; +import org.apache.arrow.vector.ValueVector; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.DataGenerator; +import org.apache.iceberg.data.DataGenerators; +import org.apache.iceberg.data.ReadFormatModelTests; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableSet; +import org.apache.iceberg.spark.SparkSchemaUtil; +import org.apache.iceberg.spark.SparkUtil; +import org.apache.iceberg.spark.data.vectorized.IcebergArrowColumnVector; +import org.apache.iceberg.types.Type; +import org.apache.spark.sql.vectorized.ColumnarBatch; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.FieldSource; + +class TestSparkVectorizedFormatModel extends ReadFormatModelTests { + + private static final FileFormat[] FILE_FORMATS = {FileFormat.ORC, FileFormat.PARQUET}; + + private static final Set SPARK_UNSUPPORTED_TYPE_IDS = + Set.of(Type.TypeID.TIME, Type.TypeID.TIMESTAMP_NANO); + + // The Arrow-based Parquet reader cannot read nested or geospatial columns; + // the vectorized ORC reader supports nested types. + private static final Set ARROW_UNSUPPORTED_TYPE_IDS = + Set.of( + Type.TypeID.STRUCT, + Type.TypeID.LIST, + Type.TypeID.MAP, + Type.TypeID.GEOMETRY, + Type.TypeID.GEOGRAPHY); + + private static final Set PARQUET_UNSUPPORTED_TYPE_IDS = + ImmutableSet.builder() + .addAll(SPARK_UNSUPPORTED_TYPE_IDS) + .addAll(ARROW_UNSUPPORTED_TYPE_IDS) + .build(); + + private static final List FORMAT_AND_GENERATOR = + Arrays.stream(FILE_FORMATS) + .flatMap( + format -> + Arrays.stream(DataGenerators.ALL) + .filter(generator -> supportsVectorized(format, generator)) + .map(generator -> Arguments.of(format, generator))) + .toList(); + + private static boolean supportsVectorized(FileFormat format, DataGenerator generator) { + return supportsGenerator(format, generator) + && (format != FileFormat.PARQUET + || !containsType(generator.schema(), ARROW_UNSUPPORTED_TYPE_IDS)); + } + + @Override + protected Collection unsupportedTypeIds(FileFormat fileFormat) { + return fileFormat == FileFormat.PARQUET + ? PARQUET_UNSUPPORTED_TYPE_IDS + : SPARK_UNSUPPORTED_TYPE_IDS; + } + + @Override + protected boolean supportsBatchReads() { + return true; + } + + @Override + protected int countRows(List batches) { + return batches.stream().mapToInt(ColumnarBatch::numRows).sum(); + } + + @Override + protected void assertContainersReused(ColumnarBatch first, ColumnarBatch second) { + assertThat(second.numCols()).isEqualTo(first.numCols()); + + for (int i = 0; i < first.numCols(); i++) { + assertThat(arrowVector(second, i)).isSameAs(arrowVector(first, i)); + } + } + + @Override + @ParameterizedTest + @FieldSource("FILE_FORMATS") + protected void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws IOException { + assumeThat(fileFormat) + .as("Vectorized Parquet reader only populates _deleted when a delete filter is set") + .isNotEqualTo(FileFormat.PARQUET); + super.testReadMetadataColumnIsDeleted(fileFormat); + } + + @Override + @ParameterizedTest + @FieldSource("FILE_FORMATS") + protected void testReadMetadataColumnRowLineage(FileFormat fileFormat) throws IOException { + assumeRowLineageSupported(fileFormat); + super.testReadMetadataColumnRowLineage(fileFormat); + } + + @Override + @ParameterizedTest + @FieldSource("FILE_FORMATS") + protected void testReadMetadataColumnRowLineageExistingValues(FileFormat fileFormat) + throws IOException { + assumeRowLineageSupported(fileFormat); + super.testReadMetadataColumnRowLineageExistingValues(fileFormat); + } + + // TODO: remove when https://github.com/apache/iceberg/issues/18543 is fixed + private static void assumeRowLineageSupported(FileFormat fileFormat) { + assumeThat(fileFormat) + .as("Vectorized ORC reader does not apply row lineage") + .isNotEqualTo(FileFormat.ORC); + } + + private ValueVector arrowVector(ColumnarBatch batch, int col) { + assertThat(batch.column(col)).isInstanceOf(IcebergArrowColumnVector.class); + + IcebergArrowColumnVector vector = (IcebergArrowColumnVector) batch.column(col); + return vector.vectorAccessor().getVector(); + } + + @Override + protected void assertContainersNotReused(ColumnarBatch first, ColumnarBatch second) { + assertThat(second.numCols()).isEqualTo(first.numCols()); + + for (int i = 0; i < first.numCols(); i++) { + assertThat(arrowVector(second, i)).isNotSameAs(arrowVector(first, i)); + } + } + + @Override + protected Object engineSchema(Schema schema) { + return SparkSchemaUtil.convert(schema); + } + + @Override + protected Object convertConstantToEngine(Type type, Object value) { + return SparkUtil.internalToSpark(type, value); + } + + @Override + protected Class engineType() { + return ColumnarBatch.class; + } + + @Override + protected ColumnarBatch convertToEngine(Record record, Schema schema) { + throw new UnsupportedOperationException("Vectorized reads are read-only"); + } + + @Override + protected void assertEquals( + Schema schema, List expected, List actual) { + throw new UnsupportedOperationException("Use assertRecordsEqual for batch reads"); + } + + @Override + protected void assertRecordsEqual( + Schema schema, List expected, CloseableIterable actual) { + int rowOffset = 0; + + for (ColumnarBatch batch : actual) { + assertThat(rowOffset + batch.numRows()).isLessThanOrEqualTo(expected.size()); + + for (int i = 0; i < batch.numRows(); i++) { + TestHelpers.assertEquals( + schema, + InternalRowConverter.convert(schema, expected.get(rowOffset + i)), + batch.getRow(i)); + } + + rowOffset += batch.numRows(); + } + + assertThat(rowOffset).isEqualTo(expected.size()); + } +}