[VL]:shuffle: extend TypeAwareCompress to INT128 (DECIMAL HUGEINT)

hhr293 · guowangy · zhulipeng · hhr293 · commit ebb225eb705e · 2026-06-23T09:27:36.000+08:00
The TAC codec from apache#11894 only covered INT64. DECIMAL(p>18) is stored as 128-bit HugeInt and previously fell through to LZ4, missing a structural compression opportunity: in OLAP workloads (TPC-H prices, taxes, ...) DECIMAL values usually fit in INT64, so the high 64 bits are 0 and the low 64 bits are narrow -- exactly the pattern FFOR exploits. Implementation: split each 16B value into lo/hi uint64 sub-streams via stride-2 gather into stack-allocated scratch buffers, then run the existing 64-bit FFOR encoder on each. Wire format reuses the 64-bit per-block (bw, count, base) header twice -- the stream is self- describing, no hi/lo length prefix needed. hi sub-streams that are all equal to base degenerate to just the 16B header (bw=0). Velox HUGEINT is mapped to tac::kUInt128; shuffle writer and frame format unchanged. Results vs LZ4 on the same int128 input: compression ratio 0.116 (TAC) vs 0.193 (LZ4) -- 40% smaller End-to-end on TPC-H SF=6000, the wins concentrate on the queries with heavy decimal-keyed shuffles: q15: shuffle size -15%, latency -8% q17: shuffle size -8%, latency -3% q18: shuffle size -8%, latency -3% Generated-by: Claude claude-opus-4-7 Co-authored-by: Guo Wangyang <wangyang.guo@intel.com> Co-authored-by: Lipeng Zhu <lipeng.zhu@intel.com> Signed-off-by: huhengrui <hengrui.hu@intel.com>
diff --git a/cpp/core/tests/FForCodecTest.cc b/cpp/core/tests/FForCodecTest.cc
@@ -643,3 +643,171 @@ TEST(TypeAwareCompressCodecTest, UnsupportedType) {
   auto result = TypeAwareCompressCodec::compress(dummy, 8, dummy, 100, kSomeUnsupportedType);
   ASSERT_FALSE(result.ok());
 }
