feat: add distributed integration test

Author: geoffreyclaude

Cargo.lock

@@ -80,7 +80,7 @@ async fn main() -> Result<()> {
 #[instrument(level = "info")]
 async fn run_otlp_example() -> Result<()> {
     // Initialize the DataFusion session context.
-    let ctx = init_session(true, true, 5, true).await?;
+    let ctx = init_session(true, true, 5, true, false).await?;
 
     // Run the SQL query with tracing enabled.
     run_traced_query(&ctx, QUERY_NAME).await?;

examples/otlp.rs

@@ -31,9 +31,19 @@ rust-version = { workspace = true }
 workspace = true
 
 [dependencies]
+arrow = "56.1"
+arrow-flight = "56.1"
+async-trait = "0.1"
 datafusion = { workspace = true, features = ["parquet", "nested_expressions"] }
+datafusion-distributed = { git = "https://github.com/datafusion-contrib/datafusion-distributed", branch = "main" }
 datafusion-tracing = { workspace = true }
+futures = { workspace = true }
+hyper-util = "0.1"
 instrumented-object-store = { workspace = true }
 object_store = { version = "0.12.1", default-features = false }
+tokio = { workspace = true, features = ["full"] }
+tokio-stream = "0.1"
+tonic = { version = "0.13", features = ["transport"] }
+tower = "0.5"
 tracing = { workspace = true }
 url = { version = "2.5" }

integration-utils/Cargo.toml

@@ -0,0 +1,5 @@
+SELECT
+    count(*),
+    "MinTemp"
+FROM weather
+GROUP BY "MinTemp"

integration-utils/queries/weather.sql

@@ -0,0 +1,98 @@
+// 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.
+//
+// This product includes software developed at Datadog (https://www.datadoghq.com/) Copyright 2025 Datadog, Inc.
+
+use arrow_flight::flight_service_client::FlightServiceClient;
+use arrow_flight::flight_service_server::FlightServiceServer;
+use async_trait::async_trait;
+use datafusion::common::DataFusionError;
+use datafusion::execution::SessionStateBuilder;
+use datafusion_distributed::{
+    ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
+    DistributedSessionBuilderContext,
+};
+use hyper_util::rt::TokioIo;
+use tonic::transport::{Endpoint, Server};
+
+const DUMMY_URL: &str = "http://localhost:50051";
+
+/// [ChannelResolver] implementation that returns gRPC clients baked by an in-memory
+/// tokio duplex rather than a TCP connection.
+#[derive(Clone)]
+pub(crate) struct InMemoryChannelResolver {
+    channel: FlightServiceClient<BoxCloneSyncChannel>,
+}
+
+impl InMemoryChannelResolver {
+    pub fn new() -> Self {
+        let (client, server) = tokio::io::duplex(1024 * 1024);
+
+        let mut client = Some(client);
+        let channel = Endpoint::try_from(DUMMY_URL)
+            .expect(
+                "Invalid dummy URL for building an endpoint. This should never happen",
+            )
+            .connect_with_connector_lazy(tower::service_fn(move |_| {
+                let client = client
+                    .take()
+                    .expect("Client taken twice. This should never happen");
+                async move { Ok::<_, std::io::Error>(TokioIo::new(client)) }
+            }));
+
+        let this = Self {
+            channel: FlightServiceClient::new(BoxCloneSyncChannel::new(channel)),
+        };
+        let this_clone = this.clone();
+
+        let endpoint =
+            ArrowFlightEndpoint::try_new(move |ctx: DistributedSessionBuilderContext| {
+                let this = this.clone();
+                async move {
+                    let builder = SessionStateBuilder::new()
+                        .with_default_features()
+                        .with_distributed_channel_resolver(this)
+                        .with_runtime_env(ctx.runtime_env.clone());
+                    Ok(builder.build())
+                }
+            })
+            .unwrap();
+
+        tokio::spawn(async move {
+            Server::builder()
+                .add_service(FlightServiceServer::new(endpoint))
+                .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
+                .await
+        });
+
+        this_clone
+    }
+}
+
+#[async_trait]
+impl ChannelResolver for InMemoryChannelResolver {
+    fn get_urls(&self) -> Result<Vec<url::Url>, DataFusionError> {
+        Ok(vec![url::Url::parse(DUMMY_URL).unwrap()])
+    }
+
+    async fn get_flight_client_for_url(
+        &self,
+        _: &url::Url,
+    ) -> Result<FlightServiceClient<BoxCloneSyncChannel>, DataFusionError> {
+        Ok(self.channel.clone())
+    }
+}

integration-utils/src/channel_resolver.rs

@@ -61,13 +61,17 @@ use datafusion::{
     error::Result, execution::SessionStateBuilder,
     physical_optimizer::PhysicalOptimizerRule, prelude::*,
 };
+use datafusion_distributed::{DistributedExt, DistributedPhysicalOptimizerRule};
 use datafusion_tracing::{
     InstrumentationOptions, instrument_with_info_spans, pretty_format_compact_batch,
 };
 use instrumented_object_store::instrument_object_store;
 use tracing::{field, info, instrument};
 use url::Url;
 
+mod channel_resolver;
+use channel_resolver::InMemoryChannelResolver;
+
 /// Executes the SQL query with instrumentation enabled, providing detailed tracing output.
 #[instrument(level = "info", skip(ctx))]
 pub async fn run_traced_query(ctx: &SessionContext, query_name: &str) -> Result<()> {
@@ -95,17 +99,23 @@ pub async fn init_session(
     record_metrics: bool,
     preview_limit: usize,
     compact_preview: bool,
+    distributed: bool,
 ) -> Result<SessionContext> {
     // Configure the session state with instrumentation for query execution.
-    let session_state = SessionStateBuilder::new()
+    let mut session_state_builder = SessionStateBuilder::new()
         .with_default_features()
-        .with_config(SessionConfig::default().with_target_partitions(8)) // Enforce target partitions to ensure consistent test results regardless of the number of CPU cores.
-        .with_physical_optimizer_rule(create_instrumentation_rule(
-            record_metrics,
-            preview_limit,
-            compact_preview,
-        ))
-        .build();
+        .with_config(SessionConfig::default().with_target_partitions(8)); // Enforce target partitions to ensure consistent test results regardless of the number of CPU cores.
+    if distributed {
+        session_state_builder = session_state_builder
+            .with_distributed_channel_resolver(InMemoryChannelResolver::new())
+            .with_physical_optimizer_rule(Arc::new(
+                DistributedPhysicalOptimizerRule::new(),
+            ));
+    }
+    session_state_builder = session_state_builder.with_physical_optimizer_rule(
+        create_instrumentation_rule(record_metrics, preview_limit, compact_preview),
+    );
+    let session_state = session_state_builder.build();
 
     let ctx = SessionContext::new_with_state(session_state);
 
@@ -153,16 +163,24 @@ pub fn create_instrumentation_rule(
     )
 }
 
-/// Returns the path to the directory containing the TPCH Parquet tables.
-pub fn tpch_tables_dir() -> PathBuf {
+/// Returns the path to the directory containing the Parquet tables.
+pub fn data_dir() -> PathBuf {
     PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("data")
 }
 
 /// Registers all TPCH Parquet tables required for executing the queries.
 #[instrument(level = "info", skip(ctx))]
 async fn register_tpch_tables(ctx: &SessionContext) -> Result<()> {
     // Construct the path to the directory containing Parquet data.
-    let data_dir = tpch_tables_dir();
+    let data_dir = data_dir();
+
+    // Register the weather table.
+    ctx.register_parquet(
+        "weather",
+        data_dir.join("weather").to_string_lossy(),
+        ParquetReadOptions::default(),
+    )
+    .await?;
 
     // Generate and register each table from Parquet files.
     // This includes all standard TPCH tables so examples/tests can rely on them.

integration-utils/src/lib.rs

@@ -56,6 +56,8 @@ struct QueryTestCase<'a> {
     ignored_preview_spans: &'a [usize],
     /// Whether to ignore the full trace in assertions.
     ignore_full_trace: bool,
+    /// Whether to run the test in distributed mode.
+    distributed: bool,
 }
 
 impl<'a> QueryTestCase<'a> {
@@ -95,6 +97,11 @@ impl<'a> QueryTestCase<'a> {
         self.ignore_full_trace = true;
         self
     }
+
+    fn distributed(mut self) -> Self {
+        self.distributed = true;
+        self
+    }
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
@@ -196,6 +203,20 @@ async fn test_topk_lineitem() -> Result<()> {
     execute_test_case("10_topk_lineitem", &QueryTestCase::new("topk_lineitem")).await
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
+async fn test_weather() -> Result<()> {
+    execute_test_case("11_weather", &QueryTestCase::new("weather")).await
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
+async fn test_distributed_weather() -> Result<()> {
+    execute_test_case(
+        "12_distributed_weather",
+        &QueryTestCase::new("weather").distributed(),
+    )
+    .await
+}
+
 /// Executes the provided [`QueryTestCase`], setting up tracing and verifying
 /// log output according to its parameters.
 async fn execute_test_case(test_name: &str, test_case: &QueryTestCase<'_>) -> Result<()> {
@@ -208,6 +229,7 @@ async fn execute_test_case(test_name: &str, test_case: &QueryTestCase<'_>) -> Re
         test_case.should_record_metrics,
         test_case.row_limit,
         test_case.use_compact_preview,
+        test_case.distributed,
     )
     .await?;