map/struct/list

facebookincubator · Oct 26, 2023 · c042e46 · c042e46
1 parent 48a8297
commit c042e46
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 12 deletions.
diff --git a/velox/dwio/common/compression/Compression.cpp b/velox/dwio/common/compression/Compression.cpp
@@ -684,7 +684,8 @@ createAsyncDecompressor(
     case CompressionKind::CompressionKind_GZIP:
       return std::make_unique<GzipIAADecompressor>(bufferSize, streamDebugInfo);
     default:
-      LOG(WARNING) << "Asynchronous mode not support for compression codec  " << kind;
+      LOG(WARNING) << "Asynchronous mode not support for compression codec  "
+                   << kind;
       return nullptr;
   }
   return nullptr;

diff --git a/velox/dwio/parquet/reader/PageReader.cpp b/velox/dwio/parquet/reader/PageReader.cpp
@@ -1177,7 +1177,7 @@ void PageReader::prefetchDictionary(const thrift::PageHeader& pageHeader) {
   return;
 }
 
-const bool FOLLY_NONNULL PageReader::iaaDecompressGzip(
+const bool PageReader::iaaDecompressGzip(
     const char* pageData,
     uint32_t compressedSize,
     uint32_t uncompressedSize,
@@ -1218,7 +1218,7 @@ const bool FOLLY_NONNULL PageReader::iaaDecompressGzip(
   return true;
 }
 
-const bool FOLLY_NONNULL PageReader::getDecompRes(int job_id) {
+const bool PageReader::getDecompRes(int job_id) {
   auto streamDebugInfo =
       fmt::format("Page Reader: Stream {}", inputStream_->getName());
   std::unique_ptr<dwio::common::compression::AsyncDecompressor> decompressor =
@@ -1228,7 +1228,7 @@ const bool FOLLY_NONNULL PageReader::getDecompRes(int job_id) {
 }
 
 PageReader::~PageReader() {
-  if (data_qpl_job_id >= 0 || dict_qpl_job_id >= 0) {
+  if (data_qpl_job_id > 0 || dict_qpl_job_id > 0) {
     auto streamDebugInfo =
         fmt::format("Page Reader: Stream {}", inputStream_->getName());
     std::unique_ptr<dwio::common::compression::AsyncDecompressor> decompressor =

diff --git a/velox/dwio/parquet/reader/PageReader.h b/velox/dwio/parquet/reader/PageReader.h
@@ -129,10 +129,10 @@ class PageReader {
   void prefetchDataPageV1(const thrift::PageHeader& pageHeader);
   void prefetchDataPageV2(const thrift::PageHeader& pageHeader);
   void prefetchDictionary(const thrift::PageHeader& pageHeader);
-  const bool FOLLY_NONNULL getDecompRes(int job_id);
+  const bool getDecompRes(int job_id);
   void prefetchNextPage();
   bool seekToPreDecompPage(int64_t row);
-  const bool FOLLY_NONNULL iaaDecompressGzip(
+  const bool iaaDecompressGzip(
       const char* FOLLY_NONNULL pageData,
       uint32_t compressedSize,
       uint32_t uncompressedSize,

diff --git a/velox/dwio/parquet/reader/RepeatedColumnReader.cpp b/velox/dwio/parquet/reader/RepeatedColumnReader.cpp
@@ -81,6 +81,21 @@ void skipUnreadLengthsAndNulls(dwio::common::SelectiveColumnReader& reader) {
   }
 }
 
+void preDecompChild(
+    dwio::common::SelectiveColumnReader* reader,
+    uint32_t index) {
+#ifdef VELOX_ENABLE_QPL
+  auto children = reader->children();
+  if (children.empty()) {
+    reader->formatData().as<ParquetData>().preDecompRowGroup(index);
+    return;
+  }
+  for (auto* child : children) {
+    preDecompChild(child, index);
+  }
+#endif
+}
+
 void enqueueChildren(
     dwio::common::SelectiveColumnReader* reader,
     uint32_t index,
@@ -135,6 +150,10 @@ void MapColumnReader::enqueueRowGroup(
   enqueueChildren(this, index, input);
 }
 
+void MapColumnReader::preDecompRowGroup(uint32_t index) {
+  preDecompChild(this, index);
+}
+
 void MapColumnReader::seekToRowGroup(uint32_t index) {
   SelectiveMapColumnReader::seekToRowGroup(index);
   readOffset_ = 0;
@@ -241,6 +260,10 @@ void ListColumnReader::enqueueRowGroup(
   enqueueChildren(this, index, input);
 }
 
+void ListColumnReader::preDecompRowGroup(uint32_t index) {
+  preDecompChild(this, index);
+}
+
 void ListColumnReader::seekToRowGroup(uint32_t index) {
   SelectiveListColumnReader::seekToRowGroup(index);
   readOffset_ = 0;

diff --git a/velox/dwio/parquet/reader/RepeatedColumnReader.h b/velox/dwio/parquet/reader/RepeatedColumnReader.h
@@ -71,6 +71,7 @@ class MapColumnReader : public dwio::common::SelectiveMapColumnReader {
   void seekToRowGroup(uint32_t index) override;
 
   void enqueueRowGroup(uint32_t index, dwio::common::BufferedInput& input);
+  void preDecompRowGroup(uint32_t index);
 
   void read(
       vector_size_t offset,
@@ -128,9 +129,7 @@ class ListColumnReader : public dwio::common::SelectiveListColumnReader {
 
   void enqueueRowGroup(uint32_t index, dwio::common::BufferedInput& input);
 
-  bool preDecompRowGroup(uint32_t index) {
-    return true;
-  }
+  void preDecompRowGroup(uint32_t index);
   void read(
       vector_size_t offset,
       RowSet rows,

diff --git a/velox/dwio/parquet/reader/StructColumnReader.cpp b/velox/dwio/parquet/reader/StructColumnReader.cpp
@@ -139,11 +139,11 @@ bool StructColumnReader::preDecompRowGroup(uint32_t index) {
       return false;
     }
     if (auto structChild = dynamic_cast<StructColumnReader*>(child)) {
-      continue;
+      structChild->preDecompRowGroup(index);
     } else if (auto listChild = dynamic_cast<ListColumnReader*>(child)) {
-      continue;
+      listChild->preDecompRowGroup(index);
     } else if (auto mapChild = dynamic_cast<MapColumnReader*>(child)) {
-      continue;
+      mapChild->preDecompRowGroup(index);
     } else {
       needPreDecomp =
           child->formatData().as<ParquetData>().preDecompRowGroup(index);