| 55 | #[async_trait] |
| 56 | impl FormatFileReader for BlobFormatReader { |
| 57 | async fn read_batch_stream( |
| 58 | &self, |
| 59 | reader: Box<dyn FileRead>, |
| 60 | file_size: u64, |
| 61 | read_fields: &[DataField], |
| 62 | _predicates: Option<&FilePredicates>, |
| 63 | batch_size: Option<usize>, |
| 64 | row_selection: Option<Vec<RowRange>>, |
| 65 | ) -> crate::Result<ArrowRecordBatchStream> { |
| 66 | validate_read_fields(read_fields)?; |
| 67 | |
| 68 | let target_schema = build_target_arrow_schema(read_fields)?; |
| 69 | let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE); |
| 70 | let blob_index = BlobFileIndex::load(reader.as_ref(), file_size).await?; |
| 71 | let mut selection = RowSelectionCursor::new(blob_index.num_rows(), row_selection)?; |
| 72 | let project_values = !read_fields.is_empty(); |
| 73 | |
| 74 | if self.descriptor_mode { |
| 75 | let file_path = self.file_path.clone(); |
| 76 | Ok(try_stream! { |
| 77 | while let Some(positions) = selection.next_batch(batch_size) { |
| 78 | let batch = if project_values { |
| 79 | build_descriptor_batch(&blob_index, &target_schema, &positions, &file_path)? |
| 80 | } else { |
| 81 | RecordBatch::try_new_with_options( |
| 82 | target_schema.clone(), |
| 83 | Vec::new(), |
| 84 | &RecordBatchOptions::new().with_row_count(Some(positions.len())), |
| 85 | ) |
| 86 | .map_err(|e| Error::UnexpectedError { |
| 87 | message: format!("Failed to build empty blob RecordBatch: {e}"), |
| 88 | source: Some(Box::new(e)), |
| 89 | })? |
| 90 | }; |
| 91 | yield batch; |
| 92 | } |
| 93 | } |
| 94 | .boxed()) |
| 95 | } else { |
| 96 | Ok(try_stream! { |
| 97 | while let Some(positions) = selection.next_batch(batch_size) { |
| 98 | let batch = read_blob_batch( |
| 99 | reader.as_ref(), |
| 100 | &blob_index, |
| 101 | &target_schema, |
| 102 | &positions, |
| 103 | project_values, |
| 104 | ).await?; |
| 105 | yield batch; |
| 106 | } |
| 107 | } |
| 108 | .boxed()) |
| 109 | } |
| 110 | } |
| 111 | } |
| 112 | |
| 113 | fn validate_read_fields(read_fields: &[DataField]) -> crate::Result<()> { |