test: Update distributed mode to support localhost workers

Author: geoffreyclaude

Cargo.lock

@@ -52,13 +52,73 @@
 use std::time::Duration;
 
 use datafusion::{common::internal_datafusion_err, error::Result};
-use integration_utils::{init_session, run_traced_query};
+use integration_utils::{DistributedMode, SessionBuilder, run_traced_query};
 use opentelemetry::{KeyValue, trace::TracerProvider};
 use opentelemetry_otlp::WithExportConfig;
 use opentelemetry_sdk::{Resource, trace::Sampler};
 use tracing::{Instrument, Level};
 use tracing_subscriber::{Registry, fmt, prelude::*};
 
+/// Internal helper macro for generating match arms using repetition.
+/// This uses Rust's macro repetition syntax ($(...)*).
+macro_rules! generate_query_span_match {
+    ($i:expr, $mode_name:expr, $query_name:expr, $($num:tt),*) => {
+        match ($i, $mode_name) {
+            $(
+                ($num, "Memory") => tracing::info_span!(
+                    concat!("tpch_query_", stringify!($num), "_memory"),
+                    query = %$query_name,
+                    query_num = $i,
+                    distributed_mode = $mode_name
+                ),
+                ($num, "Localhost") => tracing::info_span!(
+                    concat!("tpch_query_", stringify!($num), "_localhost"),
+                    query = %$query_name,
+                    query_num = $i,
+                    distributed_mode = $mode_name
+                ),
+            )*
+            _ => unreachable!("Invalid query number or mode"),
+        }
+    };
+}
+
+/// Macro to generate span creation for all query/mode combinations.
+/// This is necessary because tracing span names must be compile-time constants.
+///
+/// Usage: create_query_span!(query_num, mode_name, query_name)
+macro_rules! create_query_span {
+    ($i:expr, $mode_name:expr, $query_name:expr) => {
+        generate_query_span_match!(
+            $i,
+            $mode_name,
+            $query_name,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22
+        )
+    };
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     // Initialize tracing infrastructure and obtain a tracer provider.
@@ -74,26 +134,52 @@ async fn main() -> Result<()> {
 }
 
 async fn run_distributed_otlp_example() -> Result<()> {
-    // Loop over all 22 TPCH queries
-    for i in 1..=22 {
-        let query_name = format!("tpch/q{}", i);
-
-        // Create a new root span for each query to ensure independent traces.
-        // This span will be the root of a new trace tree.
-        let span = tracing::info_span!("tpch_query", query = %query_name, query_num = i);
-
-        // Execute the query within the new root span context.
-        async {
-            tracing::info!("Running TPCH query: {}", query_name);
-
-            // Initialize a distinct DataFusion session context for each query.
-            let ctx = init_session(false, true, 5, true, true).await?;
-
-            // Run the SQL query with tracing enabled.
-            run_traced_query(&ctx, &query_name).await
+    // Test both distributed execution modes
+    let modes = [
+        (DistributedMode::Memory, "Memory"),
+        (DistributedMode::Localhost, "Localhost"),
+    ];
+
+    for (mode, mode_name) in modes {
+        tracing::info!("Starting TPCH queries with {} distributed mode", mode_name);
+
+        // Loop over all 22 TPCH queries
+        for i in 1..=22 {
+            let query_name = format!("tpch/q{}", i);
+
+            // Create a new root span for each query to ensure independent traces.
+            // This span will be the root of a new trace tree.
+            // Note: Span names must be compile-time constants, so we use a macro to generate them.
+            let span = create_query_span!(i, mode_name, query_name);
+
+            // Execute the query within the new root span context.
+            async {
+                tracing::info!(
+                    "Running TPCH query: {} in {} mode",
+                    query_name,
+                    mode_name
+                );
+
+                // Initialize a distinct DataFusion session context for each query.
+                let ctx = SessionBuilder::new()
+                    .with_metrics()
+                    .with_preview(5)
+                    .with_compact_preview()
+                    .with_distributed_mode(mode)
+                    .build()
+                    .await?;
+
+                // Run the SQL query with tracing enabled.
+                run_traced_query(&ctx, &query_name).await
+            }
+            .instrument(span)
+            .await?;
         }
-        .instrument(span)
-        .await?;
+
+        tracing::info!(
+            "Completed all TPCH queries with {} distributed mode",
+            mode_name
+        );
     }
 
     Ok(())

examples/distributed_otlp.rs

@@ -53,7 +53,7 @@
 use std::time::Duration;
 
 use datafusion::{common::internal_datafusion_err, error::Result};
-use integration_utils::{init_session, run_traced_query};
+use integration_utils::{SessionBuilder, run_traced_query};
 use opentelemetry::{KeyValue, trace::TracerProvider};
 use opentelemetry_otlp::WithExportConfig;
 use opentelemetry_sdk::{Resource, trace::Sampler};
@@ -80,7 +80,13 @@ 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, false).await?;
+    let ctx = SessionBuilder::new()
+        .with_object_store_tracing()
+        .with_metrics()
+        .with_preview(5)
+        .with_compact_preview()
+        .build()
+        .await?;
 
     // Run the SQL query with tracing enabled.
     run_traced_query(&ctx, QUERY_NAME).await?;

examples/otlp.rs

@@ -34,13 +34,15 @@ workspace = true
 arrow = "56.1"
 arrow-flight = "56.1"
 async-trait = "0.1"
+dashmap = "6.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 }
+structopt = "0.3"
 tokio = { workspace = true, features = ["full"] }
 tokio-stream = "0.1"
 tonic = { version = "0.13", features = ["transport"] }

integration-utils/Cargo.toml

@@ -0,0 +1,77 @@
+// 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.
+
+//! Data and table utilities
+
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use datafusion::common::internal_datafusion_err;
+use datafusion::datasource::file_format::parquet::ParquetFormat;
+use datafusion::datasource::listing::ListingOptions;
+use datafusion::{error::Result, prelude::*};
+use tracing::{info, instrument};
+
+/// 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))]
+pub(crate) async fn register_tpch_tables(ctx: &SessionContext) -> Result<()> {
+    // Construct the path to the directory containing Parquet data.
+    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.
+    for table in [
+        "nation", "region", "part", "supplier", "partsupp", "customer", "orders",
+        "lineitem",
+    ] {
+        let listing_options = ListingOptions::new(Arc::new(ParquetFormat::default()));
+
+        let parquet_path = data_dir.join(table).with_extension("parquet");
+        if !parquet_path.exists() {
+            return Err(internal_datafusion_err!(
+                "Missing TPCH Parquet file: {}.\nGenerate TPCH data first by running: ./dev/generate_tpch_parquet.sh\nThis script requires 'tpchgen-cli' (install with: cargo install tpchgen-cli)",
+                parquet_path.display()
+            ));
+        }
+
+        // Generate the file path URL for the Parquet data.
+        let table_path = format!("file://{}", parquet_path.to_string_lossy());
+
+        info!("Registering table '{}' from {}", table, table_path);
+
+        // Register the table with DataFusion's session context.
+        ctx.register_listing_table(table, &table_path, listing_options, None, None)
+            .await?;
+    }
+
+    Ok(())
+}

integration-utils/src/data.rs

@@ -32,7 +32,7 @@ use tonic::transport::{Endpoint, Server};
 const DUMMY_URL: &str = "http://localhost:50051";
 const MAX_MESSAGE_SIZE: usize = 2 * 1024 * 1024 * 1024; // 2GB
 
-/// [ChannelResolver] implementation that returns gRPC clients baked by an in-memory
+/// [ChannelResolver] implementation that returns gRPC clients backed by an in-memory
 /// tokio duplex rather than a TCP connection.
 #[derive(Clone)]
 pub(crate) struct InMemoryChannelResolver {