From e2e3cd8efb9228bec8b8c5cb01c0f4016bd46839 Mon Sep 17 00:00:00 2001
From: Peng Huo <penghuo@gmail.com>
Date: Mon, 7 Nov 2022 19:24:18 -0800
Subject: [PATCH] change from splits to split in batch

Signed-off-by: Peng Huo <penghuo@gmail.com>
---
 .../java/org/opensearch/sql/executor/streaming/Batch.java   | 3 +--
 .../sql/filesystem/streaming/FileSystemStreamSource.java    | 3 +--
 .../filesystem/streaming/FileSystemStreamSourceTest.java    | 6 ++----
 3 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/core/src/main/java/org/opensearch/sql/executor/streaming/Batch.java b/core/src/main/java/org/opensearch/sql/executor/streaming/Batch.java
index 7c27ab4622..cd7d7dae5a 100644
--- a/core/src/main/java/org/opensearch/sql/executor/streaming/Batch.java
+++ b/core/src/main/java/org/opensearch/sql/executor/streaming/Batch.java
@@ -5,7 +5,6 @@
 
 package org.opensearch.sql.executor.streaming;
 
-import java.util.List;
 import lombok.Data;
 import org.opensearch.sql.storage.split.Split;
 
@@ -14,5 +13,5 @@
  */
 @Data
 public class Batch {
-  private final List<Split> splits;
+  private final Split split;
 }
diff --git a/filesystem/src/main/java/org/opensearch/sql/filesystem/streaming/FileSystemStreamSource.java b/filesystem/src/main/java/org/opensearch/sql/filesystem/streaming/FileSystemStreamSource.java
index 0a1d032c53..6a9639bdcb 100644
--- a/filesystem/src/main/java/org/opensearch/sql/filesystem/streaming/FileSystemStreamSource.java
+++ b/filesystem/src/main/java/org/opensearch/sql/filesystem/streaming/FileSystemStreamSource.java
@@ -8,7 +8,6 @@
 import com.google.common.collect.Sets;
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.HashSet;
 import java.util.Optional;
 import java.util.Set;
@@ -101,6 +100,6 @@ public Batch getBatch(Optional<Offset> start, Offset end) {
             .collect(Collectors.toSet());
 
     log.debug("fetch files {} with id from: {} to: {}.", paths, start, end);
-    return new Batch(Collections.singletonList(new FileSystemSplit(paths)));
+    return new Batch(new FileSystemSplit(paths));
   }
 }
diff --git a/filesystem/src/test/java/org/opensearch/sql/filesystem/streaming/FileSystemStreamSourceTest.java b/filesystem/src/test/java/org/opensearch/sql/filesystem/streaming/FileSystemStreamSourceTest.java
index 75c494ec8c..fba038f6a3 100644
--- a/filesystem/src/test/java/org/opensearch/sql/filesystem/streaming/FileSystemStreamSourceTest.java
+++ b/filesystem/src/test/java/org/opensearch/sql/filesystem/streaming/FileSystemStreamSourceTest.java
@@ -18,7 +18,6 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Arrays;
-import java.util.List;
 import java.util.Optional;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FileSystem;
@@ -154,10 +153,9 @@ StreamSource batchInBetween(Long start, Long end, String... uris) {
 
     private StreamSource internalBatchInBetween(
         Optional<Offset> start, Offset end, String... uris) {
-      List<Split> splits = streamSource.getBatch(start, end).getSplits();
-      assertEquals(1, splits.size());
+      Split split = streamSource.getBatch(start, end).getSplit();
       assertThat(
-          ((FileSystemSplit) splits.get(0)).getPaths(),
+          ((FileSystemSplit) split).getPaths(),
           containsInAnyOrder(
               Arrays.stream(uris)
                   .map(name -> new org.apache.hadoop.fs.Path(perTestTempDir.resolve(name).toUri()))