Fix literal rewrite rule for predicate pushdown (#809)

* Fix range union

* fix ydb/library/yql/core/extract_predicate/ut

* fix ydb/library/yql/tests/sql/yt_native_file/part6

YDBREQUESTS-3030

4 files changed, 81 insertions, 27 deletions

diff --git a/ydb/core/kqp/ut/opt/kqp_extract_predicate_unpack_ut.cpp b/ydb/core/kqp/ut/opt/kqp_extract_predicate_unpack_ut.cpp
index c1c1078932..a16eb76308 100644
--- a/ydb/core/kqp/ut/opt/kqp_extract_predicate_unpack_ut.cpp
+++ b/ydb/core/kqp/ut/opt/kqp_extract_predicate_unpack_ut.cpp
@@ -16,6 +16,13 @@ void PrepareTablesToUnpack(TSession session) {
             Value String,
             PRIMARY KEY (Key, Fk)
         );
+        CREATE TABLE `/Root/ComplexKeyNotNull` (
+            Key Int32 NOT NULL,
+            Fk Int32 NOT NULL,
+            Value String,
+            PRIMARY KEY (Key, Fk)
+        );
+
         CREATE TABLE `/Root/UintComplexKey` (
             Key UInt64,
             Fk Int64,
@@ -65,6 +72,14 @@ void PrepareTablesToUnpack(TSession session) {
             (4, 104, "Value2"),
             (5, 105, "Value3");
 
+        REPLACE INTO `/Root/ComplexKeyNotNull` (Key, Fk, Value) VALUES
+            (1, 101, "Value1"),
+            (2, 102, "Value1"),
+            (2, 103, "Value3"),
+            (3, 103, "Value2"),
+            (4, 104, "Value2"),
+            (5, 105, "Value3");
+
         REPLACE INTO `/Root/UintComplexKey` (Key, Fk, Value) VALUES
             (null, null, "NullValue"),
             (1, 101, "Value1"),
@@ -293,13 +308,24 @@ Y_UNIT_TEST(ComplexRange) {
     TestRange(
         R"(
             SELECT Key, Fk, Value FROM `/Root/ComplexKey`
+            WHERE (Key, Fk) >= (4, 104);
+        )",
+        R"([
+            [[4];[104];["Value2"]];
+            [[5];[105];["Value3"]]
+        ])",
+        2);
+
+    TestRange(
+        R"(
+            SELECT Key, Fk, Value FROM `/Root/ComplexKeyNotNull`
             WHERE (Key, Fk) >= (1, 101) AND (Key, Fk) < (4, 104);
         )",
         R"([
-            [[1];[101];["Value1"]];
-            [[2];[102];["Value1"]];
-            [[2];[103];["Value3"]];
-            [[3];[103];["Value2"]]
+            [1;101;["Value1"]];
+            [2;102;["Value1"]];
+            [2;103;["Value3"]];
+            [3;103;["Value2"]]
         ])",
         4);
 
@@ -314,6 +340,16 @@ Y_UNIT_TEST(ComplexRange) {
         1,
         2);
 
