From 646e6861f91b988753b76da5c8fe7b2ca897b203 Mon Sep 17 00:00:00 2001
From: Trianz-Akshay <108925344+Trianz-Akshay@users.noreply.github.com>
Date: Fri, 8 Mar 2024 22:22:36 +0530
Subject: [PATCH] [BUG] error occurs when using the LIMIT clause with
 AthenaGoogleBigQueryConnector (#1808)

---
 .../connectors/google/bigquery/BigQueryRecordHandler.java       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/athena-google-bigquery/src/main/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryRecordHandler.java b/athena-google-bigquery/src/main/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryRecordHandler.java
index a6b7a8cb6e..9d24184f4b 100644
--- a/athena-google-bigquery/src/main/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryRecordHandler.java
+++ b/athena-google-bigquery/src/main/java/com/amazonaws/athena/connectors/google/bigquery/BigQueryRecordHandler.java
@@ -232,7 +232,7 @@ private void getTableData(BlockSpiller spiller, ReadRecordsRequest recordsReques
                 for (ReadRowsResponse response : stream) {
                     Preconditions.checkState(response.hasArrowRecordBatch());
                     VectorSchemaRoot root = reader.processRows(response.getArrowRecordBatch());
-                    long rowLimit = (recordsRequest.getConstraints().getLimit() > 0) ? recordsRequest.getConstraints().getLimit() : root.getRowCount();
+                    long rowLimit = (recordsRequest.getConstraints().getLimit() > 0 && recordsRequest.getConstraints().getLimit() < root.getRowCount()) ? recordsRequest.getConstraints().getLimit() : root.getRowCount();
                     for (int rowIndex = 0; rowIndex < rowLimit; rowIndex++) {
                         outputResults(spiller, recordsRequest, root, rowIndex);
                     }