[fix](be) Preserve floating point values in CASE branch selection

be/benchmark/benchmark_case_expr.hpp

@@ -0,0 +1,93 @@
+// 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.
+
+#pragma once
+
+#include <benchmark/benchmark.h>
+
+#include <random>
+
+#include "core/data_type/data_type_number.h"
+#include "exprs/vcase_expr.h"
+
+namespace doris {
+
+// Keep the production result assembly visible in disassembly as well as timing it.
+template <typename IndexType, PrimitiveType PT>
+NO_INLINE ColumnPtr run_case_selection(const VCaseExpr& expr, const IndexType* indices,
+                                       std::vector<ColumnPtr>& columns, size_t rows) {
+    return expr._execute_update_result_impl<IndexType, ColumnVector<PT>>(indices, columns, rows);
+}
+
+template <typename IndexType, PrimitiveType PT>
+void BM_CaseFloatSelection(benchmark::State& state) {
+    const size_t rows = state.range(0);
+    const size_t branches = state.range(1);
+    const auto distribution = state.range(2);
+    TExprNode node;
+    node.__set_node_type(TExprNodeType::CASE_EXPR);
+    node.__set_type(DataTypeNumber<PT>().to_thrift());
+    node.__set_is_nullable(false);
+    node.case_expr.__set_has_else_expr(true);
+    VCaseExpr expr(node);
+    std::mt19937 rng(20260912);
+    std::vector<IndexType> indices(rows);
+    for (size_t row = 0; row < rows; ++row) {
+        // Interleaved, random, or 99% ELSE. All inputs are finite for before/after comparison.
+        indices[row] = static_cast<IndexType>(distribution == 0 ? row % branches
+                                              : distribution == 1
+                                                      ? rng() % branches
+                                                      : (row % 100 == 0 ? rng() % branches : 0));
+    }
+    std::vector<ColumnPtr> columns;
+    for (size_t branch = 0; branch < branches; ++branch) {
+        auto column = ColumnVector<PT>::create(rows);
+        for (size_t row = 0; row < rows; ++row) {
+            column->get_data()[row] =
+                    static_cast<typename ColumnVector<PT>::value_type>((row + branch + 1) * 0.125);
+        }
+        columns.push_back(std::move(column));
+    }
+    for (auto _ : state) {
+        auto result = run_case_selection<IndexType, PT>(expr, indices.data(), columns, rows);
+        benchmark::DoNotOptimize(result);
+    }
+    state.SetItemsProcessed(state.iterations() * rows);
+}
+
+inline void case_float_arguments(benchmark::internal::Benchmark* benchmark) {
+    for (int64_t rows : {31, 4096, 65536}) {
+        for (int64_t branches : {3, 16}) {
+            for (int64_t distribution : {0, 1, 2}) {
+                benchmark->Args({rows, branches, distribution});
+            }
+        }
+    }
+}
+
+BENCHMARK_TEMPLATE(BM_CaseFloatSelection, uint8_t, TYPE_FLOAT)->Apply(case_float_arguments);
+BENCHMARK_TEMPLATE(BM_CaseFloatSelection, uint8_t, TYPE_DOUBLE)->Apply(case_float_arguments);
+BENCHMARK_TEMPLATE(BM_CaseFloatSelection, uint16_t, TYPE_FLOAT)
+        ->Args({4096, 257, 0})
+        ->Args({4096, 257, 1})
+        ->Args({4096, 257, 2});
+BENCHMARK_TEMPLATE(BM_CaseFloatSelection, uint16_t, TYPE_DOUBLE)
+        ->Args({4096, 257, 0})
+        ->Args({4096, 257, 1})
+        ->Args({4096, 257, 2});
+
+} // namespace doris

be/benchmark/benchmark_main.cpp

@@ -26,6 +26,7 @@
 #include "benchmark_arrow_validation.hpp"
 #include "benchmark_binary_arithmetic.hpp"
 #include "benchmark_bit_pack.hpp"
+#include "benchmark_case_expr.hpp"
 #include "benchmark_column_array_view.hpp"
 #include "benchmark_column_array_view_distance.hpp"
 #include "benchmark_fastunion.hpp"

be/src/exprs/vcase_expr.h

@@ -249,12 +249,22 @@ class VCaseExpr final : public VExpr {
                             then_columns[i].get())
                             ->get_data()
                             .data();
