From e2c1f7d8876e11c28bfc4a9478080712a0868027 Mon Sep 17 00:00:00 2001
From: Bykov Ivan <i.bykov@modernsys.ru>
Date: Tue, 11 Mar 2025 14:34:21 +0500
Subject: [PATCH] Query ID Calculation Fix for DISTINCT / ORDER BY and LIMIT / 
 OFFSET

In some cases, we may encounter different query trees with the same IDs.
For two structurally similar query subnodes, the query trees may look like this:

QueryA->subNodeOne = Value X;
QueryA->subNodeTwo = NULL;
QueryB->subNodeOne = NULL;
QueryB->subNodeTwo = Value X;

When the query jumble walker calculates the query ID, it traverses
the query members from top to bottom and generates the same IDs
for these two queries because it does not change the hash value
when visiting an empty node (i.e., NULL).

There are two pairs of subnodes that can trigger this error:
- distinctClause and sortClause (both contain a list of SortGroupClauses);
- limitOffset and limitCount (both contain an int8 expression).

To fix this problem, we count NULL nodes and add that count to the query
jumble buffer every time we visit a non-NULL field.
---
 .../pg_stat_statements/expected/select.out    | 49 ++++++++++++++++++
 contrib/pg_stat_statements/sql/select.sql     | 18 +++++++
 src/backend/nodes/queryjumblefuncs.c          | 50 +++++++++++++++----
 src/include/nodes/queryjumble.h               |  3 ++
 4 files changed, 111 insertions(+), 9 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/select.out b/contrib/pg_stat_statements/expected/select.out
index dd6c756f67d..f1b71f56b7b 100644
--- a/contrib/pg_stat_statements/expected/select.out
+++ b/contrib/pg_stat_statements/expected/select.out
@@ -406,9 +406,58 @@ SELECT COUNT(*) FROM pg_stat_statements WHERE query LIKE '%SELECT GROUPING%';
      2
 (1 row)
 
+CREATE TABLE pgss_c AS
+  SELECT id FROM generate_series(1,2) AS id;
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
  t 
 ---
  t
 (1 row)
 
+-- These two queries trigger a Query ID collision in PostgreSQL versions prior to 18
+SELECT DISTINCT id FROM pgss_c;
+ id 
+----
+  2
+  1
+(2 rows)
+
+SELECT id FROM pgss_c ORDER BY id;
+ id 
+----
+  1
+  2
+(2 rows)
+
+-- These two queries trigger a Query ID collision in PostgreSQL versions prior to 18
+SELECT id FROM pgss_c LIMIT 1;
+ id 
+----
+  1
+(1 row)
+
+SELECT id FROM pgss_c OFFSET 1;
+ id 
+----
+  2
+(1 row)
+
+-- Expected four rows (Query ID collision has been fixed).
+SELECT query FROM pg_stat_statements
+  WHERE query LIKE '%pgss_c%'
+  ORDER BY query COLLATE "C";
+               query               
+-----------------------------------
+ SELECT DISTINCT id FROM pgss_c
+ SELECT id FROM pgss_c LIMIT $1
+ SELECT id FROM pgss_c OFFSET $1
+ SELECT id FROM pgss_c ORDER BY id
+(4 rows)
+
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+DROP TABLE pgss_c;
diff --git a/contrib/pg_stat_statements/sql/select.sql b/contrib/pg_stat_statements/sql/select.sql
index eb45cb81ad2..bbc37fa0005 100644
--- a/contrib/pg_stat_statements/sql/select.sql
+++ b/contrib/pg_stat_statements/sql/select.sql
@@ -146,4 +146,22 @@ SELECT (
 ) FROM (VALUES(6,7)) v3(e,f) GROUP BY ROLLUP(e,f);
 
 SELECT COUNT(*) FROM pg_stat_statements WHERE query LIKE '%SELECT GROUPING%';
+
+CREATE TABLE pgss_c AS
+  SELECT id FROM generate_series(1,2) AS id;
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+
+-- These two queries trigger a Query ID collision in PostgreSQL versions prior to 18
+SELECT DISTINCT id FROM pgss_c;
+SELECT id FROM pgss_c ORDER BY id;
+-- These two queries trigger a Query ID collision in PostgreSQL versions prior to 18
+SELECT id FROM pgss_c LIMIT 1;
+SELECT id FROM pgss_c OFFSET 1;
+
+-- Expected four rows (Query ID collision has been fixed).
+SELECT query FROM pg_stat_statements
+  WHERE query LIKE '%pgss_c%'
+  ORDER BY query COLLATE "C";
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+
+DROP TABLE pgss_c;
diff --git a/src/backend/nodes/queryjumblefuncs.c b/src/backend/nodes/queryjumblefuncs.c
index 129fb447099..cf6d43259c0 100644
--- a/src/backend/nodes/queryjumblefuncs.c
+++ b/src/backend/nodes/queryjumblefuncs.c
@@ -51,6 +51,7 @@ int			compute_query_id = COMPUTE_QUERY_ID_AUTO;
  */
 bool		query_id_enabled = false;
 
