Filtered Reader Implementation, Avro Specific Support

parquet-avro/pom.xml

@@ -16,6 +16,7 @@
   <url>https://github.com/Parquet/parquet-mr</url>
 
   <properties>
+    <avro.version>1.7.4</avro.version>
   </properties>
 
   <dependencies>
@@ -37,7 +38,7 @@
     <dependency>
       <groupId>org.apache.avro</groupId>
       <artifactId>avro</artifactId>
-      <version>1.7.3</version>
+      <version>${avro.version}</version>
     </dependency>
     <dependency>
       <groupId>org.apache.hadoop</groupId>
@@ -107,6 +108,44 @@
           </execution>
         </executions>
       </plugin>
+      <plugin>
+        <groupId>org.apache.avro</groupId>
+        <artifactId>avro-maven-plugin</artifactId>
+        <version>${avro.version}</version>
+        <executions>
+            <execution>
+                <phase>generate-test-sources</phase>
+                <goals>
+                    <goal>idl-protocol</goal>
+                </goals>
+                <configuration>
+                    <sourceDirectory>${project.basedir}/src/test/resources</sourceDirectory>
+                    <outputDirectory>${project.build.directory}/generated-test-sources</outputDirectory>
+                    <stringType>String</stringType>
+                </configuration>
+            </execution>
+        </executions>
+      </plugin>
+      <plugin>
+        <!-- Ensure that the specific classes are available during test compile but not included in jar -->
+        <groupId>org.codehaus.mojo</groupId>
+        <artifactId>build-helper-maven-plugin</artifactId>
+        <version>1.8</version>
+        <executions>
+          <execution>
+            <id>add-test-sources</id>
+            <phase>generate-test-sources</phase>
+            <goals>
+              <goal>add-test-source</goal>
+            </goals>
+            <configuration>
+              <sources>
+                <source>${project.build.directory}/generated-test-sources</source>
+              </sources>
+            </configuration>
+          </execution>
+        </executions>
+      </plugin>
     </plugins>
   </build>
 

parquet-avro/src/main/java/parquet/avro/AvroGenericRecordConverter.java → parquet-avro/src/main/java/parquet/avro/AvroIndexedRecordConverter.java

@@ -22,6 +22,9 @@
 import org.apache.avro.generic.GenericArray;
 import org.apache.avro.generic.GenericData;
 import org.apache.avro.generic.GenericRecord;
+import org.apache.avro.generic.IndexedRecord;
+import org.apache.avro.specific.SpecificData;
+import parquet.Preconditions;
 import parquet.io.api.Binary;
 import parquet.io.api.Converter;
 import parquet.io.api.GroupConverter;
@@ -30,26 +33,28 @@
 import parquet.schema.MessageType;
 import parquet.schema.Type;
 
