Merge #152300

152300: sql: use soft limit if available to decide on scan distribution r=yuzefovich a=yuzefovich

Fixes: #152295.

Release note (bug fix): In 25.1 we changed the physical planning heuristics so that large constrained table scans, estimated to scan at least 10k rows (controlled via `distribute_scan_row_count_threshold`), would force plan distribution with `distsql=auto` mode. However, if the scan has a "soft limit" we would still use the full estimate (e.g. in "estimated row count: 100 - 10,000" we would use 10k as the estimate) which can lead to distributing queries that shouldn't be, increasing the query latency. New session variable `use_soft_limit_for_distribute_scan` (default `false`) determines whether we use the soft limit when deciding that a scan is "large" or not.

Co-authored-by: Yahor Yuzefovich <[email protected]>

Author: craig[bot]

pkg/sql/distsql_physical_planner.go

@@ -717,17 +717,27 @@ func checkSupportForPlanNode(
 			// TODO(nvanbenschoten): lift this restriction.
 			return cannotDistribute, cannotDistributeRowLevelLockingErr
 		}
-
 		if n.localityOptimized {
 			// This is a locality optimized scan.
 			return cannotDistribute, localityOptimizedOpNotDistributableErr
 		}
-		// TODO(yuzefovich): consider using the soft limit in making a decision
-		// here.
 		scanRec := canDistribute
