From 4b866b04c43e607aa0432144cfca7fea2a21e105 Mon Sep 17 00:00:00 2001 From: zhangjinyuan Date: Tue, 28 Jul 2026 16:55:06 +0800 Subject: [PATCH] fix array_repeat above Presto repeat limit --- bolt/functions/lib/Repeat.cpp | 52 ++++++++++++++++++------- bolt/functions/lib/tests/RepeatTest.cpp | 17 ++++++++ 2 files changed, 54 insertions(+), 15 deletions(-) diff --git a/bolt/functions/lib/Repeat.cpp b/bolt/functions/lib/Repeat.cpp index d9dcf4e2e..e22565dc1 100644 --- a/bolt/functions/lib/Repeat.cpp +++ b/bolt/functions/lib/Repeat.cpp @@ -30,6 +30,9 @@ #include "bolt/expression/DecodedArgs.h" #include "bolt/expression/VectorFunction.h" + +#include + namespace bytedance::bolt::functions { namespace { // See documentation at https://prestodb.io/docs/current/functions/array.html @@ -37,8 +40,11 @@ class RepeatFunction : public exec::VectorFunction { public: // @param allowNegativeCount If true, negative 'count' is allowed // and treated the same as zero (Spark's behavior). - explicit RepeatFunction(bool allowNegativeCount) - : allowNegativeCount_(allowNegativeCount) {} + // @param enforceMaxResultEntries If true, enforce Presto's 10,000 entries + // limit per row. Spark's array_repeat doesn't have this limit. + explicit RepeatFunction(bool allowNegativeCount, bool enforceMaxResultEntries) + : allowNegativeCount_(allowNegativeCount), + enforceMaxResultEntries_(enforceMaxResultEntries) {} static constexpr int32_t kMaxResultEntries = 10'000; @@ -67,7 +73,10 @@ class RepeatFunction : public exec::VectorFunction { private: // Check count to make sure it is in valid range. - static int32_t checkCount(int32_t count, bool allowNegativeCount) { + static int32_t checkCount( + int32_t count, + bool allowNegativeCount, + bool enforceMaxResultEntries) { if (count < 0) { if (allowNegativeCount) { return 0; @@ -77,13 +86,23 @@ class RepeatFunction : public exec::VectorFunction { count, 0); } - BOLT_USER_CHECK_LE( - count, - kMaxResultEntries, - "Count argument of repeat function must be less than or equal to 10000"); + if (enforceMaxResultEntries) { + BOLT_USER_CHECK_LE( + count, + kMaxResultEntries, + "Count argument of repeat function must be less than or equal to 10000"); + } return count; } + static vector_size_t checkTotalCount(int64_t totalCount) { + BOLT_USER_CHECK_LE( + totalCount, + std::numeric_limits::max(), + "Total entries of repeat result exceed the maximum vector size"); + return static_cast(totalCount); + } + VectorPtr applyConstantCount( const SelectivityVector& rows, std::vector& args, @@ -99,13 +118,14 @@ class RepeatFunction : public exec::VectorFunction { } auto count = constantCount->valueAt(0); + vector_size_t totalCount; try { - count = checkCount(count, allowNegativeCount_); + count = checkCount(count, allowNegativeCount_, enforceMaxResultEntries_); + totalCount = checkTotalCount(static_cast(count) * numRows); } catch (const BoltUserError&) { context.setErrors(rows, std::current_exception()); return nullptr; } - const auto totalCount = count * numRows; // Allocate new vectors for indices, lengths and offsets. BufferPtr indices = allocateIndices(totalCount, pool); @@ -140,13 +160,14 @@ class RepeatFunction : public exec::VectorFunction { exec::EvalCtx& context) const { exec::DecodedArgs decodedArgs(rows, args, context); auto countDecoded = decodedArgs.at(1); - int32_t totalCount = 0; + int64_t totalCount = 0; context.applyToSelectedNoThrow(rows, [&](auto row) { auto count = countDecoded->isNullAt(row) ? 0 : countDecoded->valueAt(row); - count = checkCount(count, allowNegativeCount_); + count = checkCount(count, allowNegativeCount_, enforceMaxResultEntries_); totalCount += count; }); + const auto totalEntries = checkTotalCount(totalCount); const auto numRows = rows.end(); auto pool = context.pool(); @@ -160,7 +181,7 @@ class RepeatFunction : public exec::VectorFunction { } // Allocate new vectors for indices, lengths and offsets. - BufferPtr indices = allocateIndices(totalCount, pool); + BufferPtr indices = allocateIndices(totalEntries, pool); BufferPtr sizes = allocateSizes(numRows, pool); BufferPtr offsets = allocateOffsets(numRows, pool); auto rawIndices = indices->asMutable(); @@ -196,10 +217,11 @@ class RepeatFunction : public exec::VectorFunction { numRows, offsets, sizes, - BaseVector::wrapInDictionary(nullptr, indices, totalCount, args[0])); + BaseVector::wrapInDictionary(nullptr, indices, totalEntries, args[0])); } const bool allowNegativeCount_; + const bool enforceMaxResultEntries_; }; } // namespace @@ -217,14 +239,14 @@ std::shared_ptr makeRepeat( const std::string& /* name */, const std::vector& /* inputArgs */, const core::QueryConfig& /*config*/) { - return std::make_unique(false); + return std::make_unique(false, true); } std::shared_ptr makeRepeatAllowNegativeCount( const std::string& /* name */, const std::vector& /* inputArgs */, const core::QueryConfig& /*config*/) { - return std::make_unique(true); + return std::make_unique(true, false); } } // namespace bytedance::bolt::functions diff --git a/bolt/functions/lib/tests/RepeatTest.cpp b/bolt/functions/lib/tests/RepeatTest.cpp index 6d4f6e913..63f686c87 100644 --- a/bolt/functions/lib/tests/RepeatTest.cpp +++ b/bolt/functions/lib/tests/RepeatTest.cpp @@ -174,5 +174,22 @@ TEST_F(RepeatTest, repeatAllowNegativeCount) { {elementVector, countVector}, expected); } + +TEST_F(RepeatTest, repeatAllowNegativeCountAbovePrestoLimit) { + constexpr int32_t repeatCount = 10'452; + const auto elementVector = makeFlatVector({repeatCount}); + const auto countVector = makeFlatVector({repeatCount}); + const auto expected = makeArrayVector( + {std::vector(repeatCount, repeatCount)}); + + testExpression( + "repeat_allow_negative_count(C0, C1)", + {elementVector, countVector}, + expected); + testExpression( + "repeat_allow_negative_count(C0, '10452'::INTEGER)", + {elementVector}, + expected); +} } // namespace } // namespace bytedance::bolt::functions