Fixed batches being shifted by one in arrow.export (#693)

neo4j · Nov 27, 2024 · 5018823 · 5018823
1 parent 18e352d
commit 5018823
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/common/src/main/java/apoc/export/util/ProgressReporter.java b/common/src/main/java/apoc/export/util/ProgressReporter.java
@@ -95,7 +95,7 @@ public void updateRunningBatch(ProgressInfo progressInfo) {
 
     @Override
     public void done() {
-        if (totalEntities / batchSize == lastBatch) lastBatch++;
+        if (totalEntities % batchSize != 0) lastBatch++;
         updateRunningBatch(progressInfo);
         progressInfo.done(start);
         if (consumer != null) {

diff --git a/core/src/test/java/apoc/export/arrow/ArrowTest.java b/core/src/test/java/apoc/export/arrow/ArrowTest.java
@@ -147,6 +147,10 @@ private String extractFileName(Result result) {
         return result.<String>columnAs("file").next();
     }
 
+    private Long extractBatches(Result result) {
+        return result.<Long>columnAs("batches").next();
+    }
+
     private <T> T readValue(String json, Class<T> clazz) {
         if (json == null) return null;
         try {
@@ -378,6 +382,27 @@ public void testFileVolumeArrowAll() {
         db.executeTransactionally("MATCH (n:ArrowNode) DELETE n");
     }
 
+    @Test
+    public void testArrowCorrectBatchCount() {
+        // given - when
+        db.executeTransactionally("UNWIND range(0, 10000 - 1) AS id CREATE (:ArrowNode{id:id})");
+
+        Long batchesOdd = db.executeTransactionally(
+                "CALL apoc.export.arrow.query('batchCount_test.arrow', 'MATCH (n:ArrowNode) RETURN n.id AS id', {batchSize: 3}) YIELD batches",
+                Map.of(),
+                this::extractBatches);
+
+        Long batchesEven = db.executeTransactionally(
+                "CALL apoc.export.arrow.query('batchCount_test.arrow', 'MATCH (n:ArrowNode) RETURN n.id AS id', {batchSize: 2}) YIELD batches",
+                Map.of(),
+                this::extractBatches);
+
+        // then
+        assertEquals(3334, batchesOdd.longValue());
+        assertEquals(5000, batchesEven.longValue());
+        db.executeTransactionally("MATCH (n:ArrowNode) DELETE n");
+    }
+
     @Test
     public void testValidNonStorableQuery() {
         final List<byte[]> list = db.executeTransactionally(