feat: Add stripe stats pruning (#970)

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

@@ -299,22 +299,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 +476,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 +540,100 @@ 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 void mergeColumnStats(
+      ColumnStatistics& target,
+      const ColumnStatistics& source) {
+    NIMBLE_CHECK_EQ(target.getType(), 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();
+      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:
+        mergeMinMax<FloatingPointStatistics>(target, source);
+        break;
+      case StatType::STRING:
+        mergeMinMax<StringStatistics>(target, source);
+        break;
+      case StatType::DEFAULT:
+      case StatType::DEDUPLICATED:
+        break;
+    }
+  }
+
+  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 +748,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,32 @@ 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 {};
+            }
+            const auto& stats = stripeStats->toStripeColumnStatistics(
+                fileSchema_, nimbleSchema_);
+            std::vector<std::vector<std::unique_ptr<ColumnStatistics>>> result;
+            result.reserve(stats.size());
+            for (const auto& stripe : stats) {
+              auto& resultStripe = result.emplace_back();
+              resultStripe.reserve(stripe.size());
+              for (const auto& stat : stripe) {
+                resultStripe.push_back(stat->clone());
+              }
+            }
+            return result;
+          }()} {}
 
 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_;
 };

dwio/nimble/velox/selective/SelectiveNimbleReader.cpp

@@ -204,6 +204,8 @@ class SelectiveNimbleRowReader : public dwio::common::RowReader {
   // if stats are available.
   void computeStatsBasedRowSize() const;
 
+  bool shouldSkipStripe(uint32_t stripe) const;
+
   // Computes which top-level columns should use lazy I/O based on the scan
   // spec and remaining filter columns. Returns a const set used for the
   // lifetime of this reader.
@@ -272,7 +274,14 @@ int64_t SelectiveNimbleRowReader::nextRowNumber() {
           readerBase_->randomSkip()->nextSkip() >= numStripeRows) {
         readerBase_->randomSkip()->consume(numStripeRows);
         ++skippedStripes_;
-        goto advanceToNextStripe;
+        advanceToNextStripe();
+        continue;
+      }
+      if (shouldSkipStripe(currentStripe_)) {
+        maybeUpdateRandomSkip(numStripeRows);
+        ++skippedStripes_;
+        advanceToNextStripe();
+        continue;
       }
       loadCurrentStripe();
     }
@@ -284,7 +293,6 @@ int64_t SelectiveNimbleRowReader::nextRowNumber() {
       nextRowNumber_ = stripeRowOffsets_[currentStripe_] + rowInCurrentStripe_;
       return *nextRowNumber_;
     }
-  advanceToNextStripe:
     advanceToNextStripe();
   }
   // Update random skip tracker for trailing rows that were skipped due to upper
@@ -370,6 +378,60 @@ void SelectiveNimbleRowReader::computeStatsBasedRowSize() const {
   }
 }
 
+bool SelectiveNimbleRowReader::shouldSkipStripe(uint32_t stripe) const {
+  const auto& stripeStats = readerBase_->stripeColumnStats();
+  if (stripe >= stripeStats.size() || stripeStats[stripe].empty()) {
+    return false;
+  }
+  const auto& rootType = *readerBase_->fileSchemaWithId();
+  const auto& rowType = rootType.type()->asRow();
+  for (auto* childSpec : options_.scanSpec()->stableChildren()) {
+    if (!childSpec->hasFilter() || childSpec->filter() == nullptr ||
+        childSpec->isConstant() || !childSpec->readFromFile()) {
+      continue;
+    }
+    const auto columnIndex =
+        rowType.getChildIdxIfExists(childSpec->fieldName());
+    if (!columnIndex.has_value()) {
+      continue;
+    }
+    const auto& childType = rootType.childAt(columnIndex.value());
+    if (childType == nullptr) {
+      continue;
+    }
+    // Only top-level integral columns are pruned here. Use an explicit
+    // comparison rather than a switch over TypeKind to avoid -Wswitch-enum
+    // requiring every enumerator to be listed.
+    const auto kind = childType->type()->kind();
+    if (kind != TypeKind::TINYINT && kind != TypeKind::SMALLINT &&
+        kind != TypeKind::INTEGER && kind != TypeKind::BIGINT) {
+      continue;
+    }
+    // Per-stripe stats are laid out by schema type id: the writer snapshots
+    // statsCollectors_ (indexed by TypeWithId::id()) in order, so
+    // stripeStats[stripe][id] matches childType->id() just like the file-level
+    // column stats. The bound check below guards against a stats section that
+    // covers fewer columns than the current schema.
+    const auto columnId = childType->id();
+    if (columnId >= stripeStats[stripe].size() ||
+        !stripeStats[stripe][columnId]) {
+      continue;
+    }
+    const auto* integralStats =
+        stripeStats[stripe][columnId]->as<IntegralStatistics>();
+    if (integralStats == nullptr || !integralStats->getMin().has_value() ||
+        !integralStats->getMax().has_value()) {
+      continue;
+    }
+    const auto hasNull = integralStats->getNullCount() > 0;
+    if (!childSpec->filter()->testInt64Range(
+            *integralStats->getMin(), *integralStats->getMax(), hasNull)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 std::optional<size_t> SelectiveNimbleRowReader::estimatedRowSize() const {
   if (const_cast<SelectiveNimbleRowReader*>(this)->nextRowNumber() == kAtEnd) {
     return std::nullopt;