From 8ca46712bcadc15b8f94a6294b4080270c7b7936 Mon Sep 17 00:00:00 2001 From: rui-mo Date: Sun, 23 Apr 2023 12:48:54 +0000 Subject: [PATCH] add another --- velox/expression/ExprCompiler.cpp | 24 +++++++ velox/functions/FunctionRegistry.cpp | 3 +- .../prestosql/RowFunctionWithNull.cpp | 67 +++++++++++++++++++ .../GeneralFunctionsRegistration.cpp | 1 + 4 files changed, 94 insertions(+), 1 deletion(-) create mode 100644 velox/functions/prestosql/RowFunctionWithNull.cpp diff --git a/velox/expression/ExprCompiler.cpp b/velox/expression/ExprCompiler.cpp index 74b1f47d822ef..1ec8eaac1f7a8 100644 --- a/velox/expression/ExprCompiler.cpp +++ b/velox/expression/ExprCompiler.cpp @@ -37,6 +37,7 @@ using core::TypedExprPtr; const char* const kAnd = "and"; const char* const kOr = "or"; const char* const kRowConstructor = "row_constructor"; +const char* const kRowConstructorWithNull = "row_constructor_with_null"; struct ITypedExprHasher { size_t operator()(const ITypedExpr* expr) const { @@ -212,6 +213,24 @@ ExprPtr getRowConstructorExpr( trackCpuUsage); } +ExprPtr getRowConstructorWithNullExpr( + const TypePtr& type, + std::vector&& compiledChildren, + bool trackCpuUsage) { + static auto rowConstructorVectorFunction = + vectorFunctionFactories().withRLock([](auto& functionMap) { + auto functionIterator = functionMap.find(exec::kRowConstructorWithNull); + return functionIterator->second.factory(exec::kRowConstructorWithNull, {}); + }); + + return std::make_shared( + type, + std::move(compiledChildren), + rowConstructorVectorFunction, + "row_constructor_with_null", + trackCpuUsage); +} + ExprPtr getSpecialForm( const std::string& name, const TypePtr& type, @@ -222,6 +241,11 @@ ExprPtr getSpecialForm( type, std::move(compiledChildren), trackCpuUsage); } + if (name == kRowConstructorWithNull) { + return getRowConstructorWithNullExpr( + type, std::move(compiledChildren), trackCpuUsage); + } + // If we just check the output of constructSpecialForm we'll have moved // compiledChildren, and if the function isn't a special form we'll still need // compiledChildren. Splitting the check in two avoids this use after move. diff --git a/velox/functions/FunctionRegistry.cpp b/velox/functions/FunctionRegistry.cpp index 911f6e20fa9f2..13ac67eab6b03 100644 --- a/velox/functions/FunctionRegistry.cpp +++ b/velox/functions/FunctionRegistry.cpp @@ -109,7 +109,8 @@ std::shared_ptr resolveCallableSpecialForm( const std::string& functionName, const std::vector& argTypes) { // TODO Replace with struct_pack - if (functionName == "row_constructor") { + if (functionName == "row_constructor" || + functionName == "row_constructor_with_null") { auto numInput = argTypes.size(); std::vector types(numInput); std::vector names(numInput); diff --git a/velox/functions/prestosql/RowFunctionWithNull.cpp b/velox/functions/prestosql/RowFunctionWithNull.cpp new file mode 100644 index 0000000000000..fb418162570c6 --- /dev/null +++ b/velox/functions/prestosql/RowFunctionWithNull.cpp @@ -0,0 +1,67 @@ +/* + * 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/expression/Expr.h" +#include "velox/expression/VectorFunction.h" + +namespace facebook::velox::functions { +namespace { + +class RowFunctionWithNull : public exec::VectorFunction { + public: + void apply( + const SelectivityVector& rows, + std::vector& args, + const TypePtr& outputType, + exec::EvalCtx& context, + VectorPtr& result) const override { + auto argsCopy = args; + + BufferPtr nulls = + AlignedBuffer::allocate(bits::nbytes(rows.size()), context.pool(), 1); + auto* nullsPtr = nulls->asMutable(); + auto cntNull = 0; + rows.applyToSelected([&](vector_size_t i) { + bits::clearNull(nullsPtr, i); + if (!bits::isBitNull(nullsPtr, i)) { + for (size_t c = 0; c < argsCopy.size(); c++) { + auto arg = argsCopy[c].get(); + if (arg->mayHaveNulls() && arg->isNullAt(i)) { + // If any argument of the struct is null, set the struct as null. + bits::setNull(nullsPtr, i, true); + cntNull++; + break; + } + } + } + }); + + RowVectorPtr localResult = std::make_shared( + context.pool(), outputType, nulls, rows.size(), std::move(argsCopy), cntNull /*nullCount*/); + context.moveOrCopyResult(localResult, rows, result); + } + + bool isDefaultNullBehavior() const override { + return false; + } +}; +} // namespace + +VELOX_DECLARE_VECTOR_FUNCTION( + udf_concat_row_with_null, + std::vector>{}, + std::make_unique()); + +} // namespace facebook::velox::functions diff --git a/velox/functions/prestosql/registration/GeneralFunctionsRegistration.cpp b/velox/functions/prestosql/registration/GeneralFunctionsRegistration.cpp index 61df9efbd2bbd..9cb3b202a56f2 100644 --- a/velox/functions/prestosql/registration/GeneralFunctionsRegistration.cpp +++ b/velox/functions/prestosql/registration/GeneralFunctionsRegistration.cpp @@ -23,6 +23,7 @@ namespace facebook::velox::functions { void registerAllSpecialFormGeneralFunctions() { VELOX_REGISTER_VECTOR_FUNCTION(udf_in, "in"); VELOX_REGISTER_VECTOR_FUNCTION(udf_concat_row, "row_constructor"); + VELOX_REGISTER_VECTOR_FUNCTION(udf_concat_row_with_null, "row_constructor_with_null"); registerIsNullFunction("is_null"); }