From e7a9d8f2b1b42ce35b3fffe26c352dd0c9edfcbc Mon Sep 17 00:00:00 2001 From: Huaxin Gao Date: Thu, 29 Feb 2024 17:50:41 -0800 Subject: [PATCH 01/32] Iceberg/Comet integration --- build.gradle | 2 + .../comet/CometIcebergColumnReader.java | 164 ++++++++++ .../CometIcebergColumnarBatchReader.java | 303 ++++++++++++++++++ .../CometIcebergConstantColumnReader.java | 39 +++ .../comet/CometIcebergDeleteColumnReader.java | 71 ++++ .../CometIcebergPositionColumnReader.java | 62 ++++ .../vectorized/comet/CometIcebergVector.java | 157 +++++++++ .../CometIcebergVectorizedReaderBuilder.java | 142 ++++++++ 8 files changed, 940 insertions(+) create mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnReader.java create mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnarBatchReader.java create mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergConstantColumnReader.java create mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergDeleteColumnReader.java create mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergPositionColumnReader.java create mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVector.java create mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVectorizedReaderBuilder.java diff --git a/build.gradle b/build.gradle index c8bfba7967ce..127f95334e13 100644 --- a/build.gradle +++ b/build.gradle @@ -40,6 +40,7 @@ buildscript { } } +String sparkMajorVersion = '3.4' String scalaVersion = System.getProperty("scalaVersion") != null ? System.getProperty("scalaVersion") : System.getProperty("defaultScalaVersion") String sparkVersionsString = System.getProperty("sparkVersions") != null ? System.getProperty("sparkVersions") : System.getProperty("defaultSparkVersions") List sparkVersions = sparkVersionsString != null && !sparkVersionsString.isEmpty() ? sparkVersionsString.split(",") : [] @@ -787,6 +788,7 @@ project(':iceberg-parquet') { exclude group: 'org.codehaus.jackson' } + compileOnly "org.apache.comet:comet-spark-spark${sparkMajorVersion}_${scalaVersion}:0.1.0-SNAPSHOT" compileOnly libs.avro.avro compileOnly(libs.hadoop2.client) { exclude group: 'org.apache.avro', module: 'avro' diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnReader.java new file mode 100644 index 000000000000..e4cf1dcf234e --- /dev/null +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnReader.java @@ -0,0 +1,164 @@ +/* + * 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.vectorized.comet; + +import java.io.IOException; +import java.util.Map; +import org.apache.comet.parquet.AbstractColumnReader; +import org.apache.comet.parquet.ColumnReader; +import org.apache.comet.parquet.TypeUtil; +import org.apache.comet.parquet.Utils; +import org.apache.comet.vector.CometVector; +import org.apache.iceberg.parquet.VectorizedReader; +import org.apache.iceberg.spark.SparkSchemaUtil; +import org.apache.iceberg.types.Types; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.page.PageReadStore; +import org.apache.parquet.column.page.PageReader; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; + +/** + * A Iceberg Parquet column reader backed by a Boson {@link ColumnReader}. This class should be used + * together with {@link CometIcebergVector}. + * + *

Example: + * + * 

+ *   BosonIcebergColumnReader reader = ...
+ *   reader.setBatchSize(batchSize);
+ *
+ *   while (hasMoreRowsToRead) {
+ *     if (endOfRowGroup) {
+ *       reader.reset();
+ *       PageReader pageReader = ...
+ *       reader.setPageReader(pageReader);
+ *     }
+ *
+ *     int numRows = ...
+ *     BosonIcebergVector vector = reader.read(null, numRows);
+ *
+ *     // consume the vector
+ *   }
+ *
+ *   reader.close();
+ *
+ */ +@SuppressWarnings("checkstyle:VisibilityModifier") +public class CometIcebergColumnReader implements VectorizedReader { + public static final int DEFAULT_BATCH_SIZE = 5000; + + private final DataType sparkType; + protected AbstractColumnReader delegate; + private final CometIcebergVector vector; + private final ColumnDescriptor descriptor; + protected boolean initialized = false; + protected int batchSize = DEFAULT_BATCH_SIZE; + + public CometIcebergColumnReader(DataType sparkType, ColumnDescriptor descriptor) { + this.sparkType = sparkType; + this.descriptor = descriptor; + this.vector = new CometIcebergVector(sparkType, false); + } + + public CometIcebergColumnReader(Types.NestedField field) { + DataType dataType = SparkSchemaUtil.convert(field.type()); + StructField structField = new StructField(field.name(), dataType, false, Metadata.empty()); + this.sparkType = dataType; + this.descriptor = TypeUtil.convertToParquet(structField); + this.vector = new CometIcebergVector(sparkType, false); + } + + public AbstractColumnReader getDelegate() { + return delegate; + } + + /** + * This method is to initialized/reset the ColumnReader. This needs to be called for each row + * group after readNextRowGroup, so a new dictionary encoding can be set for each of the new row + * groups. + */ + public void reset() { + if (delegate != null) { + delegate.close(); + } + + delegate = Utils.getColumnReader(sparkType, descriptor, batchSize, false, false); + initialized = true; + } + + @Override + public CometIcebergVector read(CometIcebergVector reuse, int numRows) { + delegate.readBatch(numRows); + CometVector bv = delegate.currentBatch(); + if (reuse == null) reuse = vector; + reuse.setDelegate(bv); + return reuse; + } + + public ColumnDescriptor getDescriptor() { + return descriptor; + } + + public CometIcebergVector getVector() { + return vector; + } + + /** Returns the Spark data type for this column. */ + public DataType getSparkType() { + return sparkType; + } + + /** + * Set the page reader to be 'pageReader'. + * + *

NOTE: this should be called before reading a new Parquet column chunk, and after {@link + * CometIcebergColumnReader#reset} is called. + */ + public void setPageReader(PageReader pageReader) throws IOException { + reset(); + if (!initialized) { + throw new IllegalStateException("Invalid state: 'reset' should be called first"); + } + ((ColumnReader) delegate).setPageReader(pageReader); + } + + @Override + public void close() { + if (delegate != null) { + delegate.close(); + } + } + + @Override + public void setBatchSize(int size) { + // Preconditions.checkState( + // !initialized, "'reset' shouldn't be called" + " before 'setBatchSize' is called"); + this.batchSize = size; + } + + @Override + public void setRowGroupInfo( + PageReadStore pageReadStore, Map map, long size) { + throw new UnsupportedOperationException("Not supported"); + } +} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnarBatchReader.java new file mode 100644 index 000000000000..401221b53dbd --- /dev/null +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnarBatchReader.java @@ -0,0 +1,303 @@ +/* + * 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.vectorized.comet; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import org.apache.comet.parquet.AbstractColumnReader; +import org.apache.comet.parquet.BatchReader; +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.DeleteFilter; +import org.apache.iceberg.deletes.PositionDeleteIndex; +import org.apache.iceberg.parquet.VectorizedReader; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.spark.SparkSchemaUtil; +import org.apache.iceberg.util.Pair; +import org.apache.parquet.column.page.PageReadStore; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.vectorized.ColumnVector; +import org.apache.spark.sql.vectorized.ColumnarBatch; + +/** + * {@link VectorizedReader} that returns Spark's {@link ColumnarBatch} to support Spark's vectorized + * read path. The {@link ColumnarBatch} returned is created by passing in the Arrow vectors + * populated via delegated read calls to {@linkplain CometIcebergColumnReader VectorReader(s)}. + */ +@SuppressWarnings("checkstyle:VisibilityModifier") +public class CometIcebergColumnarBatchReader implements VectorizedReader { + + private final CometIcebergColumnReader[] readers; + private final boolean hasIsDeletedColumn; + private DeleteFilter deletes = null; + private long rowStartPosInBatch = 0; + private final BatchReader delegate; + + public CometIcebergColumnarBatchReader(List> readers, Schema schema) { + this.readers = + readers.stream() + .map(CometIcebergColumnReader.class::cast) + .toArray(CometIcebergColumnReader[]::new); + this.hasIsDeletedColumn = + readers.stream().anyMatch(reader -> reader instanceof CometIcebergDeleteColumnReader); + + AbstractColumnReader[] abstractColumnReaders = new AbstractColumnReader[readers.size()]; + delegate = new BatchReader(abstractColumnReaders); + delegate.setSparkSchema(SparkSchemaUtil.convert(schema)); + } + + @Override + public void setRowGroupInfo( + PageReadStore pageStore, Map metaData, long rowPosition) { + for (int i = 0; i < readers.length; i++) { + try { + if (!(readers[i] instanceof CometIcebergConstantColumnReader) + && !(readers[i] instanceof CometIcebergPositionColumnReader) + && !(readers[i] instanceof CometIcebergDeleteColumnReader)) { + readers[i].reset(); + readers[i].setPageReader( + pageStore.getPageReader(((CometIcebergColumnReader) readers[i]).getDescriptor())); + } + } catch (IOException e) { + throw new UncheckedIOException("Failed to setRowGroupInfo for Comet vectorization", e); + } + } + + for (int i = 0; i < readers.length; i++) { + delegate.getColumnReaders()[i] = ((CometIcebergColumnReader) this.readers[i]).getDelegate(); + } + + this.rowStartPosInBatch = rowPosition; + } + + public void setDeleteFilter(DeleteFilter deleteFilter) { + this.deletes = deleteFilter; + } + + @Override + public final ColumnarBatch read(ColumnarBatch reuse, int numRowsToRead) { + ColumnarBatch columnarBatch = new ColumnBatchLoader(numRowsToRead).loadDataToColumnBatch(); + rowStartPosInBatch += numRowsToRead; + return columnarBatch; + } + + @Override + public void setBatchSize(int batchSize) { + for (CometIcebergColumnReader reader : readers) { + if (reader != null) { + reader.setBatchSize(batchSize); + } + } + } + + @Override + public void close() { + for (CometIcebergColumnReader reader : readers) { + if (reader != null) { + reader.close(); + } + } + } + + private class ColumnBatchLoader { + private final int numRowsToRead; + // the rowId mapping to skip deleted rows for all column vectors inside a batch, it is null when + // there is no deletes + private int[] rowIdMapping; + // the array to indicate if a row is deleted or not, it is null when there is no "_deleted" + // metadata column + private boolean[] isDeleted; + + ColumnBatchLoader(int numRowsToRead) { + Preconditions.checkArgument( + numRowsToRead > 0, "Invalid number of rows to read: %s", numRowsToRead); + this.numRowsToRead = numRowsToRead; + if (hasIsDeletedColumn) { + isDeleted = new boolean[numRowsToRead]; + } + } + + ColumnarBatch loadDataToColumnBatch() { + int numRowsUndeleted = initRowIdMapping(); + + ColumnVector[] arrowColumnVectors = readDataToColumnVectors(); + + ColumnarBatch newColumnarBatch = new ColumnarBatch(arrowColumnVectors); + newColumnarBatch.setNumRows(numRowsUndeleted); + + if (hasEqDeletes()) { + applyEqDelete(newColumnarBatch); + } + + if (hasIsDeletedColumn && rowIdMapping != null) { + // reset the row id mapping array, so that it doesn't filter out the deleted rows + for (int i = 0; i < numRowsToRead; i++) { + rowIdMapping[i] = i; + } + newColumnarBatch.setNumRows(numRowsToRead); + } + + if (hasIsDeletedColumn) { + readDeletedColumnIfNecessary(arrowColumnVectors); + } + + return newColumnarBatch; + } + + ColumnVector[] readDataToColumnVectors() { + ColumnVector[] columnVectors = new ColumnVector[readers.length]; + // Fetch rows for all readers in the delegate + delegate.nextBatch(numRowsToRead); + for (int i = 0; i < readers.length; i++) { + CometIcebergVector bv = ((CometIcebergColumnReader) readers[i]).getVector(); + bv.setDelegate(((CometIcebergColumnReader) readers[i]).getDelegate().currentBatch()); + bv.setRowIdMapping(rowIdMapping); + columnVectors[i] = bv; + } + + return columnVectors; + } + + boolean hasEqDeletes() { + return deletes != null && deletes.hasEqDeletes(); + } + + int initRowIdMapping() { + Pair posDeleteRowIdMapping = posDelRowIdMapping(); + if (posDeleteRowIdMapping != null) { + rowIdMapping = posDeleteRowIdMapping.first(); + return posDeleteRowIdMapping.second(); + } else { + rowIdMapping = initEqDeleteRowIdMapping(); + return numRowsToRead; + } + } + + Pair posDelRowIdMapping() { + if (deletes != null && deletes.hasPosDeletes()) { + return buildPosDelRowIdMapping(deletes.deletedRowPositions()); + } else { + return null; + } + } + + /** + * Build a row id mapping inside a batch, which skips deleted rows. Here is an example of how we + * delete 2 rows in a batch with 8 rows in total. [0,1,2,3,4,5,6,7] -- Original status of the + * row id mapping array [F,F,F,F,F,F,F,F] -- Original status of the isDeleted array Position + * delete 2, 6 [0,1,3,4,5,7,-,-] -- After applying position deletes [Set Num records to 6] + * [F,F,T,F,F,F,T,F] -- After applying position deletes + * + * @param deletedRowPositions a set of deleted row positions + * @return the mapping array and the new num of rows in a batch, null if no row is deleted + */ + Pair buildPosDelRowIdMapping(PositionDeleteIndex deletedRowPositions) { + if (deletedRowPositions == null) { + return null; + } + + int[] posDelRowIdMapping = new int[numRowsToRead]; + int originalRowId = 0; + int currentRowId = 0; + while (originalRowId < numRowsToRead) { + if (!deletedRowPositions.isDeleted(originalRowId + rowStartPosInBatch)) { + posDelRowIdMapping[currentRowId] = originalRowId; + currentRowId++; + } else { + if (hasIsDeletedColumn) { + isDeleted[originalRowId] = true; + } + + deletes.incrementDeleteCount(); + } + originalRowId++; + } + + if (currentRowId == numRowsToRead) { + // there is no delete in this batch + return null; + } else { + return Pair.of(posDelRowIdMapping, currentRowId); + } + } + + int[] initEqDeleteRowIdMapping() { + int[] eqDeleteRowIdMapping = null; + if (hasEqDeletes()) { + eqDeleteRowIdMapping = new int[numRowsToRead]; + for (int i = 0; i < numRowsToRead; i++) { + eqDeleteRowIdMapping[i] = i; + } + } + + return eqDeleteRowIdMapping; + } + + /** + * Filter out the equality deleted rows. Here is an example, [0,1,2,3,4,5,6,7] -- Original + * status of the row id mapping array [F,F,F,F,F,F,F,F] -- Original status of the isDeleted + * array Position delete 2, 6 [0,1,3,4,5,7,-,-] -- After applying position deletes [Set Num + * records to 6] [F,F,T,F,F,F,T,F] -- After applying position deletes Equality delete 1 <= x <= + * 3 [0,4,5,7,-,-,-,-] -- After applying equality deletes [Set Num records to 4] + * [F,T,T,T,F,F,T,F] -- After applying equality deletes + * + * @param columnarBatch the {@link ColumnarBatch} to apply the equality delete + */ + void applyEqDelete(ColumnarBatch columnarBatch) { + Iterator it = columnarBatch.rowIterator(); + int rowId = 0; + int currentRowId = 0; + while (it.hasNext()) { + InternalRow row = it.next(); + if (deletes.eqDeletedRowFilter().test(row)) { + // the row is NOT deleted + // skip deleted rows by pointing to the next undeleted row Id + rowIdMapping[currentRowId] = rowIdMapping[rowId]; + currentRowId++; + } else { + if (hasIsDeletedColumn) { + isDeleted[rowIdMapping[rowId]] = true; + } + + deletes.incrementDeleteCount(); + } + + rowId++; + } + + columnarBatch.setNumRows(currentRowId); + } + + protected void readDeletedColumnIfNecessary(ColumnVector[] columnVectors) { + for (int i = 0; i < readers.length; i++) { + if (readers[i] instanceof CometIcebergDeleteColumnReader) { + CometIcebergDeleteColumnReader deleteColumnReader = + new CometIcebergDeleteColumnReader<>(isDeleted); + deleteColumnReader.setBatchSize(numRowsToRead); + deleteColumnReader.read(null, numRowsToRead); + columnVectors[i] = deleteColumnReader.getVector(); + } + } + } + } +} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergConstantColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergConstantColumnReader.java new file mode 100644 index 000000000000..6359a21a1211 --- /dev/null +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergConstantColumnReader.java @@ -0,0 +1,39 @@ +/* + * 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.vectorized.comet; + +import org.apache.comet.parquet.ConstantColumnReader; +import org.apache.iceberg.types.Types; + +public class CometIcebergConstantColumnReader extends CometIcebergColumnReader { + private final T value; + + public CometIcebergConstantColumnReader(T value, Types.NestedField field) { + super(field); + this.value = value; + delegate = new ConstantColumnReader(getSparkType(), getDescriptor(), value, false); + } + + @Override + public void setBatchSize(int batchSize) { + delegate.setBatchSize(batchSize); + this.batchSize = batchSize; + initialized = true; + } +} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergDeleteColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergDeleteColumnReader.java new file mode 100644 index 000000000000..8c2e5ab3f60e --- /dev/null +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergDeleteColumnReader.java @@ -0,0 +1,71 @@ +/* + * 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.vectorized.comet; + +import org.apache.comet.parquet.ConstantColumnReader; +import org.apache.comet.parquet.MetadataColumnReader; +import org.apache.comet.parquet.Native; +import org.apache.comet.parquet.TypeUtil; +import org.apache.iceberg.types.Types; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; + +public class CometIcebergDeleteColumnReader extends CometIcebergColumnReader { + public CometIcebergDeleteColumnReader(Types.NestedField field) { + super(field); + delegate = new ConstantColumnReader(getSparkType(), getDescriptor(), false, false); + } + + public CometIcebergDeleteColumnReader(boolean[] isDeleted) { + super( + DataTypes.BooleanType, + TypeUtil.convertToParquet( + new StructField("deleted", DataTypes.BooleanType, false, Metadata.empty()))); + delegate = new DeleteColumnReader(isDeleted); + } + + @Override + public void setBatchSize(int batchSize) { + delegate.setBatchSize(batchSize); + this.batchSize = batchSize; + initialized = true; + } + + private static class DeleteColumnReader extends MetadataColumnReader { + private boolean[] isDeleted; + + DeleteColumnReader(boolean[] isDeleted) { + super( + DataTypes.BooleanType, + TypeUtil.convertToParquet( + new StructField("deleted", DataTypes.BooleanType, false, Metadata.empty())), + false); + this.isDeleted = isDeleted; + } + + @Override + public void readBatch(int total) { + Native.resetBatch(nativeHandle); + Native.setIsDeleted(nativeHandle, isDeleted); + + super.readBatch(total); + } + } +} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergPositionColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergPositionColumnReader.java new file mode 100644 index 000000000000..19b569a0a229 --- /dev/null +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergPositionColumnReader.java @@ -0,0 +1,62 @@ +/* + * 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.vectorized.comet; + +import org.apache.comet.parquet.MetadataColumnReader; +import org.apache.comet.parquet.Native; +import org.apache.iceberg.types.Types; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.spark.sql.types.DataTypes; + +public class CometIcebergPositionColumnReader extends CometIcebergColumnReader { + public CometIcebergPositionColumnReader(Types.NestedField field) { + super(field); + delegate = new PositionColumnReader(getDescriptor()); + } + + @Override + public void setBatchSize(int batchSize) { + delegate.setBatchSize(batchSize); + this.batchSize = batchSize; + initialized = true; + } + + private static class PositionColumnReader extends MetadataColumnReader { + /** The current position value of the column that are used to initialize this column reader. */ + private long position; + + PositionColumnReader(ColumnDescriptor descriptor) { + this(descriptor, 0L); + } + + PositionColumnReader(ColumnDescriptor descriptor, long position) { + super(DataTypes.LongType, descriptor, false); + this.position = position; + } + + @Override + public void readBatch(int total) { + Native.resetBatch(nativeHandle); + Native.setPosition(nativeHandle, position, total); + position += total; + + super.readBatch(total); + } + } +} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVector.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVector.java new file mode 100644 index 000000000000..e8d8dad0b0a5 --- /dev/null +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVector.java @@ -0,0 +1,157 @@ +/* + * 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.vectorized.comet; + +import org.apache.comet.vector.CometDelegateVector; +import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.Decimal; +import org.apache.spark.unsafe.types.UTF8String; + +/** + * A Spark ColumnVector implementation backed by Arrow arrays. This is used by Iceberg's + * vectorization through Comet + */ +@SuppressWarnings("checkstyle:VisibilityModifier") +public class CometIcebergVector extends CometDelegateVector { + + // the rowId mapping to skip deleted rows for all column vectors inside a batch + // Here is an example: + // [0,1,2,3,4,5,6,7] -- Original status of the row id mapping array + // Position delete 2, 6 + // [0,1,3,4,5,7,-,-] -- After applying position deletes [Set Num records to 6] + // Equality delete 1 <= x <= 3 + // [0,4,5,7,-,-,-,-] -- After applying equality deletes [Set Num records to 4] + protected int[] rowIdMapping; + + public CometIcebergVector(DataType type, boolean useDecimal128) { + super(type, useDecimal128); + } + + public void setRowIdMapping(int[] rowIdMapping) { + this.rowIdMapping = rowIdMapping; + } + + @Override + public boolean isNullAt(int rowId) { + int newRowId = rowId; + if (rowIdMapping != null) { + newRowId = rowIdMapping[rowId]; + } + return super.isNullAt(newRowId); + } + + @Override + public boolean getBoolean(int rowId) { + int newRowId = rowId; + if (rowIdMapping != null) { + newRowId = rowIdMapping[rowId]; + } + return super.getBoolean(newRowId); + } + + @Override + public byte getByte(int rowId) { + int newRowId = rowId; + if (rowIdMapping != null) { + newRowId = rowIdMapping[rowId]; + } + return super.getByte(newRowId); + } + + @Override + public short getShort(int rowId) { + int newRowId = rowId; + if (rowIdMapping != null) { + newRowId = rowIdMapping[rowId]; + } + return super.getShort(newRowId); + } + + @Override + public int getInt(int rowId) { + int newRowId = rowId; + if (rowIdMapping != null) { + newRowId = rowIdMapping[rowId]; + } + return super.getInt(newRowId); + } + + @Override + public long getLong(int rowId) { + int newRowId = rowId; + if (rowIdMapping != null) { + newRowId = rowIdMapping[rowId]; + } + return super.getLong(newRowId); + } + + @Override + public float getFloat(int rowId) { + int newRowId = rowId; + if (rowIdMapping != null) { + newRowId = rowIdMapping[rowId]; + } + return super.getFloat(newRowId); + } + + @Override + public double getDouble(int rowId) { + int newRowId = rowId; + if (rowIdMapping != null) { + newRowId = rowIdMapping[rowId]; + } + return super.getDouble(newRowId); + } + + @Override + public Decimal getDecimal(int rowId, int precision, int scale) { + int newRowId = rowId; + if (isNullAt(newRowId)) { + return null; + } + if (rowIdMapping != null) { + newRowId = rowIdMapping[rowId]; + } + return super.getDecimal(newRowId, precision, scale); + } + + @Override + public UTF8String getUTF8String(int rowId) { + int newRowId = rowId; + if (isNullAt(newRowId)) { + return null; + } + if (rowIdMapping != null) { + newRowId = rowIdMapping[rowId]; + } + return super.getUTF8String(newRowId); + } + + @Override + public byte[] getBinary(int rowId) { + int newRowId = rowId; + if (isNullAt(newRowId)) { + return null; + } + if (rowIdMapping != null) { + newRowId = rowIdMapping[rowId]; + } + return super.getBinary(newRowId); + } +} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVectorizedReaderBuilder.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVectorizedReaderBuilder.java new file mode 100644 index 000000000000..53a1eaf13485 --- /dev/null +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVectorizedReaderBuilder.java @@ -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.spark.data.vectorized.comet; + +import java.util.List; +import java.util.Map; +import java.util.function.Function; +import java.util.stream.IntStream; +import org.apache.iceberg.MetadataColumns; +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.DeleteFilter; +import org.apache.iceberg.parquet.TypeWithSchemaVisitor; +import org.apache.iceberg.parquet.VectorizedReader; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.relocated.com.google.common.collect.Maps; +import org.apache.iceberg.spark.SparkSchemaUtil; +import org.apache.iceberg.types.Types; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.schema.GroupType; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.Type; +import org.apache.spark.sql.catalyst.InternalRow; + +public class CometIcebergVectorizedReaderBuilder + extends TypeWithSchemaVisitor> { + + private final MessageType parquetSchema; + private final Schema icebergSchema; + private final Map idToConstant; + private final Function>, VectorizedReader> readerFactory; + private final DeleteFilter deleteFilter; + + public CometIcebergVectorizedReaderBuilder( + Schema expectedSchema, + MessageType parquetSchema, + Map idToConstant, + Function>, VectorizedReader> readerFactory, + DeleteFilter deleteFilter) { + this.parquetSchema = parquetSchema; + this.icebergSchema = expectedSchema; + this.idToConstant = idToConstant; + this.readerFactory = readerFactory; + this.deleteFilter = deleteFilter; + } + + @Override + public VectorizedReader message( + Types.StructType expected, MessageType message, List> fieldReaders) { + GroupType groupType = message.asGroupType(); + Map> readersById = Maps.newHashMap(); + List fields = groupType.getFields(); + + IntStream.range(0, fields.size()) + .filter(pos -> fields.get(pos).getId() != null) + .forEach(pos -> readersById.put(fields.get(pos).getId().intValue(), fieldReaders.get(pos))); + + List icebergFields = + expected != null ? expected.fields() : ImmutableList.of(); + + List> reorderedFields = + Lists.newArrayListWithExpectedSize(icebergFields.size()); + + for (Types.NestedField field : icebergFields) { + int id = field.fieldId(); + VectorizedReader reader = readersById.get(id); + if (idToConstant.containsKey(id)) { + CometIcebergConstantColumnReader constantReader = + new CometIcebergConstantColumnReader<>(idToConstant.get(id), field); + reorderedFields.add(constantReader); + } else if (id == MetadataColumns.ROW_POSITION.fieldId()) { + reorderedFields.add(new CometIcebergPositionColumnReader(field)); + } else if (id == MetadataColumns.IS_DELETED.fieldId()) { + CometIcebergColumnReader deleteReader = new CometIcebergDeleteColumnReader<>(field); + reorderedFields.add(deleteReader); + } else if (reader != null) { + reorderedFields.add(reader); + } else { + CometIcebergColumnReader constantReader = + new CometIcebergConstantColumnReader<>(null, field); + reorderedFields.add(constantReader); + } + } + return vectorizedReader(reorderedFields); + } + + protected VectorizedReader vectorizedReader(List> reorderedFields) { + VectorizedReader reader = readerFactory.apply(reorderedFields); + if (deleteFilter != null) { + ((CometIcebergColumnarBatchReader) reader).setDeleteFilter(deleteFilter); + } + return reader; + } + + @Override + public VectorizedReader struct( + Types.StructType expected, GroupType groupType, List> fieldReaders) { + if (expected != null) { + throw new UnsupportedOperationException( + "Vectorized reads are not supported yet for struct fields"); + } + return null; + } + + @Override + public VectorizedReader primitive( + org.apache.iceberg.types.Type.PrimitiveType expected, PrimitiveType primitive) { + + if (primitive.getId() == null) { + return null; + } + int parquetFieldId = primitive.getId().intValue(); + ColumnDescriptor desc = parquetSchema.getColumnDescription(currentPath()); + // Nested types not yet supported for vectorized reads + if (desc.getMaxRepetitionLevel() > 0) { + return null; + } + Types.NestedField icebergField = icebergSchema.findField(parquetFieldId); + if (icebergField == null) { + return null; + } + + return new CometIcebergColumnReader(SparkSchemaUtil.convert(icebergField.type()), desc); + } +} From 566030610a6ae8ce99150f52d66ec5b6d880d230 Mon Sep 17 00:00:00 2001 From: Huaxin Gao Date: Wed, 17 Apr 2024 23:04:35 -0700 Subject: [PATCH 02/32] address comments --- .../java/org/apache/iceberg/ReaderType.java | 24 ++ .../iceberg/spark/SparkSQLProperties.java | 5 + .../vectorized/BaseColumnBatchLoader.java | 199 ++++++++++++ ...lumnReader.java => CometColumnReader.java} | 35 +- .../comet/CometColumnarBatchReader.java | 150 +++++++++ ...er.java => CometConstantColumnReader.java} | 4 +- ...ader.java => CometDeleteColumnReader.java} | 6 +- .../CometIcebergColumnarBatchReader.java | 303 ------------------ ...er.java => CometPositionColumnReader.java} | 4 +- ...metIcebergVector.java => CometVector.java} | 8 +- ...java => CometVectorizedReaderBuilder.java} | 20 +- .../iceberg/spark/source/BaseBatchReader.java | 18 +- .../iceberg/spark/source/SparkBatch.java | 4 +- .../source/SparkColumnarReaderFactory.java | 9 +- 14 files changed, 438 insertions(+), 351 deletions(-) create mode 100644 api/src/main/java/org/apache/iceberg/ReaderType.java create mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/{CometIcebergColumnReader.java => CometColumnReader.java} (82%) create mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnarBatchReader.java rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/{CometIcebergConstantColumnReader.java => CometConstantColumnReader.java} (88%) rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/{CometIcebergDeleteColumnReader.java => CometDeleteColumnReader.java} (91%) delete mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnarBatchReader.java rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/{CometIcebergPositionColumnReader.java => CometPositionColumnReader.java} (93%) rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/{CometIcebergVector.java => CometVector.java} (93%) rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/{CometIcebergVectorizedReaderBuilder.java => CometVectorizedReaderBuilder.java} (86%) diff --git a/api/src/main/java/org/apache/iceberg/ReaderType.java b/api/src/main/java/org/apache/iceberg/ReaderType.java new file mode 100644 index 000000000000..89ec0e0169b5 --- /dev/null +++ b/api/src/main/java/org/apache/iceberg/ReaderType.java @@ -0,0 +1,24 @@ +/* + * 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; + +public enum ReaderType { + ICEBERG, + COMET +} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java index 1e8c732d2d33..3f6faf984426 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java @@ -19,6 +19,7 @@ package org.apache.iceberg.spark; import java.time.Duration; +import org.apache.iceberg.ReaderType; public class SparkSQLProperties { @@ -27,6 +28,10 @@ private SparkSQLProperties() {} // Controls whether vectorized reads are enabled public static final String VECTORIZATION_ENABLED = "spark.sql.iceberg.vectorization.enabled"; + // Controls whether which reader to use for vectorization + public static final String READER_TYPE = "spark.sql.iceberg.parquet.reader-type"; + public static final ReaderType READER_TYPE_DEFAULT = ReaderType.ICEBERG; + // Controls whether reading/writing timestamps without timezones is allowed @Deprecated public static final String HANDLE_TIMESTAMP_WITHOUT_TIMEZONE = diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java new file mode 100644 index 000000000000..b33415ddfe8e --- /dev/null +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java @@ -0,0 +1,199 @@ +/* + * 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.vectorized; + +import java.util.Iterator; +import org.apache.iceberg.data.DeleteFilter; +import org.apache.iceberg.deletes.PositionDeleteIndex; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.util.Pair; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.vectorized.ColumnVector; +import org.apache.spark.sql.vectorized.ColumnarBatch; + +@SuppressWarnings("checkstyle:VisibilityModifier") +public abstract class BaseColumnBatchLoader { + protected final int numRowsToRead; + // the rowId mapping to skip deleted rows for all column vectors inside a batch, it is null when + // there is no deletes + protected int[] rowIdMapping; + // the array to indicate if a row is deleted or not, it is null when there is no "_deleted" + // metadata column + protected boolean[] isDeleted; + private final boolean hasIsDeletedColumn; + private final DeleteFilter deletes; + private final long rowStartPosInBatch; + + protected BaseColumnBatchLoader( + int numRowsToRead, + boolean hasIsDeletedColumn, + DeleteFilter deletes, + long rowStartPosInBatch) { + Preconditions.checkArgument( + numRowsToRead > 0, "Invalid number of rows to read: %s", numRowsToRead); + this.numRowsToRead = numRowsToRead; + this.hasIsDeletedColumn = hasIsDeletedColumn; + this.deletes = deletes; + this.rowStartPosInBatch = rowStartPosInBatch; + if (hasIsDeletedColumn) { + isDeleted = new boolean[numRowsToRead]; + } + } + + public ColumnarBatch loadDataToColumnBatch() { + int numRowsUndeleted = initRowIdMapping(); + + ColumnVector[] arrowColumnVectors = readDataToColumnVectors(); + + ColumnarBatch newColumnarBatch = new ColumnarBatch(arrowColumnVectors); + newColumnarBatch.setNumRows(numRowsUndeleted); + + if (hasEqDeletes()) { + applyEqDelete(newColumnarBatch); + } + + if (hasIsDeletedColumn && rowIdMapping != null) { + // reset the row id mapping array, so that it doesn't filter out the deleted rows + for (int i = 0; i < numRowsToRead; i++) { + rowIdMapping[i] = i; + } + newColumnarBatch.setNumRows(numRowsToRead); + } + + if (hasIsDeletedColumn) { + readDeletedColumnIfNecessary(arrowColumnVectors); + } + + return newColumnarBatch; + } + + protected abstract ColumnVector[] readDataToColumnVectors(); + + protected abstract void readDeletedColumnIfNecessary(ColumnVector[] arrowColumnVectors); + + boolean hasEqDeletes() { + return deletes != null && deletes.hasEqDeletes(); + } + + int initRowIdMapping() { + Pair posDeleteRowIdMapping = posDelRowIdMapping(); + if (posDeleteRowIdMapping != null) { + rowIdMapping = posDeleteRowIdMapping.first(); + return posDeleteRowIdMapping.second(); + } else { + rowIdMapping = initEqDeleteRowIdMapping(); + return numRowsToRead; + } + } + + Pair posDelRowIdMapping() { + if (deletes != null && deletes.hasPosDeletes()) { + return buildPosDelRowIdMapping(deletes.deletedRowPositions()); + } else { + return null; + } + } + + /** + * Build a row id mapping inside a batch, which skips deleted rows. Here is an example of how we + * delete 2 rows in a batch with 8 rows in total. [0,1,2,3,4,5,6,7] -- Original status of the row + * id mapping array [F,F,F,F,F,F,F,F] -- Original status of the isDeleted array Position delete 2, + * 6 [0,1,3,4,5,7,-,-] -- After applying position deletes [Set Num records to 6] [F,F,T,F,F,F,T,F] + * -- After applying position deletes + * + * @param deletedRowPositions a set of deleted row positions + * @return the mapping array and the new num of rows in a batch, null if no row is deleted + */ + Pair buildPosDelRowIdMapping(PositionDeleteIndex deletedRowPositions) { + if (deletedRowPositions == null) { + return null; + } + + int[] posDelRowIdMapping = new int[numRowsToRead]; + int originalRowId = 0; + int currentRowId = 0; + while (originalRowId < numRowsToRead) { + if (!deletedRowPositions.isDeleted(originalRowId + rowStartPosInBatch)) { + posDelRowIdMapping[currentRowId] = originalRowId; + currentRowId++; + } else { + if (hasIsDeletedColumn) { + isDeleted[originalRowId] = true; + } + + deletes.incrementDeleteCount(); + } + originalRowId++; + } + + if (currentRowId == numRowsToRead) { + // there is no delete in this batch + return null; + } else { + return Pair.of(posDelRowIdMapping, currentRowId); + } + } + + int[] initEqDeleteRowIdMapping() { + int[] eqDeleteRowIdMapping = null; + if (hasEqDeletes()) { + eqDeleteRowIdMapping = new int[numRowsToRead]; + for (int i = 0; i < numRowsToRead; i++) { + eqDeleteRowIdMapping[i] = i; + } + } + + return eqDeleteRowIdMapping; + } + + /** + * Filter out the equality deleted rows. Here is an example, [0,1,2,3,4,5,6,7] -- Original status + * of the row id mapping array [F,F,F,F,F,F,F,F] -- Original status of the isDeleted array + * Position delete 2, 6 [0,1,3,4,5,7,-,-] -- After applying position deletes [Set Num records to + * 6] [F,F,T,F,F,F,T,F] -- After applying position deletes Equality delete 1 <= x <= 3 + * [0,4,5,7,-,-,-,-] -- After applying equality deletes [Set Num records to 4] [F,T,T,T,F,F,T,F] + * -- After applying equality deletes + * + * @param columnarBatch the {@link ColumnarBatch} to apply the equality delete + */ + void applyEqDelete(ColumnarBatch columnarBatch) { + Iterator it = columnarBatch.rowIterator(); + int rowId = 0; + int currentRowId = 0; + while (it.hasNext()) { + InternalRow row = it.next(); + if (deletes.eqDeletedRowFilter().test(row)) { + // the row is NOT deleted + // skip deleted rows by pointing to the next undeleted row Id + rowIdMapping[currentRowId] = rowIdMapping[rowId]; + currentRowId++; + } else { + if (hasIsDeletedColumn) { + isDeleted[rowIdMapping[rowId]] = true; + } + + deletes.incrementDeleteCount(); + } + + rowId++; + } + + columnarBatch.setNumRows(currentRowId); + } +} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java similarity index 82% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnReader.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java index e4cf1dcf234e..a134c985912b 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java @@ -24,7 +24,6 @@ import org.apache.comet.parquet.ColumnReader; import org.apache.comet.parquet.TypeUtil; import org.apache.comet.parquet.Utils; -import org.apache.comet.vector.CometVector; import org.apache.iceberg.parquet.VectorizedReader; import org.apache.iceberg.spark.SparkSchemaUtil; import org.apache.iceberg.types.Types; @@ -38,13 +37,13 @@ import org.apache.spark.sql.types.StructField; /** - * A Iceberg Parquet column reader backed by a Boson {@link ColumnReader}. This class should be used - * together with {@link CometIcebergVector}. + * A Iceberg Parquet column reader backed by a Comet {@link ColumnReader}. This class should be used + * together with {@link CometVector}. * *

Example: * * 

- *   BosonIcebergColumnReader reader = ...
+ *   CometColumnReader reader = ...
  *   reader.setBatchSize(batchSize);
  *
  *   while (hasMoreRowsToRead) {
@@ -55,7 +54,7 @@
  *     }
  *
  *     int numRows = ...
- *     BosonIcebergVector vector = reader.read(null, numRows);
+ *     CometVector vector = reader.read(null, numRows);
  *
  *     // consume the vector
  *   }
@@ -63,29 +62,29 @@
  *   reader.close();
  *
*/ -@SuppressWarnings("checkstyle:VisibilityModifier") -public class CometIcebergColumnReader implements VectorizedReader { +@SuppressWarnings({"checkstyle:VisibilityModifier", "ParameterAssignment"}) +class CometColumnReader implements VectorizedReader { public static final int DEFAULT_BATCH_SIZE = 5000; private final DataType sparkType; protected AbstractColumnReader delegate; - private final CometIcebergVector vector; + private final CometVector vector; private final ColumnDescriptor descriptor; protected boolean initialized = false; protected int batchSize = DEFAULT_BATCH_SIZE; - public CometIcebergColumnReader(DataType sparkType, ColumnDescriptor descriptor) { + CometColumnReader(DataType sparkType, ColumnDescriptor descriptor) { this.sparkType = sparkType; this.descriptor = descriptor; - this.vector = new CometIcebergVector(sparkType, false); + this.vector = new CometVector(sparkType, false); } - public CometIcebergColumnReader(Types.NestedField field) { + CometColumnReader(Types.NestedField field) { DataType dataType = SparkSchemaUtil.convert(field.type()); StructField structField = new StructField(field.name(), dataType, false, Metadata.empty()); this.sparkType = dataType; this.descriptor = TypeUtil.convertToParquet(structField); - this.vector = new CometIcebergVector(sparkType, false); + this.vector = new CometVector(sparkType, false); } public AbstractColumnReader getDelegate() { @@ -107,10 +106,12 @@ public void reset() { } @Override - public CometIcebergVector read(CometIcebergVector reuse, int numRows) { + public CometVector read(CometVector reuse, int numRows) { delegate.readBatch(numRows); - CometVector bv = delegate.currentBatch(); - if (reuse == null) reuse = vector; + org.apache.comet.vector.CometVector bv = delegate.currentBatch(); + if (reuse == null) { + reuse = vector; + } reuse.setDelegate(bv); return reuse; } @@ -119,7 +120,7 @@ public ColumnDescriptor getDescriptor() { return descriptor; } - public CometIcebergVector getVector() { + public CometVector getVector() { return vector; } @@ -132,7 +133,7 @@ public DataType getSparkType() { * Set the page reader to be 'pageReader'. * *

NOTE: this should be called before reading a new Parquet column chunk, and after {@link - * CometIcebergColumnReader#reset} is called. + * CometColumnReader#reset} is called. */ public void setPageReader(PageReader pageReader) throws IOException { reset(); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnarBatchReader.java new file mode 100644 index 000000000000..2fd0c32f7a02 --- /dev/null +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnarBatchReader.java @@ -0,0 +1,150 @@ +/* + * 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.vectorized.comet; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.List; +import java.util.Map; +import org.apache.comet.parquet.AbstractColumnReader; +import org.apache.comet.parquet.BatchReader; +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.DeleteFilter; +import org.apache.iceberg.parquet.VectorizedReader; +import org.apache.iceberg.spark.SparkSchemaUtil; +import org.apache.iceberg.spark.data.vectorized.BaseColumnBatchLoader; +import org.apache.parquet.column.page.PageReadStore; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.vectorized.ColumnVector; +import org.apache.spark.sql.vectorized.ColumnarBatch; + +/** + * {@link VectorizedReader} that returns Spark's {@link ColumnarBatch} to support Spark's vectorized + * read path. The {@link ColumnarBatch} returned is created by passing in the Arrow vectors + * populated via delegated read calls to {@linkplain CometColumnReader VectorReader(s)}. + */ +@SuppressWarnings("checkstyle:VisibilityModifier") +public class CometColumnarBatchReader implements VectorizedReader { + + private final CometColumnReader[] readers; + private final boolean hasIsDeletedColumn; + private DeleteFilter deletes = null; + private long rowStartPosInBatch = 0; + private final BatchReader delegate; + + public CometColumnarBatchReader(List> readers, Schema schema) { + this.readers = + readers.stream().map(CometColumnReader.class::cast).toArray(CometColumnReader[]::new); + this.hasIsDeletedColumn = + readers.stream().anyMatch(reader -> reader instanceof CometDeleteColumnReader); + + AbstractColumnReader[] abstractColumnReaders = new AbstractColumnReader[readers.size()]; + delegate = new BatchReader(abstractColumnReaders); + delegate.setSparkSchema(SparkSchemaUtil.convert(schema)); + } + + @Override + public void setRowGroupInfo( + PageReadStore pageStore, Map metaData, long rowPosition) { + for (int i = 0; i < readers.length; i++) { + try { + if (!(readers[i] instanceof CometConstantColumnReader) + && !(readers[i] instanceof CometPositionColumnReader) + && !(readers[i] instanceof CometDeleteColumnReader)) { + readers[i].reset(); + readers[i].setPageReader( + pageStore.getPageReader(((CometColumnReader) readers[i]).getDescriptor())); + } + } catch (IOException e) { + throw new UncheckedIOException("Failed to setRowGroupInfo for Comet vectorization", e); + } + } + + for (int i = 0; i < readers.length; i++) { + delegate.getColumnReaders()[i] = ((CometColumnReader) this.readers[i]).getDelegate(); + } + + this.rowStartPosInBatch = rowPosition; + } + + public void setDeleteFilter(DeleteFilter deleteFilter) { + this.deletes = deleteFilter; + } + + @Override + public final ColumnarBatch read(ColumnarBatch reuse, int numRowsToRead) { + ColumnarBatch columnarBatch = new ColumnBatchLoader(numRowsToRead).loadDataToColumnBatch(); + rowStartPosInBatch += numRowsToRead; + return columnarBatch; + } + + @Override + public void setBatchSize(int batchSize) { + for (CometColumnReader reader : readers) { + if (reader != null) { + reader.setBatchSize(batchSize); + } + } + } + + @Override + public void close() { + for (CometColumnReader reader : readers) { + if (reader != null) { + reader.close(); + } + } + } + + private class ColumnBatchLoader extends BaseColumnBatchLoader { + ColumnBatchLoader(int numRowsToRead) { + super(numRowsToRead, hasIsDeletedColumn, deletes, rowStartPosInBatch); + } + + @Override + protected ColumnVector[] readDataToColumnVectors() { + ColumnVector[] columnVectors = new ColumnVector[readers.length]; + // Fetch rows for all readers in the delegate + delegate.nextBatch(numRowsToRead); + for (int i = 0; i < readers.length; i++) { + CometVector bv = readers[i].getVector(); + org.apache.comet.vector.CometVector vector = readers[i].getDelegate().currentBatch(); + bv.setDelegate(vector); + bv.setRowIdMapping(rowIdMapping); + columnVectors[i] = bv; + } + + return columnVectors; + } + + @Override + protected void readDeletedColumnIfNecessary(ColumnVector[] columnVectors) { + for (int i = 0; i < readers.length; i++) { + if (readers[i] instanceof CometDeleteColumnReader) { + CometDeleteColumnReader deleteColumnReader = new CometDeleteColumnReader<>(isDeleted); + deleteColumnReader.setBatchSize(numRowsToRead); + deleteColumnReader.read(null, numRowsToRead); + columnVectors[i] = deleteColumnReader.getVector(); + } + } + } + } +} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergConstantColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometConstantColumnReader.java similarity index 88% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergConstantColumnReader.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometConstantColumnReader.java index 6359a21a1211..11eef7ff5e0c 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergConstantColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometConstantColumnReader.java @@ -21,10 +21,10 @@ import org.apache.comet.parquet.ConstantColumnReader; import org.apache.iceberg.types.Types; -public class CometIcebergConstantColumnReader extends CometIcebergColumnReader { +class CometConstantColumnReader extends CometColumnReader { private final T value; - public CometIcebergConstantColumnReader(T value, Types.NestedField field) { + CometConstantColumnReader(T value, Types.NestedField field) { super(field); this.value = value; delegate = new ConstantColumnReader(getSparkType(), getDescriptor(), value, false); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergDeleteColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometDeleteColumnReader.java similarity index 91% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergDeleteColumnReader.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometDeleteColumnReader.java index 8c2e5ab3f60e..ae948f9eca7b 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergDeleteColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometDeleteColumnReader.java @@ -27,13 +27,13 @@ import org.apache.spark.sql.types.Metadata; import org.apache.spark.sql.types.StructField; -public class CometIcebergDeleteColumnReader extends CometIcebergColumnReader { - public CometIcebergDeleteColumnReader(Types.NestedField field) { +class CometDeleteColumnReader extends CometColumnReader { + CometDeleteColumnReader(Types.NestedField field) { super(field); delegate = new ConstantColumnReader(getSparkType(), getDescriptor(), false, false); } - public CometIcebergDeleteColumnReader(boolean[] isDeleted) { + CometDeleteColumnReader(boolean[] isDeleted) { super( DataTypes.BooleanType, TypeUtil.convertToParquet( diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnarBatchReader.java deleted file mode 100644 index 401221b53dbd..000000000000 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergColumnarBatchReader.java +++ /dev/null @@ -1,303 +0,0 @@ -/* - * 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.vectorized.comet; - -import java.io.IOException; -import java.io.UncheckedIOException; -import java.util.Iterator; -import java.util.List; -import java.util.Map; -import org.apache.comet.parquet.AbstractColumnReader; -import org.apache.comet.parquet.BatchReader; -import org.apache.iceberg.Schema; -import org.apache.iceberg.data.DeleteFilter; -import org.apache.iceberg.deletes.PositionDeleteIndex; -import org.apache.iceberg.parquet.VectorizedReader; -import org.apache.iceberg.relocated.com.google.common.base.Preconditions; -import org.apache.iceberg.spark.SparkSchemaUtil; -import org.apache.iceberg.util.Pair; -import org.apache.parquet.column.page.PageReadStore; -import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; -import org.apache.parquet.hadoop.metadata.ColumnPath; -import org.apache.spark.sql.catalyst.InternalRow; -import org.apache.spark.sql.vectorized.ColumnVector; -import org.apache.spark.sql.vectorized.ColumnarBatch; - -/** - * {@link VectorizedReader} that returns Spark's {@link ColumnarBatch} to support Spark's vectorized - * read path. The {@link ColumnarBatch} returned is created by passing in the Arrow vectors - * populated via delegated read calls to {@linkplain CometIcebergColumnReader VectorReader(s)}. - */ -@SuppressWarnings("checkstyle:VisibilityModifier") -public class CometIcebergColumnarBatchReader implements VectorizedReader { - - private final CometIcebergColumnReader[] readers; - private final boolean hasIsDeletedColumn; - private DeleteFilter deletes = null; - private long rowStartPosInBatch = 0; - private final BatchReader delegate; - - public CometIcebergColumnarBatchReader(List> readers, Schema schema) { - this.readers = - readers.stream() - .map(CometIcebergColumnReader.class::cast) - .toArray(CometIcebergColumnReader[]::new); - this.hasIsDeletedColumn = - readers.stream().anyMatch(reader -> reader instanceof CometIcebergDeleteColumnReader); - - AbstractColumnReader[] abstractColumnReaders = new AbstractColumnReader[readers.size()]; - delegate = new BatchReader(abstractColumnReaders); - delegate.setSparkSchema(SparkSchemaUtil.convert(schema)); - } - - @Override - public void setRowGroupInfo( - PageReadStore pageStore, Map metaData, long rowPosition) { - for (int i = 0; i < readers.length; i++) { - try { - if (!(readers[i] instanceof CometIcebergConstantColumnReader) - && !(readers[i] instanceof CometIcebergPositionColumnReader) - && !(readers[i] instanceof CometIcebergDeleteColumnReader)) { - readers[i].reset(); - readers[i].setPageReader( - pageStore.getPageReader(((CometIcebergColumnReader) readers[i]).getDescriptor())); - } - } catch (IOException e) { - throw new UncheckedIOException("Failed to setRowGroupInfo for Comet vectorization", e); - } - } - - for (int i = 0; i < readers.length; i++) { - delegate.getColumnReaders()[i] = ((CometIcebergColumnReader) this.readers[i]).getDelegate(); - } - - this.rowStartPosInBatch = rowPosition; - } - - public void setDeleteFilter(DeleteFilter deleteFilter) { - this.deletes = deleteFilter; - } - - @Override - public final ColumnarBatch read(ColumnarBatch reuse, int numRowsToRead) { - ColumnarBatch columnarBatch = new ColumnBatchLoader(numRowsToRead).loadDataToColumnBatch(); - rowStartPosInBatch += numRowsToRead; - return columnarBatch; - } - - @Override - public void setBatchSize(int batchSize) { - for (CometIcebergColumnReader reader : readers) { - if (reader != null) { - reader.setBatchSize(batchSize); - } - } - } - - @Override - public void close() { - for (CometIcebergColumnReader reader : readers) { - if (reader != null) { - reader.close(); - } - } - } - - private class ColumnBatchLoader { - private final int numRowsToRead; - // the rowId mapping to skip deleted rows for all column vectors inside a batch, it is null when - // there is no deletes - private int[] rowIdMapping; - // the array to indicate if a row is deleted or not, it is null when there is no "_deleted" - // metadata column - private boolean[] isDeleted; - - ColumnBatchLoader(int numRowsToRead) { - Preconditions.checkArgument( - numRowsToRead > 0, "Invalid number of rows to read: %s", numRowsToRead); - this.numRowsToRead = numRowsToRead; - if (hasIsDeletedColumn) { - isDeleted = new boolean[numRowsToRead]; - } - } - - ColumnarBatch loadDataToColumnBatch() { - int numRowsUndeleted = initRowIdMapping(); - - ColumnVector[] arrowColumnVectors = readDataToColumnVectors(); - - ColumnarBatch newColumnarBatch = new ColumnarBatch(arrowColumnVectors); - newColumnarBatch.setNumRows(numRowsUndeleted); - - if (hasEqDeletes()) { - applyEqDelete(newColumnarBatch); - } - - if (hasIsDeletedColumn && rowIdMapping != null) { - // reset the row id mapping array, so that it doesn't filter out the deleted rows - for (int i = 0; i < numRowsToRead; i++) { - rowIdMapping[i] = i; - } - newColumnarBatch.setNumRows(numRowsToRead); - } - - if (hasIsDeletedColumn) { - readDeletedColumnIfNecessary(arrowColumnVectors); - } - - return newColumnarBatch; - } - - ColumnVector[] readDataToColumnVectors() { - ColumnVector[] columnVectors = new ColumnVector[readers.length]; - // Fetch rows for all readers in the delegate - delegate.nextBatch(numRowsToRead); - for (int i = 0; i < readers.length; i++) { - CometIcebergVector bv = ((CometIcebergColumnReader) readers[i]).getVector(); - bv.setDelegate(((CometIcebergColumnReader) readers[i]).getDelegate().currentBatch()); - bv.setRowIdMapping(rowIdMapping); - columnVectors[i] = bv; - } - - return columnVectors; - } - - boolean hasEqDeletes() { - return deletes != null && deletes.hasEqDeletes(); - } - - int initRowIdMapping() { - Pair posDeleteRowIdMapping = posDelRowIdMapping(); - if (posDeleteRowIdMapping != null) { - rowIdMapping = posDeleteRowIdMapping.first(); - return posDeleteRowIdMapping.second(); - } else { - rowIdMapping = initEqDeleteRowIdMapping(); - return numRowsToRead; - } - } - - Pair posDelRowIdMapping() { - if (deletes != null && deletes.hasPosDeletes()) { - return buildPosDelRowIdMapping(deletes.deletedRowPositions()); - } else { - return null; - } - } - - /** - * Build a row id mapping inside a batch, which skips deleted rows. Here is an example of how we - * delete 2 rows in a batch with 8 rows in total. [0,1,2,3,4,5,6,7] -- Original status of the - * row id mapping array [F,F,F,F,F,F,F,F] -- Original status of the isDeleted array Position - * delete 2, 6 [0,1,3,4,5,7,-,-] -- After applying position deletes [Set Num records to 6] - * [F,F,T,F,F,F,T,F] -- After applying position deletes - * - * @param deletedRowPositions a set of deleted row positions - * @return the mapping array and the new num of rows in a batch, null if no row is deleted - */ - Pair buildPosDelRowIdMapping(PositionDeleteIndex deletedRowPositions) { - if (deletedRowPositions == null) { - return null; - } - - int[] posDelRowIdMapping = new int[numRowsToRead]; - int originalRowId = 0; - int currentRowId = 0; - while (originalRowId < numRowsToRead) { - if (!deletedRowPositions.isDeleted(originalRowId + rowStartPosInBatch)) { - posDelRowIdMapping[currentRowId] = originalRowId; - currentRowId++; - } else { - if (hasIsDeletedColumn) { - isDeleted[originalRowId] = true; - } - - deletes.incrementDeleteCount(); - } - originalRowId++; - } - - if (currentRowId == numRowsToRead) { - // there is no delete in this batch - return null; - } else { - return Pair.of(posDelRowIdMapping, currentRowId); - } - } - - int[] initEqDeleteRowIdMapping() { - int[] eqDeleteRowIdMapping = null; - if (hasEqDeletes()) { - eqDeleteRowIdMapping = new int[numRowsToRead]; - for (int i = 0; i < numRowsToRead; i++) { - eqDeleteRowIdMapping[i] = i; - } - } - - return eqDeleteRowIdMapping; - } - - /** - * Filter out the equality deleted rows. Here is an example, [0,1,2,3,4,5,6,7] -- Original - * status of the row id mapping array [F,F,F,F,F,F,F,F] -- Original status of the isDeleted - * array Position delete 2, 6 [0,1,3,4,5,7,-,-] -- After applying position deletes [Set Num - * records to 6] [F,F,T,F,F,F,T,F] -- After applying position deletes Equality delete 1 <= x <= - * 3 [0,4,5,7,-,-,-,-] -- After applying equality deletes [Set Num records to 4] - * [F,T,T,T,F,F,T,F] -- After applying equality deletes - * - * @param columnarBatch the {@link ColumnarBatch} to apply the equality delete - */ - void applyEqDelete(ColumnarBatch columnarBatch) { - Iterator it = columnarBatch.rowIterator(); - int rowId = 0; - int currentRowId = 0; - while (it.hasNext()) { - InternalRow row = it.next(); - if (deletes.eqDeletedRowFilter().test(row)) { - // the row is NOT deleted - // skip deleted rows by pointing to the next undeleted row Id - rowIdMapping[currentRowId] = rowIdMapping[rowId]; - currentRowId++; - } else { - if (hasIsDeletedColumn) { - isDeleted[rowIdMapping[rowId]] = true; - } - - deletes.incrementDeleteCount(); - } - - rowId++; - } - - columnarBatch.setNumRows(currentRowId); - } - - protected void readDeletedColumnIfNecessary(ColumnVector[] columnVectors) { - for (int i = 0; i < readers.length; i++) { - if (readers[i] instanceof CometIcebergDeleteColumnReader) { - CometIcebergDeleteColumnReader deleteColumnReader = - new CometIcebergDeleteColumnReader<>(isDeleted); - deleteColumnReader.setBatchSize(numRowsToRead); - deleteColumnReader.read(null, numRowsToRead); - columnVectors[i] = deleteColumnReader.getVector(); - } - } - } - } -} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergPositionColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometPositionColumnReader.java similarity index 93% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergPositionColumnReader.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometPositionColumnReader.java index 19b569a0a229..7a3cd7cb5789 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergPositionColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometPositionColumnReader.java @@ -24,8 +24,8 @@ import org.apache.parquet.column.ColumnDescriptor; import org.apache.spark.sql.types.DataTypes; -public class CometIcebergPositionColumnReader extends CometIcebergColumnReader { - public CometIcebergPositionColumnReader(Types.NestedField field) { +class CometPositionColumnReader extends CometColumnReader { + CometPositionColumnReader(Types.NestedField field) { super(field); delegate = new PositionColumnReader(getDescriptor()); } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVector.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVector.java similarity index 93% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVector.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVector.java index e8d8dad0b0a5..2e8b8210b8e5 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVector.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVector.java @@ -23,12 +23,8 @@ import org.apache.spark.sql.types.Decimal; import org.apache.spark.unsafe.types.UTF8String; -/** - * A Spark ColumnVector implementation backed by Arrow arrays. This is used by Iceberg's - * vectorization through Comet - */ @SuppressWarnings("checkstyle:VisibilityModifier") -public class CometIcebergVector extends CometDelegateVector { +class CometVector extends CometDelegateVector { // the rowId mapping to skip deleted rows for all column vectors inside a batch // Here is an example: @@ -39,7 +35,7 @@ public class CometIcebergVector extends CometDelegateVector { // [0,4,5,7,-,-,-,-] -- After applying equality deletes [Set Num records to 4] protected int[] rowIdMapping; - public CometIcebergVector(DataType type, boolean useDecimal128) { + CometVector(DataType type, boolean useDecimal128) { super(type, useDecimal128); } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVectorizedReaderBuilder.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVectorizedReaderBuilder.java similarity index 86% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVectorizedReaderBuilder.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVectorizedReaderBuilder.java index 53a1eaf13485..f679ea7bfa86 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometIcebergVectorizedReaderBuilder.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVectorizedReaderBuilder.java @@ -39,8 +39,7 @@ import org.apache.parquet.schema.Type; import org.apache.spark.sql.catalyst.InternalRow; -public class CometIcebergVectorizedReaderBuilder - extends TypeWithSchemaVisitor> { +public class CometVectorizedReaderBuilder extends TypeWithSchemaVisitor> { private final MessageType parquetSchema; private final Schema icebergSchema; @@ -48,7 +47,7 @@ public class CometIcebergVectorizedReaderBuilder private final Function>, VectorizedReader> readerFactory; private final DeleteFilter deleteFilter; - public CometIcebergVectorizedReaderBuilder( + public CometVectorizedReaderBuilder( Schema expectedSchema, MessageType parquetSchema, Map idToConstant, @@ -82,19 +81,18 @@ public VectorizedReader message( int id = field.fieldId(); VectorizedReader reader = readersById.get(id); if (idToConstant.containsKey(id)) { - CometIcebergConstantColumnReader constantReader = - new CometIcebergConstantColumnReader<>(idToConstant.get(id), field); + CometConstantColumnReader constantReader = + new CometConstantColumnReader<>(idToConstant.get(id), field); reorderedFields.add(constantReader); } else if (id == MetadataColumns.ROW_POSITION.fieldId()) { - reorderedFields.add(new CometIcebergPositionColumnReader(field)); + reorderedFields.add(new CometPositionColumnReader(field)); } else if (id == MetadataColumns.IS_DELETED.fieldId()) { - CometIcebergColumnReader deleteReader = new CometIcebergDeleteColumnReader<>(field); + CometColumnReader deleteReader = new CometDeleteColumnReader<>(field); reorderedFields.add(deleteReader); } else if (reader != null) { reorderedFields.add(reader); } else { - CometIcebergColumnReader constantReader = - new CometIcebergConstantColumnReader<>(null, field); + CometColumnReader constantReader = new CometConstantColumnReader<>(null, field); reorderedFields.add(constantReader); } } @@ -104,7 +102,7 @@ public VectorizedReader message( protected VectorizedReader vectorizedReader(List> reorderedFields) { VectorizedReader reader = readerFactory.apply(reorderedFields); if (deleteFilter != null) { - ((CometIcebergColumnarBatchReader) reader).setDeleteFilter(deleteFilter); + ((CometColumnarBatchReader) reader).setDeleteFilter(deleteFilter); } return reader; } @@ -137,6 +135,6 @@ public VectorizedReader primitive( return null; } - return new CometIcebergColumnReader(SparkSchemaUtil.convert(icebergField.type()), desc); + return new CometColumnReader(SparkSchemaUtil.convert(icebergField.type()), desc); } } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java index c05b694a60dc..7321b49d1c82 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java @@ -22,6 +22,7 @@ import java.util.Set; import org.apache.iceberg.FileFormat; import org.apache.iceberg.MetadataColumns; +import org.apache.iceberg.ReaderType; import org.apache.iceberg.ScanTask; import org.apache.iceberg.ScanTaskGroup; import org.apache.iceberg.Schema; @@ -39,6 +40,7 @@ abstract class BaseBatchReader extends BaseReader { private final int batchSize; + private ReaderType readerType; BaseBatchReader( Table table, @@ -51,6 +53,10 @@ abstract class BaseBatchReader extends BaseReader newBatchIterable( InputFile inputFile, FileFormat format, @@ -86,9 +92,15 @@ private CloseableIterable newParquetIterable( .project(requiredSchema) .split(start, length) .createBatchedReaderFunc( - fileSchema -> - VectorizedSparkParquetReaders.buildReader( - requiredSchema, fileSchema, idToConstant, deleteFilter)) + fileSchema -> { + if (this.readerType == ReaderType.COMET) { + return VectorizedSparkParquetReaders.buildCometReader( + requiredSchema, fileSchema, idToConstant, deleteFilter); + } else { + return VectorizedSparkParquetReaders.buildReader( + requiredSchema, fileSchema, idToConstant, deleteFilter); + } + }) .recordsPerBatch(batchSize) .filter(residual) .caseSensitive(caseSensitive()) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java index fd6783f3e1f7..5cc01ee3cf6d 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java @@ -115,11 +115,11 @@ private String[][] computePreferredLocations() { public PartitionReaderFactory createReaderFactory() { if (useParquetBatchReads()) { int batchSize = readConf.parquetBatchSize(); - return new SparkColumnarReaderFactory(batchSize); + return new SparkColumnarReaderFactory(batchSize, readConf.getReaderType()); } else if (useOrcBatchReads()) { int batchSize = readConf.orcBatchSize(); - return new SparkColumnarReaderFactory(batchSize); + return new SparkColumnarReaderFactory(batchSize, readConf.getReaderType()); } else { return new SparkRowReaderFactory(); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java index 655e20a50e11..599aece4452a 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java @@ -19,6 +19,7 @@ package org.apache.iceberg.spark.source; import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.ReaderType; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.connector.read.InputPartition; @@ -28,10 +29,12 @@ class SparkColumnarReaderFactory implements PartitionReaderFactory { private final int batchSize; + private final ReaderType readType; - SparkColumnarReaderFactory(int batchSize) { + SparkColumnarReaderFactory(int batchSize, ReaderType readType) { Preconditions.checkArgument(batchSize > 1, "Batch size must be > 1"); this.batchSize = batchSize; + this.readType = readType; } @Override @@ -49,7 +52,9 @@ public PartitionReader createColumnarReader(InputPartition inputP SparkInputPartition partition = (SparkInputPartition) inputPartition; if (partition.allTasksOfType(FileScanTask.class)) { - return new BatchDataReader(partition, batchSize); + BatchDataReader batchDataReader = new BatchDataReader(partition, batchSize); + batchDataReader.setReadType(readType); + return batchDataReader; } else { throw new UnsupportedOperationException( From 11bb4b7479cb6e5aa2548c41886fd368befe02db Mon Sep 17 00:00:00 2001 From: Huaxin Gao Date: Sun, 21 Apr 2024 10:29:33 -0700 Subject: [PATCH 03/32] address comments --- .../vectorized/BaseColumnBatchLoader.java | 24 +++++++++---------- .../comet/CometColumnarBatchReader.java | 18 ++++++++++++-- 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java index b33415ddfe8e..574a5e2ee4ea 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java @@ -56,11 +56,8 @@ protected BaseColumnBatchLoader( } } - public ColumnarBatch loadDataToColumnBatch() { - int numRowsUndeleted = initRowIdMapping(); - - ColumnVector[] arrowColumnVectors = readDataToColumnVectors(); - + protected ColumnarBatch initializeColumnBatchWithDeletions( + ColumnVector[] arrowColumnVectors, int numRowsUndeleted) { ColumnarBatch newColumnarBatch = new ColumnarBatch(arrowColumnVectors); newColumnarBatch.setNumRows(numRowsUndeleted); @@ -75,23 +72,26 @@ public ColumnarBatch loadDataToColumnBatch() { } newColumnarBatch.setNumRows(numRowsToRead); } - - if (hasIsDeletedColumn) { - readDeletedColumnIfNecessary(arrowColumnVectors); - } - return newColumnarBatch; } + /** + * This method iterates over each column reader and reads the current batch of data into the + * {@link ColumnVector}. + */ protected abstract ColumnVector[] readDataToColumnVectors(); - protected abstract void readDeletedColumnIfNecessary(ColumnVector[] arrowColumnVectors); + /** + * This method reads the current batch of data into the {@link ColumnVector}, and applies deletion + * logic, and loads data into a {@link ColumnarBatch}. + */ + public abstract ColumnarBatch loadDataToColumnBatch(); boolean hasEqDeletes() { return deletes != null && deletes.hasEqDeletes(); } - int initRowIdMapping() { + protected int initRowIdMapping() { Pair posDeleteRowIdMapping = posDelRowIdMapping(); if (posDeleteRowIdMapping != null) { rowIdMapping = posDeleteRowIdMapping.first(); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnarBatchReader.java index 2fd0c32f7a02..424f4fb53a40 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnarBatchReader.java @@ -119,6 +119,21 @@ private class ColumnBatchLoader extends BaseColumnBatchLoader { super(numRowsToRead, hasIsDeletedColumn, deletes, rowStartPosInBatch); } + @Override + public ColumnarBatch loadDataToColumnBatch() { + int numRowsUndeleted = initRowIdMapping(); + ColumnVector[] arrowColumnVectors = readDataToColumnVectors(); + + ColumnarBatch newColumnarBatch = + initializeColumnBatchWithDeletions(arrowColumnVectors, numRowsUndeleted); + + if (hasIsDeletedColumn) { + readDeletedColumnIfNecessary(arrowColumnVectors); + } + + return newColumnarBatch; + } + @Override protected ColumnVector[] readDataToColumnVectors() { ColumnVector[] columnVectors = new ColumnVector[readers.length]; @@ -135,8 +150,7 @@ protected ColumnVector[] readDataToColumnVectors() { return columnVectors; } - @Override - protected void readDeletedColumnIfNecessary(ColumnVector[] columnVectors) { + void readDeletedColumnIfNecessary(ColumnVector[] columnVectors) { for (int i = 0; i < readers.length; i++) { if (readers[i] instanceof CometDeleteColumnReader) { CometDeleteColumnReader deleteColumnReader = new CometDeleteColumnReader<>(isDeleted); From 4b0116186ec786b975b0fbe3b67d4cc9947a2b27 Mon Sep 17 00:00:00 2001 From: Huaxin Gao Date: Sun, 21 Apr 2024 10:37:17 -0700 Subject: [PATCH 04/32] remove unnecessary code --- .../iceberg/spark/data/vectorized/comet/CometColumnReader.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java index a134c985912b..a3c527e908ab 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java @@ -152,8 +152,6 @@ public void close() { @Override public void setBatchSize(int size) { - // Preconditions.checkState( - // !initialized, "'reset' shouldn't be called" + " before 'setBatchSize' is called"); this.batchSize = size; } From e73dc0ac1df5a03d258d555bc773ac258459308a Mon Sep 17 00:00:00 2001 From: Huaxin Gao Date: Thu, 25 Apr 2024 23:31:05 -0700 Subject: [PATCH 05/32] address comments --- .../apache/iceberg/spark/ParquetReaderType.java | 4 ++-- .../apache/iceberg/spark/SparkSQLProperties.java | 8 ++++---- .../data/vectorized/comet/CometColumnReader.java | 2 +- .../iceberg/spark/source/BaseBatchReader.java | 14 ++++++-------- .../iceberg/spark/source/BatchDataReader.java | 12 ++++++++---- .../apache/iceberg/spark/source/SparkBatch.java | 4 ++-- .../spark/source/SparkColumnarReaderFactory.java | 12 ++++++------ 7 files changed, 29 insertions(+), 27 deletions(-) rename api/src/main/java/org/apache/iceberg/ReaderType.java => spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java (92%) diff --git a/api/src/main/java/org/apache/iceberg/ReaderType.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java similarity index 92% rename from api/src/main/java/org/apache/iceberg/ReaderType.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java index 89ec0e0169b5..fac1604c0bd1 100644 --- a/api/src/main/java/org/apache/iceberg/ReaderType.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java @@ -16,9 +16,9 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.iceberg; +package org.apache.iceberg.spark; -public enum ReaderType { +public enum ParquetReaderType { ICEBERG, COMET } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java index 3f6faf984426..4a8909df8f5a 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java @@ -19,7 +19,6 @@ package org.apache.iceberg.spark; import java.time.Duration; -import org.apache.iceberg.ReaderType; public class SparkSQLProperties { @@ -28,9 +27,10 @@ private SparkSQLProperties() {} // Controls whether vectorized reads are enabled public static final String VECTORIZATION_ENABLED = "spark.sql.iceberg.vectorization.enabled"; - // Controls whether which reader to use for vectorization - public static final String READER_TYPE = "spark.sql.iceberg.parquet.reader-type"; - public static final ReaderType READER_TYPE_DEFAULT = ReaderType.ICEBERG; + // Controls which parquet reader to use for vectorization: either Iceberg's parquet reader or + // Comet's parquet reader + public static final String PARQUET_READER_TYPE = "spark.sql.iceberg.parquet.reader-type"; + public static final ParquetReaderType PARQUET_READER_TYPE_DEFAULT = ParquetReaderType.ICEBERG; // Controls whether reading/writing timestamps without timezones is allowed @Deprecated diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java index a3c527e908ab..de8f87eb9ca2 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java @@ -101,7 +101,7 @@ public void reset() { delegate.close(); } - delegate = Utils.getColumnReader(sparkType, descriptor, batchSize, false, false); + delegate = Utils.getColumnReader(sparkType, descriptor, batchSize, false, true); initialized = true; } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java index 7321b49d1c82..1ed24f42842d 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java @@ -22,7 +22,6 @@ import java.util.Set; import org.apache.iceberg.FileFormat; import org.apache.iceberg.MetadataColumns; -import org.apache.iceberg.ReaderType; import org.apache.iceberg.ScanTask; import org.apache.iceberg.ScanTaskGroup; import org.apache.iceberg.Schema; @@ -33,6 +32,7 @@ import org.apache.iceberg.orc.ORC; import org.apache.iceberg.parquet.Parquet; import org.apache.iceberg.relocated.com.google.common.collect.Sets; +import org.apache.iceberg.spark.ParquetReaderType; import org.apache.iceberg.spark.data.vectorized.VectorizedSparkOrcReaders; import org.apache.iceberg.spark.data.vectorized.VectorizedSparkParquetReaders; import org.apache.iceberg.types.TypeUtil; @@ -40,7 +40,7 @@ abstract class BaseBatchReader extends BaseReader { private final int batchSize; - private ReaderType readerType; + private final ParquetReaderType parquetReaderType; BaseBatchReader( Table table, @@ -48,13 +48,11 @@ abstract class BaseBatchReader extends BaseReader newBatchIterable( @@ -93,7 +91,7 @@ private CloseableIterable newParquetIterable( .split(start, length) .createBatchedReaderFunc( fileSchema -> { - if (this.readerType == ReaderType.COMET) { + if (parquetReaderType == ParquetReaderType.COMET) { return VectorizedSparkParquetReaders.buildCometReader( requiredSchema, fileSchema, idToConstant, deleteFilter); } else { diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java index f45c152203ee..7529b22eca5b 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java @@ -28,6 +28,7 @@ import org.apache.iceberg.io.CloseableIterator; import org.apache.iceberg.io.InputFile; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.spark.ParquetReaderType; import org.apache.iceberg.spark.source.metrics.TaskNumDeletes; import org.apache.iceberg.spark.source.metrics.TaskNumSplits; import org.apache.iceberg.util.SnapshotUtil; @@ -45,14 +46,16 @@ class BatchDataReader extends BaseBatchReader private final long numSplits; - BatchDataReader(SparkInputPartition partition, int batchSize) { + BatchDataReader( + SparkInputPartition partition, int batchSize, ParquetReaderType parquetReaderType) { this( partition.table(), partition.taskGroup(), SnapshotUtil.schemaFor(partition.table(), partition.branch()), partition.expectedSchema(), partition.isCaseSensitive(), - batchSize); + batchSize, + parquetReaderType); } BatchDataReader( @@ -61,8 +64,9 @@ class BatchDataReader extends BaseBatchReader Schema tableSchema, Schema expectedSchema, boolean caseSensitive, - int size) { - super(table, taskGroup, tableSchema, expectedSchema, caseSensitive, size); + int size, + ParquetReaderType parquetReaderType) { + super(table, taskGroup, tableSchema, expectedSchema, caseSensitive, size, parquetReaderType); numSplits = taskGroup.tasks().size(); LOG.debug("Reading {} file split(s) for table {}", numSplits, table.name()); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java index 5cc01ee3cf6d..3b17ca3514a4 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java @@ -115,11 +115,11 @@ private String[][] computePreferredLocations() { public PartitionReaderFactory createReaderFactory() { if (useParquetBatchReads()) { int batchSize = readConf.parquetBatchSize(); - return new SparkColumnarReaderFactory(batchSize, readConf.getReaderType()); + return new SparkColumnarReaderFactory(batchSize, readConf.parquetReaderType()); } else if (useOrcBatchReads()) { int batchSize = readConf.orcBatchSize(); - return new SparkColumnarReaderFactory(batchSize, readConf.getReaderType()); + return new SparkColumnarReaderFactory(batchSize, readConf.parquetReaderType()); } else { return new SparkRowReaderFactory(); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java index 599aece4452a..21fab7cb4205 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java @@ -19,8 +19,8 @@ package org.apache.iceberg.spark.source; import org.apache.iceberg.FileScanTask; -import org.apache.iceberg.ReaderType; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.spark.ParquetReaderType; import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.connector.read.InputPartition; import org.apache.spark.sql.connector.read.PartitionReader; @@ -29,12 +29,12 @@ class SparkColumnarReaderFactory implements PartitionReaderFactory { private final int batchSize; - private final ReaderType readType; + private final ParquetReaderType parquetReaderType; - SparkColumnarReaderFactory(int batchSize, ReaderType readType) { + SparkColumnarReaderFactory(int batchSize, ParquetReaderType readType) { Preconditions.checkArgument(batchSize > 1, "Batch size must be > 1"); this.batchSize = batchSize; - this.readType = readType; + this.parquetReaderType = readType; } @Override @@ -52,8 +52,8 @@ public PartitionReader createColumnarReader(InputPartition inputP SparkInputPartition partition = (SparkInputPartition) inputPartition; if (partition.allTasksOfType(FileScanTask.class)) { - BatchDataReader batchDataReader = new BatchDataReader(partition, batchSize); - batchDataReader.setReadType(readType); + BatchDataReader batchDataReader = + new BatchDataReader(partition, batchSize, parquetReaderType); return batchDataReader; } else { From d055ca232b9b783e9e1abbcef0d2673dbd95c8d8 Mon Sep 17 00:00:00 2001 From: Huaxin Gao Date: Tue, 30 Apr 2024 18:06:12 -0700 Subject: [PATCH 06/32] address comments --- .../apache/iceberg/spark/BatchReadConf.java | 47 +++++++++++++++++++ .../iceberg/spark/SparkSQLProperties.java | 3 +- .../{comet => }/CometColumnReader.java | 5 +- .../{comet => }/CometColumnarBatchReader.java | 5 +- .../CometConstantColumnReader.java | 2 +- .../{comet => }/CometDeleteColumnReader.java | 2 +- .../CometPositionColumnReader.java | 2 +- .../vectorized/{comet => }/CometVector.java | 2 +- .../CometVectorizedReaderBuilder.java | 4 +- .../iceberg/spark/source/BaseBatchReader.java | 16 +++---- .../iceberg/spark/source/BatchDataReader.java | 13 ++--- .../iceberg/spark/source/SparkBatch.java | 7 +-- .../source/SparkColumnarReaderFactory.java | 16 ++----- 13 files changed, 78 insertions(+), 46 deletions(-) create mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/BatchReadConf.java rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/{comet => }/CometColumnReader.java (97%) rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/{comet => }/CometColumnarBatchReader.java (96%) rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/{comet => }/CometConstantColumnReader.java (96%) rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/{comet => }/CometDeleteColumnReader.java (97%) rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/{comet => }/CometPositionColumnReader.java (97%) rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/{comet => }/CometVector.java (98%) rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/{comet => }/CometVectorizedReaderBuilder.java (97%) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/BatchReadConf.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/BatchReadConf.java new file mode 100644 index 000000000000..f014e3154653 --- /dev/null +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/BatchReadConf.java @@ -0,0 +1,47 @@ +/* + * 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; + +import java.io.Serializable; + +public class BatchReadConf implements Serializable { + + private int parquetBatchSize; + private int orcBatchSize; + private ParquetReaderType parquetReaderType; + + public BatchReadConf( + int parquetBatchSize, int orcBatchSize, ParquetReaderType parquetReaderType) { + this.parquetBatchSize = parquetBatchSize; + this.orcBatchSize = orcBatchSize; + this.parquetReaderType = parquetReaderType; + } + + public int parquetBatchSize() { + return parquetBatchSize; + } + + public int orcBatchSize() { + return orcBatchSize; + } + + public ParquetReaderType parquetReaderType() { + return parquetReaderType; + } +} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java index 4a8909df8f5a..7b172a25a888 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java @@ -27,8 +27,7 @@ private SparkSQLProperties() {} // Controls whether vectorized reads are enabled public static final String VECTORIZATION_ENABLED = "spark.sql.iceberg.vectorization.enabled"; - // Controls which parquet reader to use for vectorization: either Iceberg's parquet reader or - // Comet's parquet reader + // Controls which parquet reader to use for vectorization public static final String PARQUET_READER_TYPE = "spark.sql.iceberg.parquet.reader-type"; public static final ParquetReaderType PARQUET_READER_TYPE_DEFAULT = ParquetReaderType.ICEBERG; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java similarity index 97% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java index de8f87eb9ca2..72b4073882ce 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.iceberg.spark.data.vectorized.comet; +package org.apache.iceberg.spark.data.vectorized; import java.io.IOException; import java.util.Map; @@ -32,6 +32,7 @@ import org.apache.parquet.column.page.PageReader; import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.spark.sql.internal.SQLConf; import org.apache.spark.sql.types.DataType; import org.apache.spark.sql.types.Metadata; import org.apache.spark.sql.types.StructField; @@ -101,7 +102,7 @@ public void reset() { delegate.close(); } - delegate = Utils.getColumnReader(sparkType, descriptor, batchSize, false, true); + delegate = Utils.getColumnReader(sparkType, descriptor, batchSize, SQLConf.get()); initialized = true; } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java similarity index 96% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnarBatchReader.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index 424f4fb53a40..8269bb3fa60c 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.iceberg.spark.data.vectorized.comet; +package org.apache.iceberg.spark.data.vectorized; import java.io.IOException; import java.io.UncheckedIOException; @@ -28,7 +28,6 @@ import org.apache.iceberg.data.DeleteFilter; import org.apache.iceberg.parquet.VectorizedReader; import org.apache.iceberg.spark.SparkSchemaUtil; -import org.apache.iceberg.spark.data.vectorized.BaseColumnBatchLoader; import org.apache.parquet.column.page.PageReadStore; import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.hadoop.metadata.ColumnPath; @@ -42,7 +41,7 @@ * populated via delegated read calls to {@linkplain CometColumnReader VectorReader(s)}. */ @SuppressWarnings("checkstyle:VisibilityModifier") -public class CometColumnarBatchReader implements VectorizedReader { +class CometColumnarBatchReader implements VectorizedReader { private final CometColumnReader[] readers; private final boolean hasIsDeletedColumn; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometConstantColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java similarity index 96% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometConstantColumnReader.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java index 11eef7ff5e0c..b2f0c057ae6d 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometConstantColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.iceberg.spark.data.vectorized.comet; +package org.apache.iceberg.spark.data.vectorized; import org.apache.comet.parquet.ConstantColumnReader; import org.apache.iceberg.types.Types; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometDeleteColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java similarity index 97% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometDeleteColumnReader.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java index ae948f9eca7b..7b7a395c5a93 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometDeleteColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.iceberg.spark.data.vectorized.comet; +package org.apache.iceberg.spark.data.vectorized; import org.apache.comet.parquet.ConstantColumnReader; import org.apache.comet.parquet.MetadataColumnReader; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometPositionColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java similarity index 97% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometPositionColumnReader.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java index 7a3cd7cb5789..3c92877864f7 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometPositionColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.iceberg.spark.data.vectorized.comet; +package org.apache.iceberg.spark.data.vectorized; import org.apache.comet.parquet.MetadataColumnReader; import org.apache.comet.parquet.Native; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVector.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java similarity index 98% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVector.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java index 2e8b8210b8e5..340e0cd13696 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVector.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.iceberg.spark.data.vectorized.comet; +package org.apache.iceberg.spark.data.vectorized; import org.apache.comet.vector.CometDelegateVector; import org.apache.spark.sql.types.DataType; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVectorizedReaderBuilder.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java similarity index 97% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVectorizedReaderBuilder.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java index f679ea7bfa86..2d3c648ab3c2 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/comet/CometVectorizedReaderBuilder.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.iceberg.spark.data.vectorized.comet; +package org.apache.iceberg.spark.data.vectorized; import java.util.List; import java.util.Map; @@ -39,7 +39,7 @@ import org.apache.parquet.schema.Type; import org.apache.spark.sql.catalyst.InternalRow; -public class CometVectorizedReaderBuilder extends TypeWithSchemaVisitor> { +class CometVectorizedReaderBuilder extends TypeWithSchemaVisitor> { private final MessageType parquetSchema; private final Schema icebergSchema; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java index 1ed24f42842d..539c6408b34c 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java @@ -32,6 +32,7 @@ import org.apache.iceberg.orc.ORC; import org.apache.iceberg.parquet.Parquet; import org.apache.iceberg.relocated.com.google.common.collect.Sets; +import org.apache.iceberg.spark.BatchReadConf; import org.apache.iceberg.spark.ParquetReaderType; import org.apache.iceberg.spark.data.vectorized.VectorizedSparkOrcReaders; import org.apache.iceberg.spark.data.vectorized.VectorizedSparkParquetReaders; @@ -39,8 +40,7 @@ import org.apache.spark.sql.vectorized.ColumnarBatch; abstract class BaseBatchReader extends BaseReader { - private final int batchSize; - private final ParquetReaderType parquetReaderType; + private final BatchReadConf batchReadConf; BaseBatchReader( Table table, @@ -48,11 +48,9 @@ abstract class BaseBatchReader extends BaseReader newBatchIterable( @@ -91,7 +89,7 @@ private CloseableIterable newParquetIterable( .split(start, length) .createBatchedReaderFunc( fileSchema -> { - if (parquetReaderType == ParquetReaderType.COMET) { + if (batchReadConf.parquetReaderType() == ParquetReaderType.COMET) { return VectorizedSparkParquetReaders.buildCometReader( requiredSchema, fileSchema, idToConstant, deleteFilter); } else { @@ -99,7 +97,7 @@ private CloseableIterable newParquetIterable( requiredSchema, fileSchema, idToConstant, deleteFilter); } }) - .recordsPerBatch(batchSize) + .recordsPerBatch(batchReadConf.parquetBatchSize()) .filter(residual) .caseSensitive(caseSensitive()) // Spark eagerly consumes the batches. So the underlying memory allocated could be reused @@ -129,7 +127,7 @@ private CloseableIterable newOrcIterable( .createBatchedReaderFunc( fileSchema -> VectorizedSparkOrcReaders.buildReader(expectedSchema(), fileSchema, idToConstant)) - .recordsPerBatch(batchSize) + .recordsPerBatch(batchReadConf.orcBatchSize()) .filter(residual) .caseSensitive(caseSensitive()) .withNameMapping(nameMapping()) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java index 7529b22eca5b..31a95cf0e45d 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java @@ -28,7 +28,7 @@ import org.apache.iceberg.io.CloseableIterator; import org.apache.iceberg.io.InputFile; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; -import org.apache.iceberg.spark.ParquetReaderType; +import org.apache.iceberg.spark.BatchReadConf; import org.apache.iceberg.spark.source.metrics.TaskNumDeletes; import org.apache.iceberg.spark.source.metrics.TaskNumSplits; import org.apache.iceberg.util.SnapshotUtil; @@ -46,16 +46,14 @@ class BatchDataReader extends BaseBatchReader private final long numSplits; - BatchDataReader( - SparkInputPartition partition, int batchSize, ParquetReaderType parquetReaderType) { + BatchDataReader(SparkInputPartition partition, BatchReadConf batchReadConf) { this( partition.table(), partition.taskGroup(), SnapshotUtil.schemaFor(partition.table(), partition.branch()), partition.expectedSchema(), partition.isCaseSensitive(), - batchSize, - parquetReaderType); + batchReadConf); } BatchDataReader( @@ -64,9 +62,8 @@ class BatchDataReader extends BaseBatchReader Schema tableSchema, Schema expectedSchema, boolean caseSensitive, - int size, - ParquetReaderType parquetReaderType) { - super(table, taskGroup, tableSchema, expectedSchema, caseSensitive, size, parquetReaderType); + BatchReadConf batchReadConf) { + super(table, taskGroup, tableSchema, expectedSchema, caseSensitive, batchReadConf); numSplits = taskGroup.tasks().size(); LOG.debug("Reading {} file split(s) for table {}", numSplits, table.name()); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java index 3b17ca3514a4..790d5e83e070 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java @@ -114,12 +114,9 @@ private String[][] computePreferredLocations() { @Override public PartitionReaderFactory createReaderFactory() { if (useParquetBatchReads()) { - int batchSize = readConf.parquetBatchSize(); - return new SparkColumnarReaderFactory(batchSize, readConf.parquetReaderType()); - + return new SparkColumnarReaderFactory(readConf.batchReadConf()); } else if (useOrcBatchReads()) { - int batchSize = readConf.orcBatchSize(); - return new SparkColumnarReaderFactory(batchSize, readConf.parquetReaderType()); + return new SparkColumnarReaderFactory(readConf.batchReadConf()); } else { return new SparkRowReaderFactory(); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java index 21fab7cb4205..138ed1569263 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java @@ -20,7 +20,7 @@ import org.apache.iceberg.FileScanTask; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; -import org.apache.iceberg.spark.ParquetReaderType; +import org.apache.iceberg.spark.BatchReadConf; import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.connector.read.InputPartition; import org.apache.spark.sql.connector.read.PartitionReader; @@ -28,13 +28,10 @@ import org.apache.spark.sql.vectorized.ColumnarBatch; class SparkColumnarReaderFactory implements PartitionReaderFactory { - private final int batchSize; - private final ParquetReaderType parquetReaderType; + private final BatchReadConf batchReadConf; - SparkColumnarReaderFactory(int batchSize, ParquetReaderType readType) { - Preconditions.checkArgument(batchSize > 1, "Batch size must be > 1"); - this.batchSize = batchSize; - this.parquetReaderType = readType; + SparkColumnarReaderFactory(BatchReadConf batchReadConf) { + this.batchReadConf = batchReadConf; } @Override @@ -52,10 +49,7 @@ public PartitionReader createColumnarReader(InputPartition inputP SparkInputPartition partition = (SparkInputPartition) inputPartition; if (partition.allTasksOfType(FileScanTask.class)) { - BatchDataReader batchDataReader = - new BatchDataReader(partition, batchSize, parquetReaderType); - return batchDataReader; - + return new BatchDataReader(partition, batchReadConf); } else { throw new UnsupportedOperationException( "Unsupported task group for columnar reads: " + partition.taskGroup()); From 4f06c979ff19b64626a1397a64ac78f19218969a Mon Sep 17 00:00:00 2001 From: Huaxin Gao Date: Tue, 30 Apr 2024 23:00:56 -0700 Subject: [PATCH 07/32] remove unnecessary public --- .../iceberg/spark/data/vectorized/CometColumnarBatchReader.java | 2 +- .../spark/data/vectorized/CometVectorizedReaderBuilder.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index 8269bb3fa60c..b2ed77abe15c 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -49,7 +49,7 @@ class CometColumnarBatchReader implements VectorizedReader { private long rowStartPosInBatch = 0; private final BatchReader delegate; - public CometColumnarBatchReader(List> readers, Schema schema) { + CometColumnarBatchReader(List> readers, Schema schema) { this.readers = readers.stream().map(CometColumnReader.class::cast).toArray(CometColumnReader[]::new); this.hasIsDeletedColumn = diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java index 2d3c648ab3c2..2cc24f6ce98d 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java @@ -47,7 +47,7 @@ class CometVectorizedReaderBuilder extends TypeWithSchemaVisitor>, VectorizedReader> readerFactory; private final DeleteFilter deleteFilter; - public CometVectorizedReaderBuilder( + CometVectorizedReaderBuilder( Schema expectedSchema, MessageType parquetSchema, Map idToConstant, From d96183d87d9bf39c51097002af77a627eea9bf93 Mon Sep 17 00:00:00 2001 From: Huaxin Gao Date: Fri, 3 May 2024 16:07:04 -0700 Subject: [PATCH 08/32] address comments --- .../java/org/apache/iceberg/spark/BatchReadConf.java | 10 +++++----- .../org/apache/iceberg/spark/SparkSQLProperties.java | 2 +- .../spark/data/vectorized/CometColumnReader.java | 3 +-- .../apache/iceberg/spark/source/BaseBatchReader.java | 12 ++++++------ .../apache/iceberg/spark/source/BatchDataReader.java | 4 ++-- .../org/apache/iceberg/spark/source/SparkBatch.java | 1 + .../spark/source/SparkColumnarReaderFactory.java | 8 ++++---- 7 files changed, 20 insertions(+), 20 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/BatchReadConf.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/BatchReadConf.java index f014e3154653..2a1393b0cc0c 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/BatchReadConf.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/BatchReadConf.java @@ -22,15 +22,15 @@ public class BatchReadConf implements Serializable { - private int parquetBatchSize; - private int orcBatchSize; - private ParquetReaderType parquetReaderType; + private final int parquetBatchSize; + private final ParquetReaderType parquetReaderType; + private final int orcBatchSize; public BatchReadConf( - int parquetBatchSize, int orcBatchSize, ParquetReaderType parquetReaderType) { + int parquetBatchSize, ParquetReaderType parquetReaderType, int orcBatchSize) { this.parquetBatchSize = parquetBatchSize; - this.orcBatchSize = orcBatchSize; this.parquetReaderType = parquetReaderType; + this.orcBatchSize = orcBatchSize; } public int parquetBatchSize() { diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java index 7b172a25a888..99033c076a43 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java @@ -27,7 +27,7 @@ private SparkSQLProperties() {} // Controls whether vectorized reads are enabled public static final String VECTORIZATION_ENABLED = "spark.sql.iceberg.vectorization.enabled"; - // Controls which parquet reader to use for vectorization + // Controls which Parquet reader to use for vectorization public static final String PARQUET_READER_TYPE = "spark.sql.iceberg.parquet.reader-type"; public static final ParquetReaderType PARQUET_READER_TYPE_DEFAULT = ParquetReaderType.ICEBERG; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java index 72b4073882ce..420700900071 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java @@ -32,7 +32,6 @@ import org.apache.parquet.column.page.PageReader; import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.hadoop.metadata.ColumnPath; -import org.apache.spark.sql.internal.SQLConf; import org.apache.spark.sql.types.DataType; import org.apache.spark.sql.types.Metadata; import org.apache.spark.sql.types.StructField; @@ -102,7 +101,7 @@ public void reset() { delegate.close(); } - delegate = Utils.getColumnReader(sparkType, descriptor, batchSize, SQLConf.get()); + delegate = Utils.getColumnReader(sparkType, descriptor, batchSize, false, false); initialized = true; } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java index 539c6408b34c..a72f2cc88e46 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java @@ -40,7 +40,7 @@ import org.apache.spark.sql.vectorized.ColumnarBatch; abstract class BaseBatchReader extends BaseReader { - private final BatchReadConf batchReadConf; + private final BatchReadConf conf; BaseBatchReader( Table table, @@ -48,9 +48,9 @@ abstract class BaseBatchReader extends BaseReader newBatchIterable( @@ -89,7 +89,7 @@ private CloseableIterable newParquetIterable( .split(start, length) .createBatchedReaderFunc( fileSchema -> { - if (batchReadConf.parquetReaderType() == ParquetReaderType.COMET) { + if (conf.parquetReaderType() == ParquetReaderType.COMET) { return VectorizedSparkParquetReaders.buildCometReader( requiredSchema, fileSchema, idToConstant, deleteFilter); } else { @@ -97,7 +97,7 @@ private CloseableIterable newParquetIterable( requiredSchema, fileSchema, idToConstant, deleteFilter); } }) - .recordsPerBatch(batchReadConf.parquetBatchSize()) + .recordsPerBatch(conf.parquetBatchSize()) .filter(residual) .caseSensitive(caseSensitive()) // Spark eagerly consumes the batches. So the underlying memory allocated could be reused @@ -127,7 +127,7 @@ private CloseableIterable newOrcIterable( .createBatchedReaderFunc( fileSchema -> VectorizedSparkOrcReaders.buildReader(expectedSchema(), fileSchema, idToConstant)) - .recordsPerBatch(batchReadConf.orcBatchSize()) + .recordsPerBatch(conf.orcBatchSize()) .filter(residual) .caseSensitive(caseSensitive()) .withNameMapping(nameMapping()) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java index 31a95cf0e45d..19a672fa4171 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java @@ -46,14 +46,14 @@ class BatchDataReader extends BaseBatchReader private final long numSplits; - BatchDataReader(SparkInputPartition partition, BatchReadConf batchReadConf) { + BatchDataReader(SparkInputPartition partition, BatchReadConf conf) { this( partition.table(), partition.taskGroup(), SnapshotUtil.schemaFor(partition.table(), partition.branch()), partition.expectedSchema(), partition.isCaseSensitive(), - batchReadConf); + conf); } BatchDataReader( diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java index 790d5e83e070..e3b456fb7165 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java @@ -115,6 +115,7 @@ private String[][] computePreferredLocations() { public PartitionReaderFactory createReaderFactory() { if (useParquetBatchReads()) { return new SparkColumnarReaderFactory(readConf.batchReadConf()); + } else if (useOrcBatchReads()) { return new SparkColumnarReaderFactory(readConf.batchReadConf()); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java index 138ed1569263..a5d1699a4886 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java @@ -28,10 +28,10 @@ import org.apache.spark.sql.vectorized.ColumnarBatch; class SparkColumnarReaderFactory implements PartitionReaderFactory { - private final BatchReadConf batchReadConf; + private final BatchReadConf conf; - SparkColumnarReaderFactory(BatchReadConf batchReadConf) { - this.batchReadConf = batchReadConf; + SparkColumnarReaderFactory(BatchReadConf conf) { + this.conf = conf; } @Override @@ -49,7 +49,7 @@ public PartitionReader createColumnarReader(InputPartition inputP SparkInputPartition partition = (SparkInputPartition) inputPartition; if (partition.allTasksOfType(FileScanTask.class)) { - return new BatchDataReader(partition, batchReadConf); + return new BatchDataReader(partition, conf); } else { throw new UnsupportedOperationException( "Unsupported task group for columnar reads: " + partition.taskGroup()); From a2a37073b181ee4963205c216cba86322caf8f64 Mon Sep 17 00:00:00 2001 From: Huaxin Gao Date: Wed, 15 May 2024 23:35:05 -0700 Subject: [PATCH 09/32] address comments --- spark/v3.4/build.gradle | 2 + .../iceberg/spark/OrcBatchReadConf.java | 27 +++++++++++++ ...eadConf.java => ParquetBatchReadConf.java} | 28 +++---------- .../iceberg/spark/source/BaseBatchReader.java | 18 +++++---- .../iceberg/spark/source/BatchDataReader.java | 16 +++++--- .../iceberg/spark/source/SparkBatch.java | 39 +++++++++++++++++-- .../source/SparkColumnarReaderFactory.java | 16 +++++--- 7 files changed, 103 insertions(+), 43 deletions(-) create mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/OrcBatchReadConf.java rename spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/{BatchReadConf.java => ParquetBatchReadConf.java} (57%) diff --git a/spark/v3.4/build.gradle b/spark/v3.4/build.gradle index 32739685581b..6fdd6e3f7f4a 100644 --- a/spark/v3.4/build.gradle +++ b/spark/v3.4/build.gradle @@ -52,6 +52,8 @@ project(":iceberg-spark:iceberg-spark-${sparkMajorVersion}_${scalaVersion}") { dependencies { implementation project(path: ':iceberg-bundled-guava', configuration: 'shadow') api project(':iceberg-api') + annotationProcessor libs.immutables.value + compileOnly libs.immutables.value implementation project(':iceberg-common') implementation project(':iceberg-core') implementation project(':iceberg-data') diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/OrcBatchReadConf.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/OrcBatchReadConf.java new file mode 100644 index 000000000000..d3b339d60e3f --- /dev/null +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/OrcBatchReadConf.java @@ -0,0 +1,27 @@ +/* + * 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; + +import java.io.Serializable; +import org.immutables.value.Value; + +@Value.Immutable +public interface OrcBatchReadConf extends Serializable { + int batchSize(); +} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/BatchReadConf.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetBatchReadConf.java similarity index 57% rename from spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/BatchReadConf.java rename to spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetBatchReadConf.java index 2a1393b0cc0c..442d728d4d69 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/BatchReadConf.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetBatchReadConf.java @@ -19,29 +19,11 @@ package org.apache.iceberg.spark; import java.io.Serializable; +import org.immutables.value.Value; -public class BatchReadConf implements Serializable { +@Value.Immutable +public interface ParquetBatchReadConf extends Serializable { + int batchSize(); - private final int parquetBatchSize; - private final ParquetReaderType parquetReaderType; - private final int orcBatchSize; - - public BatchReadConf( - int parquetBatchSize, ParquetReaderType parquetReaderType, int orcBatchSize) { - this.parquetBatchSize = parquetBatchSize; - this.parquetReaderType = parquetReaderType; - this.orcBatchSize = orcBatchSize; - } - - public int parquetBatchSize() { - return parquetBatchSize; - } - - public int orcBatchSize() { - return orcBatchSize; - } - - public ParquetReaderType parquetReaderType() { - return parquetReaderType; - } + ParquetReaderType readerType(); } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java index a72f2cc88e46..780e1750a52e 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java @@ -32,7 +32,8 @@ import org.apache.iceberg.orc.ORC; import org.apache.iceberg.parquet.Parquet; import org.apache.iceberg.relocated.com.google.common.collect.Sets; -import org.apache.iceberg.spark.BatchReadConf; +import org.apache.iceberg.spark.OrcBatchReadConf; +import org.apache.iceberg.spark.ParquetBatchReadConf; import org.apache.iceberg.spark.ParquetReaderType; import org.apache.iceberg.spark.data.vectorized.VectorizedSparkOrcReaders; import org.apache.iceberg.spark.data.vectorized.VectorizedSparkParquetReaders; @@ -40,7 +41,8 @@ import org.apache.spark.sql.vectorized.ColumnarBatch; abstract class BaseBatchReader extends BaseReader { - private final BatchReadConf conf; + private final ParquetBatchReadConf parquetConf; + private final OrcBatchReadConf orcConf; BaseBatchReader( Table table, @@ -48,9 +50,11 @@ abstract class BaseBatchReader extends BaseReader newBatchIterable( @@ -89,7 +93,7 @@ private CloseableIterable newParquetIterable( .split(start, length) .createBatchedReaderFunc( fileSchema -> { - if (conf.parquetReaderType() == ParquetReaderType.COMET) { + if (parquetConf.readerType() == ParquetReaderType.COMET) { return VectorizedSparkParquetReaders.buildCometReader( requiredSchema, fileSchema, idToConstant, deleteFilter); } else { @@ -97,7 +101,7 @@ private CloseableIterable newParquetIterable( requiredSchema, fileSchema, idToConstant, deleteFilter); } }) - .recordsPerBatch(conf.parquetBatchSize()) + .recordsPerBatch(parquetConf.batchSize()) .filter(residual) .caseSensitive(caseSensitive()) // Spark eagerly consumes the batches. So the underlying memory allocated could be reused @@ -127,7 +131,7 @@ private CloseableIterable newOrcIterable( .createBatchedReaderFunc( fileSchema -> VectorizedSparkOrcReaders.buildReader(expectedSchema(), fileSchema, idToConstant)) - .recordsPerBatch(conf.orcBatchSize()) + .recordsPerBatch(orcConf.batchSize()) .filter(residual) .caseSensitive(caseSensitive()) .withNameMapping(nameMapping()) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java index 19a672fa4171..e96cbaf411ee 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java @@ -28,7 +28,8 @@ import org.apache.iceberg.io.CloseableIterator; import org.apache.iceberg.io.InputFile; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; -import org.apache.iceberg.spark.BatchReadConf; +import org.apache.iceberg.spark.OrcBatchReadConf; +import org.apache.iceberg.spark.ParquetBatchReadConf; import org.apache.iceberg.spark.source.metrics.TaskNumDeletes; import org.apache.iceberg.spark.source.metrics.TaskNumSplits; import org.apache.iceberg.util.SnapshotUtil; @@ -46,14 +47,18 @@ class BatchDataReader extends BaseBatchReader private final long numSplits; - BatchDataReader(SparkInputPartition partition, BatchReadConf conf) { + BatchDataReader( + SparkInputPartition partition, + ParquetBatchReadConf parquetBatchReadConf, + OrcBatchReadConf orcBatchReadConf) { this( partition.table(), partition.taskGroup(), SnapshotUtil.schemaFor(partition.table(), partition.branch()), partition.expectedSchema(), partition.isCaseSensitive(), - conf); + parquetBatchReadConf, + orcBatchReadConf); } BatchDataReader( @@ -62,8 +67,9 @@ class BatchDataReader extends BaseBatchReader Schema tableSchema, Schema expectedSchema, boolean caseSensitive, - BatchReadConf batchReadConf) { - super(table, taskGroup, tableSchema, expectedSchema, caseSensitive, batchReadConf); + ParquetBatchReadConf parquetConf, + OrcBatchReadConf orcConf) { + super(table, taskGroup, tableSchema, expectedSchema, caseSensitive, parquetConf, orcConf); numSplits = taskGroup.tasks().size(); LOG.debug("Reading {} file split(s) for table {}", numSplits, table.name()); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java index e3b456fb7165..3f8bba5d9825 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java @@ -28,8 +28,14 @@ import org.apache.iceberg.Schema; import org.apache.iceberg.SchemaParser; import org.apache.iceberg.Table; +import org.apache.iceberg.spark.ImmutableOrcBatchReadConf; +import org.apache.iceberg.spark.ImmutableParquetBatchReadConf; +import org.apache.iceberg.spark.OrcBatchReadConf; +import org.apache.iceberg.spark.ParquetBatchReadConf; +import org.apache.iceberg.spark.ParquetReaderType; import org.apache.iceberg.spark.SparkReadConf; import org.apache.iceberg.spark.SparkUtil; +import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types; import org.apache.spark.api.java.JavaSparkContext; import org.apache.spark.broadcast.Broadcast; @@ -113,17 +119,33 @@ private String[][] computePreferredLocations() { @Override public PartitionReaderFactory createReaderFactory() { - if (useParquetBatchReads()) { - return new SparkColumnarReaderFactory(readConf.batchReadConf()); + if (useCometBatchReads()) { + return new SparkColumnarReaderFactory( + parquetBatchReadConf(readConf.parquetBatchSize(), ParquetReaderType.COMET)); + + } else if (useParquetBatchReads()) { + return new SparkColumnarReaderFactory( + parquetBatchReadConf(readConf.parquetBatchSize(), ParquetReaderType.ICEBERG)); } else if (useOrcBatchReads()) { - return new SparkColumnarReaderFactory(readConf.batchReadConf()); + return new SparkColumnarReaderFactory(orcBatchReadConf(readConf.orcBatchSize())); } else { return new SparkRowReaderFactory(); } } + private ParquetBatchReadConf parquetBatchReadConf(int batchSize, ParquetReaderType readerType) { + return ImmutableParquetBatchReadConf.builder() + .batchSize(batchSize) + .readerType(readerType) + .build(); + } + + private OrcBatchReadConf orcBatchReadConf(int batchSize) { + return ImmutableOrcBatchReadConf.builder().batchSize(batchSize).build(); + } + // conditions for using Parquet batch reads: // - Parquet vectorization is enabled // - only primitives or metadata columns are projected @@ -152,6 +174,17 @@ private boolean supportsParquetBatchReads(Types.NestedField field) { return field.type().isPrimitiveType() || MetadataColumns.isMetadataColumn(field.fieldId()); } + private boolean useCometBatchReads() { + return readConf.parquetVectorizationEnabled() + && readConf.parquetReaderType() == ParquetReaderType.COMET + && expectedSchema.columns().stream().allMatch(this::supportsCometBatchReads) + && taskGroups.stream().allMatch(this::supportsParquetBatchReads); + } + + private boolean supportsCometBatchReads(Types.NestedField field) { + return field.type().isPrimitiveType() && !field.type().typeId().equals(Type.TypeID.UUID); + } + // conditions for using ORC batch reads: // - ORC vectorization is enabled // - all tasks are of type FileScanTask and read only ORC files with no delete files diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java index a5d1699a4886..876c0be533e0 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java @@ -20,7 +20,8 @@ import org.apache.iceberg.FileScanTask; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; -import org.apache.iceberg.spark.BatchReadConf; +import org.apache.iceberg.spark.OrcBatchReadConf; +import org.apache.iceberg.spark.ParquetBatchReadConf; import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.connector.read.InputPartition; import org.apache.spark.sql.connector.read.PartitionReader; @@ -28,10 +29,15 @@ import org.apache.spark.sql.vectorized.ColumnarBatch; class SparkColumnarReaderFactory implements PartitionReaderFactory { - private final BatchReadConf conf; + private ParquetBatchReadConf parquetConf; + private OrcBatchReadConf orcConf; - SparkColumnarReaderFactory(BatchReadConf conf) { - this.conf = conf; + SparkColumnarReaderFactory(ParquetBatchReadConf conf) { + this.parquetConf = conf; + } + + SparkColumnarReaderFactory(OrcBatchReadConf conf) { + this.orcConf = conf; } @Override @@ -49,7 +55,7 @@ public PartitionReader createColumnarReader(InputPartition inputP SparkInputPartition partition = (SparkInputPartition) inputPartition; if (partition.allTasksOfType(FileScanTask.class)) { - return new BatchDataReader(partition, conf); + return new BatchDataReader(partition, parquetConf, orcConf); } else { throw new UnsupportedOperationException( "Unsupported task group for columnar reads: " + partition.taskGroup()); From 193a85b9784194750df1a7dd63ca4f40cb58b97b Mon Sep 17 00:00:00 2001 From: Huaxin Gao Date: Fri, 17 May 2024 10:58:18 -0700 Subject: [PATCH 10/32] minor changes --- .../apache/iceberg/spark/source/SparkBatch.java | 16 +++++++--------- .../spark/source/SparkColumnarReaderFactory.java | 6 ++++-- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java index 3f8bba5d9825..11f054b11710 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java @@ -120,30 +120,28 @@ private String[][] computePreferredLocations() { @Override public PartitionReaderFactory createReaderFactory() { if (useCometBatchReads()) { - return new SparkColumnarReaderFactory( - parquetBatchReadConf(readConf.parquetBatchSize(), ParquetReaderType.COMET)); + return new SparkColumnarReaderFactory(parquetBatchReadConf(ParquetReaderType.COMET)); } else if (useParquetBatchReads()) { - return new SparkColumnarReaderFactory( - parquetBatchReadConf(readConf.parquetBatchSize(), ParquetReaderType.ICEBERG)); + return new SparkColumnarReaderFactory(parquetBatchReadConf(ParquetReaderType.ICEBERG)); } else if (useOrcBatchReads()) { - return new SparkColumnarReaderFactory(orcBatchReadConf(readConf.orcBatchSize())); + return new SparkColumnarReaderFactory(orcBatchReadConf()); } else { return new SparkRowReaderFactory(); } } - private ParquetBatchReadConf parquetBatchReadConf(int batchSize, ParquetReaderType readerType) { + private ParquetBatchReadConf parquetBatchReadConf(ParquetReaderType readerType) { return ImmutableParquetBatchReadConf.builder() - .batchSize(batchSize) + .batchSize(readConf.parquetBatchSize()) .readerType(readerType) .build(); } - private OrcBatchReadConf orcBatchReadConf(int batchSize) { - return ImmutableOrcBatchReadConf.builder().batchSize(batchSize).build(); + private OrcBatchReadConf orcBatchReadConf() { + return ImmutableOrcBatchReadConf.builder().batchSize(readConf.parquetBatchSize()).build(); } // conditions for using Parquet batch reads: diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java index 876c0be533e0..887b84fb617a 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/source/SparkColumnarReaderFactory.java @@ -29,15 +29,17 @@ import org.apache.spark.sql.vectorized.ColumnarBatch; class SparkColumnarReaderFactory implements PartitionReaderFactory { - private ParquetBatchReadConf parquetConf; - private OrcBatchReadConf orcConf; + private final ParquetBatchReadConf parquetConf; + private final OrcBatchReadConf orcConf; SparkColumnarReaderFactory(ParquetBatchReadConf conf) { this.parquetConf = conf; + this.orcConf = null; } SparkColumnarReaderFactory(OrcBatchReadConf conf) { this.orcConf = conf; + this.parquetConf = null; } @Override From e68bbb5240950e1deae1237ef2e90cd15d00198c Mon Sep 17 00:00:00 2001 From: huaxingao Date: Mon, 21 Oct 2024 10:01:51 -0700 Subject: [PATCH 11/32] update to use comet 0.3.0 --- build.gradle | 2 +- spark/v3.4/build.gradle | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/build.gradle b/build.gradle index 127f95334e13..195a97c8bb87 100644 --- a/build.gradle +++ b/build.gradle @@ -788,7 +788,7 @@ project(':iceberg-parquet') { exclude group: 'org.codehaus.jackson' } - compileOnly "org.apache.comet:comet-spark-spark${sparkMajorVersion}_${scalaVersion}:0.1.0-SNAPSHOT" + compileOnly "org.apache.datafusion:comet-spark-spark${sparkMajorVersion}_${scalaVersion}:0.3.0" compileOnly libs.avro.avro compileOnly(libs.hadoop2.client) { exclude group: 'org.apache.avro', module: 'avro' diff --git a/spark/v3.4/build.gradle b/spark/v3.4/build.gradle index 6fdd6e3f7f4a..9100f13143b1 100644 --- a/spark/v3.4/build.gradle +++ b/spark/v3.4/build.gradle @@ -79,6 +79,8 @@ project(":iceberg-spark:iceberg-spark-${sparkMajorVersion}_${scalaVersion}") { exclude group: 'org.roaringbitmap' } + compileOnly "org.apache.datafusion:comet-spark-spark${sparkMajorVersion}_${scalaVersion}:0.3.0" + implementation libs.parquet.column implementation libs.parquet.hadoop @@ -191,6 +193,7 @@ project(":iceberg-spark:iceberg-spark-extensions-${sparkMajorVersion}_${scalaVer testImplementation libs.avro.avro testImplementation libs.parquet.hadoop testImplementation libs.junit.vintage.engine + testImplementation "org.apache.datafusion:comet-spark-spark${sparkMajorVersion}_${scalaVersion}:0.3.0" // Required because we remove antlr plugin dependencies from the compile configuration, see note above runtimeOnly libs.antlr.runtime From f2fbb3c6597be00d57d6e623aad50dbb7c342f8b Mon Sep 17 00:00:00 2001 From: huaxingao Date: Mon, 21 Oct 2024 11:01:39 -0700 Subject: [PATCH 12/32] use the new Comet Utils.getColumnReader method --- .baseline/checkstyle/checkstyle-suppressions.xml | 3 +++ .../iceberg/spark/data/vectorized/CometColumnReader.java | 7 +++++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/.baseline/checkstyle/checkstyle-suppressions.xml b/.baseline/checkstyle/checkstyle-suppressions.xml index 60b0681a687d..0db6ef1e2893 100644 --- a/.baseline/checkstyle/checkstyle-suppressions.xml +++ b/.baseline/checkstyle/checkstyle-suppressions.xml @@ -48,4 +48,7 @@ + + + diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java index 420700900071..aaacc517ee21 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java @@ -24,6 +24,8 @@ import org.apache.comet.parquet.ColumnReader; import org.apache.comet.parquet.TypeUtil; import org.apache.comet.parquet.Utils; +import org.apache.comet.shaded.arrow.c.CometSchemaImporter; +import org.apache.comet.shaded.arrow.memory.RootAllocator; import org.apache.iceberg.parquet.VectorizedReader; import org.apache.iceberg.spark.SparkSchemaUtil; import org.apache.iceberg.types.Types; @@ -101,7 +103,9 @@ public void reset() { delegate.close(); } - delegate = Utils.getColumnReader(sparkType, descriptor, batchSize, false, false); + CometSchemaImporter importer = new CometSchemaImporter(new RootAllocator()); + + delegate = Utils.getColumnReader(sparkType, descriptor, importer, batchSize, false, false); initialized = true; } @@ -136,7 +140,6 @@ public DataType getSparkType() { * CometColumnReader#reset} is called. */ public void setPageReader(PageReader pageReader) throws IOException { - reset(); if (!initialized) { throw new IllegalStateException("Invalid state: 'reset' should be called first"); } From fa0ee52dcd244682b44557e7841e2534792de5d8 Mon Sep 17 00:00:00 2001 From: huaxingao Date: Mon, 21 Oct 2024 12:02:49 -0700 Subject: [PATCH 13/32] change PARQUET_READER_TYPE_DEFAULT to Comet to test CometReader --- .../main/java/org/apache/iceberg/spark/SparkSQLProperties.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java index 99033c076a43..733744ade421 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java @@ -29,7 +29,7 @@ private SparkSQLProperties() {} // Controls which Parquet reader to use for vectorization public static final String PARQUET_READER_TYPE = "spark.sql.iceberg.parquet.reader-type"; - public static final ParquetReaderType PARQUET_READER_TYPE_DEFAULT = ParquetReaderType.ICEBERG; + public static final ParquetReaderType PARQUET_READER_TYPE_DEFAULT = ParquetReaderType.COMET; // Controls whether reading/writing timestamps without timezones is allowed @Deprecated From 5514f869c19722cc8c55acc01dbd67202556e24d Mon Sep 17 00:00:00 2001 From: huaxingao Date: Mon, 21 Oct 2024 13:22:14 -0700 Subject: [PATCH 14/32] Ignore SmokeTest#testGettingStarted for now --- .../integration/java/org/apache/iceberg/spark/SmokeTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/spark/v3.4/spark-runtime/src/integration/java/org/apache/iceberg/spark/SmokeTest.java b/spark/v3.4/spark-runtime/src/integration/java/org/apache/iceberg/spark/SmokeTest.java index 252793c7b8a7..03b8bd273fa6 100644 --- a/spark/v3.4/spark-runtime/src/integration/java/org/apache/iceberg/spark/SmokeTest.java +++ b/spark/v3.4/spark-runtime/src/integration/java/org/apache/iceberg/spark/SmokeTest.java @@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.analysis.NoSuchTableException; import org.junit.Assert; import org.junit.Before; +import org.junit.Ignore; import org.junit.Test; public class SmokeTest extends SparkExtensionsTestBase { @@ -44,7 +45,7 @@ public void dropTable() { // Run through our Doc's Getting Started Example // TODO Update doc example so that it can actually be run, modifications were required for this // test suite to run - @Test + @Ignore public void testGettingStarted() throws IOException { // Creating a table sql("CREATE TABLE %s (id bigint, data string) USING iceberg", tableName); From 1eb40e5fcbbbe34f3f5af1eeaad34a9a6e35b26f Mon Sep 17 00:00:00 2001 From: huaxingao Date: Tue, 3 Dec 2024 19:28:50 -0800 Subject: [PATCH 15/32] rebase --- .../apache/iceberg/spark/SparkReadConf.java | 8 ++ .../data/vectorized/CometColumnReader.java | 3 - .../vectorized/CometColumnarBatchReader.java | 7 +- .../spark/data/vectorized/CometVector.java | 79 ++++++------------- 4 files changed, 34 insertions(+), 63 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java index fdc9347bc3d1..c547153e4d33 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java @@ -359,4 +359,12 @@ public boolean reportColumnStats() { .defaultValue(SparkSQLProperties.REPORT_COLUMN_STATS_DEFAULT) .parse(); } + + public ParquetReaderType parquetReaderType() { + return confParser + .enumConf(ParquetReaderType::valueOf) + .sessionConf(SparkSQLProperties.PARQUET_READER_TYPE) + .defaultValue(SparkSQLProperties.PARQUET_READER_TYPE_DEFAULT) + .parse(); + } } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java index aaacc517ee21..16ad3bee28d3 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java @@ -113,9 +113,6 @@ public void reset() { public CometVector read(CometVector reuse, int numRows) { delegate.readBatch(numRows); org.apache.comet.vector.CometVector bv = delegate.currentBatch(); - if (reuse == null) { - reuse = vector; - } reuse.setDelegate(bv); return reuse; } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index b2ed77abe15c..14b523ca1c6e 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -69,8 +69,7 @@ public void setRowGroupInfo( && !(readers[i] instanceof CometPositionColumnReader) && !(readers[i] instanceof CometDeleteColumnReader)) { readers[i].reset(); - readers[i].setPageReader( - pageStore.getPageReader(((CometColumnReader) readers[i]).getDescriptor())); + readers[i].setPageReader(pageStore.getPageReader(readers[i].getDescriptor())); } } catch (IOException e) { throw new UncheckedIOException("Failed to setRowGroupInfo for Comet vectorization", e); @@ -78,7 +77,7 @@ public void setRowGroupInfo( } for (int i = 0; i < readers.length; i++) { - delegate.getColumnReaders()[i] = ((CometColumnReader) this.readers[i]).getDelegate(); + delegate.getColumnReaders()[i] = this.readers[i].getDelegate(); } this.rowStartPosInBatch = rowPosition; @@ -154,7 +153,7 @@ void readDeletedColumnIfNecessary(ColumnVector[] columnVectors) { if (readers[i] instanceof CometDeleteColumnReader) { CometDeleteColumnReader deleteColumnReader = new CometDeleteColumnReader<>(isDeleted); deleteColumnReader.setBatchSize(numRowsToRead); - deleteColumnReader.read(null, numRowsToRead); + deleteColumnReader.read(deleteColumnReader.getVector(), numRowsToRead); columnVectors[i] = deleteColumnReader.getVector(); } } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java index 340e0cd13696..c944a3ec52fc 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java @@ -45,109 +45,76 @@ public void setRowIdMapping(int[] rowIdMapping) { @Override public boolean isNullAt(int rowId) { - int newRowId = rowId; - if (rowIdMapping != null) { - newRowId = rowIdMapping[rowId]; - } - return super.isNullAt(newRowId); + return super.isNullAt(mapRowId(rowId)); } @Override public boolean getBoolean(int rowId) { - int newRowId = rowId; - if (rowIdMapping != null) { - newRowId = rowIdMapping[rowId]; - } - return super.getBoolean(newRowId); + return super.getBoolean(mapRowId(rowId)); } @Override public byte getByte(int rowId) { - int newRowId = rowId; - if (rowIdMapping != null) { - newRowId = rowIdMapping[rowId]; - } - return super.getByte(newRowId); + return super.getByte(mapRowId(rowId)); } @Override public short getShort(int rowId) { - int newRowId = rowId; - if (rowIdMapping != null) { - newRowId = rowIdMapping[rowId]; - } - return super.getShort(newRowId); + return super.getShort(mapRowId(rowId)); } @Override public int getInt(int rowId) { - int newRowId = rowId; - if (rowIdMapping != null) { - newRowId = rowIdMapping[rowId]; - } - return super.getInt(newRowId); + return super.getInt(mapRowId(rowId)); } @Override public long getLong(int rowId) { - int newRowId = rowId; - if (rowIdMapping != null) { - newRowId = rowIdMapping[rowId]; - } - return super.getLong(newRowId); + return super.getLong(mapRowId(rowId)); } @Override public float getFloat(int rowId) { - int newRowId = rowId; - if (rowIdMapping != null) { - newRowId = rowIdMapping[rowId]; - } - return super.getFloat(newRowId); + return super.getFloat(mapRowId(rowId)); } @Override public double getDouble(int rowId) { - int newRowId = rowId; - if (rowIdMapping != null) { - newRowId = rowIdMapping[rowId]; - } - return super.getDouble(newRowId); + return super.getDouble(mapRowId(rowId)); } @Override public Decimal getDecimal(int rowId, int precision, int scale) { - int newRowId = rowId; - if (isNullAt(newRowId)) { + if (isNullAt(rowId)) { return null; } - if (rowIdMapping != null) { - newRowId = rowIdMapping[rowId]; - } - return super.getDecimal(newRowId, precision, scale); + + return super.getDecimal(mapRowId(rowId), precision, scale); } @Override public UTF8String getUTF8String(int rowId) { - int newRowId = rowId; - if (isNullAt(newRowId)) { + if (isNullAt(rowId)) { return null; } - if (rowIdMapping != null) { - newRowId = rowIdMapping[rowId]; - } - return super.getUTF8String(newRowId); + + return super.getUTF8String(mapRowId(rowId)); } @Override public byte[] getBinary(int rowId) { - int newRowId = rowId; - if (isNullAt(newRowId)) { + if (isNullAt(rowId)) { return null; } + + return super.getBinary(mapRowId(rowId)); + } + + private int mapRowId(int rowId) { if (rowIdMapping != null) { - newRowId = rowIdMapping[rowId]; + return rowIdMapping[rowId]; } - return super.getBinary(newRowId); + + return rowId; } } From 83196f6a6aa0e78bc1e8e510c247d14361c8775b Mon Sep 17 00:00:00 2001 From: huaxingao Date: Wed, 4 Dec 2024 11:22:02 -0800 Subject: [PATCH 16/32] add setRowGroupInfo(PageReadStore pageStore, Map metaData) --- .../vectorized/CometColumnarBatchReader.java | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index 14b523ca1c6e..dc66a1ec4b69 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -63,11 +63,17 @@ class CometColumnarBatchReader implements VectorizedReader { @Override public void setRowGroupInfo( PageReadStore pageStore, Map metaData, long rowPosition) { + setRowGroupInfo(pageStore, metaData); + } + + @Override + public void setRowGroupInfo( + PageReadStore pageStore, Map metaData) { for (int i = 0; i < readers.length; i++) { try { if (!(readers[i] instanceof CometConstantColumnReader) - && !(readers[i] instanceof CometPositionColumnReader) - && !(readers[i] instanceof CometDeleteColumnReader)) { + && !(readers[i] instanceof CometPositionColumnReader) + && !(readers[i] instanceof CometDeleteColumnReader)) { readers[i].reset(); readers[i].setPageReader(pageStore.getPageReader(readers[i].getDescriptor())); } @@ -80,7 +86,13 @@ public void setRowGroupInfo( delegate.getColumnReaders()[i] = this.readers[i].getDelegate(); } - this.rowStartPosInBatch = rowPosition; + this.rowStartPosInBatch = + pageStore + .getRowIndexOffset() + .orElseThrow( + () -> + new IllegalArgumentException( + "PageReadStore does not contain row index offset")); } public void setDeleteFilter(DeleteFilter deleteFilter) { From d1c6a14f4ae645eb12516d813287b3ed9cd81e2a Mon Sep 17 00:00:00 2001 From: huaxingao Date: Wed, 4 Dec 2024 11:41:54 -0800 Subject: [PATCH 17/32] formatting --- .../vectorized/CometColumnarBatchReader.java | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index dc66a1ec4b69..dd1f0b234892 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -68,12 +68,12 @@ public void setRowGroupInfo( @Override public void setRowGroupInfo( - PageReadStore pageStore, Map metaData) { + PageReadStore pageStore, Map metaData) { for (int i = 0; i < readers.length; i++) { try { if (!(readers[i] instanceof CometConstantColumnReader) - && !(readers[i] instanceof CometPositionColumnReader) - && !(readers[i] instanceof CometDeleteColumnReader)) { + && !(readers[i] instanceof CometPositionColumnReader) + && !(readers[i] instanceof CometDeleteColumnReader)) { readers[i].reset(); readers[i].setPageReader(pageStore.getPageReader(readers[i].getDescriptor())); } @@ -87,12 +87,12 @@ public void setRowGroupInfo( } this.rowStartPosInBatch = - pageStore - .getRowIndexOffset() - .orElseThrow( - () -> - new IllegalArgumentException( - "PageReadStore does not contain row index offset")); + pageStore + .getRowIndexOffset() + .orElseThrow( + () -> + new IllegalArgumentException( + "PageReadStore does not contain row index offset")); } public void setDeleteFilter(DeleteFilter deleteFilter) { From 0eb4ce73f84207c5fe994fe4cbee176fb4e6ed70 Mon Sep 17 00:00:00 2001 From: huaxingao Date: Wed, 4 Dec 2024 14:08:00 -0800 Subject: [PATCH 18/32] ignore a few tests for now --- .../apache/iceberg/spark/source/TestDataFrameWriterV2.java | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestDataFrameWriterV2.java b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestDataFrameWriterV2.java index 47a0e87b9398..fe027dc90686 100644 --- a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestDataFrameWriterV2.java +++ b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestDataFrameWriterV2.java @@ -41,6 +41,7 @@ import org.junit.After; import org.junit.Assert; import org.junit.Before; +import org.junit.Ignore; import org.junit.Test; public class TestDataFrameWriterV2 extends SparkTestBaseWithCatalog { @@ -214,7 +215,7 @@ public void testWriteWithCaseSensitiveOption() throws NoSuchTableException, Pars Assert.assertEquals(4, fields.size()); } - @Test + @Ignore public void testMergeSchemaIgnoreCastingLongToInt() throws Exception { sql( "ALTER TABLE %s SET TBLPROPERTIES ('%s'='true')", @@ -254,7 +255,7 @@ public void testMergeSchemaIgnoreCastingLongToInt() throws Exception { assertThat(idField.type().typeId()).isEqualTo(Type.TypeID.LONG); } - @Test + @Ignore public void testMergeSchemaIgnoreCastingDoubleToFloat() throws Exception { removeTables(); sql("CREATE TABLE %s (id double, data string) USING iceberg", tableName); @@ -296,7 +297,7 @@ public void testMergeSchemaIgnoreCastingDoubleToFloat() throws Exception { assertThat(idField.type().typeId()).isEqualTo(Type.TypeID.DOUBLE); } - @Test + @Ignore public void testMergeSchemaIgnoreCastingDecimalToDecimalWithNarrowerPrecision() throws Exception { removeTables(); sql("CREATE TABLE %s (id decimal(6,2), data string) USING iceberg", tableName); From b9ca9f3e4535bbe128670311eb20858f8ff48270 Mon Sep 17 00:00:00 2001 From: huaxingao Date: Thu, 26 Dec 2024 10:35:00 -0800 Subject: [PATCH 19/32] remove comet dependency in build.gradle --- build.gradle | 2 -- 1 file changed, 2 deletions(-) diff --git a/build.gradle b/build.gradle index 195a97c8bb87..c8bfba7967ce 100644 --- a/build.gradle +++ b/build.gradle @@ -40,7 +40,6 @@ buildscript { } } -String sparkMajorVersion = '3.4' String scalaVersion = System.getProperty("scalaVersion") != null ? System.getProperty("scalaVersion") : System.getProperty("defaultScalaVersion") String sparkVersionsString = System.getProperty("sparkVersions") != null ? System.getProperty("sparkVersions") : System.getProperty("defaultSparkVersions") List sparkVersions = sparkVersionsString != null && !sparkVersionsString.isEmpty() ? sparkVersionsString.split(",") : [] @@ -788,7 +787,6 @@ project(':iceberg-parquet') { exclude group: 'org.codehaus.jackson' } - compileOnly "org.apache.datafusion:comet-spark-spark${sparkMajorVersion}_${scalaVersion}:0.3.0" compileOnly libs.avro.avro compileOnly(libs.hadoop2.client) { exclude group: 'org.apache.avro', module: 'avro' From d552d4aa7141776529819410c21c7aa03346bf12 Mon Sep 17 00:00:00 2001 From: huaxingao Date: Thu, 26 Dec 2024 12:30:16 -0800 Subject: [PATCH 20/32] Trigger Build From 46d0170835366b1054ee06efcae31d3f21690f32 Mon Sep 17 00:00:00 2001 From: huaxingao Date: Sun, 29 Dec 2024 09:40:00 -0800 Subject: [PATCH 21/32] add ColumnarBatchUtil --- .../vectorized/BaseColumnBatchLoader.java | 199 ------------------ .../vectorized/ColumnVectorWithFilter.java | 10 +- .../vectorized/CometColumnarBatchReader.java | 57 +++-- 3 files changed, 45 insertions(+), 221 deletions(-) delete mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java deleted file mode 100644 index 574a5e2ee4ea..000000000000 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/BaseColumnBatchLoader.java +++ /dev/null @@ -1,199 +0,0 @@ -/* - * 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.vectorized; - -import java.util.Iterator; -import org.apache.iceberg.data.DeleteFilter; -import org.apache.iceberg.deletes.PositionDeleteIndex; -import org.apache.iceberg.relocated.com.google.common.base.Preconditions; -import org.apache.iceberg.util.Pair; -import org.apache.spark.sql.catalyst.InternalRow; -import org.apache.spark.sql.vectorized.ColumnVector; -import org.apache.spark.sql.vectorized.ColumnarBatch; - -@SuppressWarnings("checkstyle:VisibilityModifier") -public abstract class BaseColumnBatchLoader { - protected final int numRowsToRead; - // the rowId mapping to skip deleted rows for all column vectors inside a batch, it is null when - // there is no deletes - protected int[] rowIdMapping; - // the array to indicate if a row is deleted or not, it is null when there is no "_deleted" - // metadata column - protected boolean[] isDeleted; - private final boolean hasIsDeletedColumn; - private final DeleteFilter deletes; - private final long rowStartPosInBatch; - - protected BaseColumnBatchLoader( - int numRowsToRead, - boolean hasIsDeletedColumn, - DeleteFilter deletes, - long rowStartPosInBatch) { - Preconditions.checkArgument( - numRowsToRead > 0, "Invalid number of rows to read: %s", numRowsToRead); - this.numRowsToRead = numRowsToRead; - this.hasIsDeletedColumn = hasIsDeletedColumn; - this.deletes = deletes; - this.rowStartPosInBatch = rowStartPosInBatch; - if (hasIsDeletedColumn) { - isDeleted = new boolean[numRowsToRead]; - } - } - - protected ColumnarBatch initializeColumnBatchWithDeletions( - ColumnVector[] arrowColumnVectors, int numRowsUndeleted) { - ColumnarBatch newColumnarBatch = new ColumnarBatch(arrowColumnVectors); - newColumnarBatch.setNumRows(numRowsUndeleted); - - if (hasEqDeletes()) { - applyEqDelete(newColumnarBatch); - } - - if (hasIsDeletedColumn && rowIdMapping != null) { - // reset the row id mapping array, so that it doesn't filter out the deleted rows - for (int i = 0; i < numRowsToRead; i++) { - rowIdMapping[i] = i; - } - newColumnarBatch.setNumRows(numRowsToRead); - } - return newColumnarBatch; - } - - /** - * This method iterates over each column reader and reads the current batch of data into the - * {@link ColumnVector}. - */ - protected abstract ColumnVector[] readDataToColumnVectors(); - - /** - * This method reads the current batch of data into the {@link ColumnVector}, and applies deletion - * logic, and loads data into a {@link ColumnarBatch}. - */ - public abstract ColumnarBatch loadDataToColumnBatch(); - - boolean hasEqDeletes() { - return deletes != null && deletes.hasEqDeletes(); - } - - protected int initRowIdMapping() { - Pair posDeleteRowIdMapping = posDelRowIdMapping(); - if (posDeleteRowIdMapping != null) { - rowIdMapping = posDeleteRowIdMapping.first(); - return posDeleteRowIdMapping.second(); - } else { - rowIdMapping = initEqDeleteRowIdMapping(); - return numRowsToRead; - } - } - - Pair posDelRowIdMapping() { - if (deletes != null && deletes.hasPosDeletes()) { - return buildPosDelRowIdMapping(deletes.deletedRowPositions()); - } else { - return null; - } - } - - /** - * Build a row id mapping inside a batch, which skips deleted rows. Here is an example of how we - * delete 2 rows in a batch with 8 rows in total. [0,1,2,3,4,5,6,7] -- Original status of the row - * id mapping array [F,F,F,F,F,F,F,F] -- Original status of the isDeleted array Position delete 2, - * 6 [0,1,3,4,5,7,-,-] -- After applying position deletes [Set Num records to 6] [F,F,T,F,F,F,T,F] - * -- After applying position deletes - * - * @param deletedRowPositions a set of deleted row positions - * @return the mapping array and the new num of rows in a batch, null if no row is deleted - */ - Pair buildPosDelRowIdMapping(PositionDeleteIndex deletedRowPositions) { - if (deletedRowPositions == null) { - return null; - } - - int[] posDelRowIdMapping = new int[numRowsToRead]; - int originalRowId = 0; - int currentRowId = 0; - while (originalRowId < numRowsToRead) { - if (!deletedRowPositions.isDeleted(originalRowId + rowStartPosInBatch)) { - posDelRowIdMapping[currentRowId] = originalRowId; - currentRowId++; - } else { - if (hasIsDeletedColumn) { - isDeleted[originalRowId] = true; - } - - deletes.incrementDeleteCount(); - } - originalRowId++; - } - - if (currentRowId == numRowsToRead) { - // there is no delete in this batch - return null; - } else { - return Pair.of(posDelRowIdMapping, currentRowId); - } - } - - int[] initEqDeleteRowIdMapping() { - int[] eqDeleteRowIdMapping = null; - if (hasEqDeletes()) { - eqDeleteRowIdMapping = new int[numRowsToRead]; - for (int i = 0; i < numRowsToRead; i++) { - eqDeleteRowIdMapping[i] = i; - } - } - - return eqDeleteRowIdMapping; - } - - /** - * Filter out the equality deleted rows. Here is an example, [0,1,2,3,4,5,6,7] -- Original status - * of the row id mapping array [F,F,F,F,F,F,F,F] -- Original status of the isDeleted array - * Position delete 2, 6 [0,1,3,4,5,7,-,-] -- After applying position deletes [Set Num records to - * 6] [F,F,T,F,F,F,T,F] -- After applying position deletes Equality delete 1 <= x <= 3 - * [0,4,5,7,-,-,-,-] -- After applying equality deletes [Set Num records to 4] [F,T,T,T,F,F,T,F] - * -- After applying equality deletes - * - * @param columnarBatch the {@link ColumnarBatch} to apply the equality delete - */ - void applyEqDelete(ColumnarBatch columnarBatch) { - Iterator it = columnarBatch.rowIterator(); - int rowId = 0; - int currentRowId = 0; - while (it.hasNext()) { - InternalRow row = it.next(); - if (deletes.eqDeletedRowFilter().test(row)) { - // the row is NOT deleted - // skip deleted rows by pointing to the next undeleted row Id - rowIdMapping[currentRowId] = rowIdMapping[rowId]; - currentRowId++; - } else { - if (hasIsDeletedColumn) { - isDeleted[rowIdMapping[rowId]] = true; - } - - deletes.incrementDeleteCount(); - } - - rowId++; - } - - columnarBatch.setNumRows(currentRowId); - } -} diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java index ab0d652321d3..f343478268d4 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java @@ -24,13 +24,21 @@ import org.apache.spark.unsafe.types.UTF8String; public class ColumnVectorWithFilter extends IcebergArrowColumnVector { - private final int[] rowIdMapping; + private int[] rowIdMapping; public ColumnVectorWithFilter(VectorHolder holder, int[] rowIdMapping) { super(holder); this.rowIdMapping = rowIdMapping; } + public ColumnVectorWithFilter(VectorHolder holder) { + super(holder); + } + + public void setRowIdMapping(int[] rowIdMapping) { + this.rowIdMapping = rowIdMapping; + } + @Override public boolean isNullAt(int rowId) { return nullabilityHolder().isNullAt(rowIdMapping[rowId]) == 1; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index dd1f0b234892..cd5b1663b904 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -27,6 +27,7 @@ import org.apache.iceberg.Schema; import org.apache.iceberg.data.DeleteFilter; import org.apache.iceberg.parquet.VectorizedReader; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import org.apache.iceberg.spark.SparkSchemaUtil; import org.apache.parquet.column.page.PageReadStore; import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; @@ -124,27 +125,53 @@ public void close() { } } - private class ColumnBatchLoader extends BaseColumnBatchLoader { + private class ColumnBatchLoader { + private final int numRowsToRead; + // the rowId mapping to skip deleted rows for all column vectors inside a batch, it is null when + // there is no deletes + private int[] rowIdMapping; + // the array to indicate if a row is deleted or not, it is null when there is no "_deleted" + // metadata column + private boolean[] isDeleted; + ColumnBatchLoader(int numRowsToRead) { - super(numRowsToRead, hasIsDeletedColumn, deletes, rowStartPosInBatch); + Preconditions.checkArgument( + numRowsToRead > 0, "Invalid number of rows to read: %s", numRowsToRead); + this.numRowsToRead = numRowsToRead; + if (hasIsDeletedColumn) { + isDeleted = new boolean[numRowsToRead]; + } } - @Override - public ColumnarBatch loadDataToColumnBatch() { - int numRowsUndeleted = initRowIdMapping(); + ColumnarBatch loadDataToColumnBatch() { ColumnVector[] arrowColumnVectors = readDataToColumnVectors(); + ColumnarBatch columnarBatch = new ColumnarBatch(arrowColumnVectors); + ColumnarBatchUtil.applyDeletesToColumnarBatch( + columnarBatch, deletes, isDeleted, numRowsToRead, rowStartPosInBatch, hasIsDeletedColumn); - ColumnarBatch newColumnarBatch = - initializeColumnBatchWithDeletions(arrowColumnVectors, numRowsUndeleted); + if (hasIsDeletedColumn) { + // reset the row id mapping array, so that it doesn't filter out the deleted rows + ColumnarBatchUtil.resetRowIdMapping(columnarBatch, numRowsToRead); + } if (hasIsDeletedColumn) { readDeletedColumnIfNecessary(arrowColumnVectors); } - return newColumnarBatch; + return columnarBatch; + } + + void readDeletedColumnIfNecessary(ColumnVector[] columnVectors) { + for (int i = 0; i < readers.length; i++) { + if (readers[i] instanceof CometDeleteColumnReader) { + CometDeleteColumnReader deleteColumnReader = new CometDeleteColumnReader<>(isDeleted); + deleteColumnReader.setBatchSize(numRowsToRead); + deleteColumnReader.read(deleteColumnReader.getVector(), numRowsToRead); + columnVectors[i] = deleteColumnReader.getVector(); + } + } } - @Override protected ColumnVector[] readDataToColumnVectors() { ColumnVector[] columnVectors = new ColumnVector[readers.length]; // Fetch rows for all readers in the delegate @@ -153,22 +180,10 @@ protected ColumnVector[] readDataToColumnVectors() { CometVector bv = readers[i].getVector(); org.apache.comet.vector.CometVector vector = readers[i].getDelegate().currentBatch(); bv.setDelegate(vector); - bv.setRowIdMapping(rowIdMapping); columnVectors[i] = bv; } return columnVectors; } - - void readDeletedColumnIfNecessary(ColumnVector[] columnVectors) { - for (int i = 0; i < readers.length; i++) { - if (readers[i] instanceof CometDeleteColumnReader) { - CometDeleteColumnReader deleteColumnReader = new CometDeleteColumnReader<>(isDeleted); - deleteColumnReader.setBatchSize(numRowsToRead); - deleteColumnReader.read(deleteColumnReader.getVector(), numRowsToRead); - columnVectors[i] = deleteColumnReader.getVector(); - } - } - } } } From 9db707d3568cc2e5941d14b93fa98c5df10cbdf9 Mon Sep 17 00:00:00 2001 From: huaxingao Date: Sat, 25 Jan 2025 19:24:05 -0800 Subject: [PATCH 22/32] rebase --- .../iceberg/spark/ParquetReaderType.java | 12 +++ .../data/vectorized/CometColumnReader.java | 30 +------ .../vectorized/CometColumnarBatchReader.java | 81 ++++++++++--------- .../vectorized/CometConstantColumnReader.java | 3 +- .../vectorized/CometDeleteColumnReader.java | 9 +-- .../vectorized/CometPositionColumnReader.java | 6 +- .../spark/data/vectorized/CometVector.java | 11 +-- .../CometVectorizedReaderBuilder.java | 9 ++- .../VectorizedSparkParquetReaders.java | 17 ++++ .../iceberg/spark/source/TestParquetScan.java | 12 +++ 10 files changed, 105 insertions(+), 85 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java index fac1604c0bd1..b0b197377338 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java @@ -18,7 +18,19 @@ */ package org.apache.iceberg.spark; +/** Enumerates the types of Parquet readers. */ public enum ParquetReaderType { + /** ICEBERG type utilizes the Parquet reader from Apache Iceberg. */ ICEBERG, + + /** + * COMET type changes the Parquet reader to the Apache DataFusion Comet Parquet reader. Comet + * Parquet reader performs I/O and decompression in the JVM but decodes in native to improve + * performance. Additionally, Comet will convert Spark's physical plan into a native physical plan + * and execute this plan natively. + * + *

TODO: Implement {@link org.apache.comet.parquet.SupportsComet} in SparkScan to convert Spark + * physical plan to native physical plan for native execution. + */ COMET } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java index 16ad3bee28d3..231a6ed32a88 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java @@ -38,37 +38,12 @@ import org.apache.spark.sql.types.Metadata; import org.apache.spark.sql.types.StructField; -/** - * A Iceberg Parquet column reader backed by a Comet {@link ColumnReader}. This class should be used - * together with {@link CometVector}. - * - *

Example: - * - * 

- *   CometColumnReader reader = ...
- *   reader.setBatchSize(batchSize);
- *
- *   while (hasMoreRowsToRead) {
- *     if (endOfRowGroup) {
- *       reader.reset();
- *       PageReader pageReader = ...
- *       reader.setPageReader(pageReader);
- *     }
- *
- *     int numRows = ...
- *     CometVector vector = reader.read(null, numRows);
- *
- *     // consume the vector
- *   }
- *
- *   reader.close();
- *
- */ @SuppressWarnings({"checkstyle:VisibilityModifier", "ParameterAssignment"}) class CometColumnReader implements VectorizedReader { public static final int DEFAULT_BATCH_SIZE = 5000; private final DataType sparkType; + // the delegated column reader from Comet side protected AbstractColumnReader delegate; private final CometVector vector; private final ColumnDescriptor descriptor; @@ -94,7 +69,7 @@ public AbstractColumnReader getDelegate() { } /** - * This method is to initialized/reset the ColumnReader. This needs to be called for each row + * This method is to initialized/reset the CometColumnReader. This needs to be called for each row * group after readNextRowGroup, so a new dictionary encoding can be set for each of the new row * groups. */ @@ -145,6 +120,7 @@ public void setPageReader(PageReader pageReader) throws IOException { @Override public void close() { + // close reader on native side if (delegate != null) { delegate.close(); } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index cd5b1663b904..d11cd5a4cc37 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -29,6 +29,7 @@ import org.apache.iceberg.parquet.VectorizedReader; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import org.apache.iceberg.spark.SparkSchemaUtil; +import org.apache.iceberg.util.Pair; import org.apache.parquet.column.page.PageReadStore; import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.hadoop.metadata.ColumnPath; @@ -39,7 +40,7 @@ /** * {@link VectorizedReader} that returns Spark's {@link ColumnarBatch} to support Spark's vectorized * read path. The {@link ColumnarBatch} returned is created by passing in the Arrow vectors - * populated via delegated read calls to {@linkplain CometColumnReader VectorReader(s)}. + * populated via delegated read calls to {@link CometColumnReader VectorReader(s)}. */ @SuppressWarnings("checkstyle:VisibilityModifier") class CometColumnarBatchReader implements VectorizedReader { @@ -48,6 +49,7 @@ class CometColumnarBatchReader implements VectorizedReader { private final boolean hasIsDeletedColumn; private DeleteFilter deletes = null; private long rowStartPosInBatch = 0; + // The delegated batch reader on Comet side private final BatchReader delegate; CometColumnarBatchReader(List> readers, Schema schema) { @@ -126,64 +128,67 @@ public void close() { } private class ColumnBatchLoader { - private final int numRowsToRead; - // the rowId mapping to skip deleted rows for all column vectors inside a batch, it is null when - // there is no deletes - private int[] rowIdMapping; - // the array to indicate if a row is deleted or not, it is null when there is no "_deleted" - // metadata column - private boolean[] isDeleted; + private final int batchSize; ColumnBatchLoader(int numRowsToRead) { Preconditions.checkArgument( numRowsToRead > 0, "Invalid number of rows to read: %s", numRowsToRead); - this.numRowsToRead = numRowsToRead; - if (hasIsDeletedColumn) { - isDeleted = new boolean[numRowsToRead]; - } + this.batchSize = numRowsToRead; } ColumnarBatch loadDataToColumnBatch() { ColumnVector[] arrowColumnVectors = readDataToColumnVectors(); - ColumnarBatch columnarBatch = new ColumnarBatch(arrowColumnVectors); - ColumnarBatchUtil.applyDeletesToColumnarBatch( - columnarBatch, deletes, isDeleted, numRowsToRead, rowStartPosInBatch, hasIsDeletedColumn); - + int numLiveRows = batchSize; if (hasIsDeletedColumn) { - // reset the row id mapping array, so that it doesn't filter out the deleted rows - ColumnarBatchUtil.resetRowIdMapping(columnarBatch, numRowsToRead); + boolean[] isDeleted = + ColumnarBatchUtil.buildIsDeleted( + arrowColumnVectors, deletes, rowStartPosInBatch, batchSize); + readDeletedColumn(arrowColumnVectors, isDeleted); + } else { + Pair pair = + ColumnarBatchUtil.buildRowIdMapping( + arrowColumnVectors, deletes, rowStartPosInBatch, batchSize); + if (pair != null) { + int[] rowIdMapping = pair.first(); + numLiveRows = pair.second(); + for (int i = 0; i < arrowColumnVectors.length; i++) { + ((CometVector) arrowColumnVectors[i]).setRowIdMapping(rowIdMapping); + } + } } - if (hasIsDeletedColumn) { - readDeletedColumnIfNecessary(arrowColumnVectors); + if (deletes != null && deletes.hasEqDeletes()) { + arrowColumnVectors = ColumnarBatchUtil.removeExtraColumns(deletes, arrowColumnVectors); } - return columnarBatch; + ColumnarBatch newColumnarBatch = new ColumnarBatch(arrowColumnVectors); + newColumnarBatch.setNumRows(numLiveRows); + return newColumnarBatch; } - void readDeletedColumnIfNecessary(ColumnVector[] columnVectors) { - for (int i = 0; i < readers.length; i++) { - if (readers[i] instanceof CometDeleteColumnReader) { - CometDeleteColumnReader deleteColumnReader = new CometDeleteColumnReader<>(isDeleted); - deleteColumnReader.setBatchSize(numRowsToRead); - deleteColumnReader.read(deleteColumnReader.getVector(), numRowsToRead); - columnVectors[i] = deleteColumnReader.getVector(); - } - } - } - - protected ColumnVector[] readDataToColumnVectors() { - ColumnVector[] columnVectors = new ColumnVector[readers.length]; + ColumnVector[] readDataToColumnVectors() { + CometVector[] columnVectors = new CometVector[readers.length]; // Fetch rows for all readers in the delegate - delegate.nextBatch(numRowsToRead); + delegate.nextBatch(batchSize); for (int i = 0; i < readers.length; i++) { - CometVector bv = readers[i].getVector(); + columnVectors[i] = readers[i].getVector(); + columnVectors[i].resetRowIdMapping(); org.apache.comet.vector.CometVector vector = readers[i].getDelegate().currentBatch(); - bv.setDelegate(vector); - columnVectors[i] = bv; + columnVectors[i].setDelegate(vector); } return columnVectors; } + + void readDeletedColumn(ColumnVector[] columnVectors, boolean[] isDeleted) { + for (int i = 0; i < readers.length; i++) { + if (readers[i] instanceof CometDeleteColumnReader) { + CometDeleteColumnReader deleteColumnReader = new CometDeleteColumnReader<>(isDeleted); + deleteColumnReader.setBatchSize(batchSize); + deleteColumnReader.read(deleteColumnReader.getVector(), batchSize); + columnVectors[i] = deleteColumnReader.getVector(); + } + } + } } } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java index b2f0c057ae6d..29d87914c9d9 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java @@ -22,11 +22,10 @@ import org.apache.iceberg.types.Types; class CometConstantColumnReader extends CometColumnReader { - private final T value; CometConstantColumnReader(T value, Types.NestedField field) { super(field); - this.value = value; + // use delegate to set constant value on the native side to be consumed by native execution. delegate = new ConstantColumnReader(getSparkType(), getDescriptor(), value, false); } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java index 7b7a395c5a93..f476f1056558 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java @@ -22,6 +22,7 @@ import org.apache.comet.parquet.MetadataColumnReader; import org.apache.comet.parquet.Native; import org.apache.comet.parquet.TypeUtil; +import org.apache.iceberg.MetadataColumns; import org.apache.iceberg.types.Types; import org.apache.spark.sql.types.DataTypes; import org.apache.spark.sql.types.Metadata; @@ -34,10 +35,7 @@ class CometDeleteColumnReader extends CometColumnReader { } CometDeleteColumnReader(boolean[] isDeleted) { - super( - DataTypes.BooleanType, - TypeUtil.convertToParquet( - new StructField("deleted", DataTypes.BooleanType, false, Metadata.empty()))); + super(MetadataColumns.IS_DELETED); delegate = new DeleteColumnReader(isDeleted); } @@ -55,7 +53,7 @@ private static class DeleteColumnReader extends MetadataColumnReader { super( DataTypes.BooleanType, TypeUtil.convertToParquet( - new StructField("deleted", DataTypes.BooleanType, false, Metadata.empty())), + new StructField("_deleted", DataTypes.BooleanType, false, Metadata.empty())), false); this.isDeleted = isDeleted; } @@ -63,6 +61,7 @@ private static class DeleteColumnReader extends MetadataColumnReader { @Override public void readBatch(int total) { Native.resetBatch(nativeHandle); + // set isDeleted on the native side to be consumed by native execution Native.setIsDeleted(nativeHandle, isDeleted); super.readBatch(total); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java index 3c92877864f7..d5a6c586ede3 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java @@ -42,17 +42,13 @@ private static class PositionColumnReader extends MetadataColumnReader { private long position; PositionColumnReader(ColumnDescriptor descriptor) { - this(descriptor, 0L); - } - - PositionColumnReader(ColumnDescriptor descriptor, long position) { super(DataTypes.LongType, descriptor, false); - this.position = position; } @Override public void readBatch(int total) { Native.resetBatch(nativeHandle); + // set position on the native side to be consumed by native execution Native.setPosition(nativeHandle, position, total); position += total; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java index c944a3ec52fc..ceed84fcbcdd 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java @@ -26,13 +26,6 @@ @SuppressWarnings("checkstyle:VisibilityModifier") class CometVector extends CometDelegateVector { - // the rowId mapping to skip deleted rows for all column vectors inside a batch - // Here is an example: - // [0,1,2,3,4,5,6,7] -- Original status of the row id mapping array - // Position delete 2, 6 - // [0,1,3,4,5,7,-,-] -- After applying position deletes [Set Num records to 6] - // Equality delete 1 <= x <= 3 - // [0,4,5,7,-,-,-,-] -- After applying equality deletes [Set Num records to 4] protected int[] rowIdMapping; CometVector(DataType type, boolean useDecimal128) { @@ -43,6 +36,10 @@ public void setRowIdMapping(int[] rowIdMapping) { this.rowIdMapping = rowIdMapping; } + public void resetRowIdMapping() { + this.rowIdMapping = null; + } + @Override public boolean isNullAt(int rowId) { return super.isNullAt(mapRowId(rowId)); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java index 2cc24f6ce98d..d36f1a727477 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java @@ -91,9 +91,16 @@ public VectorizedReader message( reorderedFields.add(deleteReader); } else if (reader != null) { reorderedFields.add(reader); - } else { + } else if (field.initialDefault() != null) { + CometColumnReader constantReader = + new CometConstantColumnReader<>(field.initialDefault(), field); + reorderedFields.add(constantReader); + } else if (field.isOptional()) { CometColumnReader constantReader = new CometConstantColumnReader<>(null, field); reorderedFields.add(constantReader); + } else { + throw new IllegalArgumentException( + String.format("Missing required field: %s", field.name())); } } return vectorizedReader(reorderedFields); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkParquetReaders.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkParquetReaders.java index 636ad3be7dcc..b523bc5bff11 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkParquetReaders.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkParquetReaders.java @@ -70,6 +70,23 @@ public static ColumnarBatchReader buildReader( deleteFilter)); } + public static CometColumnarBatchReader buildCometReader( + Schema expectedSchema, + MessageType fileSchema, + Map idToConstant, + DeleteFilter deleteFilter) { + return (CometColumnarBatchReader) + TypeWithSchemaVisitor.visit( + expectedSchema.asStruct(), + fileSchema, + new CometVectorizedReaderBuilder( + expectedSchema, + fileSchema, + idToConstant, + readers -> new CometColumnarBatchReader(readers, expectedSchema), + deleteFilter)); + } + // enables unsafe memory access to avoid costly checks to see if index is within bounds // as long as it is not configured explicitly (see BoundsChecking in Arrow) private static void enableUnsafeMemoryAccess() { diff --git a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestParquetScan.java b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestParquetScan.java index 6b9ec85b7f0b..dbbb52794555 100644 --- a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestParquetScan.java +++ b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestParquetScan.java @@ -19,6 +19,7 @@ package org.apache.iceberg.spark.source; import static org.apache.iceberg.Files.localOutput; +import static org.apache.iceberg.spark.SparkSQLProperties.PARQUET_READER_TYPE; import static org.assertj.core.api.Assumptions.assumeThat; import java.io.File; @@ -35,14 +36,25 @@ import org.apache.iceberg.TableProperties; import org.apache.iceberg.io.FileAppender; import org.apache.iceberg.parquet.Parquet; +import org.apache.iceberg.spark.ParquetReaderType; import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.sql.SparkSession; +import org.junit.jupiter.api.BeforeAll; public class TestParquetScan extends ScanTestBase { protected boolean vectorized() { return false; } + @BeforeAll + public static void startSpark() { + ScanTestBase.spark = SparkSession.builder().master("local[2]").getOrCreate(); + ScanTestBase.spark.conf().set(PARQUET_READER_TYPE, ParquetReaderType.ICEBERG.toString()); + ScanTestBase.sc = JavaSparkContext.fromSparkContext(spark.sparkContext()); + } + @Override protected void configureTable(Table table) { table From d61325b70ac50911e6a61ef7738a0311533f2ab6 Mon Sep 17 00:00:00 2001 From: huaxingao Date: Mon, 27 Jan 2025 17:37:35 -0800 Subject: [PATCH 23/32] rebase --- .../vectorized/ColumnVectorWithFilter.java | 107 +++++++++++++----- .../data/vectorized/ColumnarBatchReader.java | 39 ++++--- .../vectorized/CometColumnarBatchReader.java | 32 +++--- .../vectorized/CometDeleteColumnReader.java | 12 +- .../vectorized/CometPositionColumnReader.java | 2 +- .../vectorized/IcebergArrowColumnVector.java | 6 - 6 files changed, 123 insertions(+), 75 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java index f343478268d4..77becc4ab808 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java @@ -18,86 +18,133 @@ */ package org.apache.iceberg.spark.data.vectorized; -import org.apache.iceberg.arrow.vectorized.VectorHolder; import org.apache.spark.sql.types.Decimal; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.vectorized.ColumnVector; import org.apache.spark.sql.vectorized.ColumnarArray; +import org.apache.spark.sql.vectorized.ColumnarMap; import org.apache.spark.unsafe.types.UTF8String; -public class ColumnVectorWithFilter extends IcebergArrowColumnVector { - private int[] rowIdMapping; +/** + * A column vector implementation that applies row-level filtering. + * + *

This class wraps an existing column vector and uses a row ID mapping array to remap row + * indices during data access. Each method that retrieves data for a specific row translates the + * provided row index using the mapping array, effectively filtering the original data to only + * expose the live subset of rows. This approach allows efficient row-level filtering without + * modifying the underlying data. + */ +public class ColumnVectorWithFilter extends ColumnVector { + private final ColumnVector delegate; + private final int[] rowIdMapping; + private volatile ColumnVectorWithFilter[] children = null; - public ColumnVectorWithFilter(VectorHolder holder, int[] rowIdMapping) { - super(holder); + public ColumnVectorWithFilter(ColumnVector delegate, int[] rowIdMapping) { + super(delegate.dataType()); + this.delegate = delegate; this.rowIdMapping = rowIdMapping; } - public ColumnVectorWithFilter(VectorHolder holder) { - super(holder); + @Override + public void close() { + delegate.close(); } - public void setRowIdMapping(int[] rowIdMapping) { - this.rowIdMapping = rowIdMapping; + @Override + public boolean hasNull() { + return delegate.hasNull(); + } + + @Override + public int numNulls() { + // computing the actual number of nulls with rowIdMapping is expensive + // it is OK to overestimate and return the number of nulls in the original vector + return delegate.numNulls(); } @Override public boolean isNullAt(int rowId) { - return nullabilityHolder().isNullAt(rowIdMapping[rowId]) == 1; + return delegate.isNullAt(rowIdMapping[rowId]); } @Override public boolean getBoolean(int rowId) { - return accessor().getBoolean(rowIdMapping[rowId]); + return delegate.getBoolean(rowIdMapping[rowId]); + } + + @Override + public byte getByte(int rowId) { + return delegate.getByte(rowIdMapping[rowId]); + } + + @Override + public short getShort(int rowId) { + return delegate.getShort(rowIdMapping[rowId]); } @Override public int getInt(int rowId) { - return accessor().getInt(rowIdMapping[rowId]); + return delegate.getInt(rowIdMapping[rowId]); } @Override public long getLong(int rowId) { - return accessor().getLong(rowIdMapping[rowId]); + return delegate.getLong(rowIdMapping[rowId]); } @Override public float getFloat(int rowId) { - return accessor().getFloat(rowIdMapping[rowId]); + return delegate.getFloat(rowIdMapping[rowId]); } @Override public double getDouble(int rowId) { - return accessor().getDouble(rowIdMapping[rowId]); + return delegate.getDouble(rowIdMapping[rowId]); } @Override public ColumnarArray getArray(int rowId) { - if (isNullAt(rowId)) { - return null; - } - return accessor().getArray(rowIdMapping[rowId]); + return delegate.getArray(rowIdMapping[rowId]); + } + + @Override + public ColumnarMap getMap(int rowId) { + return delegate.getMap(rowIdMapping[rowId]); } @Override public Decimal getDecimal(int rowId, int precision, int scale) { - if (isNullAt(rowId)) { - return null; - } - return accessor().getDecimal(rowIdMapping[rowId], precision, scale); + return delegate.getDecimal(rowIdMapping[rowId], precision, scale); } @Override public UTF8String getUTF8String(int rowId) { - if (isNullAt(rowId)) { - return null; - } - return accessor().getUTF8String(rowIdMapping[rowId]); + return delegate.getUTF8String(rowIdMapping[rowId]); } @Override public byte[] getBinary(int rowId) { - if (isNullAt(rowId)) { - return null; + return delegate.getBinary(rowIdMapping[rowId]); + } + + @Override + public ColumnVector getChild(int ordinal) { + if (children == null) { + synchronized (this) { + if (children == null) { + if (dataType() instanceof StructType) { + StructType structType = (StructType) dataType(); + this.children = new ColumnVectorWithFilter[structType.length()]; + for (int index = 0; index < structType.length(); index++) { + children[index] = new ColumnVectorWithFilter(delegate.getChild(index), rowIdMapping); + } + } else { + throw new UnsupportedOperationException("Unsupported nested type: " + dataType()); + } + } + } } - return accessor().getBinary(rowIdMapping[rowId]); + + return children[ordinal]; } } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnarBatchReader.java index c65c24d02f59..2123939399cb 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnarBatchReader.java @@ -94,43 +94,42 @@ private class ColumnBatchLoader { } ColumnarBatch loadDataToColumnBatch() { - ColumnVector[] arrowColumnVectors = readDataToColumnVectors(); + ColumnVector[] vectors = readDataToColumnVectors(); int numLiveRows = batchSize; + if (hasIsDeletedColumn) { - boolean[] isDeleted = - ColumnarBatchUtil.buildIsDeleted( - arrowColumnVectors, deletes, rowStartPosInBatch, batchSize); - for (int i = 0; i < arrowColumnVectors.length; i++) { - ColumnVector vector = arrowColumnVectors[i]; + boolean[] isDeleted = buildIsDeleted(vectors); + for (ColumnVector vector : vectors) { if (vector instanceof DeletedColumnVector) { ((DeletedColumnVector) vector).setValue(isDeleted); } } } else { - Pair pair = - ColumnarBatchUtil.buildRowIdMapping( - arrowColumnVectors, deletes, rowStartPosInBatch, batchSize); + Pair pair = buildRowIdMapping(vectors); if (pair != null) { int[] rowIdMapping = pair.first(); numLiveRows = pair.second(); - for (int i = 0; i < arrowColumnVectors.length; i++) { - ColumnVector vector = arrowColumnVectors[i]; - if (vector instanceof IcebergArrowColumnVector) { - arrowColumnVectors[i] = - new ColumnVectorWithFilter( - ((IcebergArrowColumnVector) vector).vector(), rowIdMapping); - } + for (int i = 0; i < vectors.length; i++) { + vectors[i] = new ColumnVectorWithFilter(vectors[i], rowIdMapping); } } } if (deletes != null && deletes.hasEqDeletes()) { - arrowColumnVectors = ColumnarBatchUtil.removeExtraColumns(deletes, arrowColumnVectors); + vectors = ColumnarBatchUtil.removeExtraColumns(deletes, vectors); } - ColumnarBatch newColumnarBatch = new ColumnarBatch(arrowColumnVectors); - newColumnarBatch.setNumRows(numLiveRows); - return newColumnarBatch; + ColumnarBatch batch = new ColumnarBatch(vectors); + batch.setNumRows(numLiveRows); + return batch; + } + + private boolean[] buildIsDeleted(ColumnVector[] vectors) { + return ColumnarBatchUtil.buildIsDeleted(vectors, deletes, rowStartPosInBatch, batchSize); + } + + private Pair buildRowIdMapping(ColumnVector[] vectors) { + return ColumnarBatchUtil.buildRowIdMapping(vectors, deletes, rowStartPosInBatch, batchSize); } ColumnVector[] readDataToColumnVectors() { diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index d11cd5a4cc37..42c4b7662e30 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -137,33 +137,37 @@ private class ColumnBatchLoader { } ColumnarBatch loadDataToColumnBatch() { - ColumnVector[] arrowColumnVectors = readDataToColumnVectors(); + ColumnVector[] vectors = readDataToColumnVectors(); int numLiveRows = batchSize; if (hasIsDeletedColumn) { - boolean[] isDeleted = - ColumnarBatchUtil.buildIsDeleted( - arrowColumnVectors, deletes, rowStartPosInBatch, batchSize); - readDeletedColumn(arrowColumnVectors, isDeleted); + boolean[] isDeleted = buildIsDeleted(vectors); + readDeletedColumn(vectors, isDeleted); } else { - Pair pair = - ColumnarBatchUtil.buildRowIdMapping( - arrowColumnVectors, deletes, rowStartPosInBatch, batchSize); + Pair pair = buildRowIdMapping(vectors); if (pair != null) { int[] rowIdMapping = pair.first(); numLiveRows = pair.second(); - for (int i = 0; i < arrowColumnVectors.length; i++) { - ((CometVector) arrowColumnVectors[i]).setRowIdMapping(rowIdMapping); + for (int i = 0; i < vectors.length; i++) { + ((CometVector) vectors[i]).setRowIdMapping(rowIdMapping); } } } if (deletes != null && deletes.hasEqDeletes()) { - arrowColumnVectors = ColumnarBatchUtil.removeExtraColumns(deletes, arrowColumnVectors); + vectors = ColumnarBatchUtil.removeExtraColumns(deletes, vectors); } - ColumnarBatch newColumnarBatch = new ColumnarBatch(arrowColumnVectors); - newColumnarBatch.setNumRows(numLiveRows); - return newColumnarBatch; + ColumnarBatch batch = new ColumnarBatch(vectors); + batch.setNumRows(numLiveRows); + return batch; + } + + private boolean[] buildIsDeleted(ColumnVector[] vectors) { + return ColumnarBatchUtil.buildIsDeleted(vectors, deletes, rowStartPosInBatch, batchSize); + } + + private Pair buildRowIdMapping(ColumnVector[] vectors) { + return ColumnarBatchUtil.buildRowIdMapping(vectors, deletes, rowStartPosInBatch, batchSize); } ColumnVector[] readDataToColumnVectors() { diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java index f476f1056558..e40774613c2e 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java @@ -18,7 +18,6 @@ */ package org.apache.iceberg.spark.data.vectorized; -import org.apache.comet.parquet.ConstantColumnReader; import org.apache.comet.parquet.MetadataColumnReader; import org.apache.comet.parquet.Native; import org.apache.comet.parquet.TypeUtil; @@ -31,7 +30,7 @@ class CometDeleteColumnReader extends CometColumnReader { CometDeleteColumnReader(Types.NestedField field) { super(field); - delegate = new ConstantColumnReader(getSparkType(), getDescriptor(), false, false); + delegate = new DeleteColumnReader(); } CometDeleteColumnReader(boolean[] isDeleted) { @@ -49,12 +48,17 @@ public void setBatchSize(int batchSize) { private static class DeleteColumnReader extends MetadataColumnReader { private boolean[] isDeleted; - DeleteColumnReader(boolean[] isDeleted) { + DeleteColumnReader() { super( DataTypes.BooleanType, TypeUtil.convertToParquet( new StructField("_deleted", DataTypes.BooleanType, false, Metadata.empty())), - false); + false /* useDecimal128 = false */); + this.isDeleted = new boolean[0]; + } + + DeleteColumnReader(boolean[] isDeleted) { + this(); this.isDeleted = isDeleted; } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java index d5a6c586ede3..31ccc591a7ad 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java @@ -42,7 +42,7 @@ private static class PositionColumnReader extends MetadataColumnReader { private long position; PositionColumnReader(ColumnDescriptor descriptor) { - super(DataTypes.LongType, descriptor, false); + super(DataTypes.LongType, descriptor, false /* useDecimal128 = false */); } @Override diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/IcebergArrowColumnVector.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/IcebergArrowColumnVector.java index afd5810bca2a..38ec3a0e838c 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/IcebergArrowColumnVector.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/IcebergArrowColumnVector.java @@ -39,17 +39,11 @@ public class IcebergArrowColumnVector extends ColumnVector { private final ArrowVectorAccessor accessor; private final NullabilityHolder nullabilityHolder; - private final VectorHolder holder; public IcebergArrowColumnVector(VectorHolder holder) { super(SparkSchemaUtil.convert(holder.icebergType())); this.nullabilityHolder = holder.nullabilityHolder(); this.accessor = ArrowVectorAccessors.getVectorAccessor(holder); - this.holder = holder; - } - - public VectorHolder vector() { - return holder; } protected ArrowVectorAccessor accessor() { From e173dd32e50ee107e402b630cb944dd0aa16385b Mon Sep 17 00:00:00 2001 From: huaxingao Date: Tue, 28 Jan 2025 14:38:58 -0800 Subject: [PATCH 24/32] convert constant value to Spark format --- .../vectorized/CometConstantColumnReader.java | 25 ++++++++++++++++++- .../iceberg/spark/source/TestParquetScan.java | 12 --------- 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java index 29d87914c9d9..30d48a709c3e 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java @@ -18,15 +18,23 @@ */ package org.apache.iceberg.spark.data.vectorized; +import java.math.BigDecimal; import org.apache.comet.parquet.ConstantColumnReader; import org.apache.iceberg.types.Types; +import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Decimal; +import org.apache.spark.sql.types.DecimalType; +import org.apache.spark.unsafe.types.UTF8String; class CometConstantColumnReader extends CometColumnReader { CometConstantColumnReader(T value, Types.NestedField field) { super(field); // use delegate to set constant value on the native side to be consumed by native execution. - delegate = new ConstantColumnReader(getSparkType(), getDescriptor(), value, false); + delegate = + new ConstantColumnReader( + getSparkType(), getDescriptor(), convertToSparkValue(value), false); } @Override @@ -35,4 +43,19 @@ public void setBatchSize(int batchSize) { this.batchSize = batchSize; initialized = true; } + + private Object convertToSparkValue(T value) { + DataType dataType = getSparkType(); + if (dataType == DataTypes.StringType) { + return UTF8String.fromString((String) value); + } else if (dataType instanceof DecimalType) { + return Decimal.apply((BigDecimal) value); + } else if (dataType == DataTypes.BinaryType) { + // Iceberg default value should always use HeapBufferBuffer, so calling ByteBuffer.array() + // should be safe. + return ((java.nio.ByteBuffer) value).array(); + } else { + return value; + } + } } diff --git a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestParquetScan.java b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestParquetScan.java index dbbb52794555..6b9ec85b7f0b 100644 --- a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestParquetScan.java +++ b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestParquetScan.java @@ -19,7 +19,6 @@ package org.apache.iceberg.spark.source; import static org.apache.iceberg.Files.localOutput; -import static org.apache.iceberg.spark.SparkSQLProperties.PARQUET_READER_TYPE; import static org.assertj.core.api.Assumptions.assumeThat; import java.io.File; @@ -36,25 +35,14 @@ import org.apache.iceberg.TableProperties; import org.apache.iceberg.io.FileAppender; import org.apache.iceberg.parquet.Parquet; -import org.apache.iceberg.spark.ParquetReaderType; import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; -import org.apache.spark.api.java.JavaSparkContext; -import org.apache.spark.sql.SparkSession; -import org.junit.jupiter.api.BeforeAll; public class TestParquetScan extends ScanTestBase { protected boolean vectorized() { return false; } - @BeforeAll - public static void startSpark() { - ScanTestBase.spark = SparkSession.builder().master("local[2]").getOrCreate(); - ScanTestBase.spark.conf().set(PARQUET_READER_TYPE, ParquetReaderType.ICEBERG.toString()); - ScanTestBase.sc = JavaSparkContext.fromSparkContext(spark.sparkContext()); - } - @Override protected void configureTable(Table table) { table From a6b15d3bbbe09b42a4f9694c5e49aedfd9b4aeb2 Mon Sep 17 00:00:00 2001 From: huaxingao Date: Tue, 28 Jan 2025 15:37:25 -0800 Subject: [PATCH 25/32] check type before casting --- .../spark/data/vectorized/CometConstantColumnReader.java | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java index 30d48a709c3e..bf0b594cd5b4 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java @@ -19,6 +19,7 @@ package org.apache.iceberg.spark.data.vectorized; import java.math.BigDecimal; +import java.nio.ByteBuffer; import org.apache.comet.parquet.ConstantColumnReader; import org.apache.iceberg.types.Types; import org.apache.spark.sql.types.DataType; @@ -46,11 +47,11 @@ public void setBatchSize(int batchSize) { private Object convertToSparkValue(T value) { DataType dataType = getSparkType(); - if (dataType == DataTypes.StringType) { + if (dataType == DataTypes.StringType && value instanceof String) { return UTF8String.fromString((String) value); - } else if (dataType instanceof DecimalType) { + } else if (dataType instanceof DecimalType && value instanceof BigDecimal) { return Decimal.apply((BigDecimal) value); - } else if (dataType == DataTypes.BinaryType) { + } else if (dataType == DataTypes.BinaryType && value instanceof ByteBuffer) { // Iceberg default value should always use HeapBufferBuffer, so calling ByteBuffer.array() // should be safe. return ((java.nio.ByteBuffer) value).array(); From 77775a3f136567e957cc410640cb0ae0062f85aa Mon Sep 17 00:00:00 2001 From: huaxingao Date: Tue, 28 Jan 2025 19:11:38 -0800 Subject: [PATCH 26/32] address comments --- .../iceberg/spark/ParquetReaderType.java | 32 +++++++++++++++++-- .../apache/iceberg/spark/SparkReadConf.java | 2 +- .../iceberg/spark/SparkSQLProperties.java | 2 +- .../data/vectorized/CometColumnReader.java | 13 ++++---- .../vectorized/CometColumnarBatchReader.java | 9 +++--- .../vectorized/CometConstantColumnReader.java | 5 ++- .../vectorized/CometPositionColumnReader.java | 2 +- 7 files changed, 45 insertions(+), 20 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java index b0b197377338..e2db3a39dc67 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java @@ -18,10 +18,12 @@ */ package org.apache.iceberg.spark; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; + /** Enumerates the types of Parquet readers. */ public enum ParquetReaderType { - /** ICEBERG type utilizes the Parquet reader from Apache Iceberg. */ - ICEBERG, + /** ICEBERG type utilizes the built-in Parquet reader. */ + ICEBERG("iceberg"), /** * COMET type changes the Parquet reader to the Apache DataFusion Comet Parquet reader. Comet @@ -32,5 +34,29 @@ public enum ParquetReaderType { *

TODO: Implement {@link org.apache.comet.parquet.SupportsComet} in SparkScan to convert Spark * physical plan to native physical plan for native execution. */ - COMET + COMET("comet"); + + private final String parquetReaderType; + + ParquetReaderType(String readerType) { + this.parquetReaderType = readerType; + } + + public static ParquetReaderType fromName(String parquetReaderType) { + Preconditions.checkArgument(parquetReaderType != null, "Parquet reader type is null"); + + if (ICEBERG.parquetReaderType().equalsIgnoreCase(parquetReaderType)) { + return ICEBERG; + + } else if (COMET.parquetReaderType().equalsIgnoreCase(parquetReaderType)) { + return COMET; + + } else { + throw new IllegalArgumentException("Unknown parquet reader type: " + parquetReaderType); + } + } + + public String parquetReaderType() { + return parquetReaderType; + } } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java index c547153e4d33..607ec4650e55 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java @@ -362,7 +362,7 @@ public boolean reportColumnStats() { public ParquetReaderType parquetReaderType() { return confParser - .enumConf(ParquetReaderType::valueOf) + .enumConf(ParquetReaderType::fromName) .sessionConf(SparkSQLProperties.PARQUET_READER_TYPE) .defaultValue(SparkSQLProperties.PARQUET_READER_TYPE_DEFAULT) .parse(); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java index 733744ade421..87daef4c0d5b 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java @@ -27,7 +27,7 @@ private SparkSQLProperties() {} // Controls whether vectorized reads are enabled public static final String VECTORIZATION_ENABLED = "spark.sql.iceberg.vectorization.enabled"; - // Controls which Parquet reader to use for vectorization + // Controls which Parquet reader implementation to use public static final String PARQUET_READER_TYPE = "spark.sql.iceberg.parquet.reader-type"; public static final ParquetReaderType PARQUET_READER_TYPE_DEFAULT = ParquetReaderType.COMET; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java index 231a6ed32a88..b7b23a8941d0 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java @@ -27,6 +27,7 @@ import org.apache.comet.shaded.arrow.c.CometSchemaImporter; import org.apache.comet.shaded.arrow.memory.RootAllocator; import org.apache.iceberg.parquet.VectorizedReader; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import org.apache.iceberg.spark.SparkSchemaUtil; import org.apache.iceberg.types.Types; import org.apache.parquet.column.ColumnDescriptor; @@ -38,7 +39,7 @@ import org.apache.spark.sql.types.Metadata; import org.apache.spark.sql.types.StructField; -@SuppressWarnings({"checkstyle:VisibilityModifier", "ParameterAssignment"}) +@SuppressWarnings("checkstyle:VisibilityModifier") class CometColumnReader implements VectorizedReader { public static final int DEFAULT_BATCH_SIZE = 5000; @@ -92,16 +93,16 @@ public CometVector read(CometVector reuse, int numRows) { return reuse; } - public ColumnDescriptor getDescriptor() { + public ColumnDescriptor descriptor() { return descriptor; } - public CometVector getVector() { + public CometVector vector() { return vector; } /** Returns the Spark data type for this column. */ - public DataType getSparkType() { + public DataType sparkType() { return sparkType; } @@ -112,9 +113,7 @@ public DataType getSparkType() { * CometColumnReader#reset} is called. */ public void setPageReader(PageReader pageReader) throws IOException { - if (!initialized) { - throw new IllegalStateException("Invalid state: 'reset' should be called first"); - } + Preconditions.checkState(initialized, "Invalid state: 'reset' should be called first"); ((ColumnReader) delegate).setPageReader(pageReader); } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index 42c4b7662e30..c472af80ed75 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -78,7 +78,7 @@ public void setRowGroupInfo( && !(readers[i] instanceof CometPositionColumnReader) && !(readers[i] instanceof CometDeleteColumnReader)) { readers[i].reset(); - readers[i].setPageReader(pageStore.getPageReader(readers[i].getDescriptor())); + readers[i].setPageReader(pageStore.getPageReader(readers[i].descriptor())); } } catch (IOException e) { throw new UncheckedIOException("Failed to setRowGroupInfo for Comet vectorization", e); @@ -139,6 +139,7 @@ private class ColumnBatchLoader { ColumnarBatch loadDataToColumnBatch() { ColumnVector[] vectors = readDataToColumnVectors(); int numLiveRows = batchSize; + if (hasIsDeletedColumn) { boolean[] isDeleted = buildIsDeleted(vectors); readDeletedColumn(vectors, isDeleted); @@ -175,7 +176,7 @@ ColumnVector[] readDataToColumnVectors() { // Fetch rows for all readers in the delegate delegate.nextBatch(batchSize); for (int i = 0; i < readers.length; i++) { - columnVectors[i] = readers[i].getVector(); + columnVectors[i] = readers[i].vector(); columnVectors[i].resetRowIdMapping(); org.apache.comet.vector.CometVector vector = readers[i].getDelegate().currentBatch(); columnVectors[i].setDelegate(vector); @@ -189,8 +190,8 @@ void readDeletedColumn(ColumnVector[] columnVectors, boolean[] isDeleted) { if (readers[i] instanceof CometDeleteColumnReader) { CometDeleteColumnReader deleteColumnReader = new CometDeleteColumnReader<>(isDeleted); deleteColumnReader.setBatchSize(batchSize); - deleteColumnReader.read(deleteColumnReader.getVector(), batchSize); - columnVectors[i] = deleteColumnReader.getVector(); + deleteColumnReader.read(deleteColumnReader.vector(), batchSize); + columnVectors[i] = deleteColumnReader.vector(); } } } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java index bf0b594cd5b4..0baafe33c9bd 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java @@ -34,8 +34,7 @@ class CometConstantColumnReader extends CometColumnReader { super(field); // use delegate to set constant value on the native side to be consumed by native execution. delegate = - new ConstantColumnReader( - getSparkType(), getDescriptor(), convertToSparkValue(value), false); + new ConstantColumnReader(sparkType(), descriptor(), convertToSparkValue(value), false); } @Override @@ -46,7 +45,7 @@ public void setBatchSize(int batchSize) { } private Object convertToSparkValue(T value) { - DataType dataType = getSparkType(); + DataType dataType = sparkType(); if (dataType == DataTypes.StringType && value instanceof String) { return UTF8String.fromString((String) value); } else if (dataType instanceof DecimalType && value instanceof BigDecimal) { diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java index 31ccc591a7ad..fc5109984ebe 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java @@ -27,7 +27,7 @@ class CometPositionColumnReader extends CometColumnReader { CometPositionColumnReader(Types.NestedField field) { super(field); - delegate = new PositionColumnReader(getDescriptor()); + delegate = new PositionColumnReader(descriptor()); } @Override From 4bf5cbfe2f09b13038dcb43c2ed53790a87cb27c Mon Sep 17 00:00:00 2001 From: huaxingao Date: Tue, 28 Jan 2025 22:10:53 -0800 Subject: [PATCH 27/32] address comments --- .../iceberg/spark/ParquetReaderType.java | 33 +++++-------------- .../apache/iceberg/spark/SparkReadConf.java | 2 +- .../data/vectorized/CometColumnReader.java | 25 ++++++++++---- .../vectorized/CometColumnarBatchReader.java | 10 +++--- .../vectorized/CometConstantColumnReader.java | 10 +++--- .../vectorized/CometDeleteColumnReader.java | 10 +++--- .../vectorized/CometPositionColumnReader.java | 8 ++--- .../spark/source/TestSparkReaderDeletes.java | 12 ++++--- 8 files changed, 54 insertions(+), 56 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java index e2db3a39dc67..d9742c048251 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ParquetReaderType.java @@ -23,7 +23,7 @@ /** Enumerates the types of Parquet readers. */ public enum ParquetReaderType { /** ICEBERG type utilizes the built-in Parquet reader. */ - ICEBERG("iceberg"), + ICEBERG, /** * COMET type changes the Parquet reader to the Apache DataFusion Comet Parquet reader. Comet @@ -34,29 +34,14 @@ public enum ParquetReaderType { *

TODO: Implement {@link org.apache.comet.parquet.SupportsComet} in SparkScan to convert Spark * physical plan to native physical plan for native execution. */ - COMET("comet"); - - private final String parquetReaderType; - - ParquetReaderType(String readerType) { - this.parquetReaderType = readerType; - } - - public static ParquetReaderType fromName(String parquetReaderType) { - Preconditions.checkArgument(parquetReaderType != null, "Parquet reader type is null"); - - if (ICEBERG.parquetReaderType().equalsIgnoreCase(parquetReaderType)) { - return ICEBERG; - - } else if (COMET.parquetReaderType().equalsIgnoreCase(parquetReaderType)) { - return COMET; - - } else { - throw new IllegalArgumentException("Unknown parquet reader type: " + parquetReaderType); + COMET; + + public static ParquetReaderType fromString(String typeAsString) { + Preconditions.checkArgument(typeAsString != null, "Parquet reader type is null"); + try { + return ParquetReaderType.valueOf(typeAsString.toUpperCase()); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("Unknown parquet reader type: " + typeAsString); } } - - public String parquetReaderType() { - return parquetReaderType; - } } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java index 607ec4650e55..5799799eaf16 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java @@ -362,7 +362,7 @@ public boolean reportColumnStats() { public ParquetReaderType parquetReaderType() { return confParser - .enumConf(ParquetReaderType::fromName) + .enumConf(ParquetReaderType::fromString) .sessionConf(SparkSQLProperties.PARQUET_READER_TYPE) .defaultValue(SparkSQLProperties.PARQUET_READER_TYPE_DEFAULT) .parse(); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java index b7b23a8941d0..d89731896b6e 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java @@ -39,17 +39,16 @@ import org.apache.spark.sql.types.Metadata; import org.apache.spark.sql.types.StructField; -@SuppressWarnings("checkstyle:VisibilityModifier") class CometColumnReader implements VectorizedReader { public static final int DEFAULT_BATCH_SIZE = 5000; private final DataType sparkType; // the delegated column reader from Comet side - protected AbstractColumnReader delegate; + private AbstractColumnReader delegate; private final CometVector vector; private final ColumnDescriptor descriptor; - protected boolean initialized = false; - protected int batchSize = DEFAULT_BATCH_SIZE; + private boolean initialized = false; + private int batchSize = DEFAULT_BATCH_SIZE; CometColumnReader(DataType sparkType, ColumnDescriptor descriptor) { this.sparkType = sparkType; @@ -65,10 +64,22 @@ class CometColumnReader implements VectorizedReader { this.vector = new CometVector(sparkType, false); } - public AbstractColumnReader getDelegate() { + public AbstractColumnReader delegate() { return delegate; } + void setDelegate(AbstractColumnReader delegate) { + this.delegate = delegate; + } + + void setInitialized(boolean initialized) { + this.initialized = initialized; + } + + public int batchSize() { + return batchSize; + } + /** * This method is to initialized/reset the CometColumnReader. This needs to be called for each row * group after readNextRowGroup, so a new dictionary encoding can be set for each of the new row @@ -81,8 +92,8 @@ public void reset() { CometSchemaImporter importer = new CometSchemaImporter(new RootAllocator()); - delegate = Utils.getColumnReader(sparkType, descriptor, importer, batchSize, false, false); - initialized = true; + this.delegate = Utils.getColumnReader(sparkType, descriptor, importer, batchSize, false, false); + this.initialized = true; } @Override diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index c472af80ed75..f766e5c728ea 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -47,10 +47,10 @@ class CometColumnarBatchReader implements VectorizedReader { private final CometColumnReader[] readers; private final boolean hasIsDeletedColumn; - private DeleteFilter deletes = null; - private long rowStartPosInBatch = 0; // The delegated batch reader on Comet side private final BatchReader delegate; + private DeleteFilter deletes = null; + private long rowStartPosInBatch = 0; CometColumnarBatchReader(List> readers, Schema schema) { this.readers = @@ -59,7 +59,7 @@ class CometColumnarBatchReader implements VectorizedReader { readers.stream().anyMatch(reader -> reader instanceof CometDeleteColumnReader); AbstractColumnReader[] abstractColumnReaders = new AbstractColumnReader[readers.size()]; - delegate = new BatchReader(abstractColumnReaders); + this.delegate = new BatchReader(abstractColumnReaders); delegate.setSparkSchema(SparkSchemaUtil.convert(schema)); } @@ -86,7 +86,7 @@ public void setRowGroupInfo( } for (int i = 0; i < readers.length; i++) { - delegate.getColumnReaders()[i] = this.readers[i].getDelegate(); + delegate.getColumnReaders()[i] = this.readers[i].delegate(); } this.rowStartPosInBatch = @@ -178,7 +178,7 @@ ColumnVector[] readDataToColumnVectors() { for (int i = 0; i < readers.length; i++) { columnVectors[i] = readers[i].vector(); columnVectors[i].resetRowIdMapping(); - org.apache.comet.vector.CometVector vector = readers[i].getDelegate().currentBatch(); + org.apache.comet.vector.CometVector vector = readers[i].delegate().currentBatch(); columnVectors[i].setDelegate(vector); } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java index 0baafe33c9bd..0316427bb070 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java @@ -33,15 +33,15 @@ class CometConstantColumnReader extends CometColumnReader { CometConstantColumnReader(T value, Types.NestedField field) { super(field); // use delegate to set constant value on the native side to be consumed by native execution. - delegate = - new ConstantColumnReader(sparkType(), descriptor(), convertToSparkValue(value), false); + setDelegate( + new ConstantColumnReader(sparkType(), descriptor(), convertToSparkValue(value), false)); } @Override public void setBatchSize(int batchSize) { - delegate.setBatchSize(batchSize); - this.batchSize = batchSize; - initialized = true; + delegate().setBatchSize(batchSize); + setBatchSize(batchSize); + setInitialized(true); } private Object convertToSparkValue(T value) { diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java index e40774613c2e..a2be29a4a954 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java @@ -30,19 +30,19 @@ class CometDeleteColumnReader extends CometColumnReader { CometDeleteColumnReader(Types.NestedField field) { super(field); - delegate = new DeleteColumnReader(); + setDelegate(new DeleteColumnReader()); } CometDeleteColumnReader(boolean[] isDeleted) { super(MetadataColumns.IS_DELETED); - delegate = new DeleteColumnReader(isDeleted); + setDelegate(new DeleteColumnReader(isDeleted)); } @Override public void setBatchSize(int batchSize) { - delegate.setBatchSize(batchSize); - this.batchSize = batchSize; - initialized = true; + delegate().setBatchSize(batchSize); + setBatchSize(batchSize); + setInitialized(true); } private static class DeleteColumnReader extends MetadataColumnReader { diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java index fc5109984ebe..0d6ee114cc4d 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java @@ -27,14 +27,14 @@ class CometPositionColumnReader extends CometColumnReader { CometPositionColumnReader(Types.NestedField field) { super(field); - delegate = new PositionColumnReader(descriptor()); + setDelegate(new PositionColumnReader(descriptor())); } @Override public void setBatchSize(int batchSize) { - delegate.setBatchSize(batchSize); - this.batchSize = batchSize; - initialized = true; + delegate().setBatchSize(batchSize); + setBatchSize(batchSize); + setInitialized(true); } private static class PositionColumnReader extends MetadataColumnReader { diff --git a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java index dda49b49465c..e5c8dd73922e 100644 --- a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java +++ b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java @@ -19,6 +19,7 @@ package org.apache.iceberg.spark.source; import static org.apache.hadoop.hive.conf.HiveConf.ConfVars.METASTOREURIS; +import static org.apache.iceberg.spark.SparkSQLProperties.PARQUET_READER_TYPE; import static org.apache.iceberg.spark.source.SparkSQLExecutionHelper.lastExecutedMetricValue; import static org.apache.iceberg.types.Types.NestedField.required; import static org.assertj.core.api.Assertions.assertThat; @@ -109,12 +110,12 @@ public class TestSparkReaderDeletes extends DeleteReadTests { @Parameters(name = "fileFormat = {0}, formatVersion = {1}, vectorized = {2}, planningMode = {3}") public static Object[][] parameters() { return new Object[][] { - new Object[] {FileFormat.PARQUET, 2, false, PlanningMode.DISTRIBUTED}, + // new Object[] {FileFormat.PARQUET, 2, false, PlanningMode.DISTRIBUTED}, new Object[] {FileFormat.PARQUET, 2, true, PlanningMode.LOCAL}, - new Object[] {FileFormat.ORC, 2, false, PlanningMode.DISTRIBUTED}, - new Object[] {FileFormat.AVRO, 2, false, PlanningMode.LOCAL}, - new Object[] {FileFormat.PARQUET, 3, false, PlanningMode.DISTRIBUTED}, - new Object[] {FileFormat.PARQUET, 3, true, PlanningMode.LOCAL}, + // new Object[] {FileFormat.ORC, 2, false, PlanningMode.DISTRIBUTED}, + // new Object[] {FileFormat.AVRO, 2, false, PlanningMode.LOCAL}, + // new Object[] {FileFormat.PARQUET, 3, false, PlanningMode.DISTRIBUTED}, + // new Object[] {FileFormat.PARQUET, 3, true, PlanningMode.LOCAL}, }; } @@ -131,6 +132,7 @@ public static void startMetastoreAndSpark() { .config("spark.ui.liveUpdate.period", 0) .config(SQLConf.PARTITION_OVERWRITE_MODE().key(), "dynamic") .config("spark.hadoop." + METASTOREURIS.varname, hiveConf.get(METASTOREURIS.varname)) + // .config(PARQUET_READER_TYPE, "iceberg") .enableHiveSupport() .getOrCreate(); From 8f34742e0fb4bacba49bb1678b56c9006feaffbe Mon Sep 17 00:00:00 2001 From: huaxingao Date: Tue, 28 Jan 2025 22:16:29 -0800 Subject: [PATCH 28/32] remove un-intended change in test --- .../iceberg/spark/source/TestSparkReaderDeletes.java | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java index e5c8dd73922e..dda49b49465c 100644 --- a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java +++ b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java @@ -19,7 +19,6 @@ package org.apache.iceberg.spark.source; import static org.apache.hadoop.hive.conf.HiveConf.ConfVars.METASTOREURIS; -import static org.apache.iceberg.spark.SparkSQLProperties.PARQUET_READER_TYPE; import static org.apache.iceberg.spark.source.SparkSQLExecutionHelper.lastExecutedMetricValue; import static org.apache.iceberg.types.Types.NestedField.required; import static org.assertj.core.api.Assertions.assertThat; @@ -110,12 +109,12 @@ public class TestSparkReaderDeletes extends DeleteReadTests { @Parameters(name = "fileFormat = {0}, formatVersion = {1}, vectorized = {2}, planningMode = {3}") public static Object[][] parameters() { return new Object[][] { - // new Object[] {FileFormat.PARQUET, 2, false, PlanningMode.DISTRIBUTED}, + new Object[] {FileFormat.PARQUET, 2, false, PlanningMode.DISTRIBUTED}, new Object[] {FileFormat.PARQUET, 2, true, PlanningMode.LOCAL}, - // new Object[] {FileFormat.ORC, 2, false, PlanningMode.DISTRIBUTED}, - // new Object[] {FileFormat.AVRO, 2, false, PlanningMode.LOCAL}, - // new Object[] {FileFormat.PARQUET, 3, false, PlanningMode.DISTRIBUTED}, - // new Object[] {FileFormat.PARQUET, 3, true, PlanningMode.LOCAL}, + new Object[] {FileFormat.ORC, 2, false, PlanningMode.DISTRIBUTED}, + new Object[] {FileFormat.AVRO, 2, false, PlanningMode.LOCAL}, + new Object[] {FileFormat.PARQUET, 3, false, PlanningMode.DISTRIBUTED}, + new Object[] {FileFormat.PARQUET, 3, true, PlanningMode.LOCAL}, }; } @@ -132,7 +131,6 @@ public static void startMetastoreAndSpark() { .config("spark.ui.liveUpdate.period", 0) .config(SQLConf.PARTITION_OVERWRITE_MODE().key(), "dynamic") .config("spark.hadoop." + METASTOREURIS.varname, hiveConf.get(METASTOREURIS.varname)) - // .config(PARQUET_READER_TYPE, "iceberg") .enableHiveSupport() .getOrCreate(); From 0d9e974f32d525754d8e1e796205e3bba752ea67 Mon Sep 17 00:00:00 2001 From: huaxingao Date: Wed, 29 Jan 2025 15:54:32 -0800 Subject: [PATCH 29/32] address comments --- .../data/vectorized/CometColumnReader.java | 26 ++-- .../vectorized/CometColumnarBatchReader.java | 14 +-- .../vectorized/CometConstantColumnReader.java | 2 +- .../vectorized/CometDeleteColumnReader.java | 2 +- .../vectorized/CometPositionColumnReader.java | 2 +- .../spark/data/vectorized/CometVector.java | 117 ------------------ 6 files changed, 18 insertions(+), 145 deletions(-) delete mode 100644 spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java index d89731896b6e..9bea276324b4 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java @@ -38,22 +38,22 @@ import org.apache.spark.sql.types.DataType; import org.apache.spark.sql.types.Metadata; import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.vectorized.ColumnVector; -class CometColumnReader implements VectorizedReader { +class CometColumnReader implements VectorizedReader { public static final int DEFAULT_BATCH_SIZE = 5000; + private final ColumnDescriptor descriptor; private final DataType sparkType; + // the delegated column reader from Comet side private AbstractColumnReader delegate; - private final CometVector vector; - private final ColumnDescriptor descriptor; private boolean initialized = false; private int batchSize = DEFAULT_BATCH_SIZE; CometColumnReader(DataType sparkType, ColumnDescriptor descriptor) { this.sparkType = sparkType; this.descriptor = descriptor; - this.vector = new CometVector(sparkType, false); } CometColumnReader(Types.NestedField field) { @@ -61,7 +61,6 @@ class CometColumnReader implements VectorizedReader { StructField structField = new StructField(field.name(), dataType, false, Metadata.empty()); this.sparkType = dataType; this.descriptor = TypeUtil.convertToParquet(structField); - this.vector = new CometVector(sparkType, false); } public AbstractColumnReader delegate() { @@ -96,22 +95,10 @@ public void reset() { this.initialized = true; } - @Override - public CometVector read(CometVector reuse, int numRows) { - delegate.readBatch(numRows); - org.apache.comet.vector.CometVector bv = delegate.currentBatch(); - reuse.setDelegate(bv); - return reuse; - } - public ColumnDescriptor descriptor() { return descriptor; } - public CometVector vector() { - return vector; - } - /** Returns the Spark data type for this column. */ public DataType sparkType() { return sparkType; @@ -146,4 +133,9 @@ public void setRowGroupInfo( PageReadStore pageReadStore, Map map, long size) { throw new UnsupportedOperationException("Not supported"); } + + @Override + public ColumnVector read(ColumnVector reuse, int numRowsToRead) { + throw new UnsupportedOperationException("Not supported"); + } } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index f766e5c728ea..de0e3ed40309 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -47,6 +47,7 @@ class CometColumnarBatchReader implements VectorizedReader { private final CometColumnReader[] readers; private final boolean hasIsDeletedColumn; + // The delegated batch reader on Comet side private final BatchReader delegate; private DeleteFilter deletes = null; @@ -149,7 +150,7 @@ ColumnarBatch loadDataToColumnBatch() { int[] rowIdMapping = pair.first(); numLiveRows = pair.second(); for (int i = 0; i < vectors.length; i++) { - ((CometVector) vectors[i]).setRowIdMapping(rowIdMapping); + vectors[i] = new ColumnVectorWithFilter(vectors[i], rowIdMapping); } } } @@ -172,14 +173,11 @@ private Pair buildRowIdMapping(ColumnVector[] vectors) { } ColumnVector[] readDataToColumnVectors() { - CometVector[] columnVectors = new CometVector[readers.length]; + ColumnVector[] columnVectors = new ColumnVector[readers.length]; // Fetch rows for all readers in the delegate delegate.nextBatch(batchSize); for (int i = 0; i < readers.length; i++) { - columnVectors[i] = readers[i].vector(); - columnVectors[i].resetRowIdMapping(); - org.apache.comet.vector.CometVector vector = readers[i].delegate().currentBatch(); - columnVectors[i].setDelegate(vector); + columnVectors[i] = readers[i].delegate().currentBatch(); } return columnVectors; @@ -190,8 +188,8 @@ void readDeletedColumn(ColumnVector[] columnVectors, boolean[] isDeleted) { if (readers[i] instanceof CometDeleteColumnReader) { CometDeleteColumnReader deleteColumnReader = new CometDeleteColumnReader<>(isDeleted); deleteColumnReader.setBatchSize(batchSize); - deleteColumnReader.read(deleteColumnReader.vector(), batchSize); - columnVectors[i] = deleteColumnReader.vector(); + deleteColumnReader.delegate().readBatch(batchSize); + columnVectors[i] = deleteColumnReader.delegate().currentBatch(); } } } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java index 0316427bb070..0b461e8e2f41 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java @@ -39,8 +39,8 @@ class CometConstantColumnReader extends CometColumnReader { @Override public void setBatchSize(int batchSize) { + super.setBatchSize(batchSize); delegate().setBatchSize(batchSize); - setBatchSize(batchSize); setInitialized(true); } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java index a2be29a4a954..d834d1f899a2 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometDeleteColumnReader.java @@ -40,8 +40,8 @@ class CometDeleteColumnReader extends CometColumnReader { @Override public void setBatchSize(int batchSize) { + super.setBatchSize(batchSize); delegate().setBatchSize(batchSize); - setBatchSize(batchSize); setInitialized(true); } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java index 0d6ee114cc4d..ac68352fd29a 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometPositionColumnReader.java @@ -32,8 +32,8 @@ class CometPositionColumnReader extends CometColumnReader { @Override public void setBatchSize(int batchSize) { + super.setBatchSize(batchSize); delegate().setBatchSize(batchSize); - setBatchSize(batchSize); setInitialized(true); } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java deleted file mode 100644 index ceed84fcbcdd..000000000000 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVector.java +++ /dev/null @@ -1,117 +0,0 @@ -/* - * 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.vectorized; - -import org.apache.comet.vector.CometDelegateVector; -import org.apache.spark.sql.types.DataType; -import org.apache.spark.sql.types.Decimal; -import org.apache.spark.unsafe.types.UTF8String; - -@SuppressWarnings("checkstyle:VisibilityModifier") -class CometVector extends CometDelegateVector { - - protected int[] rowIdMapping; - - CometVector(DataType type, boolean useDecimal128) { - super(type, useDecimal128); - } - - public void setRowIdMapping(int[] rowIdMapping) { - this.rowIdMapping = rowIdMapping; - } - - public void resetRowIdMapping() { - this.rowIdMapping = null; - } - - @Override - public boolean isNullAt(int rowId) { - return super.isNullAt(mapRowId(rowId)); - } - - @Override - public boolean getBoolean(int rowId) { - return super.getBoolean(mapRowId(rowId)); - } - - @Override - public byte getByte(int rowId) { - return super.getByte(mapRowId(rowId)); - } - - @Override - public short getShort(int rowId) { - return super.getShort(mapRowId(rowId)); - } - - @Override - public int getInt(int rowId) { - return super.getInt(mapRowId(rowId)); - } - - @Override - public long getLong(int rowId) { - return super.getLong(mapRowId(rowId)); - } - - @Override - public float getFloat(int rowId) { - return super.getFloat(mapRowId(rowId)); - } - - @Override - public double getDouble(int rowId) { - return super.getDouble(mapRowId(rowId)); - } - - @Override - public Decimal getDecimal(int rowId, int precision, int scale) { - if (isNullAt(rowId)) { - return null; - } - - return super.getDecimal(mapRowId(rowId), precision, scale); - } - - @Override - public UTF8String getUTF8String(int rowId) { - if (isNullAt(rowId)) { - return null; - } - - return super.getUTF8String(mapRowId(rowId)); - } - - @Override - public byte[] getBinary(int rowId) { - if (isNullAt(rowId)) { - return null; - } - - return super.getBinary(mapRowId(rowId)); - } - - private int mapRowId(int rowId) { - if (rowIdMapping != null) { - return rowIdMapping[rowId]; - } - - return rowId; - } -} From 46dd439d3c684e8f550d1d2d3fe527df8b5a1171 Mon Sep 17 00:00:00 2001 From: huaxingao Date: Wed, 29 Jan 2025 17:12:06 -0800 Subject: [PATCH 30/32] address comments --- .../spark/data/vectorized/CometColumnReader.java | 15 ++++++++++----- .../data/vectorized/CometColumnarBatchReader.java | 8 +++++++- .../vectorized/CometConstantColumnReader.java | 4 ++++ 3 files changed, 21 insertions(+), 6 deletions(-) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java index 9bea276324b4..54576aef14c9 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java @@ -41,15 +41,17 @@ import org.apache.spark.sql.vectorized.ColumnVector; class CometColumnReader implements VectorizedReader { - public static final int DEFAULT_BATCH_SIZE = 5000; + // use the Comet default batch size + public static final int DEFAULT_BATCH_SIZE = 8192; private final ColumnDescriptor descriptor; private final DataType sparkType; - // the delegated column reader from Comet side + // The delegated ColumnReader from Comet side private AbstractColumnReader delegate; private boolean initialized = false; private int batchSize = DEFAULT_BATCH_SIZE; + private CometSchemaImporter importer; CometColumnReader(DataType sparkType, ColumnDescriptor descriptor) { this.sparkType = sparkType; @@ -89,8 +91,7 @@ public void reset() { delegate.close(); } - CometSchemaImporter importer = new CometSchemaImporter(new RootAllocator()); - + this.importer = new CometSchemaImporter(new RootAllocator()); this.delegate = Utils.getColumnReader(sparkType, descriptor, importer, batchSize, false, false); this.initialized = true; } @@ -117,7 +118,11 @@ public void setPageReader(PageReader pageReader) throws IOException { @Override public void close() { - // close reader on native side + // close resources on native side + if (importer != null) { + importer.close(); + } + if (delegate != null) { delegate.close(); } diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java index de0e3ed40309..1440e5d1d3f7 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java @@ -48,7 +48,13 @@ class CometColumnarBatchReader implements VectorizedReader { private final CometColumnReader[] readers; private final boolean hasIsDeletedColumn; - // The delegated batch reader on Comet side + // The delegated BatchReader on the Comet side does the real work of loading a batch of rows. + // The Comet BatchReader contains an array of ColumnReader. There is no need to explicitly call + // ColumnReader.readBatch; instead, BatchReader.nextBatch will be called, which underneath calls + // ColumnReader.readBatch. The only exception is DeleteColumnReader, because at the time of + // calling BatchReader.nextBatch, the isDeleted value is not yet available, so + // DeleteColumnReader.readBatch must be called explicitly later, after the isDeleted value is + // available. private final BatchReader delegate; private DeleteFilter deletes = null; private long rowStartPosInBatch = 0; diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java index 0b461e8e2f41..c665002e8f66 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometConstantColumnReader.java @@ -46,11 +46,15 @@ public void setBatchSize(int batchSize) { private Object convertToSparkValue(T value) { DataType dataType = sparkType(); + // Match the value to Spark internal type if necessary if (dataType == DataTypes.StringType && value instanceof String) { + // the internal type for StringType is UTF8String return UTF8String.fromString((String) value); } else if (dataType instanceof DecimalType && value instanceof BigDecimal) { + // the internal type for DecimalType is Decimal return Decimal.apply((BigDecimal) value); } else if (dataType == DataTypes.BinaryType && value instanceof ByteBuffer) { + // the internal type for DecimalType is byte[] // Iceberg default value should always use HeapBufferBuffer, so calling ByteBuffer.array() // should be safe. return ((java.nio.ByteBuffer) value).array(); From 10901b08685f7221bdf081d0863d530dae72760c Mon Sep 17 00:00:00 2001 From: huaxingao Date: Wed, 29 Jan 2025 17:43:03 -0800 Subject: [PATCH 31/32] close importer in reset --- .../iceberg/spark/data/vectorized/CometColumnReader.java | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java index 54576aef14c9..4794863ab1bf 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnReader.java @@ -87,6 +87,10 @@ public int batchSize() { * groups. */ public void reset() { + if (importer != null) { + importer.close(); + } + if (delegate != null) { delegate.close(); } From dae79ad0343fd748f8887a7758ec23f56321d7b4 Mon Sep 17 00:00:00 2001 From: huaxingao Date: Fri, 31 Jan 2025 14:17:45 -0800 Subject: [PATCH 32/32] revert to iceberg reader --- .../java/org/apache/iceberg/spark/SmokeTest.java | 3 +-- .../java/org/apache/iceberg/spark/SparkSQLProperties.java | 2 +- .../apache/iceberg/spark/source/TestDataFrameWriterV2.java | 7 +++---- 3 files changed, 5 insertions(+), 7 deletions(-) diff --git a/spark/v3.4/spark-runtime/src/integration/java/org/apache/iceberg/spark/SmokeTest.java b/spark/v3.4/spark-runtime/src/integration/java/org/apache/iceberg/spark/SmokeTest.java index 03b8bd273fa6..252793c7b8a7 100644 --- a/spark/v3.4/spark-runtime/src/integration/java/org/apache/iceberg/spark/SmokeTest.java +++ b/spark/v3.4/spark-runtime/src/integration/java/org/apache/iceberg/spark/SmokeTest.java @@ -28,7 +28,6 @@ import org.apache.spark.sql.catalyst.analysis.NoSuchTableException; import org.junit.Assert; import org.junit.Before; -import org.junit.Ignore; import org.junit.Test; public class SmokeTest extends SparkExtensionsTestBase { @@ -45,7 +44,7 @@ public void dropTable() { // Run through our Doc's Getting Started Example // TODO Update doc example so that it can actually be run, modifications were required for this // test suite to run - @Ignore + @Test public void testGettingStarted() throws IOException { // Creating a table sql("CREATE TABLE %s (id bigint, data string) USING iceberg", tableName); diff --git a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java index 87daef4c0d5b..0ca12369bd64 100644 --- a/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java +++ b/spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java @@ -29,7 +29,7 @@ private SparkSQLProperties() {} // Controls which Parquet reader implementation to use public static final String PARQUET_READER_TYPE = "spark.sql.iceberg.parquet.reader-type"; - public static final ParquetReaderType PARQUET_READER_TYPE_DEFAULT = ParquetReaderType.COMET; + public static final ParquetReaderType PARQUET_READER_TYPE_DEFAULT = ParquetReaderType.ICEBERG; // Controls whether reading/writing timestamps without timezones is allowed @Deprecated diff --git a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestDataFrameWriterV2.java b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestDataFrameWriterV2.java index fe027dc90686..47a0e87b9398 100644 --- a/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestDataFrameWriterV2.java +++ b/spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestDataFrameWriterV2.java @@ -41,7 +41,6 @@ import org.junit.After; import org.junit.Assert; import org.junit.Before; -import org.junit.Ignore; import org.junit.Test; public class TestDataFrameWriterV2 extends SparkTestBaseWithCatalog { @@ -215,7 +214,7 @@ public void testWriteWithCaseSensitiveOption() throws NoSuchTableException, Pars Assert.assertEquals(4, fields.size()); } - @Ignore + @Test public void testMergeSchemaIgnoreCastingLongToInt() throws Exception { sql( "ALTER TABLE %s SET TBLPROPERTIES ('%s'='true')", @@ -255,7 +254,7 @@ public void testMergeSchemaIgnoreCastingLongToInt() throws Exception { assertThat(idField.type().typeId()).isEqualTo(Type.TypeID.LONG); } - @Ignore + @Test public void testMergeSchemaIgnoreCastingDoubleToFloat() throws Exception { removeTables(); sql("CREATE TABLE %s (id double, data string) USING iceberg", tableName); @@ -297,7 +296,7 @@ public void testMergeSchemaIgnoreCastingDoubleToFloat() throws Exception { assertThat(idField.type().typeId()).isEqualTo(Type.TypeID.DOUBLE); } - @Ignore + @Test public void testMergeSchemaIgnoreCastingDecimalToDecimalWithNarrowerPrecision() throws Exception { removeTables(); sql("CREATE TABLE %s (id decimal(6,2), data string) USING iceberg", tableName);