From 21e41fbbe9391039cf4c0e65ba4fd57a1f90330d Mon Sep 17 00:00:00 2001
From: Walt Boettge <wboettge@amazon.com>
Date: Fri, 3 Feb 2023 09:42:16 -0800
Subject: [PATCH] Rewrite scope_identity on some index lookups (#94)

Currently, a problem arises when using scope_identity to lookup on
an index of identities. Babelfish implemented identity columns using INT
sequences, but scope_identity returns a numeric type. Numeric has higher
precedence than int, so the index column is implicitly cast to a
numeric. However, this means queries that should use an index lookup
instead use a sequential scan, significantly degrading performance.

With this change, if a call to scope_identity or
babelfish_get_last_identity_numeric (used with @@identity) is found in a
where clause compared to an integer, the function is replaced with
babelfish_get_last_identity, which returns an int so the index is used.

Task: BABEL-3384

Signed-off-by: Walt Boettge <wboettge@amazon.com>
---
 src/postgres/src/backend/parser/parse_expr.c | 51 ++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/src/postgres/src/backend/parser/parse_expr.c b/src/postgres/src/backend/parser/parse_expr.c
index a7e8d776a7ef..7d4ec452a594 100644
--- a/src/postgres/src/backend/parser/parse_expr.c
+++ b/src/postgres/src/backend/parser/parse_expr.c
@@ -910,6 +910,54 @@ exprIsNullConstant(Node *arg)
 	return false;
 }
 
+/*
+ * Helper function for replacing certain calls to scope_identity with
+ * babelfish_get_last_identity, only when using TSQL.
+ * 
+ * Fixes an issue where if scope_identity is compared to an identity index.
+ * scope_identity returns a numeric type, and needs to be rewritten to the
+ * underlying function that returns an int type so the index is used.
+ */
+static void
+rewrite_scope_identity_call(ParseState *pstate, Node **lexpr, Node **rexpr)
+{
+	Var       *col_expr;
+	FuncExpr  *func_expr;
+	FuncCall  *new_call;
+	char      *func_name;
+
+	if (sql_dialect != SQL_DIALECT_TSQL)
+		return;
+	if (!(*lexpr) || !(*rexpr))
+		return;
+
+	if ((IsA(*lexpr, Var) && IsA(*rexpr, FuncExpr)))
+	{
+		col_expr = (Var*) *lexpr;
+		func_expr = (FuncExpr*) *rexpr;
+	}
+	else if (IsA(*lexpr, FuncExpr) && IsA(*rexpr, Var))
+	{
+		col_expr = (Var*) *rexpr;
+		func_expr = (FuncExpr*) *lexpr;
+	}
+	else
+		return;
+
+	func_name = get_func_name(func_expr->funcid);
+	if (strcmp(func_name, "babelfish_get_last_identity_numeric") != 0 &&
+			strcmp(func_name, "scope_identity") != 0)
+		return;
+	if (col_expr->vartype != INT4OID)
+		return;
+
+	new_call = makeFuncCall(list_make1(makeString("babelfish_get_last_identity")), NULL, COERCE_EXPLICIT_CALL, -1);
+	if (IsA(*rexpr, FuncExpr))
+		*rexpr = transformFuncCall(pstate, new_call);
+	else
+		*lexpr = transformFuncCall(pstate, new_call);
+}
+
 static Node *
 transformAExprOp(ParseState *pstate, A_Expr *a)
 {
@@ -988,6 +1036,9 @@ transformAExprOp(ParseState *pstate, A_Expr *a)
 		lexpr = transformExprRecurse(pstate, lexpr);
 		rexpr = transformExprRecurse(pstate, rexpr);
 
+		/* So that where clauses with scope_identity use an index */
+		rewrite_scope_identity_call(pstate, &lexpr, &rexpr);
+
 		result = (Node *) make_op(pstate,
 								  a->name,
 								  lexpr,