Add variant spec version to Variant type

Author: aihuaxu

parquet-column/src/main/java/org/apache/parquet/schema/LogicalTypeAnnotation.java

@@ -43,6 +43,8 @@
 import org.apache.parquet.Preconditions;
 
 public abstract class LogicalTypeAnnotation {
+  public static final byte VARIANT_SPEC_VERSION = 1;
+
   enum LogicalTypeToken {
     MAP {
       @Override
@@ -60,9 +62,8 @@ protected LogicalTypeAnnotation fromString(List<String> params) {
       @Override
       protected LogicalTypeAnnotation fromString(List<String> params) {
         Preconditions.checkArgument(
-            params.isEmpty(), "Expecting 0 parameter for variant logical type, got %d", params.size());
-
-        return variantType();
+            params.size() == 1, "Expecting 1 parameter for variant logical type, got %d", params.size());
+        return variantType(Byte.parseByte(params.get(0)));
       }
     },
     STRING {
@@ -278,8 +279,8 @@ public static ListLogicalTypeAnnotation listType() {
     return ListLogicalTypeAnnotation.INSTANCE;
   }
 
-  public static VariantLogicalTypeAnnotation variantType() {
-    return VariantLogicalTypeAnnotation.INSTANCE;
+  public static VariantLogicalTypeAnnotation variantType(byte specificationVersion) {
+    return new VariantLogicalTypeAnnotation(specificationVersion);
   }
 
   public static EnumLogicalTypeAnnotation enumType() {
@@ -1142,9 +1143,11 @@ public int hashCode() {
   }
 
   public static class VariantLogicalTypeAnnotation extends LogicalTypeAnnotation {
-    private static final VariantLogicalTypeAnnotation INSTANCE = new VariantLogicalTypeAnnotation();
+    private byte specificationVersion;
 
-    private VariantLogicalTypeAnnotation() {}
+    private VariantLogicalTypeAnnotation(byte specificationVersion) {
+      this.specificationVersion = specificationVersion;
+    }
 
     @Override
     public OriginalType toOriginalType() {
@@ -1161,6 +1164,24 @@ public <T> Optional<T> accept(LogicalTypeAnnotationVisitor<T> logicalTypeAnnotat
     LogicalTypeToken getType() {
       return LogicalTypeToken.VARIANT;
     }
+
+    public byte getSpecificationVersion() {
+      return this.specificationVersion;
+    }
+
+    @Override
+    protected String typeParametersAsString() {
+      return "(" + specificationVersion + ")";
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      if (!(obj instanceof VariantLogicalTypeAnnotation)) {
+        return false;
+      }
+      VariantLogicalTypeAnnotation other = (VariantLogicalTypeAnnotation) obj;
+      return specificationVersion == other.specificationVersion;
+    }
   }
 
   /**

parquet-column/src/main/java/org/apache/parquet/schema/MessageTypeParser.java

@@ -124,17 +124,30 @@ private static void addGroupType(Tokenizer st, Repetition r, GroupBuilder<?> bui
       t = st.nextToken();
       if (isLogicalType(t)) {
         LogicalTypeAnnotation.LogicalTypeToken logicalType = LogicalTypeAnnotation.LogicalTypeToken.valueOf(t);
-        LogicalTypeAnnotation logicalTypeAnnotation = logicalType.fromString(new ArrayList<>());
+        t = st.nextToken();
+        List<String> tokens = new ArrayList<>();
+        if ("(".equals(t)) {
+          while (!")".equals(t)) {
+            if (!(",".equals(t) || "(".equals(t) || ")".equals(t))) {
+              tokens.add(t);
+            }
+            t = st.nextToken();
+          }
+          t = st.nextToken();
+        }
+
+        LogicalTypeAnnotation logicalTypeAnnotation = logicalType.fromString(tokens);
         childBuilder.as(logicalTypeAnnotation);
         annotation = logicalTypeAnnotation.toString();
       } else {
         // Try to parse as OriginalType
         OriginalType originalType = OriginalType.valueOf(t);
         childBuilder.as(originalType);
         annotation = originalType.toString();
+        t = st.nextToken();
       }
 
-      check(st.nextToken(), ")", "logical type ended by )", st);
+      check(t, ")", "logical type ended by )", st);
       t = st.nextToken();
     }
     if (t.equals("=")) {

parquet-column/src/test/java/org/apache/parquet/parser/TestParquetParser.java

@@ -452,15 +452,15 @@ public void testEmbeddedAnnotations() {
   @Test
   public void testVARIANTAnnotation() {
     String message = "message Message {\n"
-        + "  required group aVariant (VARIANT) {\n"
+        + "  required group aVariant (VARIANT(2)) {\n"
         + "     required binary metadata;\n"
         + "     required binary value;\n"
         + "  }\n"
         + "}\n";
 
     MessageType expected = buildMessage()
         .requiredGroup()
-        .as(LogicalTypeAnnotation.variantType())
+        .as(LogicalTypeAnnotation.variantType((byte) 2))
         .required(BINARY)
         .named("metadata")
         .required(BINARY)

parquet-column/src/test/java/org/apache/parquet/schema/TestTypeBuilders.java

@@ -20,6 +20,7 @@
 
 import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.MICROS;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.MILLIS;
+import static org.apache.parquet.schema.LogicalTypeAnnotation.VARIANT_SPEC_VERSION;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.timestampType;
 import static org.apache.parquet.schema.OriginalType.BSON;
 import static org.apache.parquet.schema.OriginalType.DATE;
@@ -1421,15 +1422,15 @@ public void testVariantLogicalType() {
     GroupType variantExpected = new GroupType(
         REQUIRED,
         name,
-        LogicalTypeAnnotation.variantType(),
+        LogicalTypeAnnotation.variantType(VARIANT_SPEC_VERSION),
         new PrimitiveType(REQUIRED, BINARY, "metadata"),
         new PrimitiveType(REQUIRED, BINARY, "value"));
 
     GroupType variantActual = Types.buildGroup(REQUIRED)
         .addFields(
             Types.required(BINARY).named("metadata"),
             Types.required(BINARY).named("value"))
-        .as(LogicalTypeAnnotation.variantType())
+        .as(LogicalTypeAnnotation.variantType(VARIANT_SPEC_VERSION))
         .named(name);
 
     assertEquals(variantExpected, variantActual);
@@ -1441,7 +1442,7 @@ public void testVariantLogicalTypeWithShredded() {
     GroupType variantExpected = new GroupType(
         REQUIRED,
         name,
-        LogicalTypeAnnotation.variantType(),
+        LogicalTypeAnnotation.variantType(VARIANT_SPEC_VERSION),
         new PrimitiveType(REQUIRED, BINARY, "metadata"),
         new PrimitiveType(OPTIONAL, BINARY, "value"),
         new PrimitiveType(OPTIONAL, BINARY, "typed_value", LogicalTypeAnnotation.stringType()));
@@ -1453,7 +1454,7 @@ public void testVariantLogicalTypeWithShredded() {
             Types.optional(BINARY)
                 .as(LogicalTypeAnnotation.stringType())
                 .named("typed_value"))
-        .as(LogicalTypeAnnotation.variantType())
+        .as(LogicalTypeAnnotation.variantType(VARIANT_SPEC_VERSION))
         .named(name);
 
     assertEquals(variantExpected, variantActual);

parquet-column/src/test/java/org/apache/parquet/schema/TestTypeBuildersWithLogicalTypes.java

@@ -21,6 +21,7 @@
 import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.MICROS;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.MILLIS;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.NANOS;
+import static org.apache.parquet.schema.LogicalTypeAnnotation.VARIANT_SPEC_VERSION;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.bsonType;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.dateType;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.decimalType;
@@ -481,7 +482,7 @@ public void testVariantLogicalType() {
     GroupType variant = new GroupType(
         REQUIRED,
         name,
-        LogicalTypeAnnotation.variantType(),
+        LogicalTypeAnnotation.variantType(VARIANT_SPEC_VERSION),
         Types.required(BINARY).named("metadata"),
         Types.required(BINARY).named("value"));
 
@@ -504,7 +505,7 @@ public void testVariantLogicalTypeWithShredded() {
     GroupType variant = new GroupType(
         REQUIRED,
         name,
-        LogicalTypeAnnotation.variantType(),
+        LogicalTypeAnnotation.variantType(VARIANT_SPEC_VERSION),
         Types.required(BINARY).named("metadata"),
         Types.optional(BINARY).named("value"),
         Types.optional(BINARY).as(LogicalTypeAnnotation.stringType()).named("typed_value"));
@@ -522,6 +523,7 @@ public void testVariantLogicalTypeWithShredded() {
     assertNull(annotation.toOriginalType());
     assertTrue(annotation instanceof LogicalTypeAnnotation.VariantLogicalTypeAnnotation);
   }
+
   /**
    * A convenience method to avoid a large number of @Test(expected=...) tests
    *

parquet-format-structures/src/main/java/org/apache/parquet/format/LogicalTypes.java

@@ -32,6 +32,12 @@ public static LogicalType DECIMAL(int scale, int precision) {
     return LogicalType.DECIMAL(new DecimalType(scale, precision));
   }
 
+  public static LogicalType VARIANT(byte specificationVersion) {
+    VariantType type = new VariantType();
+    type.setSpecification_version(specificationVersion);
+    return LogicalType.VARIANT(type);
+  }
+
   public static final LogicalType UTF8 = LogicalType.STRING(new StringType());
   public static final LogicalType MAP = LogicalType.MAP(new MapType());
   public static final LogicalType LIST = LogicalType.LIST(new ListType());
@@ -54,11 +60,4 @@ public static LogicalType DECIMAL(int scale, int precision) {
   public static final LogicalType BSON = LogicalType.BSON(new BsonType());
   public static final LogicalType FLOAT16 = LogicalType.FLOAT16(new Float16Type());
   public static final LogicalType UUID = LogicalType.UUID(new UUIDType());
-  public static final LogicalType VARIANT = LogicalType.VARIANT(variant());
-
-  private static final VariantType variant() {
-    VariantType type = new VariantType();
-    type.setSpecification_version((byte) 1);
-    return type;
-  }
 }

parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java

@@ -103,6 +103,7 @@
 import org.apache.parquet.format.Type;
 import org.apache.parquet.format.TypeDefinedOrder;
 import org.apache.parquet.format.Uncompressed;
+import org.apache.parquet.format.VariantType;
 import org.apache.parquet.format.XxHash;
 import org.apache.parquet.hadoop.metadata.BlockMetaData;
 import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
@@ -517,7 +518,7 @@ public Optional<LogicalType> visit(LogicalTypeAnnotation.IntervalLogicalTypeAnno
 
     @Override
     public Optional<LogicalType> visit(LogicalTypeAnnotation.VariantLogicalTypeAnnotation variantLogicalType) {
-      return of(LogicalTypes.VARIANT);
+      return of(LogicalTypes.VARIANT(variantLogicalType.getSpecificationVersion()));
     }
   }
 
@@ -1183,7 +1184,8 @@ LogicalTypeAnnotation getLogicalTypeAnnotation(LogicalType type) {
       case FLOAT16:
         return LogicalTypeAnnotation.float16Type();
       case VARIANT:
-        return LogicalTypeAnnotation.variantType();
+        VariantType variant = type.getVARIANT();
+        return LogicalTypeAnnotation.variantType(variant.getSpecification_version());
       default:
         throw new RuntimeException("Unknown logical type " + type);
     }

parquet-hadoop/src/test/java/org/apache/parquet/format/converter/TestParquetMetadataConverter.java

@@ -29,6 +29,7 @@
 import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.MICROS;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.MILLIS;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit.NANOS;
+import static org.apache.parquet.schema.LogicalTypeAnnotation.VARIANT_SPEC_VERSION;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.bsonType;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.dateType;
 import static org.apache.parquet.schema.LogicalTypeAnnotation.decimalType;
@@ -1594,7 +1595,7 @@ public void testMapConvertedTypeReadWrite() throws Exception {
   public void testVariantLogicalType() {
     MessageType expected = Types.buildMessage()
         .requiredGroup()
-        .as(variantType())
+        .as(variantType(VARIANT_SPEC_VERSION))
         .required(PrimitiveTypeName.BINARY)
         .named("metadata")
         .required(PrimitiveTypeName.BINARY)
@@ -1606,6 +1607,9 @@ public void testVariantLogicalType() {
     List<SchemaElement> parquetSchema = parquetMetadataConverter.toParquetSchema(expected);
     MessageType schema = parquetMetadataConverter.fromParquetSchema(parquetSchema, null);
     assertEquals(expected, schema);
+    LogicalTypeAnnotation logicalType = schema.getType("v").getLogicalTypeAnnotation();
+    assertEquals(LogicalTypeAnnotation.variantType(VARIANT_SPEC_VERSION), logicalType);
+    assertEquals(VARIANT_SPEC_VERSION, ((LogicalTypeAnnotation.VariantLogicalTypeAnnotation) logicalType).getSpecificationVersion());
   }
 
   private void verifyMapMessageType(final MessageType messageType, final String keyValueName) throws IOException {