From 48615d8326b7077586137374c7e56d3611ce40d3 Mon Sep 17 00:00:00 2001 From: Alfred Xu Date: Mon, 24 Aug 2026 16:12:54 +0900 Subject: [PATCH 1/2] feat(cudf): support Spark xxhash64 and runtime bloom filters - Register xxhash64_with_seed and might_contain in the Spark cuDF expression adapter. - Add reduce-only bloom_filter_agg support for raw and intermediate/final aggregation. - Cover expression selection, xxhash64 parity, might_contain, and Bloom-filter aggregation. - Include the CudfReduce parenthesis correction required for GCC compilation. - Fuse 34d5c2212dfe956396babe1cfd88575884930371 and 2842252ee9d1217c35f61d84888c618bbd4e757a on dev 28798f62306ed2bf7db4b2161b91f72ceabb3c15. --- .../cudf/exec/CudfAggregation.cpp | 14 +- .../experimental/cudf/exec/CudfAggregation.h | 5 + velox/experimental/cudf/exec/CudfReduce.cpp | 322 +++++++++++++++++- velox/experimental/cudf/exec/CudfReduce.h | 5 +- .../cudf/exec/SparkAggregateFunctions.cpp | 42 +++ .../cudf/expression/CMakeLists.txt | 2 + .../cudf/expression/SparkFunctions.cpp | 48 +++ .../sparksql/MightContainFunction.cpp | 200 +++++++++++ .../sparksql/MightContainFunction.h | 43 +++ .../expression/sparksql/XxHash64Function.cpp | 90 +++++ .../expression/sparksql/XxHash64Function.h | 40 +++ .../ExpressionEvaluatorSelectionTest.cpp | 35 ++ .../cudf/tests/sparksql/AggregationTest.cpp | 82 +++++ .../cudf/tests/sparksql/FilterProjectTest.cpp | 48 +++ 14 files changed, 970 insertions(+), 6 deletions(-) create mode 100644 velox/experimental/cudf/expression/sparksql/MightContainFunction.cpp create mode 100644 velox/experimental/cudf/expression/sparksql/MightContainFunction.h create mode 100644 velox/experimental/cudf/expression/sparksql/XxHash64Function.cpp create mode 100644 velox/experimental/cudf/expression/sparksql/XxHash64Function.h diff --git a/velox/experimental/cudf/exec/CudfAggregation.cpp b/velox/experimental/cudf/exec/CudfAggregation.cpp index 12ecef7b877..5856893e7d9 100644 --- a/velox/experimental/cudf/exec/CudfAggregation.cpp +++ b/velox/experimental/cudf/exec/CudfAggregation.cpp @@ -206,6 +206,12 @@ std::vector resolveAggregateInfos( const auto isDecimalAggregate = aggregate.rawInputTypes.size() == 1 && aggregate.rawInputTypes[0]->isDecimal(); + std::vector extraInputs; + if (aggregate.call->inputs().size() > 1) { + extraInputs.assign( + aggregate.call->inputs().begin() + 1, aggregate.call->inputs().end()); + } + params.emplace_back( companionStep, aggregate.call->name(), @@ -215,7 +221,8 @@ std::vector resolveAggregateInfos( isCountFunctionName(aggregate.call->name()) ? std::make_optional(getCountInputKind(aggregate, constants[i])) : std::nullopt, - isDecimalAggregate); + isDecimalAggregate, + std::move(extraInputs)); } return params; } @@ -279,6 +286,11 @@ AggregationInputChannels buildAggregationInputChannels( } else if ( auto constant = dynamic_cast(arg.get())) { + if (!aggInputs.empty()) { + // Extra constant arguments (bloom_filter_agg estimatedNumItems and + // numBits) are not aggregation input columns. + continue; + } result.constants[i] = constant->toConstantVector(operatorCtx.pool()); aggInputs.push_back(fallbackChannel); } else { diff --git a/velox/experimental/cudf/exec/CudfAggregation.h b/velox/experimental/cudf/exec/CudfAggregation.h index 727b042b776..be7959728c9 100644 --- a/velox/experimental/cudf/exec/CudfAggregation.h +++ b/velox/experimental/cudf/exec/CudfAggregation.h @@ -19,11 +19,13 @@ #include "velox/experimental/cudf/expression/ExpressionEvaluator.h" #include "velox/experimental/cudf/vector/CudfVector.h" +#include "velox/core/ITypedExpr.h" #include "velox/exec/Operator.h" #include "velox/expression/FunctionSignature.h" #include #include +#include namespace facebook::velox::cudf_velox { @@ -67,6 +69,9 @@ struct ResolvedAggregateInfo { // Routing keys off the function family, not the physical batch type (which is // VARBINARY/STRING on intermediate and final steps). bool isDecimalAggregate; + // Extra call arguments after the first aggregation input, e.g. constant + // estimatedNumItems/numBits for bloom_filter_agg. + std::vector extraInputs; }; // Parse aggregate inputs from the aggregation node and resolve companion steps, diff --git a/velox/experimental/cudf/exec/CudfReduce.cpp b/velox/experimental/cudf/exec/CudfReduce.cpp index fe6dc2489b8..d332dbc0ff0 100644 --- a/velox/experimental/cudf/exec/CudfReduce.cpp +++ b/velox/experimental/cudf/exec/CudfReduce.cpp @@ -25,20 +25,33 @@ #include "velox/experimental/cudf/exec/VeloxCudfInterop.h" #include "velox/experimental/cudf/expression/ExpressionEvaluator.h" +#include "velox/common/base/BloomFilter.h" +#include "velox/core/Expressions.h" +#include "velox/core/QueryConfig.h" #include "velox/exec/Aggregate.h" #include "velox/exec/AggregateFunctionRegistry.h" #include "velox/exec/Task.h" #include "velox/expression/Expr.h" +#include "velox/functions/sparksql/SparkQueryConfig.h" #include "velox/type/Type.h" +#include "velox/vector/SimpleVector.h" #include #include +#include #include #include +#include #include #include +#include #include +#include + +#include +#include + namespace { using namespace facebook::velox; @@ -721,8 +734,299 @@ struct ApproxDistinctAggregator : ReduceAggregator { std::int32_t precision_; }; +int64_t constantInt64Value(const core::TypedExprPtr& expr) { + VELOX_CHECK( + expr->isConstantKind(), "bloom_filter_agg extra arg must be constant"); + const auto* constant = expr->asUnchecked(); + VELOX_USER_CHECK( + !constant->isNull(), "bloom_filter_agg extra argument cannot be null"); + + auto readFromVector = [](const VectorPtr& vec) -> int64_t { + switch (vec->typeKind()) { + case TypeKind::TINYINT: + return vec->as>()->valueAt(0); + case TypeKind::SMALLINT: + return vec->as>()->valueAt(0); + case TypeKind::INTEGER: + return vec->as>()->valueAt(0); + case TypeKind::BIGINT: + return vec->as>()->valueAt(0); + default: + VELOX_FAIL( + "bloom_filter_agg extra argument must be integer, got {}", + vec->type()->toString()); + } + }; + + if (constant->hasValueVector()) { + return readFromVector(constant->valueVector()); + } + switch (constant->type()->kind()) { + case TypeKind::TINYINT: + return constant->value().value(); + case TypeKind::SMALLINT: + return constant->value().value(); + case TypeKind::INTEGER: + return constant->value().value(); + case TypeKind::BIGINT: + return constant->value().value(); + default: + VELOX_FAIL( + "bloom_filter_agg extra argument must be integer, got {}", + constant->type()->toString()); + } +} + +std::unique_ptr makeHostBytesStringColumn( + const std::string& data, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + const auto size = static_cast(data.size()); + cudf::size_type offsets[2] = {0, size}; + rmm::device_buffer offsetsDevice{2 * sizeof(cudf::size_type), stream, mr}; + CUDF_CUDA_TRY(cudaMemcpyAsync( + offsetsDevice.data(), + offsets, + 2 * sizeof(cudf::size_type), + cudaMemcpyHostToDevice, + stream.value())); + + rmm::device_buffer chars{data.size(), stream, mr}; + if (!data.empty()) { + CUDF_CUDA_TRY(cudaMemcpyAsync( + chars.data(), + data.data(), + data.size(), + cudaMemcpyHostToDevice, + stream.value())); + } + stream.synchronize(); + + auto offsetsColumn = std::make_unique( + cudf::data_type{cudf::type_id::INT32}, + 2, + std::move(offsetsDevice), + rmm::device_buffer{}, + 0); + return cudf::make_strings_column( + 1, std::move(offsetsColumn), std::move(chars), 0, rmm::device_buffer{}); +} + +std::unique_ptr makeSerializedBloomColumn( + const BloomFilter<>& bloom, + bool initialized, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + if (!initialized) { + cudf::string_scalar nullScalar("", false, stream, mr); + return cudf::make_column_from_scalar(nullScalar, 1, stream, mr); + } + std::string data; + data.resize(bloom.serializedSize()); + bloom.serialize(data.data()); + return makeHostBytesStringColumn(data, stream, mr); +} + +struct BloomFilterAggAggregator : ReduceAggregator { + BloomFilterAggAggregator( + core::AggregationNode::Step step, + uint32_t inputIndex, + VectorPtr constant, + const TypePtr& resultType, + std::vector extraInputs, + const core::QueryConfig* queryConfig) + : ReduceAggregator(step, inputIndex, std::move(constant), resultType), + extraInputs_(std::move(extraInputs)) { + initCapacity(queryConfig); + } + + std::unique_ptr doReduce( + cudf::table_view const& input, + TypePtr const& /* outputType */, + vector_size_t inputRowCount, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) override { + if (exec::isRawInput(step)) { + return reduceRaw(input, inputRowCount, stream, mr); + } + return reduceMerge(input, stream, mr); + } + + private: + void initCapacity(const core::QueryConfig* queryConfig) { + int64_t defaultExpectedNumItems = 1'000'000; + int64_t defaultNumBits = 8'388'608; + int64_t maxNumBits = 67'108'864; + int64_t maxNumItems = 4'000'000; + if (queryConfig != nullptr) { + functions::sparksql::SparkQueryConfig spark{*queryConfig}; + defaultExpectedNumItems = spark.bloomFilterExpectedNumItems(); + defaultNumBits = spark.bloomFilterNumBits(); + maxNumBits = spark.bloomFilterMaxNumBits(); + maxNumItems = spark.bloomFilterMaxNumItems(); + } + + int64_t estimatedNumItems; + int64_t numBits; + if (extraInputs_.size() >= 2) { + estimatedNumItems = constantInt64Value(extraInputs_[0]); + numBits = constantInt64Value(extraInputs_[1]); + } else if (extraInputs_.size() == 1) { + estimatedNumItems = constantInt64Value(extraInputs_[0]); + numBits = BloomFilter<>::optimalNumOfBits(estimatedNumItems, maxNumItems); + } else { + estimatedNumItems = defaultExpectedNumItems; + numBits = defaultNumBits; + } + VELOX_USER_CHECK_GT( + estimatedNumItems, 0, "estimatedNumItems must be positive"); + VELOX_USER_CHECK_GT(numBits, 0, "numBits must be positive"); + capacity_ = static_cast(std::min(numBits, maxNumBits) / 16); + } + + std::unique_ptr reduceRaw( + cudf::table_view const& input, + vector_size_t inputRowCount, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + BloomFilter<> bloom; + bool initialized = false; + + auto insertValue = [&](int64_t value) { + if (!initialized) { + bloom.reset(capacity_); + initialized = true; + } + bloom.insert(folly::hasher()(value)); + }; + + if (constant != nullptr) { + VELOX_USER_CHECK( + !constant->isNullAt(0), + "First argument of bloom_filter_agg cannot be null"); + VELOX_CHECK_EQ(constant->typeKind(), TypeKind::BIGINT); + if (inputRowCount > 0) { + insertValue(constant->as>()->valueAt(0)); + } + return makeSerializedBloomColumn(bloom, initialized, stream, mr); + } + + VELOX_CHECK_GT(input.num_columns(), inputIndex); + auto inputCol = input.column(inputIndex); + const auto numRows = inputCol.size(); + if (numRows == 0) { + return makeSerializedBloomColumn(bloom, false, stream, mr); + } + + VELOX_CHECK( + inputCol.type().id() == cudf::type_id::INT64 || + inputCol.type().id() == cudf::type_id::UINT64, + "bloom_filter_agg raw input must be BIGINT"); + + std::vector hostValues(static_cast(numRows)); + CUDF_CUDA_TRY(cudaMemcpyAsync( + hostValues.data(), + inputCol.data(), + static_cast(numRows) * sizeof(int64_t), + cudaMemcpyDeviceToHost, + stream.value())); + + std::vector hostNulls; + if (inputCol.nullable() && inputCol.null_mask() != nullptr) { + const auto bytes = cudf::bitmask_allocation_size_bytes(numRows); + hostNulls.resize(bytes / sizeof(cudf::bitmask_type)); + CUDF_CUDA_TRY(cudaMemcpyAsync( + hostNulls.data(), + inputCol.null_mask(), + bytes, + cudaMemcpyDeviceToHost, + stream.value())); + } + stream.synchronize(); + + for (cudf::size_type i = 0; i < numRows; ++i) { + if (!hostNulls.empty() && !cudf::bit_is_set(hostNulls.data(), i)) { + VELOX_USER_FAIL("First argument of bloom_filter_agg cannot be null"); + } + insertValue(hostValues[static_cast(i)]); + } + return makeSerializedBloomColumn(bloom, initialized, stream, mr); + } + + std::unique_ptr reduceMerge( + cudf::table_view const& input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + VELOX_CHECK_GT(input.num_columns(), inputIndex); + auto sketchColumn = input.column(inputIndex); + BloomFilter<> bloom; + bool initialized = false; + if (sketchColumn.size() == 0) { + return makeSerializedBloomColumn(bloom, false, stream, mr); + } + + auto stringsCol = cudf::strings_column_view(sketchColumn); + auto offsetsCol = stringsCol.offsets(); + auto charsPtr = stringsCol.chars_begin(stream); + const auto numOffsets = sketchColumn.size() + 1; + std::vector hostOffsets(static_cast(numOffsets)); + CUDF_CUDA_TRY(cudaMemcpyAsync( + hostOffsets.data(), + offsetsCol.begin(), + static_cast(numOffsets) * sizeof(cudf::size_type), + cudaMemcpyDeviceToHost, + stream.value())); + + std::vector hostNulls; + if (sketchColumn.nullable() && sketchColumn.null_mask() != nullptr) { + const auto bytes = + cudf::bitmask_allocation_size_bytes(sketchColumn.size()); + hostNulls.resize(bytes / sizeof(cudf::bitmask_type)); + CUDF_CUDA_TRY(cudaMemcpyAsync( + hostNulls.data(), + sketchColumn.null_mask(), + bytes, + cudaMemcpyDeviceToHost, + stream.value())); + } + stream.synchronize(); + + const auto charsSize = + static_cast(hostOffsets.back() - hostOffsets.front()); + std::vector hostChars(charsSize); + if (charsSize > 0) { + CUDF_CUDA_TRY(cudaMemcpyAsync( + hostChars.data(), + charsPtr + hostOffsets.front(), + charsSize, + cudaMemcpyDeviceToHost, + stream.value())); + stream.synchronize(); + } + + const auto base = hostOffsets.front(); + for (cudf::size_type i = 0; i < sketchColumn.size(); ++i) { + if (!hostNulls.empty() && !cudf::bit_is_set(hostNulls.data(), i)) { + continue; + } + const auto start = hostOffsets[static_cast(i)] - base; + const auto end = hostOffsets[static_cast(i) + 1] - base; + if (end <= start) { + continue; + } + bloom.merge(hostChars.data() + start); + initialized = true; + } + return makeSerializedBloomColumn(bloom, initialized, stream, mr); + } + + std::vector extraInputs_; + int32_t capacity_{0}; +}; + std::unique_ptr createReduceAggregator( - const ResolvedAggregateInfo& p) { + const ResolvedAggregateInfo& p, + const core::QueryConfig* queryConfig) { auto const& kind = p.kind; auto prefix = cudf_velox::CudfConfig::getInstance().functionNamePrefix; if (kind.rfind(prefix + "sum", 0) == 0) { @@ -752,6 +1056,14 @@ std::unique_ptr createReduceAggregator( } else if (kind.rfind(prefix + "approx_distinct", 0) == 0) { return std::make_unique( p.companionStep, p.inputIndex, p.constant, p.resultType); + } else if (kind.rfind(prefix + "bloom_filter_agg", 0) == 0) { + return std::make_unique( + p.companionStep, + p.inputIndex, + p.constant, + p.resultType, + p.extraInputs, + queryConfig); } else { VELOX_NYI("Reduce aggregation not yet supported, kind: {}", kind); } @@ -765,14 +1077,15 @@ std::vector> toReduceAggregators( core::AggregationNode const& aggregationNode, core::AggregationNode::Step step, TypePtr const& outputType, - std::vector const& constants) { + std::vector const& constants, + const core::QueryConfig* queryConfig) { auto params = resolveAggregateInfos(aggregationNode, step, outputType, constants); std::vector> aggregators; aggregators.reserve(params.size()); for (const auto& p : params) { - aggregators.push_back(createReduceAggregator(p)); + aggregators.push_back(createReduceAggregator(p, queryConfig)); } return aggregators; } @@ -871,7 +1184,8 @@ void CudfReduce::initialize() { *aggregationNode_, aggregationNode_->step(), outputType_, - aggregationInput.constants); + aggregationInput.constants, + &operatorCtx_->driverCtx()->queryConfig()); aggregationNode_.reset(); } diff --git a/velox/experimental/cudf/exec/CudfReduce.h b/velox/experimental/cudf/exec/CudfReduce.h index 9d569fbdb25..9fd61ea63fa 100644 --- a/velox/experimental/cudf/exec/CudfReduce.h +++ b/velox/experimental/cudf/exec/CudfReduce.h @@ -18,6 +18,8 @@ #include "velox/experimental/cudf/exec/CudfAggregation.h" #include "velox/experimental/cudf/exec/CudfOperator.h" +#include "velox/core/QueryConfig.h" + namespace facebook::velox::cudf_velox { struct ReduceAggregator { @@ -51,7 +53,8 @@ std::vector> toReduceAggregators( core::AggregationNode const& aggregationNode, core::AggregationNode::Step step, TypePtr const& outputType, - std::vector const& constants); + std::vector const& constants, + const core::QueryConfig* queryConfig = nullptr); bool canReduceBeEvaluatedByCudf( const core::AggregationNode& aggregationNode, diff --git a/velox/experimental/cudf/exec/SparkAggregateFunctions.cpp b/velox/experimental/cudf/exec/SparkAggregateFunctions.cpp index 619e798b412..5798cb178f9 100644 --- a/velox/experimental/cudf/exec/SparkAggregateFunctions.cpp +++ b/velox/experimental/cudf/exec/SparkAggregateFunctions.cpp @@ -125,6 +125,48 @@ void registerSparkAggregateFunctions(const std::string& prefix) { prefix + "collect_list", core::AggregationNode::Step::kFinal, collectionMergeSignature); + + // Spark runtime bloom filters are uncorrelated scalar subqueries, so the + // native path is a global CudfReduce. Do not advertise these signatures in + // the groupby registry: grouped bloom_filter_agg is not implemented. + auto bloomRawOneArg = FunctionSignatureBuilder() + .returnType("varbinary") + .argumentType("bigint") + .build(); + auto bloomRawTwoArg = FunctionSignatureBuilder() + .returnType("varbinary") + .argumentType("bigint") + .constantArgumentType("bigint") + .build(); + auto bloomRawThreeArg = FunctionSignatureBuilder() + .returnType("varbinary") + .argumentType("bigint") + .constantArgumentType("bigint") + .constantArgumentType("bigint") + .build(); + auto bloomMerge = FunctionSignatureBuilder() + .returnType("varbinary") + .argumentType("varbinary") + .build(); + for (auto step : { + core::AggregationNode::Step::kPartial, + core::AggregationNode::Step::kSingle, + }) { + appendReduceAggregationFunctionForStep( + prefix + "bloom_filter_agg", step, bloomRawOneArg); + appendReduceAggregationFunctionForStep( + prefix + "bloom_filter_agg", step, bloomRawTwoArg); + appendReduceAggregationFunctionForStep( + prefix + "bloom_filter_agg", step, bloomRawThreeArg); + } + appendReduceAggregationFunctionForStep( + prefix + "bloom_filter_agg", + core::AggregationNode::Step::kIntermediate, + bloomMerge); + appendReduceAggregationFunctionForStep( + prefix + "bloom_filter_agg", + core::AggregationNode::Step::kFinal, + bloomMerge); } } // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/expression/CMakeLists.txt b/velox/experimental/cudf/expression/CMakeLists.txt index 4047806c9e4..76a3b7e0be3 100644 --- a/velox/experimental/cudf/expression/CMakeLists.txt +++ b/velox/experimental/cudf/expression/CMakeLists.txt @@ -37,7 +37,9 @@ add_library( SparkFunctions.cpp sparksql/DateAddFunction.cpp sparksql/HashFunction.cpp + sparksql/MightContainFunction.cpp sparksql/SubStringFunction.cpp + sparksql/XxHash64Function.cpp SubfieldFiltersToAst.cpp ) diff --git a/velox/experimental/cudf/expression/SparkFunctions.cpp b/velox/experimental/cudf/expression/SparkFunctions.cpp index 9d8b01f02ee..cd527d773ed 100644 --- a/velox/experimental/cudf/expression/SparkFunctions.cpp +++ b/velox/experimental/cudf/expression/SparkFunctions.cpp @@ -20,7 +20,9 @@ #include "velox/experimental/cudf/expression/SparkFunctions.h" #include "velox/experimental/cudf/expression/sparksql/DateAddFunction.h" #include "velox/experimental/cudf/expression/sparksql/HashFunction.h" +#include "velox/experimental/cudf/expression/sparksql/MightContainFunction.h" #include "velox/experimental/cudf/expression/sparksql/SubStringFunction.h" +#include "velox/experimental/cudf/expression/sparksql/XxHash64Function.h" #include "velox/expression/FunctionSignature.h" #include "velox/functions/sparksql/SparkQueryConfig.h" @@ -177,6 +179,52 @@ void registerSparkFunctions(const std::string& prefix) { true, sparksql::HashFunction::canEvaluate); + registerCudfFunction( + prefix + "xxhash64_with_seed", + [](const std::string&, + const core::TypedExprPtr& expr, + memory::MemoryPool* pool) { + return std::make_shared(expr, pool); + }, + {FunctionSignatureBuilder() + .returnType("bigint") + .constantArgumentType("bigint") + .argumentType("any") + .variableArity("any") + .build()}, + true, + sparksql::XxHash64Function::canEvaluate); + + registerCudfFunction( + prefix + "might_contain", + [](const std::string&, + const core::TypedExprPtr& expr, + memory::MemoryPool* pool) { + return std::make_shared(expr, pool); + }, + {FunctionSignatureBuilder() + .returnType("boolean") + .constantArgumentType("varbinary") + .argumentType("tinyint") + .build(), + FunctionSignatureBuilder() + .returnType("boolean") + .constantArgumentType("varbinary") + .argumentType("smallint") + .build(), + FunctionSignatureBuilder() + .returnType("boolean") + .constantArgumentType("varbinary") + .argumentType("integer") + .build(), + FunctionSignatureBuilder() + .returnType("boolean") + .constantArgumentType("varbinary") + .argumentType("bigint") + .build()}, + true, + sparksql::MightContainFunction::canEvaluate); + registerCudfFunction( prefix + "date_add", [](const std::string&, diff --git a/velox/experimental/cudf/expression/sparksql/MightContainFunction.cpp b/velox/experimental/cudf/expression/sparksql/MightContainFunction.cpp new file mode 100644 index 00000000000..d5f7fd3c443 --- /dev/null +++ b/velox/experimental/cudf/expression/sparksql/MightContainFunction.cpp @@ -0,0 +1,200 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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 "velox/experimental/cudf/expression/sparksql/MightContainFunction.h" + +#include "velox/common/base/BloomFilter.h" +#include "velox/common/memory/Memory.h" +#include "velox/core/Expressions.h" +#include "velox/vector/BaseVector.h" +#include "velox/vector/SimpleVector.h" + +#include +#include +#include +#include + +#include + +namespace facebook::velox::cudf_velox::sparksql { +namespace { + +template +void copyKeysAsInt64( + const cudf::column_view& inputView, + std::vector& hostKeys, + rmm::cuda_stream_view stream) { + std::vector tmp(static_cast(inputView.size())); + if (inputView.size() > 0) { + CUDF_CUDA_TRY(cudaMemcpyAsync( + tmp.data(), + inputView.data(), + static_cast(inputView.size()) * sizeof(T), + cudaMemcpyDeviceToHost, + stream.value())); + stream.synchronize(); + } + for (cudf::size_type i = 0; i < inputView.size(); ++i) { + hostKeys[static_cast(i)] = static_cast(tmp[i]); + } +} + +void copyKeysToHost( + const cudf::column_view& inputView, + std::vector& hostKeys, + rmm::cuda_stream_view stream) { + switch (inputView.type().id()) { + case cudf::type_id::INT64: + if (inputView.size() > 0) { + CUDF_CUDA_TRY(cudaMemcpyAsync( + hostKeys.data(), + inputView.data(), + static_cast(inputView.size()) * sizeof(int64_t), + cudaMemcpyDeviceToHost, + stream.value())); + } + break; + case cudf::type_id::INT32: + copyKeysAsInt64(inputView, hostKeys, stream); + break; + case cudf::type_id::INT16: + copyKeysAsInt64(inputView, hostKeys, stream); + break; + case cudf::type_id::INT8: + copyKeysAsInt64(inputView, hostKeys, stream); + break; + case cudf::type_id::UINT64: + copyKeysAsInt64(inputView, hostKeys, stream); + break; + case cudf::type_id::UINT32: + copyKeysAsInt64(inputView, hostKeys, stream); + break; + case cudf::type_id::UINT16: + copyKeysAsInt64(inputView, hostKeys, stream); + break; + case cudf::type_id::UINT8: + copyKeysAsInt64(inputView, hostKeys, stream); + break; + default: + VELOX_FAIL( + "might_contain hash input must be an integer type; saw cudf type_id={}", + static_cast(inputView.type().id())); + } +} + +} // namespace + +bool MightContainFunction::canEvaluate(const core::TypedExprPtr& expr) { + return expr->inputs().size() == 2 && expr->inputs()[0]->isConstantKind(); +} + +MightContainFunction::MightContainFunction( + const core::TypedExprPtr& expr, + memory::MemoryPool* pool) { + VELOX_CHECK_EQ( + expr->inputs().size(), 2, "might_contain expects exactly 2 inputs"); + VELOX_CHECK( + expr->inputs()[0]->isConstantKind(), + "might_contain bloom filter must be a constant"); + const auto* bloomExpr = + expr->inputs()[0]->asUnchecked(); + const auto bloomValue = bloomExpr->hasValueVector() + ? bloomExpr->valueVector() + : bloomExpr->toConstantVector(pool); + if (bloomValue->isNullAt(0)) { + bloomIsNull_ = true; + return; + } + auto serialized = bloomValue->as>()->valueAt(0); + serialized_.assign(serialized.data(), serialized.size()); +} + +ColumnOrView MightContainFunction::eval( + std::vector& inputColumns, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const { + VELOX_CHECK_EQ( + inputColumns.size(), + 1, + "might_contain receives 1 column input; bloom filter is literal"); + auto inputView = asView(inputColumns[0]); + const auto numRows = inputView.size(); + + if (bloomIsNull_) { + return cudf::make_numeric_column( + cudf::data_type{cudf::type_id::BOOL8}, + numRows, + cudf::mask_state::ALL_NULL, + stream, + mr); + } + + std::vector hostKeys(static_cast(numRows)); + if (numRows > 0) { + copyKeysToHost(inputView, hostKeys, stream); + } + + std::vector hostNulls; + const cudf::bitmask_type* deviceNulls = inputView.null_mask(); + if (deviceNulls != nullptr) { + const auto bytes = cudf::bitmask_allocation_size_bytes(numRows); + hostNulls.resize(bytes / sizeof(cudf::bitmask_type)); + CUDF_CUDA_TRY(cudaMemcpyAsync( + hostNulls.data(), + deviceNulls, + bytes, + cudaMemcpyDeviceToHost, + stream.value())); + } + stream.synchronize(); + + std::vector hostResult(static_cast(numRows), 0); + for (cudf::size_type i = 0; i < numRows; ++i) { + if (!hostNulls.empty() && !cudf::bit_is_set(hostNulls.data(), i)) { + continue; + } + const uint64_t hashed = folly::hasher()(hostKeys[i]); + hostResult[static_cast(i)] = + BloomFilter<>::mayContain(serialized_.c_str(), hashed) ? 1u : 0u; + } + + rmm::device_buffer data( + static_cast(numRows) * sizeof(uint8_t), stream, mr); + if (numRows > 0) { + CUDF_CUDA_TRY(cudaMemcpyAsync( + data.data(), + hostResult.data(), + static_cast(numRows) * sizeof(uint8_t), + cudaMemcpyHostToDevice, + stream.value())); + } + + rmm::device_buffer nullMask{}; + cudf::size_type nullCount = 0; + if (deviceNulls != nullptr) { + nullMask = rmm::device_buffer( + deviceNulls, cudf::bitmask_allocation_size_bytes(numRows), stream, mr); + nullCount = inputView.null_count(); + } + + return std::make_unique( + cudf::data_type{cudf::type_id::BOOL8}, + numRows, + std::move(data), + std::move(nullMask), + nullCount); +} + +} // namespace facebook::velox::cudf_velox::sparksql diff --git a/velox/experimental/cudf/expression/sparksql/MightContainFunction.h b/velox/experimental/cudf/expression/sparksql/MightContainFunction.h new file mode 100644 index 00000000000..e1d36b617f2 --- /dev/null +++ b/velox/experimental/cudf/expression/sparksql/MightContainFunction.h @@ -0,0 +1,43 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ +#pragma once + +#include "velox/experimental/cudf/expression/ExpressionEvaluator.h" + +namespace facebook::velox::cudf_velox::sparksql { + +/// Spark might_contain(serializedBloom, hashKey). The serialized Bloom filter +/// must be a constant (after runtime-bloom subquery materialization). Probe +/// keys are copied to the host and tested with Velox BloomFilter::mayContain. +class MightContainFunction : public CudfFunction { + public: + static bool canEvaluate(const core::TypedExprPtr& expr); + + MightContainFunction( + const core::TypedExprPtr& expr, + memory::MemoryPool* pool); + + ColumnOrView eval( + std::vector& inputColumns, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override; + + private: + std::string serialized_; + bool bloomIsNull_{false}; +}; + +} // namespace facebook::velox::cudf_velox::sparksql diff --git a/velox/experimental/cudf/expression/sparksql/XxHash64Function.cpp b/velox/experimental/cudf/expression/sparksql/XxHash64Function.cpp new file mode 100644 index 00000000000..b5c95056831 --- /dev/null +++ b/velox/experimental/cudf/expression/sparksql/XxHash64Function.cpp @@ -0,0 +1,90 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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 "velox/experimental/cudf/CudfConfig.h" +#include "velox/experimental/cudf/expression/sparksql/XxHash64Function.h" + +#include "velox/common/memory/Memory.h" +#include "velox/core/Expressions.h" +#include "velox/vector/BaseVector.h" +#include "velox/vector/SimpleVector.h" + +#include +#include +#include +#include + +namespace facebook::velox::cudf_velox::sparksql { +namespace { + +cudf::table_view convertToTableView(std::vector& inputColumns) { + std::vector columns; + columns.reserve(inputColumns.size()); + for (auto& col : inputColumns) { + columns.push_back(asView(col)); + } + return cudf::table_view(columns); +} + +} // namespace + +bool XxHash64Function::canEvaluate(const core::TypedExprPtr& expr) { + if (expr->inputs().size() < 2) { + return false; + } + + // Multi-column xxhash64_with_seed runs on GPU via cuDF xxhash_64, which + // combines columns with a constant seed instead of Spark's iterative + // seed-chaining (rapidsai/cudf#21720). When CPU fallback is allowed, reject + // this shape so it stays on CPU; when fallback is disabled, preserve the + // existing forced-GPU behavior used by hash_with_seed. + const bool hasMultipleDataColumns = expr->inputs().size() > 2; + return !hasMultipleDataColumns || !CudfConfig::getInstance().allowCpuFallback; +} + +XxHash64Function::XxHash64Function( + const core::TypedExprPtr& expr, + memory::MemoryPool* pool) { + VELOX_CHECK_GE( + expr->inputs().size(), 2, "xxhash64 expects at least 2 inputs"); + VELOX_CHECK( + expr->inputs()[0]->isConstantKind(), "xxhash64 seed must be a constant"); + const auto* seedExpr = + expr->inputs()[0]->asUnchecked(); + VELOX_CHECK(!seedExpr->isNull(), "xxhash64 seed must be non-null"); + const auto vec = seedExpr->hasValueVector() + ? seedExpr->valueVector() + : seedExpr->toConstantVector(pool); + seedValue_ = vec->as>()->valueAt(0); +} + +ColumnOrView XxHash64Function::eval( + std::vector& inputColumns, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const { + VELOX_CHECK(!inputColumns.empty()); + auto inputTableView = convertToTableView(inputColumns); + auto hashes = cudf::hashing::xxhash_64( + inputTableView, static_cast(seedValue_), stream, mr); + // Spark hashes a null input to the seed and never returns null. + if (hashes->nullable() && hashes->null_count() > 0) { + cudf::numeric_scalar replacement( + static_cast(seedValue_), true, stream, mr); + hashes = cudf::replace_nulls(hashes->view(), replacement, stream, mr); + } + return hashes; +} + +} // namespace facebook::velox::cudf_velox::sparksql diff --git a/velox/experimental/cudf/expression/sparksql/XxHash64Function.h b/velox/experimental/cudf/expression/sparksql/XxHash64Function.h new file mode 100644 index 00000000000..60c03d014a4 --- /dev/null +++ b/velox/experimental/cudf/expression/sparksql/XxHash64Function.h @@ -0,0 +1,40 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ +#pragma once + +#include "velox/experimental/cudf/expression/ExpressionEvaluator.h" + +namespace facebook::velox::cudf_velox::sparksql { + +/// Spark xxhash64_with_seed(seed, ...). Computes xxHash64 over the remaining +/// arguments using the constant BIGINT seed. Null inputs hash to the seed so +/// the result is never null, matching Spark defaultNullBehavior(false). +class XxHash64Function : public CudfFunction { + public: + static bool canEvaluate(const core::TypedExprPtr& expr); + + XxHash64Function(const core::TypedExprPtr& expr, memory::MemoryPool* pool); + + ColumnOrView eval( + std::vector& inputColumns, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override; + + private: + int64_t seedValue_; +}; + +} // namespace facebook::velox::cudf_velox::sparksql diff --git a/velox/experimental/cudf/tests/ExpressionEvaluatorSelectionTest.cpp b/velox/experimental/cudf/tests/ExpressionEvaluatorSelectionTest.cpp index 295a4f1aae3..2845cef2de3 100644 --- a/velox/experimental/cudf/tests/ExpressionEvaluatorSelectionTest.cpp +++ b/velox/experimental/cudf/tests/ExpressionEvaluatorSelectionTest.cpp @@ -32,6 +32,7 @@ #include "velox/functions/sparksql/SparkQueryConfig.h" #include "velox/functions/sparksql/registration/Register.h" #include "velox/type/Type.h" +#include "velox/type/Variant.h" #include #include @@ -644,6 +645,40 @@ TEST_F(CudfExpressionSelectionTest, signatureVarargsHashWithSeed) { } } +TEST_F(CudfExpressionSelectionTest, signatureVarargsXxHash64WithSeed) { + facebook::velox::functions::sparksql::registerFunctions(); + + auto multiCol = optimizeTypedExpr( + "xxhash64_with_seed(42, a, b)", + rowType_, + queryCtx_.get(), + execCtx_.get()); + ASSERT_FALSE(canExprRunOnGpu(multiCol, queryCtx_.get(), pool_.get())); + + auto singleCol = optimizeTypedExpr( + "xxhash64_with_seed(42, a)", rowType_, queryCtx_.get(), execCtx_.get()); + ASSERT_TRUE(canExprRunOnGpu(singleCol, queryCtx_.get(), pool_.get())); +} + +TEST_F(CudfExpressionSelectionTest, signatureMightContain) { + facebook::velox::functions::sparksql::registerFunctions(); + + auto bloom = std::make_shared( + VARBINARY(), variant::binary(std::string(16, '\0'))); + auto key = std::make_shared(BIGINT(), "a"); + auto expr = std::make_shared( + BOOLEAN(), std::vector{bloom, key}, "might_contain"); + ASSERT_TRUE(canExprRunOnGpu(expr, queryCtx_.get(), pool_.get())); + + auto nonConstBloom = + std::make_shared(VARBINARY(), "bloom"); + auto nonConstExpr = std::make_shared( + BOOLEAN(), + std::vector{nonConstBloom, key}, + "might_contain"); + ASSERT_FALSE(canExprRunOnGpu(nonConstExpr, queryCtx_.get(), pool_.get())); +} + TEST_F(CudfExpressionSelectionTest, signatureTypeVariableCoalesce) { // OK: same type BIGINT auto ok1 = optimizeTypedExpr( diff --git a/velox/experimental/cudf/tests/sparksql/AggregationTest.cpp b/velox/experimental/cudf/tests/sparksql/AggregationTest.cpp index 935d605ce11..e2666ec5f44 100644 --- a/velox/experimental/cudf/tests/sparksql/AggregationTest.cpp +++ b/velox/experimental/cudf/tests/sparksql/AggregationTest.cpp @@ -14,14 +14,23 @@ * limitations under the License. */ +#include "velox/experimental/cudf/CudfConfig.h" #include "velox/experimental/cudf/exec/AggregationRegistry.h" #include "velox/experimental/cudf/exec/SparkAggregateFunctions.h" #include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/experimental/cudf/expression/SparkFunctions.h" +#include "velox/common/base/BloomFilter.h" #include "velox/common/base/tests/GTestUtils.h" +#include "velox/exec/tests/utils/AssertQueryBuilder.h" #include "velox/exec/tests/utils/PlanBuilder.h" #include "velox/functions/lib/aggregates/tests/utils/AggregationTestBase.h" #include "velox/functions/sparksql/aggregates/Register.h" +#include "velox/functions/sparksql/registration/Register.h" +#include "velox/vector/SimpleVector.h" + +#include +#include using namespace facebook::velox::exec::test; using namespace facebook::velox::functions::aggregate::test; @@ -32,12 +41,14 @@ class AggregationTest : public AggregationTestBase { protected: void SetUp() override { AggregationTestBase::SetUp(); + functions::sparksql::registerFunctions(""); functions::aggregate::sparksql::registerAggregateFunctions(""); filesystems::registerLocalFileSystem(); // After register supports function prefix, we could register the function // with spark_ to align with sparksql AverageAggregationTest, the // function name overwrite may not work well in some condition. cudf_velox::registerCudf(); + cudf_velox::registerSparkFunctions(""); cudf_velox::registerSparkAggregateFunctions(""); } @@ -45,6 +56,18 @@ class AggregationTest : public AggregationTestBase { cudf_velox::unregisterCudf(); cudf_velox::unregisterAggregateFunctions(); } + + VectorPtr getSerializedBloomFilter(int32_t capacity) { + BloomFilter<> bloomFilter; + bloomFilter.reset(capacity); + for (auto i = 0; i < 9; ++i) { + bloomFilter.insert(folly::hasher()(i)); + } + std::string data; + data.resize(bloomFilter.serializedSize()); + bloomFilter.serialize(data.data()); + return makeConstant(StringView(data), 1, VARBINARY()); + } }; TEST_F(AggregationTest, sumReal) { @@ -127,4 +150,63 @@ TEST_F(AggregationTest, groupedCollectListOfRow) { assertQuery(partialFinal, expected); } +TEST_F(AggregationTest, bloomFilterAgg) { + auto& cudfConfig = cudf_velox::CudfConfig::getInstance(); + const auto previousFallback = cudfConfig.allowCpuFallback; + auto restoreFallback = + folly::makeGuard([&] { cudfConfig.allowCpuFallback = previousFallback; }); + cudfConfig.allowCpuFallback = false; + auto vectors = makeRowVector({makeFlatVector( + 100, [](vector_size_t row) { return row % 9; })}); + auto expected = makeRowVector({getSerializedBloomFilter(11)}); + auto plan = PlanBuilder() + .values({vectors}) + .partialAggregation({}, {"bloom_filter_agg(c0, 5, 64)"}) + .finalAggregation() + .planNode(); + assertQuery(plan, expected); + + auto single = PlanBuilder() + .values({vectors}) + .singleAggregation({}, {"bloom_filter_agg(c0, 5, 64)"}) + .planNode(); + assertQuery(single, expected); +} + +TEST_F(AggregationTest, bloomFilterAggFromXxHash64) { + auto& cudfConfig = cudf_velox::CudfConfig::getInstance(); + const auto previousFallback = cudfConfig.allowCpuFallback; + auto restoreFallback = + folly::makeGuard([&] { cudfConfig.allowCpuFallback = previousFallback; }); + cudfConfig.allowCpuFallback = false; + auto keys = makeFlatVector({1, 2, 3, 4, 5, 1, 2}); + auto input = makeRowVector({keys}); + auto plan = PlanBuilder() + .values({input}) + .project({"xxhash64_with_seed(cast(42 as bigint), c0) AS h"}) + .singleAggregation({}, {"bloom_filter_agg(h, 5, 64)"}) + .planNode(); + auto gpu = AssertQueryBuilder(plan).copyResults(pool()); + ASSERT_EQ(gpu->size(), 1); + ASSERT_FALSE(gpu->childAt(0)->isNullAt(0)); + + auto containPlan = + PlanBuilder() + .values({input}) + .project({"xxhash64_with_seed(cast(42 as bigint), c0) AS h"}) + .planNode(); + auto hashes = AssertQueryBuilder(containPlan).copyResults(pool()); + auto serialized = gpu->childAt(0)->as>()->valueAt(0); + std::string data(serialized.data(), serialized.size()); + for (auto i = 0; i < hashes->size(); ++i) { + auto hash = hashes->childAt(0)->as>()->valueAt(i); + ASSERT_TRUE( + BloomFilter<>::mayContain( + data.c_str(), folly::hasher()(hash))); + } + ASSERT_FALSE( + BloomFilter<>::mayContain( + data.c_str(), folly::hasher()(int64_t{123456789}))); +} + } // namespace facebook::velox::exec::sparksql::test diff --git a/velox/experimental/cudf/tests/sparksql/FilterProjectTest.cpp b/velox/experimental/cudf/tests/sparksql/FilterProjectTest.cpp index ea23dabe418..01a5ecfea03 100644 --- a/velox/experimental/cudf/tests/sparksql/FilterProjectTest.cpp +++ b/velox/experimental/cudf/tests/sparksql/FilterProjectTest.cpp @@ -21,15 +21,21 @@ #include "velox/experimental/cudf/tests/CudfFunctionBaseTest.h" #include "velox/experimental/cudf/tests/utils/ExpressionTestUtil.h" +#include "velox/common/base/BloomFilter.h" #include "velox/common/base/tests/GTestUtils.h" +#include "velox/core/Expressions.h" #include "velox/dwio/common/tests/utils/BatchMaker.h" #include "velox/exec/tests/utils/AssertQueryBuilder.h" #include "velox/exec/tests/utils/OperatorTestBase.h" #include "velox/exec/tests/utils/PlanBuilder.h" +#include "velox/expression/Expr.h" #include "velox/functions/prestosql/ArrayConstructor.h" #include "velox/functions/sparksql/registration/Register.h" #include "velox/parse/TypeResolver.h" #include "velox/type/TimestampConversion.h" +#include "velox/type/Variant.h" + +#include using namespace facebook::velox::exec::test; using namespace facebook::velox; @@ -155,6 +161,48 @@ TEST_F(CudfFilterProjectTest, hashWithSeed) { facebook::velox::test::assertEqualVectors(expected, hashResults); } +TEST_F(CudfFilterProjectTest, xxhash64WithSeed) { + auto input = makeRowVector({ + makeNullableFlatVector( + {INT64_MAX, INT64_MIN, 0, 42, std::nullopt}), + }); + assertExpressionMatchesCpu( + "xxhash64_with_seed(cast(42 as bigint), c0)", input, input->rowType()); +} + +TEST_F(CudfFilterProjectTest, mightContain) { + constexpr int32_t kSize = 10; + BloomFilter<> bloomFilter; + bloomFilter.reset(kSize); + for (auto i = 0; i < kSize; ++i) { + bloomFilter.insert(folly::hasher()(i)); + } + std::string serialized; + serialized.resize(bloomFilter.serializedSize()); + bloomFilter.serialize(serialized.data()); + + auto keys = + makeNullableFlatVector({0, 1, 9, 10, 123451, std::nullopt}); + auto input = makeRowVector({keys}); + auto bloomExpr = std::make_shared( + VARBINARY(), variant::binary(serialized)); + auto keyExpr = std::make_shared(BIGINT(), "c0"); + auto expr = std::make_shared( + BOOLEAN(), + std::vector{bloomExpr, keyExpr}, + "might_contain"); + + auto gpu = evaluate(expr, input); + + auto selected = SelectivityVector(keys->size()); + exec::ExprSet cpuExpr({expr}, &execCtx_); + auto data = makeRowVector({keys}); + exec::EvalCtx evalCtx(&execCtx_, &cpuExpr, data.get()); + std::vector cpuResults(1); + cpuExpr.eval(selected, evalCtx, cpuResults); + facebook::velox::test::assertEqualVectors(cpuResults[0], gpu); +} + TEST_F(CudfFilterProjectTest, sparkExpressionParity) { auto input = makeRowVector({ makeNullableFlatVector({1, std::nullopt, -7, 42, 0, 9}), From 668f0bc6ca37d1e8456c8389b347735040a1fb37 Mon Sep 17 00:00:00 2001 From: Alfred Xu Date: Tue, 25 Aug 2026 11:30:11 +0900 Subject: [PATCH 2/2] fix(cudf): reject multi-column Spark xxhash64 - Reject multi-column xxhash64_with_seed regardless of CPU fallback because cuDF does not implement Spark iterative seed chaining. - Preserve the single-column GPU path required by the Q20 runtime Bloom-filter workflow. - Cover supported single-column and rejected multi-column expression selection. Fixes #101 --- .../expression/sparksql/XxHash64Function.cpp | 12 +++++------- .../ExpressionEvaluatorSelectionTest.cpp | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/velox/experimental/cudf/expression/sparksql/XxHash64Function.cpp b/velox/experimental/cudf/expression/sparksql/XxHash64Function.cpp index b5c95056831..5a2aa9e7085 100644 --- a/velox/experimental/cudf/expression/sparksql/XxHash64Function.cpp +++ b/velox/experimental/cudf/expression/sparksql/XxHash64Function.cpp @@ -13,7 +13,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#include "velox/experimental/cudf/CudfConfig.h" #include "velox/experimental/cudf/expression/sparksql/XxHash64Function.h" #include "velox/common/memory/Memory.h" @@ -45,13 +44,12 @@ bool XxHash64Function::canEvaluate(const core::TypedExprPtr& expr) { return false; } - // Multi-column xxhash64_with_seed runs on GPU via cuDF xxhash_64, which - // combines columns with a constant seed instead of Spark's iterative - // seed-chaining (rapidsai/cudf#21720). When CPU fallback is allowed, reject - // this shape so it stays on CPU; when fallback is disabled, preserve the - // existing forced-GPU behavior used by hash_with_seed. + // Multi-column xxhash64_with_seed cannot run on GPU via cuDF xxhash_64: + // cuDF combines columns with a constant seed instead of Spark's iterative + // seed-chaining (rapidsai/cudf#21720). Always reject this shape so callers + // either use CPU fallback or fail closed instead of returning wrong hashes. const bool hasMultipleDataColumns = expr->inputs().size() > 2; - return !hasMultipleDataColumns || !CudfConfig::getInstance().allowCpuFallback; + return !hasMultipleDataColumns; } XxHash64Function::XxHash64Function( diff --git a/velox/experimental/cudf/tests/ExpressionEvaluatorSelectionTest.cpp b/velox/experimental/cudf/tests/ExpressionEvaluatorSelectionTest.cpp index 2845cef2de3..c86ea5abef5 100644 --- a/velox/experimental/cudf/tests/ExpressionEvaluatorSelectionTest.cpp +++ b/velox/experimental/cudf/tests/ExpressionEvaluatorSelectionTest.cpp @@ -22,6 +22,7 @@ #include "velox/experimental/cudf/expression/JitExpression.h" #include "velox/experimental/cudf/expression/PrestoFunctions.h" #include "velox/experimental/cudf/expression/SparkFunctions.h" +#include "velox/experimental/cudf/expression/sparksql/XxHash64Function.h" #include "velox/experimental/cudf/tests/utils/ExpressionTestUtil.h" #include "velox/common/memory/Memory.h" @@ -142,6 +143,24 @@ TEST_F(CudfExpressionSelectionTest, sparkExpressionCoverage) { ASSERT_NE(createCudfExpression(arrayExpr, rowType_, pool_.get()), nullptr); } +TEST_F(CudfExpressionSelectionTest, sparkXxHash64RejectsMultipleColumns) { + auto singleColumn = optimizeTypedExpr( + "xxhash64_with_seed(cast(42 as bigint), a)", + rowType_, + queryCtx_.get(), + execCtx_.get(), + {.parseIntegerAsBigint = false, .functionPrefix = ""}); + EXPECT_TRUE(sparksql::XxHash64Function::canEvaluate(singleColumn)); + + auto multipleColumns = optimizeTypedExpr( + "xxhash64_with_seed(cast(42 as bigint), a, b)", + rowType_, + queryCtx_.get(), + execCtx_.get(), + {.parseIntegerAsBigint = false, .functionPrefix = ""}); + EXPECT_FALSE(sparksql::XxHash64Function::canEvaluate(multipleColumns)); +} + TEST_F(CudfExpressionSelectionTest, multiBranchSwitch) { auto expr = optimizeTypedExpr( "CASE WHEN a = 1 THEN name "