apache
diff --git a/‎Cargo.lock‎
Lines changed: 321 additions & 0 deletions b/‎Cargo.lock‎
Lines changed: 321 additions & 0 deletions
diff --git a/‎Cargo.toml‎
Lines changed: 2 additions & 0 deletions b/‎Cargo.toml‎
Lines changed: 2 additions & 0 deletions
diff --git a/‎crates/core/Cargo.toml‎
Lines changed: 2 additions & 0 deletions b/‎crates/core/Cargo.toml‎
Lines changed: 2 additions & 0 deletions
diff --git a/‎crates/core/src/context.rs‎
Lines changed: 30 additions & 4 deletions b/‎crates/core/src/context.rs‎
Lines changed: 30 additions & 4 deletions
diff --git a/‎crates/core/src/distributed_worker.rs‎
Lines changed: 206 additions & 0 deletions b/‎crates/core/src/distributed_worker.rs‎
Lines changed: 206 additions & 0 deletions
diff --git a/‎crates/core/src/distributed_worker_resolver.rs‎
Lines changed: 84 additions & 0 deletions b/‎crates/core/src/distributed_worker_resolver.rs‎
Lines changed: 84 additions & 0 deletions
@@ -32,6 +32,7 @@ resolver = "3"
 
 [workspace.dependencies]
 tokio = { version = "1.52" }
+tonic = { version = "0.14", features = ["transport"] }
 pyo3 = { version = "0.28" }
 pyo3-async-runtimes = { version = "0.28" }
 pyo3-log = "0.13.3"
@@ -50,6 +51,7 @@ datafusion-functions-aggregate = { version = "54" }
 datafusion-functions-window = { version = "54" }
 datafusion-spark = { version = "54" }
 datafusion-expr = { version = "54" }
+datafusion-distributed = { version = "2" }
 prost = "0.14.3"
 serde_json = "1"
 uuid = { version = "1.23" }
 
@@ -40,6 +40,7 @@ tokio = { workspace = true, features = [
   "rt-multi-thread",
   "sync",
 ] }