+static inline Size CompressJumble(unsigned char *jumble, Size jumble_len);
 static void AppendJumble(JumbleState *jstate,
 						 const unsigned char *item, Size size);
 static void RecordConstLocation(JumbleState *jstate, int location);
@@ -118,6 +119,7 @@ JumbleQuery(Query *query)
 		palloc(jstate->clocations_buf_size * sizeof(LocationLen));
 	jstate->clocations_count = 0;
 	jstate->highest_extern_param_id = 0;
+	jstate->null_sequence_number = 0;
 
 	/* Compute query ID and mark the Query node with it */
 	_jumbleNode(jstate, (Node *) query);
@@ -153,6 +155,26 @@ EnableQueryId(void)
 		query_id_enabled = true;
 }
 
+
+/*
+ * CompressJumble: Check whether the jumble buffer has overflowed.
+ * If it has, calculate the hash and store it in buffer.
+ */
+static inline Size
+CompressJumble(unsigned char *jumble, Size jumble_len)
+{
+	if (unlikely(jumble_len >= JUMBLE_SIZE))
+	{
+		uint64		start_hash;
+
+		start_hash = DatumGetUInt64(hash_any_extended(jumble,
+													  JUMBLE_SIZE, 0));
+		memcpy(jumble, &start_hash, sizeof(start_hash));
+		return sizeof(start_hash);
+	}
+	return jumble_len;
+}
+
 /*
  * AppendJumble: Append a value that is substantive in a given query to
  * the current jumble.
@@ -172,21 +194,28 @@ AppendJumble(JumbleState *jstate, const unsigned char *item, Size size)
 	{
 		Size		part_size;
 
-		if (jumble_len >= JUMBLE_SIZE)
-		{
-			uint64		start_hash;
-
-			start_hash = DatumGetUInt64(hash_any_extended(jumble,
-														  JUMBLE_SIZE, 0));
-			memcpy(jumble, &start_hash, sizeof(start_hash));
-			jumble_len = sizeof(start_hash);
-		}
+		jumble_len = CompressJumble(jumble, jumble_len);
 		part_size = Min(size, JUMBLE_SIZE - jumble_len);
 		memcpy(jumble + jumble_len, item, part_size);
 		jumble_len += part_size;
 		item += part_size;
 		size -= part_size;
 	}
+
+	/*
+	 * To prevent Query ID collisions, add to the query jumble information the
+	 * count of visited NULL nodes since the last non-NULL node was hashed.
+	 * This approach consumes less jumble memory than simply adding a NULL
+	 * marker for every NULL node field.
+	 */
+	if (jstate->null_sequence_number)
+	{
+		jumble_len = CompressJumble(jumble, jumble_len);
+		jumble[jumble_len] = jstate->null_sequence_number & 0xFF;
+		jstate->null_sequence_number = 0;
+		++jumble_len;
+	}
+
 	jstate->jumble_len = jumble_len;
 }
 
@@ -238,7 +267,10 @@ _jumbleNode(JumbleState *jstate, Node *node)
 	Node	   *expr = node;
 
 	if (expr == NULL)
+	{
+		++jstate->null_sequence_number;
 		return;
+	}
 
 	/* Guard against stack overflow due to overly complex expressions */
 	check_stack_depth();
diff --git a/src/include/nodes/queryjumble.h b/src/include/nodes/queryjumble.h
index f1c55c8067f..81dba02f08c 100644
--- a/src/include/nodes/queryjumble.h
+++ b/src/include/nodes/queryjumble.h
@@ -48,6 +48,9 @@ typedef struct JumbleState
 
 	/* highest Param id we've seen, in order to start normalization correctly */
 	int			highest_extern_param_id;
+
+	/* count of visited NULL nodes since the last non-NULL node was hashed */
+	int			null_sequence_number;
 } JumbleState;
 
 /* Values for the compute_query_id GUC */
-- 
2.39.5

