cherry pick of 16153 (#5384)

go/vt/vtexplain/testdata/multi-output/selectsharded-output.txt

@@ -128,7 +128,7 @@ select name from user where id not in (select id from t1) /* non-correlated subq
 ----------------------------------------------------------------------
 select name from user where exists (select id from t1) /* non-correlated subquery as EXISTS */
 
-1 ks_unsharded/-: select 1 from t1 limit 10001 /* non-correlated subquery as EXISTS */
+1 ks_unsharded/-: select 1 from t1 limit 1 /* non-correlated subquery as EXISTS */
 2 ks_sharded/-40: select `name` from `user` where 1 limit 10001 /* non-correlated subquery as EXISTS */
 2 ks_sharded/40-80: select `name` from `user` where 1 limit 10001 /* non-correlated subquery as EXISTS */
 2 ks_sharded/80-c0: select `name` from `user` where 1 limit 10001 /* non-correlated subquery as EXISTS */

go/vt/vtgate/planbuilder/operators/query_planning.go

@@ -19,6 +19,7 @@ package operators
 import (
  "fmt"
  "io"
+ "strconv"
 
  "vitess.io/vitess/go/vt/sqlparser"
  "vitess.io/vitess/go/vt/vterrors"
@@ -244,8 +245,170 @@ func tryPushLimit(ctx *plancontext.PlanningContext, in *Limit) (Operator, *Apply
  }
 
  return src, Rewrote(fmt.Sprintf("push limit to %s of apply join", side))
+ case *Limit:
+ combinedLimit := mergeLimits(in.AST, src.AST)
+ if combinedLimit == nil {
+ break
+ }
+ // we can remove the other LIMIT
+ in.AST = combinedLimit
+ in.Source = src.Source
+ return in, Rewrote("merged two limits")
+
+ }
+ return setUpperLimit(in)
+}
+
+func mergeLimits(l1, l2 *sqlparser.Limit) *sqlparser.Limit {
+ // To merge two relational LIMIT operators with LIMIT and OFFSET, we need to combine their
+ // LIMIT and OFFSET values appropriately.
+ // Let's denote the first LIMIT operator as LIMIT_1 with LIMIT_1 and OFFSET_1,
+ // and the second LIMIT operator as LIMIT_2 with LIMIT_2 and OFFSET_2.
+ // The second LIMIT operator receives the output of the first LIMIT operator, meaning the first LIMIT and
+ // OFFSET are applied first, and then the second LIMIT and OFFSET are applied to the resulting subset.
+ //
+ // The goal is to determine the effective combined LIMIT and OFFSET values when applying these two operators sequentially.
+ //
+ // Combined Offset:
+ // The combined offset (OFFSET_combined) is the sum of the two offsets because you need to skip OFFSET_1 rows first,
+ // and then apply the second offset OFFSET_2 to the result.
+ // OFFSET_combined = OFFSET_1 + OFFSET_2
+
+ // Combined Limit:
+ // The combined limit (LIMIT_combined) needs to account for both limits. The effective limit should not exceed the rows returned by the first limit,
+ // so it is the minimum of the remaining rows after the first offset and the second limit.
+ // LIMIT_combined = min(LIMIT_2, LIMIT_1 - OFFSET_2)
+
+ // Note: If LIMIT_1 - OFFSET_2 is negative or zero, it means there are no rows left to limit, so LIMIT_combined should be zero.
+
+ // Example:
+ // First LIMIT operator: LIMIT 10 OFFSET 5 (LIMIT_1 = 10, OFFSET_1 = 5)
+ // Second LIMIT operator: LIMIT 7 OFFSET 3 (LIMIT_2 = 7, OFFSET_2 = 3)
+
+ // Calculations:
+ // Combined OFFSET:
+ // OFFSET_combined = 5 + 3 = 8
+
+ // Combined LIMIT:
+ // remaining rows after OFFSET_2 = 10 - 3 = 7
+ // LIMIT_combined = min(7, 7) = 7
+
+ // So, the combined result would be:
+ // LIMIT 7 OFFSET 8
+
+ // This method ensures that the final combined LIMIT and OFFSET correctly reflect the sequential application of the two original operators.
+ combinedLimit, failed := mergeLimitExpressions(l1.Rowcount, l2.Rowcount, l2.Offset)
+ if failed {
+ return nil
+ }
+ combinedOffset, failed := mergeOffsetExpressions(l1.Offset, l2.Offset)
+ if failed {
+ return nil
+ }
+
+ return &sqlparser.Limit{
+ Offset: combinedOffset,
+ Rowcount: combinedLimit,
+ }
+}
+
+func mergeOffsetExpressions(e1, e2 sqlparser.Expr) (expr sqlparser.Expr, failed bool) {
+ switch {
+ case e1 == nil && e2 == nil:
+ return nil, false
+ case e1 == nil:
+ return e2, false
+ case e2 == nil:
+ return e1, false
  default:
- return setUpperLimit(in)
+ v1str, ok := e1.(*sqlparser.Literal)
+ if !ok {
+ return nil, true
+ }
+ v2str, ok := e2.(*sqlparser.Literal)
+ if !ok {
+ return nil, true
+ }
+ v1, _ := strconv.Atoi(v1str.Val)
+ v2, _ := strconv.Atoi(v2str.Val)
+ return sqlparser.NewIntLiteral(strconv.Itoa(v1 + v2)), false
+ }
+}
+
+// mergeLimitExpressions merges two LIMIT expressions with an OFFSET expression.
+// l1: First LIMIT expression.
+// l2: Second LIMIT expression.
+// off2: Second OFFSET expression.
+// Returns the merged LIMIT expression and a boolean indicating if the merge failed.
+func mergeLimitExpressions(l1, l2, off2 sqlparser.Expr) (expr sqlparser.Expr, failed bool) {
+ switch {
+ // If both limits are nil, there's nothing to merge, return nil without failure.
+ case l1 == nil && l2 == nil:
+ return nil, false
+
+ // If the first limit is nil, the second limit determines the final limit.
+ case l1 == nil:
+ return l2, false
+
+ // If the second limit is nil, calculate the remaining limit after applying the offset to the first limit.
+ case l2 == nil:
+ if off2 == nil {
+ // No offset, so the first limit is used directly.
+ return l1, false
+ }
+ off2, ok := off2.(*sqlparser.Literal)
+ if !ok {
+ // If the offset is not a literal, fail the merge.
+ return nil, true
+ }
+ lim1str, ok := l1.(*sqlparser.Literal)
+ if !ok {
+ // If the first limit is not a literal, return the first limit without failing.
+ return nil, false
+ }
+ // Calculate the remaining limit after the offset.
+ off2int, _ := strconv.Atoi(off2.Val)
+ l1int, _ := strconv.Atoi(lim1str.Val)
+ lim := l1int - off2int
+ if lim < 0 {
+ lim = 0
+ }
+ return sqlparser.NewIntLiteral(strconv.Itoa(lim)), false
+
+ default:
+ v1str, ok1 := l1.(*sqlparser.Literal)
+ if ok1 && v1str.Val == "1" {
+ // If the first limit is "1", it dominates, so return it.
+ return l1, false
+ }
+ v2str, ok2 := l2.(*sqlparser.Literal)
+ if ok2 && v2str.Val == "1" {
+ // If the second limit is "1", it dominates, so return it.
+ return l2, false
+ }
+ if !ok1 || !ok2 {
+ // If either limit is not a literal, fail the merge.
+ return nil, true
+ }
+
+ var off2int int
+ if off2 != nil {
+ off2, ok := off2.(*sqlparser.Literal)
+ if !ok {
+ // If the offset is not a literal, fail the merge.
+ return nil, true
+ }
+ off2int, _ = strconv.Atoi(off2.Val)
+ }
+
+ v1, _ := strconv.Atoi(v1str.Val)
+ v2, _ := strconv.Atoi(v2str.Val)
+ lim := min(v2, v1-off2int)
+ if lim < 0 {
+ // If the combined limit is negative, set it to zero.
+ lim = 0
+ }
+ return sqlparser.NewIntLiteral(strconv.Itoa(lim)), false
  }
 }
 

