1 files changed, 4 insertions, 2 deletions
diff --git a/sql/catalyst/src/test/java/org/apache/spark/sql/catalyst/expressions/RowBasedKeyValueBatchSuite.java b/sql/catalyst/src/test/java/org/apache/spark/sql/catalyst/expressions/RowBasedKeyValueBatchSuite.java
index 0dd129cea7..fb3dbe8ed1 100644
--- a/sql/catalyst/src/test/java/org/apache/spark/sql/catalyst/expressions/RowBasedKeyValueBatchSuite.java
+++ b/sql/catalyst/src/test/java/org/apache/spark/sql/catalyst/expressions/RowBasedKeyValueBatchSuite.java
@@ -338,15 +338,17 @@ public class RowBasedKeyValueBatchSuite {
 
   @Test
   public void appendRowUntilExceedingPageSize() throws Exception {
+    // Use default size or spark.buffer.pageSize if specified
+    int pageSizeToUse = (int) memoryManager.pageSizeBytes();
     RowBasedKeyValueBatch batch = RowBasedKeyValueBatch.allocate(keySchema,
-            valueSchema, taskMemoryManager, 64 * 1024 * 1024); //enough capacity
+            valueSchema, taskMemoryManager, pageSizeToUse); //enough capacity
     try {
       UnsafeRow key = makeKeyRow(1, "A");
       UnsafeRow value = makeValueRow(1, 1);
       int recordLength = 8 + key.getSizeInBytes() + value.getSizeInBytes() + 8;
       int totalSize = 4;
       int numRows = 0;
-      while (totalSize + recordLength < 64 * 1024 * 1024) { // default page size
+      while (totalSize + recordLength < pageSizeToUse) {
         appendRow(batch, key, value);
         totalSize += recordLength;
         numRows++;