Spark test

pvary · pvary · commit 2efac10d7894 · 2025-02-11T09:17:55.000+01:00
diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnarBatchReader.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnarBatchReader.java
@@ -23,7 +23,7 @@
 import org.apache.iceberg.arrow.vectorized.BaseBatchReader;
 import org.apache.iceberg.arrow.vectorized.VectorizedArrowReader;
 import org.apache.iceberg.arrow.vectorized.VectorizedArrowReader.DeletedVectorReader;
-import org.apache.iceberg.data.DeleteFilter;
+import org.apache.iceberg.io.datafile.DeleteFilter;
 import org.apache.iceberg.parquet.VectorizedReader;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.util.Pair;
diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnarBatchUtil.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnarBatchUtil.java
@@ -20,8 +20,8 @@
 
 import java.util.Arrays;
 import java.util.function.Predicate;
-import org.apache.iceberg.data.DeleteFilter;
 import org.apache.iceberg.deletes.PositionDeleteIndex;
+import org.apache.iceberg.io.datafile.DeleteFilter;
 import org.apache.iceberg.util.Pair;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.vectorized.ColumnVector;
diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometColumnarBatchReader.java
@@ -25,7 +25,7 @@
 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.io.datafile.DeleteFilter;
 import org.apache.iceberg.parquet.VectorizedReader;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.spark.SparkSchemaUtil;
diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/CometVectorizedReaderBuilder.java
@@ -24,7 +24,7 @@
 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.io.datafile.DeleteFilter;
 import org.apache.iceberg.parquet.TypeWithSchemaVisitor;
 import org.apache.iceberg.parquet.VectorizedReader;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkParquetReaders.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkParquetReaders.java
@@ -24,7 +24,7 @@
 import org.apache.arrow.vector.NullCheckingForGet;
 import org.apache.iceberg.Schema;
 import org.apache.iceberg.arrow.vectorized.VectorizedReaderBuilder;
-import org.apache.iceberg.data.DeleteFilter;
+import org.apache.iceberg.io.datafile.DeleteFilter;
 import org.apache.iceberg.parquet.TypeWithSchemaVisitor;
 import org.apache.iceberg.parquet.VectorizedReader;
 import org.apache.iceberg.spark.SparkUtil;
diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java
@@ -19,25 +19,27 @@
 package org.apache.iceberg.spark.source;
 
 import java.util.Map;
-import java.util.Set;
+import org.apache.iceberg.ContentScanTask;
 import org.apache.iceberg.FileFormat;
-import org.apache.iceberg.MetadataColumns;
 import org.apache.iceberg.ScanTask;
 import org.apache.iceberg.ScanTaskGroup;
 import org.apache.iceberg.Schema;
 import org.apache.iceberg.Table;
-import org.apache.iceberg.expressions.Expression;
 import org.apache.iceberg.io.CloseableIterable;
 import org.apache.iceberg.io.InputFile;
+import org.apache.iceberg.io.datafile.DataFileServiceRegistry;
+import org.apache.iceberg.io.datafile.DeleteFilter;
+import org.apache.iceberg.io.datafile.ReaderBuilder;
+import org.apache.iceberg.io.datafile.ReaderService;
+import org.apache.iceberg.io.datafile.ServiceBase;
 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.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;
