From f4665176bd053413db3a40bd2796d23819197b8a Mon Sep 17 00:00:00 2001 From: rui-mo Date: Thu, 7 Sep 2023 08:56:12 +0800 Subject: [PATCH] Add RowConstructorWithNull expression --- velox/expression/CMakeLists.txt | 1 + velox/expression/RegisterSpecialForm.cpp | 4 ++ velox/expression/RowConstructorWithNull.cpp | 66 +++++++++++++++++ velox/expression/RowConstructorWithNull.h | 44 ++++++++++++ velox/functions/prestosql/CMakeLists.txt | 1 + .../prestosql/RowFunctionWithNull.cpp | 72 +++++++++++++++++++ .../GeneralFunctionsRegistration.cpp | 2 + 7 files changed, 190 insertions(+) create mode 100644 velox/expression/RowConstructorWithNull.cpp create mode 100644 velox/expression/RowConstructorWithNull.h create mode 100644 velox/functions/prestosql/RowFunctionWithNull.cpp diff --git a/velox/expression/CMakeLists.txt b/velox/expression/CMakeLists.txt index 10c850bf66a1..a246533ea4c1 100644 --- a/velox/expression/CMakeLists.txt +++ b/velox/expression/CMakeLists.txt @@ -36,6 +36,7 @@ add_library( VectorFunction.cpp RegisterSpecialForm.cpp RowConstructor.cpp + RowConstructorWithNull.cpp SimpleFunctionRegistry.cpp SpecialFormRegistry.cpp SwitchExpr.cpp diff --git a/velox/expression/RegisterSpecialForm.cpp b/velox/expression/RegisterSpecialForm.cpp index 1de29c58382f..016acdfaaacb 100644 --- a/velox/expression/RegisterSpecialForm.cpp +++ b/velox/expression/RegisterSpecialForm.cpp @@ -22,6 +22,7 @@ #include "velox/expression/ConjunctExpr.h" #include "velox/expression/FunctionCallToSpecialForm.h" #include "velox/expression/RowConstructor.h" +#include "velox/expression/RowConstructorWithNull.h" #include "velox/expression/SpecialFormRegistry.h" #include "velox/expression/SwitchExpr.h" #include "velox/expression/TryExpr.h" @@ -47,5 +48,8 @@ void registerFunctionCallToSpecialForms() { registerFunctionCallToSpecialForm( RowConstructorCallToSpecialForm::kRowConstructor, std::make_unique()); + registerFunctionCallToSpecialForm( + RowConstructorWithNullCallToSpecialForm::kRowConstructorWithNull, + std::make_unique()); } } // namespace facebook::velox::exec diff --git a/velox/expression/RowConstructorWithNull.cpp b/velox/expression/RowConstructorWithNull.cpp new file mode 100644 index 000000000000..f20938aafb72 --- /dev/null +++ b/velox/expression/RowConstructorWithNull.cpp @@ -0,0 +1,66 @@ +/* + * 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/RowConstructorWithNull.h" +#include "velox/expression/VectorFunction.h" + +namespace facebook::velox::exec { + +TypePtr RowConstructorWithNullCallToSpecialForm::resolveType( + const std::vector& argTypes) { + auto numInput = argTypes.size(); + std::vector names(numInput); + std::vector types(numInput); + for (auto i = 0; i < numInput; i++) { + types[i] = argTypes[i]; + names[i] = fmt::format("c{}", i + 1); + } + return ROW(std::move(names), std::move(types)); +} + +ExprPtr RowConstructorWithNullCallToSpecialForm::constructSpecialForm( + const std::string& name, + const TypePtr& type, + std::vector&& compiledChildren, + bool trackCpuUsage, + const core::QueryConfig& config) { + auto rowConstructorVectorFunction = + vectorFunctionFactories().withRLock([&config, &name](auto& functionMap) { + auto functionIterator = functionMap.find(name); + return functionIterator->second.factory(name, {}, config); + }); + + return std::make_shared( + type, + std::move(compiledChildren), + rowConstructorVectorFunction, + name, + trackCpuUsage); +} + +ExprPtr RowConstructorWithNullCallToSpecialForm::constructSpecialForm( + const TypePtr& type, + std::vector&& compiledChildren, + bool trackCpuUsage, + const core::QueryConfig& config) { + return constructSpecialForm( + kRowConstructorWithNull, + type, + std::move(compiledChildren), + trackCpuUsage, + config); +} +} // namespace facebook::velox::exec diff --git a/velox/expression/RowConstructorWithNull.h b/velox/expression/RowConstructorWithNull.h new file mode 100644 index 000000000000..b9993514c828 --- /dev/null +++ b/velox/expression/RowConstructorWithNull.h @@ -0,0 +1,44 @@ +/* + * 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/expression/FunctionCallToSpecialForm.h" +#include "velox/expression/SpecialForm.h" + +namespace facebook::velox::exec { +class RowConstructorWithNullCallToSpecialForm + : public FunctionCallToSpecialForm { + public: + TypePtr resolveType(const std::vector& argTypes) override; + + ExprPtr constructSpecialForm( + const TypePtr& type, + std::vector&& compiledChildren, + bool trackCpuUsage, + const core::QueryConfig& config) override; + + static constexpr const char* kRowConstructorWithNull = + "row_constructor_with_null"; + + protected: + ExprPtr constructSpecialForm( + const std::string& name, + const TypePtr& type, + std::vector&& compiledChildren, + bool trackCpuUsage, + const core::QueryConfig& config); +}; +} // namespace facebook::velox::exec diff --git a/velox/functions/prestosql/CMakeLists.txt b/velox/functions/prestosql/CMakeLists.txt index e2e820452936..dade3728cf2a 100644 --- a/velox/functions/prestosql/CMakeLists.txt +++ b/velox/functions/prestosql/CMakeLists.txt @@ -46,6 +46,7 @@ add_library( Repeat.cpp Reverse.cpp RowFunction.cpp + RowFunctionWithNull.cpp Sequence.cpp SimpleComparisonMatcher.cpp Slice.cpp diff --git a/velox/functions/prestosql/RowFunctionWithNull.cpp b/velox/functions/prestosql/RowFunctionWithNull.cpp new file mode 100644 index 000000000000..fb650de6d3ca --- /dev/null +++ b/velox/functions/prestosql/RowFunctionWithNull.cpp @@ -0,0 +1,72 @@ +/* + * 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 row = std::make_shared( + context.pool(), + outputType, + nulls, + rows.size(), + std::move(argsCopy), + cntNull /*nullCount*/); + context.moveOrCopyResult(row, 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 4dfe6a8082a5..eb045d5707dd 100644 --- a/velox/functions/prestosql/registration/GeneralFunctionsRegistration.cpp +++ b/velox/functions/prestosql/registration/GeneralFunctionsRegistration.cpp @@ -25,6 +25,8 @@ void registerAllSpecialFormGeneralFunctions() { exec::registerFunctionCallToSpecialForms(); 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"); }