Support filter out strip by provided range

datafusion-contrib · Sep 29, 2024 · b17dc02 · b17dc02
1 parent 9e5cfbf
commit b17dc02
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 13 deletions.
diff --git a/src/arrow_reader.rs b/src/arrow_reader.rs
@@ -39,7 +39,7 @@ pub struct ArrowReaderBuilder<R> {
     pub(crate) batch_size: usize,
     pub(crate) projection: ProjectionMask,
     pub(crate) schema_ref: Option<SchemaRef>,
-    pub(crate) range: Option<Range<usize>>,
+    pub(crate) file_byte_range: Option<Range<usize>>,
 }
 
 impl<R> ArrowReaderBuilder<R> {
@@ -50,7 +50,7 @@ impl<R> ArrowReaderBuilder<R> {
             batch_size: DEFAULT_BATCH_SIZE,
             projection: ProjectionMask::all(),
             schema_ref: None,
-            range: None,
+            file_byte_range: None,
         }
     }
 
@@ -73,8 +73,9 @@ impl<R> ArrowReaderBuilder<R> {
         self
     }
 
-    pub fn with_range(mut self, range: Range<usize>) -> Self {
-        self.range = Some(range);
+    /// Specifies a range of file bytes that will read the strips offset within this range
+    pub fn with_file_byte_range(mut self, range: Range<usize>) -> Self {
+        self.file_byte_range = Some(range);
         self
     }
 
@@ -116,7 +117,7 @@ impl<R: ChunkReader> ArrowReaderBuilder<R> {
             file_metadata: self.file_metadata,
             projected_data_type,
             stripe_index: 0,
-            range: self.range,
+            file_byte_range: self.file_byte_range,
         };
         ArrowReader {
             cursor,
@@ -185,18 +186,18 @@ pub(crate) struct Cursor<R> {
     pub file_metadata: Arc<FileMetadata>,
     pub projected_data_type: RootDataType,
     pub stripe_index: usize,
-    pub range: Option<Range<usize>>,
+    pub file_byte_range: Option<Range<usize>>,
 }
 
 impl<R: ChunkReader> Cursor<R> {
     fn get_stripe_metadatas(&self) -> Vec<StripeMetadata> {
-        if let Some(range) = self.range.clone() {
+        if let Some(range) = self.file_byte_range.clone() {
             self.file_metadata
                 .stripe_metadatas()
                 .iter()
                 .filter(|info| {
                     let offset = info.offset() as usize;
-                    !(offset < range.start || offset >= range.end)
+                    range.contains(&offset)
                 })
                 .map(|info| info.to_owned())
                 .collect::<Vec<_>>()

diff --git a/src/async_arrow_reader.rs b/src/async_arrow_reader.rs
@@ -104,9 +104,9 @@ impl<R: AsyncChunkReader + 'static> StripeFactory<R> {
             .cloned();
 
         if let Some(info) = info {
-            if let Some(range) = self.inner.range.clone() {
+            if let Some(range) = self.inner.file_byte_range.clone() {
                 let offset = info.offset() as usize;
-                if offset < range.start || offset >= range.end {
+                if !range.contains(&offset) {
                     self.inner.stripe_index += 1;
                     return Ok((self, None));
                 }
@@ -221,7 +221,7 @@ impl<R: AsyncChunkReader + 'static> ArrowReaderBuilder<R> {
             file_metadata: self.file_metadata,
             projected_data_type,
             stripe_index: 0,
-            range: self.range,
+            file_byte_range: self.file_byte_range,
         };
         ArrowStreamReader::new(cursor, self.batch_size, schema_ref)
     }

diff --git a/tests/basic/main.rs b/tests/basic/main.rs
@@ -58,7 +58,7 @@ async fn new_arrow_stream_reader_range(
     ArrowReaderBuilder::try_new_async(f)
         .await
         .unwrap()
-        .with_range(range)
+        .with_file_byte_range(range)
         .build_async()
 }
 
@@ -71,7 +71,7 @@ fn new_arrow_reader_range(path: &str, range: Range<usize>) -> ArrowReader<File>
     let f = File::open(path).expect("no file found");
     ArrowReaderBuilder::try_new(f)
         .unwrap()
-        .with_range(range)
+        .with_file_byte_range(range)
         .build()
 }