This is an automated email from the ASF dual-hosted git repository.
zhouyuan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 9a2a8be7e8 [VL] extend TypeAwareCompress to INT128 in shuffle (#12299)
9a2a8be7e8 is described below
commit 9a2a8be7e8a7661d94fdaab79747709163ef108e
Author: hu hengrui <[email protected]>
AuthorDate: Tue Jun 23 22:15:16 2026 +0800
[VL] extend TypeAwareCompress to INT128 in shuffle (#12299)
* [VL]:shuffle: extend TypeAwareCompress to INT128 (DECIMAL HUGEINT)
The TAC codec from #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 <[email protected]>
Co-authored-by: Lipeng Zhu <[email protected]>
Signed-off-by: huhengrui <[email protected]>
* format: fix clang-format violations in PR #12299
clang-format-15 -Werror flagged four spots introduced by 27c53a85
(extend TypeAwareCompress to INT128). No semantic change.
- ffor.hpp:354 writeHeader call: argument-pack break style.
- ffor.hpp:394 trailing comment missing space after //.
- ffor.hpp:473 std::memcpy fits on one line under 120-col limit.
- ffor.hpp:667 dst64 assignment fits on one line.
To be squashed into 27c53a85 before merge.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
* fix: bound 64-bit FFOR decompress by output buffer capacity
decompress64(input, inputSize, output) has no way to know the caller's
output buffer capacity, so a corrupt or truncated stream header can
drive the decoder to write past the end of `output` before the caller
can detect a value-count mismatch.
Two specific OOB paths exposed today:
- Tail block: `count` is read straight from the wire (`inPtr[1]`).
Previously gated only by `inPtr + tailBytes <= inEnd`, so a header
with count=0xFF copies 255*8 = 2040 bytes into output regardless
of how much room the caller actually has.
- Regular block: `blockVals = inPtr[1] * kLanes` is gated only by
`<= kMaxValuesPerBlock`. A maximum-sized block is happily decoded
even when the caller's buffer is much smaller, overflowing output.
Either case is a heap/stack OOB write triggered by malformed shuffle
bytes. The 128-bit path already guards against this via outValuesMax
(ffor.hpp L634/L642); this brings the 64-bit path to parity.
Changes:
- decompress64{,Impl}: add `outputSize`, derive
`outValuesMax = outputSize / sizeof(uint64_t)`, and reject any
tail `count` or block `blockVals` that would exceed
`outValuesMax - nDecoded`.
- FForCodec::decompress: forward the existing `outputSize` arg
(previously dropped at the call site) into ffor::decompress64.
- Tests: update the 5 existing decompress64 call sites to pass the
output buffer capacity.
No behavior change for well-formed streams. On malformed input, the
loop stops cleanly at the first oversized header instead of writing
past `output`; callers still detect the short result via the existing
nDecoded vs expected-value-count check in TypeAwareCompressCodec.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
Co-authored-by: Guo Wangyang <[email protected]>
Co-authored-by: Lipeng Zhu <[email protected]>
Signed-off-by: huhengrui <[email protected]>
---------
Signed-off-by: huhengrui <[email protected]>
Co-authored-by: Guo Wangyang <[email protected]>
Co-authored-by: Lipeng Zhu <[email protected]>
Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>
---
cpp/core/tests/FForCodecTest.cc | 178 ++++++++++++++-
cpp/core/utils/tac/FForCodec.cc | 49 +++-
cpp/core/utils/tac/FForCodec.h | 23 +-
cpp/core/utils/tac/TypeAwareCompressCodec.cc | 64 ++++--
cpp/core/utils/tac/TypeAwareCompressCodec.h | 4 +-
cpp/core/utils/tac/ffor.hpp | 323 +++++++++++++++++++++------
cpp/velox/shuffle/VeloxTypeAwareCompress.h | 2 +
7 files changed, 557 insertions(+), 86 deletions(-)
diff --git a/cpp/core/tests/FForCodecTest.cc b/cpp/core/tests/FForCodecTest.cc
index a723322211..d4ef8aa4b9 100644
--- a/cpp/core/tests/FForCodecTest.cc
+++ b/cpp/core/tests/FForCodecTest.cc
@@ -85,7 +85,7 @@ void compressRoundtrip(const uint64_t* data, size_t num) {
size_t written = compress64(data, num, buf.data());
std::vector<uint64_t> decoded(num);
- size_t nDecoded = decompress64(buf.data(), written, decoded.data());
+ size_t nDecoded = decompress64(buf.data(), written, decoded.data(),
decoded.size() * sizeof(uint64_t));
ASSERT_EQ(nDecoded, num);
for (size_t i = 0; i < num; ++i) {
@@ -402,7 +402,7 @@ TEST(FForTest, Compress64MisalignedOutput) {
size_t written = compress64(data.data(), data.size(), out);
std::vector<uint64_t> decoded(256);
- size_t n = decompress64(out, written, decoded.data());
+ size_t n = decompress64(out, written, decoded.data(), decoded.size() *
sizeof(uint64_t));
ASSERT_EQ(n, size_t(256));
for (size_t i = 0; i < 256; ++i) {
ASSERT_EQ(decoded[i], data[i]) << "offset=" << offset << " i=" << i;
@@ -422,7 +422,7 @@ TEST(FForTest, Compress64MisalignedInput) {
size_t written = compress64(misalignedInput, 256, comp.data());
std::vector<uint64_t> decoded(256);
- size_t n = decompress64(comp.data(), written, decoded.data());
+ size_t n = decompress64(comp.data(), written, decoded.data(),
decoded.size() * sizeof(uint64_t));
ASSERT_EQ(n, size_t(256));
for (size_t i = 0; i < 256; ++i) {
ASSERT_EQ(decoded[i], raw[i]) << "offset=" << offset << " i=" << i;
@@ -438,7 +438,7 @@ TEST(FForTest, Decompress64MisalignedOutput) {
std::vector<uint8_t> outBuf(256 * sizeof(uint64_t) + 16);
for (size_t offset = 0; offset < 8; ++offset) {
auto* misalignedOutput = reinterpret_cast<uint64_t*>(outBuf.data() +
offset);
- size_t n = decompress64(comp.data(), written, misalignedOutput);
+ size_t n = decompress64(comp.data(), written, misalignedOutput, 256 *
sizeof(uint64_t));
ASSERT_EQ(n, size_t(256));
for (size_t i = 0; i < 256; ++i) {
uint64_t val;
@@ -463,7 +463,7 @@ TEST(FForTest, Compress64AllMisaligned) {
size_t written = compress64(inPtr, 256, compBuf.data() + compOff);
auto* outPtr = reinterpret_cast<uint64_t*>(outBuf.data() + outOff);
- size_t n = decompress64(compBuf.data() + compOff, written, outPtr);
+ size_t n = decompress64(compBuf.data() + compOff, written, outPtr, 256
* sizeof(uint64_t));
ASSERT_EQ(n, size_t(256));
for (size_t i = 0; i < 256; ++i) {
uint64_t val;
@@ -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
index ec079662f9..f31eb1a36c 100644
--- 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);
@@ -53,7 +54,51 @@ FForCodec::decompress(const uint8_t* input, int64_t
inputSize, uint8_t* output,
return arrow::Status::Invalid("FForCodec: output size ", outputSize, " is
not a multiple of 8.");
}
- auto nDecoded = ffor::decompress64(input, inputSize,
reinterpret_cast<uint64_t*>(output));
+ auto nDecoded =
+ ffor::decompress64(input, inputSize,
reinterpret_cast<uint64_t*>(output), static_cast<size_t>(outputSize));
+ 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);
}
diff --git a/cpp/core/utils/tac/FForCodec.h b/cpp/core/utils/tac/FForCodec.h
index b91a13860a..a176e00f33 100644
--- 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
index 2362f999b5..507aa972f6 100644
--- 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
index 6955525d2e..635c36ee78 100644
--- 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
index 0d632efff5..d659d3bb35 100644
--- a/cpp/core/utils/tac/ffor.hpp
+++ b/cpp/core/utils/tac/ffor.hpp
@@ -50,6 +50,7 @@
#pragma once
+#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
@@ -59,6 +60,13 @@
namespace gluten {
namespace ffor {
+// Byte order (applies to the 128-bit codec below): this codec round-trips
+// data as native uint64 reads/writes against the lo (offset 0) and hi
+// (offset 8) halves of each 128-bit value (DECIMAL128's in-memory layout in
+// Velox). Producer and consumer share byte order by virtue of running in
+// the same Spark cluster -- LZ4 and any other shuffle codec carry the same
+// implicit assumption -- so no explicit endian guard is needed here.
+
static constexpr unsigned kLanes = 4;
// Compile-time mask for a given bit width.
@@ -335,15 +343,55 @@ inline constexpr size_t compress64Bound(size_t num) {
return (nBlocks + 1) * kHeaderSize + num * sizeof(uint64_t);
}
+// Encode one block: write header + bit-packed payload into `out`.
+// Returns the number of bytes written. `out` must be 8-byte aligned;
+// callers whose output buffer is unaligned stage through a local scratch.
+// Shared by the 64-bit and 128-bit codecs.
+inline size_t encodeBlock(const uint64_t* src, size_t blockVals, uint64_t*
out) {
+ uint64_t base;
+ unsigned bw;
+ analyze(src, blockVals, base, bw);
+ writeHeader(
+ reinterpret_cast<uint8_t*>(out), static_cast<uint8_t>(bw),
static_cast<uint8_t>(blockVals / kLanes), base);
+ const size_t compWords = compressedWords(blockVals, bw);
+ encodeRt(src, out + 2, base, blockVals, bw);
+ return (2 + compWords) * sizeof(uint64_t);
+}
+
+// Decode one block of `blockVals` values from `in` into `dst`.
+// Returns the number of bytes consumed, or 0 on failure (corrupt header or
+// payload that would read past inEnd). `in` must be 8-byte aligned;
+// callers whose input buffer is unaligned stage through a local scratch.
+// Shared by the 64-bit and 128-bit codecs.
+inline size_t decodeBlock(const uint64_t* in, size_t inBytes, size_t
blockVals, uint64_t* dst) {
+ if (inBytes < kHeaderSize) {
+ return 0;
+ }
+ uint8_t bw;
+ uint8_t count;
+ uint64_t base;
+ readHeader(reinterpret_cast<const uint8_t*>(in), bw, count, base);
+ if (bw > 64 || bw == kBwTailMarker || static_cast<size_t>(count) * kLanes !=
blockVals) {
+ return 0;
+ }
+ const size_t compWords = compressedWords(blockVals, bw);
+ const size_t totalBytes = (2 + compWords) * sizeof(uint64_t);
+ if (totalBytes > inBytes) {
+ return 0;
+ }
+ decodeRt(in + 2, dst, base, blockVals, bw);
+ return totalBytes;
+}
+
// Template-based compress/decompress with alignment dispatch.
// InAligned: true if input (const uint64_t*) is 8-byte aligned.
// OutAligned: true if output (uint8_t*) is 8-byte aligned.
-// When aligned, encode_rt/decode_rt work on pointers directly.
-// When not aligned, a per-block aligned temp buffer is used.
+// encodeBlock requires aligned uint64_t* output; when the caller's output
+// buffer is unaligned, we stage through tmpOut and memcpy back.
template <bool InAligned, bool OutAligned>
inline size_t compress64Impl(const uint64_t* input, size_t num, uint8_t*
output) {
alignas(64) uint64_t tmpIn[kMaxValuesPerBlock];
- alignas(64) uint64_t tmpOut[kMaxValuesPerBlock];
+ alignas(64) uint64_t tmpOut[kMaxValuesPerBlock + 2]; // header(2 words) +
payload
uint8_t* outPtr = output;
size_t remaining = num;
@@ -355,35 +403,21 @@ inline size_t compress64Impl(const uint64_t* input,
size_t num, uint8_t* output)
blockVals = kMaxValuesPerBlock;
}
- // Analyze — read input via memcpy if unaligned.
- const uint64_t* analyzeSrc;
+ const uint64_t* src;
if constexpr (InAligned) {
- analyzeSrc = inPtr;
+ src = inPtr;
} else {
std::memcpy(tmpIn, inPtr, blockVals * sizeof(uint64_t));
- analyzeSrc = tmpIn;
+ src = tmpIn;
}
- uint64_t base;
- unsigned bw;
- analyze(analyzeSrc, blockVals, base, bw);
-
- writeHeader(outPtr, static_cast<uint8_t>(bw),
static_cast<uint8_t>(blockVals / kLanes), base);
- outPtr += kHeaderSize;
-
- size_t compN = compressedWords(blockVals, bw);
- size_t compBytes = compN * sizeof(uint64_t);
-
- // Encode: pick aligned src/dst.
- const uint64_t* encIn = InAligned ? inPtr : tmpIn;
- uint64_t* encOut = OutAligned ? reinterpret_cast<uint64_t*>(outPtr) :
tmpOut;
-
- encodeRt(encIn, encOut, base, blockVals, bw);
-
- if constexpr (!OutAligned) {
- std::memcpy(outPtr, tmpOut, compBytes);
+ if constexpr (OutAligned) {
+ outPtr += encodeBlock(src, blockVals,
reinterpret_cast<uint64_t*>(outPtr));
+ } else {
+ const size_t produced = encodeBlock(src, blockVals, tmpOut);
+ std::memcpy(outPtr, tmpOut, produced);
+ outPtr += produced;
}
- outPtr += compBytes;
inPtr += blockVals;
remaining -= blockVals;
@@ -418,62 +452,53 @@ inline size_t compress64(const uint64_t* input, size_t
num, uint8_t* output) {
}
// Template-based decompress with alignment dispatch.
+// decodeBlock requires aligned uint64_t* input; when the caller's input
+// buffer is unaligned, we stage through tmpIn and memcpy in.
template <bool InAligned, bool OutAligned>
-inline size_t decompress64Impl(const uint8_t* input, size_t inputSize,
uint64_t* output) {
- alignas(64) uint64_t tmpIn[kMaxValuesPerBlock];
+inline size_t decompress64Impl(const uint8_t* input, size_t inputSize,
uint64_t* output, size_t outputSize) {
+ alignas(64) uint64_t tmpIn[kMaxValuesPerBlock + 2];
alignas(64) uint64_t tmpOut[kMaxValuesPerBlock];
const uint8_t* inPtr = input;
const uint8_t* inEnd = input + inputSize;
+ const size_t outValuesMax = outputSize / sizeof(uint64_t);
size_t nDecoded = 0;
while (inPtr + kHeaderSize <= inEnd) {
- uint8_t bw;
- uint8_t count;
- uint64_t base;
- readHeader(inPtr, bw, count, base);
- inPtr += kHeaderSize;
-
- if (bw == kBwTailMarker) {
- if (count > 0) {
- // memcpy handles any alignment, no special case needed.
- std::memcpy(reinterpret_cast<uint8_t*>(output) + nDecoded *
sizeof(uint64_t), inPtr, count * sizeof(uint64_t));
+ if (inPtr[0] == kBwTailMarker) {
+ const uint8_t count = inPtr[1];
+ inPtr += kHeaderSize;
+ const size_t tailBytes = static_cast<size_t>(count) * sizeof(uint64_t);
+ if (count > 0 && inPtr + tailBytes <= inEnd && count <= outValuesMax -
nDecoded) {
+ std::memcpy(reinterpret_cast<uint8_t*>(output) + nDecoded *
sizeof(uint64_t), inPtr, tailBytes);
nDecoded += count;
}
break;
}
-
- size_t blockVals = static_cast<size_t>(count) * kLanes;
- size_t compBytes = compressedWords(blockVals, bw) * sizeof(uint64_t);
-
- if (inPtr + compBytes > inEnd) {
+ const size_t blockVals = static_cast<size_t>(inPtr[1]) * kLanes;
+ if (blockVals == 0 || blockVals > kMaxValuesPerBlock || blockVals >
outValuesMax - nDecoded) {
break;
}
+ const size_t remaining = static_cast<size_t>(inEnd - inPtr);
+ uint64_t* decDst = OutAligned ? output + nDecoded : tmpOut;
- // Decode: pick aligned src/dst.
- const uint64_t* decIn;
+ size_t consumed;
if constexpr (InAligned) {
- decIn = reinterpret_cast<const uint64_t*>(inPtr);
+ consumed = decodeBlock(reinterpret_cast<const uint64_t*>(inPtr),
remaining, blockVals, decDst);
} else {
- std::memcpy(tmpIn, inPtr, compBytes);
- decIn = tmpIn;
+ const size_t n = std::min(remaining, sizeof(tmpIn));
+ std::memcpy(tmpIn, inPtr, n);
+ consumed = decodeBlock(tmpIn, n, blockVals, decDst);
}
-
- uint64_t* decOut;
- if constexpr (OutAligned) {
- decOut = output + nDecoded;
- } else {
- decOut = tmpOut;
+ if (consumed == 0) {
+ break;
}
-
- decodeRt(decIn, decOut, base, blockVals, bw);
+ inPtr += consumed;
if constexpr (!OutAligned) {
std::memcpy(
reinterpret_cast<uint8_t*>(output) + nDecoded * sizeof(uint64_t),
tmpOut, blockVals * sizeof(uint64_t));
}
-
- inPtr += compBytes;
nDecoded += blockVals;
}
@@ -481,19 +506,191 @@ inline size_t decompress64Impl(const uint8_t* input,
size_t inputSize, uint64_t*
}
// Runtime dispatch.
-inline size_t decompress64(const uint8_t* input, size_t inputSize, uint64_t*
output) {
+inline size_t decompress64(const uint8_t* input, size_t inputSize, uint64_t*
output, size_t outputSize) {
bool inOk = (reinterpret_cast<uintptr_t>(input) % alignof(uint64_t) == 0);
bool outOk = (reinterpret_cast<uintptr_t>(output) % alignof(uint64_t) == 0);
if (inOk && outOk) {
- return decompress64Impl<true, true>(input, inputSize, output);
+ return decompress64Impl<true, true>(input, inputSize, output, outputSize);
+ }
+ if (inOk && !outOk) {
+ return decompress64Impl<true, false>(input, inputSize, output, outputSize);
+ }
+ if (!inOk && outOk) {
+ return decompress64Impl<false, true>(input, inputSize, output, outputSize);
+ }
+ return decompress64Impl<false, false>(input, inputSize, output, outputSize);
+}
+
+//
=============================================================================
+// 128-bit codec.
+//
+// Each 128-bit value occupies a 16B slot (lo at offset 0, hi at offset 8 --
+// the DECIMAL128 / __int128_t layout used by Velox). Per block, the lo and
+// hi halves are gathered into two stack scratches and each is fed through
+// the 64-bit FFOR encoder. Reads/writes go through native uint64, so the
+// codec is byte-order agnostic as long as producer and consumer agree.
+//
+// Wire format per block: [hdr][lo payload][hdr][hi payload]
+// followed by one tail block (kBwTailMarker) carrying the remaining 16B
+// values raw.
+//
=============================================================================
+
+inline constexpr size_t compress128Bound(size_t numValues) {
+ // Two 64-bit streams (lo + hi), worst case each =
compress64Bound(numValues).
+ return 2 * compress64Bound(numValues);
+}
+
+// 128-bit compress. See encodeBlock for InAligned/OutAligned semantics.
+template <bool InAligned, bool OutAligned>
+inline size_t compress128Impl(const uint8_t* input, size_t numValues, uint8_t*
output) {
+ alignas(64) uint64_t loBuffer[kMaxValuesPerBlock];
+ alignas(64) uint64_t hiBuffer[kMaxValuesPerBlock];
+ alignas(64) uint64_t tmpIn[kMaxValuesPerBlock * 2];
+ alignas(64) uint64_t tmpOut[kMaxValuesPerBlock + 2]; // header(2 words) +
payload
+
+ uint8_t* outPtr = output;
+ size_t remaining = numValues;
+ const uint8_t* inPtr = input;
+
+ while (remaining >= kLanes) {
+ size_t blockVals = remaining - (remaining % kLanes);
+ if (blockVals > kMaxValuesPerBlock) {
+ blockVals = kMaxValuesPerBlock;
+ }
+
+ const uint64_t* in64;
+ if constexpr (InAligned) {
+ in64 = reinterpret_cast<const uint64_t*>(inPtr);
+ } else {
+ std::memcpy(tmpIn, inPtr, blockVals * sizeof(__int128_t));
+ in64 = tmpIn;
+ }
+ for (size_t j = 0; j < blockVals; ++j) {
+ loBuffer[j] = in64[j * 2];
+ hiBuffer[j] = in64[j * 2 + 1];
+ }
+
+ if constexpr (OutAligned) {
+ outPtr += encodeBlock(loBuffer, blockVals,
reinterpret_cast<uint64_t*>(outPtr));
+ outPtr += encodeBlock(hiBuffer, blockVals,
reinterpret_cast<uint64_t*>(outPtr));
+ } else {
+ size_t n = encodeBlock(loBuffer, blockVals, tmpOut);
+ std::memcpy(outPtr, tmpOut, n);
+ outPtr += n;
+ n = encodeBlock(hiBuffer, blockVals, tmpOut);
+ std::memcpy(outPtr, tmpOut, n);
+ outPtr += n;
+ }
+
+ inPtr += blockVals * sizeof(__int128_t);
+ remaining -= blockVals;
+ }
+
+ // Tail: one header + remaining values copied raw.
+ writeHeader(outPtr, kBwTailMarker, static_cast<uint8_t>(remaining), 0);
+ outPtr += kHeaderSize;
+ if (remaining > 0) {
+ std::memcpy(outPtr, inPtr, remaining * sizeof(__int128_t));
+ outPtr += remaining * sizeof(__int128_t);
+ }
+ return static_cast<size_t>(outPtr - output);
+}
+
+// Runtime dispatch — check alignment once, pick the right template.
+inline size_t compress128(const uint8_t* input, size_t numValues, uint8_t*
output) {
+ const bool inOk = (reinterpret_cast<uintptr_t>(input) % alignof(uint64_t) ==
0);
+ const bool outOk = (reinterpret_cast<uintptr_t>(output) % alignof(uint64_t)
== 0);
+ if (inOk && outOk) {
+ return compress128Impl<true, true>(input, numValues, output);
+ }
+ if (inOk && !outOk) {
+ return compress128Impl<true, false>(input, numValues, output);
+ }
+ if (!inOk && outOk) {
+ return compress128Impl<false, true>(input, numValues, output);
+ }
+ return compress128Impl<false, false>(input, numValues, output);
+}
+
+// 128-bit decompress. decodeBlock requires aligned uint64_t* input; when
+// the caller's input buffer is unaligned, we stage through tmpIn.
+template <bool InAligned, bool OutAligned>
+inline size_t decompress128Impl(const uint8_t* input, size_t inputSize,
uint8_t* output, size_t outputSize) {
+ alignas(64) uint64_t loBuffer[kMaxValuesPerBlock];
+ alignas(64) uint64_t hiBuffer[kMaxValuesPerBlock];
+ alignas(64) uint64_t tmpIn[kMaxValuesPerBlock + 2];
+ alignas(64) uint64_t tmpOut[kMaxValuesPerBlock * 2];
+
+ const uint8_t* inPtr = input;
+ const uint8_t* inEnd = input + inputSize;
+ const size_t outValuesMax = outputSize / sizeof(__int128_t);
+ size_t nDecoded = 0;
+
+ while (inPtr + kHeaderSize <= inEnd) {
+ if (inPtr[0] == kBwTailMarker) {
+ const uint8_t count = inPtr[1];
+ inPtr += kHeaderSize;
+ const size_t tailBytes = static_cast<size_t>(count) * sizeof(__int128_t);
+ if (inPtr + tailBytes > inEnd || count > outValuesMax - nDecoded) {
+ break;
+ }
+ std::memcpy(output + nDecoded * sizeof(__int128_t), inPtr, tailBytes);
+ nDecoded += count;
+ break;
+ }
+ const size_t blockVals = static_cast<size_t>(inPtr[1]) * kLanes;
+ if (blockVals == 0 || blockVals > kMaxValuesPerBlock || blockVals >
outValuesMax - nDecoded) {
+ break;
+ }
+
+ auto decodeOne = [&](uint64_t* dst) -> bool {
+ const size_t remaining = static_cast<size_t>(inEnd - inPtr);
+ size_t consumed;
+ if constexpr (InAligned) {
+ consumed = decodeBlock(reinterpret_cast<const uint64_t*>(inPtr),
remaining, blockVals, dst);
+ } else {
+ const size_t n = std::min(remaining, sizeof(tmpIn));
+ std::memcpy(tmpIn, inPtr, n);
+ consumed = decodeBlock(tmpIn, n, blockVals, dst);
+ }
+ if (consumed == 0) {
+ return false;
+ }
+ inPtr += consumed;
+ return true;
+ };
+ if (!decodeOne(loBuffer) || !decodeOne(hiBuffer)) {
+ break;
+ }
+
+ uint64_t* dst64 = OutAligned ? reinterpret_cast<uint64_t*>(output +
nDecoded * sizeof(__int128_t)) : tmpOut;
+ for (size_t j = 0; j < blockVals; ++j) {
+ dst64[j * 2] = loBuffer[j];
+ dst64[j * 2 + 1] = hiBuffer[j];
+ }
+ if constexpr (!OutAligned) {
+ std::memcpy(output + nDecoded * sizeof(__int128_t), tmpOut, blockVals *
sizeof(__int128_t));
+ }
+ nDecoded += blockVals;
+ }
+
+ return nDecoded;
+}
+
+// Runtime dispatch — check alignment once, pick the right template.
+inline size_t decompress128(const uint8_t* input, size_t inputSize, uint8_t*
output, size_t outputSize) {
+ const bool inOk = (reinterpret_cast<uintptr_t>(input) % alignof(uint64_t) ==
0);
+ const bool outOk = (reinterpret_cast<uintptr_t>(output) % alignof(uint64_t)
== 0);
+ if (inOk && outOk) {
+ return decompress128Impl<true, true>(input, inputSize, output, outputSize);
}
if (inOk && !outOk) {
- return decompress64Impl<true, false>(input, inputSize, output);
+ return decompress128Impl<true, false>(input, inputSize, output,
outputSize);
}
if (!inOk && outOk) {
- return decompress64Impl<false, true>(input, inputSize, output);
+ return decompress128Impl<false, true>(input, inputSize, output,
outputSize);
}
- return decompress64Impl<false, false>(input, inputSize, output);
+ return decompress128Impl<false, false>(input, inputSize, output, outputSize);
}
} // namespace ffor
diff --git a/cpp/velox/shuffle/VeloxTypeAwareCompress.h
b/cpp/velox/shuffle/VeloxTypeAwareCompress.h
index fbbabc5c06..370ee92151 100644
--- a/cpp/velox/shuffle/VeloxTypeAwareCompress.h
+++ b/cpp/velox/shuffle/VeloxTypeAwareCompress.h
@@ -28,6 +28,8 @@ inline int8_t veloxTypeToTacType(facebook::velox::TypeKind
kind) {
switch (kind) {
case facebook::velox::TypeKind::BIGINT:
return tac::kUInt64;
+ case facebook::velox::TypeKind::HUGEINT:
+ return tac::kUInt128;
default:
return tac::kUnsupported;
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
