Fix Parquet collection wrapper extraction

Author: xiangfu0

pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetAvroRecordExtractor.java

@@ -37,6 +37,11 @@ protected Object transformValue(Object value, Schema.Field field) {
     return handleDeprecatedTypes(convert(value), field);
   }
 
+  @Override
+  protected Object[] convertMultiValue(Object value) {
+    return ParquetRecordExtractorUtils.unwrapListElementMaps(super.convertMultiValue(value));
+  }
+
   Object handleDeprecatedTypes(Object value, Schema.Field field) {
     Schema.Type avroColumnType = field.schema().getType();
     if (avroColumnType == Schema.Type.UNION) {

pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractor.java

@@ -161,6 +161,10 @@ private Object extractValue(Group from, int fieldIndex, Type fieldType, int inde
       if (logicalTypeAnnotation instanceof LogicalTypeAnnotation.ListLogicalTypeAnnotation) {
         return extractList(group);
       }
+      if (logicalTypeAnnotation instanceof LogicalTypeAnnotation.MapLogicalTypeAnnotation
+          || logicalTypeAnnotation instanceof LogicalTypeAnnotation.MapKeyValueTypeAnnotation) {
+        return ParquetRecordExtractorUtils.unwrapMapKeyValues(extractMap(group));
+      }
       return extractMap(group);
     }
     return null;
@@ -182,9 +186,9 @@ public Object[] extractList(Group group) {
       return new Object[0];
     }
     if (numValues == 1 && array[0] instanceof Object[]) {
-      return (Object[]) array[0];
+      return ParquetRecordExtractorUtils.unwrapListElementMaps((Object[]) array[0]);
     }
-    return array;
+    return ParquetRecordExtractorUtils.unwrapListElementMaps(array);
   }
 
   public Map<String, Object> extractMap(Group group) {

pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetRecordExtractorUtils.java

@@ -0,0 +1,112 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.plugin.inputformat.parquet;
+
+import java.util.Collection;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+
+/**
+ * Shared helpers for normalizing Parquet list wrapper records while extracting Pinot rows.
+ *
+ * <p>This class is stateless and thread-safe.
+ */
+final class ParquetRecordExtractorUtils {
+  private static final String LIST_ELEMENT_FIELD_NAME = "element";
+  private static final String MAP_ENTRIES_FIELD_NAME = "key_value";
+  private static final String MAP_KEY_FIELD_NAME = "key";
+  private static final String MAP_VALUE_FIELD_NAME = "value";
+
+  private ParquetRecordExtractorUtils() {
+  }
+
+  static Object[] unwrapListElementMaps(Object[] values) {
+    if (values.length == 0) {
+      return values;
+    }
+
+    boolean hasElementMap = false;
+    for (Object value : values) {
+      if (value == null) {
+        continue;
+      }
+      if (!(value instanceof Map)) {
+        return values;
+      }
+      Map<?, ?> map = (Map<?, ?>) value;
+      if (map.size() != 1 || !map.containsKey(LIST_ELEMENT_FIELD_NAME)) {
+        return values;
+      }
+      hasElementMap = true;
+    }
+    if (!hasElementMap) {
+      return values;
+    }
+
+    Object[] unwrappedValues = new Object[values.length];
+    for (int i = 0; i < values.length; i++) {
+      Object value = values[i];
+      unwrappedValues[i] = value == null ? null : ((Map<?, ?>) value).get(LIST_ELEMENT_FIELD_NAME);
+    }
+    return unwrappedValues;
+  }
+
+  static Map<String, Object> unwrapMapKeyValues(Map<String, Object> map) {
+    if (map.size() != 1 || !map.containsKey(MAP_ENTRIES_FIELD_NAME)) {
+      return map;
+    }
+
+    Object entries = map.get(MAP_ENTRIES_FIELD_NAME);
+    Object[] entryArray = toEntryArray(entries);
+    if (entryArray == null || entryArray.length == 0) {
+      return new LinkedHashMap<>();
+    }
+
+    Map<String, Object> unwrappedMap = new LinkedHashMap<>();
+    for (Object entry : entryArray) {
+      if (!(entry instanceof Map)) {
+        return map;
+      }
+      Map<?, ?> entryMap = (Map<?, ?>) entry;
+      if (!entryMap.containsKey(MAP_KEY_FIELD_NAME) || !entryMap.containsKey(MAP_VALUE_FIELD_NAME)) {
+        return map;
+      }
+      Object key = entryMap.get(MAP_KEY_FIELD_NAME);
+      if (key == null) {
+        return map;
+      }
+      unwrappedMap.put(key.toString(), entryMap.get(MAP_VALUE_FIELD_NAME));
+    }
+    return unwrappedMap;
+  }
+
+  private static Object[] toEntryArray(Object entries) {
+    if (entries instanceof Object[]) {
+      return (Object[]) entries;
+    }
+    if (entries instanceof Collection) {
+      return ((Collection<?>) entries).toArray();
+    }
+    if (entries != null) {
+      return new Object[]{entries};
+    }
+    return null;
+  }
+}