+    TestRange(
+        R"(
+            SELECT Key, Fk, Value FROM `/Root/ComplexKey`
+            WHERE (Key < 2) OR (Key = 2 AND Fk > 102)
+        )",
+        R"([
+            [[1];[101];["Value1"]];[[2];[103];["Value3"]]
+        ])",
+        2,
+        2);
 }
 
 Y_UNIT_TEST(SqlIn) {
diff --git a/ydb/library/yql/core/extract_predicate/extract_predicate_impl.cpp b/ydb/library/yql/core/extract_predicate/extract_predicate_impl.cpp
index c170edad57..5fc49d6cc4 100644
--- a/ydb/library/yql/core/extract_predicate/extract_predicate_impl.cpp
+++ b/ydb/library/yql/core/extract_predicate/extract_predicate_impl.cpp
@@ -430,8 +430,8 @@ TExprNode::TPtr ExpandTupleBinOp(const TExprNode& node, TExprContext& ctx) {
     if (node.IsCallable({"<=", ">="})) {
         return ctx.Builder(node.Pos())
             .Callable("Or")
-                .Add(0, ctx.RenameNode(node, "=="))
-                .Add(1, ctx.RenameNode(node, node.IsCallable("<=") ? "<" : ">"))
+                .Add(0, ctx.RenameNode(node, node.IsCallable("<=") ? "<" : ">"))
+                .Add(1, ctx.RenameNode(node, "=="))
             .Seal()
             .Build();
     }
@@ -1542,15 +1542,15 @@ TMaybe<TRangeBoundHint> CompareBounds(
             return Nothing();
         }
         if (auto cmp = TryCompareColumns(hint1.Columns[i], hint2.Columns[i])) {
-            if ((cmp < 0) == min) {
+            if (cmp == 0) {
+                continue;
+            } else if ((cmp < 0) == min) {
                 hint = hint1;
             } else {
                 hint = hint2;
             }
 
-            if (cmp != 0) {
-                break;
-            }
+            break;
         } else {
             return Nothing();
         }
@@ -1599,7 +1599,7 @@ TMaybe<TRangeHint> RangeHintExtend(const TMaybe<TRangeHint>& hint1, size_t hint1
     }
 }
 
-bool IsValid(const TRangeBoundHint& left, const TRangeBoundHint& right, bool acceptExclusivePoint = true) {
+bool IsValid(const TRangeBoundHint& left, const TRangeBoundHint& right) {
     for (size_t i = 0; ; ++i) {
         if (i >= left.Columns.size() || i >= right.Columns.size()) {
             // ok, we have +-inf and sure that it's valid
@@ -1608,15 +1608,12 @@ bool IsValid(const TRangeBoundHint& left, const TRangeBoundHint& right, bool acc
         auto cmp = TryCompareColumns(left.Columns[i], right.Columns[i]);
         if (!cmp) {
             return false;
-        } else {
-            if (*cmp < 0) {
-                return true;
-            } else if (*cmp > 0) {
-                return false;
-            }
+        } else if (*cmp < 0) {
+            return true;
+        } else if (*cmp > 0) {
+            return false;
         }
     }
-    return acceptExclusivePoint || left.Inclusive || right.Inclusive;
 }
 
 TMaybe<TRangeHint> RangeHintUnion(const TRangeHint& hint1, const TRangeHint& hint2) {
@@ -1630,11 +1627,33 @@ TMaybe<TRangeHint> RangeHintUnion(const TRangeHint& hint1, const TRangeHint& hin
     if (!left || !right || !intersection) {
         return Nothing();
     }
-    if (IsValid(intersection->Left, intersection->Right, false)) {
-        return TRangeHint{.Left = std::move(*left), .Right = std::move(*right)};
-    } else {
-        return Nothing();
+
+    { // check if there is no gap between ranges
+        for (size_t i = 0; ; ++i) {
+            bool leftFinished = i >= intersection->Left.Columns.size();
+            bool rightFinished = i >= intersection->Right.Columns.size();
+            if (leftFinished || rightFinished) {
+                if (leftFinished && intersection->Left.Inclusive) {
+                    break;
+                }
+                if (rightFinished && intersection->Right.Inclusive) {
+                    break;
+                }
+                return {};
+            }
+
+            auto cmp = TryCompareColumns(intersection->Left.Columns[i], intersection->Right.Columns[i]);
+            if (!cmp) {
+                return {};
+            } else if (*cmp < 0) {
+                break;
+            } else if (*cmp > 0) {
+                return {};
+            }
+        }
     }
+
+    return TRangeHint{.Left = std::move(*left), .Right = std::move(*right)};
 }
 
 TMaybe<TRangeHint> RangeHintUnion(const TMaybe<TRangeHint>& hint1, const TMaybe<TRangeHint>& hint2) {
diff --git a/ydb/library/yql/core/extract_predicate/ut/extract_predicate_ut.cpp b/ydb/library/yql/core/extract_predicate/ut/extract_predicate_ut.cpp
index 4b4b5021c6..39f6907c4e 100644
--- a/ydb/library/yql/core/extract_predicate/ut/extract_predicate_ut.cpp
+++ b/ydb/library/yql/core/extract_predicate/ut/extract_predicate_ut.cpp
@@ -523,8 +523,7 @@ Y_UNIT_TEST_SUITE(TYqlExtractPredicate) {
             "(let $2 (StructType '('\"x\" $1) '('\"y\" $1)))\n"
             "(let $3 (Int32 '1))\n"
             "(let $4 (Int32 '\"2\"))\n"
-            "(return (RangeOr (RangeAnd (Range $2 (lambda '($5) (== (Member $5 '\"x\") $3))) (Range $2 (lambda '($6) (== (Member $6 '\"y\") $4)))) (Range $2 (lambda '($7) (< (Member $7 '\"x\") $3))) (RangeAnd (Range $2 (lambda '($8) (== (Memb"
-            "er $8 '\"x\") $3))) (Range $2 (lambda '($9) (< (Member $9 '\"y\") $4))))))\n"
+            "(return (RangeOr (Range $2 (lambda '($5) (< (Member $5 '\"x\") $3))) (RangeAnd (Range $2 (lambda '($6) (== (Member $6 '\"x\") $3))) (Range $2 (lambda '($7) (< (Member $7 '\"y\") $4)))) (RangeAnd (Range $2 (lambda '($8) (== (Member $8 '\"x\") $3))) (Range $2 (lambda '($9) (== (Member $9 '\"y\") $4))))))\n"
             ")\n";
 
         TExprContext exprCtx;
diff --git a/ydb/library/yql/tests/sql/yt_native_file/part6/canondata/result.json b/ydb/library/yql/tests/sql/yt_native_file/part6/canondata/result.json
index a4d78e9530..74db01dec8 100644
--- a/ydb/library/yql/tests/sql/yt_native_file/part6/canondata/result.json
+++ b/ydb/library/yql/tests/sql/yt_native_file/part6/canondata/result.json
@@ -669,9 +669,9 @@
     ],
     "test.test[compute_range-tuples_compare-default.txt-Debug]": [
         {
-            "checksum": "962da2716dd68cffc53234286a48d7c7",
-            "size": 4596,
-            "uri": "https://{canondata_backend}/995452/3862ff9a9bced9f226a84891d10e8a5df01e3ec0/resource.tar.gz#test.test_compute_range-tuples_compare-default.txt-Debug_/opt.yql"
+            "checksum": "1214969b596fc6210776601e78bd0c24",
+            "size": 4603,
+            "uri": "https://{canondata_backend}/1784117/8b994472d7c14f553ca7722e774cd8684ad93d6a/resource.tar.gz#test.test_compute_range-tuples_compare-default.txt-Debug_/opt.yql"
         }
     ],
     "test.test[compute_range-tuples_compare-default.txt-Plan]": [