Fix ORC related failed UT

oap-project · Nov 21, 2023 · 0a7bb2c · 0a7bb2c
1 parent 8d8b7d5
commit 0a7bb2c
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 2 deletions.
diff --git a/velox/connectors/hive/HiveConnector.cpp b/velox/connectors/hive/HiveConnector.cpp
@@ -150,6 +150,7 @@ std::unique_ptr<core::PartitionFunction> HivePartitionFunctionSpec::create(
 void HiveConnectorFactory::initialize() {
   static bool once = []() {
     dwio::common::registerFileSinks();
+    dwrf::registerOrcReaderFactory();
     dwrf::registerDwrfReaderFactory();
     dwrf::registerDwrfWriterFactory();
 // Meta's buck build system needs this check.

diff --git a/velox/dwio/dwrf/common/FileMetadata.h b/velox/dwio/dwrf/common/FileMetadata.h
@@ -426,7 +426,8 @@ class FooterWrapper : public ProtoWrapperBase {
 
   // TODO: ORC has not supported column statistics yet
   int statisticsSize() const {
-    return format_ == DwrfFormat::kDwrf ? dwrfPtr()->statistics_size() : 0;
+    return format_ == DwrfFormat::kDwrf ? dwrfPtr()->statistics_size()
+                                        : orcPtr()->statistics_size();
   }
 
   const ::google::protobuf::RepeatedPtrField<
@@ -438,7 +439,6 @@ class FooterWrapper : public ProtoWrapperBase {
 
   const ::facebook::velox::dwrf::proto::ColumnStatistics& statistics(
       int index) const {
-    VELOX_CHECK_EQ(format_, DwrfFormat::kDwrf);
     return dwrfPtr()->statistics(index);
   }
 

diff --git a/velox/dwio/dwrf/reader/DwrfReader.cpp b/velox/dwio/dwrf/reader/DwrfReader.cpp
@@ -1108,4 +1108,12 @@ void unregisterDwrfReaderFactory() {
   dwio::common::unregisterReaderFactory(dwio::common::FileFormat::DWRF);
 }
 
+void registerOrcReaderFactory() {
+  dwio::common::registerReaderFactory(std::make_shared<OrcReaderFactory>());
+}
+
+void unregisterOrcReaderFactory() {
+  dwio::common::unregisterReaderFactory(dwio::common::FileFormat::ORC);
+}
+
 } // namespace facebook::velox::dwrf
diff --git a/velox/dwio/dwrf/reader/DwrfReader.h b/velox/dwio/dwrf/reader/DwrfReader.h
@@ -365,8 +365,23 @@ class DwrfReaderFactory : public dwio::common::ReaderFactory {
   }
 };
 
+class OrcReaderFactory : public dwio::common::ReaderFactory {
+ public:
+  OrcReaderFactory() : ReaderFactory(dwio::common::FileFormat::ORC) {}
+
+  std::unique_ptr<dwio::common::Reader> createReader(
+      std::unique_ptr<dwio::common::BufferedInput> input,
+      const dwio::common::ReaderOptions& options) override {
+    return DwrfReader::create(std::move(input), options);
+  }
+};
+
 void registerDwrfReaderFactory();
 
 void unregisterDwrfReaderFactory();
 
+void registerOrcReaderFactory();
+
+void unregisterOrcReaderFactory();
+
 } // namespace facebook::velox::dwrf