feat: Add stripe stats pruning

Summary:
Add a stripe-level vectorized stats optional section for Nimble files and use it in the selective reader to conservatively skip stripes for integral filters when min/max proves no rows can match.

The writer snapshots per-stripe column stats and reconstructs file-level stats by merging those snapshots, preserving a fallback path for files that do not emit stripe snapshots. The reader preloads and deserializes the new optional stats section while treating absent stats as a normal full-scan fallback.

Deduplicated columns are now handled correctly across stripes. mergeColumnStats merges into the wrapped base statistics (which back the delegating getters) and accumulates the dedup counters, instead of writing to unused fields. The per-stripe ancestor deduplication wrapping applied during finalization is undone between stripes (releasing the base collector) so every stripe finalizes from an identical collector structure; without this, stripe 1+ skipped the child logical-size rollup and the reconstructed file-level stats were wrong (tripping the raw-size consistency check). Genuine deduplicated nodes configured at init stay wrapped.

Differential Revision: D112143149

Author: apurva-meta

dwio/nimble/tablet/Constants.h

@@ -34,6 +34,7 @@ constexpr std::string_view kMetadataSection = "columnar.metadata";
 constexpr std::string_view kStatsSection = "columnar.stats";
 constexpr std::string_view kVectorizedStatsSection =
     "columnar.vectorized_stats";
+constexpr std::string_view kStripeStatsSection = "columnar.stripe_stats";
 constexpr std::string_view kClusterIndexSection = "columnar.cluster.index";
 constexpr std::string_view kChunkIndexSection = "columnar.chunk.index";
 constexpr std::string_view kHashIndexSection = "columnar.hash.index";

dwio/nimble/tablet/TabletReader.cpp