-            if constexpr (std::is_same_v<ColumnType, ColumnDate> ||
-                          std::is_same_v<ColumnType, ColumnDateTime> ||
-                          std::is_same_v<ColumnType, ColumnDateV2> ||
-                          std::is_same_v<ColumnType, ColumnDateTimeV2> ||
-                          std::is_same_v<ColumnType, ColumnTimeStampNs> ||
-                          std::is_same_v<ColumnType, ColumnTimeStampTz>) {
+            if constexpr (std::is_same_v<ColumnType, ColumnFloat32> ||
+                          std::is_same_v<ColumnType, ColumnFloat64>) {
+                // Arithmetic masking propagates unselected NaN/Infinity and loses signed zero.
+                // Conditional stores also let the compiler vectorize without loading from a
+                // selected source/destination pointer, as a ternary assignment can do.
+                for (size_t row_idx = 0; row_idx < rows_count; row_idx++) {
+                    if (then_idx[row_idx] == i) {
+                        result_raw_data[row_idx] = column_raw_data[row_idx];
+                    }
+                }
+            } else if constexpr (std::is_same_v<ColumnType, ColumnDate> ||
+                                 std::is_same_v<ColumnType, ColumnDateTime> ||
+                                 std::is_same_v<ColumnType, ColumnDateV2> ||
+                                 std::is_same_v<ColumnType, ColumnDateTimeV2> ||
+                                 std::is_same_v<ColumnType, ColumnTimeStampNs> ||
+                                 std::is_same_v<ColumnType, ColumnTimeStampTz>) {
                 for (int row_idx = 0; row_idx < rows_count; row_idx++) {
                     result_raw_data[row_idx] = (then_idx[row_idx] == i) ? column_raw_data[row_idx]
                                                                         : result_raw_data[row_idx];

be/test/exprs/vcase_expr_test.cpp

@@ -0,0 +1,120 @@
+// 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/vcase_expr.h"
+
+#include <gtest/gtest.h>
+
+#include <array>
+#include <bit>
+#include <limits>
+#include <type_traits>
+
+#include "core/data_type/data_type_number.h"
+
+namespace doris {
+
+template <typename Index, PrimitiveType PT>
+struct CaseFloatTypes {
+    using IndexType = Index;
+    using ColumnType = ColumnVector<PT>;
+    using DataType = DataTypeNumber<PT>;
+};
+
+using CaseFloatTestTypes =
+        ::testing::Types<CaseFloatTypes<uint8_t, TYPE_FLOAT>, CaseFloatTypes<uint8_t, TYPE_DOUBLE>,
+                         CaseFloatTypes<uint16_t, TYPE_FLOAT>,
+                         CaseFloatTypes<uint16_t, TYPE_DOUBLE>>;
+
+template <typename T>
+class VCaseFloatTest : public ::testing::Test {
+protected:
+    using Index = typename T::IndexType;
+    using Column = typename T::ColumnType;
+    using Value = typename Column::value_type;
+    using Bits = std::conditional_t<sizeof(Value) == 4, uint32_t, uint64_t>;
+
+    void check_selection(size_t rows, size_t branches, bool constant) {
+        SCOPED_TRACE(::testing::Message()
+                     << "rows=" << rows << " branches=" << branches << " constant=" << constant);
+        TExprNode node;
+        node.__set_node_type(TExprNodeType::CASE_EXPR);
+        node.__set_type(typename T::DataType().to_thrift());
+        node.__set_is_nullable(false);
+        node.case_expr.__set_has_else_expr(true);
+        VCaseExpr expr(node);
+        // Arithmetic masking corrupts unselected infinities/NaNs, and adding to +0 loses -0.
+        const std::array<Value, 9> values = {Value(1.25),
+                                             Value(-2.5),
+                                             Value(0.0),
+                                             Value(-0.0),
+                                             std::numeric_limits<Value>::infinity(),
+                                             -std::numeric_limits<Value>::infinity(),
+                                             std::numeric_limits<Value>::quiet_NaN(),
+                                             std::numeric_limits<Value>::denorm_min(),
+                                             std::numeric_limits<Value>::max()};
+        std::vector<Index> indices(rows);
+        for (size_t row = 0; row < rows; ++row) {
+            indices[row] = row % branches;
+        }
+        std::vector<ColumnPtr> columns;
+        for (size_t branch = 0; branch < branches; ++branch) {
+            auto column = Column::create(constant ? 1 : rows);
+            for (size_t row = 0; row < column->size(); ++row) {
+                column->get_data()[row] = values[(row / branches + branch) % values.size()];
+            }
+            if (constant) {
+                columns.push_back(ColumnConst::create(std::move(column), rows));
+            } else {
+                columns.push_back(std::move(column));
+            }
+        }
+        auto result = expr.template _execute_update_result_impl<Index, Column>(indices.data(),
+                                                                               columns, rows);
+        const auto& actual = assert_cast<const Column&>(*result).get_data();
+        ASSERT_EQ(actual.size(), rows);
+        for (size_t row = 0; row < rows; ++row) {
+            const auto expected =
+                    values[((constant ? 0 : row / branches) + indices[row]) % values.size()];
+            // Compare bits to include NaN payloads, signed zero and subnormal values.
+            ASSERT_EQ(std::bit_cast<Bits>(actual[row]), std::bit_cast<Bits>(expected)) << row;
+        }
+    }
+};
+
+TYPED_TEST_SUITE(VCaseFloatTest, CaseFloatTestTypes);
+
+TYPED_TEST(VCaseFloatTest, NonFiniteValuesAndVectorTails) {
+    for (size_t rows : {0, 1, 3, 7, 8, 15, 16, 31, 32, 33, 4095, 4096, 4099}) {
+        this->check_selection(rows, 9, false);
+    }
+}
+
+TYPED_TEST(VCaseFloatTest, ConstantBranches) {
+    for (size_t rows : {1, 31, 4099}) {
+        this->check_selection(rows, 9, true);
+    }
+}
+
+TYPED_TEST(VCaseFloatTest, MaximumAndWideBranchIndices) {
+    // 255 columns still use uint8_t; 257 columns exercise indices beyond the uint8_t range.
+    const size_t branches = sizeof(typename TypeParam::IndexType) == 1 ? 255 : 257;
+    this->check_selection(4099, branches, false);
+    this->check_selection(4099, branches, true);
+}
+
+} // namespace doris

regression-test/data/query_p0/sql_functions/conditional_functions/test_case_float_nonfinite.out

@@ -0,0 +1,81 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !double_false_then --
+0	1	1	1
+1	Infinity	-Infinity	NaN
+2	2	2	2
+
+-- !double_false_else --
+0	1	1	1
+1	2	2	2
+2	Infinity	-Infinity	NaN
+
+-- !double_false_nullable --
+0	1
+1	Infinity
+2	\N
+
+-- !double_false_batches --
+1	1
+2	4097
+Infinity	1
+
+-- !float_false_then --
+0	1.0	1.0	1.0
+1	Infinity	-Infinity	NaN
+2	2.0	2.0	2.0
+
+-- !float_false_else --
+0	1.0	1.0	1.0
+1	2.0	2.0	2.0
+2	Infinity	-Infinity	NaN
+
+-- !float_false_nullable --
+0	1.0
+1	Infinity
+2	\N
+
+-- !float_false_batches --
+1.0	1
+2.0	4097
+Infinity	1
+
+-- !double_true_then --
+0	1	1	1
+1	Infinity	-Infinity	NaN
+2	2	2	2
+
+-- !double_true_else --
+0	1	1	1
+1	2	2	2
+2	Infinity	-Infinity	NaN
+
+-- !double_true_nullable --
+0	1
+1	Infinity
+2	\N
+
+-- !double_true_batches --
+1	1
+2	4097
+Infinity	1
+
+-- !float_true_then --
+0	1.0	1.0	1.0
+1	Infinity	-Infinity	NaN
+2	2.0	2.0	2.0
+
+-- !float_true_else --
+0	1.0	1.0	1.0
+1	2.0	2.0	2.0
+2	Infinity	-Infinity	NaN
+
+-- !float_true_nullable --
+0	1.0
+1	Infinity
+2	\N
+
+-- !float_true_batches --
+1.0	1
+2.0	4097
+Infinity	1
+

regression-test/suites/query_p0/sql_functions/conditional_functions/test_case_float_nonfinite.groovy

@@ -0,0 +1,64 @@
+// 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.
+
+suite("test_case_float_nonfinite") {
+    for (def shortCircuit : [false, true]) {
+        sql "set short_circuit_evaluation = ${shortCircuit}"
+        for (def type : ["double", "float"]) {
+            // Every row overflows from finite operands. Keeping number in the expression
+            // prevents constant folding and makes unselected branches contain infinity too.
+            def infinity = "cast((cast(number as double) + cast(1e308 as double)) * cast(1e308 as double) as ${type})"
+            def nan = "cast((${infinity}) - (${infinity}) as ${type})"
+            "qt_${type}_${shortCircuit}_then" """
+                select number,
+                       case when number = 0 then cast(1 as ${type})
+                            when number = 1 then ${infinity} else cast(2 as ${type}) end,
+                       case when number = 0 then cast(1 as ${type})
+                            when number = 1 then cast(-(${infinity}) as ${type}) else cast(2 as ${type}) end,
+                       case when number = 0 then cast(1 as ${type})
+                            when number = 1 then ${nan} else cast(2 as ${type}) end
+                from numbers("number" = "3") order by number
+            """
+            "qt_${type}_${shortCircuit}_else" """
+                select number,
+                       case when number = 0 then cast(1 as ${type})
+                            when number = 1 then cast(2 as ${type}) else ${infinity} end,
+                       case when number = 0 then cast(1 as ${type})
+                            when number = 1 then cast(2 as ${type}) else cast(-(${infinity}) as ${type}) end,
+                       case when number = 0 then cast(1 as ${type})
+                            when number = 1 then cast(2 as ${type}) else ${nan} end
+                from numbers("number" = "3") order by number
+            """
+            "qt_${type}_${shortCircuit}_nullable" """
+                select number,
+                       case when number = 0 then cast(1 as ${type})
+                            when number = 1 then ${infinity} end
+                from numbers("number" = "3") order by number
+            """
+            // The overflowing branch is evaluated in the first batch. Finite rows must
+            // have the same result there and in the tail batch where it is never selected.
+            "qt_${type}_${shortCircuit}_batches" """
+                select result, count(*) from (
+                    select case when number = 0 then cast(1 as ${type})
+                                when number = 1 then ${infinity}
+                                else cast(2 as ${type}) end as result
+                    from numbers("number" = "4099")
+                ) t group by result order by result
+            """
+        }
+    }
+}