ccciudatu commented on code in PR #43:
URL: https://github.com/apache/datafusion-ray/pull/43#discussion_r1838036671
##########
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:
It's a bit lighter (hardly relevant in this scenario) and it's immutable by
design (arguably also irrelevant, though I personally prefer it when I can
choose). I think I got this habit after watching [this
video](https://www.youtube.com/watch?v=A4cKi7PTJSs). :)
But I don't mind switching to `Vec` if you think that would improve
readability.
--
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
