invertedidx: add convenience function for extracting index column var…

…iable

Release note: None

pkg/sql/opt/invertedidx/inverted_index_expr.go

@@ -497,3 +497,21 @@ func extractInvertedFilterCondition(
 		return filterPlanner.extractInvertedFilterConditionFromLeaf(evalCtx, filterCond)
 	}
 }
+
+// indexColumnVariable returns a variable expression (a type-asserted form of e)
+// and ok=true if e is a variable expression that corresponds to the inverted
+// index column.
+func indexColumnVariable(
+	tabID opt.TableID, index cat.Index, e opt.Expr,
+) (_ *memo.VariableExpr, ok bool) {
+	variable, ok := e.(*memo.VariableExpr)
+	if !ok {
+		return nil, false
+	}
+	invertedIndexCol := tabID.ColumnID(index.VirtualInvertedColumn().InvertedSourceColumnOrdinal())
+	if variable.Col != invertedIndexCol {
+		// The column does not match the index column.
+		return nil, false
+	}
+	return variable, true
+}

pkg/sql/opt/invertedidx/json_array.go

@@ -58,13 +58,8 @@ func (j *jsonOrArrayJoinPlanner) canExtractJSONOrArrayJoinCondition(
 ) bool {
 	// The first argument should be a variable corresponding to the index
 	// column.
-	variable, ok := left.(*memo.VariableExpr)
+	variable, ok := indexColumnVariable(j.tabID, j.index, left)
 	if !ok {
-		return false
-	}
-	if variable.Col != j.tabID.ColumnID(
-		j.index.VirtualInvertedColumn().InvertedSourceColumnOrdinal(),
-	) {
 		// The column does not match the index column.
 		return false
 	}
@@ -300,16 +295,10 @@ func (j *jsonOrArrayFilterPlanner) extractJSONOrArrayContainsCondition(
 ) invertedexpr.InvertedExpression {
 	// The first argument should be a variable corresponding to the index
 	// column.
-	variable, ok := left.(*memo.VariableExpr)
+	variable, ok := indexColumnVariable(j.tabID, j.index, left)
 	if !ok {
 		return invertedexpr.NonInvertedColExpression{}
 	}
-	if variable.Col != j.tabID.ColumnID(
-		j.index.VirtualInvertedColumn().InvertedSourceColumnOrdinal(),
-	) {
-		// The column does not match the index column.
-		return invertedexpr.NonInvertedColExpression{}
-	}
 
 	// The second argument should be a constant.
 	if !memo.CanExtractConstDatum(right) {