+
+TEST(TypeAwareCompressCodecTest, Int128Supported) {
+  ASSERT_TRUE(TypeAwareCompressCodec::support(tac::kUInt128));
+}
+
+TEST(TypeAwareCompressCodecTest, Int128NarrowRoundtrip) {
+  // 1024 INT128 values where the high 64 bits are 0 (typical DECIMAL fitting
+  // in int64) and the low 64 bits range narrowly.  Both hi and lo sub-streams
+  // should compress extremely well.
+  constexpr int64_t kNumValues = 1024;
+  std::vector<uint8_t> data(kNumValues * sizeof(__int128_t), 0);
+  for (int64_t i = 0; i < kNumValues; ++i) {
+    uint64_t lo = 1000000ULL + static_cast<uint64_t>(i);
+    std::memcpy(data.data() + i * sizeof(__int128_t), &lo, sizeof(uint64_t));
+  }
+
+  const int64_t inputSize = data.size();
+  auto maxLen = TypeAwareCompressCodec::maxCompressedLen(inputSize, tac::kUInt128);
+  ASSERT_GT(maxLen, 0);
+
+  std::vector<uint8_t> compressed(maxLen);
+  auto compResult =
+      TypeAwareCompressCodec::compress(data.data(), inputSize, compressed.data(), compressed.size(), tac::kUInt128);
+  ASSERT_TRUE(compResult.ok()) << compResult.status().ToString();
+  const int64_t compressedSize = *compResult;
+
+  ASSERT_LT(compressedSize, inputSize / 2);
+
+  std::vector<uint8_t> decoded(inputSize, 0xff);
+  auto decResult = TypeAwareCompressCodec::decompress(compressed.data(), compressedSize, decoded.data(), inputSize);
+  ASSERT_TRUE(decResult.ok()) << decResult.status().ToString();
+  ASSERT_EQ(*decResult, compressedSize);
+  ASSERT_EQ(0, std::memcmp(decoded.data(), data.data(), inputSize));
+}
+
+TEST(TypeAwareCompressCodecTest, Int128HighEntropyRoundtrip) {
+  // High-entropy 128-bit values: both halves are random.  Verifies round-trip
+  // correctness on a workload FFor cannot compress effectively.
+  constexpr int64_t kNumValues = 1024;
+  std::vector<uint8_t> data(kNumValues * sizeof(__int128_t));
+  std::mt19937_64 rng(42);
+  for (int64_t i = 0; i < kNumValues; ++i) {
+    uint64_t lo = rng();
+    uint64_t hi = rng();
+    std::memcpy(data.data() + i * sizeof(__int128_t), &lo, sizeof(uint64_t));
+    std::memcpy(data.data() + i * sizeof(__int128_t) + 8, &hi, sizeof(uint64_t));
+  }
+
+  const int64_t inputSize = data.size();
+  auto maxLen = TypeAwareCompressCodec::maxCompressedLen(inputSize, tac::kUInt128);
+  std::vector<uint8_t> compressed(maxLen);
+  auto compResult =
+      TypeAwareCompressCodec::compress(data.data(), inputSize, compressed.data(), compressed.size(), tac::kUInt128);
+  ASSERT_TRUE(compResult.ok()) << compResult.status().ToString();
+
+  std::vector<uint8_t> decoded(inputSize);
+  auto decResult = TypeAwareCompressCodec::decompress(compressed.data(), *compResult, decoded.data(), inputSize);
+  ASSERT_TRUE(decResult.ok()) << decResult.status().ToString();
+  ASSERT_EQ(*decResult, *compResult);
+  ASSERT_EQ(0, std::memcmp(decoded.data(), data.data(), inputSize));
+}
+
+namespace {
+
+// Helper: build kNumValues 128-bit values whose lo halves form a narrow
+// arithmetic sequence and hi halves are zero (typical "DECIMAL fits in int64"
+// pattern).  Returns the raw byte buffer.
+std::vector<uint8_t> makeNarrowInt128(int64_t kNumValues) {
+  std::vector<uint8_t> data(kNumValues * sizeof(__int128_t), 0);
+  for (int64_t i = 0; i < kNumValues; ++i) {
+    uint64_t lo = 1000000ULL + static_cast<uint64_t>(i);
+    std::memcpy(data.data() + i * sizeof(__int128_t), &lo, sizeof(uint64_t));
+  }
+  return data;
+}
+
+} // namespace
+
+TEST(TypeAwareCompressCodecTest, Int128TailRoundtrip) {
+  // Element count that is NOT a multiple of kLanes (=4): 1027 = 256*4 + 3.
+  // Forces compress128 to emit 1 full block plus a 3-element tail block,
+  // exercising the kBwTailMarker path in both compress and decompress.
+  constexpr int64_t kNumValues = 1027;
+  auto data = makeNarrowInt128(kNumValues);
+
+  const int64_t inputSize = data.size();
+  auto maxLen = TypeAwareCompressCodec::maxCompressedLen(inputSize, tac::kUInt128);
+  std::vector<uint8_t> compressed(maxLen);
+  auto compResult =
+      TypeAwareCompressCodec::compress(data.data(), inputSize, compressed.data(), compressed.size(), tac::kUInt128);
+  ASSERT_TRUE(compResult.ok()) << compResult.status().ToString();
+  const int64_t compressedSize = *compResult;
+
+  std::vector<uint8_t> decoded(inputSize, 0xff);
+  auto decResult = TypeAwareCompressCodec::decompress(compressed.data(), compressedSize, decoded.data(), inputSize);
+  ASSERT_TRUE(decResult.ok()) << decResult.status().ToString();
+  ASSERT_EQ(0, std::memcmp(decoded.data(), data.data(), inputSize));
+}
+
+TEST(TypeAwareCompressCodecTest, Int128MisalignedRoundtrip) {
+  // Drive compress128/decompress128 with input AND output pointers offset by
+  // 1 byte so neither is uint64-aligned.  Exercises the <false, false>
+  // template instantiation, which is otherwise dead code under the natural
+  // alignment of std::vector<uint8_t>::data().
+  constexpr int64_t kNumValues = 1024;
+  auto src = makeNarrowInt128(kNumValues);
+  const int64_t inputSize = src.size();
+
+  // Stage src into a buffer whose payload starts at a 1-byte-offset address.
+  std::vector<uint8_t> inBuf(inputSize + 1, 0);
+  std::memcpy(inBuf.data() + 1, src.data(), inputSize);
+
+  auto maxLen = TypeAwareCompressCodec::maxCompressedLen(inputSize, tac::kUInt128);
+  std::vector<uint8_t> compBuf(maxLen + 1, 0);
+  // Skip the first byte to misalign the compressed-data start as well.
+  auto compResult =
+      TypeAwareCompressCodec::compress(inBuf.data() + 1, inputSize, compBuf.data() + 1, maxLen, tac::kUInt128);
+  ASSERT_TRUE(compResult.ok()) << compResult.status().ToString();
+  const int64_t compressedSize = *compResult;
+
+  std::vector<uint8_t> outBuf(inputSize + 1, 0xff);
+  auto decResult = TypeAwareCompressCodec::decompress(compBuf.data() + 1, compressedSize, outBuf.data() + 1, inputSize);
+  ASSERT_TRUE(decResult.ok()) << decResult.status().ToString();
+  ASSERT_EQ(0, std::memcmp(outBuf.data() + 1, src.data(), inputSize));
+}
+
+TEST(TypeAwareCompressCodecTest, Int128TruncatedInputRejected) {
+  // Compress a normal stream, then hand decompress() a prefix that drops the
+  // terminator tail header plus the final block's hi sub-header, leaving the
+  // truncation point inside the final block's lo payload.  decompress128Impl
+  // must detect the short read via decodeBlock's bounds check and return a
+  // value count below the expected number of 128-bit values, which
+  // TypeAwareCompressCodec turns into an Invalid status.
+  constexpr int64_t kNumValues = 1024;
+  auto data = makeNarrowInt128(kNumValues);
+  const int64_t inputSize = data.size();
+  auto maxLen = TypeAwareCompressCodec::maxCompressedLen(inputSize, tac::kUInt128);
+  std::vector<uint8_t> compressed(maxLen);
+  auto compResult =
+      TypeAwareCompressCodec::compress(data.data(), inputSize, compressed.data(), compressed.size(), tac::kUInt128);
+  ASSERT_TRUE(compResult.ok()) << compResult.status().ToString();
+  const int64_t compressedSize = *compResult;
+  // 64 bytes = 16B tail header + 16B hi sub-header + 32B into the final lo
+  // payload.  Smaller drops can land on a clean tail boundary and silently
+  // succeed.
+  ASSERT_GT(compressedSize, 64);
+
+  std::vector<uint8_t> decoded(inputSize, 0);
+  auto decResult =
+      TypeAwareCompressCodec::decompress(compressed.data(), compressedSize - 64, decoded.data(), inputSize);
+  ASSERT_FALSE(decResult.ok());
+}
+
+TEST(TypeAwareCompressCodecTest, Int128OutputBufferTooSmallRejected) {
+  // compress128 must reject an output buffer smaller than maxCompressedLen,
+  // since the codec writes worst-case-sized headers up-front and only knows
+  // the final length after the analyze step.
+  constexpr int64_t kNumValues = 64;
+  auto data = makeNarrowInt128(kNumValues);
+  const int64_t inputSize = data.size();
+  auto maxLen = TypeAwareCompressCodec::maxCompressedLen(inputSize, tac::kUInt128);
+  ASSERT_GT(maxLen, 0);
+
+  std::vector<uint8_t> tooSmall(maxLen - 1);
+  auto result =
+      TypeAwareCompressCodec::compress(data.data(), inputSize, tooSmall.data(), tooSmall.size(), tac::kUInt128);
+  ASSERT_FALSE(result.ok());
+}
diff --git a/cpp/core/utils/tac/FForCodec.cc b/cpp/core/utils/tac/FForCodec.cc
@@ -37,7 +37,8 @@ FForCodec::compress(const uint8_t* input, int64_t inputSize, uint8_t* output, in
   size_t numValues = inputSize / sizeof(uint64_t);
   auto maxLen = static_cast<int64_t>(ffor::compress64Bound(numValues));
   if (outputSize < maxLen) {
-    return arrow::Status::Invalid("FForCodec: output buffer too small.");
+    return arrow::Status::Invalid(
+        "FForCodec: output buffer too small (need ", maxLen, " bytes, got ", outputSize, ").");
   }
 
   auto written = ffor::compress64(reinterpret_cast<const uint64_t*>(input), numValues, output);
@@ -57,4 +58,47 @@ FForCodec::decompress(const uint8_t* input, int64_t inputSize, uint8_t* output,
   return static_cast<int64_t>(nDecoded);
 }
 
+int64_t FForCodec::maxCompressedLength128(int64_t inputSize) {
+  if (inputSize % sizeof(__int128_t) != 0) {
+    return 0;
+  }
+  size_t numValues = inputSize / sizeof(__int128_t);
+  return static_cast<int64_t>(ffor::compress128Bound(numValues));
+}
+
+arrow::Result<int64_t>
+FForCodec::compress128(const uint8_t* input, int64_t inputSize, uint8_t* output, int64_t outputSize) {
+  if (inputSize == 0) {
+    return 0;
+  }
+  if (inputSize % sizeof(__int128_t) != 0) {
+    return arrow::Status::Invalid(
+        "FForCodec: input size ", inputSize, " is not a multiple of ", sizeof(__int128_t), ".");
+  }
+
+  size_t numValues = inputSize / sizeof(__int128_t);
+  auto maxLen = static_cast<int64_t>(ffor::compress128Bound(numValues));
+  if (outputSize < maxLen) {
+    return arrow::Status::Invalid(
+        "FForCodec: output buffer too small for 128-bit compression (need ", maxLen, " bytes, got ", outputSize, ").");
+  }
+
+  auto written = ffor::compress128(input, numValues, output);
+  return static_cast<int64_t>(written);
+}
+
+arrow::Result<int64_t>
+FForCodec::decompress128(const uint8_t* input, int64_t inputSize, uint8_t* output, int64_t outputSize) {
+  if (outputSize == 0) {
+    return 0;
+  }
+  if (outputSize % sizeof(__int128_t) != 0) {
+    return arrow::Status::Invalid(
+        "FForCodec: output size ", outputSize, " is not a multiple of ", sizeof(__int128_t), ".");
+  }
+
+  auto nDecoded = ffor::decompress128(input, inputSize, output, static_cast<size_t>(outputSize));
+  return static_cast<int64_t>(nDecoded);
+}
+
 } // namespace gluten
