This is an automated email from the ASF dual-hosted git repository.
jinchengchenghh 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 9e52cbda69 [VL] Fix host to GPU conversion in CudfVectorStream (#13067)
9e52cbda69 is described below
commit 9e52cbda69a83fec2a768385ede2bffaece30e72
Author: Reema <[email protected]>
AuthorDate: Fri Oct 23 11:37:15 2026 +0300
[VL] Fix host to GPU conversion in CudfVectorStream (#13067)
---
cpp/velox/operators/plannodes/CudfVectorStream.h | 39 +++-
cpp/velox/operators/plannodes/RowVectorStream.cc | 7 +-
cpp/velox/tests/CMakeLists.txt | 2 +
cpp/velox/tests/CudfValueStreamTest.cc | 206 +++++++++++++++++++++
cpp/velox/tests/ValueStreamDynamicFilterTest.cc | 32 ++++
.../GlutenBroadcastValueStreamTestSupport.scala | 29 +++
.../GlutenBroadcastValueStreamTestSupport.scala | 27 +++
.../GlutenBroadcastValueStreamTestSupport.scala | 29 +++
.../gluten/utils/velox/VeloxTestSettings.scala | 1 +
.../org/apache/spark/sql/GlutenSQLQuerySuite.scala | 52 ++++++
.../GlutenBroadcastValueStreamSuite.scala | 101 ++++++++++
11 files changed, 519 insertions(+), 6 deletions(-)
diff --git a/cpp/velox/operators/plannodes/CudfVectorStream.h
b/cpp/velox/operators/plannodes/CudfVectorStream.h
index 492b8556be..cdde9d8d48 100644
--- a/cpp/velox/operators/plannodes/CudfVectorStream.h
+++ b/cpp/velox/operators/plannodes/CudfVectorStream.h
@@ -24,6 +24,7 @@
#include "velox/exec/Operator.h"
#include "velox/exec/Task.h"
#include "velox/experimental/cudf/exec/CudfOperator.h"
+#include "velox/experimental/cudf/exec/GpuResources.h"
#include "velox/experimental/cudf/exec/Utilities.h"
#include "velox/experimental/cudf/exec/VeloxCudfInterop.h"
#include "velox/experimental/cudf/vector/CudfVector.h"
@@ -115,13 +116,41 @@ class CudfVectorStream : public CudfVectorStreamBase {
VELOX_DCHECK(vp != nullptr);
auto cudfVector =
std::dynamic_pointer_cast<facebook::velox::cudf_velox::CudfVector>(vp);
if (cudfVector == nullptr) {
- // The vector may comes from BroadcastExchange, in this case, it's not a
CudfVector.
- vp->setType(outputType_);
- return vp;
+ // BroadcastExchange may return a host RowVector; upload it for GPU
operators.
+ auto stream =
facebook::velox::cudf_velox::cudfGlobalStreamPool().get_stream();
+ if (outputType_->size() == 0) {
+ // Preserve row count because zero-column cuDF tables cannot store it.
+ return std::make_shared<facebook::velox::cudf_velox::CudfVector>(
+ vp->pool(), outputType_, vp->size(),
std::make_unique<cudf::table>(), stream);
+ }
+ // Broadcast preprojection can append join-key columns beyond the
stream's output schema.
+ VELOX_CHECK_GE(
+ vp->childrenSize(),
+ outputType_->size(),
+ "Value stream batch has fewer columns than the declared output
type");
+ std::vector<facebook::velox::VectorPtr> children(
+ vp->children().begin(), vp->children().begin() +
outputType_->size());
+ auto host = std::make_shared<facebook::velox::RowVector>(
+ vp->pool(), outputType_, facebook::velox::BufferPtr(0), vp->size(),
std::move(children));
+ auto table = facebook::velox::cudf_velox::with_arrow::toCudfTable(
+ host, host->pool(), stream,
facebook::velox::cudf_velox::get_output_mr());
+ return std::make_shared<facebook::velox::cudf_velox::CudfVector>(
+ vp->pool(), outputType_, vp->size(), std::move(table), stream);
+ }
+ // GPU broadcast deserialization can retain the same extra join-key
columns.
+ VELOX_CHECK_GE(
+ cudfVector->getTableView().num_columns(),
+ outputType_->size(),
+ "Value stream batch has fewer columns than the declared output type");
+ auto stream = cudfVector->stream();
+ auto table = cudfVector->release();
+ if (table->num_columns() > outputType_->size()) {
+ auto columns = table->release();
+ columns.resize(outputType_->size());
+ table = std::make_unique<cudf::table>(std::move(columns));
}
- VELOX_CHECK_NOT_NULL(cudfVector);
return std::make_shared<facebook::velox::cudf_velox::CudfVector>(
- vp->pool(), outputType_, vp->size(), cudfVector->release(),
cudfVector->stream());
+ vp->pool(), outputType_, vp->size(), std::move(table), stream);
}
};
diff --git a/cpp/velox/operators/plannodes/RowVectorStream.cc
b/cpp/velox/operators/plannodes/RowVectorStream.cc
index 4119eff037..38872620cc 100644
--- a/cpp/velox/operators/plannodes/RowVectorStream.cc
+++ b/cpp/velox/operators/plannodes/RowVectorStream.cc
@@ -102,8 +102,13 @@ facebook::velox::RowVectorPtr RowVectorStream::next() {
const std::shared_ptr<VeloxColumnarBatch>& vb =
VeloxColumnarBatch::from(pool_, cb);
auto vp = vb->getRowVector();
VELOX_DCHECK(vp != nullptr);
+ // Broadcast preprojection can append join-key columns beyond the stream's
output schema.
+ VELOX_CHECK_GE(
+ vp->childrenSize(), outputType_->size(), "Value stream batch has fewer
columns than the declared output type");
+ std::vector<facebook::velox::VectorPtr> children(
+ vp->children().begin(), vp->children().begin() + outputType_->size());
return std::make_shared<facebook::velox::RowVector>(
- vp->pool(), outputType_, facebook::velox::BufferPtr(0), vp->size(),
vp->children());
+ vp->pool(), outputType_, facebook::velox::BufferPtr(0), vp->size(),
std::move(children));
}
ValueStreamDataSource::ValueStreamDataSource(
diff --git a/cpp/velox/tests/CMakeLists.txt b/cpp/velox/tests/CMakeLists.txt
index df00a8db15..a5d7329830 100644
--- a/cpp/velox/tests/CMakeLists.txt
+++ b/cpp/velox/tests/CMakeLists.txt
@@ -106,6 +106,8 @@ if(ENABLE_GPU)
add_velox_gpu_test(velox_gpu_shuffle_writer_test SOURCES
VeloxGpuShuffleWriterTest.cc)
add_velox_gpu_test(velox_gpu_lock_test SOURCES GpuLockTest.cc)
+ add_velox_gpu_test(velox_cudf_value_stream_test SOURCES
+ CudfValueStreamTest.cc)
endif()
set(VELOX_TEST_COMMON_SRCS JsonToProtoConverter.cc FilePathGenerator.cc)
diff --git a/cpp/velox/tests/CudfValueStreamTest.cc
b/cpp/velox/tests/CudfValueStreamTest.cc
new file mode 100644
index 0000000000..83a2de76fe
--- /dev/null
+++ b/cpp/velox/tests/CudfValueStreamTest.cc
@@ -0,0 +1,206 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 "operators/plannodes/CudfVectorStream.h"
+#include "velox/common/base/tests/GTestUtils.h"
+#include "velox/experimental/cudf/CudfConfig.h"
+#include "velox/experimental/cudf/exec/ToCudf.h"
+#include "velox/vector/tests/utils/VectorTestBase.h"
+
+using namespace facebook::velox;
+using namespace facebook::velox::exec;
+
+namespace facebook::velox::test {
+namespace {
+
+class TestBatchIterator final : public gluten::ColumnarBatchIterator {
+ public:
+ explicit TestBatchIterator(std::vector<RowVectorPtr> batches) :
batches_(std::move(batches)) {}
+
+ std::shared_ptr<gluten::ColumnarBatch> next() override {
+ if (index_ == batches_.size()) {
+ return nullptr;
+ }
+ return std::make_shared<gluten::VeloxColumnarBatch>(batches_[index_++]);
+ }
+
+ private:
+ std::vector<RowVectorPtr> batches_;
+ size_t index_ = 0;
+};
+
+class CudfValueStreamTest : public ::testing::Test, public VectorTestBase {
+ protected:
+ static void SetUpTestSuite() {
+
memory::MemoryManager::testingSetInstance(memory::MemoryManager::Options{});
+ cudf_velox::registerCudf();
+
Operator::registerOperator(std::make_unique<gluten::CudfVectorStreamOperatorTranslator>());
+ }
+
+ static void TearDownTestSuite() {
+ cudf_velox::unregisterCudf();
+ Operator::unregisterAllOperators();
+ }
+
+ std::shared_ptr<Task> makeTask(std::vector<RowVectorPtr> batches, const
RowTypePtr& outputType) {
+ auto iterator =
std::make_shared<gluten::ResultIterator>(std::make_unique<TestBatchIterator>(std::move(batches)));
+ auto node = std::make_shared<gluten::CudfValueStreamNode>("cudf-vs",
outputType, std::move(iterator));
+ // Inspect the source's GPU output before the driver adapter adds
CudfToVelox.
+ auto queryCtx =
+ core::QueryCtx::create(nullptr,
core::QueryConfig{{{cudf_velox::CudfConfig::kCudfEnabled, "false"}}});
+ return Task::create("test-cudf-value-stream", core::PlanFragment{node}, 0,
queryCtx, Task::ExecutionMode::kSerial);
+ }
+
+ void assertGpuValues(const RowVectorPtr& expected, const
cudf_velox::CudfVectorPtr& actual) {
+ auto host = cudf_velox::with_arrow::toVeloxColumn(
+ actual->getTableView(), pool(), asRowType(expected->type()), "",
actual->stream(), cudf_velox::get_output_mr());
+ assertEqualVectors(expected, host);
+ }
+};
+
+TEST_F(CudfValueStreamTest, convertsHostBatchToCudfVector) {
+ auto batch = makeRowVector({"id"}, {makeNullableFlatVector<int64_t>({10,
std::nullopt, 30})});
+ auto task = makeTask({batch}, asRowType(batch->type()));
+
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ auto output = task->next(&future);
+ ASSERT_NE(output, nullptr);
+ auto cudfVector = std::dynamic_pointer_cast<cudf_velox::CudfVector>(output);
+ ASSERT_NE(cudfVector, nullptr);
+ EXPECT_EQ(cudfVector->type(), batch->type());
+ EXPECT_EQ(cudfVector->size(), 3);
+ EXPECT_EQ(cudfVector->getTableView().num_rows(), 3);
+ EXPECT_EQ(cudfVector->getTableView().num_columns(), 1);
+ assertGpuValues(batch, cudfVector);
+ EXPECT_EQ(task->next(&future), nullptr);
+}
+
+TEST_F(CudfValueStreamTest, trimsExtraHostColumns) {
+ auto batch =
+ makeRowVector({"id", "col_0"}, {makeFlatVector<int64_t>({10, 20, 30}),
makeFlatVector<int64_t>({101, 202, 303})});
+ auto outputType = ROW({"id"}, {BIGINT()});
+ auto expected = makeRowVector({"id"}, {batch->childAt(0)});
+ auto task = makeTask({batch}, outputType);
+
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ auto output = task->next(&future);
+ ASSERT_NE(output, nullptr);
+ auto cudfVector = std::dynamic_pointer_cast<cudf_velox::CudfVector>(output);
+ ASSERT_NE(cudfVector, nullptr);
+ EXPECT_EQ(cudfVector->type(), outputType);
+ EXPECT_EQ(cudfVector->size(), 3);
+ EXPECT_EQ(cudfVector->getTableView().num_columns(), 1);
+ assertGpuValues(expected, cudfVector);
+ EXPECT_EQ(task->next(&future), nullptr);
+}
+
+TEST_F(CudfValueStreamTest, preservesZeroColumnRowCount) {
+ auto outputType = ROW({}, {});
+ auto batch = makeRowVector(outputType, 3);
+ auto task = makeTask({batch}, outputType);
+
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ auto output = task->next(&future);
+ ASSERT_NE(output, nullptr);
+ auto cudfVector = std::dynamic_pointer_cast<cudf_velox::CudfVector>(output);
+ ASSERT_NE(cudfVector, nullptr);
+ EXPECT_EQ(cudfVector->type(), outputType);
+ EXPECT_EQ(cudfVector->size(), 3);
+ EXPECT_EQ(cudfVector->getTableView().num_columns(), 0);
+ EXPECT_EQ(task->next(&future), nullptr);
+}
+
+TEST_F(CudfValueStreamTest, rejectsHostBatchWithMissingColumns) {
+ for (int32_t numColumns : {0, 1}) {
+ SCOPED_TRACE(numColumns);
+ auto batch = numColumns == 0 ? makeRowVector(ROW({}, {}), 3)
+ : makeRowVector({"id"},
{makeFlatVector<int64_t>({10, 20, 30})});
+ auto task = makeTask({batch}, ROW({"id", "value"}, {BIGINT(), BIGINT()}));
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ VELOX_ASSERT_THROW(task->next(&future), "Value stream batch has fewer
columns than the declared output type");
+ }
+}
+
+TEST_F(CudfValueStreamTest, loadsLazyHostColumns) {
+ auto lazy = makeLazyFlatVector<int64_t>(3, [](vector_size_t row) { return
(row + 1) * 10; });
+ auto batch = makeRowVector({"id"}, {lazy});
+ auto expected = makeRowVector({"id"}, {makeFlatVector<int64_t>({10, 20,
30})});
+ auto task = makeTask({batch}, asRowType(batch->type()));
+ EXPECT_FALSE(lazy->isLoaded());
+
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ auto output = task->next(&future);
+ ASSERT_NE(output, nullptr);
+ auto cudfVector = std::dynamic_pointer_cast<cudf_velox::CudfVector>(output);
+ ASSERT_NE(cudfVector, nullptr);
+ EXPECT_TRUE(lazy->isLoaded());
+ assertGpuValues(expected, cudfVector);
+ EXPECT_EQ(task->next(&future), nullptr);
+}
+
+TEST_F(CudfValueStreamTest, preservesGpuBatchAndStream) {
+ auto host = makeRowVector({"id"}, {makeNullableFlatVector<int64_t>({10,
std::nullopt, 30})});
+ auto outputType = asRowType(host->type());
+ auto stream = cudf_velox::cudfGlobalStreamPool().get_stream();
+ auto table = cudf_velox::with_arrow::toCudfTable(host, pool(), stream,
cudf_velox::get_output_mr());
+ const auto* data = table->view().column(0).head<int64_t>();
+ auto batch = std::make_shared<cudf_velox::CudfVector>(pool(), outputType,
host->size(), std::move(table), stream);
+ auto task = makeTask({batch}, outputType);
+
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ auto output = task->next(&future);
+ ASSERT_NE(output, nullptr);
+ auto cudfVector = std::dynamic_pointer_cast<cudf_velox::CudfVector>(output);
+ ASSERT_NE(cudfVector, nullptr);
+ EXPECT_EQ(cudfVector->type(), outputType);
+ EXPECT_EQ(cudfVector->size(), host->size());
+ EXPECT_EQ(cudfVector->stream(), stream);
+ EXPECT_EQ(cudfVector->getTableView().column(0).head<int64_t>(), data);
+ assertGpuValues(host, cudfVector);
+ EXPECT_EQ(task->next(&future), nullptr);
+}
+
+TEST_F(CudfValueStreamTest, trimsExtraGpuColumns) {
+ auto host = makeRowVector(
+ {"id", "col_0"},
+ {makeNullableFlatVector<int64_t>({10, std::nullopt, 30}),
makeFlatVector<int64_t>({11, 21, 31})});
+ auto outputType = ROW({"renamed_id"}, {BIGINT()});
+ auto expected = makeRowVector({"renamed_id"}, {host->childAt(0)});
+ auto stream = cudf_velox::cudfGlobalStreamPool().get_stream();
+ auto table = cudf_velox::with_arrow::toCudfTable(host, pool(), stream,
cudf_velox::get_output_mr());
+ const auto* data = table->view().column(0).head<int64_t>();
+ auto batch = std::make_shared<cudf_velox::CudfVector>(pool(), host->type(),
host->size(), std::move(table), stream);
+ auto task = makeTask({batch}, outputType);
+
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ auto output = task->next(&future);
+ ASSERT_NE(output, nullptr);
+ auto cudfVector = std::dynamic_pointer_cast<cudf_velox::CudfVector>(output);
+ ASSERT_NE(cudfVector, nullptr);
+ EXPECT_EQ(cudfVector->type(), outputType);
+ EXPECT_EQ(cudfVector->size(), host->size());
+ EXPECT_EQ(cudfVector->stream(), stream);
+ ASSERT_EQ(cudfVector->getTableView().num_columns(), 1);
+ EXPECT_EQ(cudfVector->getTableView().column(0).head<int64_t>(), data);
+ assertGpuValues(expected, cudfVector);
+ EXPECT_EQ(task->next(&future), nullptr);
+}
+
+} // namespace
+} // namespace facebook::velox::test
diff --git a/cpp/velox/tests/ValueStreamDynamicFilterTest.cc
b/cpp/velox/tests/ValueStreamDynamicFilterTest.cc
index 5c1f156140..899214e935 100644
--- a/cpp/velox/tests/ValueStreamDynamicFilterTest.cc
+++ b/cpp/velox/tests/ValueStreamDynamicFilterTest.cc
@@ -116,6 +116,38 @@ TEST_F(ValueStreamDynamicFilterTest,
noFilterPassesAllRows) {
ASSERT_EQ(ids, (std::vector<int64_t>{10, 20, 30}));
}
+TEST_F(ValueStreamDynamicFilterTest, trimsExtraTrailingColumns) {
+ auto batch = makeRowVector(
+ {"id", "col_0"},
+ {makeNullableFlatVector<int64_t>({10, std::nullopt, 30}),
makeFlatVector<int64_t>({11, 21, 31})});
+ auto originalType = batch->type();
+ auto outputType = ROW({"renamed_id"}, {BIGINT()});
+ auto expected = makeRowVector({"renamed_id"},
{makeNullableFlatVector<int64_t>({10, std::nullopt, 30})});
+ auto scanNode = makeTableScanNode("vs-extra-columns", outputType);
+
+ auto queryCtx = core::QueryCtx::create();
+ auto task =
+ Task::create("test-extra-columns", core::PlanFragment{scanNode}, 0,
queryCtx, Task::ExecutionMode::kSerial);
+ task->addSplit(scanNode->id(), Split{makeSplit({batch})});
+ task->noMoreSplits(scanNode->id());
+
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ auto output = task->next(&future);
+ ASSERT_NE(output, nullptr);
+ EXPECT_EQ(output->type(), outputType);
+ EXPECT_EQ(output->size(), 3);
+ ASSERT_EQ(output->childrenSize(), 1);
+ assertEqualVectors(expected, output);
+
+ // Broadcast batches can be shared; trimming must not modify the source
batch.
+ EXPECT_EQ(batch->type(), originalType);
+ EXPECT_EQ(batch->size(), 3);
+ ASSERT_EQ(batch->childrenSize(), 2);
+ assertEqualVectors(expected->childAt(0), batch->childAt(0));
+ assertEqualVectors(makeFlatVector<int64_t>({11, 21, 31}), batch->childAt(1));
+ EXPECT_EQ(task->next(&future), nullptr);
+}
+
TEST_F(ValueStreamDynamicFilterTest, emptySplitProducesNoRows) {
auto outputType = ROW({"id"}, {BIGINT()});
auto scanNode = makeTableScanNode("vs-empty-split", outputType);
diff --git
a/gluten-ut/spark35/src/test/backends-bolt/org/apache/gluten/GlutenBroadcastValueStreamTestSupport.scala
b/gluten-ut/spark35/src/test/backends-bolt/org/apache/gluten/GlutenBroadcastValueStreamTestSupport.scala
new file mode 100644
index 0000000000..ad8b557cef
--- /dev/null
+++
b/gluten-ut/spark35/src/test/backends-bolt/org/apache/gluten/GlutenBroadcastValueStreamTestSupport.scala
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.
+ */
+package org.apache.gluten
+
+import org.apache.gluten.columnarbatch.ColumnarBatches
+import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators
+
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+trait GlutenBroadcastValueStreamTestSupport {
+
+ protected def loadBroadcastBatch(batch: ColumnarBatch): ColumnarBatch = {
+ ColumnarBatches.load(ArrowBufferAllocators.contextInstance(), batch)
+ }
+}
diff --git
a/gluten-ut/spark35/src/test/backends-clickhouse/org/apache/gluten/GlutenBroadcastValueStreamTestSupport.scala
b/gluten-ut/spark35/src/test/backends-clickhouse/org/apache/gluten/GlutenBroadcastValueStreamTestSupport.scala
new file mode 100644
index 0000000000..90980977dc
--- /dev/null
+++
b/gluten-ut/spark35/src/test/backends-clickhouse/org/apache/gluten/GlutenBroadcastValueStreamTestSupport.scala
@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.
+ */
+package org.apache.gluten
+
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+trait GlutenBroadcastValueStreamTestSupport {
+
+ protected def loadBroadcastBatch(batch: ColumnarBatch): ColumnarBatch = {
+ // ClickHouse column vectors already expose Spark data types.
+ batch
+ }
+}
diff --git
a/gluten-ut/spark35/src/test/backends-velox/org/apache/gluten/GlutenBroadcastValueStreamTestSupport.scala
b/gluten-ut/spark35/src/test/backends-velox/org/apache/gluten/GlutenBroadcastValueStreamTestSupport.scala
new file mode 100644
index 0000000000..ad8b557cef
--- /dev/null
+++
b/gluten-ut/spark35/src/test/backends-velox/org/apache/gluten/GlutenBroadcastValueStreamTestSupport.scala
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.
+ */
+package org.apache.gluten
+
+import org.apache.gluten.columnarbatch.ColumnarBatches
+import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators
+
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+trait GlutenBroadcastValueStreamTestSupport {
+
+ protected def loadBroadcastBatch(batch: ColumnarBatch): ColumnarBatch = {
+ ColumnarBatches.load(ArrowBufferAllocators.contextInstance(), batch)
+ }
+}
diff --git
a/gluten-ut/spark35/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
b/gluten-ut/spark35/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
index a16b97c357..08c584714a 100644
---
a/gluten-ut/spark35/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
+++
b/gluten-ut/spark35/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
@@ -507,6 +507,7 @@ class VeloxTestSettings extends BackendTestSettings {
enableSuite[GlutenOuterJoinSuiteForceShjOff]
enableSuite[FallbackStrategiesSuite]
enableSuite[GlutenBroadcastExchangeSuite]
+ enableSuite[GlutenBroadcastValueStreamSuite]
enableSuite[GlutenLocalBroadcastExchangeSuite]
enableSuite[GlutenCoalesceShufflePartitionsSuite]
// Rewrite for Gluten. Change details are in the inline comments in
individual tests.
diff --git
a/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/GlutenSQLQuerySuite.scala
b/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/GlutenSQLQuerySuite.scala
index 8f397c517e..e3269994d8 100644
---
a/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/GlutenSQLQuerySuite.scala
+++
b/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/GlutenSQLQuerySuite.scala
@@ -16,6 +16,10 @@
*/
package org.apache.spark.sql
+import org.apache.gluten.config.GlutenConfig
+import org.apache.gluten.execution.{BroadcastHashJoinExecTransformerBase,
TakeOrderedAndProjectExecTransformer}
+import org.apache.gluten.utils.BackendTestUtils
+
import org.apache.spark.SparkException
import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
import org.apache.spark.sql.internal.SQLConf
@@ -23,6 +27,54 @@ import org.apache.spark.sql.internal.SQLConf
class GlutenSQLQuerySuite extends SQLQuerySuite with GlutenSQLTestsTrait {
import testImplicits._
+ testGluten("broadcast expression join followed by aggregate and top N") {
+ withTempPath {
+ dir =>
+ val leftPath = s"${dir.getCanonicalPath}/left"
+ val rightPath = s"${dir.getCanonicalPath}/right"
+ Seq((1, 10), (2, 20), (3, 30), (4, 40), (2, 5))
+ .toDF("id", "value")
+ .repartition(2)
+ .write
+ .parquet(leftPath)
+ Seq(0, 1, 2).toDF("id").write.parquet(rightPath)
+
+ withTempView("top_n_left", "top_n_right") {
+ spark.read.parquet(leftPath).createOrReplaceTempView("top_n_left")
+ spark.read.parquet(rightPath).createOrReplaceTempView("top_n_right")
+
+ val isVelox = BackendTestUtils.isVeloxBackendLoaded()
+ val cudfModes = if (isVelox) Seq(false, true) else Seq(false)
+ for (cudfEnabled <- cudfModes; buildOnce <- Seq(true, false)) {
+ withClue(s"cuDF=$cudfEnabled, buildOnce=$buildOnce: ") {
+ withSQLConf(
+ // Exercise CPU execution with both cuDF broadcast settings on
CPU CI.
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "134217728",
+ GlutenConfig.COLUMNAR_CUDF_ENABLED.key -> cudfEnabled.toString,
+ "spark.gluten.velox.buildHashTableOncePerExecutor.enabled" ->
buildOnce.toString,
+ "spark.gluten.velox.offHeapBroadcastBuildRelation.enabled" ->
"false"
+ ) {
+ // Exercise both broadcast expression-key projection and the
exchange before TopN.
+ val df = sql("""
+ |SELECT /*+ BROADCAST(r) */ l.id, SUM(l.value)
AS total
+ |FROM top_n_left l JOIN top_n_right r ON l.id =
r.id + 1
+ |GROUP BY l.id
+ |ORDER BY total DESC, l.id
+ |LIMIT 2
+ |""".stripMargin)
+ val plan = getExecutedPlan(df)
+
assert(plan.exists(_.isInstanceOf[BroadcastHashJoinExecTransformerBase]))
+
assert(plan.exists(_.isInstanceOf[TakeOrderedAndProjectExecTransformer]))
+ checkAnswer(df, Seq(Row(3, 30L), Row(2, 25L)))
+ }
+ }
+ }
+ }
+ }
+ }
+
testGluten("SPARK-28156: self-join should not miss cached view") {
withTable("table1") {
withView("table1_vw") {
diff --git
a/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/execution/GlutenBroadcastValueStreamSuite.scala
b/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/execution/GlutenBroadcastValueStreamSuite.scala
new file mode 100644
index 0000000000..dfc1871c69
--- /dev/null
+++
b/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/execution/GlutenBroadcastValueStreamSuite.scala
@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.
+ */
+package org.apache.spark.sql.execution
+
+import org.apache.gluten.GlutenBroadcastValueStreamTestSupport
+import org.apache.gluten.config.GlutenConfig
+import org.apache.gluten.execution.{BroadcastHashJoinExecTransformerBase,
TakeOrderedAndProjectExecTransformer}
+
+import org.apache.spark.sql.{GlutenSQLTestsTrait, Row}
+import org.apache.spark.sql.execution.joins.BuildSideRelation
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.task.TaskResources
+
+class GlutenBroadcastValueStreamSuite
+ extends GlutenSQLTestsTrait
+ with GlutenBroadcastValueStreamTestSupport {
+ import testImplicits._
+
+ testGluten("broadcast expression join preserves batch schema") {
+ withTempPath {
+ dir =>
+ val leftPath = s"${dir.getCanonicalPath}/left"
+ val rightPath = s"${dir.getCanonicalPath}/right"
+ Seq((1, 10), (2, 20), (3, 30), (4, 40), (2, 5))
+ .toDF("id", "value")
+ .repartition(2)
+ .write
+ .parquet(leftPath)
+ Seq(0, 1, 2).toDF("id").write.parquet(rightPath)
+
+ withTempView("top_n_left", "top_n_right") {
+ spark.read.parquet(leftPath).createOrReplaceTempView("top_n_left")
+ spark.read.parquet(rightPath).createOrReplaceTempView("top_n_right")
+
+ for (cudfEnabled <- Seq(false, true); buildOnce <- Seq(true, false))
{
+ withClue(s"cuDF=$cudfEnabled, buildOnce=$buildOnce: ") {
+ withSQLConf(
+ // Exercise CPU execution with both cuDF broadcast settings on
CPU CI.
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "134217728",
+ GlutenConfig.COLUMNAR_CUDF_ENABLED.key -> cudfEnabled.toString,
+ "spark.gluten.velox.buildHashTableOncePerExecutor.enabled" ->
buildOnce.toString,
+ "spark.gluten.velox.offHeapBroadcastBuildRelation.enabled" ->
"false"
+ ) {
+ // Exercise both broadcast expression-key projection and the
exchange before TopN.
+ val df = sql("""
+ |SELECT /*+ BROADCAST(r) */ l.id, SUM(l.value)
AS total
+ |FROM top_n_left l JOIN top_n_right r ON l.id =
r.id + 1
+ |GROUP BY l.id
+ |ORDER BY total DESC, l.id
+ |LIMIT 2
+ |""".stripMargin)
+ val plan = getExecutedPlan(df)
+
assert(plan.exists(_.isInstanceOf[BroadcastHashJoinExecTransformerBase]))
+
assert(plan.exists(_.isInstanceOf[TakeOrderedAndProjectExecTransformer]))
+ checkAnswer(df, Seq(Row(3, 30L), Row(2, 25L)))
+
+ val exchanges = plan.collect { case e:
ColumnarBroadcastExchangeExec => e }
+ assert(exchanges.size == 1)
+ val exchange = exchanges.head
+ val relation =
exchange.doExecuteBroadcast[BuildSideRelation]().value
+ val outputTypes = exchange.output.map(_.dataType)
+ // Broadcast preprojection appends the expression key in
either cuDF mode.
+ val expectedColumns = outputTypes.size + (if (buildOnce) 1
else 0)
+
+ TaskResources.runUnsafe {
+ // The iterator recycles loaded batches and closes its
serializer when drained.
+ val batches = relation.asReadOnlyCopy().deserialized
+ var rows = 0
+ while (batches.hasNext) {
+ val batch = batches.next()
+ assert(batch.numCols() == expectedColumns)
+ val loaded = loadBroadcastBatch(batch)
+ val types = (0 until loaded.numCols()).map(i =>
loaded.column(i).dataType())
+ assert(types.take(outputTypes.size) == outputTypes)
+ rows += loaded.numRows()
+ }
+ assert(rows == 3)
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
