Move batch size to config

Author: pvary

arrow/src/main/java/org/apache/iceberg/arrow/vectorized/ArrowReader.java

@@ -123,8 +123,7 @@ public class ArrowReader extends CloseableGroup {
   private final EncryptionManager encryption;
   private final int batchSize;
   private final boolean reuseContainers;
-
-  public static final String ARROW_OBJECT_MODEL = "arrow";
+  private static final String ARROW_OBJECT_MODEL = "arrow";
 
   public static void register() {
     ObjectModelRegistry.registerReader(
@@ -344,7 +343,7 @@ CloseableIterator<ColumnarBatch> open(FileScanTask task) {
             ObjectModelRegistry.readBuilder(FileFormat.PARQUET, ARROW_OBJECT_MODEL, location)
                 .project(expectedSchema)
                 .split(task.start(), task.length())
-                .recordsPerBatch(batchSize)
+                .set(ReadBuilder.RECORDS_PER_BATCH_KEY, String.valueOf(batchSize))
                 .filter(task.residual(), filterCaseSensitive);
 
         if (reuseContainers) {

core/src/main/java/org/apache/iceberg/io/ReadBuilder.java

@@ -30,6 +30,9 @@
  * @param <R> type of the reader
  */
 public interface ReadBuilder<R extends ReadBuilder<R>> {
+  /** The key for the batch size in case of vectorized reads. */
+  String RECORDS_PER_BATCH_KEY = "iceberg.records-per-batch";
+
   /**
    * Restricts the read to the given range: [start, start + length).
    *
@@ -90,11 +93,6 @@ default R reuseContainers(boolean newReuseContainers) {
     return (R) this;
   }
 
-  /** Sets the batch size for vectorized readers. */
-  default R recordsPerBatch(int numRowsPerBatch) {
-    throw new UnsupportedOperationException("Not supported");
-  }
-
   /**
    * Accessors for constant field values. Used for calculating values in the result which are coming
    * from metadata, and not coming from the data files themselves. The keys of the map are the

orc/src/main/java/org/apache/iceberg/orc/ORC.java

@@ -927,9 +927,10 @@ public ReadBuilder constantFieldAccessors(Map<Integer, ?> newConstantFieldAccess
       return this;
     }
 
-    @Override
+    @Deprecated
     public ReadBuilder recordsPerBatch(int numRecordsPerBatch) {
       this.recordsPerBatch = numRecordsPerBatch;
+      set(RECORDS_PER_BATCH_KEY, String.valueOf(numRecordsPerBatch));
       return this;
     }
 
@@ -969,7 +970,7 @@ public <D> CloseableIterable<D> build() {
           filterCaseSensitive,
           filter,
           batchReader,
-          recordsPerBatch);
+          conf.getInt(RECORDS_PER_BATCH_KEY, recordsPerBatch));
     }
   }
 

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

@@ -1429,9 +1429,10 @@ public B reuseContainers(boolean newReuseContainers) {
       return (B) this;
     }
 
-    @Override
+    @Deprecated
     public B recordsPerBatch(int numRowsPerBatch) {
       this.maxRecordsPerBatch = numRowsPerBatch;
+      set(RECORDS_PER_BATCH_KEY, String.valueOf(numRowsPerBatch));
       return (B) this;
     }
 
@@ -1593,7 +1594,9 @@ private <D> CloseableIterable<D> buildFunctionBasedReader(ParquetReadOptions opt
             filter,
             reuseContainers,
             filterCaseSensitive,
-            maxRecordsPerBatch);
+            properties.containsKey(RECORDS_PER_BATCH_KEY)
+                ? Integer.parseInt(properties.get(RECORDS_PER_BATCH_KEY))
+                : maxRecordsPerBatch);
       } else {
         Function<MessageType, ParquetValueReader<?>> readBuilder =
             readerFuncWithSchema != null

spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/BaseBatchReader.java

@@ -78,12 +78,13 @@ protected CloseableIterable<ColumnarBatch> newBatchIterable(
     if (parquetConf != null) {
       readBuilder =
           readBuilder
-              .recordsPerBatch(parquetConf.batchSize())
+              .set(ReadBuilder.RECORDS_PER_BATCH_KEY, String.valueOf(parquetConf.batchSize()))
               .set(
                   VectorizedSparkParquetReaders.PARQUET_READER_TYPE,
                   parquetConf.readerType().name());
     } else if (orcConf != null) {
-      readBuilder = readBuilder.recordsPerBatch(orcConf.batchSize());
+      readBuilder =
+          readBuilder.set(ReadBuilder.RECORDS_PER_BATCH_KEY, String.valueOf(orcConf.batchSize()));
     }
 
     if (readBuilder instanceof Parquet.SupportsDeleteFilter<?>) {