-import org.apache.iceberg.types.TypeUtil;
+import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.vectorized.ColumnarBatch;
 
 abstract class BaseBatchReader<T extends ScanTask> extends BaseReader<ColumnarBatch, T> {
@@ -58,83 +60,109 @@ abstract class BaseBatchReader<T extends ScanTask> extends BaseReader<ColumnarBa
   }
 
   protected CloseableIterable<ColumnarBatch> newBatchIterable(
-      InputFile inputFile,
-      FileFormat format,
-      long start,
-      long length,
-      Expression residual,
-      Map<Integer, ?> idToConstant,
-      SparkDeleteFilter deleteFilter) {
-    switch (format) {
-      case PARQUET:
-        return newParquetIterable(inputFile, start, length, residual, idToConstant, deleteFilter);
+      InputFile inputFile, ContentScanTask<?> task, Table table, SparkDeleteFilter deleteFilter) {
+    ReaderBuilder<?> readerBuilder =
+        DataFileServiceRegistry.read(
+                task.file().format(),
+                InternalRow.class.getName(),
+                parquetConf != null ? parquetConf.readerType().name() : null,
+                inputFile,
+                task,
+                expectedSchema(),
+                table,
+                deleteFilter)
+            .split(task.start(), task.length())
+            .filter(task.residual())
+            .caseSensitive(caseSensitive())
+            // Spark eagerly consumes the batches. So the underlying memory allocated could be
+            // reused
+            // without worrying about subsequent reads clobbering over each other. This improves
+            // read performance as every batch read doesn't have to pay the cost of allocating
+            // memory.
+            .reuseContainers()
+            .withNameMapping(nameMapping());
+    if (parquetConf != null) {
+      readerBuilder = readerBuilder.recordsPerBatch(parquetConf.batchSize());
+    } else if (orcConf != null) {
+      readerBuilder = readerBuilder.recordsPerBatch(orcConf.batchSize());
+    }
+
+    return readerBuilder.build();
+  }
 
-      case ORC:
-        return newOrcIterable(inputFile, start, length, residual, idToConstant);
+  public static class IcebergParquetReaderService extends ServiceBase implements ReaderService {
+    @SuppressWarnings("checkstyle:RedundantModifier")
+    public IcebergParquetReaderService() {
+      super(FileFormat.PARQUET, InternalRow.class.getName(), ParquetReaderType.ICEBERG.name());
+    }
 
-      default:
-        throw new UnsupportedOperationException(
-            "Format: " + format + " not supported for batched reads");
+    @Override
+    public ReaderBuilder<?> builder(
+        InputFile inputFile,
+        ContentScanTask<?> task,
+        Schema readSchema,
+        Table table,
+        DeleteFilter<?> deleteFilter) {
+      // get required schema if there are deletes
+      Schema requiredSchema = deleteFilter != null ? deleteFilter.requiredSchema() : readSchema;
+      return Parquet.read(inputFile)
+          .project(requiredSchema)
+          .createBatchedReaderFunc(
+              fileSchema ->
+                  VectorizedSparkParquetReaders.buildReader(
+                      requiredSchema,
+                      fileSchema,
+                      constantsMap(task, readSchema, table),
+                      (DeleteFilter<InternalRow>) deleteFilter));
     }
   }
 
-  private CloseableIterable<ColumnarBatch> newParquetIterable(
-      InputFile inputFile,
-      long start,
-      long length,
-      Expression residual,
-      Map<Integer, ?> idToConstant,
-      SparkDeleteFilter deleteFilter) {
-    // get required schema if there are deletes
-    Schema requiredSchema = deleteFilter != null ? deleteFilter.requiredSchema() : expectedSchema();
+  public static class CometParquetReaderService extends ServiceBase implements ReaderService {
+    @SuppressWarnings("checkstyle:RedundantModifier")
+    public CometParquetReaderService() {
+      super(FileFormat.PARQUET, InternalRow.class.getName(), ParquetReaderType.COMET.name());
+    }
 
-    return Parquet.read(inputFile)
-        .project(requiredSchema)
-        .split(start, length)
-        .createBatchedReaderFunc(
-            fileSchema -> {
-              if (parquetConf.readerType() == ParquetReaderType.COMET) {
-                return VectorizedSparkParquetReaders.buildCometReader(
-                    requiredSchema, fileSchema, idToConstant, deleteFilter);
-              } else {
-                return VectorizedSparkParquetReaders.buildReader(
-                    requiredSchema, fileSchema, idToConstant, deleteFilter);
-              }
-            })
-        .recordsPerBatch(parquetConf.batchSize())
-        .filter(residual)
-        .caseSensitive(caseSensitive())
-        // Spark eagerly consumes the batches. So the underlying memory allocated could be reused
-        // without worrying about subsequent reads clobbering over each other. This improves
-        // read performance as every batch read doesn't have to pay the cost of allocating memory.
-        .reuseContainers()
-        .withNameMapping(nameMapping())
-        .build();
+    @Override
+    public ReaderBuilder<?> builder(
+        InputFile inputFile,
+        ContentScanTask<?> task,
+        Schema readSchema,
+        Table table,
+        DeleteFilter<?> deleteFilter) {
+      // get required schema if there are deletes
+      Schema requiredSchema = deleteFilter != null ? deleteFilter.requiredSchema() : readSchema;
+      return Parquet.read(inputFile)
+          .project(requiredSchema)
+          .createBatchedReaderFunc(
+              fileSchema ->
+                  VectorizedSparkParquetReaders.buildCometReader(
+                      requiredSchema,
+                      fileSchema,
+                      constantsMap(task, readSchema, table),
+                      (DeleteFilter<InternalRow>) deleteFilter));
+    }
   }
 
-  private CloseableIterable<ColumnarBatch> newOrcIterable(
-      InputFile inputFile,
-      long start,
-      long length,
-      Expression residual,
-      Map<Integer, ?> idToConstant) {
-    Set<Integer> constantFieldIds = idToConstant.keySet();
-    Set<Integer> metadataFieldIds = MetadataColumns.metadataFieldIds();
-    Sets.SetView<Integer> constantAndMetadataFieldIds =
-        Sets.union(constantFieldIds, metadataFieldIds);
-    Schema schemaWithoutConstantAndMetadataFields =
-        TypeUtil.selectNot(expectedSchema(), constantAndMetadataFieldIds);
+  public static class ORCReaderService extends ServiceBase implements ReaderService {
+    @SuppressWarnings("checkstyle:RedundantModifier")
+    public ORCReaderService() {
+      super(FileFormat.ORC, InternalRow.class.getName());
+    }
 
-    return ORC.read(inputFile)
-        .project(schemaWithoutConstantAndMetadataFields)
-        .split(start, length)
-        .createBatchedReaderFunc(
-            fileSchema ->
-                VectorizedSparkOrcReaders.buildReader(expectedSchema(), fileSchema, idToConstant))
-        .recordsPerBatch(orcConf.batchSize())
-        .filter(residual)
-        .caseSensitive(caseSensitive())
-        .withNameMapping(nameMapping())
-        .build();
+    @Override
+    public ReaderBuilder<?> builder(
+        InputFile inputFile,
+        ContentScanTask<?> task,
+        Schema readSchema,
+        Table table,
+        DeleteFilter<?> deleteFilter) {
+      Map<Integer, ?> idToConstant = constantsMap(task, readSchema, table);
+      return ORC.read(inputFile)
+          .project(ORC.schemaWithoutConstantAndMetadataFields(readSchema, idToConstant))
+          .createBatchedReaderFunc(
+              fileSchema ->
+                  VectorizedSparkOrcReaders.buildReader(readSchema, fileSchema, idToConstant));
+    }
   }
 }
diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/BaseReader.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/BaseReader.java
@@ -182,8 +182,13 @@ private Map<String, InputFile> inputFiles() {
   }
 
   protected Map<Integer, ?> constantsMap(ContentScanTask<?> task, Schema readSchema) {
+    return constantsMap(task, readSchema, table);
+  }
+
+  protected static Map<Integer, ?> constantsMap(
+      ContentScanTask<?> task, Schema readSchema, Table tableToRead) {
     if (readSchema.findField(MetadataColumns.PARTITION_COLUMN_ID) != null) {
-      StructType partitionType = Partitioning.partitionType(table);
+      StructType partitionType = Partitioning.partitionType(tableToRead);
       return PartitionUtil.constantsMap(task, partitionType, SparkUtil::internalToSpark);
     } else {
       return PartitionUtil.constantsMap(task, SparkUtil::internalToSpark);
diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/BatchDataReader.java
@@ -18,7 +18,6 @@
  */
 package org.apache.iceberg.spark.source;
 
-import java.util.Map;
 import java.util.stream.Stream;
 import org.apache.iceberg.ContentFile;
 import org.apache.iceberg.FileScanTask;
@@ -95,8 +94,6 @@ protected CloseableIterator<ColumnarBatch> open(FileScanTask task) {
     // update the current file for Spark's filename() function
     InputFileBlockHolder.set(filePath, task.start(), task.length());
 
-    Map<Integer, ?> idToConstant = constantsMap(task, expectedSchema());
-
     InputFile inputFile = getInputFile(filePath);
     Preconditions.checkNotNull(inputFile, "Could not find InputFile associated with FileScanTask");
 
@@ -105,14 +102,6 @@ protected CloseableIterator<ColumnarBatch> open(FileScanTask task) {
             ? null
             : new SparkDeleteFilter(filePath, task.deletes(), counter(), false);
 
-    return newBatchIterable(
-            inputFile,
-            task.file().format(),
-            task.start(),
-            task.length(),
-            task.residual(),
-            idToConstant,
-            deleteFilter)
-        .iterator();
+    return newBatchIterable(inputFile, task, table(), deleteFilter).iterator();
   }
 }
diff --git a/spark/v3.5/spark/src/main/resources/META-INF/services/org.apache.iceberg.io.datafile.ReaderService b/spark/v3.5/spark/src/main/resources/META-INF/services/org.apache.iceberg.io.datafile.ReaderService
@@ -0,0 +1,22 @@
+#
+# 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.
+#
+
+org.apache.iceberg.spark.source.BaseBatchReader$IcebergParquetReaderService
+org.apache.iceberg.spark.source.BaseBatchReader$CometParquetReaderService
+org.apache.iceberg.spark.source.BaseBatchReader$ORCReaderService
diff --git a/spark/v3.5/spark/src/test/resources/META-INF/services/org.apache.iceberg.io.datafile.ReaderService b/spark/v3.5/spark/src/test/resources/META-INF/services/org.apache.iceberg.io.datafile.ReaderService
@@ -0,0 +1,22 @@
+#
+# 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.
+#
+
+org.apache.iceberg.spark.source.BaseBatchReader$IcebergParquetReaderService
+org.apache.iceberg.spark.source.BaseBatchReader$CometParquetReaderService
+org.apache.iceberg.spark.source.BaseBatchReader$ORCReaderService
