apache / apache/datafusion

Pass SessionConfig extensions to ObjectStore used during physical plan execution

Open
#14,804 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

### Is your feature request related to a problem or challenge?

In https://github.com/apache/arrow-rs-object-store/issues/17 I've described a general need for the `ObjectStore` trait to be able to support passing contextual data to custom implementations. In https://github.com/apache/arrow-rs/pull/7160 I have implemented and approach to this by providing the ability for `GetOptions` to store opaque instances of values indexed by their `TypeId`, [similar to what is possible in datafusion with `SessionConfig`](https://docs.rs/datafusion/latest/datafusion/prelude/struct.SessionConfig.html#method.with_extension).

This issue is about taking incorporating this new behavior(s) in `ObjectStore` and incorporating it here in datafusion such that the custom data on a `SessionConfig` is passed on when creating `GetOptions`s instances for retrieving files from an object store.

### Describe the solution you'd like

I think the simplest approach here would be one where we create a new `ObjectStore` implementation during query execution that looks something like:

```
struct ContextualizedObjectStore {
inner: Arc,
extensions: object_store::Extensions,
}
```

We would then have a `get_opts` method for the `ObjectStore` impl trait that looks something like:

```
async fn get_opts(
&self,
location: &Path,
mut options: GetOptions,
) -> object_store::Result {
options.extensions = self.extensions.clone();
self.inner.get_opts(location, options).await
}
```

Initializing instances of this new type as a wrapper around whatever given `Arc` is available would look something like:

```
let object_store = context
.runtime_env()
.object_store(&self.object_store_url)
.map(|inner: Arc| -> Arc {
Arc::new(ContextualizedObjectStore::new(
inner,
context.session_config().clone_extensions_for_object_store(),
))
});
```

With this approach, whenever the resulting `Arc` is used to retrieve a file from object store, the underlying implementation would have access to the `object_store::Extensions` created from the `SessionConfig` extensions.

### Describe alternatives you've considered

This is covered in https://github.com/apache/arrow-rs-object-store/issues/17 and https://github.com/apache/arrow-rs/issues/7135.

Basically, there are two alternative directions:

* Update the `ObjectStore` API by providing optional trait methods that take an actual context type that can carry custom/extension data.
* Considered by maintainers to be too heavy-handed.
* Don't do anything.
* This means for my use case, we wouldn't be able to properly parent tracing spans for object store accesses that happen during query execution.

### Additional context

_No response_

Contributor guide

Open the contributing guide

Research direction

Start by tracing physical plan execution where ObjectStore instances create GetOptions, then read the SessionConfig extension APIs and the linked object-store context discussion. Done means extensions from SessionConfig reach GetOptions used for object-store retrieval without changing the underlying custom ObjectStore behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
data-engineering
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.