Skip to content

Commit

Permalink
resolve comments
Browse files Browse the repository at this point in the history
  • Loading branch information
duanmeng committed Mar 21, 2023
1 parent 0134a1e commit 83dcfab
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 92 deletions.
4 changes: 2 additions & 2 deletions velox/docs/functions/presto/array.rst
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,11 @@ Array Functions

.. function:: any_match(array(T), function(T, boolean)) → boolean

Returns whether any elements of an array match the given predicate.
Returns whether at least one element of an array match the given predicate.

Returns true if one or more elements match the predicate;
Returns false if none of the elements matches (a special case is when the array is empty);
Returns NULL NULL if the predicate function returns NULL for one or more elements and false for all other elements.
Returns NULL if the predicate function returns NULL for one or more elements and false for all other elements.
Throws an exception if the predicate fails for one or more elements and returns false or NULL for the rest.

.. function:: array_average(array(double)) -> double
Expand Down
119 changes: 29 additions & 90 deletions velox/functions/prestosql/ArrayMatch.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
namespace facebook::velox::functions {
namespace {

template <bool match>
class MatchFunction : public exec::VectorFunction {
private:
void apply(
Expand Down Expand Up @@ -90,24 +91,12 @@ class MatchFunction : public exec::VectorFunction {

protected:
static FOLLY_ALWAYS_INLINE bool hasError(
const ErrorVectorPtr& elementErrors,
const ErrorVectorPtr& errors,
int idx) {
return elementErrors && idx < elementErrors->size() &&
!elementErrors->isNullAt(idx);
return errors && idx < errors->size() &&
!errors->isNullAt(idx);
}

private:
virtual void applyInternal(
FlatVector<bool>* flatResult,
exec::EvalCtx& context,
vector_size_t row,
const vector_size_t* offsets,
const vector_size_t* sizes,
const ErrorVectorPtr& elementErrors,
const exec::LocalDecodedVector& bitsDecoder) const = 0;
};

class AllMatchFunction : public MatchFunction {
private:
void applyInternal(
FlatVector<bool>* flatResult,
Expand All @@ -116,10 +105,14 @@ class AllMatchFunction : public MatchFunction {
const vector_size_t* offsets,
const vector_size_t* sizes,
const ErrorVectorPtr& elementErrors,
const exec::LocalDecodedVector& bitsDecoder) const override {
const exec::LocalDecodedVector& bitsDecoder) const {
auto size = sizes[row];
auto offset = offsets[row];
auto allMatch = true;

// All or any match.
// All: early exit on non-match. Initial value = true.
// Any: early exit on match. Initial value = false.
bool result = match;
auto hasNull = false;
std::exception_ptr errorPtr{nullptr};
for (auto i = 0; i < size; ++i) {
Expand All @@ -132,105 +125,51 @@ class AllMatchFunction : public MatchFunction {

if (bitsDecoder->isNullAt(idx)) {
hasNull = true;
} else if (!bitsDecoder->valueAt<bool>(idx)) {
allMatch = false;
} else if (bitsDecoder->valueAt<bool>(idx) == !match) {
result = !result;
break;
}
}

// Errors for individual array elements should be suppressed only if the
// outcome can be decided by some other array element, e.g. if there is
// another element that returns 'false' for the predicate.
if (!allMatch) {
flatResult->set(row, false);
// another element that returns 'true' for the predicate.
if (result != match) {
flatResult->set(row, !match);
} else if (errorPtr) {
context.setError(row, errorPtr);
} else if (hasNull) {
flatResult->setNull(row, true);
} else {
flatResult->set(row, true);
flatResult->set(row, match);
}
}

public:
static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
// array(T), function(T) -> boolean
return {exec::FunctionSignatureBuilder()
.typeVariable("T")
.returnType("boolean")
.argumentType("array(T)")
.argumentType("function(T, boolean)")
.build()};
}
};

class AnyMatchFunction : public MatchFunction {
private:
void applyInternal(
FlatVector<bool>* flatResult,
exec::EvalCtx& context,
vector_size_t row,
const vector_size_t* offsets,
const vector_size_t* sizes,
const ErrorVectorPtr& elementErrors,
const exec::LocalDecodedVector& bitsDecoder) const override {
auto size = sizes[row];
auto offset = offsets[row];
auto nonMatch = true;
auto hasNull = false;
std::exception_ptr errorPtr{nullptr};
for (auto i = 0; i < size; ++i) {
auto idx = offset + i;
if (hasError(elementErrors, idx)) {
errorPtr = *std::static_pointer_cast<std::exception_ptr>(
elementErrors->valueAt(idx));
continue;
}
class AllMatchFunction : public MatchFunction<true> {};

if (bitsDecoder->isNullAt(idx)) {
hasNull = true;
} else if (bitsDecoder->valueAt<bool>(idx)) {
nonMatch = false;
break;
}
}
class AnyMatchFunction : public MatchFunction<false> {};

// Errors for individual array elements should be suppressed only if the
// outcome can be decided by some other array element, e.g. if there is
// another element that returns 'true' for the predicate.
if (!nonMatch) {
flatResult->set(row, true);
} else if (errorPtr) {
context.setError(row, errorPtr);
} else if (hasNull) {
flatResult->setNull(row, true);
} else {
flatResult->set(row, false);
}
}

public:
static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
// array(T), function(T) -> boolean
return {exec::FunctionSignatureBuilder()
.typeVariable("T")
.returnType("boolean")
.argumentType("array(T)")
.argumentType("function(T, boolean)")
.build()};
}
};
std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
// array(T), function(T) -> boolean
return {exec::FunctionSignatureBuilder()
.typeVariable("T")
.returnType("boolean")
.argumentType("array(T)")
.argumentType("function(T, boolean)")
.build()};
}

} // namespace

VELOX_DECLARE_VECTOR_FUNCTION(
udf_all_match,
AllMatchFunction::signatures(),
signatures(),
std::make_unique<AllMatchFunction>());

VELOX_DECLARE_VECTOR_FUNCTION(
udf_any_match,
AnyMatchFunction::signatures(),
signatures(),
std::make_unique<AnyMatchFunction>());

} // namespace facebook::velox::functions

0 comments on commit 83dcfab

Please sign in to comment.