diff --git a/cpp/velox/operators/plannodes/RowVectorStream.cc b/cpp/velox/operators/plannodes/RowVectorStream.cc index 4baa1f84712..dc2ef251463 100644 --- a/cpp/velox/operators/plannodes/RowVectorStream.cc +++ b/cpp/velox/operators/plannodes/RowVectorStream.cc @@ -99,9 +99,14 @@ std::shared_ptr RowVectorStream::nextInternal() { facebook::velox::RowVectorPtr RowVectorStream::next() { auto cb = nextInternal(); + if (cb == nullptr || cb->numRows() == 0) { + return nullptr; + } const std::shared_ptr& vb = VeloxColumnarBatch::from(pool_, cb); auto vp = vb->getRowVector(); - VELOX_DCHECK(vp != nullptr); + if (vp == nullptr || vp->size() == 0) { + return nullptr; + } return std::make_shared( vp->pool(), outputType_, facebook::velox::BufferPtr(0), vp->size(), vp->children()); } @@ -136,47 +141,49 @@ void ValueStreamDataSource::addSplit(std::shared_ptr ValueStreamDataSource::next( uint64_t size, facebook::velox::ContinueFuture& future) { - // Try to get current iterator if we don't have one - while (!currentIterator_) { - if (pendingIterators_.empty()) { - // No more iterators to process - return nullptr; - } - - // Get next RowVectorStream from queue - currentIterator_ = pendingIterators_.front(); - pendingIterators_.erase(pendingIterators_.begin()); - } + for (;;) { + // Try to get current iterator if we don't have one. + while (!currentIterator_) { + if (pendingIterators_.empty()) { + // No more iterators to process. Return an engaged null RowVectorPtr to + // tell TableScan the current split is finished, not connector-blocked. + return facebook::velox::RowVectorPtr(nullptr); + } - // Check if current stream has more data - if (!currentIterator_->hasNext()) { - // Current stream exhausted, try next one - currentIterator_ = nullptr; - return next(size, future); // Recursively try next stream - } + // Get next RowVectorStream from queue. + currentIterator_ = pendingIterators_.front(); + pendingIterators_.erase(pendingIterators_.begin()); + } - // Get next batch from current stream (RowVectorStream handles conversion) - auto rowVector = currentIterator_->next(); + // Check if current stream has more data. + if (!currentIterator_->hasNext()) { + currentIterator_ = nullptr; + continue; + } - if (!rowVector) { - currentIterator_ = nullptr; - return next(size, future); // Recursively try next stream - } + // Get next batch from current stream. Null or zero-row batches are skipped + // in-loop: a long run of empty batches must not grow the native stack, so + // no recursion here. + auto rowVector = currentIterator_->next(); + if (!rowVector) { + continue; + } - // Update metrics - completedRows_ += rowVector->size(); - completedBytes_ += rowVector->estimateFlatSize(); + // Update metrics. + completedRows_ += rowVector->size(); + completedBytes_ += rowVector->estimateFlatSize(); - // Apply dynamic filters if any have been pushed down. - if (!dynamicFilters_.empty()) { - rowVector = applyDynamicFilters(rowVector); - if (!rowVector) { - // All rows filtered out, try next batch. - return next(size, future); + // Apply dynamic filters if any have been pushed down. + if (!dynamicFilters_.empty()) { + rowVector = applyDynamicFilters(rowVector); + if (!rowVector) { + // All rows filtered out, try next batch. + continue; + } } - } - return rowVector; + return rowVector; + } } facebook::velox::RowVectorPtr ValueStreamDataSource::applyDynamicFilters(const facebook::velox::RowVectorPtr& input) { @@ -278,4 +285,4 @@ void ValueStreamDataSource::applyFilterOnColumn( rows.updateBounds(); } -} // namespace gluten \ No newline at end of file +} // namespace gluten diff --git a/cpp/velox/tests/ValueStreamDynamicFilterTest.cc b/cpp/velox/tests/ValueStreamDynamicFilterTest.cc index f9ba13bce6d..53c57549d4d 100644 --- a/cpp/velox/tests/ValueStreamDynamicFilterTest.cc +++ b/cpp/velox/tests/ValueStreamDynamicFilterTest.cc @@ -18,7 +18,10 @@ #include #include "memory/VeloxColumnarBatch.h" +#include "operators/functions/RegistrationAllFunctions.h" #include "operators/plannodes/RowVectorStream.h" +#include "velox/exec/tests/utils/PlanBuilder.h" +#include "velox/parse/TypeResolver.h" #include "velox/type/Filter.h" #include "velox/vector/DecodedVector.h" #include "velox/vector/FlatVector.h" @@ -51,6 +54,8 @@ class ValueStreamDynamicFilterTest : public ::testing::Test, public VectorTestBa protected: static void SetUpTestCase() { memory::MemoryManager::testingSetInstance(memory::MemoryManager::Options{}); + gluten::registerAllFunctions(); + parse::registerTypeResolver(); } void SetUp() override { @@ -116,6 +121,58 @@ TEST_F(ValueStreamDynamicFilterTest, noFilterPassesAllRows) { ASSERT_EQ(ids, (std::vector{10, 20, 30})); } +TEST_F(ValueStreamDynamicFilterTest, skipsEmptyBatchesWithinSplit) { + auto empty = makeRowVector({"id"}, {makeFlatVector(std::vector{})}); + auto batch = makeRowVector({"id"}, {makeFlatVector({40, 50})}); + auto outputType = asRowType(batch->type()); + auto scanNode = makeTableScanNode("vs-empty", outputType); + + auto queryCtx = core::QueryCtx::create(); + auto task = + Task::create("test-empty-batches", core::PlanFragment{scanNode}, 0, queryCtx, Task::ExecutionMode::kSerial); + + task->addSplit(scanNode->id(), Split{makeSplit({empty, batch, empty})}); + task->noMoreSplits(scanNode->id()); + + auto ids = readAllInt64(task.get()); + ASSERT_EQ(ids, (std::vector{40, 50})); +} + +TEST_F(ValueStreamDynamicFilterTest, emptySplitProducesNoRows) { + auto outputType = ROW({"id"}, {BIGINT()}); + auto scanNode = makeTableScanNode("vs-empty-split", outputType); + + auto queryCtx = core::QueryCtx::create(); + auto task = Task::create("test-empty-split", core::PlanFragment{scanNode}, 0, queryCtx, Task::ExecutionMode::kSerial); + + task->addSplit(scanNode->id(), Split{makeSplit({})}); + task->noMoreSplits(scanNode->id()); + + auto ids = readAllInt64(task.get()); + ASSERT_TRUE(ids.empty()); +} + +TEST_F(ValueStreamDynamicFilterTest, filterProjectOverValueStream) { + auto batch = makeRowVector({"id"}, {makeFlatVector({1, 2, 3})}); + auto outputType = asRowType(batch->type()); + auto scanNode = makeTableScanNode("vs-filter-project", outputType); + auto planNodeIdGenerator = std::make_shared(); + auto planNode = exec::test::PlanBuilder(scanNode, planNodeIdGenerator, pool_.get()) + .filter("greaterthan(id, 1)") + .project({"add(id, 1)"}) + .planNode(); + + auto queryCtx = core::QueryCtx::create(); + auto task = + Task::create("test-filter-project", core::PlanFragment{planNode}, 0, queryCtx, Task::ExecutionMode::kSerial); + + task->addSplit(scanNode->id(), Split{makeSplit({batch})}); + task->noMoreSplits(scanNode->id()); + + auto ids = readAllInt64(task.get()); + ASSERT_EQ(ids, (std::vector{3, 4})); +} + // Test that filtering works when filter is injected after first batch. TEST_F(ValueStreamDynamicFilterTest, filterBigintRange) { auto batch1 = makeRowVector({"id"}, {makeFlatVector({1, 2, 3, 4, 5})});