Merge #144703

144703: logical: improve crud writer sql queries r=jeffswenson a=jeffswenson

This change includes two improvements to the SQL queries used by the
crud writer:
1. Where clauses now use '=' when comparing to the primary key. This
   works around an issue with generic query plans in the optimizer that
   caused `INSERT` and `UPDATE` to plan full table scans.
2. The read refresh statement includes computed columns in the primary
   key lookup. This was needed to support computed virtual and stored
   columns in the crud writer. Virtual columns in the primary key are
   still prohibited by schema validation because the classic SQL writer
   does not support them.

Epic: [CRDB-48647](https://cockroachlabs.atlassian.net/browse/CRDB-48647)
Release note: none

Co-authored-by: Jeff Swenson <[email protected]>

Author: craig[bot]

pkg/crosscluster/logical/BUILD.bazel

@@ -173,6 +173,7 @@ go_test(
         "//pkg/sql/execinfra",
         "//pkg/sql/execinfrapb",
         "//pkg/sql/isql",
+        "//pkg/sql/parser/statements",
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/randgen",
@@ -184,6 +185,7 @@ go_test(
         "//pkg/sql/sqlclustersettings",
         "//pkg/sql/sqltestutils",
         "//pkg/sql/stats",
+        "//pkg/sql/types",
         "//pkg/testutils",
         "//pkg/testutils/datapathutils",
         "//pkg/testutils/jobutils",

pkg/crosscluster/logical/event_decoder.go

@@ -76,7 +76,7 @@ func newEventDecoder(
 				return err
 			}
 
-			columns := getPhysicalColumnsSchema(descriptor)
+			columns := getColumnSchema(descriptor)
 			columnNames := make([]string, 0, len(columns))
 			for _, column := range columns {
 				columnNames = append(columnNames, column.column.GetName())

pkg/crosscluster/logical/logical_replication_job_test.go

@@ -196,7 +196,7 @@ func TestOptimsitcInsertCorruption(t *testing.T) {
 
 	dbSource.Exec(t, "SET CLUSTER SETTING logical_replication.consumer.try_optimistic_insert.enabled = true")
 
-	createTableStmt := `CREATE TABLE computed_cols (
+	createTableStmt := `CREATE TABLE indexed (
 		a INT,
 		b INT,
 		c INT,
@@ -205,35 +205,35 @@ func TestOptimsitcInsertCorruption(t *testing.T) {
 	dbSource.Exec(t, createTableStmt)
 	dbDest.Exec(t, createTableStmt)
 
-	createIdxStmt := `CREATE INDEX c ON computed_cols (c)`
+	createIdxStmt := `CREATE INDEX c ON indexed (c)`
 	dbSource.Exec(t, createIdxStmt)
 	dbDest.Exec(t, createIdxStmt)
 
 	// Insert initial data into destination that should be overwritten. This
 	// gives more opportunities for corruption.
-	dbDest.Exec(t, "INSERT INTO computed_cols (a, b, c) VALUES (1, 2, 3), (3, 4, 5)")
-	dbSource.Exec(t, "INSERT INTO computed_cols (a, b, c) VALUES (1, 2, 6), (3, 4, 7)")
+	dbDest.Exec(t, "INSERT INTO indexed (a, b, c) VALUES (1, 2, 3), (3, 4, 5)")
+	dbSource.Exec(t, "INSERT INTO indexed (a, b, c) VALUES (1, 2, 6), (3, 4, 7)")
 
 	// Create logical replication stream from source to destination
 	sourceURL := replicationtestutils.GetExternalConnectionURI(t, s, s, serverutils.DBName("a"))
 	var jobID jobspb.JobID
 	dbDest.QueryRow(t,
-		"CREATE LOGICAL REPLICATION STREAM FROM TABLE computed_cols ON $1 INTO TABLE computed_cols",
+		"CREATE LOGICAL REPLICATION STREAM FROM TABLE indexed ON $1 INTO TABLE indexed",
 		sourceURL.String(),
 	).Scan(&jobID)
 
 	WaitUntilReplicatedTime(t, s.Clock().Now(), dbDest, jobID)
 
-	dbSource.CheckQueryResults(t, "SELECT * FROM computed_cols", [][]string{
+	dbSource.CheckQueryResults(t, "SELECT * FROM indexed", [][]string{
 		{"1", "2", "6"},
 		{"3", "4", "7"},
 	})
-	dbDest.CheckQueryResults(t, "SELECT * FROM computed_cols", [][]string{
+	dbDest.CheckQueryResults(t, "SELECT * FROM indexed", [][]string{
 		{"1", "2", "6"},
 		{"3", "4", "7"},
 	})
 
-	compareReplicatedTables(t, s, "a", "b", "computed_cols", dbSource, dbDest)
+	compareReplicatedTables(t, s, "a", "b", "indexed", dbSource, dbDest)
 }
 
 type fatalDLQ struct{ *testing.T }
@@ -1473,6 +1473,64 @@ func TestTombstoneUpdate(t *testing.T) {
 	dst.CheckQueryResults(t, "SELECT * FROM tab WHERE pk = 1", [][]string{})
 }
 
+func TestReplicateComputedColumns(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	skip.UnderDeadlock(t)
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+
+	server, s, dbSource, dbDest := setupLogicalTestServer(t, ctx, testClusterBaseClusterArgs, 1)
+	defer server.Stopper().Stop(ctx)
+
+	// Note: the primary key is basicaly (a, b), but with a computed column that
+	// is d=a+b.
+	createTableStmt := `CREATE TABLE computed_cols (
+		a INT,
+		b INT,
+		c INT,
+		d INT AS (a + b) STORED,
+		cIsOdd BOOLEAN AS (c % 2 = 1) VIRTUAL,
+		PRIMARY KEY (a, d)
+	)`
+	dbSource.Exec(t, createTableStmt)
+	dbDest.Exec(t, createTableStmt)
+
+	var writer string
+	dbSource.QueryRow(t, "SHOW CLUSTER SETTING logical_replication.consumer.immediate_mode_writer").Scan(&writer)
+	if writer != "legacy-kv" {
+		createIdxStmt := `CREATE INDEX c ON computed_cols (c) WHERE cIsOdd`
+		dbSource.Exec(t, createIdxStmt)
+		dbDest.Exec(t, createIdxStmt)
+	}
+
+	// Insert initial data into destination that should be overwritten. This
+	// tests the conflict case which is more error prone.
+	dbDest.Exec(t, "INSERT INTO computed_cols (a, b, c) VALUES (1, 2, 3), (3, 4, 5)")
+	dbSource.Exec(t, "INSERT INTO computed_cols (a, b, c) VALUES (1, 2, 6), (3, 4, 7)")
+
+	// Create logical replication stream from source to destination
+	sourceURL := replicationtestutils.GetExternalConnectionURI(t, s, s, serverutils.DBName("a"))
+	var jobID jobspb.JobID
+	dbDest.QueryRow(t,
+		"CREATE LOGICAL REPLICATION STREAM FROM TABLE computed_cols ON $1 INTO TABLE computed_cols",
+		sourceURL.String(),
+	).Scan(&jobID)
+
+	WaitUntilReplicatedTime(t, s.Clock().Now(), dbDest, jobID)
+
+	dbSource.CheckQueryResults(t, "SELECT * FROM computed_cols", [][]string{
+		{"1", "2", "6", "3", "false"},
+		{"3", "4", "7", "7", "true"},
+	})
+	dbDest.CheckQueryResults(t, "SELECT * FROM computed_cols", [][]string{
+		{"1", "2", "6", "3", "false"},
+		{"3", "4", "7", "7", "true"},
+	})
+
+	compareReplicatedTables(t, s, "a", "b", "computed_cols", dbSource, dbDest)
+}
+
 func TestForeignKeyConstraints(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	skip.UnderDeadlock(t)

pkg/crosscluster/logical/replication_statements.go

@@ -20,40 +20,50 @@ import (
 type columnSchema struct {
 	column       catalog.Column
 	isPrimaryKey bool
+	isComputed   bool
 }
 
-// getPhysicalColumns returns the list of columns that are part of the table's
-// primary key and value.
-func getPhysicalColumns(table catalog.TableDescriptor) []catalog.Column {
-	columns := table.AllColumns()
-	result := make([]catalog.Column, 0, len(columns))
-	for _, col := range columns {
-		if !col.IsComputed() && !col.IsVirtual() && !col.IsSystemColumn() {
-			result = append(result, col)
-		}
-	}
-	return result
-}
-
-func getPhysicalColumnsSchema(table catalog.TableDescriptor) []columnSchema {
-	columns := getPhysicalColumns(table)
-	primaryIdx := table.GetPrimaryIndex()
-
+// getColumnSchema returns the list of all columns that is decoded by the CRUD
+// writer. It returns columns in column key order. The crud writer passes
+// around a tree.Datums for each row where column[i] is the column definition
+// for datums[i].
+//
+// A column is decoded by crud LDR writer if either are true:
+//  1. The column is part of the primary key. Every primary key column is needed
+//     to perform read refreshes.
+//  2. The column is not a computed column. If a column is not computed, it must
+//     be included in update and insert statements.
+//
+// Notably, this excludes computed columns that are not part of the primary key
+// and system columns like crdb_internal_mvcc_timestamp.
+func getColumnSchema(table catalog.TableDescriptor) []columnSchema {
 	// Create a map of column ID to column for fast lookup
+	primaryIdx := table.GetPrimaryIndex()
 	isPrimaryKey := make(map[catid.ColumnID]bool)
 	for _, col := range primaryIdx.IndexDesc().KeyColumnIDs {
 		isPrimaryKey[col] = true
 	}
 
-	cols := make([]columnSchema, 0, len(columns))
+	columns := table.AllColumns()
+	result := make([]columnSchema, 0, len(columns))
 	for _, col := range columns {
-		cols = append(cols, columnSchema{
+		if col.IsSystemColumn() {
+			continue
+		}
+
+		isComputed := col.IsComputed()
+		if isComputed && !isPrimaryKey[col.GetID()] {
+			continue
+		}
+
+		result = append(result, columnSchema{
 			column:       col,
 			isPrimaryKey: isPrimaryKey[col.GetID()],
+			isComputed:   isComputed,
 		})
 	}
 
-	return cols
+	return result
 }
 
 // newTypedPlaceholder creates a placeholder with the appropriate type for a column.
@@ -77,20 +87,26 @@ func newTypedPlaceholder(idx int, col catalog.Column) (*tree.CastExpr, error) {
 func newInsertStatement(
 	table catalog.TableDescriptor,
 ) (statements.Statement[tree.Statement], error) {
-	columns := getPhysicalColumns(table)
+	columns := getColumnSchema(table)
 
-	columnNames := make(tree.NameList, len(columns))
+	columnNames := make(tree.NameList, 0, len(columns))
+	parameters := make(tree.Exprs, 0, len(columns))
 	for i, col := range columns {
-		columnNames[i] = tree.Name(col.GetName())
-	}
+		// NOTE: this consumes a placholder ID because its part of the tree.Datums,
+		// but it doesn't show up in the query because computed columns are not
+		// needed for insert statements.
+		if col.isComputed {
+			continue
+		}
 
-	parameters := make(tree.Exprs, len(columnNames))
-	for i, col := range columns {
 		var err error
-		parameters[i], err = newTypedPlaceholder(i+1, col)
+		parameter, err := newTypedPlaceholder(i+1, col.column)
 		if err != nil {
 			return statements.Statement[tree.Statement]{}, err
 		}
+
+		columnNames = append(columnNames, tree.Name(col.column.GetName()))
+		parameters = append(parameters, parameter)
 	}
 
 	parameterValues := &tree.ValuesClause{
@@ -119,16 +135,31 @@ func newInsertStatement(
 // newMatchesLastRow creates a WHERE clause for matching all columns of a row.
 // It returns a tree.Expr that compares each column to a placeholder parameter.
 // Parameters are ordered by column ID, starting from startParamIdx.
-func newMatchesLastRow(columns []catalog.Column, startParamIdx int) (tree.Expr, error) {
+func newMatchesLastRow(columns []columnSchema, startParamIdx int) (tree.Expr, error) {
 	var whereClause tree.Expr
 	for i, col := range columns {
-		placeholder, err := newTypedPlaceholder(startParamIdx+i, col)
+		if col.isComputed {
+			// Skip computed columns since they are not needed to fully specify the
+			// contents of the row.
+			continue
+		}
+
+		placeholder, err := newTypedPlaceholder(startParamIdx+i, col.column)
 		if err != nil {
 			return nil, err
 		}
+
+		eq := treecmp.MakeComparisonOperator(treecmp.IsNotDistinctFrom)
+		if col.isPrimaryKey {
+			// Generic query planner plans full table scans if we use `IS NOT
+			// DISTINCT FROM`. Use `=` operator for primary key columns since they
+			// are guaranteed to be non-NULL.
+			eq = treecmp.MakeComparisonOperator(treecmp.EQ)
+		}
+
 		colExpr := &tree.ComparisonExpr{
-			Operator: treecmp.MakeComparisonOperator(treecmp.IsNotDistinctFrom),
-			Left:     &tree.ColumnItem{ColumnName: tree.Name(col.GetName())},
+			Operator: eq,
+			Left:     &tree.ColumnItem{ColumnName: tree.Name(col.column.GetName())},
 			Right:    placeholder,
 		}
 
@@ -153,31 +184,37 @@ func newMatchesLastRow(columns []catalog.Column, startParamIdx int) (tree.Expr,
 func newUpdateStatement(
 	table catalog.TableDescriptor,
 ) (statements.Statement[tree.Statement], error) {
-	columns := getPhysicalColumns(table)
+	columns := getColumnSchema(table)
 
 	// Create WHERE clause for matching the previous row values
 	whereClause, err := newMatchesLastRow(columns, 1)
 	if err != nil {
 		return statements.Statement[tree.Statement]{}, err
 	}
 
-	exprs := make(tree.UpdateExprs, len(columns))
+	exprs := make(tree.UpdateExprs, 0, len(columns))
 	for i, col := range columns {
-		nameNode := tree.Name(col.GetName())
+		if col.isComputed {
+			// Skip computed columns since they are not needed to fully specify the
+			// contents of the row.
+			continue
+		}
+
+		nameNode := tree.Name(col.column.GetName())
 		names := tree.NameList{nameNode}
 
 		// Create a placeholder for the new value (len(columns)+i+1) since we
 		// use 1-indexed placeholders and the first len(columns) placeholders
 		// are for the where clause.
-		placeholder, err := newTypedPlaceholder(len(columns)+i+1, col)
+		placeholder, err := newTypedPlaceholder(len(columns)+i+1, col.column)
 		if err != nil {
 			return statements.Statement[tree.Statement]{}, err
 		}
 
-		exprs[i] = &tree.UpdateExpr{
+		exprs = append(exprs, &tree.UpdateExpr{
 			Names: names,
 			Expr:  placeholder,
-		}
+		})
 	}
 
 	// Create the final update statement
@@ -203,7 +240,7 @@ func newUpdateStatement(
 func newDeleteStatement(
 	table catalog.TableDescriptor,
 ) (statements.Statement[tree.Statement], error) {
-	columns := getPhysicalColumns(table)
+	columns := getColumnSchema(table)
 
 	// Create WHERE clause for matching the row to delete
 	whereClause, err := newMatchesLastRow(columns, 1)
@@ -250,7 +287,7 @@ func newDeleteStatement(
 func newBulkSelectStatement(
 	table catalog.TableDescriptor,
 ) (statements.Statement[tree.Statement], error) {
-	cols := getPhysicalColumnsSchema(table)
+	cols := getColumnSchema(table)
 	primaryKeyColumns := make([]catalog.Column, 0, len(cols))
 	for _, col := range cols {
 		if col.isPrimaryKey {