From 3d78113c2297dadeab90029dc9193001ce42b50e Mon Sep 17 00:00:00 2001
From: evenyag <realevenyag@gmail.com>
Date: Fri, 14 Mar 2025 00:35:19 +0800
Subject: [PATCH] feat: reduce rows returned

---
 src/mito2/src/read/dedup.rs | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/mito2/src/read/dedup.rs b/src/mito2/src/read/dedup.rs
index 563b590b7f..ff47b38711 100644
--- a/src/mito2/src/read/dedup.rs
+++ b/src/mito2/src/read/dedup.rs
@@ -602,11 +602,17 @@ impl TagOnlyReader {
 impl BatchReader for TagOnlyReader {
     async fn next_batch(&mut self) -> Result<Option<Batch>> {
         while let Some(batch) = self.source.next_batch().await? {
+            if batch.is_empty() {
+                // Ensure that the batch is not empty before proceeding.
+                continue;
+            }
+
             if let Some(to_return) = self.to_return.take() {
                 if to_return.primary_key() != batch.primary_key() {
                     self.to_return = Some(batch);
                     // A new key, store the batch and returns the previous one.
-                    return Ok(Some(to_return));
+                    // Safety: The batch is not empty, so it has at least one row.
+                    return Ok(Some(to_return.slice(0, 1)));
                 } else {
                     // The same key, override the batch.
                     self.to_return = Some(batch);