Support struct reading with different schemas

rui-mo · rui-mo · commit c03152bfbead · 2023-08-29T20:03:12.000+08:00
diff --git a/velox/dwio/common/SelectiveStructColumnReader.cpp b/velox/dwio/common/SelectiveStructColumnReader.cpp
@@ -128,8 +128,7 @@ void SelectiveStructColumnReaderBase::read(
     activeRows = outputRows_;
   }
 
-  auto& childSpecs = scanSpec_->children();
-  VELOX_CHECK(!childSpecs.empty());
+  auto& childSpecs = scanSpec_->stableChildren();
   for (size_t i = 0; i < childSpecs.size(); ++i) {
     auto& childSpec = childSpecs[i];
     if (isChildConstant(*childSpec)) {
@@ -214,7 +213,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 {
@@ -294,7 +293,6 @@ void setNullField(vector_size_t size, VectorPtr& field) {
 void SelectiveStructColumnReaderBase::getValues(
     RowSet rows,
     VectorPtr* result) {
-  VELOX_CHECK(!scanSpec_->children().empty());
   VELOX_CHECK(
       *result != nullptr,
       "SelectiveStructColumnReaderBase expects a non-null result");
@@ -329,7 +327,7 @@ void SelectiveStructColumnReaderBase::getValues(
     resultRow->clearNulls(0, rows.size());
   }
   bool lazyPrepared = false;
-  for (auto& childSpec : scanSpec_->children()) {
+  for (auto& childSpec : scanSpec_->stableChildren()) {
     if (!childSpec->projectOut()) {
       continue;
     }
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
@@ -36,7 +36,9 @@ namespace facebook::velox::parquet {
 std::unique_ptr<dwio::common::SelectiveColumnReader> ParquetColumnReader::build(
     const std::shared_ptr<const dwio::common::TypeWithId>& dataType,
     ParquetParams& params,
-    common::ScanSpec& scanSpec) {
+    common::ScanSpec& scanSpec,
+    const TypePtr& outputType,
+    memory::MemoryPool& pool) {
   auto colName = scanSpec.fieldName();
 
   switch (dataType->type()->kind()) {
@@ -56,17 +58,34 @@ std::unique_ptr<dwio::common::SelectiveColumnReader> ParquetColumnReader::build(
           dataType->type(), dataType, params, scanSpec);
 
     case TypeKind::ROW:
-      return std::make_unique<StructColumnReader>(dataType, params, scanSpec);
+      return std::make_unique<StructColumnReader>(
+          dataType,
+          params,
+          scanSpec,
+          outputType ? asRowType(outputType) : nullptr,
+          pool);
 
     case TypeKind::VARBINARY:
     case TypeKind::VARCHAR:
       return std::make_unique<StringColumnReader>(dataType, params, scanSpec);
 
     case TypeKind::ARRAY:
-      return std::make_unique<ListColumnReader>(dataType, params, scanSpec);
+      return std::make_unique<ListColumnReader>(
+          dataType,
+          params,
+          scanSpec,
+          outputType ? std::dynamic_pointer_cast<const ArrayType>(outputType)
+                     : nullptr,
+          pool);
 
     case TypeKind::MAP:
-      return std::make_unique<MapColumnReader>(dataType, params, scanSpec);
+      return std::make_unique<MapColumnReader>(
+          dataType,
+          params,
+          scanSpec,
+          outputType ? std::dynamic_pointer_cast<const MapType>(outputType)
+                     : nullptr,
+          pool);
 
     case TypeKind::BOOLEAN:
       return std::make_unique<BooleanColumnReader>(dataType, params, scanSpec);
diff --git a/velox/dwio/parquet/reader/ParquetColumnReader.h b/velox/dwio/parquet/reader/ParquetColumnReader.h
@@ -44,6 +44,8 @@ class ParquetColumnReader {
   static std::unique_ptr<dwio::common::SelectiveColumnReader> build(
       const std::shared_ptr<const dwio::common::TypeWithId>& dataType,
       ParquetParams& params,
-      common::ScanSpec& scanSpec);
+      common::ScanSpec& scanSpec,
+      const TypePtr& outputType,
+      memory::MemoryPool& pool);
 };
 } // namespace facebook::velox::parquet
diff --git a/velox/dwio/parquet/reader/ParquetReader.cpp b/velox/dwio/parquet/reader/ParquetReader.cpp
@@ -85,6 +85,16 @@ class ReaderBase {
       int32_t rowGroupIndex,
       const dwio::common::TypeWithId& type) const;
 
+  /// @brief Convert the names of row type to lower case when
+  /// fileColumnNamesReadAsLowerCase is true.
+  /// @param rowTypePtr the input row type.
+  /// @param fileColumnNamesReadAsLowerCase whether to convert names into lower
+  /// case.
+  /// @return row type with names converted.
+  static std::shared_ptr<const RowType> convertRowTypeNames(
+      const RowTypePtr& rowTypePtr,
+      bool fileColumnNamesReadAsLowerCase);
+
  private:
   // Reads and parses file footer.
   void loadFileMetaData();
@@ -547,22 +557,39 @@ TypePtr ReaderBase::convertType(
   }
 }
 
+std::shared_ptr<const RowType> ReaderBase::convertRowTypeNames(
+    const RowTypePtr& rowTypePtr,
+    bool fileColumnNamesReadAsLowerCase) {
+  if (!fileColumnNamesReadAsLowerCase) {
+    return rowTypePtr;
+  }
+  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);
+  }
+  return TypeFactory<TypeKind::ROW>::create(std::move(names), std::move(types));
+}
+
 std::shared_ptr<const RowType> ReaderBase::createRowType(
     std::vector<std::shared_ptr<const ParquetTypeWithId::TypeWithId>> children,
     bool fileColumnNamesReadAsLowerCase) {
   std::vector<std::string> childNames;
+  childNames.reserve(children.size());
   std::vector<TypePtr> childTypes;
-  for (auto& child : children) {
-    auto childName =
-        std::static_pointer_cast<const ParquetTypeWithId>(child)->name_;
-    if (fileColumnNamesReadAsLowerCase) {
-      folly::toLowerAscii(childName);
-    }
-    childNames.push_back(std::move(childName));
-    childTypes.push_back(child->type());
+  childTypes.reserve(children.size());
+  for (const auto& child : children) {
+    childNames.emplace_back(
+        std::static_pointer_cast<const ParquetTypeWithId>(child)->name_);
+    childTypes.emplace_back(child->type());
   }
-  return TypeFactory<TypeKind::ROW>::create(
-      std::move(childNames), std::move(childTypes));
+  return convertRowTypeNames(
+      TypeFactory<TypeKind::ROW>::create(
+          std::move(childNames), std::move(childTypes)),
+      fileColumnNamesReadAsLowerCase);
 }
 
 void ReaderBase::scheduleRowGroups(
@@ -640,7 +667,11 @@ ParquetRowReader::ParquetRowReader(
   columnReader_ = ParquetColumnReader::build(
       readerBase_->schemaWithId(), // Id is schema id
       params,
-      *options_.getScanSpec());
+      *options_.getScanSpec(),
+      ReaderBase::convertRowTypeNames(
+          asRowType(options_.getSelector()->getSchemaWithId()->type()),
+          readerBase_->isFileColumnNamesReadAsLowerCase()),
+      pool_);
 
   filterRowGroups();
   if (!rowGroupIds_.empty()) {
diff --git a/velox/dwio/parquet/reader/RepeatedColumnReader.cpp b/velox/dwio/parquet/reader/RepeatedColumnReader.cpp
@@ -31,6 +31,9 @@ PageReader* FOLLY_NULLABLE readLeafRepDefs(
       return nullptr;
     }
     auto pageReader = reader->formatData().as<ParquetData>().reader();
+    if (pageReader == nullptr) {
+      return nullptr;
+    }
     pageReader->decodeRepDefs(numTop);
     return pageReader;
   }
@@ -110,18 +113,28 @@ void ensureRepDefs(
 MapColumnReader::MapColumnReader(
     std::shared_ptr<const dwio::common::TypeWithId> requestedType,
     ParquetParams& params,
-    common::ScanSpec& scanSpec)
+    common::ScanSpec& scanSpec,
+    const std::shared_ptr<const MapType>& outputType,
+    memory::MemoryPool& pool)
     : dwio::common::SelectiveMapColumnReader(
           requestedType,
           requestedType,
           params,
           scanSpec) {
   auto& keyChildType = requestedType->childAt(0);
   auto& elementChildType = requestedType->childAt(1);
-  keyReader_ =
-      ParquetColumnReader::build(keyChildType, params, *scanSpec.children()[0]);
+  keyReader_ = ParquetColumnReader::build(
+      keyChildType,
+      params,
+      *scanSpec.children()[0],
+      outputType ? outputType->keyType() : nullptr,
+      pool);
   elementReader_ = ParquetColumnReader::build(
-      elementChildType, params, *scanSpec.children()[1]);
+      elementChildType,
+      params,
+      *scanSpec.children()[1],
+      outputType ? outputType->valueType() : nullptr,
+      pool);
   reinterpret_cast<const ParquetTypeWithId*>(requestedType.get())
       ->makeLevelInfo(levelInfo_);
   children_ = {keyReader_.get(), elementReader_.get()};
@@ -218,15 +231,21 @@ void MapColumnReader::filterRowGroups(
 ListColumnReader::ListColumnReader(
     std::shared_ptr<const dwio::common::TypeWithId> requestedType,
     ParquetParams& params,
-    common::ScanSpec& scanSpec)
+    common::ScanSpec& scanSpec,
+    const std::shared_ptr<const ArrayType>& outputType,
+    memory::MemoryPool& pool)
     : dwio::common::SelectiveListColumnReader(
           requestedType,
           requestedType,
           params,
           scanSpec) {
   auto& childType = requestedType->childAt(0);
-  child_ =
-      ParquetColumnReader::build(childType, params, *scanSpec.children()[0]);
+  child_ = ParquetColumnReader::build(
+      childType,
+      params,
+      *scanSpec.children()[0],
+      outputType ? outputType->elementType() : nullptr,
+      pool);
   reinterpret_cast<const ParquetTypeWithId*>(requestedType.get())
       ->makeLevelInfo(levelInfo_);
   children_ = {child_.get()};
diff --git a/velox/dwio/parquet/reader/RepeatedColumnReader.h b/velox/dwio/parquet/reader/RepeatedColumnReader.h
@@ -58,7 +58,9 @@ class MapColumnReader : public dwio::common::SelectiveMapColumnReader {
   MapColumnReader(
       std::shared_ptr<const dwio::common::TypeWithId> requestedType,
       ParquetParams& params,
-      common::ScanSpec& scanSpec);
+      common::ScanSpec& scanSpec,
+      const std::shared_ptr<const MapType>& outputType,
+      memory::MemoryPool& pool);
 
   void prepareRead(
       vector_size_t offset,
@@ -113,7 +115,9 @@ class ListColumnReader : public dwio::common::SelectiveListColumnReader {
   ListColumnReader(
       std::shared_ptr<const dwio::common::TypeWithId> requestedType,
       ParquetParams& params,
-      common::ScanSpec& scanSpec);
+      common::ScanSpec& scanSpec,
+      const std::shared_ptr<const ArrayType>& outputType,
+      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
@@ -22,18 +22,47 @@ namespace facebook::velox::parquet {
 StructColumnReader::StructColumnReader(
     const std::shared_ptr<const dwio::common::TypeWithId>& dataType,
     ParquetParams& params,
-    common::ScanSpec& scanSpec)
+    common::ScanSpec& scanSpec,
+    const RowTypePtr& outputType,
+    memory::MemoryPool& pool)
     : SelectiveStructColumnReader(dataType, dataType, params, scanSpec) {
   auto& childSpecs = scanSpec_->stableChildren();
+  std::vector<int> missingFields;
   for (auto i = 0; i < childSpecs.size(); ++i) {
     if (childSpecs[i]->isConstant()) {
       continue;
     }
-    auto childDataType = fileType_->childByName(childSpecs[i]->fieldName());
+    const auto& fieldName = childSpecs[i]->fieldName();
+    if (outputType && !fileType_->containsChild(fieldName)) {
+      missingFields.emplace_back(i);
+      continue;
+    }
+    auto childDataType = fileType_->childByName(fieldName);
 
-    addChild(ParquetColumnReader::build(childDataType, params, *childSpecs[i]));
+    addChild(ParquetColumnReader::build(
+        childDataType,
+        params,
+        *childSpecs[i],
+        outputType ? outputType->findChild(fieldName) : nullptr,
+        pool));
     childSpecs[i]->setSubscript(children_.size() - 1);
   }
+
+  if (outputType) {
+    // 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(outputType, 1, &pool));
+    } else {
+      // Set null constant for the missing child field of output type.
+      for (int channel : missingFields) {
+        childSpecs[channel]->setConstantValue(BaseVector::createNullConstant(
+            outputType->findChild(childSpecs[channel]->fieldName()), 1, &pool));
+      }
+    }
+  }
+
   auto type = reinterpret_cast<const ParquetTypeWithId*>(fileType_.get());
   if (type->parent()) {
     levelMode_ = reinterpret_cast<const ParquetTypeWithId*>(fileType_.get())
@@ -43,7 +72,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
@@ -26,7 +26,9 @@ class StructColumnReader : public dwio::common::SelectiveStructColumnReader {
   StructColumnReader(
       const std::shared_ptr<const dwio::common::TypeWithId>& dataType,
       ParquetParams& params,
-      common::ScanSpec& scanSpec);
+      common::ScanSpec& scanSpec,
+      const RowTypePtr& outputType,
+      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
