From 60c5e148e2ee1c30a1a28eb8fcc1fc42af6bf8d2 Mon Sep 17 00:00:00 2001
From: Drew Kimball <drewk@cockroachlabs.com>
Date: Mon, 8 Jan 2024 03:38:42 -0700
Subject: [PATCH] opt: add method to add strict dependency to FuncDepSet

This commit adds a new method, `AddStrictDependency`, to `FuncDepSet`.
This will be used in the following commit to add a dependency between
a Window operator's partition columns and its functions. Existing
methods don't work for this because the dependency is not a key.

Epic: None

Release note: None
---
 pkg/sql/opt/props/func_dep.go           |  6 +++
 pkg/sql/opt/props/func_dep_rand_test.go | 63 +++++++++++++++++++++++++
 2 files changed, 69 insertions(+)

diff --git a/pkg/sql/opt/props/func_dep.go b/pkg/sql/opt/props/func_dep.go
index 0b8ac2aa3e28..136084c3a977 100644
--- a/pkg/sql/opt/props/func_dep.go
+++ b/pkg/sql/opt/props/func_dep.go
@@ -929,6 +929,12 @@ func (f *FuncDepSet) AddEquivalency(a, b opt.ColumnID) {
 	f.tryToReduceKey(opt.ColSet{} /* notNullCols */)
 }
 
+// AddStrictDependency adds a new strict dependency to the set.
+func (f *FuncDepSet) AddStrictDependency(from, to opt.ColSet) {
+	f.addDependency(from, to, true /* strict */, false /* equiv */)
+	f.tryToReduceKey(opt.ColSet{} /* notNullCols */)
+}
+
 // AddConstants adds a strict FD to the set that declares each given column as
 // having the same constant value for all rows. If a column is nullable, then
 // its value may be NULL, but then the column must be NULL for all rows. For
diff --git a/pkg/sql/opt/props/func_dep_rand_test.go b/pkg/sql/opt/props/func_dep_rand_test.go
index 36af140e00de..729c36d5da25 100644
--- a/pkg/sql/opt/props/func_dep_rand_test.go
+++ b/pkg/sql/opt/props/func_dep_rand_test.go
@@ -690,6 +690,67 @@ func (o *addSynthOp) ApplyToFDs(fd FuncDepSet) FuncDepSet {
 	return out
 }
 
+// addStrictDepOp is a test operation corresponding to AddStrictDependency.
+type addStrictDepOp struct {
+	from, to opt.ColSet
+}
+
+func genAddStrictDep(minCols, maxCols int) testOpGenerator {
+	return func(tc *testConfig) testOp {
+		from := tc.randColSet(minCols, maxCols)
+		to := tc.randColSet(minCols, maxCols)
+		from.DifferenceWith(to)
+		return &addStrictDepOp{
+			from: from,
+			to:   to,
+		}
+	}
+}
+
+func (o *addStrictDepOp) String() string {
+	return fmt.Sprintf("AddStrictDependency(%s, %s)", o.from, o.to)
+}
+
+func (o *addStrictDepOp) FilterRelation(tr testRelation) testRelation {
+	// Filter out rows where the from->to FD doesn't hold. The code here parallels
+	// that in testRelation.checkKey.
+	//
+	// We split the rows into groups (keyed on the `from` columns), picking the
+	// first row in each group as the "representative" of that group. All other
+	// rows in the group are checked against the representative row.
+	var out testRelation
+	m := make(map[rowKey]testRow)
+	perm := rand.Perm(len(tr))
+	for _, rowIdx := range perm {
+		r := tr[rowIdx]
+		k, _ := r.key(o.from)
+		if first, ok := m[k]; ok {
+			shouldFilter := false
+			for col, ok := o.to.Next(0); ok; col, ok = o.to.Next(col + 1) {
+				if first.value(col) != r.value(col) {
+					// Filter out row.
+					shouldFilter = true
+					break
+				}
+			}
+			if shouldFilter {
+				continue
+			}
+		} else {
+			m[k] = r
+		}
+		out = append(out, r)
+	}
+	return out
+}
+
+func (o *addStrictDepOp) ApplyToFDs(fd FuncDepSet) FuncDepSet {
+	var out FuncDepSet
+	out.CopyFrom(&fd)
+	out.AddStrictDependency(o.from, o.to)
+	return out
+}
+
 // testState corresponds to a chain of applied test operations. The head of a
 // testStates chain has no parent and no op and just corresponds to the initial
 // (empty) FDs and test relation.
@@ -803,6 +864,7 @@ func TestFuncDepOpsRandom(t *testing.T) {
 				genAddConst(1 /* minCols */, 3 /* maxCols */),
 				genAddEquiv(),
 				genAddSynth(0 /* minCols */, 3 /* maxCols */),
+				genAddStrictDep(0 /* minCols */, 3 /* maxCols */),
 			},
 		},
 
@@ -819,6 +881,7 @@ func TestFuncDepOpsRandom(t *testing.T) {
 				genAddConst(1 /* minCols */, 4 /* maxCols */),
 				genAddEquiv(),
 				genAddSynth(0 /* minCols */, 3 /* maxCols */),
+				genAddStrictDep(0 /* minCols */, 3 /* maxCols */),
 			},
 		},
 	}