diff --git a/cpp/core/utils/tac/FForCodec.h b/cpp/core/utils/tac/FForCodec.h
@@ -23,11 +23,12 @@
 
 namespace gluten {
 
-// FFOR (Frame-of-Reference) codec for uint64_t data using 4-lane layout.
-// Used for INT64/UINT64 columns in shuffle.
+// FFOR (Frame-of-Reference) codec for uint64_t / 128-bit data using 4-lane layout.
+// Used for INT64/UINT64 and INT128/HUGEINT (DECIMAL) columns in shuffle.
 class FForCodec {
  public:
   // Returns the maximum compressed size in bytes for the given input size.
+  // Input is treated as a stream of uint64 values; size must be a multiple of 8.
   static int64_t maxCompressedLength(int64_t inputSize);
 
   // Compress uint64_t data.
@@ -40,6 +41,24 @@ class FForCodec {
   // Returns the number of uint64_t values decoded.
   static arrow::Result<int64_t>
   decompress(const uint8_t* input, int64_t inputSize, uint8_t* output, int64_t outputSize);
+
+  // Worst-case compressed size for a stream of 128-bit values.
+  // inputSize must be a multiple of sizeof(__int128_t); returns 0 otherwise.
+  static int64_t maxCompressedLength128(int64_t inputSize);
+
+  // Compress a stream of 128-bit values whose in-memory layout consists of
+  // two 64-bit halves: the low 8B at offset 0 and the high 8B at offset 8
+  // (DECIMAL128 in Velox). Internally splits each value into hi/lo uint64
+  // sub-streams per block and runs the 64-bit FFOR encoder on each.
+  // inputSize must be a multiple of sizeof(__int128_t); returns the number
+  // of compressed bytes written to output.
+  static arrow::Result<int64_t>
+  compress128(const uint8_t* input, int64_t inputSize, uint8_t* output, int64_t outputSize);
+
+  // Decompress data produced by compress128().  outputSize must be a multiple
+  // of sizeof(__int128_t); returns the number of 128-bit values decoded.
+  static arrow::Result<int64_t>
+  decompress128(const uint8_t* input, int64_t inputSize, uint8_t* output, int64_t outputSize);
 };
 
 } // namespace gluten
diff --git a/cpp/core/utils/tac/TypeAwareCompressCodec.cc b/cpp/core/utils/tac/TypeAwareCompressCodec.cc
@@ -21,14 +21,18 @@
 namespace gluten {
 
 bool TypeAwareCompressCodec::support(int8_t tacType) {
-  return tacType == tac::kUInt64;
+  return tacType == tac::kUInt64 || tacType == tac::kUInt128;
 }
 
 int64_t TypeAwareCompressCodec::maxCompressedLen(int64_t inputLen, int8_t tacType) {
-  if (!support(tacType)) {
-    return 0;
+  switch (tacType) {
+    case tac::kUInt64:
+      return kPayloadHeaderSize + FForCodec::maxCompressedLength(inputLen);
+    case tac::kUInt128:
+      return kPayloadHeaderSize + FForCodec::maxCompressedLength128(inputLen);
+    default:
+      return 0;
   }
-  return kPayloadHeaderSize + FForCodec::maxCompressedLength(inputLen);
 }
 
 arrow::Result<int64_t> TypeAwareCompressCodec::compress(
@@ -50,10 +54,21 @@ arrow::Result<int64_t> TypeAwareCompressCodec::compress(
   auto* out = output;
   *out++ = static_cast<uint8_t>(CodecId::kFFor);
   *out++ = static_cast<uint8_t>(tacType);
+  int64_t availableOutput = outputLen - kPayloadHeaderSize;
 
-  auto availableOutput = outputLen - kPayloadHeaderSize;
-  ARROW_ASSIGN_OR_RAISE(auto compressedLen, FForCodec::compress(input, inputLen, out, availableOutput));
-
+  int64_t compressedLen = 0;
+  switch (tacType) {
+    case tac::kUInt64: {
+      ARROW_ASSIGN_OR_RAISE(compressedLen, FForCodec::compress(input, inputLen, out, availableOutput));
+      break;
+    }
+    case tac::kUInt128: {
+      ARROW_ASSIGN_OR_RAISE(compressedLen, FForCodec::compress128(input, inputLen, out, availableOutput));
+      break;
+    }
+    default:
+      return arrow::Status::Invalid("Unsupported tac type in compress: ", static_cast<int>(tacType));
+  }
   return kPayloadHeaderSize + compressedLen;
 }
 
@@ -65,18 +80,41 @@ TypeAwareCompressCodec::decompress(const uint8_t* input, int64_t inputLen, uint8
 
   auto* in = input;
   auto codecId = static_cast<CodecId>(*in++);
-  [[maybe_unused]] auto tacType = *in++;
+  auto tacType = static_cast<int8_t>(*in++);
   auto dataLen = inputLen - kPayloadHeaderSize;
 
   switch (codecId) {
-    case CodecId::kFFor: {
-      ARROW_ASSIGN_OR_RAISE(auto nDecoded, FForCodec::decompress(in, dataLen, output, outputLen));
-      (void)nDecoded;
-      return inputLen;
-    }
+    case CodecId::kFFor:
+      break;
     default:
       return arrow::Status::Invalid("Unknown type-aware codec ID: ", static_cast<int>(codecId));
   }
+
+  int64_t nDecoded = 0;
+  int64_t valueSize = 0;
+  const char* typeName = nullptr;
+  switch (tacType) {
+    case tac::kUInt64: {
+      ARROW_ASSIGN_OR_RAISE(nDecoded, FForCodec::decompress(in, dataLen, output, outputLen));
+      valueSize = sizeof(uint64_t);
+      typeName = "uint64";
+      break;
+    }
+    case tac::kUInt128: {
+      ARROW_ASSIGN_OR_RAISE(nDecoded, FForCodec::decompress128(in, dataLen, output, outputLen));
+      valueSize = 2 * sizeof(uint64_t);
+      typeName = "uint128";
+      break;
+    }
+    default:
+      return arrow::Status::Invalid("Unknown tac type in decompress: ", static_cast<int>(tacType));
+  }
+  const int64_t expected = outputLen / valueSize;
+  if (nDecoded != expected) {
+    return arrow::Status::Invalid(
+        "TAC decompress ", typeName, " value count mismatch: expected ", expected, " got ", nDecoded);
+  }
+  return inputLen;
 }
 
 } // namespace gluten
diff --git a/cpp/core/utils/tac/TypeAwareCompressCodec.h b/cpp/core/utils/tac/TypeAwareCompressCodec.h
@@ -30,6 +30,7 @@ namespace tac {
 enum TacDataType : int8_t {
   kUnsupported = -1, // Not compressible by TAC.
   kUInt64 = 0, // 8-byte unsigned integer (also used for int64, double, date64).
+  kUInt128 = 1, // 16-byte unsigned integer (used for HUGEINT / DECIMAL128).
 };
 
 } // namespace tac
@@ -38,7 +39,8 @@ enum TacDataType : int8_t {
 /// compression algorithm based on the data type of the buffer.
 ///
 /// Currently supported:
-///   kUInt64 -> FFor (Frame-of-Reference + Bit-Packing) for uint64_t streams.
+///   kUInt64  -> FFOR (Frame-of-Reference + Bit-Packing) for uint64_t streams.
+///   kUInt128 -> FFOR applied to lo/hi uint64 sub-streams of 128-bit values.
 ///
 /// The compressed wire format is self-describing: decompress() does not need
 /// a type hint because codec ID and element width are embedded in the header.
diff --git a/cpp/core/utils/tac/ffor.hpp b/cpp/core/utils/tac/ffor.hpp
diff --git a/cpp/velox/shuffle/VeloxTypeAwareCompress.h b/cpp/velox/shuffle/VeloxTypeAwareCompress.h