go/vt/vtgate/planbuilder/operators/query_planning_test.go

@@ -0,0 +1,185 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed 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.
+*/
+
+package operators
+
+import (
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+
+ "vitess.io/vitess/go/vt/sqlparser"
+)
+
+func TestMergeOffsetExpressions(t *testing.T) {
+ tests := []struct {
+ name string
+ offset1 sqlparser.Expr
+ offset2 sqlparser.Expr
+ expectedExpr sqlparser.Expr
+ expectedFailed bool
+ }{
+ {
+ name: "both offsets are integers",
+ offset1: sqlparser.NewIntLiteral("5"),
+ offset2: sqlparser.NewIntLiteral("3"),
+ expectedExpr: sqlparser.NewIntLiteral("8"),
+ expectedFailed: false,
+ },
+ {
+ name: "first offset is nil",
+ offset1: nil,
+ offset2: sqlparser.NewIntLiteral("3"),
+ expectedExpr: sqlparser.NewIntLiteral("3"),
+ expectedFailed: false,
+ },
+ {
+ name: "second offset is nil",
+ offset1: sqlparser.NewIntLiteral("5"),
+ offset2: nil,
+ expectedExpr: sqlparser.NewIntLiteral("5"),
+ expectedFailed: false,
+ },
+ {
+ name: "both offsets are nil",
+ offset1: nil,
+ offset2: nil,
+ expectedExpr: nil,
+ expectedFailed: false,
+ },
+ {
+ name: "first offset is argument",
+ offset1: sqlparser.NewArgument("offset1"),
+ offset2: sqlparser.NewIntLiteral("3"),
+ expectedExpr: nil,
+ expectedFailed: true,
+ },
+ {
+ name: "second offset is argument",
+ offset1: sqlparser.NewIntLiteral("5"),
+ offset2: sqlparser.NewArgument("offset2"),
+ expectedExpr: nil,
+ expectedFailed: true,
+ },
+ {
+ name: "both offsets are arguments",
+ offset1: sqlparser.NewArgument("offset1"),
+ offset2: sqlparser.NewArgument("offset2"),
+ expectedExpr: nil,
+ expectedFailed: true,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ expr, failed := mergeOffsetExpressions(tt.offset1, tt.offset2)
+ assert.Equal(t, tt.expectedExpr, expr)
+ assert.Equal(t, tt.expectedFailed, failed, "failed")
+ })
+ }
+}
+
+func TestMergeLimitExpressions(t *testing.T) {
+ tests := []struct {
+ name string
+ limit1 sqlparser.Expr
+ limit2 sqlparser.Expr
+ offset2 sqlparser.Expr
+ expectedExpr sqlparser.Expr
+ expectedFailed bool
+ }{
+ {
+ name: "valid limits and offset",
+ limit1: sqlparser.NewIntLiteral("10"),
+ limit2: sqlparser.NewIntLiteral("7"),
+ offset2: sqlparser.NewIntLiteral("3"),
+ expectedExpr: sqlparser.NewIntLiteral("7"),
+ expectedFailed: false,
+ },
+ {
+ name: "remaining rows after offset2 is zero",
+ limit1: sqlparser.NewIntLiteral("3"),
+ limit2: sqlparser.NewIntLiteral("7"),
+ offset2: sqlparser.NewIntLiteral("5"),
+ expectedExpr: sqlparser.NewIntLiteral("0"),
+ expectedFailed: false,
+ },
+ {
+ name: "first limit is nil",
+ limit1: nil,
+ limit2: sqlparser.NewIntLiteral("7"),
+ offset2: sqlparser.NewIntLiteral("3"),
+ expectedExpr: sqlparser.NewIntLiteral("7"),
+ expectedFailed: false,
+ },
+ {
+ name: "second limit is nil",
+ limit1: sqlparser.NewIntLiteral("10"),
+ limit2: nil,
+ offset2: sqlparser.NewIntLiteral("3"),
+ expectedExpr: sqlparser.NewIntLiteral("7"),
+ expectedFailed: false,
+ },
+ {
+ name: "offset2 is nil",
+ limit1: sqlparser.NewIntLiteral("10"),
+ limit2: sqlparser.NewIntLiteral("7"),
+ offset2: nil,
+ expectedExpr: sqlparser.NewIntLiteral("7"),
+ expectedFailed: false,
+ },
+ {
+ name: "first limit is argument",
+ limit1: sqlparser.NewArgument("limit1"),
+ limit2: sqlparser.NewIntLiteral("7"),
+ offset2: sqlparser.NewIntLiteral("3"),
+ expectedExpr: nil,
+ expectedFailed: true,
+ },
+ {
+ name: "second limit is argument",
+ limit1: sqlparser.NewIntLiteral("10"),
+ limit2: sqlparser.NewArgument("limit2"),
+ offset2: sqlparser.NewIntLiteral("3"),
+ expectedExpr: nil,
+ expectedFailed: true,
+ },
+ {
+ name: "offset2 is argument",
+ limit1: sqlparser.NewIntLiteral("10"),
+ limit2: sqlparser.NewIntLiteral("7"),
+ offset2: sqlparser.NewArgument("offset2"),
+ expectedExpr: nil,
+ expectedFailed: true,
+ },
+ {
+ name: "all are arguments",
+ limit1: sqlparser.NewArgument("limit1"),
+ limit2: sqlparser.NewArgument("limit2"),
+ offset2: sqlparser.NewArgument("offset2"),
+ expectedExpr: nil,
+ expectedFailed: true,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ expr, failed := mergeLimitExpressions(tt.limit1, tt.limit2, tt.offset2)
+ assert.Equal(t, tt.expectedExpr, expr)
+ assert.Equal(t, tt.expectedFailed, failed, "failed")
+ })
+ }
+}