From 4df01b19452111e09f5aceabfcebf8e38b365f4d Mon Sep 17 00:00:00 2001
From: Michael Buckley <michaelabuckley@gmail.com>
Date: Sat, 28 Sep 2024 17:26:41 -0400
Subject: [PATCH] Fix result set parsing when coord-sort is present.

---
 .../builder/sql/SearchQueryExecutor.java      | 50 +++++++++++++------
 1 file changed, 35 insertions(+), 15 deletions(-)

diff --git a/hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/builder/sql/SearchQueryExecutor.java b/hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/builder/sql/SearchQueryExecutor.java
index 49cdd701531..f636ab7eb4c 100644
--- a/hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/builder/sql/SearchQueryExecutor.java
+++ b/hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/builder/sql/SearchQueryExecutor.java
@@ -144,21 +144,7 @@ public class SearchQueryExecutor implements ISearchQueryExecutor {
 				if (myResultSet == null || !myResultSet.hasNext()) {
 					myNext = NO_MORE;
 				} else {
-					Object nextRow = Objects.requireNonNull(myResultSet.next());
-					// We should typically get two columns back, the first is the partition ID and the second
-					// is the resource ID. But if we're doing a count query, we'll get a single column in an array
-					// or maybe even just a single non array value depending on how the platform handles it.
-					if (nextRow instanceof Number) {
-						myNext = ((Number) nextRow).longValue();
-					} else {
-						Object[] nextRowAsArray = (Object[]) nextRow;
-						if (nextRowAsArray.length == 1) {
-							myNext = (Long) nextRowAsArray[0];
-						} else {
-							Integer nextPartitionId = (Integer) nextRowAsArray[0];
-							myNext = (Long) nextRowAsArray[1];
-						}
-					}
+					myNext = getNextPid(myResultSet);
 				}
 
 			} catch (Exception e) {
@@ -169,6 +155,40 @@ public class SearchQueryExecutor implements ISearchQueryExecutor {
 		}
 	}
 
+	private long getNextPid(ScrollableResultsIterator<Object> theResultSet) {
+		Object nextRow = Objects.requireNonNull(theResultSet.next());
+		// We should typically get two columns back, the first is the partition ID and the second
+		// is the resource ID. But if we're doing a count query, we'll get a single column in an array
+		// or maybe even just a single non array value depending on how the platform handles it.
+		if (nextRow instanceof Number) {
+			return ((Number) nextRow).longValue();
+		} else {
+			Object[] nextRowAsArray = (Object[]) nextRow;
+			if (nextRowAsArray.length == 1) {
+				return (Long) nextRowAsArray[0];
+			} else {
+				int i;
+				// TODO MB add a strategy object to GeneratedSql to describe the result set.
+				// or make SQE generic
+				// Comment to reviewer: this will be cleaner with the next
+				// merge from ja_20240718_pk_schema_selector
+
+				// We have some cases to distinguish:
+				// - res_id
+				// - count
+				// - partition_id, res_id
+				// - res_id, coord-dist
+				// - partition_id, res_id, coord-dist
+				// Assume res_id is first Long in row, and is in first two columns
+				if (nextRowAsArray[0] instanceof Long) {
+					return (long) nextRowAsArray[0];
+				} else {
+					return (long) nextRowAsArray[1];
+				}
+			}
+		}
+	}
+
 	public static SearchQueryExecutor emptyExecutor() {
 		return NO_VALUE_EXECUTOR;
 	}