commit 472eebb09d035e5073c384c319ba549f74d79691
author belugabehr <12578579+belugabehr@users.noreply.github.com> Thu Oct 22 10:57:05 2020 -0400
committer GitHub <noreply@github.com> Thu Oct 22 16:57:05 2020 +0200
tree 2eed4ee233cb4fa50d6fae05535f8037fd38f2e6
parent 14b7a8be8ac98c7f96fc3fd2d5c21cbddf736b64

PARQUET-1914: Allow ProtoParquetReader To Support InputFile (#817)

parquet-protobuf/src/main/java/org/apache/parquet/proto/ProtoParquetReader.java

diff --git a/parquet-protobuf/src/main/java/org/apache/parquet/proto/ProtoParquetReader.java b/parquet-protobuf/src/main/java/org/apache/parquet/proto/ProtoParquetReader.java
index 73ddec2..a864474 100644
--- a/parquet-protobuf/src/main/java/org/apache/parquet/proto/ProtoParquetReader.java
+++ b/parquet-protobuf/src/main/java/org/apache/parquet/proto/ProtoParquetReader.java
@@ -20,21 +20,26 @@
 
 import java.io.IOException;
 
-import com.google.protobuf.MessageOrBuilder;
-
 import org.apache.hadoop.fs.Path;
-
 import org.apache.parquet.filter.UnboundRecordFilter;
 import org.apache.parquet.hadoop.ParquetReader;
+import org.apache.parquet.hadoop.api.ReadSupport;
+import org.apache.parquet.io.InputFile;
+
+import com.google.protobuf.MessageOrBuilder;
 
 /**
  * Read Protobuf records from a Parquet file.
  */
-public class ProtoParquetReader<T extends MessageOrBuilder> extends ParquetReader<T> {
+public class ProtoParquetReader<T extends MessageOrBuilder>
+    extends ParquetReader<T> {
 
-  @SuppressWarnings("unchecked")
-  public static <T> Builder<T> builder(Path file) {
-    return ParquetReader.builder(new ProtoReadSupport(), file);
+  public static <T> ParquetReader.Builder<T> builder(Path file) {
+    return new ProtoParquetReader.Builder<T>(file);
+  }
+
+  public static <T> ParquetReader.Builder<T> builder(InputFile file) {
+    return new ProtoParquetReader.Builder<T>(file);
   }
 
   /**
@@ -59,4 +64,21 @@
   public ProtoParquetReader(Path file, UnboundRecordFilter recordFilter) throws IOException {
     super(file, new ProtoReadSupport(), recordFilter);
   }
+
+  private static class Builder<T> extends ParquetReader.Builder<T> {
+
+    protected Builder(InputFile file) {
+      super(file);
+    }
+
+    protected Builder(Path path) {
+      super(path);
+    }
+
+    @SuppressWarnings({ "rawtypes", "unchecked" })
+    @Override
+    protected ReadSupport<T> getReadSupport() {
+      return new ProtoReadSupport();
+    }
+  }
 }

parquet-protobuf/src/test/java/org/apache/parquet/proto/TestUtils.java

diff --git a/parquet-protobuf/src/test/java/org/apache/parquet/proto/TestUtils.java b/parquet-protobuf/src/test/java/org/apache/parquet/proto/TestUtils.java
index 2fbd2a4..38256f4 100644
--- a/parquet-protobuf/src/test/java/org/apache/parquet/proto/TestUtils.java
+++ b/parquet-protobuf/src/test/java/org/apache/parquet/proto/TestUtils.java
@@ -22,8 +22,12 @@
 import com.google.protobuf.Message;
 import com.google.protobuf.MessageOrBuilder;
 import com.twitter.elephantbird.util.Protobufs;
+
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.Path;
 import org.apache.parquet.hadoop.ParquetReader;
+import org.apache.parquet.hadoop.util.HadoopInputFile;
+import org.apache.parquet.io.InputFile;
 
 import java.io.File;
 import java.io.IOException;
@@ -170,7 +174,8 @@
    * @return List of protobuf messages for the given type.
    */
   public static <T extends MessageOrBuilder> List<T> readMessages(Path file, Class<T> messageClass) throws IOException {
-    ParquetReader.Builder readerBuilder = ProtoParquetReader.builder(file);
+    InputFile inputFile = HadoopInputFile.fromPath(file, new Configuration());
+    ParquetReader.Builder readerBuilder = ProtoParquetReader.builder(inputFile);
     if (messageClass != null) {
       readerBuilder.set(ProtoReadSupport.PB_CLASS, messageClass.getName()).build();
     }