jayzhan211 commented on code in PR #43: URL: https://github.com/apache/datafusion-ray/pull/43#discussion_r1838005895
########## src/ext.rs: ########## @@ -0,0 +1,126 @@ +use async_trait::async_trait; +use datafusion::common::DataFusionError; +use datafusion::common::Result; +use datafusion::execution::FunctionRegistry; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use datafusion_python::context::PySessionContext; +use datafusion_python::utils::wait_for_future; +use pyo3::{pyfunction, PyResult, Python}; +use std::collections::HashMap; +use std::fmt::Debug; +use std::sync::{Arc, OnceLock}; + +mod built_in; + +#[cfg(feature = "flight-sql-tables")] +mod flight; + +/// Creates a datafusion session context preconfigured with the enabled extensions +/// that will register additional table providers, catalogs etc. +/// If no extensions are required, the plain `datafusion.SessionContext()` will work just fine. +/// # Arguments +/// * `settings` - dictionary containing extension-specific key/value config options +#[pyfunction] +pub fn extended_session_context( + settings: HashMap<String, String>, + py: Python, +) -> PyResult<PySessionContext> { + let future_context = Extensions::session_context(&settings); + let ctx = wait_for_future(py, future_context)?; + Ok(ctx.into()) +} + +/// Allows third party table/catalog providers, object stores, etc. +/// to be registered with the DataFusion context. +#[async_trait] +trait Extension: Debug + Send + Sync + 'static { + /// SessionContext initialization, using the provided key/value settings if needed. + /// Declared async to allow implementers to perform network or other I/O operations. + async fn init(&self, ctx: &SessionContext, settings: &HashMap<String, String>) -> Result<()> { + let _ = ctx; + let _ = settings; + Ok(()) + } + + /// Codecs for the custom physical plan nodes created by this extension, if any. + fn codecs(&self) -> Vec<Box<dyn PhysicalExtensionCodec>> { + vec![] + } +} + +/// A composite extension registry for enabled extensions. +#[derive(Debug)] +pub(crate) struct Extensions(Box<[Box<dyn Extension>]>); Review Comment: Why not `Vec<Box<dyn Extension>>`? What is the advantage of `Box<[T]>`? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For additional commands, e-mail: github-h...@datafusion.apache.org