@@ -74,7 +74,9 @@ TabletReader::Options TabletReader::configureOptions(
   Options tabletOptions;
   tabletOptions.maxFooterIoBytes = options.footerSpeculativeIoSize();
   tabletOptions.preloadOptionalSections = {
-      std::string(kSchemaSection), std::string(kVectorizedStatsSection)};
+      std::string(kSchemaSection),
+      std::string(kVectorizedStatsSection),
+      std::string(kStripeStatsSection)};
   tabletOptions.loadClusterIndex = options.loadClusterIndex();
   if (tabletOptions.loadClusterIndex) {
     tabletOptions.preloadOptionalSections.emplace_back(kClusterIndexSection);

dwio/nimble/velox/FieldWriter.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <cmath>
 #include <set>
 
 #include <folly/Executor.h>
@@ -299,22 +300,19 @@ class FieldWriterContext {
   }
 
   inline std::vector<ColumnStatistics*> columnStats() {
-    if (!statsFinalized_) {
-      return {};
-    }
-    // TODO: add a build method for this pattern.
     std::vector<ColumnStatistics*> statsViews;
-    statsViews.reserve(statsCollectors_.size());
-    for (auto& collector : statsCollectors_) {
-      // FIXME: don't need this branching.
-      statsViews.push_back(
-          collector->shared()
-              ? collector->as<SharedStatisticsCollector>()->getBaseStatistics()
-              : collector->getStatsView());
+    statsViews.reserve(fileStats_.size());
+    for (auto& stat : fileStats_) {
+      statsViews.push_back(stat.get());
     }
     return statsViews;
   }
 
+  inline const std::vector<std::vector<std::unique_ptr<ColumnStatistics>>>&
+  stripeStats() const {
+    return stripeStats_;
+  }
+
   void handleFlatmapFieldAddEvent(
       const TypeBuilder& typeBuilder,
       std::string_view flatmapKey,
@@ -479,10 +477,40 @@ class FieldWriterContext {
   // 1. roll up logical and physical sizes
   // 2. wrap all ancestors of deduplicated types
   // 3. backfill both known and newly wrapped deduplicated stats
-  void finalizeStatsCollectors() {
+  void finalizeStripeStatsCollectors() {
     NIMBLE_CHECK(!statsFinalized_);
     finalizeStatsCollector(schemaWithId_);
     statsFinalized_ = true;
+    snapshotStripeStats();
+    resetStatsCollectors();
+  }
+
+  void finalizeFileStatsFromStripes() {
+    fileStats_.clear();
+    if (stripeStats_.empty()) {
+      if (!statsFinalized_) {
+        finalizeStatsCollector(schemaWithId_);
+        statsFinalized_ = true;
+      }
+      fileStats_.reserve(statsCollectors_.size());
+      for (auto& collector : statsCollectors_) {
+        auto* stats = collector->shared()
+            ? collector->as<SharedStatisticsCollector>()->getBaseStatistics()
+            : collector->getStatsView();
+        fileStats_.push_back(stats->clone());
+      }
+      return;
+    }
+    fileStats_.reserve(stripeStats_.front().size());
+    for (const auto& stat : stripeStats_.front()) {
+      fileStats_.push_back(stat->clone());
+    }
+    for (size_t stripe = 1; stripe < stripeStats_.size(); ++stripe) {
+      NIMBLE_CHECK_EQ(stripeStats_[stripe].size(), fileStats_.size());
+      for (size_t column = 0; column < fileStats_.size(); ++column) {
+        mergeColumnStats(*fileStats_[column], *stripeStats_[stripe][column]);
+      }
+    }
   }
 
  protected:
@@ -513,6 +541,157 @@ class FieldWriterContext {
       [](TypeBuilder&, uint32_t) {}};
 
  private:
+  // Merges min/max from source into target for a scalar statistics subclass
+  // (Integral/FloatingPoint/String). Keeps the widest range across stripes.
+  template <typename StatsT>
+  static void mergeMinMax(
+      ColumnStatistics& target,
+      const ColumnStatistics& source) {
+    auto* targetStats = target.as<StatsT>();
+    const auto* sourceStats = source.as<StatsT>();
+    if (sourceStats->getMin().has_value() &&
+        (!targetStats->getMin().has_value() ||
+         *sourceStats->getMin() < *targetStats->getMin())) {
+      targetStats->min_ = sourceStats->getMin();
+    }
+    if (sourceStats->getMax().has_value() &&
+        (!targetStats->getMax().has_value() ||
+         *sourceStats->getMax() > *targetStats->getMax())) {
+      targetStats->max_ = sourceStats->getMax();
+    }
+  }
+
+  static bool shouldUseSourceFloatingPointMin(
+      const FloatingPointStatistics& target,
+      const FloatingPointStatistics& source) {
+    if (!source.getMin().has_value()) {
+      return false;
+    }
+    if (!target.getMin().has_value()) {
+      return true;
+    }
+    if (std::isnan(*target.getMin())) {
+      return false;
+    }
+    if (std::isnan(*source.getMin())) {
+      return true;
+    }
+    return *source.getMin() < *target.getMin();
+  }
+
+  static bool shouldUseSourceFloatingPointMax(
+      const FloatingPointStatistics& target,
+      const FloatingPointStatistics& source) {
+    if (!source.getMax().has_value()) {
+      return false;
+    }
+    if (!target.getMax().has_value()) {
+      return true;
+    }
+    if (std::isnan(*target.getMax())) {
+      return false;
+    }
+    if (std::isnan(*source.getMax())) {
+      return true;
+    }
+    return *source.getMax() > *target.getMax();
+  }
+
+  static void mergeFloatingPointMinMax(
+      ColumnStatistics& target,
+      const ColumnStatistics& source) {
+    auto* targetStats = target.as<FloatingPointStatistics>();
+    const auto* sourceStats = source.as<FloatingPointStatistics>();
+    if (shouldUseSourceFloatingPointMin(*targetStats, *sourceStats)) {
+      targetStats->min_ = sourceStats->getMin();
+    }
+    if (shouldUseSourceFloatingPointMax(*targetStats, *sourceStats)) {
+      targetStats->max_ = sourceStats->getMax();
+    }
+  }
+
+  static void mergeColumnStats(
+      ColumnStatistics& target,
+      const ColumnStatistics& source) {
+    NIMBLE_CHECK(
+        target.getType() == source.getType(),
+        "Merging stats with mismatched types: {} vs {}.",
+        static_cast<uint8_t>(target.getType()),
+        static_cast<uint8_t>(source.getType()));
+    // Deduplicated stats delegate their base metrics (value/null count, sizes)
+    // to the wrapped base statistics, so merge into that base rather than
+    // target's own (unused) member fields, and accumulate the dedup-specific
+    // counters separately.
+    if (target.getType() == StatType::DEDUPLICATED) {
+      auto* targetStats = target.as<DeduplicatedColumnStatistics>();
+      const auto* sourceStats = source.as<DeduplicatedColumnStatistics>();
+      targetStats->dedupedCount_ += sourceStats->getDedupedCount();
+      targetStats->dedupedLogicalSize_ += sourceStats->getDedupedLogicalSize();
+      auto* targetBase = targetStats->getBaseStatistics();
+      const auto* sourceBase = sourceStats->getBaseStatistics();
+      NIMBLE_CHECK(
+          (targetBase == nullptr) == (sourceBase == nullptr),
+          "Mismatched deduplicated stats structure during merge.");
+      if (targetBase != nullptr && sourceBase != nullptr) {
+        mergeColumnStats(*targetBase, *sourceBase);
+      }
+      return;
+    }
+
+    target.valueCount_ += source.getValueCount();
+    target.nullCount_ += source.getNullCount();
+    target.logicalSize_ += source.getLogicalSize();
+    target.physicalSize_ += source.getPhysicalSize();
+    switch (target.getType()) {
+      case StatType::INTEGRAL:
+        mergeMinMax<IntegralStatistics>(target, source);
+        break;
+      case StatType::FLOATING_POINT:
+        mergeFloatingPointMinMax(target, source);
+        break;
+      case StatType::STRING:
+        mergeMinMax<StringStatistics>(target, source);
+        break;
+      case StatType::DEFAULT:
+        break;
+      case StatType::DEDUPLICATED:
+        NIMBLE_UNREACHABLE("Deduplicated stats are handled above.");
+    }
+  }
+
+  void snapshotStripeStats() {
+    auto& stripe = stripeStats_.emplace_back();
+    stripe.reserve(statsCollectors_.size());
+    for (auto& collector : statsCollectors_) {
+      auto* stats = collector->shared()
+          ? collector->as<SharedStatisticsCollector>()->getBaseStatistics()
+          : collector->getStatsView();
+      stripe.push_back(stats->clone());
+    }
+  }
+
+  void resetStatsCollectors() {
+    // finalizeStatsCollector() wraps ancestors of deduplicated nodes into a
+    // DeduplicatedStatisticsCollector in place. That wrapping must be undone
+    // between stripes; otherwise the next stripe's finalization would take the
+    // "already deduplicated" branch and skip rolling child logical/physical
+    // sizes up into the ancestor, corrupting the file-level stats
+    // reconstruction. Genuine deduplicated nodes (configured at init via
+    // dictionary array / deduplicated map ids) must remain wrapped.
+    for (uint32_t id = 0; id < statsCollectors_.size(); ++id) {
+      auto& collector = statsCollectors_[id];
+      if (collector->getType() == StatType::DEDUPLICATED &&
+          !hasDictionaryArrayNodeId(id) && !hasDeduplicatedMapNodeId(id)) {
+        collector = collector->asChecked<DeduplicatedStatisticsCollector>()
+                        ->releaseBaseCollector();
+      }
+    }
+    for (auto& collector : statsCollectors_) {
+      collector->reset();
+    }
+    statsFinalized_ = false;
+  }
+
   void finalizeStatsCollector(
       const std::shared_ptr<const velox::dwio::common::TypeWithId>& type) {
     auto statsCollector = statsCollectors_[type->id()].get();
@@ -627,6 +806,8 @@ class FieldWriterContext {
   std::vector<std::pair<uint32_t, std::unique_ptr<StreamData>>> streams_;
   std::shared_ptr<const velox::dwio::common::TypeWithId> schemaWithId_;
   std::vector<std::unique_ptr<StatisticsCollector>> statsCollectors_;
+  std::vector<std::vector<std::unique_ptr<ColumnStatistics>>> stripeStats_;
+  std::vector<std::unique_ptr<ColumnStatistics>> fileStats_;
   bool statsFinalized_{false};
 };
 

dwio/nimble/velox/VeloxWriter.cpp

@@ -954,6 +954,7 @@ void VeloxWriter::writeMetadata() {
 }
 
 void VeloxWriter::writeColumnStats() {
+  context_->finalizeFileStatsFromStripes();
   // When enableStatsConsistencyCheck is true, verify that fileRawSize
   // (accumulated via RawSizeUtils) matches the root column statistics.
   if (context_->options().enableStatsConsistencyCheck) {
@@ -969,6 +970,12 @@ void VeloxWriter::writeColumnStats() {
     Buffer buffer{*encodingMemoryPool_};
     tabletWriter_->writeOptionalSection(
         std::string(kVectorizedStatsSection), fileStats.serialize(buffer));
+    VectorizedStripeStats stripeStats{
+        context_->stripeStats(), encodingMemoryPool_.get()};
+    Buffer stripeStatsBuffer{*encodingMemoryPool_};
+    tabletWriter_->writeOptionalSection(
+        std::string(kStripeStatsSection),
+        stripeStats.serialize(stripeStatsBuffer));
   } else {
     flatbuffers::FlatBufferBuilder builder;
     builder.Finish(
@@ -1086,9 +1093,6 @@ void VeloxWriter::close() {
     try {
       writeStripe();
       rootWriter_->close();
-      if (context_->options().enableStatsCollection) {
-        context_->finalizeStatsCollectors();
-      }
 
       writeMetadata();
       if (context_->options().enableStatsCollection) {
@@ -1553,6 +1557,9 @@ bool VeloxWriter::writeStripe() {
   } else {
     writeStreams();
   }
+  if (context_->options().enableStatsCollection) {
+    context_->finalizeStripeStatsCollectors();
+  }
 
   uint64_t stripeSize{0};
   {

dwio/nimble/velox/selective/ReaderBase.cpp

@@ -153,7 +153,22 @@ ReaderBase::ReaderBase(
           return {};
         }
         return fileStats->toColumnStatistics(fileSchema_, nimbleSchema_);
-      }()} {}
+      }()},
+      stripeColumnStats_{
+          [&]() -> std::vector<std::vector<std::unique_ptr<ColumnStatistics>>> {
+            auto statsSection =
+                tablet_->loadOptionalSection(std::string(kStripeStatsSection));
+            if (!statsSection.has_value()) {
+              return {};
+            }
+            auto stripeStats = VectorizedStripeStats::deserialize(
+                statsSection->content(), *pool_);
+            if (!stripeStats) {
+              return {};
+            }
+            return stripeStats->takeStripeColumnStatistics(
+                fileSchema_, nimbleSchema_);
+          }()} {}
 
 void LazyInput::load() {
   if (!loaded_) {

dwio/nimble/velox/selective/ReaderBase.h

@@ -118,6 +118,11 @@ class ReaderBase {
     return fileColumnStats_;
   }
 
+  const std::vector<std::vector<std::unique_ptr<ColumnStatistics>>>&
+  stripeColumnStats() const {
+    return stripeColumnStats_;
+  }
+
  private:
   ReaderBase(
       std::unique_ptr<velox::dwio::common::BufferedInput> input,
@@ -138,6 +143,8 @@ class ReaderBase {
   // File-level column statistics deserialized from the vectorized stats
   // optional section at construction.
   const std::vector<std::unique_ptr<ColumnStatistics>> fileColumnStats_;
+  const std::vector<std::vector<std::unique_ptr<ColumnStatistics>>>
+      stripeColumnStats_;
   mutable std::shared_ptr<const velox::dwio::common::TypeWithId>
       fileSchemaWithId_;
 };