Method consume_batch

parquet/src/arrow/array_reader/cached_array_reader.rs:270–340

(&mut self)

    }

    fn consume_batch(&mut self) -> Result<ArrayRef> {
        let row_count = self.selections.len();
        if row_count == 0 {
            return Ok(new_empty_array(self.inner.get_data_type()));
        }

        let start_position = self.outer_position - row_count;

        let selection_buffer = self.selections.finish();

        let start_batch = start_position / self.batch_size;
        let end_batch = (start_position + row_count - 1) / self.batch_size;

        let mut selected_arrays = Vec::new();
        for batch_id in start_batch..=end_batch {
            let batch_start = batch_id * self.batch_size;
            let batch_end = batch_start + self.batch_size - 1;
            let batch_id = self.get_batch_id_from_position(batch_start);

            // Calculate the overlap between the start_position and the batch
            let overlap_start = start_position.max(batch_start);
            let overlap_end = (start_position + row_count - 1).min(batch_end);

            if overlap_start > overlap_end {
                continue;
            }

            let selection_start = overlap_start - start_position;
            let selection_length = overlap_end - overlap_start + 1;
            let mask = selection_buffer.slice(selection_start, selection_length);

            if mask.count_set_bits() == 0 {
                continue;
            }

            let mask_array = BooleanArray::from(mask);
            // Read from local cache instead of shared cache to avoid cache eviction issues
            let cached = self
                .local_cache
                .get(&batch_id)
                .expect("data must be already cached in the read_records call, this is a bug");
            let cached = cached.slice(overlap_start - batch_start, selection_length);
            let filtered = arrow_select::filter::filter(&cached, &mask_array)?;
            selected_arrays.push(filtered);
        }

        self.selections = BooleanBufferBuilder::new(0);

        // Only remove batches from local buffer that are completely behind current position
        // Keep the current batch and any future batches as they might still be needed
        let current_batch_id = self.get_batch_id_from_position(self.outer_position);
        self.local_cache
            .retain(|batch_id, _| batch_id.val >= current_batch_id.val);

        // For consumers, cleanup batches that have been completely consumed
        // This reduces the memory usage of the shared cache
        if self.role == CacheRole::Consumer {
            self.cleanup_consumed_batches();