From 26a6e1b94648a0e5debcf662929aee1fc3eeb058 Mon Sep 17 00:00:00 2001
From: Jiquan Long <jiquan.long@zilliz.com>
Date: Fri, 16 Jun 2023 09:56:39 +0800
Subject: [PATCH] Fix BinaryRange expression on integer overflow case (#24920)

Signed-off-by: longjiquan <jiquan.long@zilliz.com>
---
 internal/core/src/query/PlanProto.cpp         | 18 +++++--------
 .../src/query/visitors/ExecExprVisitor.cpp    | 25 +++++++++++++++++--
 internal/core/unittest/test_query.cpp         |  2 +-
 3 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/internal/core/src/query/PlanProto.cpp b/internal/core/src/query/PlanProto.cpp
index bca61c934..b8bd6981e 100644
--- a/internal/core/src/query/PlanProto.cpp
+++ b/internal/core/src/query/PlanProto.cpp
@@ -342,22 +342,16 @@ ProtoParser::ParseBinaryRangeExpr(const proto::plan::BinaryRangeExpr& expr_pb) {
                 return ExtractBinaryRangeExprImpl<bool>(
                     field_id, data_type, expr_pb);
             }
-            case DataType::INT8: {
-                return ExtractBinaryRangeExprImpl<int8_t>(
-                    field_id, data_type, expr_pb);
-            }
-            case DataType::INT16: {
-                return ExtractBinaryRangeExprImpl<int16_t>(
-                    field_id, data_type, expr_pb);
-            }
-            case DataType::INT32: {
-                return ExtractBinaryRangeExprImpl<int32_t>(
-                    field_id, data_type, expr_pb);
-            }
+
+            // see also: https://github.com/milvus-io/milvus/issues/23646.
+            case DataType::INT8:
+            case DataType::INT16:
+            case DataType::INT32:
             case DataType::INT64: {
                 return ExtractBinaryRangeExprImpl<int64_t>(
                     field_id, data_type, expr_pb);
             }
+
             case DataType::FLOAT: {
                 return ExtractBinaryRangeExprImpl<float>(
                     field_id, data_type, expr_pb);
diff --git a/internal/core/src/query/visitors/ExecExprVisitor.cpp b/internal/core/src/query/visitors/ExecExprVisitor.cpp
index 1a5930652..62362c3c9 100644
--- a/internal/core/src/query/visitors/ExecExprVisitor.cpp
+++ b/internal/core/src/query/visitors/ExecExprVisitor.cpp
@@ -939,12 +939,33 @@ ExecExprVisitor::ExecBinaryRangeVisitorDispatcher(BinaryRangeExpr& expr_raw)
 
     bool lower_inclusive = expr.lower_inclusive_;
     bool upper_inclusive = expr.upper_inclusive_;
-    IndexInnerType val1 = expr.lower_value_;
-    IndexInnerType val2 = expr.upper_value_;
+
+    // see also: https://github.com/milvus-io/milvus/issues/23646.
+    typedef std::conditional_t<std::is_integral_v<IndexInnerType>,
+                               int64_t,
+                               IndexInnerType>
+        HighPrecisionType;
+
+    auto val1 = static_cast<HighPrecisionType>(expr.lower_value_);
+    auto val2 = static_cast<HighPrecisionType>(expr.upper_value_);
 
     auto index_func = [&](Index* index) {
+        if constexpr (std::is_integral_v<T>) {
+            if (gt_ub<T>(val1)) {
+                return TargetBitmap(index->Size(), false);
+            } else if (lt_lb<T>(val1)) {
+                val1 = std::numeric_limits<T>::min();
+            }
+
+            if (gt_ub<T>(val2)) {
+                val2 = std::numeric_limits<T>::max();
+            } else if (lt_lb<T>(val2)) {
+                return TargetBitmap(index->Size(), false);
+            }
+        }
         return index->Range(val1, lower_inclusive, val2, upper_inclusive);
     };
+
     if (lower_inclusive && upper_inclusive) {
         auto elem_func = [val1, val2](MayConstRef<T> x) {
             return (val1 <= x && x <= val2);
diff --git a/internal/core/unittest/test_query.cpp b/internal/core/unittest/test_query.cpp
index 8830f0827..708159d68 100644
--- a/internal/core/unittest/test_query.cpp
+++ b/internal/core/unittest/test_query.cpp
@@ -701,7 +701,7 @@ TEST(Query, FillSegment) {
     auto dataset = DataGen(schema, N);
     const auto std_vec = dataset.get_col<int64_t>(FieldId(101));  // ids field
     const auto std_vfloat_vec =
-        dataset.get_col<float>(FieldId(100));    // vector field
+        dataset.get_col<float>(FieldId(100));  // vector field
     const auto std_i32_vec =
         dataset.get_col<int32_t>(FieldId(102));  // scalar field
 
-- 
GitLab