+tonic = { workspace = true }
 pyo3 = { workspace = true, features = [
   "extension-module",
   "generate-import-lib",
@@ -54,6 +55,7 @@ datafusion-substrait = { workspace = true, optional = true }
 datafusion-proto = { workspace = true }
 datafusion-ffi = { workspace = true }
 datafusion-spark = { workspace = true }
+datafusion-distributed = { workspace = true }
 prost = { workspace = true } # keep in line with `datafusion-substrait`
 serde_json = { workspace = true }
 uuid = { workspace = true, features = ["v4"] }
 
@@ -43,10 +43,11 @@ use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool, Unboun
 use datafusion::execution::options::{ArrowReadOptions, ReadOptions};
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::execution::session_state::SessionStateBuilder;
-use datafusion::execution::{FunctionRegistry, TaskContextProvider};
+use datafusion::execution::{FunctionRegistry, SessionState, TaskContextProvider};
 use datafusion::prelude::{
     AvroReadOptions, CsvReadOptions, DataFrame, JsonReadOptions, ParquetReadOptions,
 };
+use datafusion_distributed::{DistributedConfig, DistributedExt, SessionStateBuilderExt};
 use datafusion_ffi::catalog_provider::FFI_CatalogProvider;
 use datafusion_ffi::catalog_provider_list::FFI_CatalogProviderList;
 use datafusion_ffi::config::extension_options::FFI_ExtensionOptions;
@@ -78,6 +79,7 @@ use crate::common::data_type::PyScalarValue;
 use crate::common::df_schema::PyDFSchema;
 use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
+use crate::distributed_worker_resolver::PyWorkerResolver;
 use crate::errors::{
     PyDataFusionError, PyDataFusionResult, from_datafusion_error, py_datafusion_err,
 };
@@ -219,6 +221,15 @@ impl PySessionConfig {
 
         Ok(Self::from(config))
     }
+
+    #[pyo3(signature = (worker_resolver))]
+    fn with_distributed(&self, worker_resolver: PyWorkerResolver) -> Self {
+        let config = self
+            .config
+            .clone()
+            .with_distributed_worker_resolver(worker_resolver);
+        Self::from(config)
+    }
 }
 
 /// Runtime options for a SessionContext
@@ -392,13 +403,20 @@ impl PySessionContext {
         } else {
             RuntimeEnvBuilder::default()
         };
+        let distributed = DistributedConfig::from_config_options(config.options()).is_ok();
+
         let runtime = Arc::new(runtime_env_builder.build()?);
-        let session_state = SessionStateBuilder::new()
+        let mut builder = SessionStateBuilder::new()
             .with_config(config)
             .with_runtime_env(runtime)
             .with_default_features()
-            .with_analyzer_rule(Arc::new(crate::analyzer::ResolveLambdaVariables::new()))
-            .build();
+            .with_analyzer_rule(Arc::new(crate::analyzer::ResolveLambdaVariables::new()));
+
+        if distributed {
+            builder = builder.with_distributed_planner();
+        }
+
+        let session_state = builder.build();
         let ctx = Arc::new(SessionContext::new_with_state(session_state));
         Ok(PySessionContext {
             ctx,
@@ -1430,6 +1448,14 @@ impl PySessionContext {
 }
 
 impl PySessionContext {
+    pub(crate) fn from_session_state(session_state: SessionState) -> Self {
+        Self {
+            ctx: Arc::new(SessionContext::new_with_state(session_state)),
+            logical_codec: Arc::new(PythonLogicalCodec::default()),
+            physical_codec: Arc::new(PythonPhysicalCodec::default()),
+        }
+    }
+
     async fn _table(&self, name: &str) -> datafusion::common::Result<DataFrame> {
         self.ctx.table(name).await
     }
 
@@ -0,0 +1,206 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::HashMap;
+use std::net::SocketAddr;
+
+use async_trait::async_trait;
+use datafusion::common::{DataFusionError, Result as DataFusionResult};
+use datafusion::execution::{SessionState, SessionStateBuilder};
+use datafusion_distributed::{Worker, WorkerQueryContext, WorkerSessionBuilder};
+use datafusion_python_util::wait_for_future;
+use pyo3::Borrowed;
+use pyo3::exceptions::{PyRuntimeError, PyTypeError};
+use pyo3::prelude::*;
+use tonic::transport::Server;
+
+use crate::context::PySessionContext;
+use crate::errors::{PyDataFusionError, PyDataFusionResult};
+
+#[pyclass(
+    from_py_object,
+    frozen,
+    name = "Worker",
+    module = "datafusion",
+    subclass
+)]
+#[derive(Clone)]
+pub struct PyWorker {
+    worker: Worker,
+}
+
+#[pymethods]
+impl PyWorker {
+    #[new]
+    fn new() -> Self {
+        Self {
+            worker: Worker::default(),
+        }
+    }
+
+    #[staticmethod]
+    fn from_session_builder(session_builder: PyWorkerSessionBuilder) -> Self {
+        Self {
+            worker: Worker::from_session_builder(session_builder),
+        }
+    }
+
+    fn with_version(&self, version: String) -> Self {
+        Self {
+            worker: self.worker.clone().with_version(version),
+        }
+    }
+
+    fn with_max_message_size(&self, size: usize) -> Self {
+        Self {
+            worker: self.worker.clone().with_max_message_size(size),
+        }
+    }
+
+    #[pyo3(signature = (host = "127.0.0.1", port = 50051))]
+    fn serve(&self, py: Python<'_>, host: &str, port: u16) -> PyDataFusionResult<()> {
+        let addr = parse_socket_addr(host, port)?;
+        let worker = self.worker.clone();
+        wait_for_future(py, serve_worker(worker, addr))?.map_err(PyDataFusionError::from)
+    }
+
+    #[pyo3(signature = (host = "127.0.0.1", port = 50051))]
+    fn serve_async<'py>(
+        &self,
+        py: Python<'py>,
+        host: &str,
+        port: u16,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let addr = parse_socket_addr(host, port)?;
+        let worker = self.worker.clone();
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            serve_worker(worker, addr)
+                .await
+                .map_err(PyDataFusionError::from)?;
+            Ok(())
+        })
+    }
+}
+
+#[pyclass(name = "WorkerQueryContext", module = "datafusion", subclass)]
+pub struct PyWorkerQueryContext {
+    builder: Option<SessionStateBuilder>,
+    headers: HashMap<String, String>,
+}
+
+impl PyWorkerQueryContext {
+    fn new(ctx: WorkerQueryContext) -> Self {
+        let headers = ctx
+            .headers
+            .iter()
+            .map(|(name, value)| {
+                (
+                    name.as_str().to_owned(),
+                    value.to_str().unwrap_or_default().to_owned(),
+                )
+            })
+            .collect();
+
+        Self {
+            builder: Some(ctx.builder),
+            headers,
+        }
+    }
+}
+
+#[pymethods]
+impl PyWorkerQueryContext {
+    fn session_context(mut slf: PyRefMut<'_, Self>) -> PyResult<PySessionContext> {
+        let builder = slf.builder.take().ok_or_else(|| {
+            PyRuntimeError::new_err("WorkerQueryContext.session_context() can only be called once")
+        })?;
+        Ok(PySessionContext::from_session_state(builder.build()))
+    }
+
+    #[getter]
+    fn headers(&self) -> HashMap<String, String> {
+        self.headers.clone()
+    }
+}
+
+pub(crate) struct PyWorkerSessionBuilder {
+    callback: Py<PyAny>,
+}
+
+impl FromPyObject<'_, '_> for PyWorkerSessionBuilder {
+    type Error = PyErr;
+
+    fn extract(obj: Borrowed<'_, '_, PyAny>) -> Result<Self, Self::Error> {
+        if !obj.is_callable() {
+            return Err(PyTypeError::new_err(
+                "Expected worker session builder to be callable",
+            ));
+        }
+
+        Ok(Self {
+            callback: obj.to_owned().unbind(),
+        })
+    }
+}
+
+#[async_trait]
+impl WorkerSessionBuilder for PyWorkerSessionBuilder {
+    async fn build_session_state(
+        &self,
+        ctx: WorkerQueryContext,
+    ) -> Result<SessionState, DataFusionError> {
+        Python::attach(|py| -> PyResult<SessionState> {
+            let ctx = Py::new(py, PyWorkerQueryContext::new(ctx))?;
+            let result = self.callback.call1(py, (ctx,))?;
+            let session_context = extract_session_context(result.bind(py))?;
+            Ok(session_context.ctx.state())
+        })
+        .map_err(|error| DataFusionError::External(Box::new(error)))
+    }
+}
+
+fn extract_session_context(obj: &Bound<'_, PyAny>) -> PyResult<PySessionContext> {
+    if let Ok(session_context) = obj.extract::<PySessionContext>() {
+        return Ok(session_context);
+    }
+
+    if let Ok(ctx_attr) = obj.getattr("ctx") {
+        if let Ok(session_context) = ctx_attr.extract::<PySessionContext>() {
+            return Ok(session_context);
+        }
+    }
+
+    Err(PyTypeError::new_err(
+        "WorkerSessionBuilder.build_session_state() must return a datafusion.SessionContext",
+    ))
+}
+
+fn parse_socket_addr(host: &str, port: u16) -> PyDataFusionResult<SocketAddr> {
+    format!("{host}:{port}").parse().map_err(|error| {
+        PyDataFusionError::Common(format!(
+            "invalid worker bind address {host}:{port}: {error}"
+        ))
+    })
+}
+
+async fn serve_worker(worker: Worker, addr: SocketAddr) -> DataFusionResult<()> {
+    Server::builder()
+        .add_service(worker.into_worker_server())
+        .serve(addr)
+        .await
+        .map_err(|error| DataFusionError::External(Box::new(error)))
+}
@@ -0,0 +1,84 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use datafusion::common::DataFusionError;
+use datafusion_distributed::WorkerResolver;
+use pyo3::Borrowed;
+use pyo3::exceptions::{PyTypeError, PyValueError};
+use pyo3::prelude::*;
+use pyo3::types::PyString;
+use url::Url;
+
+pub(crate) struct PyWorkerResolver {
+    get_urls: Py<PyAny>,
+}
+
+impl FromPyObject<'_, '_> for PyWorkerResolver {
+    type Error = PyErr;
+
+    fn extract(obj: Borrowed<'_, '_, PyAny>) -> Result<Self, Self::Error> {
+        let get_urls = obj.getattr("get_urls")?;
+        if !get_urls.is_callable() {
+            return Err(PyTypeError::new_err(
+                "Expected worker_resolver.get_urls to be callable",
+            ));
+        }
+
+        Ok(Self {
+            get_urls: get_urls.unbind(),
+        })
+    }
+}
+
+struct WorkerUrls(Vec<Url>);
+
+impl FromPyObject<'_, '_> for WorkerUrls {
+    type Error = PyErr;
+
+    fn extract(obj: Borrowed<'_, '_, PyAny>) -> Result<Self, Self::Error> {
+        if obj.is_instance_of::<PyString>() {
+            return Err(PyTypeError::new_err(
+                "WorkerResolver.get_urls() must return an iterable of URL strings, not a string",
+            ));
+        }
+
+        let mut parsed_urls = Vec::new();
+        for url in obj.try_iter()? {
+            let url = url?;
+            let url = url.extract::<String>()?;
+            let parsed_url = Url::parse(&url).map_err(|error| {
+                PyValueError::new_err(format!(
+                    "WorkerResolver.get_urls() returned invalid URL {url:?}: {error}"
+                ))
+            })?;
+            parsed_urls.push(parsed_url);
+        }
+
+        Ok(Self(parsed_urls))
+    }
+}
+
+impl WorkerResolver for PyWorkerResolver {
+    fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
+        Python::attach(|py| -> PyResult<Vec<Url>> {
+            let urls = self.get_urls.call0(py)?;
+            let urls = urls.extract::<WorkerUrls>(py)?;
+            Ok(urls.0)
+        })
+        .map_err(|error| DataFusionError::External(Box::new(error)))
+    }
+}