-		if n.estimatedRowCount != 0 && n.estimatedRowCount >= sd.DistributeScanRowCountThreshold {
-			log.VEventf(ctx, 2, "large scan recommends plan distribution")
-			scanRec = shouldDistribute
+		if n.estimatedRowCount != 0 {
+			var suffix string
+			estimate := n.estimatedRowCount
+			if n.softLimit != 0 && sd.UseSoftLimitForDistributeScan {
+				estimate = uint64(n.softLimit)
+				suffix = " (using soft limit)"
+			}
+			if estimate >= sd.DistributeScanRowCountThreshold {
+				log.VEventf(ctx, 2, "large scan recommends plan distribution%s", suffix)
+				scanRec = shouldDistribute
+			} else if n.softLimit != 0 && n.estimatedRowCount >= sd.DistributeScanRowCountThreshold {
+				log.VEventf(
+					ctx, 2, `estimated row count would consider the scan "large" `+
+						`while soft limit hint makes it "small"`,
+				)
+			}
 		}
 		if n.isFull && (n.estimatedRowCount == 0 || sd.AlwaysDistributeFullScans) {
 			// In the absence of table stats, we default to always distributing

pkg/sql/exec_util.go

@@ -3645,6 +3645,10 @@ func (m *sessionDataMutator) SetAlwaysDistributeFullScans(val bool) {
 	m.data.AlwaysDistributeFullScans = val
 }
 
+func (m *sessionDataMutator) SetUseSoftLimitForDistributeScan(val bool) {
+	m.data.UseSoftLimitForDistributeScan = val
+}
+
 func (m *sessionDataMutator) SetDistributeJoinRowCountThreshold(val uint64) {
 	m.data.DistributeJoinRowCountThreshold = val
 }

pkg/sql/logictest/testdata/logic_test/information_schema

@@ -4162,6 +4162,7 @@ use_cputs_on_non_unique_indexes                                  off
 use_improved_routine_dependency_tracking                         on
 use_pre_25_2_variadic_builtins                                   off
 use_proc_txn_control_extended_protocol_fix                       on
+use_soft_limit_for_distribute_scan                               off
 variable_inequality_lookup_join_enabled                          on
 vector_search_beam_size                                          32
 vector_search_rerank_multiplier                                  50

pkg/sql/logictest/testdata/logic_test/pg_catalog

@@ -3176,6 +3176,7 @@ use_declarative_schema_changer                                   on
 use_improved_routine_dependency_tracking                         on                  NULL      NULL        NULL        string
 use_pre_25_2_variadic_builtins                                   off                 NULL      NULL        NULL        string
 use_proc_txn_control_extended_protocol_fix                       on                  NULL      NULL        NULL        string
+use_soft_limit_for_distribute_scan                               off                 NULL      NULL        NULL        string
 variable_inequality_lookup_join_enabled                          on                  NULL      NULL        NULL        string
 vector_search_beam_size                                          32                  NULL      NULL        NULL        string
 vector_search_rerank_multiplier                                  50                  NULL      NULL        NULL        string
@@ -3418,6 +3419,7 @@ use_declarative_schema_changer                                   on
 use_improved_routine_dependency_tracking                         on                  NULL  user     NULL      on                  on
 use_pre_25_2_variadic_builtins                                   off                 NULL  user     NULL      off                 off
 use_proc_txn_control_extended_protocol_fix                       on                  NULL  user     NULL      on                  on
+use_soft_limit_for_distribute_scan                               off                 NULL  user     NULL      off                 off
 variable_inequality_lookup_join_enabled                          on                  NULL  user     NULL      on                  on
 vector_search_beam_size                                          32                  NULL  user     NULL      32                  32
 vector_search_rerank_multiplier                                  50                  NULL  user     NULL      50                  50
@@ -3652,6 +3654,7 @@ use_declarative_schema_changer                                   NULL    NULL
 use_improved_routine_dependency_tracking                         NULL    NULL     NULL     NULL        NULL
 use_pre_25_2_variadic_builtins                                   NULL    NULL     NULL     NULL        NULL
 use_proc_txn_control_extended_protocol_fix                       NULL    NULL     NULL     NULL        NULL
+use_soft_limit_for_distribute_scan                               NULL    NULL     NULL     NULL        NULL
 variable_inequality_lookup_join_enabled                          NULL    NULL     NULL     NULL        NULL
 vector_search_beam_size                                          NULL    NULL     NULL     NULL        NULL
 vector_search_rerank_multiplier                                  NULL    NULL     NULL     NULL        NULL

pkg/sql/logictest/testdata/logic_test/show_source

@@ -253,6 +253,7 @@ use_declarative_schema_changer                                   on
 use_improved_routine_dependency_tracking                         on
 use_pre_25_2_variadic_builtins                                   off
 use_proc_txn_control_extended_protocol_fix                       on
+use_soft_limit_for_distribute_scan                               off
 variable_inequality_lookup_join_enabled                          on
 vector_search_beam_size                                          32
 vector_search_rerank_multiplier                                  50

pkg/sql/opt/exec/execbuilder/testdata/distsql_auto_mode

@@ -280,3 +280,83 @@ query T
 SELECT info FROM [EXPLAIN SELECT a FROM abc INNER LOOKUP JOIN kv ON b = k] WHERE info LIKE 'distribution%'
 ----
 distribution: full
+
+subtest regression_152295
+
+statement ok
+CREATE TABLE a (
+  i INT PRIMARY KEY,
+  j INT
+)
+
+statement ok
+CREATE TABLE b (
+  k INT PRIMARY KEY
+)
+
+statement ok
+ALTER TABLE a SPLIT AT SELECT i FROM generate_series(1, 9) AS g(i)
+
+retry
+statement ok
+ALTER TABLE a EXPERIMENTAL_RELOCATE SELECT ARRAY[i%5+1], i FROM generate_series(0, 9) AS g(i)
+
+statement ok
+ALTER TABLE b SPLIT AT SELECT i FROM generate_series(1, 9) AS g(i)
+
+retry
+statement ok
+ALTER TABLE b EXPERIMENTAL_RELOCATE SELECT ARRAY[i%5+1], i FROM generate_series(0, 9) AS g(i)
+
+statement ok
+ALTER TABLE a INJECT STATISTICS '[
+  {
+    "columns": ["i"],
+    "created_at": "2018-01-01 1:00:00.00000+00:00",
+    "row_count": 100000,
+    "distinct_count": 100000
+  }
+]'
+
+statement ok
+SET use_soft_limit_for_distribute_scan = true
+
+# We choose to not distribute this query since the constrained scan has a soft
+# limit hint of 100 that is below the distribute scan threshold of 10k (even
+# though the "estimated row count" is 33,334).
+query T retry
+SELECT info FROM [EXPLAIN SELECT * FROM a INNER LOOKUP JOIN b ON k = j AND i < 10000 LIMIT 1]
+  WHERE info LIKE 'distribution%' OR info LIKE '%estimated row count%'
+----
+distribution: local
+          estimated row count: 100 - 33,334 (33% of the table; stats collected <hidden> ago)
+
+statement ok
+SET distribute_scan_row_count_threshold = 10
+
+# But now the soft limit hint exceeds the threshold - we should distribute.
+query T
+SELECT info FROM [EXPLAIN SELECT * FROM a INNER LOOKUP JOIN b ON k = j AND i < 10000 LIMIT 1]
+  WHERE info LIKE 'distribution%' OR info LIKE '%estimated row count%'
+----
+distribution: full
+          estimated row count: 100 - 33,334 (33% of the table; stats collected <hidden> ago)
+
+statement ok
+RESET distribute_scan_row_count_threshold
+
+statement ok
+SET use_soft_limit_for_distribute_scan = false
+
+# Now we don't look at the soft limit hint - we should distribute.
+query T
+SELECT info FROM [EXPLAIN SELECT * FROM a INNER LOOKUP JOIN b ON k = j AND i < 10000 LIMIT 1]
+  WHERE info LIKE 'distribution%' OR info LIKE '%estimated row count%'
+----
+distribution: full
+          estimated row count: 100 - 33,334 (33% of the table; stats collected <hidden> ago)
+
+statement ok
+RESET use_soft_limit_for_distribute_scan;
+
+subtest end

pkg/sql/sessiondatapb/local_only_session_data.proto

@@ -732,6 +732,10 @@ message LocalOnlySessionData {
   // during query optimization. The names of the rules can be found in the
   // opt/rule_name.og.go file.
   repeated string disable_optimizer_rules = 184;
+  // UseSoftLimitForDistributeScan, if set, means that we'll use the soft limit
+  // hint - if available - when comparing against
+  // DistributeScanRowCountThreshold.
+  bool use_soft_limit_for_distribute_scan = 185;
 
   ///////////////////////////////////////////////////////////////////////////
   // WARNING: consider whether a session parameter you're adding needs to  //

pkg/sql/vars.go

@@ -737,6 +737,23 @@ var varGen = map[string]sessionVar{
 		GlobalDefault: globalFalse,
 	},
 
+	// CockroachDB extension.
+	`use_soft_limit_for_distribute_scan`: {
+		GetStringVal: makePostgresBoolGetStringValFn(`use_soft_limit_for_distribute_scan`),
+		Set: func(_ context.Context, m sessionDataMutator, s string) error {
+			b, err := paramparse.ParseBoolVar("use_soft_limit_for_distribute_scan", s)
+			if err != nil {
+				return err
+			}
+			m.SetUseSoftLimitForDistributeScan(b)
+			return nil
+		},
+		Get: func(evalCtx *extendedEvalContext, _ *kv.Txn) (string, error) {
+			return formatBoolAsPostgresSetting(evalCtx.SessionData().UseSoftLimitForDistributeScan), nil
+		},
+		GlobalDefault: globalFalse,
+	},
+
 	// CockroachDB extension.
 	`distribute_join_row_count_threshold`: {
 		Get: func(evalCtx *extendedEvalContext, _ *kv.Txn) (string, error) {