doris/be/test/exprs/vcondition_expr_test.cpp

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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 "exprs/vcondition_expr.h"

#include <gen_cpp/Exprs_types.h>
#include <gen_cpp/Types_types.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include <cmath>
#include <limits>
#include <memory>
#include <string>
#include <vector>

#include "core/assert_cast.h"
#include "core/column/column_nullable.h"
#include "core/column/column_vector.h"
#include "core/data_type/data_type_date_or_datetime_v2.h"
#include "core/data_type/data_type_nullable.h"
#include "core/data_type/data_type_number.h"
#include "core/data_type/data_type_timestamp_ns.h"
#include "exprs/vexpr.h"
#include "exprs/vexpr_context.h"

namespace doris {

// Build a minimal TExprNode as the input of a VConditionExpr.
// Only fields required by the VExpr base ctor (so that create_data_type works) are set.
static TExprNode make_function_node(const std::string& fn_name, TPrimitiveType::type ptype,
                                    bool is_nullable, int scale = -1) {
    TExprNode node;
    node.node_type = TExprNodeType::FUNCTION_CALL;
    node.num_children = 0;
    node.__set_is_nullable(is_nullable);

    TTypeDesc type_desc;
    TTypeNode type_node;
    type_node.type = TTypeNodeType::SCALAR;
    TScalarType scalar_type;
    scalar_type.__set_type(ptype);
    if (scale >= 0) {
        scalar_type.__set_scale(scale);
    }
    type_node.__set_scalar_type(scalar_type);
    type_desc.types.push_back(type_node);
    node.__set_type(type_desc);

    TFunction fn;
    TFunctionName function_name;
    function_name.function_name = fn_name;
    fn.name = function_name;
    node.__set_fn(fn);

    return node;
}

static TExprNode make_coalesce_node(TPrimitiveType::type ptype, bool is_nullable, int scale = -1) {
    return make_function_node("coalesce", ptype, is_nullable, scale);
}

// Mock child expression: returns the pre-injected ColumnPtr / DataTypePtr to the parent expr.
// Behavior is modeled after MockVExprForTryCast in try_cast_expr_test.cpp.
class MockChildVExpr : public VExpr {
public:
    MockChildVExpr(ColumnPtr column, DataTypePtr type)
            : _column(std::move(column)), _type(std::move(type)) {}

    MOCK_CONST_METHOD0(clone, VExprSPtr());

    const std::string& expr_name() const override { return _name; }

    Status execute(VExprContext* context, Block* block, int* result_column_id) const override {
        block->insert({_column, _type, "mock_child"});
        *result_column_id = static_cast<int>(block->columns()) - 1;
        return Status::OK();
    }

    Status execute_column_impl(VExprContext* /*context*/, const Block* /*block*/,
                               const Selector* /*selector*/, size_t /*count*/,
                               ColumnPtr& result_column) const override {
        result_column = _column;
        return Status::OK();
    }

