[VL] parquet file metadata columns support in velox (apache#3870)

[VL] parquet file metadata columns support in velox. Co-authored-by: Zhen Li <[email protected]>
bigo-sg · Oct 8, 2024 · e306bb4 · e306bb4
1 parent 362f4e7
commit e306bb4
Show file tree

Hide file tree

Showing 32 changed files with 591 additions and 48 deletions.
diff --git a/...nds-clickhouse/src/main/scala/io/glutenproject/backendsapi/clickhouse/CHIteratorApi.scala b/...nds-clickhouse/src/main/scala/io/glutenproject/backendsapi/clickhouse/CHIteratorApi.scala
@@ -75,7 +75,8 @@ class CHIteratorApi extends IteratorApi with Logging with LogLevelUtil {
   override def genSplitInfo(
       partition: InputPartition,
       partitionSchema: StructType,
-      fileFormat: ReadFileFormat): SplitInfo = {
+      fileFormat: ReadFileFormat,
+      metadataColumnNames: Seq[String]): SplitInfo = {
     partition match {
       case p: GlutenMergeTreePartition =>
         val partLists = new JArrayList[String]()
@@ -128,6 +129,7 @@ class CHIteratorApi extends IteratorApi with Logging with LogLevelUtil {
           starts,
           lengths,
           partitionColumns,
+          new JArrayList[JMap[String, String]](),
           fileFormat,
           preferredLocations.toList.asJava)
       case _ =>

diff --git a/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/IteratorApiImpl.scala b/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/IteratorApiImpl.scala
@@ -20,6 +20,7 @@ import io.glutenproject.GlutenNumaBindingInfo
 import io.glutenproject.backendsapi.IteratorApi
 import io.glutenproject.execution._
 import io.glutenproject.metrics.IMetrics
+import io.glutenproject.sql.shims.SparkShimLoader
 import io.glutenproject.substrait.plan.PlanNode
 import io.glutenproject.substrait.rel.{LocalFilesBuilder, LocalFilesNode, SplitInfo}
 import io.glutenproject.substrait.rel.LocalFilesNode.ReadFileFormat
@@ -53,11 +54,12 @@ class IteratorApiImpl extends IteratorApi with Logging {
   override def genSplitInfo(
       partition: InputPartition,
       partitionSchema: StructType,
-      fileFormat: ReadFileFormat): SplitInfo = {
+      fileFormat: ReadFileFormat,
+      metadataColumnNames: Seq[String]): SplitInfo = {
     partition match {
       case f: FilePartition =>
-        val (paths, starts, lengths, partitionColumns) =
-          constructSplitInfo(partitionSchema, f.files)
+        val (paths, starts, lengths, partitionColumns, metadataColumns) =
+          constructSplitInfo(partitionSchema, f.files, metadataColumnNames)
         val preferredLocations =
           SoftAffinity.getFilePartitionLocations(f)
         LocalFilesBuilder.makeLocalFiles(
@@ -66,6 +68,7 @@ class IteratorApiImpl extends IteratorApi with Logging {
           starts,
           lengths,
           partitionColumns,
+          metadataColumns,
           fileFormat,
           preferredLocations.toList.asJava)
       case _ =>
@@ -92,11 +95,15 @@ class IteratorApiImpl extends IteratorApi with Logging {
     }
   }
 
-  private def constructSplitInfo(schema: StructType, files: Array[PartitionedFile]) = {
+  private def constructSplitInfo(
+      schema: StructType,
+      files: Array[PartitionedFile],
+      metadataColumnNames: Seq[String]) = {
     val paths = new JArrayList[String]()
     val starts = new JArrayList[JLong]
     val lengths = new JArrayList[JLong]()
     val partitionColumns = new JArrayList[JMap[String, String]]
+    var metadataColumns = new JArrayList[JMap[String, String]]
     files.foreach {
       file =>
         // The "file.filePath" in PartitionedFile is not the original encoded path, so the decoded
@@ -106,7 +113,9 @@ class IteratorApiImpl extends IteratorApi with Logging {
             .decode(file.filePath.toString, StandardCharsets.UTF_8.name()))
         starts.add(JLong.valueOf(file.start))
         lengths.add(JLong.valueOf(file.length))
-
+        val metadataColumn =
+          SparkShimLoader.getSparkShims.generateMetadataColumns(file, metadataColumnNames)
+        metadataColumns.add(metadataColumn)
         val partitionColumn = new JHashMap[String, String]()
         for (i <- 0 until file.partitionValues.numFields) {
           val partitionColumnValue = if (file.partitionValues.isNullAt(i)) {
@@ -131,7 +140,7 @@ class IteratorApiImpl extends IteratorApi with Logging {
         }
         partitionColumns.add(partitionColumn)
     }
-    (paths, starts, lengths, partitionColumns)
+    (paths, starts, lengths, partitionColumns, metadataColumns)
   }
 
   override def injectWriteFilesTempPath(path: String): Unit = {

diff --git a/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/VeloxBackend.scala b/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/VeloxBackend.scala
@@ -248,6 +248,8 @@ object BackendSettings extends BackendSettingsApi {
     }
   }
 
+  override def supportNativeMetadataColumns(): Boolean = true
+
   override def supportExpandExec(): Boolean = true
 
   override def supportSortExec(): Boolean = true

diff --git a/cpp/velox/compute/VeloxPlanConverter.cc b/cpp/velox/compute/VeloxPlanConverter.cc
@@ -60,6 +60,7 @@ std::shared_ptr<SplitInfo> parseScanSplitInfo(
   splitInfo->starts.reserve(fileList.size());
   splitInfo->lengths.reserve(fileList.size());
   splitInfo->partitionColumns.reserve(fileList.size());
+  splitInfo->metadataColumns.reserve(fileList.size());
   for (const auto& file : fileList) {
     // Expect all Partitions share the same index.
     splitInfo->partitionIndex = file.partition_index();
@@ -70,6 +71,12 @@ std::shared_ptr<SplitInfo> parseScanSplitInfo(
     }
     splitInfo->partitionColumns.emplace_back(partitionColumnMap);
 
+    std::unordered_map<std::string, std::string> metadataColumnMap;
+    for (const auto& metadataColumn : file.metadata_columns()) {
+      metadataColumnMap[metadataColumn.key()] = metadataColumn.value();
+    }
+    splitInfo->metadataColumns.emplace_back(metadataColumnMap);
+
     splitInfo->paths.emplace_back(file.uri_file());
     splitInfo->starts.emplace_back(file.start());
     splitInfo->lengths.emplace_back(file.length());

diff --git a/cpp/velox/compute/WholeStageResultIterator.cc b/cpp/velox/compute/WholeStageResultIterator.cc
@@ -145,15 +145,28 @@ WholeStageResultIterator::WholeStageResultIterator(
     const auto& lengths = scanInfo->lengths;
     const auto& format = scanInfo->format;
     const auto& partitionColumns = scanInfo->partitionColumns;
+    const auto& metadataColumns = scanInfo->metadataColumns;
 
     std::vector<std::shared_ptr<velox::connector::ConnectorSplit>> connectorSplits;
     connectorSplits.reserve(paths.size());
     for (int idx = 0; idx < paths.size(); idx++) {
       auto partitionColumn = partitionColumns[idx];
+      auto metadataColumn = metadataColumns[idx];
       std::unordered_map<std::string, std::optional<std::string>> partitionKeys;
       constructPartitionColumns(partitionKeys, partitionColumn);
       auto split = std::make_shared<velox::connector::hive::HiveConnectorSplit>(
-          kHiveConnectorId, paths[idx], format, starts[idx], lengths[idx], partitionKeys);
+          kHiveConnectorId,
+          paths[idx],
+          format,
+          starts[idx],
+          lengths[idx],
+          partitionKeys,
+          std::nullopt,
+          std::unordered_map<std::string, std::string>(),
+          nullptr,
+          std::unordered_map<std::string, std::string>(),
+          0,
+          metadataColumn);
       connectorSplits.emplace_back(split);
     }
 

diff --git a/cpp/velox/substrait/SubstraitParser.cc b/cpp/velox/substrait/SubstraitParser.cc
@@ -104,31 +104,41 @@ std::vector<TypePtr> SubstraitParser::parseNamedStruct(const ::substrait::NamedS
   return typeList;
 }
 
-std::vector<bool> SubstraitParser::parsePartitionColumns(const ::substrait::NamedStruct& namedStruct) {
+void SubstraitParser::parsePartitionAndMetadataColumns(
+    const ::substrait::NamedStruct& namedStruct,
+    std::vector<bool>& isPartitionColumns,
+    std::vector<bool>& isMetadataColumns) {
   const auto& columnsTypes = namedStruct.column_types();
-  std::vector<bool> isPartitionColumns;
   if (columnsTypes.size() == 0) {
-    // Regard all columns as non-partitioned columns.
+    // Regard all columns as regular columns.
     isPartitionColumns.resize(namedStruct.names().size(), false);
-    return isPartitionColumns;
+    isMetadataColumns.resize(namedStruct.names().size(), false);
+    return;
   } else {
     VELOX_CHECK_EQ(columnsTypes.size(), namedStruct.names().size(), "Wrong size for column types and column names.");
   }
 
   isPartitionColumns.reserve(columnsTypes.size());
+  isMetadataColumns.reserve(columnsTypes.size());
   for (const auto& columnType : columnsTypes) {
     switch (columnType) {
       case ::substrait::NamedStruct::NORMAL_COL:
         isPartitionColumns.emplace_back(false);
+        isMetadataColumns.emplace_back(false);
         break;
       case ::substrait::NamedStruct::PARTITION_COL:
         isPartitionColumns.emplace_back(true);
+        isMetadataColumns.emplace_back(false);
+        break;
+      case ::substrait::NamedStruct::METADATA_COL:
+        isPartitionColumns.emplace_back(false);
+        isMetadataColumns.emplace_back(true);
         break;
       default:
         VELOX_FAIL("Unspecified column type.");
     }
   }
-  return isPartitionColumns;
+  return;
 }
 
 int32_t SubstraitParser::parseReferenceSegment(const ::substrait::Expression::ReferenceSegment& refSegment) {

diff --git a/cpp/velox/substrait/SubstraitParser.h b/cpp/velox/substrait/SubstraitParser.h
@@ -41,8 +41,11 @@ class SubstraitParser {
       const ::substrait::NamedStruct& namedStruct,
       bool asLowerCase = false);
 
-  /// Used to parse partition columns from Substrait NamedStruct.
-  static std::vector<bool> parsePartitionColumns(const ::substrait::NamedStruct& namedStruct);
+  /// Used to parse partition & metadata columns from Substrait NamedStruct.
+  static void parsePartitionAndMetadataColumns(
+      const ::substrait::NamedStruct& namedStruct,
+      std::vector<bool>& isPartitionColumns,
+      std::vector<bool>& isMetadataColumns);
 
   /// Parse Substrait Type to Velox type.
   static facebook::velox::TypePtr parseType(const ::substrait::Type& substraitType, bool asLowerCase = false);

diff --git a/cpp/velox/substrait/SubstraitToVeloxPlan.cc b/cpp/velox/substrait/SubstraitToVeloxPlan.cc
@@ -596,11 +596,12 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
   std::vector<std::string> tableColumnNames;
   std::vector<std::string> partitionedKey;
   std::vector<bool> isPartitionColumns;
+  std::vector<bool> isMetadataColumns;
   tableColumnNames.reserve(writeRel.table_schema().names_size());
 
   VELOX_CHECK(writeRel.has_table_schema(), "WriteRel should have the table schema to store the column information");
   const auto& tableSchema = writeRel.table_schema();
-  isPartitionColumns = SubstraitParser::parsePartitionColumns(tableSchema);
+  SubstraitParser::parsePartitionAndMetadataColumns(tableSchema, isPartitionColumns, isMetadataColumns);
 
   for (const auto& name : tableSchema.names()) {
     tableColumnNames.emplace_back(name);
@@ -1040,6 +1041,7 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
   std::vector<std::string> colNameList;
   std::vector<TypePtr> veloxTypeList;
   std::vector<bool> isPartitionColumns;
+  std::vector<bool> isMetadataColumns;
   // Convert field names into lower case when not case-sensitive.
   std::shared_ptr<const facebook::velox::Config> veloxCfg =
       std::make_shared<const facebook::velox::core::MemConfigMutable>(confMap_);
@@ -1055,7 +1057,7 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
       colNameList.emplace_back(fieldName);
     }
     veloxTypeList = SubstraitParser::parseNamedStruct(baseSchema, asLowerCase);
-    isPartitionColumns = SubstraitParser::parsePartitionColumns(baseSchema);
+    SubstraitParser::parsePartitionAndMetadataColumns(baseSchema, isPartitionColumns, isMetadataColumns);
   }
 
   // Do not hard-code connector ID and allow for connectors other than Hive.
@@ -1110,8 +1112,13 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
   std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>> assignments;
   for (int idx = 0; idx < colNameList.size(); idx++) {
     auto outName = SubstraitParser::makeNodeName(planNodeId_, idx);
-    auto columnType = isPartitionColumns[idx] ? connector::hive::HiveColumnHandle::ColumnType::kPartitionKey
-                                              : connector::hive::HiveColumnHandle::ColumnType::kRegular;
+    auto columnType = connector::hive::HiveColumnHandle::ColumnType::kRegular;
+    if (isPartitionColumns[idx]) {
+      columnType = connector::hive::HiveColumnHandle::ColumnType::kPartitionKey;
+    }
+    if (isMetadataColumns[idx]) {
+      columnType = connector::hive::HiveColumnHandle::ColumnType::kSynthesized;
+    }
     assignments[outName] = std::make_shared<connector::hive::HiveColumnHandle>(
         colNameList[idx], columnType, veloxTypeList[idx], veloxTypeList[idx]);
     outNames.emplace_back(outName);

diff --git a/cpp/velox/substrait/SubstraitToVeloxPlan.h b/cpp/velox/substrait/SubstraitToVeloxPlan.h
@@ -36,6 +36,9 @@ struct SplitInfo {
   /// The partition columns associated with partitioned table.
   std::vector<std::unordered_map<std::string, std::string>> partitionColumns;
 
+  /// The metadata columns associated with partitioned table.
+  std::vector<std::unordered_map<std::string, std::string>> metadataColumns;
+
   /// The file paths to be scanned.
   std::vector<std::string> paths;
 

diff --git a/cpp/velox/substrait/SubstraitToVeloxPlanValidator.cc b/cpp/velox/substrait/SubstraitToVeloxPlanValidator.cc
@@ -366,7 +366,9 @@ bool SubstraitToVeloxPlanValidator::validate(const ::substrait::WriteRel& writeR
   // Validate partition key type.
   if (writeRel.has_table_schema()) {
     const auto& tableSchema = writeRel.table_schema();
-    auto isPartitionColumns = SubstraitParser::parsePartitionColumns(tableSchema);
+    std::vector<bool> isMetadataColumns;
+    std::vector<bool> isPartitionColumns;
+    SubstraitParser::parsePartitionAndMetadataColumns(tableSchema, isPartitionColumns, isMetadataColumns);
     for (auto i = 0; i < types.size(); i++) {
       if (isPartitionColumns[i]) {
         switch (types[i]->kind()) {

diff --git a/gluten-core/src/main/java/io/glutenproject/substrait/rel/LocalFilesBuilder.java b/gluten-core/src/main/java/io/glutenproject/substrait/rel/LocalFilesBuilder.java
@@ -28,10 +28,18 @@ public static LocalFilesNode makeLocalFiles(
       List<Long> starts,
       List<Long> lengths,
       List<Map<String, String>> partitionColumns,
+      List<Map<String, String>> metadataColumns,
       LocalFilesNode.ReadFileFormat fileFormat,
       List<String> preferredLocations) {
     return new LocalFilesNode(
-        index, paths, starts, lengths, partitionColumns, fileFormat, preferredLocations);
+        index,
+        paths,
+        starts,
+        lengths,
+        partitionColumns,
+        metadataColumns,
+        fileFormat,
+        preferredLocations);
   }
 
   public static LocalFilesNode makeLocalFiles(String iterPath) {

diff --git a/gluten-core/src/main/java/io/glutenproject/substrait/rel/LocalFilesNode.java b/gluten-core/src/main/java/io/glutenproject/substrait/rel/LocalFilesNode.java
@@ -35,6 +35,7 @@ public class LocalFilesNode implements SplitInfo {
   private final List<Long> starts = new ArrayList<>();
   private final List<Long> lengths = new ArrayList<>();
   private final List<Map<String, String>> partitionColumns = new ArrayList<>();
+  private final List<Map<String, String>> metadataColumns = new ArrayList<>();
   private final List<String> preferredLocations = new ArrayList<>();
 
   // The format of file to read.
@@ -60,6 +61,7 @@ public enum ReadFileFormat {
       List<Long> starts,
       List<Long> lengths,
       List<Map<String, String>> partitionColumns,
+      List<Map<String, String>> metadataColumns,
       ReadFileFormat fileFormat,
       List<String> preferredLocations) {
     this.index = index;
@@ -68,6 +70,7 @@ public enum ReadFileFormat {
     this.lengths.addAll(lengths);
     this.fileFormat = fileFormat;
     this.partitionColumns.addAll(partitionColumns);
+    this.metadataColumns.addAll(metadataColumns);
     this.preferredLocations.addAll(preferredLocations);
   }
 
@@ -141,7 +144,22 @@ public ReadRel.LocalFiles toProtobuf() {
       }
       fileBuilder.setLength(lengths.get(i));
       fileBuilder.setStart(starts.get(i));
-
+      if (!metadataColumns.isEmpty()) {
+        Map<String, String> metadataColumn = metadataColumns.get(i);
+        if (!metadataColumn.isEmpty()) {
+          metadataColumn.forEach(
+              (key, value) -> {
+                ReadRel.LocalFiles.FileOrFiles.metadataColumn.Builder mcBuilder =
+                    ReadRel.LocalFiles.FileOrFiles.metadataColumn.newBuilder();
+                mcBuilder.setKey(key).setValue(value);
+                fileBuilder.addMetadataColumns(mcBuilder.build());
+              });
+        }
+      } else {
+        ReadRel.LocalFiles.FileOrFiles.metadataColumn.Builder mcBuilder =
+            ReadRel.LocalFiles.FileOrFiles.metadataColumn.newBuilder();
+        fileBuilder.addMetadataColumns(mcBuilder.build());
+      }
       NamedStruct namedStruct = buildNamedStruct();
       fileBuilder.setSchema(namedStruct);
 

diff --git a/gluten-core/src/main/resources/substrait/proto/substrait/algebra.proto b/gluten-core/src/main/resources/substrait/proto/substrait/algebra.proto
@@ -169,6 +169,12 @@ message ReadRel {
 
       /// File schema
       NamedStruct schema = 17;
+
+      message metadataColumn {
+        string key = 1;
+        string value = 2;
+      }
+      repeated metadataColumn metadata_columns = 18;
     }
   }
 }

diff --git a/gluten-core/src/main/resources/substrait/proto/substrait/type.proto b/gluten-core/src/main/resources/substrait/proto/substrait/type.proto
@@ -237,5 +237,6 @@ message NamedStruct {
   enum ColumnType {
     NORMAL_COL = 0;
     PARTITION_COL = 1;
+    METADATA_COL = 2;
   }
 }
diff --git a/gluten-core/src/main/scala/io/glutenproject/backendsapi/BackendSettingsApi.scala b/gluten-core/src/main/scala/io/glutenproject/backendsapi/BackendSettingsApi.scala
@@ -41,6 +41,7 @@ trait BackendSettingsApi {
       fields: Array[StructField],
       bucketSpec: Option[BucketSpec],
       options: Map[String, String]): ValidationResult = ValidationResult.ok
+  def supportNativeMetadataColumns(): Boolean = false
   def supportExpandExec(): Boolean = false
   def supportSortExec(): Boolean = false
   def supportSortMergeJoinExec(): Boolean = true

diff --git a/gluten-core/src/main/scala/io/glutenproject/backendsapi/IteratorApi.scala b/gluten-core/src/main/scala/io/glutenproject/backendsapi/IteratorApi.scala
@@ -36,7 +36,8 @@ trait IteratorApi {
   def genSplitInfo(
       partition: InputPartition,
       partitionSchema: StructType,
-      fileFormat: ReadFileFormat): SplitInfo
+      fileFormat: ReadFileFormat,
+      metadataColumnNames: Seq[String]): SplitInfo
 
   /** Generate native row partition. */
   def genPartitions(
-Original file line number
+Diff line change
@@ Expand Up / @@ -248,6 +248,8 @@ object BackendSettings extends BackendSettingsApi { @@
         }
       }
+      override def supportNativeMetadataColumns(): Boolean = true
       override def supportExpandExec(): Boolean = true
       override def supportSortExec(): Boolean = true
@@ Expand Down @@