Skip to content

Commit 2d96876

Browse files
committed
Wire up with_new_state with DataSource (apache#20718)
## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes apache#123` indicates that this PR will close issue apache#123. --> - Closes #. ## Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> `ExecutionPlan::with_new_state()` allows devs to inject custom information in their nodes that they can use for tweaking their `ExecutionPlan` implementations. This mechanism does not work today if the `ExecutionPlan` is a `DataSourceExec`, as this one does not implement the `with_new_state()` method from `ExecutionPlan`. In order to let people use this also for their own `DataSource` implementations, this PR adds this method to it. ## What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> Propagates `with_new_state()` to the `DataSource` trait, so that custom `DataSourceExec` can also benefit from it. ## Are these changes tested? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> As it's just method plumbing, my impression is that having a test for it is overkill. ## Are there any user-facing changes? Users can now implement their `with_new_state()` also in `DataSource`, not only in `ExecutionPlan` <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> (cherry picked from commit 46ac990)
1 parent 42ecad5 commit 2d96876

1 file changed

Lines changed: 28 additions & 0 deletions

File tree

datafusion/datasource/src/source.rs

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,22 @@ pub trait DataSource: Send + Sync + Debug {
211211
fn with_preserve_order(&self, _preserve_order: bool) -> Option<Arc<dyn DataSource>> {
212212
None
213213
}
214+
215+
/// Injects arbitrary run-time state into this DataSource, returning a new instance
216+
/// that incorporates that state *if* it is relevant to the concrete DataSource implementation.
217+
///
218+
/// This is a generic entry point: the `state` can be any type wrapped in
219+
/// `Arc<dyn Any + Send + Sync>`. A data source that cares about the state should
220+
/// down-cast it to the concrete type it expects and, if successful, return a
221+
/// modified copy of itself that captures the provided value. If the state is
222+
/// not applicable, the default behaviour is to return `None` so that parent
223+
/// nodes can continue propagating the attempt further down the plan tree.
224+
fn with_new_state(
225+
&self,
226+
_state: Arc<dyn Any + Send + Sync>,
227+
) -> Option<Arc<dyn DataSource>> {
228+
None
229+
}
214230
}
215231

216232
/// [`ExecutionPlan`] that reads one or more files
@@ -407,6 +423,18 @@ impl ExecutionPlan for DataSourceExec {
407423
as Arc<dyn ExecutionPlan>
408424
})
409425
}
426+
427+
fn with_new_state(
428+
&self,
429+
state: Arc<dyn Any + Send + Sync>,
430+
) -> Option<Arc<dyn ExecutionPlan>> {
431+
self.data_source
432+
.with_new_state(state)
433+
.map(|new_data_source| {
434+
Arc::new(self.clone().with_data_source(new_data_source))
435+
as Arc<dyn ExecutionPlan>
436+
})
437+
}
410438
}
411439

412440
impl DataSourceExec {

0 commit comments

Comments
 (0)