polish MemoryStream related code (#14537)

apache · Feb 7, 2025 · c0d4ae7 · c0d4ae7
1 parent a9fb58c
commit c0d4ae7
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 7 deletions.
diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs
@@ -726,7 +726,7 @@ pub struct MemoryStream {
     projection: Option<Vec<usize>>,
     /// Index into the data
     index: usize,
-    /// The remaining number of rows to return
+    /// The remaining number of rows to return. If None, all rows are returned
     fetch: Option<usize>,
 }
 
@@ -778,11 +778,9 @@ impl Stream for MemoryStream {
             None => batch.clone(),
         };
 
-        if self.fetch.is_none() {
+        let Some(&fetch) = self.fetch.as_ref() else {
             return Poll::Ready(Some(Ok(batch)));
-        }
-
-        let fetch = self.fetch.unwrap();
+        };
         if fetch == 0 {
             return Poll::Ready(None);
         }

diff --git a/datafusion/sqllogictest/test_files/string/dictionary_utf8.slt b/datafusion/sqllogictest/test_files/string/dictionary_utf8.slt
@@ -67,4 +67,4 @@ statement ok
 drop table test_substr_base;
 
 statement ok
-drop table test_datetime_base;
+drop table test_datetime_base;
diff --git a/datafusion/sqllogictest/test_files/string/large_string.slt b/datafusion/sqllogictest/test_files/string/large_string.slt
@@ -75,4 +75,4 @@ statement ok
 drop table test_substr_base;
 
 statement ok
-drop table test_datetime_base;
+drop table test_datetime_base;