    DataTypePtr execute_type(const Block* /*block*/) const override { return _type; }

private:
    ColumnPtr _column;
    DataTypePtr _type;
    std::string _name = "mock_child";
};

// Helper: build a nullable Float64 column from a list of (value, is_null) pairs.
static ColumnPtr make_nullable_float64_column(const std::vector<std::pair<double, bool>>& values) {
    auto nested = ColumnFloat64::create();
    auto null_map = ColumnUInt8::create();
    for (auto& [v, is_null] : values) {
        nested->insert_value(v);
        null_map->insert_value(is_null ? 1 : 0);
    }
    return ColumnNullable::create(std::move(nested), std::move(null_map));
}

// Helper: build a non-nullable Float64 column from a list of doubles.
static ColumnPtr make_float64_column(const std::vector<double>& values) {
    auto col = ColumnFloat64::create();
    for (auto v : values) {
        col->insert_value(v);
    }
    return col;
}

// Helper: build a non-nullable Boolean column from a list of 0/1 values.
static ColumnPtr make_bool_column(const std::vector<UInt8>& values) {
    auto col = ColumnUInt8::create();
    for (auto v : values) {
        col->insert_value(v);
    }
    return col;
}

// Helper: extract the Float64 value at `row` from the result column,
// handling both nullable and non-nullable cases.
static double get_float64_value(const ColumnPtr& column, size_t row, bool* is_null = nullptr) {
    if (const auto* nullable = check_and_get_column<ColumnNullable>(column.get())) {
        if (is_null) {
            *is_null = nullable->is_null_at(row);
        }
        const auto& nested =
                assert_cast<const ColumnFloat64&>(nullable->get_nested_column()).get_data();
        return nested[row];
    }
    if (is_null) {
        *is_null = false;
    }
    return assert_cast<const ColumnFloat64*>(column.get())->get_data()[row];
}

class VConditionExprCoalesceTest : public ::testing::Test {};

// Targets the fix: after some rows of col0 are filled, NaN/Inf values in the
// same rows of later columns must not pollute the already-filled result via
// "value * 0" arithmetic.
// Input:
//   col0 (nullable, double): [1.0, NULL]
//   col1 (non-nullable, double): [NaN, 100.0]
// Expected coalesce result: [1.0, 100.0]
TEST_F(VConditionExprCoalesceTest, Float64_NaN_NotPolluteResult) {
    auto coalesce_node = make_coalesce_node(TPrimitiveType::DOUBLE, /*is_nullable=*/true);
    auto coalesce_expr = VectorizedCoalesceExpr::create_shared(coalesce_node);

    // _data_type is already set to Nullable(Float64) by the base ctor; reassert it explicitly
    // through the public data_type() accessor (the underlying field is protected).
    coalesce_expr->data_type() =
            std::make_shared<DataTypeNullable>(std::make_shared<DataTypeFloat64>());

    const double kNaN = std::numeric_limits<double>::quiet_NaN();
    auto col0 = make_nullable_float64_column({{1.0, false}, {0.0, true}});
    auto col1 = make_float64_column({kNaN, 100.0});

    auto child0 = std::make_shared<MockChildVExpr>(
            col0, std::make_shared<DataTypeNullable>(std::make_shared<DataTypeFloat64>()));
    auto child1 = std::make_shared<MockChildVExpr>(col1, std::make_shared<DataTypeFloat64>());
    coalesce_expr->add_child(child0);
    coalesce_expr->add_child(child1);

    VExprContext context(coalesce_expr);
    ColumnPtr result;
    auto st = coalesce_expr->execute_column_impl(&context, /*block=*/nullptr,
                                                 /*selector=*/nullptr, /*count=*/2, result);
    ASSERT_TRUE(st.ok()) << st.to_string();
    ASSERT_TRUE(result.get() != nullptr);
    ASSERT_EQ(result->size(), 2);

    bool is_null = false;
    auto v0 = get_float64_value(result, 0, &is_null);
    EXPECT_FALSE(is_null);
    EXPECT_FALSE(std::isnan(v0)) << "row0 must not be polluted by NaN, got " << v0;
    EXPECT_DOUBLE_EQ(v0, 1.0);

    auto v1 = get_float64_value(result, 1, &is_null);
    EXPECT_FALSE(is_null);
    EXPECT_DOUBLE_EQ(v1, 100.0);
}

// Same shape as above, but the polluting value in the later column is +Inf:
//   col0 (nullable, double): [2.0, NULL]
//   col1 (non-nullable, double): [+Inf, 200.0]
// Expected coalesce result: [2.0, 200.0]
TEST_F(VConditionExprCoalesceTest, Float64_Inf_NotPolluteResult) {
    auto coalesce_node = make_coalesce_node(TPrimitiveType::DOUBLE, /*is_nullable=*/true);
    auto coalesce_expr = VectorizedCoalesceExpr::create_shared(coalesce_node);
    coalesce_expr->data_type() =
            std::make_shared<DataTypeNullable>(std::make_shared<DataTypeFloat64>());

    const double kInf = std::numeric_limits<double>::infinity();
    auto col0 = make_nullable_float64_column({{2.0, false}, {0.0, true}});
    auto col1 = make_float64_column({kInf, 200.0});

    coalesce_expr->add_child(std::make_shared<MockChildVExpr>(
            col0, std::make_shared<DataTypeNullable>(std::make_shared<DataTypeFloat64>())));
    coalesce_expr->add_child(
            std::make_shared<MockChildVExpr>(col1, std::make_shared<DataTypeFloat64>()));

    VExprContext context(coalesce_expr);
    ColumnPtr result;
    auto st = coalesce_expr->execute_column_impl(&context, /*block=*/nullptr,
                                                 /*selector=*/nullptr, /*count=*/2, result);
    ASSERT_TRUE(st.ok()) << st.to_string();
    ASSERT_EQ(result->size(), 2);

    bool is_null = false;
    auto v0 = get_float64_value(result, 0, &is_null);
    EXPECT_FALSE(is_null);
    EXPECT_FALSE(std::isnan(v0)) << "row0 must not be polluted by Inf*0, got " << v0;
    EXPECT_TRUE(std::isfinite(v0));
    EXPECT_DOUBLE_EQ(v0, 2.0);

    auto v1 = get_float64_value(result, 1, &is_null);
    EXPECT_FALSE(is_null);
    EXPECT_DOUBLE_EQ(v1, 200.0);
}

// Verify that when NaN itself is the valid non-null value picked by coalesce,
// the result must preserve the NaN (the fix should not over-correct).
//   col0 (nullable, double): [NaN, NULL]
//   col1 (non-nullable, double): [10.0, 20.0]
// Expected: row0 takes col0's NaN (non-null, value is NaN); row1 takes col1's 20.0.
TEST_F(VConditionExprCoalesceTest, Float64_NaN_PreservedWhenSelected) {
    auto coalesce_node = make_coalesce_node(TPrimitiveType::DOUBLE, /*is_nullable=*/true);
    auto coalesce_expr = VectorizedCoalesceExpr::create_shared(coalesce_node);
    coalesce_expr->data_type() =
            std::make_shared<DataTypeNullable>(std::make_shared<DataTypeFloat64>());

    const double kNaN = std::numeric_limits<double>::quiet_NaN();
    auto col0 = make_nullable_float64_column({{kNaN, false}, {0.0, true}});
    auto col1 = make_float64_column({10.0, 20.0});

    coalesce_expr->add_child(std::make_shared<MockChildVExpr>(
            col0, std::make_shared<DataTypeNullable>(std::make_shared<DataTypeFloat64>())));
    coalesce_expr->add_child(
            std::make_shared<MockChildVExpr>(col1, std::make_shared<DataTypeFloat64>()));

    VExprContext context(coalesce_expr);
    ColumnPtr result;
    auto st = coalesce_expr->execute_column_impl(&context, /*block=*/nullptr,
                                                 /*selector=*/nullptr, /*count=*/2, result);
    ASSERT_TRUE(st.ok()) << st.to_string();
    ASSERT_EQ(result->size(), 2);

    bool is_null = false;
    auto v0 = get_float64_value(result, 0, &is_null);
    EXPECT_FALSE(is_null);
    EXPECT_TRUE(std::isnan(v0)) << "row0 should preserve NaN from col0";

    auto v1 = get_float64_value(result, 1, &is_null);
    EXPECT_FALSE(is_null);
    EXPECT_DOUBLE_EQ(v1, 20.0);
}

// Float32 variant: same shape as the Float64 NaN case.
TEST_F(VConditionExprCoalesceTest, Float32_NaN_NotPolluteResult) {
    auto coalesce_node = make_coalesce_node(TPrimitiveType::FLOAT, /*is_nullable=*/true);
    auto coalesce_expr = VectorizedCoalesceExpr::create_shared(coalesce_node);
    coalesce_expr->data_type() =
            std::make_shared<DataTypeNullable>(std::make_shared<DataTypeFloat32>());

    const float kNaN = std::numeric_limits<float>::quiet_NaN();

    auto nested0 = ColumnFloat32::create();
    nested0->insert_value(1.5f);
    nested0->insert_value(0.0f);
    auto null_map0 = ColumnUInt8::create();
    null_map0->insert_value(0);
    null_map0->insert_value(1);
    ColumnPtr col0 = ColumnNullable::create(std::move(nested0), std::move(null_map0));

    auto nested1 = ColumnFloat32::create();
    nested1->insert_value(kNaN);
    nested1->insert_value(7.5f);
    ColumnPtr col1 = std::move(nested1);

    coalesce_expr->add_child(std::make_shared<MockChildVExpr>(
            col0, std::make_shared<DataTypeNullable>(std::make_shared<DataTypeFloat32>())));
    coalesce_expr->add_child(
            std::make_shared<MockChildVExpr>(col1, std::make_shared<DataTypeFloat32>()));

    VExprContext context(coalesce_expr);
    ColumnPtr result;
    auto st = coalesce_expr->execute_column_impl(&context, /*block=*/nullptr,
                                                 /*selector=*/nullptr, /*count=*/2, result);
    ASSERT_TRUE(st.ok()) << st.to_string();
    ASSERT_EQ(result->size(), 2);

    auto get_f32 = [&](size_t row, bool* is_null) -> float {
        if (const auto* nullable = check_and_get_column<ColumnNullable>(result.get())) {
            *is_null = nullable->is_null_at(row);
            return assert_cast<const ColumnFloat32&>(nullable->get_nested_column()).get_data()[row];
        }
        *is_null = false;
        return assert_cast<const ColumnFloat32*>(result.get())->get_data()[row];
    };

    bool is_null = false;
    float v0 = get_f32(0, &is_null);
    EXPECT_FALSE(is_null);
    EXPECT_FALSE(std::isnan(v0)) << "row0 must not be polluted by NaN, got " << v0;
    EXPECT_FLOAT_EQ(v0, 1.5f);

    float v1 = get_f32(1, &is_null);
    EXPECT_FALSE(is_null);
    EXPECT_FLOAT_EQ(v1, 7.5f);
}

// Regression: verify that the non-floating (integer) path is unaffected by the Float branch change.
TEST_F(VConditionExprCoalesceTest, Int32_NormalPathStillWorks) {
    auto coalesce_node = make_coalesce_node(TPrimitiveType::INT, /*is_nullable=*/true);
    auto coalesce_expr = VectorizedCoalesceExpr::create_shared(coalesce_node);
    coalesce_expr->data_type() =
            std::make_shared<DataTypeNullable>(std::make_shared<DataTypeInt32>());

    auto nested0 = ColumnInt32::create();
    nested0->insert_value(11);
    nested0->insert_value(0);
    auto null_map0 = ColumnUInt8::create();
    null_map0->insert_value(0);
    null_map0->insert_value(1);
    ColumnPtr col0 = ColumnNullable::create(std::move(nested0), std::move(null_map0));

    auto nested1 = ColumnInt32::create();
    nested1->insert_value(99);
    nested1->insert_value(22);
    ColumnPtr col1 = std::move(nested1);

    coalesce_expr->add_child(std::make_shared<MockChildVExpr>(
            col0, std::make_shared<DataTypeNullable>(std::make_shared<DataTypeInt32>())));
    coalesce_expr->add_child(
            std::make_shared<MockChildVExpr>(col1, std::make_shared<DataTypeInt32>()));

    VExprContext context(coalesce_expr);
    ColumnPtr result;
    auto st = coalesce_expr->execute_column_impl(&context, /*block=*/nullptr,
                                                 /*selector=*/nullptr, /*count=*/2, result);
    ASSERT_TRUE(st.ok()) << st.to_string();
    ASSERT_EQ(result->size(), 2);

    auto get_int = [&](size_t row, bool* is_null) -> int32_t {
        if (const auto* nullable = check_and_get_column<ColumnNullable>(result.get())) {
            *is_null = nullable->is_null_at(row);
            return assert_cast<const ColumnInt32&>(nullable->get_nested_column()).get_data()[row];
        }
        *is_null = false;
        return assert_cast<const ColumnInt32*>(result.get())->get_data()[row];
    };

    bool is_null = false;
    EXPECT_EQ(get_int(0, &is_null), 11);
    EXPECT_FALSE(is_null);
    EXPECT_EQ(get_int(1, &is_null), 22);
    EXPECT_FALSE(is_null);
}

TEST_F(VConditionExprCoalesceTest, TimeStampNs) {
    auto coalesce_node =
            make_coalesce_node(TPrimitiveType::TIMESTAMP_NS, /*is_nullable=*/true, /*scale=*/9);
    auto coalesce_expr = VectorizedCoalesceExpr::create_shared(coalesce_node);
    auto nano_type = std::make_shared<DataTypeTimeStampNs>();
    coalesce_expr->data_type() = std::make_shared<DataTypeNullable>(nano_type);

    auto nested0 = ColumnTimeStampNs::create();
    nested0->insert_value(TimeStampNsValue(11));
    nested0->insert_value(TimeStampNsValue(0));
    nested0->insert_value(TimeStampNsValue(-17));
    nested0->insert_value(TimeStampNsValue(0));
    auto null_map0 = ColumnUInt8::create();
    null_map0->insert_value(0);
    null_map0->insert_value(1);
    null_map0->insert_value(0);
    null_map0->insert_value(1);
    ColumnPtr col0 = ColumnNullable::create(std::move(nested0), std::move(null_map0));

    auto col1 = ColumnTimeStampNs::create();
    col1->insert_value(TimeStampNsValue(99));
    col1->insert_value(TimeStampNsValue(22));
    col1->insert_value(TimeStampNsValue(std::numeric_limits<int64_t>::max()));
    col1->insert_value(TimeStampNsValue(std::numeric_limits<int64_t>::min()));

    coalesce_expr->add_child(
            std::make_shared<MockChildVExpr>(col0, std::make_shared<DataTypeNullable>(nano_type)));
    coalesce_expr->add_child(std::make_shared<MockChildVExpr>(std::move(col1), nano_type));

    VExprContext context(coalesce_expr);
    ColumnPtr result;
    auto st = coalesce_expr->execute_column_impl(&context, /*block=*/nullptr,
                                                 /*selector=*/nullptr, /*count=*/4, result);
    ASSERT_TRUE(st.ok()) << st.to_string();
    ASSERT_EQ(result->size(), 4);

    const auto& nullable = assert_cast<const ColumnNullable&>(*result);
    const auto& values =
            assert_cast<const ColumnTimeStampNs&>(nullable.get_nested_column()).get_data();
    EXPECT_FALSE(nullable.is_null_at(0));
    EXPECT_EQ(values[0].epoch_nanos(), 11);
    EXPECT_FALSE(nullable.is_null_at(1));
    EXPECT_EQ(values[1].epoch_nanos(), 22);
    EXPECT_FALSE(nullable.is_null_at(2));
    EXPECT_EQ(values[2].epoch_nanos(), -17);
    EXPECT_FALSE(nullable.is_null_at(3));
    EXPECT_EQ(values[3].epoch_nanos(), std::numeric_limits<int64_t>::min());
}

class VConditionExprIfTest : public ::testing::Test {};

// Same shape as `IF(NULLIF(b, p), f, b)`: NULLIF wraps b's column itself as the nested
// column of its Nullable(Boolean) result. IF treats a NULL condition as false; that
// normalization must not be written into the shared nested column, otherwise the else
// branch (and every other user of b in the block) reads polluted values.
// Input:
//   b (non-nullable bool): [1, 1]
//   cond = Nullable(nested = b's column, null_map = [0, 1]), logically [true, NULL]
//   f (non-nullable bool): [0, 0]
// Expected IF(cond, f, b): [0, 1]; b (and cond's nested column) must stay [1, 1].
TEST_F(VConditionExprIfTest, NullableCondition_NotPolluteSharedNestedColumn) {
    auto if_node = make_function_node("if", TPrimitiveType::BOOLEAN, /*is_nullable=*/false);
    auto if_expr = VectorizedIfExpr::create_shared(if_node);
    auto bool_type = std::make_shared<DataTypeUInt8>();
    if_expr->data_type() = bool_type;

    ColumnPtr b_column = make_bool_column({1, 1});
    ColumnPtr f_column = make_bool_column({0, 0});
    ColumnPtr cond_column = ColumnNullable::create(b_column, make_bool_column({0, 1}));

    if_expr->add_child(std::make_shared<MockChildVExpr>(
            cond_column, std::make_shared<DataTypeNullable>(bool_type)));
    if_expr->add_child(std::make_shared<MockChildVExpr>(f_column, bool_type));
    if_expr->add_child(std::make_shared<MockChildVExpr>(b_column, bool_type));

    VExprContext context(if_expr);
    ColumnPtr result;
    auto st = if_expr->execute_column_impl(&context, /*block=*/nullptr, /*selector=*/nullptr,
                                           /*count=*/2, result);
    ASSERT_TRUE(st.ok()) << st.to_string();
    ASSERT_TRUE(result.get() != nullptr);
    const auto& result_data = assert_cast<const ColumnUInt8&>(*result).get_data();
    ASSERT_EQ(result_data.size(), 2);
    EXPECT_EQ(result_data[0], 0);
    EXPECT_EQ(result_data[1], 1);

    const auto& b_data = assert_cast<const ColumnUInt8&>(*b_column).get_data();
    EXPECT_EQ(b_data[0], 1);
    EXPECT_EQ(b_data[1], 1);
    const auto& cond_nested_data =
            assert_cast<const ColumnUInt8&>(
                    assert_cast<const ColumnNullable&>(*cond_column).get_nested_column())
                    .get_data();
    EXPECT_EQ(cond_nested_data[0], 1);
    EXPECT_EQ(cond_nested_data[1], 1);
}

} // namespace doris