Support struct column reading with different schemas

rui-mo · rui-mo · commit e847a3bf3183 · 2023-11-22T13:49:45.000+08:00
diff --git a/velox/connectors/hive/SplitReader.cpp b/velox/connectors/hive/SplitReader.cpp
@@ -220,9 +220,18 @@ std::vector<TypePtr> SplitReader::adaptColumns(
     } else {
       auto fileTypeIdx = fileType->getChildIdxIfExists(fieldName);
       if (!fileTypeIdx.has_value()) {
-        // Column is missing. Most likely due to schema evolution.
-        VELOX_CHECK(tableSchema);
-        setNullConstantValue(childSpec, tableSchema->findChild(fieldName));
+        // If field name exists in the user-specified output type,
+        // set the column as null constant.
+        // Related PR: https://github.com/facebookincubator/velox/pull/6427.
+        auto outputTypeIdx = readerOutputType_->getChildIdxIfExists(fieldName);
+        if (outputTypeIdx.has_value()) {
+          setNullConstantValue(
+              childSpec, readerOutputType_->childAt(outputTypeIdx.value()));
+        } else {
+          // Column is missing. Most likely due to schema evolution.
+          VELOX_CHECK(tableSchema);
+          setNullConstantValue(childSpec, tableSchema->findChild(fieldName));
+        }
       } else {
         // Column no longer missing, reset constant value set on the spec.
         childSpec->setConstantValue(nullptr);
diff --git a/velox/dwio/common/SelectiveStructColumnReader.cpp b/velox/dwio/common/SelectiveStructColumnReader.cpp
@@ -133,7 +133,6 @@ void SelectiveStructColumnReaderBase::read(
   }
 
   auto& childSpecs = scanSpec_->children();
-  VELOX_CHECK(!childSpecs.empty());
   for (size_t i = 0; i < childSpecs.size(); ++i) {
     auto& childSpec = childSpecs[i];
     if (isChildConstant(*childSpec)) {
@@ -218,7 +217,7 @@ bool SelectiveStructColumnReaderBase::isChildConstant(
        fileType_->type()->kind() !=
            TypeKind::MAP && // If this is the case it means this is a flat map,
                             // so it can't have "missing" fields.
-       childSpec.channel() >= fileType_->size());
+       !fileType_->containsChild(childSpec.fieldName()));
 }
 
 namespace {
@@ -302,7 +301,6 @@ void setNullField(
 void SelectiveStructColumnReaderBase::getValues(
     RowSet rows,
     VectorPtr* result) {
-  VELOX_CHECK(!scanSpec_->children().empty());
   VELOX_CHECK_NOT_NULL(
       *result, "SelectiveStructColumnReaderBase expects a non-null result");
   VELOX_CHECK(
diff --git a/velox/dwio/common/TypeWithId.h b/velox/dwio/common/TypeWithId.h
@@ -59,6 +59,11 @@ class TypeWithId : public velox::Tree<std::shared_ptr<const TypeWithId>> {
 
   const std::shared_ptr<const TypeWithId>& childAt(uint32_t idx) const override;
 
+  bool containsChild(const std::string& name) const {
+    VELOX_CHECK_EQ(type_->kind(), velox::TypeKind::ROW);
+    return type_->as<velox::TypeKind::ROW>().containsChild(name);
+  }
+
   const std::shared_ptr<const TypeWithId>& childByName(
       const std::string& name) const {
     VELOX_CHECK_EQ(type_->kind(), velox::TypeKind::ROW);
diff --git a/velox/dwio/parquet/reader/ParquetColumnReader.cpp b/velox/dwio/parquet/reader/ParquetColumnReader.cpp
@@ -37,7 +37,8 @@ std::unique_ptr<dwio::common::SelectiveColumnReader> ParquetColumnReader::build(
     const std::shared_ptr<const dwio::common::TypeWithId>& requestedType,
     const std::shared_ptr<const dwio::common::TypeWithId>& fileType,
     ParquetParams& params,
-    common::ScanSpec& scanSpec) {
+    common::ScanSpec& scanSpec,
+    memory::MemoryPool& pool) {
   auto colName = scanSpec.fieldName();
 
   switch (fileType->type()->kind()) {
@@ -58,19 +59,19 @@ std::unique_ptr<dwio::common::SelectiveColumnReader> ParquetColumnReader::build(
 
     case TypeKind::ROW:
       return std::make_unique<StructColumnReader>(
-          requestedType, fileType, params, scanSpec);
+          requestedType, fileType, params, scanSpec, pool);
 
     case TypeKind::VARBINARY:
     case TypeKind::VARCHAR:
       return std::make_unique<StringColumnReader>(fileType, params, scanSpec);
 
     case TypeKind::ARRAY:
       return std::make_unique<ListColumnReader>(
-          requestedType, fileType, params, scanSpec);
+          requestedType, fileType, params, scanSpec, pool);
 
     case TypeKind::MAP:
       return std::make_unique<MapColumnReader>(
-          requestedType, fileType, params, scanSpec);
+          requestedType, fileType, params, scanSpec, pool);
 
     case TypeKind::BOOLEAN:
       return std::make_unique<BooleanColumnReader>(
diff --git a/velox/dwio/parquet/reader/ParquetColumnReader.h b/velox/dwio/parquet/reader/ParquetColumnReader.h
@@ -45,6 +45,7 @@ class ParquetColumnReader {
       const std::shared_ptr<const dwio::common::TypeWithId>& requestedType,
       const std::shared_ptr<const dwio::common::TypeWithId>& fileType,
       ParquetParams& params,
-      common::ScanSpec& scanSpec);
+      common::ScanSpec& scanSpec,
+      memory::MemoryPool& pool);
 };
 } // namespace facebook::velox::parquet
diff --git a/velox/dwio/parquet/reader/ParquetReader.cpp b/velox/dwio/parquet/reader/ParquetReader.cpp
@@ -84,6 +84,11 @@ class ReaderBase {
   /// the data still exists in the buffered inputs.
   bool isRowGroupBuffered(int32_t rowGroupIndex) const;
 
+  static std::shared_ptr<const dwio::common::TypeWithId> createTypeWithId(
+      const std::shared_ptr<const dwio::common::TypeWithId>& inputType,
+      const RowTypePtr& rowTypePtr,
+      bool fileColumnNamesReadAsLowerCase);
+
  private:
   // Reads and parses file footer.
   void loadFileMetaData();
@@ -564,6 +569,33 @@ std::shared_ptr<const RowType> ReaderBase::createRowType(
       std::move(childNames), std::move(childTypes));
 }
 
+std::shared_ptr<const dwio::common::TypeWithId> ReaderBase::createTypeWithId(
+    const std::shared_ptr<const dwio::common::TypeWithId>& inputType,
+    const RowTypePtr& rowTypePtr,
+    bool fileColumnNamesReadAsLowerCase) {
+  if (!fileColumnNamesReadAsLowerCase) {
+    return inputType;
+  }
+  std::vector<std::string> names;
+  names.reserve(rowTypePtr->names().size());
+  std::vector<TypePtr> types = rowTypePtr->children();
+  for (const auto& name : rowTypePtr->names()) {
+    std::string childName = name;
+    folly::toLowerAscii(childName);
+    names.emplace_back(childName);
+  }
+  auto convertedType =
+      TypeFactory<TypeKind::ROW>::create(std::move(names), std::move(types));
+
+  auto children = inputType->getChildren();
+  return std::make_shared<const dwio::common::TypeWithId>(
+      convertedType,
+      std::move(children),
+      inputType->id(),
+      inputType->maxId(),
+      inputType->column());
+}
+
 void ReaderBase::scheduleRowGroups(
     const std::vector<uint32_t>& rowGroupIds,
     int32_t currentGroup,
@@ -630,13 +662,19 @@ ParquetRowReader::ParquetRowReader(
     return; // TODO
   }
   ParquetParams params(pool_, columnReaderStats_, readerBase_->fileMetaData());
-  auto columnSelector = std::make_shared<ColumnSelector>(
-      ColumnSelector::apply(options_.getSelector(), readerBase_->schema()));
+  // ColumnSelector::apply does not work for schema pruning case.
+  auto columnSelector = options_.getSelector() == nullptr
+      ? std::make_shared<ColumnSelector>(ColumnSelector(readerBase_->schema()))
+      : options_.getSelector();
   columnReader_ = ParquetColumnReader::build(
-      columnSelector->getSchemaWithId(),
+      ReaderBase::createTypeWithId(
+          columnSelector->getSchemaWithId(),
+          asRowType(columnSelector->getSchemaWithId()->type()),
+          readerBase_->isFileColumnNamesReadAsLowerCase()),
       readerBase_->schemaWithId(), // Id is schema id
       params,
-      *options_.getScanSpec());
+      *options_.getScanSpec(),
+      pool_);
 
   filterRowGroups();
   if (!rowGroupIds_.empty()) {
diff --git a/velox/dwio/parquet/reader/RepeatedColumnReader.cpp b/velox/dwio/parquet/reader/RepeatedColumnReader.cpp
@@ -33,6 +33,9 @@ PageReader* FOLLY_NULLABLE readLeafRepDefs(
       return nullptr;
     }
     auto pageReader = reader->formatData().as<ParquetData>().reader();
+    if (pageReader == nullptr) {
+      return nullptr;
+    }
     pageReader->decodeRepDefs(numTop);
     return pageReader;
   }
@@ -113,7 +116,8 @@ MapColumnReader::MapColumnReader(
     const std::shared_ptr<const dwio::common::TypeWithId>& requestedType,
     const std::shared_ptr<const dwio::common::TypeWithId>& fileType,
     ParquetParams& params,
-    common::ScanSpec& scanSpec)
+    common::ScanSpec& scanSpec,
+    memory::MemoryPool& pool)
     : dwio::common::SelectiveMapColumnReader(
           requestedType,
           fileType,
@@ -123,9 +127,17 @@ MapColumnReader::MapColumnReader(
   auto& keyChildType = requestedType->childAt(0);
   auto& elementChildType = requestedType->childAt(1);
   keyReader_ = ParquetColumnReader::build(
-      keyChildType, fileType_->childAt(0), params, *scanSpec.children()[0]);
+      keyChildType,
+      fileType_->childAt(0),
+      params,
+      *scanSpec.children()[0],
+      pool);
   elementReader_ = ParquetColumnReader::build(
-      elementChildType, fileType_->childAt(1), params, *scanSpec.children()[1]);
+      elementChildType,
+      fileType_->childAt(1),
+      params,
+      *scanSpec.children()[1],
+      pool);
   reinterpret_cast<const ParquetTypeWithId*>(fileType.get())
       ->makeLevelInfo(levelInfo_);
   children_ = {keyReader_.get(), elementReader_.get()};
@@ -223,15 +235,16 @@ ListColumnReader::ListColumnReader(
     const std::shared_ptr<const dwio::common::TypeWithId>& requestedType,
     const std::shared_ptr<const dwio::common::TypeWithId>& fileType,
     ParquetParams& params,
-    common::ScanSpec& scanSpec)
+    common::ScanSpec& scanSpec,
+    memory::MemoryPool& pool)
     : dwio::common::SelectiveListColumnReader(
           requestedType,
           fileType,
           params,
           scanSpec) {
   auto& childType = requestedType->childAt(0);
   child_ = ParquetColumnReader::build(
-      childType, fileType_->childAt(0), params, *scanSpec.children()[0]);
+      childType, fileType_->childAt(0), params, *scanSpec.children()[0], pool);
   reinterpret_cast<const ParquetTypeWithId*>(fileType.get())
       ->makeLevelInfo(levelInfo_);
   children_ = {child_.get()};
diff --git a/velox/dwio/parquet/reader/RepeatedColumnReader.h b/velox/dwio/parquet/reader/RepeatedColumnReader.h
@@ -59,7 +59,8 @@ class MapColumnReader : public dwio::common::SelectiveMapColumnReader {
       const std::shared_ptr<const dwio::common::TypeWithId>& requestedType,
       const std::shared_ptr<const dwio::common::TypeWithId>& fileType,
       ParquetParams& params,
-      common::ScanSpec& scanSpec);
+      common::ScanSpec& scanSpec,
+      memory::MemoryPool& pool);
 
   void prepareRead(
       vector_size_t offset,
@@ -115,7 +116,8 @@ class ListColumnReader : public dwio::common::SelectiveListColumnReader {
       const std::shared_ptr<const dwio::common::TypeWithId>& requestedType,
       const std::shared_ptr<const dwio::common::TypeWithId>& fileType,
       ParquetParams& params,
-      common::ScanSpec& scanSpec);
+      common::ScanSpec& scanSpec,
+      memory::MemoryPool& pool);
 
   void prepareRead(
       vector_size_t offset,
diff --git a/velox/dwio/parquet/reader/StructColumnReader.cpp b/velox/dwio/parquet/reader/StructColumnReader.cpp
@@ -30,21 +30,46 @@ StructColumnReader::StructColumnReader(
     const std::shared_ptr<const dwio::common::TypeWithId>& requestedType,
     const std::shared_ptr<const dwio::common::TypeWithId>& fileType,
     ParquetParams& params,
-    common::ScanSpec& scanSpec)
+    common::ScanSpec& scanSpec,
+    memory::MemoryPool& pool)
     : SelectiveStructColumnReader(requestedType, fileType, params, scanSpec) {
   auto& childSpecs = scanSpec_->stableChildren();
+  std::vector<int> missingFields;
   for (auto i = 0; i < childSpecs.size(); ++i) {
     auto childSpec = childSpecs[i];
     if (childSpecs[i]->isConstant()) {
       continue;
     }
-    auto childFileType = fileType_->childByName(childSpec->fieldName());
-    auto childRequestedType =
-        requestedType_->childByName(childSpec->fieldName());
+    const auto& fieldName = childSpec->fieldName();
+    if (!fileType_->containsChild(fieldName)) {
+      missingFields.emplace_back(i);
+      continue;
+    }
+    auto childFileType = fileType_->childByName(fieldName);
+    auto childRequestedType = requestedType_->childByName(fieldName);
     addChild(ParquetColumnReader::build(
-        childRequestedType, childFileType, params, *childSpec));
+        childRequestedType, childFileType, params, *childSpec, pool));
     childSpecs[i]->setSubscript(children_.size() - 1);
   }
+
+  if (missingFields.size() > 0) {
+    // Set the struct as null if all the children fields in the output type are
+    // missing and the number of child fields is more than one.
+    if (childSpecs.size() > 1 && missingFields.size() == childSpecs.size()) {
+      scanSpec_->setConstantValue(
+          BaseVector::createNullConstant(requestedType_->type(), 1, &pool));
+    } else {
+      // Set null constant for the missing child field of output type.
+      for (int channel : missingFields) {
+        childSpecs[channel]->setConstantValue(BaseVector::createNullConstant(
+            requestedType_->childByName(childSpecs[channel]->fieldName())
+                ->type(),
+            1,
+            &pool));
+      }
+    }
+  }
+
   auto type = reinterpret_cast<const ParquetTypeWithId*>(fileType_.get());
   if (type->parent()) {
     levelMode_ = reinterpret_cast<const ParquetTypeWithId*>(fileType_.get())
@@ -54,7 +79,10 @@ StructColumnReader::StructColumnReader(
     // this and the child.
     auto child = childForRepDefs_;
     for (;;) {
-      assert(child);
+      if (child == nullptr) {
+        levelMode_ = LevelMode::kNulls;
+        break;
+      }
       if (child->fileType().type()->kind() == TypeKind::ARRAY ||
           child->fileType().type()->kind() == TypeKind::MAP) {
         levelMode_ = LevelMode::kStructOverLists;
diff --git a/velox/dwio/parquet/reader/StructColumnReader.h b/velox/dwio/parquet/reader/StructColumnReader.h
@@ -36,7 +36,8 @@ class StructColumnReader : public dwio::common::SelectiveStructColumnReader {
       const std::shared_ptr<const dwio::common::TypeWithId>& requestedType,
       const std::shared_ptr<const dwio::common::TypeWithId>& fileType,
       ParquetParams& params,
-      common::ScanSpec& scanSpec);
+      common::ScanSpec& scanSpec,
+      memory::MemoryPool& pool);
 
   void read(vector_size_t offset, RowSet rows, const uint64_t* incomingNulls)
       override;
diff --git a/velox/dwio/parquet/tests/examples/contacts.parquet b/velox/dwio/parquet/tests/examples/contacts.parquet
diff --git a/velox/dwio/parquet/tests/reader/ParquetTableScanTest.cpp b/velox/dwio/parquet/tests/reader/ParquetTableScanTest.cpp
