Using Parquet.icebergSchema instead of keeping track of it in the ParquetFormatModel.ReadBuilderWrapper

Author: pvary

parquet/src/main/java/org/apache/iceberg/parquet/Parquet.java

@@ -1174,7 +1174,7 @@ public static class ReadBuilder implements InternalData.ReadBuilder {
     private Schema schema = null;
     private Expression filter = null;
     private ReadSupport<?> readSupport = null;
-    private Function<MessageType, VectorizedReader<?>> batchedReaderFunc = null;
+    private BatchReaderFunction batchedReaderFunc = null;
     private ReaderFunction readerFunction = null;
     private boolean filterRecords = true;
     private boolean caseSensitive = true;
@@ -1240,6 +1240,50 @@ public ReaderFunction withSchema(Schema expectedSchema) {
       }
     }
 
+    public interface BatchReaderFunction {
+      Function<MessageType, VectorizedReader<?>> apply();
+
+      default BatchReaderFunction withSchema(Schema schema) {
+        return this;
+      }
+    }
+
+    private static class UnaryBatchReaderFunction implements BatchReaderFunction {
+      private final Function<MessageType, VectorizedReader<?>> readerFunc;
+
+      UnaryBatchReaderFunction(Function<MessageType, VectorizedReader<?>> readerFunc) {
+        this.readerFunc = readerFunc;
+      }
+
+      @Override
+      public Function<MessageType, VectorizedReader<?>> apply() {
+        return readerFunc;
+      }
+    }
+
+    private static class BinaryBatchReaderFunction implements BatchReaderFunction {
+      private final BiFunction<Schema, MessageType, VectorizedReader<?>> readerFuncWithSchema;
+      private Schema schema;
+
+      BinaryBatchReaderFunction(
+          BiFunction<Schema, MessageType, VectorizedReader<?>> readerFuncWithSchema) {
+        this.readerFuncWithSchema = readerFuncWithSchema;
+      }
+
+      @Override
+      public Function<MessageType, VectorizedReader<?>> apply() {
+        Preconditions.checkArgument(
+            schema != null, "Schema must be set for 2-argument reader function");
+        return messageType -> readerFuncWithSchema.apply(schema, messageType);
+      }
+
+      @Override
+      public BinaryBatchReaderFunction withSchema(Schema expectedSchema) {
+        this.schema = expectedSchema;
+        return this;
+      }
+    }
+
     private ReadBuilder(InputFile file) {
       this.file = file;
     }
@@ -1314,14 +1358,27 @@ public ReadBuilder createReaderFunc(
       return this;
     }
 
-    public ReadBuilder createBatchedReaderFunc(Function<MessageType, VectorizedReader<?>> func) {
+    public ReadBuilder createBatchedReaderFunc(
+        Function<MessageType, VectorizedReader<?>> newReaderFunction) {
+      Preconditions.checkArgument(
+          this.batchedReaderFunc == null,
+          "Cannot set batched reader function: batched reader function already set");
+      Preconditions.checkArgument(
+          this.readerFunction == null,
+          "Cannot set batched reader function: ReaderFunction already set");
+      this.batchedReaderFunc = new UnaryBatchReaderFunction(newReaderFunction);
+      return this;
+    }
+
+    public ReadBuilder createBatchedReaderFunc(
+        BiFunction<Schema, MessageType, VectorizedReader<?>> newReaderFunction) {
       Preconditions.checkArgument(
           this.batchedReaderFunc == null,
           "Cannot set batched reader function: batched reader function already set");
       Preconditions.checkArgument(
           this.readerFunction == null,
           "Cannot set batched reader function: ReaderFunction already set");
-      this.batchedReaderFunc = func;
+      this.batchedReaderFunc = new BinaryBatchReaderFunction(newReaderFunction);
       return this;
     }
 
@@ -1441,11 +1498,13 @@ public <D> CloseableIterable<D> build() {
         }
 
         if (batchedReaderFunc != null) {
+          Function<MessageType, VectorizedReader<?>> readBuilder =
+              batchedReaderFunc.withSchema(schema).apply();
           return new VectorizedParquetReader<>(
               file,
               schema,
               options,
-              batchedReaderFunc,
+              readBuilder,
               mapping,
               filter,
               reuseContainers,

parquet/src/main/java/org/apache/iceberg/parquet/ParquetFormatModel.java

@@ -245,7 +245,6 @@ private static class ReadBuilderWrapper<D, S> implements ReadBuilder<D, S> {
     private final ReaderFunction<D> readerFunction;
     private final BatchReaderFunction<D> batchReaderFunction;
     private final Map<String, String> config = Maps.newHashMap();
-    private Schema icebergSchema;
     private Map<Integer, ?> idToConstant = ImmutableMap.of();
 
     private ReadBuilderWrapper(
@@ -265,7 +264,6 @@ public ReadBuilder<D, S> split(long newStart, long newLength) {
 
     @Override
     public ReadBuilder<D, S> project(Schema schema) {
-      this.icebergSchema = schema;
       internal.project(schema);
       return this;
     }
@@ -318,12 +316,13 @@ public CloseableIterable<D> build() {
       if (readerFunction != null) {
         return internal
             .createReaderFunc(
-                messageType -> readerFunction.read(icebergSchema, messageType, idToConstant))
+                (icebergSchema, messageType) ->
+                    readerFunction.read(icebergSchema, messageType, idToConstant))
             .build();
       } else if (batchReaderFunction != null) {
         return internal
             .createBatchedReaderFunc(
-                messageType ->
+                (icebergSchema, messageType) ->
                     batchReaderFunction.read(icebergSchema, messageType, idToConstant, config))
             .build();
       } else {