remove second object store req by padding each initial range request

Author: friendlymatthew

src/parallel_csv_opener.rs

@@ -142,6 +142,7 @@ impl<D: CsvDecoder + Clone> FileSource for ParallelCsvSource<D> {
         Ok(Arc::new(ParallelCsvOpener {
             schema: self.schema.clone(),
             batch_size: self.batch_size,
+            boundaries: self.boundaries.clone(),
             state: self.state.clone(),
             decoder: Arc::new(self.decoder.clone()),
             object_store,
@@ -189,9 +190,12 @@ impl<D: CsvDecoder + Clone> fmt::Debug for ParallelCsvSource<D> {
     }
 }
 
+const PADDING_BYTES: usize = 64 * 1024;
+
 struct ParallelCsvOpener {
     schema: SchemaRef,
     batch_size: usize,
+    boundaries: Arc<[usize]>,
     state: Arc<SharedState>,
     decoder: Arc<dyn CsvDecoder>,
     object_store: Arc<dyn ObjectStore>,
@@ -205,6 +209,7 @@ impl FileOpener for ParallelCsvOpener {
         let state = self.state.clone();
         let schema = self.schema.clone();
         let batch_size = self.batch_size;
+        let boundaries = self.boundaries.clone();
         let decoder = self.decoder.clone();
         let object_store = self.object_store.clone();
         let file_path = self.file_path.clone();
@@ -214,26 +219,35 @@ impl FileOpener for ParallelCsvOpener {
             coord.partition_range(partition_idx)
         };
 
+        let num_partitions = boundaries.len() - 1;
+        let file_end = boundaries[num_partitions];
+
+        let fetch_start = start.saturating_sub(PADDING_BYTES);
+        let fetch_end = (end + PADDING_BYTES).min(file_end);
+
         Ok(Box::pin(async move {
-            // phase 1: read byte range and scan for quote parity
+            // phase 1: fetch a range with some padding on both sides
             let result = object_store
                 .get_opts(
                     &file_path,
-                    GetOptions::new().with_range(Some(start as u64..end as u64)),
+                    GetOptions::new().with_range(Some(fetch_start as u64..fetch_end as u64)),
                 )
                 .await
                 .map_err(|e| DataFusionError::External(Box::new(e)))?;
 
-            let bytes = result
+            let full_bytes = result
                 .bytes()
                 .await
                 .map_err(|e| DataFusionError::External(Box::new(e)))?;
 
+            // phase 2: classify only the core [start, end) range
+            let core_offset = start - fetch_start;
+            let core_len = end - start;
             let classify_result = {
-                let bytes = bytes.clone();
+                let core_bytes = full_bytes.slice(core_offset..core_offset + core_len);
                 tokio::task::spawn_blocking(move || {
-                    let mut scanner = SubPartition::new(64 * 1024);
-                    scanner.ingest(&bytes);
+                    let mut scanner = SubPartition::new(PADDING_BYTES);
+                    scanner.ingest(&core_bytes);
                     scanner.finish()
                 })
                 .await
@@ -244,41 +258,40 @@ impl FileOpener for ParallelCsvOpener {
                 let mut results = state.classify_results.lock().unwrap();
                 results[partition_idx] = Some(classify_result);
             }
+
             state.barrier.wait().await;
 
-            // phase 2: resolve quotes and produce aligned ranges
+            // phase 3: resolve quotes and produce aligned ranges
             let resolved = state
                 .resolved
                 .get_or_init(|| async {
-                    let mut coordinator = state.coordinator.lock().unwrap();
+                    let mut resolver = state.coordinator.lock().unwrap();
                     {
                         let mut results = state.classify_results.lock().unwrap();
                         for (i, result) in results.iter_mut().enumerate() {
-                            coordinator.submit(i, result.take().unwrap());
+                            resolver.submit(i, result.take().unwrap());
                         }
                     }
-                    coordinator.resolve()
+                    resolver.resolve()
                 })
                 .await;
 
-            // phase 3: read aligned range from object store, decode
+            // phase 4: slice to get the aligned range
             let split = &resolved[partition_idx];
-
-            // note: this is the second object store request per parallel branch :P
-            let opts = GetOptions::new().with_range(Some(split.start as u64..split.end as u64));
-            let result = object_store
-                .get_opts(&file_path, opts)
-                .await
-                .map_err(|e| DataFusionError::External(Box::new(e)))?;
-
-            let aligned_bytes = result
-                .bytes()
-                .await
-                .map_err(|e| DataFusionError::External(Box::new(e)))?;
-
-            let batches = decoder
-                .decode(schema.clone(), batch_size, &aligned_bytes)
-                .map_err(DataFusionError::from)?;
+            let local_start = split
+                .start
+                .checked_sub(fetch_start)
+                .expect("resolved start before fetch_start");
+            let local_end = split.end - fetch_start;
+
+            let aligned_bytes = full_bytes.slice(local_start..local_end);
+
+            let batches = tokio::task::spawn_blocking(move || {
+                decoder.decode(schema, batch_size, &aligned_bytes)
+            })
+            .await
+            .map_err(|e| DataFusionError::External(Box::new(e)))?
+            .map_err(DataFusionError::from)?;
 
             let stream = futures::stream::iter(batches.into_iter().map(Ok::<_, DataFusionError>));
 

src/partition.rs

@@ -211,7 +211,7 @@ impl PartitionResolver {
     }
 }
 
-fn find_first_set_bit(bitsets: &[u64]) -> Option<usize> {
+pub fn find_first_set_bit(bitsets: &[u64]) -> Option<usize> {
     bitsets
         .iter()
         .enumerate()