Skip to content

Commit

Permalink
add another
Browse files Browse the repository at this point in the history
  • Loading branch information
rui-mo committed Apr 23, 2023
1 parent fed7b95 commit 8ca4671
Show file tree
Hide file tree
Showing 4 changed files with 94 additions and 1 deletion.
24 changes: 24 additions & 0 deletions velox/expression/ExprCompiler.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -212,6 +213,24 @@ ExprPtr getRowConstructorExpr(
trackCpuUsage);
}

ExprPtr getRowConstructorWithNullExpr(
const TypePtr& type,
std::vector<ExprPtr>&& 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<Expr>(
type,
std::move(compiledChildren),
rowConstructorVectorFunction,
"row_constructor_with_null",
trackCpuUsage);
}

ExprPtr getSpecialForm(
const std::string& name,
const TypePtr& type,
Expand All @@ -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.
Expand Down
3 changes: 2 additions & 1 deletion velox/functions/FunctionRegistry.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,8 @@ std::shared_ptr<const Type> resolveCallableSpecialForm(
const std::string& functionName,
const std::vector<TypePtr>& 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<TypePtr> types(numInput);
std::vector<std::string> names(numInput);
Expand Down
67 changes: 67 additions & 0 deletions velox/functions/prestosql/RowFunctionWithNull.cpp
Original file line number Diff line number Diff line change
@@ -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<VectorPtr>& args,
const TypePtr& outputType,
exec::EvalCtx& context,
VectorPtr& result) const override {
auto argsCopy = args;

BufferPtr nulls =
AlignedBuffer::allocate<char>(bits::nbytes(rows.size()), context.pool(), 1);
auto* nullsPtr = nulls->asMutable<uint64_t>();
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<RowVector>(
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::shared_ptr<exec::FunctionSignature>>{},
std::make_unique<RowFunctionWithNull>());

} // namespace facebook::velox::functions
Original file line number Diff line number Diff line change
Expand Up @@ -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");
}

Expand Down

0 comments on commit 8ca4671

Please sign in to comment.