feat: Add kCompactRaw raw encoding support and simplify deserialization API (facebookincubator#633)

Summary:

CONTEXT: The nimble serializer/deserializer needed several improvements:
(1) Stream sizes in the kCompactRaw trailer only supported Trivial and Varint
encoding. Stream sizes tend to be similar across streams, making delta encoding
a natural fit. However, in practice delta-varint provides marginal savings over
plain varint for stream sizes since the values vary enough that deltas don't
compress much better.
(2) DeserializerOptions required callers to pass a full SerializationVersion,
but the deserializer auto-detects version from the header byte — only a bool
flag is needed.
(3) StreamData took two separate bool flags (encodingEnabled, useVarintRowCount)
instead of deriving them from the version.
(4) Several helper functions were missing for version classification.

WHAT:
- Add Zigzag.h utility (branchless zigzag encode/decode for int32/int64) and
  Delta encoding support for kCompactRaw stream sizes. Delta encoding stores
  the first value as-is, then zigzag+varint encodes consecutive deltas.
  Wire: [0x08][count:varint][first:varint][zigzag(delta_1):varint]...[trailer_size:u32]

- Simplify DeserializerOptions: replace `std::optional<SerializationVersion> version`
  with `bool hasHeader`. When true, auto-detect version from the first byte of
  serialized data. When false, assume kLegacy (no version header).

- Simplify StreamData: take SerializationVersion instead of two bool flags
  (encodingEnabled, useVarintRowCount), derive both internally.

- Add version helper functions in Options.h:
  - `nonLegacyFormat()`: returns true for kCompact/kCompactRaw/kTabletRaw
  - `isRawFormat()`: returns true for kCompactRaw/kTabletRaw
  - `isTabletRawFormat()`: returns true for kTabletRaw only

- Add validation in Serializer constructor: reject explicit kLegacy and kTabletRaw
  versions, validate streamSizesEncodingType requires non-legacy version.

- Extract `StreamDataReader::readVersion()` method for version header parsing.

- Add fuzz tests for mixed-version serialization and projection (cycles through
  kCompact, kCompactRaw, kCompactRaw+Delta) with decode verification.

- SST writer benchmark: remove sst_partitioned_index flag, change
  nimble_stream_sizes_encoding default from Trivial to Delta.

Reviewed By: srsuryadev, jiahaol-work, tanjialiang

Differential Revision: D98941222

Author: xiaoxmeng

dwio/nimble/common/Zigzag.h

@@ -0,0 +1,52 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cstdint>
+
+// ZigZag encoding maps signed integers to unsigned integers so that values
+// with small absolute value have small encoded values, making them efficient
+// for varint encoding.
+//
+//   zigzagEncode32(0)  == 0
+//   zigzagEncode32(-1) == 1
+//   zigzagEncode32(1)  == 2
+//   zigzagEncode32(-2) == 3
+//   ...
+//
+// In general:
+//   if x >= 0, zigzagEncode32(x) == 2*x
+//   if x <  0, zigzagEncode32(x) == -2*x - 1
+
+namespace facebook::nimble::zigzag {
+
+inline constexpr uint32_t zigzagEncode32(int32_t val) noexcept {
+  return static_cast<uint32_t>((val << 1) ^ (val >> 31));
+}
+
+inline constexpr int32_t zigzagDecode32(uint32_t val) noexcept {
+  return static_cast<int32_t>((val >> 1) ^ -(val & 1));
+}
+
+inline constexpr uint64_t zigzagEncode64(int64_t val) noexcept {
+  return static_cast<uint64_t>((val << 1) ^ (val >> 63));
+}
+
+inline constexpr int64_t zigzagDecode64(uint64_t val) noexcept {
+  return static_cast<int64_t>((val >> 1) ^ -(val & 1));
+}
+
+} // namespace facebook::nimble::zigzag

dwio/nimble/common/tests/ZigzagTests.cpp

@@ -0,0 +1,87 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <gtest/gtest.h>
+
+#include <limits>
+
+#include "dwio/nimble/common/Zigzag.h"
+
+using namespace ::facebook::nimble::zigzag;
+
+TEST(ZigzagTest, encode32) {
+  EXPECT_EQ(zigzagEncode32(0), 0);
+  EXPECT_EQ(zigzagEncode32(-1), 1);
+  EXPECT_EQ(zigzagEncode32(1), 2);
+  EXPECT_EQ(zigzagEncode32(-2), 3);
+  EXPECT_EQ(zigzagEncode32(2), 4);
+  EXPECT_EQ(zigzagEncode32(-3), 5);
+  EXPECT_EQ(zigzagEncode32(std::numeric_limits<int32_t>::max()), 4294967294U);
+  EXPECT_EQ(zigzagEncode32(std::numeric_limits<int32_t>::min()), 4294967295U);
+}
+
+TEST(ZigzagTest, decode32) {
+  EXPECT_EQ(zigzagDecode32(0), 0);
+  EXPECT_EQ(zigzagDecode32(1), -1);
+  EXPECT_EQ(zigzagDecode32(2), 1);
+  EXPECT_EQ(zigzagDecode32(3), -2);
+  EXPECT_EQ(zigzagDecode32(4), 2);
+  EXPECT_EQ(zigzagDecode32(5), -3);
+  EXPECT_EQ(zigzagDecode32(4294967294U), std::numeric_limits<int32_t>::max());
+  EXPECT_EQ(zigzagDecode32(4294967295U), std::numeric_limits<int32_t>::min());
+}
+
+TEST(ZigzagTest, roundTrip32) {
+  for (int32_t val :
+       {0,
+        1,
+        -1,
+        100,
+        -100,
+        1'000'000,
+        -1'000'000,
+        std::numeric_limits<int32_t>::max(),
+        std::numeric_limits<int32_t>::min()}) {
+    EXPECT_EQ(zigzagDecode32(zigzagEncode32(val)), val);
+  }
+}
+
+TEST(ZigzagTest, encode64) {
+  EXPECT_EQ(zigzagEncode64(0), 0);
+  EXPECT_EQ(zigzagEncode64(-1), 1);
+  EXPECT_EQ(zigzagEncode64(1), 2);
+  EXPECT_EQ(zigzagEncode64(-2), 3);
+  EXPECT_EQ(
+      zigzagEncode64(std::numeric_limits<int64_t>::max()),
+      std::numeric_limits<uint64_t>::max() - 1);
+  EXPECT_EQ(
+      zigzagEncode64(std::numeric_limits<int64_t>::min()),
+      std::numeric_limits<uint64_t>::max());
+}
+
+TEST(ZigzagTest, roundTrip64) {
+  for (int64_t val :
+       {int64_t{0},
+        int64_t{1},
+        int64_t{-1},
+        int64_t{100},
+        int64_t{-100},
+        int64_t{1'000'000'000'000},
+        int64_t{-1'000'000'000'000},
+        std::numeric_limits<int64_t>::max(),
+        std::numeric_limits<int64_t>::min()}) {
+    EXPECT_EQ(zigzagDecode64(zigzagEncode64(val)), val);
+  }
+}

dwio/nimble/serializer/Deserializer.cpp

@@ -91,14 +91,10 @@ class DeserializerImpl : public Decoder {
   DeserializerImpl(
       const Type* type,
       bool inMapStream,
-      bool enableEncoding,
-      bool useVarintRowCount,
       velox::memory::MemoryPool* pool)
       : type_{type},
         pool_{pool},
         inMapStream_{inMapStream},
-        encodingEnabled_{enableEncoding},
-        useVarintRowCount_{useVarintRowCount},
         scalarKind_{getScalarKindForType(*type)},
         typeStorageWidth_{getTypeStorageWidth(*type)} {}
 
@@ -156,19 +152,18 @@ class DeserializerImpl : public Decoder {
   }
 
   // Add data starting at the given row offset.
-  void addBatch(uint32_t rowOffset, std::string_view data) {
+  // version: the auto-detected serialization version, used to determine
+  // encoding enabled and varint row count settings.
+  void addBatch(
+      uint32_t rowOffset,
+      std::string_view data,
+      SerializationVersion version) {
     if (data.empty()) {
       return;
     }
     batchSegments_.emplace_back(
         BatchSegment{
-            rowOffset,
-            serde::StreamData(
-                scalarKind_,
-                encodingEnabled_,
-                useVarintRowCount_,
-                data,
-                pool_)});
+            rowOffset, serde::StreamData(scalarKind_, version, data, pool_)});
   }
 
   // Record a segment where this key is present in every row (in-map stream
@@ -499,10 +494,6 @@ class DeserializerImpl : public Decoder {
   // True for inMap streams (fills with 'false' when missing), false for nulls
   // streams (fills with 'true' when missing).
   const bool inMapStream_;
-  // True when nimble encoding is used for scalar streams.
-  const bool encodingEnabled_;
-  // Whether encoding headers use varint (true) or fixed u32 (false) row counts.
-  const bool useVarintRowCount_;
   // Cached from type at construction to avoid per-call dispatch.
   const ScalarKind scalarKind_;
   const uint32_t typeStorageWidth_;
@@ -639,15 +630,10 @@ Deserializer::Deserializer(
 void Deserializer::createDeserializersForType(
     const Type& type,
     uint32_t depth) {
-  const bool enableEncoding = options_.enableEncoding();
-  const bool useVarintRowCount =
-      !isTabletRawFormat(options_.serializationVersion());
   deserializerMap_[getMainDescriptor(type).offset()] =
       std::make_unique<DeserializerImpl>(
           &type,
           /*inMapStream=*/false,
-          enableEncoding,
-          useVarintRowCount,
           pool_);
   // FlatMap is only supported at depth 1 (top-level columns). FlatMap keys can
   // vary across batches, causing gaps in nulls/inMap streams. Gap detection is
@@ -661,8 +647,6 @@ void Deserializer::createDeserializersForType(
       deserializerMap_[inMapOffset] = std::make_unique<DeserializerImpl>(
           &type,
           /*inMapStream=*/true,
-          enableEncoding,
-          useVarintRowCount,
           pool_);
       inMapChildTypes_[inMapOffset] = flatMap.childAt(i).get();
     }
@@ -691,6 +675,7 @@ void Deserializer::deserialize(
   serde::StreamDataReader reader{pool_, options_};
   for (auto sv : data) {
     const auto batchRows = reader.initialize(sv);
+    const auto version = reader.version();
     // Reset present tracking from previous batch.
     if (hasInMapChildren && !inMapPresentOffsetsList_.empty()) {
       for (auto off : inMapPresentOffsetsList_) {
@@ -707,7 +692,7 @@ void Deserializer::deserialize(
         auto* decoder = deserializers_[offset];
         if (decoder != nullptr) {
           DeserializerImpl::toDecoderImpl(decoder)->addBatch(
-              rowOffset, streamData);
+              rowOffset, streamData, version);
         }
       }
     });

dwio/nimble/serializer/DeserializerImpl.cpp

@@ -30,14 +30,13 @@ namespace facebook::nimble::serde {
 
 StreamData::StreamData(
     ScalarKind kind,
-    bool encodingEnabled,
-    bool useVarintRowCount,
+    SerializationVersion version,
     std::string_view data,
     velox::memory::MemoryPool* pool)
     : kind_{kind},
       pool_{pool},
-      encodingEnabled_{encodingEnabled},
-      useVarintRowCount_{useVarintRowCount} {
+      encodingEnabled_{nonLegacyFormat(version)},
+      useVarintRowCount_{!isTabletRawFormat(version)} {
   NIMBLE_CHECK_NOT_NULL(pool_, "Memory pool required for encoding");
   init(data);
 }
@@ -63,11 +62,12 @@ uint32_t StreamData::decodeStrings(uint32_t count, std::string_view* output) {
   return index;
 }
 
-void StreamData::reset(std::string_view data, bool encodingEnabled) {
+void StreamData::reset(std::string_view data, SerializationVersion version) {
   readRows_ = 0;
   encoding_.reset();
   stringBuffers_.clear();
-  encodingEnabled_ = encodingEnabled;
+  encodingEnabled_ = nonLegacyFormat(version);
+  useVarintRowCount_ = !isTabletRawFormat(version);
   // Re-initialize with new data.
   init(data);
 }
@@ -207,39 +207,37 @@ StreamDataReader::StreamDataReader(
 uint32_t StreamDataReader::initialize(std::string_view data) {
   pos_ = data.data();
   end_ = data.end();
-  if (options_.hasVersionHeader()) {
-    const auto version = static_cast<SerializationVersion>(*pos_);
-    NIMBLE_CHECK(
-        version == SerializationVersion::kLegacy ||
-            version == SerializationVersion::kCompact ||
-            version == SerializationVersion::kCompactRaw ||
-            version == SerializationVersion::kTabletRaw,
-        "Unsupported version {}",
-        static_cast<uint8_t>(version));
-    // Verify the version read from serialized data matches options.
-    NIMBLE_CHECK_EQ(
-        version,
-        *options_.version,
-        "Version mismatch: data has version {}, options expect {}",
-        version,
-        *options_.version);
-    ++pos_;
-  }
+  readVersion();
   // All non-legacy formats use varint row count.
-  if (usesVarintRowCount(options_.serializationVersion())) {
+  if (usesVarintRowCount(version_)) {
     return varint::readVarint32(&pos_);
   }
   return encoding::readUint32(pos_);
 }
 
+void StreamDataReader::readVersion() {
+  if (!options_.hasHeader) {
+    version_ = SerializationVersion::kLegacy;
+    return;
+  }
+  version_ = static_cast<SerializationVersion>(*pos_);
+  NIMBLE_CHECK(
+      version_ == SerializationVersion::kLegacy ||
+          version_ == SerializationVersion::kCompact ||
+          version_ == SerializationVersion::kCompactRaw ||
+          version_ == SerializationVersion::kTabletRaw,
+      "Unsupported version {}",
+      static_cast<uint8_t>(version_));
+  ++pos_;
+}
+
 void StreamDataReader::iterateStreams(
     const std::function<void(uint32_t offset, std::string_view data)>&
         callback) {
-  if (options_.enableEncoding()) {
+  if (nonLegacyFormat(version_)) {
     // kCompact/kCompactRaw/kTabletRaw format: read stream sizes from trailer.
-    const auto streamSizes =
-        detail::readStreamSizes(end_, options_.serializationVersion(), pool_);
-    const bool tabletRaw = isTabletRawFormat(options_.serializationVersion());
+    const auto streamSizes = detail::readStreamSizes(end_, version_, pool_);
+    const bool tabletRaw = isTabletRawFormat(version_);
 
     for (uint32_t i = 0; i < streamSizes.size(); ++i) {
       std::string_view streamData(pos_, streamSizes[i]);