feat(nimble): Add per-chunk null count (#992)

Summary:
Pull Request resolved: #992

The first real per-chunk statistic. Chunk stats are still purely positional;
this records how many nulls each chunk holds, so a later reader can skip chunks
on null-related predicates and we can eventually retire the read-time
byte-peeking row estimates in ChunkedDecoder.

- The writer counts nulls per chunk during encode, only when chunk stats are
  enabled. The fleet default is off, so no writer pays for this yet.
- Stored as a new field appended to the existing chunk-stats flatbuffer. Files
  written before this simply lack it, and the reader reports "unknown" rather
  than a wrong value.
- Reader accessor StreamIndex::chunkNullCount returns the count, or nullopt for
  pre-stats files.

Rationale: https://docs.google.com/document/d/1cTCXje3RkPdhwFGK-FTY5muy4Cky2qnfnP-e_VFkjQA/edit?tab=t.0#heading=h.rpljwjiiglku

Differential Revision: D111913795

Author: prashantgolash

dwio/nimble/index/ChunkStatsGroup.cpp

@@ -144,6 +144,23 @@ ChunkLocation StreamIndex::lookupChunk(uint32_t rowId) const {
       chunkIndex, streamOffset, nextOffset - streamOffset, rowOffset};
 }
 
+std::optional<uint32_t> StreamIndex::chunkNullCount(uint32_t chunkIndex) const {
+  const auto* root = asFlatBuffersRoot<serialization::StripeChunkStats>(
+      chunkStats_->metadata_->content());
+
+  const auto* nullCounts = root->stream_chunk_null_counts();
+  if (nullCounts == nullptr) {
+    // Absent in files written before per-chunk null statistics were added; the
+    // null count is unknown.
+    return std::nullopt;
+  }
+  // The array is present, so chunkIndex (derived from ChunkLocation) must be in
+  // range. An out-of-range index is a programmer error or malformed metadata,
+  // not a legacy file, so fail loudly rather than masking it as "unknown".
+  NIMBLE_CHECK_LT(chunkIndex, nullCounts->size());
+  return nullCounts->Get(chunkIndex);
+}
+
 uint32_t StreamIndex::rowCount() const {
   if (endChunkOffset_ == startChunkOffset_) {
     return 0;

dwio/nimble/index/ChunkStatsGroup.h

@@ -17,6 +17,7 @@
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 
 #include "dwio/nimble/common/Exceptions.h"
 #include "dwio/nimble/index/IndexTypes.h"
@@ -101,6 +102,11 @@ class StreamIndex {
   /// Lookup chunk by row ID within the stream's row range.
   ChunkLocation lookupChunk(uint32_t rowId) const;
 
+  /// Returns the per-chunk null-value count for the chunk at the given absolute
+  /// position (ChunkLocation::chunkIndex), or std::nullopt when per-chunk null
+  /// statistics are absent (files written before chunk statistics were added).
+  std::optional<uint32_t> chunkNullCount(uint32_t chunkIndex) const;
+
   /// Returns the total number of rows in this stream.
   uint32_t rowCount() const;
 

dwio/nimble/index/tests/ClusterIndexTestUtils.cpp

@@ -103,7 +103,9 @@ std::vector<Chunk> createChunks(
     auto pos = buffer.reserve(spec.size);
     std::memset(pos, 'X', spec.size);
     chunks.push_back(
-        {.rowCount = spec.rowCount, .content = {{pos, spec.size}}});
+        {.rowCount = spec.rowCount,
+         .nullCount = spec.nullCount,
+         .content = {{pos, spec.size}}});
   }
   return chunks;
 }
@@ -167,6 +169,7 @@ StreamStats ChunkStatsTestHelper::streamStats(uint32_t streamId) const {
   NIMBLE_CHECK_NOT_NULL(chunkRows);
   const auto* chunkOffsets = root->stream_chunk_offsets();
   NIMBLE_CHECK_NOT_NULL(chunkOffsets);
+  const auto* chunkNullCounts = root->stream_chunk_null_counts();
 
   const uint32_t stripeCount = chunkStats_->stripeCount_;
 
@@ -186,6 +189,9 @@ StreamStats ChunkStatsTestHelper::streamStats(uint32_t streamId) const {
       const uint32_t chunkIndex = startChunkIndex + i;
       stats.chunkRows.push_back(chunkRows->Get(chunkIndex));
       stats.chunkOffsets.push_back(chunkOffsets->Get(chunkIndex));
+      if (chunkNullCounts != nullptr) {
+        stats.chunkNullCounts.push_back(chunkNullCounts->Get(chunkIndex));
+      }
     }
   }
 

dwio/nimble/index/tests/ClusterIndexTestUtils.h

@@ -134,12 +134,16 @@ struct StreamStats {
   std::vector<uint32_t> chunkRows;
   /// Byte offset of each chunk within its stream.
   std::vector<uint32_t> chunkOffsets;
+  /// Per-chunk null-value count for this stream. Empty when per-chunk null
+  /// statistics are absent from the section.
+  std::vector<uint32_t> chunkNullCounts;
 };
 
 /// Specification for a single chunk in a stream (for test data creation).
 struct ChunkSpec {
-  uint32_t rowCount;
-  uint32_t size;
+  uint32_t rowCount{};
+  uint32_t size{};
+  uint32_t nullCount{0};
 };
 
 /// Specification for a single stream (for test data creation).

dwio/nimble/tablet/Chunk.h

@@ -27,6 +27,10 @@ namespace facebook::nimble {
 struct Chunk {
   uint32_t rowCount{0};
 
+  /// Number of null values in this chunk. Only populated when chunk statistics
+  /// are enabled (VeloxWriterOptions::enableChunkIndex); left at 0 otherwise.
+  uint32_t nullCount{0};
+
   /// The encoded and compressed data content of this chunk, stored as a vector
   /// of string views. Each string_view points to a buffer containing a portion
   /// of the chunk's data. Multiple buffers may be used for large chunks.

dwio/nimble/tablet/ChunkStats.fbs

@@ -39,6 +39,10 @@ table StripeChunkStats {
   stream_chunk_rows:[uint32];
   /// Byte offset of each chunk within its stream.
   stream_chunk_offsets:[uint32];
+  /// Per-chunk null-value count, parallel to stream_chunk_rows/offsets.
+  /// Appended field: absent in files written before per-chunk null statistics
+  /// were added, in which case readers treat the null count as unknown.
+  stream_chunk_null_counts:[uint32];
 }
 
 /// Root table for the chunk stats optional section ("columnar.chunk.stats").

dwio/nimble/tablet/ChunkStatsWriter.cpp

@@ -66,6 +66,7 @@ void ChunkStatsWriter::addStream(
     accumulatedRows += chunk.rowCount;
     index.chunkRows.emplace_back(accumulatedRows);
     index.chunkOffsets.emplace_back(accumulatedOffset);
+    index.chunkNullCounts.emplace_back(chunk.nullCount);
     accumulatedOffset += chunk.contentSize();
     ++index.chunkCount;
   }
@@ -123,8 +124,10 @@ void ChunkStatsWriter::writeGroup(
 
   std::vector<uint32_t> flattenedChunkRows;
   std::vector<uint32_t> flattenedChunkOffsets;
+  std::vector<uint32_t> flattenedChunkNullCounts;
   flattenedChunkRows.reserve(accumulatedChunkCount);
   flattenedChunkOffsets.reserve(accumulatedChunkCount);
+  flattenedChunkNullCounts.reserve(accumulatedChunkCount);
 
   for (const auto& stripe : groupIndex_->stripes) {
     for (size_t streamId = 0; streamId < streamCount; ++streamId) {
@@ -138,6 +141,10 @@ void ChunkStatsWriter::writeGroup(
             flattenedChunkOffsets.end(),
             stream.chunkOffsets.begin(),
             stream.chunkOffsets.end());
+        flattenedChunkNullCounts.insert(
+            flattenedChunkNullCounts.end(),
+            stream.chunkNullCounts.begin(),
+            stream.chunkNullCounts.end());
       }
     }
   }
@@ -148,7 +155,8 @@ void ChunkStatsWriter::writeGroup(
       static_cast<uint32_t>(streamCount),
       builder.CreateVector(flattenedStreamChunkCounts),
       builder.CreateVector(flattenedChunkRows),
-      builder.CreateVector(flattenedChunkOffsets));
+      builder.CreateVector(flattenedChunkOffsets),
+      builder.CreateVector(flattenedChunkNullCounts));
   builder.Finish(chunkIndex);
 
   chunkIndexSections_.push_back(createMetadataSection(asView(builder)));

dwio/nimble/tablet/ChunkStatsWriter.h

@@ -82,6 +82,8 @@ class ChunkStatsWriter {
     std::vector<uint32_t> chunkRows;
     // Byte offsets of each chunk within the stream.
     std::vector<uint32_t> chunkOffsets;
+    // Per-chunk null-value count (statistic used for chunk skipping).
+    std::vector<uint32_t> chunkNullCounts;
     // Number of chunks in this stripe for this stream.
     uint32_t chunkCount{0};
   };

dwio/nimble/tablet/tests/ChunkStatsWriterTest.cpp

@@ -165,6 +165,45 @@ TEST_F(ChunkStatsWriterTest, singleStripe) {
   EXPECT_EQ(r11.rowOffset, 60);
 }
 
+TEST_F(ChunkStatsWriterTest, perChunkNullCounts) {
+  ChunkStatsWriter writer(*pool_);
+  Buffer buffer{*pool_};
+  TestChunkFileIndex fileIndex;
+
+  // 1 stripe, 2 streams. ChunkSpec = {rowCount, size, nullCount}.
+  // Stream 0: 3 chunks with null counts 3, 0, 5.
+  // Stream 1: 2 chunks with null counts 0, 40 (a fully-null chunk).
+  writer.newStripe(2);
+  auto chunks0 = createChunks(buffer, {{30, 10, 3}, {45, 15, 0}, {25, 8, 5}});
+  auto chunks1 = createChunks(buffer, {{60, 20, 0}, {40, 12, 40}});
+  writer.addStream(0, chunks0);
+  writer.addStream(1, chunks1);
+
+  writer.writeGroup(2, 1, createMetadataSectionCallback(fileIndex));
+  writer.writeRoot(writeRootCallback(fileIndex));
+
+  auto chunkIndex = loadChunkIndex(fileIndex.groupMetadataSections[0], 0, 1);
+
+  // Null counts round-trip through the flatbuffer in flattened order.
+  ChunkStatsTestHelper helper(chunkIndex.get());
+  EXPECT_EQ(
+      helper.streamStats(0).chunkNullCounts, (std::vector<uint32_t>{3, 0, 5}));
+  EXPECT_EQ(
+      helper.streamStats(1).chunkNullCounts, (std::vector<uint32_t>{0, 40}));
+
+  // Public reader accessor: lookupChunk() yields the absolute chunk index,
+  // which chunkNullCount() maps to the per-chunk null statistic.
+  auto stream0 = chunkIndex->createStreamIndex(0, 0, 33);
+  ASSERT_NE(stream0, nullptr);
+  EXPECT_EQ(stream0->chunkNullCount(stream0->lookupChunk(0).chunkIndex), 3);
+  EXPECT_EQ(stream0->chunkNullCount(stream0->lookupChunk(30).chunkIndex), 0);
+  EXPECT_EQ(stream0->chunkNullCount(stream0->lookupChunk(75).chunkIndex), 5);
+
+  auto stream1 = chunkIndex->createStreamIndex(0, 1, 32);
+  ASSERT_NE(stream1, nullptr);
+  EXPECT_EQ(stream1->chunkNullCount(stream1->lookupChunk(60).chunkIndex), 40);
+}
+
 TEST_F(ChunkStatsWriterTest, multipleStripesInSingleGroup) {
   ChunkStatsWriter writer(*pool_);
   Buffer buffer{*pool_};

dwio/nimble/velox/VeloxWriter.cpp

@@ -15,6 +15,7 @@
  */
 #include "dwio/nimble/velox/VeloxWriter.h"
 
+#include <algorithm>
 #include <memory>
 #include <optional>
 #include <unordered_map>
@@ -1453,6 +1454,14 @@ uint32_t VeloxWriter::encodeChunk(
   }
   uint32_t chunkBytes{0};
   chunk.rowCount = chunkView.rowCount();
+  // Collect the per-chunk null-count statistic only when chunk statistics are
+  // enabled, so writers with chunk stats off pay no extra cost. Streams with no
+  // nulls keep the default 0 without scanning.
+  if (context_->options().enableChunkIndex && chunkView.hasNulls()) {
+    const auto nonNulls = chunkView.nonNulls();
+    chunk.nullCount = static_cast<uint32_t>(
+        std::count(nonNulls.begin(), nonNulls.end(), false));
+  }
   ChunkedStreamWriter chunkWriter{
       *encodingBuffer_, context_->options().chunkCompression};
   for (auto& buffer : chunkWriter.encode(encoded)) {