Integrate datafusion-distributed with Python

Cargo.toml

@@ -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" }

crates/core/Cargo.toml

@@ -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"] }

crates/core/src/context.rs

@@ -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
  }

crates/core/src/distributed_worker.rs

@@ -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")
+ && 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)))
+}

crates/core/src/distributed_worker_resolver.rs

@@ -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)))
+ }
+}