-class AvroGenericRecordConverter extends GroupConverter {
+class AvroIndexedRecordConverter extends GroupConverter {
 
   private final ParentValueContainer parent;
-  protected GenericData.Record currentRecord;
+  protected IndexedRecord currentRecord;
   private final Converter[] converters;
 
   private final GroupType parquetSchema;
   private final Schema avroSchema;
+  private final Class<? extends IndexedRecord> specificClass;
 
-  public AvroGenericRecordConverter(MessageType parquetSchema, Schema avroSchema) {
+  public AvroIndexedRecordConverter(MessageType parquetSchema, Schema avroSchema) {
     this(null, parquetSchema, avroSchema);
   }
 
-  public AvroGenericRecordConverter(ParentValueContainer parent, GroupType
+  public AvroIndexedRecordConverter(ParentValueContainer parent, GroupType
       parquetSchema, Schema avroSchema) {
     this.parent = parent;
     this.parquetSchema = parquetSchema;
     this.avroSchema = avroSchema;
     int schemaSize = parquetSchema.getFieldCount();
     this.converters = new Converter[schemaSize];
+    this.specificClass = SpecificData.get().getClass(avroSchema);
     int index = 0; // parquet ignores Avro nulls, so index may differ
     for (int avroIndex = 0; avroIndex < avroSchema.getFields().size(); avroIndex++) {
       Schema.Field field = avroSchema.getFields().get(avroIndex);
@@ -62,7 +67,7 @@ public AvroGenericRecordConverter(ParentValueContainer parent, GroupType
       converters[index] = newConverter(fieldSchema, type, new ParentValueContainer() {
         @Override
         void add(Object value) {
-          AvroGenericRecordConverter.this.set(finalAvroIndex, value);
+          AvroIndexedRecordConverter.this.set(finalAvroIndex, value);
         }
       });
       index++;
@@ -86,13 +91,15 @@ private static Converter newConverter(Schema schema, Type type,
     } else if (schema.getType().equals(Schema.Type.STRING)) {
       return new FieldStringConverter(parent);
     } else if (schema.getType().equals(Schema.Type.RECORD)) {
-      return new AvroGenericRecordConverter(parent, type.asGroupType(), schema);
+      return new AvroIndexedRecordConverter(parent, type.asGroupType(), schema);
     } else if (schema.getType().equals(Schema.Type.ENUM)) {
-      return new FieldStringConverter(parent);
+      return new FieldEnumConverter(parent,schema);
     } else if (schema.getType().equals(Schema.Type.ARRAY)) {
-      return new GenericArrayConverter(parent, type, schema);
+      return new AvroArrayConverter(parent, type, schema);
     } else if (schema.getType().equals(Schema.Type.MAP)) {
       return new MapConverter(parent, type, schema);
+    } else if (schema.getType().equals(Schema.Type.UNION)) {
+      return new AvroUnionConverter(parent, type, schema);
     } else if (schema.getType().equals(Schema.Type.FIXED)) {
       return new FieldFixedConverter(parent, schema);
     }
@@ -111,7 +118,10 @@ public Converter getConverter(int fieldIndex) {
 
   @Override
   public void start() {
-    this.currentRecord = new GenericData.Record(avroSchema);
+    // Should do the right thing whether it is generic or specific
+    this.currentRecord = (this.specificClass == null) ?
+        new GenericData.Record(avroSchema) :
+        (IndexedRecord)SpecificData.newInstance(specificClass, avroSchema);
   }
 
   @Override
@@ -121,7 +131,7 @@ public void end() {
     }
   }
 
-  GenericRecord getCurrentRecord() {
+  IndexedRecord getCurrentRecord() {
     return currentRecord;
   }
 
@@ -239,6 +249,26 @@ final public void addBinary(Binary value) {
 
   }
 
+  static final class FieldEnumConverter extends PrimitiveConverter {
+
+    private final ParentValueContainer parent;
+    private final Class<? extends Enum> enumClass;
+
+    public FieldEnumConverter(ParentValueContainer parent, Schema enumSchema) {
+      this.parent = parent;
+      this.enumClass = SpecificData.get().getClass(enumSchema);
+    }
+
+    @Override
+    final public void addBinary(Binary value) {
+      Object enumValue = value.toStringUsingUTF8();
+      if (enumClass != null) {
+        enumValue = (Enum.valueOf(enumClass,(String)enumValue));
+      }
+      parent.add(enumValue);
+    }
+  }
+
   static final class FieldFixedConverter extends PrimitiveConverter {
 
     private final ParentValueContainer parent;
@@ -256,14 +286,14 @@ final public void addBinary(Binary value) {
 
   }
 
-  static final class GenericArrayConverter<T> extends GroupConverter {
+  static final class AvroArrayConverter<T> extends GroupConverter {
 
     private final ParentValueContainer parent;
     private final Schema avroSchema;
     private final Converter converter;
     private GenericArray<T> array;
 
-    public GenericArrayConverter(ParentValueContainer parent, Type parquetSchema,
+    public AvroArrayConverter(ParentValueContainer parent, Type parquetSchema,
         Schema avroSchema) {
       this.parent = parent;
       this.avroSchema = avroSchema;
@@ -294,6 +324,51 @@ public void end() {
     }
   }
 
+  static final class AvroUnionConverter<T> extends GroupConverter {
+
+    private final ParentValueContainer parent;
+    private final Converter[] memberConverters;
+    private Object memberValue = null;
+
+    public AvroUnionConverter(ParentValueContainer parent, Type parquetSchema,
+                              Schema avroSchema) {
+      this.parent = parent;
+      GroupType parquetGroup = parquetSchema.asGroupType();
+      this.memberConverters = new Converter[ parquetGroup.getFieldCount()];
+
+      int parquetIndex = 0;
+      for (int index = 0; index < avroSchema.getTypes().size(); index++) {
+        Schema memberSchema = avroSchema.getTypes().get(index);
+        if (!memberSchema.getType().equals(Schema.Type.NULL)) {
+          Type memberType = parquetGroup.getType(parquetIndex);
+          memberConverters[parquetIndex] = newConverter(memberSchema, memberType, new ParentValueContainer() {
+            @Override
+            void add(Object value) {
+              Preconditions.checkArgument(memberValue==null, "Union is resolving to more than one type");
+              memberValue = value;
+            }
+          });
+          parquetIndex++; // Note for nulls the parquetIndex id not increased
+        }
+      }
+    }
+
+    @Override
+    public Converter getConverter(int fieldIndex) {
+      return memberConverters[fieldIndex];
+    }
+
+    @Override
+    public void start() {
+      memberValue = null;
+    }
+
+    @Override
+    public void end() {
+      parent.add(memberValue);
+    }
+  }
+
   static final class MapConverter<V> extends GroupConverter {
 
     private final ParentValueContainer parent;

parquet-avro/src/main/java/parquet/avro/AvroParquetInputFormat.java

@@ -15,14 +15,14 @@
  */
 package parquet.avro;
 
-import org.apache.avro.generic.GenericRecord;
+import org.apache.avro.generic.IndexedRecord;
 import parquet.avro.AvroReadSupport;
 import parquet.hadoop.ParquetInputFormat;
 
 /**
  * A Hadoop {@link org.apache.hadoop.mapreduce.InputFormat} for Parquet files.
  */
-public class AvroParquetInputFormat extends ParquetInputFormat<GenericRecord> {
+public class AvroParquetInputFormat extends ParquetInputFormat<IndexedRecord> {
   public AvroParquetInputFormat() {
     super(AvroReadSupport.class);
   }

parquet-avro/src/main/java/parquet/avro/AvroParquetOutputFormat.java

@@ -16,7 +16,7 @@
 package parquet.avro;
 
 import org.apache.avro.Schema;
-import org.apache.avro.generic.GenericRecord;
+import org.apache.avro.generic.IndexedRecord;
 import org.apache.hadoop.mapreduce.Job;
 import parquet.avro.AvroWriteSupport;
 import parquet.hadoop.ParquetOutputFormat;
@@ -25,7 +25,7 @@
 /**
  * A Hadoop {@link org.apache.hadoop.mapreduce.OutputFormat} for Parquet files.
  */
-public class AvroParquetOutputFormat extends ParquetOutputFormat<GenericRecord> {
+public class AvroParquetOutputFormat extends ParquetOutputFormat<IndexedRecord> {
 
   public static void setSchema(Job job, Schema schema) {
     AvroWriteSupport.setSchema(ContextUtil.getConfiguration(job), schema);

parquet-avro/src/main/java/parquet/avro/AvroParquetReader.java

@@ -16,16 +16,23 @@
 package parquet.avro;
 
 import java.io.IOException;
+
+import org.apache.avro.generic.IndexedRecord;
 import org.apache.hadoop.fs.Path;
+import parquet.filter.UnboundRecordFilter;
 import parquet.hadoop.ParquetReader;
 import parquet.hadoop.api.ReadSupport;
 
 /**
  * Read Avro records from a Parquet file.
  */
-public class AvroParquetReader<T> extends ParquetReader<T> {
+public class AvroParquetReader<T extends IndexedRecord> extends ParquetReader<T> {
 
   public AvroParquetReader(Path file) throws IOException {
     super(file, (ReadSupport<T>) new AvroReadSupport());
   }
+
+  public AvroParquetReader(Path file, UnboundRecordFilter recordFilter ) throws IOException {
+    super(file, (ReadSupport<T>) new AvroReadSupport(), recordFilter);
+  }
 }

parquet-avro/src/main/java/parquet/avro/AvroParquetWriter.java

@@ -17,6 +17,7 @@
 
 import java.io.IOException;
 import org.apache.avro.Schema;
+import org.apache.avro.generic.IndexedRecord;
 import org.apache.hadoop.fs.Path;
 import parquet.hadoop.ParquetWriter;
 import parquet.hadoop.api.WriteSupport;
@@ -25,7 +26,7 @@
 /**
  * Write Avro records to a Parquet file.
  */
-public class AvroParquetWriter<T> extends ParquetWriter<T> {
+public class AvroParquetWriter<T extends IndexedRecord> extends ParquetWriter<T> {
 
   /** Create a new {@link AvroParquetWriter}.
    *
@@ -39,21 +40,38 @@
   public AvroParquetWriter(Path file, Schema avroSchema,
       CompressionCodecName compressionCodecName, int blockSize,
       int pageSize) throws IOException {
+    super(file, (WriteSupport<T>)new AvroWriteSupport(new AvroSchemaConverter().convert(avroSchema), avroSchema),
+	      compressionCodecName, blockSize, pageSize);
+  }
+
+  /** Create a new {@link AvroParquetWriter}.
+   *
+   * @param file The file name to write to.
+   * @param avroSchema The schema to write with.
+   * @param compressionCodecName Compression code to use, or CompressionCodecName.UNCOMPRESSED
+   * @param blockSize HDFS block size
+   * @param pageSize See parquet write up. Blocks are subdivided into pages for alignment and other purposes.
+   * @param enableDictionary Whether to use a dictionary to compress columns.
+   * @throws IOException
+   */
+  public AvroParquetWriter(Path file, Schema avroSchema,
+                           CompressionCodecName compressionCodecName, int blockSize,
+                           int pageSize, boolean enableDictionary) throws IOException {
     super(file, (WriteSupport<T>)
-	new AvroWriteSupport(new AvroSchemaConverter().convert(avroSchema), avroSchema),
-	compressionCodecName, blockSize, pageSize);
+        new AvroWriteSupport(new AvroSchemaConverter().convert(avroSchema),avroSchema),
+        compressionCodecName, blockSize, pageSize, enableDictionary, false);
   }
 
   /** Create a new {@link AvroParquetWriter}. The default block size is 50 MB.The default
    *  page size is 1 MB.  Default compression is no compression. (Inherited from {@link ParquetWriter})
    *
-   * @param file
-   * @param avroSchema
+   * @param file The file name to write to.
+   * @param avroSchema The schema to write with.
    * @throws IOException
    */
   public AvroParquetWriter(Path file, Schema avroSchema) throws IOException {
     this(file, avroSchema, CompressionCodecName.UNCOMPRESSED,
-	 DEFAULT_BLOCK_SIZE, DEFAULT_PAGE_SIZE);
+	  DEFAULT_BLOCK_SIZE, DEFAULT_PAGE_SIZE);
   }
 
 }