diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index 414ce601763..03b68354f81 100644 --- a/rust/lance/src/io/exec/filtered_read.rs +++ b/rust/lance/src/io/exec/filtered_read.rs @@ -459,6 +459,7 @@ impl FilteredReadStream { options: FilteredReadOptions, metrics: &ExecutionPlanMetricsSet, plan: FilteredReadInternalPlan, + cached_fragments: Option>>, ) -> DataFusionResult { let global_metrics = Arc::new(FilteredReadGlobalMetrics::new(metrics)); @@ -482,23 +483,31 @@ impl FilteredReadStream { io_parallelism ); - // Ideally we don't need to collect here but if we don't we get "implementation of FnOnce is - // not general enough" false positives from rustc - let frag_futs = fragments - .iter() - .map(|frag| { - Result::Ok(Self::load_fragment( - dataset.clone(), - frag.clone(), - options.with_deleted_rows, - )) - }) - .collect::>(); - let loaded_fragments = futures::stream::iter(frag_futs) - // Cannot use unordered because we need to populate logical_offset based on user-provided order - .try_buffered(io_parallelism) - .try_collect::>() - .await?; + let loaded_fragments = match cached_fragments { + // Planning already loaded every fragment; don't repeat it + Some(loaded) => loaded, + None => { + // Ideally we don't need to collect here but if we don't we get "implementation of FnOnce is + // not general enough" false positives from rustc + let frag_futs = fragments + .iter() + .map(|frag| { + Result::Ok(Self::load_fragment( + dataset.clone(), + frag.clone(), + options.with_deleted_rows, + )) + }) + .collect::>(); + Arc::new( + futures::stream::iter(frag_futs) + // Cannot use unordered because we need to populate logical_offset based on user-provided order + .try_buffered(io_parallelism) + .try_collect::>() + .await?, + ) + } + }; let output_schema = public_blob_v2_binary_projection_schema(&options.projection); @@ -1677,6 +1686,21 @@ impl FilteredReadOptions { /// In the future, we may introduce high-level cardinality statistics similar to those used by query /// engines like Postgres. This might allow us to know, without executing an index search, that a scan /// would be better. In that case we accept the force_scan hint to skip the index search. +/// A resolved plan together with the fragment metadata loaded to compute it, +/// so stream construction reuses the load instead of repeating it. The +/// distributed path (`with_plan`, plan computed on another node) carries no +/// fragments and the stream loads them itself. +struct PlannedRead { + plan: FilteredReadInternalPlan, + loaded_fragments: Option>>, +} + +impl std::fmt::Debug for PlannedRead { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PlannedRead").finish() + } +} + #[derive(Debug)] pub struct FilteredReadExec { dataset: Arc, @@ -1685,7 +1709,7 @@ pub struct FilteredReadExec { metrics: ExecutionPlanMetricsSet, index_input: Option>, // Precomputed internal plan - plan: Arc>, + plan: Arc>, // When execute is first called we will initialize the FilteredReadStream. In order to support // multiple partitions, each partition will share the stream. running_stream: Arc>>, @@ -1819,6 +1843,9 @@ impl FilteredReadExec { /// Set the pre-computed plan for execution pub async fn with_plan(self, plan: FilteredReadPlan) -> Result { + // Distributed path: the plan was computed elsewhere (e.g. on a query + // node), so no fragment metadata is cached; stream construction + // loads it itself. let mut rows = BTreeMap::new(); for (fragment_id, selection) in plan.rows.iter() { let ranges = match selection { @@ -1846,7 +1873,10 @@ impl FilteredReadExec { scan_range_after_filter: plan.scan_range_after_filter, }; let plan_cell = Arc::new(OnceCell::new()); - let _ = plan_cell.set(internal_plan); + let _ = plan_cell.set(PlannedRead { + plan: internal_plan, + loaded_fragments: None, + }); Ok(Self { plan: plan_cell, ..self @@ -1855,13 +1885,13 @@ impl FilteredReadExec { /// Get or create the internal plan async fn get_or_create_plan_impl<'a>( - plan_cell: &'a OnceCell, + plan_cell: &'a OnceCell, dataset: Arc, options: &FilteredReadOptions, index_input: Option<&Arc>, partition: usize, ctx: Arc, - ) -> Result<&'a FilteredReadInternalPlan> { + ) -> Result<&'a PlannedRead> { plan_cell .get_or_try_init(|| async { // Execute index if present @@ -1899,12 +1929,14 @@ impl FilteredReadExec { .try_collect::>() .await?; - // Plan the scan - Ok(FilteredReadStream::plan_scan( - &loaded_fragments, - &evaluated_index, - options, - )) + // Plan the scan; keep the loaded fragments so stream + // construction doesn't reload every fragment + let plan = + FilteredReadStream::plan_scan(&loaded_fragments, &evaluated_index, options); + Ok(PlannedRead { + plan, + loaded_fragments: Some(Arc::new(loaded_fragments)), + }) }) .await } @@ -1920,7 +1952,7 @@ impl FilteredReadExec { ctx, ) .await?; - Ok(internal_plan.to_external_plan()) + Ok(internal_plan.plan.to_external_plan()) } fn obtain_stream( @@ -1956,7 +1988,7 @@ impl FilteredReadExec { let inner = if let Some(running_stream) = &*running_stream { running_stream.get_stream(&metrics, partition) } else { - let plan = Self::get_or_create_plan_impl( + let planned = Self::get_or_create_plan_impl( &plan_cell, dataset.clone(), &options, @@ -1966,10 +1998,15 @@ impl FilteredReadExec { ) .await .map_err(|e| DataFusionError::External(e.into()))?; - let new_running_stream = - FilteredReadStream::try_new(dataset, options, &metrics, plan.clone()) - .await - .map_err(|e| DataFusionError::External(e.into()))?; + let new_running_stream = FilteredReadStream::try_new( + dataset, + options, + &metrics, + planned.plan.clone(), + planned.loaded_fragments.clone(), + ) + .await + .map_err(|e| DataFusionError::External(e.into()))?; let first_stream = new_running_stream.get_stream(&metrics, partition); *running_stream = Some(new_running_stream); first_stream @@ -2032,7 +2069,7 @@ impl FilteredReadExec { /// Return the pre-computed plan if one exists, without triggering initialization. pub fn plan(&self) -> Option { - self.plan.get().map(|p| p.to_external_plan()) + self.plan.get().map(|p| p.plan.to_external_plan()) } }