From 77cd83cb3173a5db6ccdabc602115acc1578b776 Mon Sep 17 00:00:00 2001
From: Shri Harish Saravanan
 <114852752+shriharish1310@users.noreply.github.com>
Date: Mon, 3 Aug 2026 09:39:36 -0500
Subject: [PATCH] Test bulk import row and sort key types

---
 .../runner/BulkImportJobDriverIT.java         | 74 +++++++++++++++++--
 1 file changed, 67 insertions(+), 7 deletions(-)

diff --git a/java/bulk-import/bulk-import-runner/src/test/java/sleeper/bulkimport/runner/BulkImportJobDriverIT.java b/java/bulk-import/bulk-import-runner/src/test/java/sleeper/bulkimport/runner/BulkImportJobDriverIT.java
index c4aacc1fcd0..e5abda4e48b 100644
--- a/java/bulk-import/bulk-import-runner/src/test/java/sleeper/bulkimport/runner/BulkImportJobDriverIT.java
+++ b/java/bulk-import/bulk-import-runner/src/test/java/sleeper/bulkimport/runner/BulkImportJobDriverIT.java
@@ -51,6 +51,7 @@
 import sleeper.core.schema.type.ListType;
 import sleeper.core.schema.type.LongType;
 import sleeper.core.schema.type.MapType;
+import sleeper.core.schema.type.PrimitiveType;
 import sleeper.core.schema.type.StringType;
 import sleeper.core.statestore.FileReference;
 import sleeper.core.statestore.StateStore;
@@ -107,14 +108,28 @@
 
 class BulkImportJobDriverIT {
 
-    private static Stream<Arguments> getStreamOfBulkImportJobRunners() {
+    private static Stream<Named<BulkImportJobRunner>> getBulkImportJobRunners() {
         return Stream.of(
-                Arguments.of(Named.of("BulkImportJobDataframeDriver",
-                        (BulkImportJobRunner) BulkImportJobDataframeDriver::createFileReferences)),
-                Arguments.of(Named.of("BulkImportJobRDDDriver",
-                        (BulkImportJobRunner) BulkImportJobRDDDriver::createFileReferences)),
-                Arguments.of(Named.of("BulkImportDataframeLocalSortDriver",
-                        (BulkImportJobRunner) BulkImportDataframeLocalSortDriver::createFileReferences)));
+                Named.of("BulkImportJobDataframeDriver",
+                        (BulkImportJobRunner) BulkImportJobDataframeDriver::createFileReferences),
+                Named.of("BulkImportJobRDDDriver",
+                        (BulkImportJobRunner) BulkImportJobRDDDriver::createFileReferences),
+                Named.of("BulkImportDataframeLocalSortDriver",
+                        (BulkImportJobRunner) BulkImportDataframeLocalSortDriver::createFileReferences));
+    }
+
+    private static Stream<Arguments> getStreamOfBulkImportJobRunners() {
+        return getBulkImportJobRunners().map(Arguments::of);
+    }
+
+    private static Stream<Arguments> getStreamOfBulkImportJobRunnersAndKeyTypes() {
+        return getBulkImportJobRunners().flatMap(runner -> Stream.of(
+                Arguments.of(runner, Named.of("LongType", new KeyTypeTestData(
+                        new LongType(), 1L, 2L))),
+                Arguments.of(runner, Named.of("StringType", new KeyTypeTestData(
+                        new StringType(), "A", "B"))),
+                Arguments.of(runner, Named.of("ByteArrayType", new KeyTypeTestData(
+                        new ByteArrayType(), new byte[]{1}, new byte[]{2})))));
     }
 
     @TempDir
@@ -193,6 +208,25 @@ void shouldImportDataSinglePartition(BulkImportJobRunner runner) throws Exceptio
                         ingestFinishedStatus(summary(startTime, endTime, 200, 200), 1))));
     }
 
+    @ParameterizedTest
+    @MethodSource("getStreamOfBulkImportJobRunnersAndKeyTypes")
+    void shouldImportDataWithSupportedRowAndSortKeyTypes(
+            BulkImportJobRunner runner, KeyTypeTestData keyType) throws Exception {
+        // Given
+        tableProperties.setSchema(getSchemaWithKeyType(keyType.type()));
+        update(stateStore()).initialise(tableProperties);
+        List<Row> rows = getRows(keyType);
+        writeRowsToFile(rows, dataDir + "/import/a.parquet");
+
+        // When
+        BulkImportJob job = jobForTable().id("my-job")
+                .files(List.of(dataDir + "/import/a.parquet")).build();
+        runJob(runner, job);
+
+        // Then
+        assertThat(readRowsInPartitionTreeOrder()).isEqualTo(sorted(rows));
+    }
+
     @ParameterizedTest
     @MethodSource("getStreamOfBulkImportJobRunners")
     void shouldImportDataSinglePartitionIdenticalRowKeyDifferentSortKeys(BulkImportJobRunner runner) throws Exception {
@@ -458,6 +492,29 @@ private static Schema getSchema() {
                 .build();
     }
 
+    private static Schema getSchemaWithKeyType(PrimitiveType keyType) {
+        return Schema.builder()
+                .rowKeyFields(new Field("key", keyType))
+                .sortKeyFields(new Field("sort", keyType))
+                .valueFields(new Field("value", new StringType()))
+                .build();
+    }
+
+    private static List<Row> getRows(KeyTypeTestData keyType) {
+        return List.of(
+                row(keyType.higherValue(), keyType.higherValue(), "higher key"),
+                row(keyType.lowerValue(), keyType.higherValue(), "higher sort"),
+                row(keyType.lowerValue(), keyType.lowerValue(), "lower sort"));
+    }
+
+    private static Row row(Object key, Object sort, String value) {
+        Row row = new Row();
+        row.put("key", key);
+        row.put("sort", sort);
+        row.put("value", value);
+        return row;
+    }
+
     private static List<Row> getRows() {
         List<Row> rows = new ArrayList<>(200);
         for (int i = 0; i < 100; i++) {
@@ -571,4 +628,7 @@ private String createDir(String name) {
         }
         return path.toString();
     }
+
+    private record KeyTypeTestData(PrimitiveType type, Object lowerValue, Object higherValue) {
+    }
 }