Spark: Encapsulate parquet objects for Comet

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

@@ -127,6 +127,7 @@
 
 public class Parquet {
   private static final Logger LOG = LoggerFactory.getLogger(Parquet.class);
+  private static final String VECTORIZED_READER_FACTORY = "read.parquet.vectorized-reader.factory";
 
   private Parquet() {}
 
@@ -1378,6 +1379,22 @@ public ReadBuilder setCustomType(int fieldId, Class<? extends StructLike> struct
       return this;
     }
 
+    /**
+     * Sets the vectorized reader factory class to use for reading Parquet files.
+     *
+     * @param factoryClassName fully qualified class name of the VectorizedParquetReaderFactory
+     *     implementation, or null to use the default reader
+     * @return this builder for method chaining
+     */
+    public ReadBuilder vectorizedReaderFactory(String factoryClassName) {
+      if (factoryClassName != null) {
+        this.properties.put(VECTORIZED_READER_FACTORY, factoryClassName);
+      } else {
+        this.properties.remove(VECTORIZED_READER_FACTORY);
+      }
+      return this;
+    }
+
     public ReadBuilder withFileEncryptionKey(ByteBuffer encryptionKey) {
       this.fileEncryptionKey = encryptionKey;
       return this;
@@ -1389,7 +1406,7 @@ public ReadBuilder withAADPrefix(ByteBuffer aadPrefix) {
     }
 
     @Override
-    @SuppressWarnings({"unchecked", "checkstyle:CyclomaticComplexity"})
+    @SuppressWarnings({"unchecked", "checkstyle:CyclomaticComplexity", "MethodLength"})
     public <D> CloseableIterable<D> build() {
       FileDecryptionProperties fileDecryptionProperties = null;
       if (fileEncryptionKey != null) {
@@ -1442,6 +1459,29 @@ public <D> CloseableIterable<D> build() {
         }
 
         if (batchedReaderFunc != null) {
+          // Try to load custom vectorized reader factory from properties
+          String factoryName = properties.get(VECTORIZED_READER_FACTORY);
+
+          if (factoryName != null) {
+            LOG.info("Loading custom vectorized reader factory: {}", factoryName);
+            VectorizedParquetReaderFactory factory = loadReaderFactory(factoryName);
+            if (factory != null) {
+              return factory.createReader(
+                  VectorizedParquetReaderFactory.ReaderParams.builder(
+                          file, schema, options, batchedReaderFunc)
+                      .nameMapping(mapping)
+                      .filter(filter)
+                      .reuseContainers(reuseContainers)
+                      .caseSensitive(caseSensitive)
+                      .maxRecordsPerBatch(maxRecordsPerBatch)
+                      .properties(properties)
+                      .split(start, length)
+                      .encryption(fileEncryptionKey, fileAADPrefix)
+                      .build());
+            }
+          }
+
+          // Fall back to default VectorizedParquetReader
           return new VectorizedParquetReader<>(
               file,
               schema,
@@ -1537,6 +1577,21 @@ public <D> CloseableIterable<D> build() {
     }
   }
 
+  private static VectorizedParquetReaderFactory loadReaderFactory(String className) {
+    try {
+      Class<?> factoryClass = Class.forName(className);
+      if (VectorizedParquetReaderFactory.class.isAssignableFrom(factoryClass)) {
+        return (VectorizedParquetReaderFactory) factoryClass.getDeclaredConstructor().newInstance();
+      } else {
+        LOG.warn("Class {} does not implement VectorizedParquetReaderFactory interface", className);
+        return null;
+      }
+    } catch (Exception e) {
+      LOG.warn("Failed to instantiate vectorized reader factory: {}", className, e);
+      return null;
+    }
+  }
+
   private static class ParquetReadBuilder<T> extends ParquetReader.Builder<T> {
     private Schema schema = null;
     private ReadSupport<T> readSupport = null;

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

@@ -43,7 +43,7 @@
  *
  * @param <T> type of value to read
  */
-class ReadConf<T> {
+public class ReadConf<T> {
   private final ParquetFileReader reader;
   private final InputFile file;
   private final ParquetReadOptions options;
@@ -60,7 +60,7 @@ class ReadConf<T> {
   private final List<Map<ColumnPath, ColumnChunkMetaData>> columnChunkMetaDataForRowGroups;
 
   @SuppressWarnings("unchecked")
-  ReadConf(
+  public ReadConf(
       InputFile file,
       ParquetReadOptions options,
       Schema expectedSchema,
@@ -146,7 +146,7 @@ private ReadConf(ReadConf<T> toCopy) {
     this.columnChunkMetaDataForRowGroups = toCopy.columnChunkMetaDataForRowGroups;
   }
 
-  ParquetFileReader reader() {
+  public ParquetFileReader reader() {
     if (reader != null) {
       reader.setRequestedSchema(projection);
       return reader;
@@ -157,35 +157,43 @@ ParquetFileReader reader() {
     return newReader;
   }
 
-  ParquetValueReader<T> model() {
+  public InputFile file() {
+    return file;
+  }
+
+  public MessageType projection() {
+    return projection;
+  }
+
+  public ParquetValueReader<T> model() {
     return model;
   }
 
-  VectorizedReader<T> vectorizedModel() {
+  public VectorizedReader<T> vectorizedModel() {
     return vectorizedModel;
   }
 
-  boolean[] shouldSkip() {
+  public boolean[] shouldSkip() {
     return shouldSkip;
   }
 
-  long totalValues() {
+  public long totalValues() {
     return totalValues;
   }
 
-  boolean reuseContainers() {
+  public boolean reuseContainers() {
     return reuseContainers;
   }
 
-  Integer batchSize() {
+  public Integer batchSize() {
     return batchSize;
   }
 
-  List<Map<ColumnPath, ColumnChunkMetaData>> columnChunkMetadataForRowGroups() {
+  public List<Map<ColumnPath, ColumnChunkMetaData>> columnChunkMetadataForRowGroups() {
     return columnChunkMetaDataForRowGroups;
   }
 
-  ReadConf<T> copy() {
+  public ReadConf<T> copy() {
     return new ReadConf<>(this);
   }