Split header and data size in encoding size estimation

sdruzkin · facebook-github-bot · commit 74840cc8568b · 2025-10-12T19:53:37.000-07:00
Summary: Split header and data size in encoding size estimation so that the read factor can only be applied to the compressible data. Current implementation give funky estimates for small sized data.

Differential Revision: D84458948
diff --git a/dwio/nimble/encodings/EncodingSelectionPolicy.h b/dwio/nimble/encodings/EncodingSelectionPolicy.h
@@ -186,9 +186,10 @@ class ManualEncodingSelectionPolicy : public EncodingSelectionPolicy<T> {
       // We use read factor weights to raise/lower the favorability of each
       // encoding.
       auto readFactor = pair.second;
-      auto cost = size.value() * readFactor;
+      auto cost = size.value().cost(readFactor);
       NIMBLE_SELECTION_LOG(
-          YELLOW << "Encoding: " << encodingType << ", Size: " << size.value()
+          YELLOW << "Encoding: " << encodingType
+                 << ", Size: " << size.value().size()
                  << ", Factor: " << readFactor << ", Cost: " << cost);
       if (cost < minCost) {
         minCost = cost;
diff --git a/dwio/nimble/encodings/EncodingSizeEstimation.h b/dwio/nimble/encodings/EncodingSizeEstimation.h
@@ -22,27 +22,40 @@
 #include "dwio/nimble/common/Exceptions.h"
 #include "dwio/nimble/common/FixedBitArray.h"
 #include "dwio/nimble/common/Types.h"
+#include "dwio/nimble/encodings/Statistics.h"
 
 namespace facebook::nimble {
 namespace detail {
 
+struct SizeEstimation {
+  uint64_t headerSize;
+  uint64_t dataSize;
+
+  uint64_t cost(double readFactor) const {
+    return headerSize + dataSize * readFactor;
+  }
+
+  uint64_t size() const {
+    return headerSize + dataSize;
+  }
+};
+
 // This class is meant to quickly estimate the size of encoded data using a
 // given encoding type. It does a lot of assumptions, and it is not meant to be
 // 100% accurate.
 template <typename T, bool FixedByteWidth>
 struct EncodingSizeEstimation {
   using physicalType = typename TypeTraits<T>::physicalType;
 
-  static std::optional<uint64_t> estimateNumericSize(
+  static std::optional<SizeEstimation> estimateNumericSize(
       const EncodingType encodingType,
       const uint64_t entryCount,
       const Statistics<physicalType>& statistics) {
     switch (encodingType) {
       case EncodingType::Constant: {
         return statistics.uniqueCounts().size() == 1
-            ? std::optional<uint64_t>{getEncodingOverhead<
-                  EncodingType::Constant,
-                  physicalType>()}
+            ? std::optional<
+                  SizeEstimation>{{getEncodingOverhead<EncodingType::Constant, physicalType>(), 0}}
             : std::nullopt;
       }
       case EncodingType::MainlyConstant: {
@@ -71,25 +84,26 @@ struct EncodingSizeEstimation {
         // stored bit packed.
         const auto uncommonIndicesSize =
             bitPackedBytes(0, entryCount, uncommonCount);
-        uint32_t overhead =
+        const uint32_t overhead =
             getEncodingOverhead<EncodingType::MainlyConstant, physicalType>() +
             // Overhead for storing uncommon values
             getEncodingOverhead<EncodingType::FixedBitWidth, physicalType>() +
             // Overhead for storing uncommon bitmap
             getEncodingOverhead<EncodingType::SparseBool, bool>() +
             getEncodingOverhead<EncodingType::FixedBitWidth, uint32_t>();
-        return overhead + sizeof(physicalType) + uncommonValueSize +
-            uncommonIndicesSize;
+        return {
+            {overhead + sizeof(physicalType),
+             uncommonValueSize + uncommonIndicesSize}};
       }
       case EncodingType::Trivial: {
-        return getEncodingOverhead<EncodingType::Trivial, physicalType>() +
-            (entryCount * sizeof(physicalType));
+        return {
+            {getEncodingOverhead<EncodingType::Trivial, physicalType>(),
+             entryCount * sizeof(physicalType)}};
       }
       case EncodingType::FixedBitWidth: {
-        return getEncodingOverhead<
-                   EncodingType::FixedBitWidth,
-                   physicalType>() +
-            bitPackedBytes(statistics.min(), statistics.max(), entryCount);
+        return {
+            {getEncodingOverhead<EncodingType::FixedBitWidth, physicalType>(),
+             bitPackedBytes(statistics.min(), statistics.max(), entryCount)}};
       }
       case EncodingType::Dictionary: {
         // Assumptions:
@@ -100,13 +114,13 @@ struct EncodingSizeEstimation {
             bitPackedBytes(0, statistics.uniqueCounts().size(), entryCount);
         const uint64_t alphabetSize =
             statistics.uniqueCounts().size() * sizeof(physicalType);
-        uint32_t overhead =
+        const uint32_t overhead =
             getEncodingOverhead<EncodingType::Dictionary, physicalType>() +
             // Alphabet overhead
             getEncodingOverhead<EncodingType::Trivial, physicalType>() +
             // Indices overhead
             getEncodingOverhead<EncodingType::FixedBitWidth, uint32_t>();
-        return overhead + alphabetSize + indicesSize;
+        return {{overhead, alphabetSize + indicesSize}};
       }
       case EncodingType::RLE: {
         // Assumptions:
@@ -122,13 +136,13 @@ struct EncodingSizeEstimation {
             statistics.minRepeat(),
             statistics.maxRepeat(),
             statistics.consecutiveRepeatCount());
-        uint32_t overhead =
+        const uint32_t overhead =
             getEncodingOverhead<EncodingType::RLE, physicalType>() +
             // Overhead of run values
             getEncodingOverhead<EncodingType::FixedBitWidth, physicalType>() +
             // Overhead of run lengths
             getEncodingOverhead<EncodingType::FixedBitWidth, uint32_t>();
-        return overhead + runValuesSize + runLengthsSize;
+        return {{overhead, runValuesSize + runLengthsSize}};
       }
       case EncodingType::Varint: {
         // Note: the condition below actually support floating point numbers as
@@ -145,8 +159,9 @@ struct EncodingSizeEstimation {
               [&i](const uint64_t sum, const uint64_t bucketSize) {
                 return sum + (bucketSize * (++i));
               });
-          return getEncodingOverhead<EncodingType::Varint, physicalType>() +
-              dataSize;
+          return {
+              {getEncodingOverhead<EncodingType::Varint, physicalType>(),
+               dataSize}};
         } else {
           return std::nullopt;
         }
@@ -157,16 +172,15 @@ struct EncodingSizeEstimation {
     }
   }
 
-  static std::optional<uint64_t> estimateBoolSize(
+  static std::optional<SizeEstimation> estimateBoolSize(
       const EncodingType encodingType,
       const size_t entryCount,
       const Statistics<physicalType>& statistics) {
     switch (encodingType) {
       case EncodingType::Constant: {
         return statistics.uniqueCounts().size() == 1
-            ? std::optional<uint64_t>{getEncodingOverhead<
-                  EncodingType::Constant,
-                  physicalType>()}
+            ? std::optional<
+                  SizeEstimation>{{getEncodingOverhead<EncodingType::Constant, physicalType>(), 0}}
             : std::nullopt;
       }
       case EncodingType::SparseBool: {
@@ -177,16 +191,18 @@ struct EncodingSizeEstimation {
         const auto exceptionCount = std::min(
             statistics.uniqueCounts().at(true),
             statistics.uniqueCounts().at(false));
-        uint32_t overhead =
+        const uint32_t overhead =
             getEncodingOverhead<EncodingType::SparseBool, physicalType>() +
             // Overhead for storing exception indices
             getEncodingOverhead<EncodingType::FixedBitWidth, uint32_t>();
-        return overhead + sizeof(bool) +
-            bitPackedBytes(0, entryCount, exceptionCount);
+        return {
+            {overhead + sizeof(bool),
+             bitPackedBytes(0, entryCount, exceptionCount)}};
       }
       case EncodingType::Trivial: {
-        return getEncodingOverhead<EncodingType::Trivial, physicalType>() +
-            FixedBitArray::bufferSize(entryCount, 1);
+        return {
+            {getEncodingOverhead<EncodingType::Trivial, physicalType>(),
+             FixedBitArray::bufferSize(entryCount, 1)}};
       }
       case EncodingType::RLE: {
         // Assumptions:
@@ -197,29 +213,28 @@ struct EncodingSizeEstimation {
             statistics.minRepeat(),
             statistics.maxRepeat(),
             statistics.consecutiveRepeatCount());
-        uint32_t overhead =
+        const uint32_t overhead =
             getEncodingOverhead<EncodingType::RLE, physicalType>() +
             // Overhead of run lengths
             getEncodingOverhead<EncodingType::FixedBitWidth, uint32_t>();
-        return overhead + sizeof(bool) + runLengthsSize;
+        return {{overhead + sizeof(bool), runLengthsSize}};
       }
       default: {
         return std::nullopt;
       }
     }
   }
 
-  static std::optional<uint64_t> estimateStringSize(
+  static std::optional<SizeEstimation> estimateStringSize(
       const EncodingType encodingType,
       const size_t entryCount,
       const Statistics<std::string_view>& statistics) {
     const uint32_t maxStringSize = statistics.max().size();
     switch (encodingType) {
       case EncodingType::Constant: {
         return statistics.uniqueCounts().size() == 1
-            ? std::optional<uint64_t>{getEncodingOverhead<
-                  EncodingType::Constant,
-                  physicalType>(maxStringSize)}
+            ? std::optional<
+                  SizeEstimation>{{getEncodingOverhead<EncodingType::Constant, physicalType>(maxStringSize), 0}}
             : std::nullopt;
       }
       case EncodingType::MainlyConstant: {
@@ -264,7 +279,7 @@ struct EncodingSizeEstimation {
         // stored bit packed.
         const auto uncommonIndicesSize =
             bitPackedBytes(0, entryCount, uncommonCount);
-        uint32_t overhead =
+        const uint32_t overhead =
             getEncodingOverhead<EncodingType::MainlyConstant, physicalType>(
                 maxUniqueCount->first.size()) +
             // Overhead for storing uncommon values
@@ -273,17 +288,18 @@ struct EncodingSizeEstimation {
             // Overhead for storing uncommon bitmap
             getEncodingOverhead<EncodingType::SparseBool, bool>();
 
-        return overhead + alphabetSize + uncommonIndicesSize;
+        return {{overhead, alphabetSize + uncommonIndicesSize}};
       }
       case EncodingType::Trivial: {
         // We assume string lengths will be stored bit packed.
-        return getEncodingOverhead<EncodingType::Trivial, physicalType>(
-                   maxStringSize) +
-            statistics.totalStringsLength() +
-            bitPackedBytes(
-                   statistics.min().size(),
-                   statistics.max().size(),
-                   entryCount);
+        return {
+            {getEncodingOverhead<EncodingType::Trivial, physicalType>(
+                 maxStringSize),
+             statistics.totalStringsLength() +
+                 bitPackedBytes(
+                     statistics.min().size(),
+                     statistics.max().size(),
+                     entryCount)}};
       }
       case EncodingType::Dictionary: {
         // Assumptions:
@@ -305,23 +321,23 @@ struct EncodingSizeEstimation {
             bitPackedBytes(statistics.min().size(),
                            statistics.max().size(),
                            statistics.uniqueCounts().size());
-        uint32_t overhead =
+        const uint32_t overhead =
             getEncodingOverhead<EncodingType::Dictionary, physicalType>(
                 maxStringSize) +
             // Alphabet overhead
             getEncodingOverhead<EncodingType::Trivial, physicalType>(
                 maxStringSize) +
             // Indices overhead
             getEncodingOverhead<EncodingType::FixedBitWidth, uint32_t>();
-        return overhead + alphabetSize + indicesSize;
+        return {{overhead, alphabetSize + indicesSize}};
       }
       case EncodingType::RLE: {
         // Assumptions:
         // Run values are stored using dictionary (and inside, trivial +
         // bit-packing). Run lengths are stored using bit-packing (with bit
         // width needed to store max repetition count).
 
-        uint64_t runValuesSize =
+        const uint64_t runValuesSize =
             // (unique) strings blob size
             std::accumulate(
                 statistics.uniqueCounts().cbegin(),
@@ -344,23 +360,23 @@ struct EncodingSizeEstimation {
             statistics.minRepeat(),
             statistics.maxRepeat(),
             statistics.consecutiveRepeatCount());
-        uint32_t overhead =
+        const uint32_t overhead =
             getEncodingOverhead<EncodingType::RLE, physicalType>() +
             // Overhead of run values
             getEncodingOverhead<EncodingType::Dictionary, physicalType>() +
             getEncodingOverhead<EncodingType::Trivial, physicalType>() +
             getEncodingOverhead<EncodingType::FixedBitWidth, uint32_t>() +
             // Overhead of run lengths
             getEncodingOverhead<EncodingType::FixedBitWidth, uint32_t>();
-        return overhead + runValuesSize + runLengthsSize;
+        return {{overhead, runValuesSize + runLengthsSize}};
       }
       default: {
         return std::nullopt;
       }
     }
   }
 
-  static std::optional<uint64_t> estimateSize(
+  static std::optional<SizeEstimation> estimateSize(
       const EncodingType encodingType,
       const size_t entryCount,
       const Statistics<physicalType>& statistics) {
diff --git a/dwio/nimble/encodings/tests/EncodingSelectionTests.cpp b/dwio/nimble/encodings/tests/EncodingSelectionTests.cpp
@@ -134,7 +134,7 @@ void verifySizeEstimate(
           encodingTypeForEstimation,
           values.size(),
           nimble::Statistics<T>::create(values));
-  EXPECT_EQ(estimatedSize, expectedEstimatedSize);
+  EXPECT_EQ(estimatedSize.value().size(), expectedEstimatedSize);
 }
 
 template <typename T>
@@ -180,7 +180,7 @@ void test(std::span<const T> values, std::vector<EncodingDetails> expected) {
     LOG(INFO) << "Expected: " << expected[i].encodingType << "<"
               << expected[i].dataType << ">[" << expected[i].nestedEncodingName
               << ":" << expected[i].level << "]";
-    LOG(INFO) << "Actual: " << actual[i].encodingType << "<"
+    LOG(INFO) << "Actual:   " << actual[i].encodingType << "<"
               << actual[i].dataType << ">[" << actual[i].nestedEncodingName
               << ":" << actual[i].level << "]";
     EXPECT_EQ(expected[i].encodingType, actual[i].encodingType);
@@ -435,7 +435,7 @@ TYPED_TEST(EncodingSelectionNumericTests, SelectRunLength) {
 
   if constexpr (
       nimble::isFloatingPointType<T>() || std::is_same_v<int32_t, T> ||
-      sizeof(T) > 4) {
+      sizeof(T) >= 4) {
     // Floating point types and big types prefer storing the run values as
     // dictionary
     test<T>(
@@ -994,3 +994,18 @@ TEST(EncodingSelectionTests, TestNullable) {
 
   LOG(INFO) << "Final size: " << serialized.size();
 }
+
+TEST(EncodingSelectionTests, TestSizeEstimateCost) {
+  std::vector<uint8_t> values{1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
+  auto estimatedSize =
+      nimble::detail::EncodingSizeEstimation<uint8_t, false>::estimateSize(
+          nimble::EncodingType::Trivial,
+          values.size(),
+          nimble::Statistics<uint8_t>::create(values))
+          .value();
+  EXPECT_EQ(estimatedSize.headerSize, 7);
+  EXPECT_EQ(estimatedSize.dataSize, 10);
+  EXPECT_EQ(estimatedSize.size(), 17);
+  EXPECT_EQ(estimatedSize.cost(10), 107);
+  EXPECT_EQ(estimatedSize.cost(0), 7);
+}
