From 481aa92b448d835cc2fa9e349281db045953e5d1 Mon Sep 17 00:00:00 2001
From: Mikko Kortelainen <kortemik@users.noreply.github.com>
Date: Tue, 14 Apr 2026 15:49:17 +0300
Subject: [PATCH] Revert "remove collectAsList call (#13)"

This reverts commit 9f6132bbe4913187bd7a407df4c1a843c5cbfe86.
---
 .../java/com/teragrep/functions/dpf_02/BatchCollect.java   | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/main/java/com/teragrep/functions/dpf_02/BatchCollect.java b/src/main/java/com/teragrep/functions/dpf_02/BatchCollect.java
index 1af66d6..4368803 100644
--- a/src/main/java/com/teragrep/functions/dpf_02/BatchCollect.java
+++ b/src/main/java/com/teragrep/functions/dpf_02/BatchCollect.java
@@ -132,13 +132,14 @@ public void collect(Dataset<Row> batchDF, Long batchId, List<AbstractStep> postB
         if (!skipLimiting && limit > 0) {
             orderedDs = orderedDs.limit(limit);
         }
-
+        List<Row> collected = orderedDs.collectAsList();
+        Dataset<Row> createdDsFromCollected = SparkSession.builder().getOrCreate().createDataFrame(collected, this.inputSchema);
         Dataset<Row> current;
         if (this.savedDs == null) {
-            current = orderedDs;
+            current = createdDsFromCollected;
         }
         else {
-            current = savedDs.union(orderedDs);
+            current = savedDs.union(createdDsFromCollected);
         }
 
         current = orderDataset(current);