test: add new distributed TPCH example

Author: geoffreyclaude

dev/generate_tpch_parquet.sh

@@ -33,10 +33,10 @@ if ! command -v tpchgen-cli >/dev/null 2>&1; then
   exit 1
 fi
 
-echo "Generating TPCH Parquet data (SF=0.1) into $OUT_DIR"
+echo "Generating TPCH Parquet data (SF=1) into $OUT_DIR"
 
 # Generate all tables, single file per table, Parquet format
-tpchgen-cli --scale-factor 0.1 --format parquet --output-dir "$OUT_DIR"
+tpchgen-cli --scale-factor 1 --format parquet --output-dir "$OUT_DIR"
 
 echo "TPCH Parquet generation complete. Files in $OUT_DIR:"
 ls -1 "$OUT_DIR" | sed 's/^/ - /'

examples/Cargo.toml

@@ -35,6 +35,10 @@ workspace = true
 name = "otlp"
 path = "otlp.rs"
 
+[[example]]
+name = "distributed_otlp"
+path = "distributed_otlp.rs"
+
 [dependencies]
 datafusion = { workspace = true, features = ["parquet", "nested_expressions"] }
 integration-utils = { path = "../integration-utils" }

examples/distributed_otlp.rs

@@ -0,0 +1,146 @@
+// 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.
+
+//! # Distributed OTLP Example
+//!
+//! This example demonstrates OpenTelemetry integration with DataFusion for distributed tracing.
+//! It shows how to:
+//!
+//! - Configure OpenTelemetry with an OTLP exporter
+//! - Set up service metadata and tracing layers
+//! - Instrument DataFusion query execution for distributed plans
+//! - Execute multiple queries with separate contexts
+//!
+//! ## Prerequisites
+//!
+//! Before running this example, you'll need an OpenTelemetry collector. For local development,
+//! Jaeger is recommended:
+//!
+//! ```bash
+//! docker run --rm --name jaeger \
+//!   -p 16686:16686 \
+//!   -p 4317:4317 \
+//!   jaegertracing/jaeger:2.7.0
+//! ```
+//!
+//! After starting Jaeger, you can view traces at http://localhost:16686
+//!
+//! ## Running the Example
+//!
+//! ```bash
+//! cargo run --example distributed_otlp
+//! ```
+//!
+//! This example executes all 22 TPCH benchmark queries with distributed query execution enabled.
+
+use std::time::Duration;
+
+use datafusion::{common::internal_datafusion_err, error::Result};
+use integration_utils::{init_session, 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::*};
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    // Initialize tracing infrastructure and obtain a tracer provider.
+    let tracer_provider = init_tracing()?;
+
+    // Run the example under the root span.
+    run_distributed_otlp_example().await?;
+
+    // Properly shutdown tracing to ensure all data is flushed.
+    tracer_provider
+        .shutdown()
+        .map_err(|e| internal_datafusion_err!("Tracer shutdown error: {}", e))
+}
+
+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
+        }
+        .instrument(span)
+        .await?;
+    }
+
+    Ok(())
+}
+
+/// Initializes OpenTelemetry and tracing infrastructure to enable tracing of query execution.
+fn init_tracing() -> Result<opentelemetry_sdk::trace::SdkTracerProvider> {
+    // Set service metadata for tracing.
+    let resource = Resource::builder()
+        .with_attribute(KeyValue::new("service.name", "datafusion-tracing"))
+        .build();
+
+    // Configure an OTLP exporter to send tracing data.
+    let exporter = opentelemetry_otlp::SpanExporter::builder()
+        .with_tonic()
+        .with_endpoint("http://localhost:4317") // Endpoint for OTLP collector.
+        .with_timeout(Duration::from_secs(10))
+        .build()
+        .map_err(|e| internal_datafusion_err!("OTLP exporter error: {}", e))?;
+
+    // Create a tracer provider configured with the exporter and sampling strategy.
+    let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
+        .with_batch_exporter(exporter)
+        .with_resource(resource)
+        .with_sampler(Sampler::AlwaysOn)
+        .build();
+
+    // Obtain a tracer instance for recording tracing information.
+    let tracer = tracer_provider.tracer("datafusion-tracing-query");
+
+    // Create a telemetry layer using the tracer to collect and filter tracing data at INFO level.
+    let telemetry_layer = tracing_opentelemetry::layer()
+        .with_tracer(tracer)
+        .with_filter(tracing::level_filters::LevelFilter::INFO);
+
+    // Create a formatting layer to output logs to stdout, including thread IDs and names.
+    let fmt_layer = fmt::layer()
+        .with_thread_ids(true)
+        .with_thread_names(true)
+        .with_writer(std::io::stdout.with_max_level(Level::INFO));
+
+    // Combine the telemetry and formatting layers into a tracing subscriber and initialize it.
+    Registry::default()
+        .with(telemetry_layer)
+        .with(fmt_layer)
+        .init();
+
+    // Return the configured tracer provider
+    Ok(tracer_provider)
+}

integration-utils/queries/tpch/q1.sql

@@ -0,0 +1,21 @@
+select
+    l_returnflag,
+    l_linestatus,
+    sum(l_quantity) as sum_qty,
+    sum(l_extendedprice) as sum_base_price,
+    sum(l_extendedprice * (1 - l_discount)) as sum_disc_price,
+    sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge,
+    avg(l_quantity) as avg_qty,
+    avg(l_extendedprice) as avg_price,
+    avg(l_discount) as avg_disc,
+    count(*) as count_order
+from
+    lineitem
+where
+        l_shipdate <= date '1998-09-02'
+group by
+    l_returnflag,
+    l_linestatus
+order by
+    l_returnflag,
+    l_linestatus;

integration-utils/queries/tpch/q10.sql

@@ -0,0 +1,31 @@
+select
+    c_custkey,
+    c_name,
+    sum(l_extendedprice * (1 - l_discount)) as revenue,
+    c_acctbal,
+    n_name,
+    c_address,
+    c_phone,
+    c_comment
+from
+    customer,
+    orders,
+    lineitem,
+    nation
+where
+        c_custkey = o_custkey
+  and l_orderkey = o_orderkey
+  and o_orderdate >= date '1993-10-01'
+  and o_orderdate < date '1994-01-01'
+  and l_returnflag = 'R'
+  and c_nationkey = n_nationkey
+group by
+    c_custkey,
+    c_name,
+    c_acctbal,
+    c_phone,
+    n_name,
+    c_address,
+    c_comment
+order by
+    revenue desc;

integration-utils/queries/tpch/q11.sql

@@ -0,0 +1,27 @@
+select
+    ps_partkey,
+    sum(ps_supplycost * ps_availqty) as value
+from
+    partsupp,
+    supplier,
+    nation
+where
+    ps_suppkey = s_suppkey
+  and s_nationkey = n_nationkey
+  and n_name = 'GERMANY'
+group by
+    ps_partkey having
+    sum(ps_supplycost * ps_availqty) > (
+    select
+    sum(ps_supplycost * ps_availqty) * 0.0001
+    from
+    partsupp,
+    supplier,
+    nation
+    where
+    ps_suppkey = s_suppkey
+                  and s_nationkey = n_nationkey
+                  and n_name = 'GERMANY'
+    )
+order by
+    value desc;

integration-utils/queries/tpch/q12.sql

@@ -0,0 +1,30 @@
+select
+    l_shipmode,
+    sum(case
+            when o_orderpriority = '1-URGENT'
+                or o_orderpriority = '2-HIGH'
+                then 1
+            else 0
+        end) as high_line_count,
+    sum(case
+            when o_orderpriority <> '1-URGENT'
+                and o_orderpriority <> '2-HIGH'
+                then 1
+            else 0
+        end) as low_line_count
+from
+    lineitem
+        join
+    orders
+    on
+            l_orderkey = o_orderkey
+where
+        l_shipmode in ('MAIL', 'SHIP')
+  and l_commitdate < l_receiptdate
+  and l_shipdate < l_commitdate
+  and l_receiptdate >= date '1994-01-01'
+  and l_receiptdate < date '1995-01-01'
+group by
+    l_shipmode
+order by
+    l_shipmode;

integration-utils/queries/tpch/q13.sql

@@ -0,0 +1,20 @@
+select
+    c_count,
+    count(*) as custdist
+from
+    (
+        select
+            c_custkey,
+            count(o_orderkey)
+        from
+            customer left outer join orders on
+                        c_custkey = o_custkey
+                    and o_comment not like '%special%requests%'
+        group by
+            c_custkey
+    ) as c_orders (c_custkey, c_count)
+group by
+    c_count
+order by
+    custdist desc,
+    c_count desc;

integration-utils/queries/tpch/q14.sql

@@ -0,0 +1,13 @@
+select
+            100.00 * sum(case
+                             when p_type like 'PROMO%'
+                                 then l_extendedprice * (1 - l_discount)
+                             else 0
+            end) / sum(l_extendedprice * (1 - l_discount)) as promo_revenue
+from
+    lineitem,
+    part
+where
+        l_partkey = p_partkey
+  and l_shipdate >= date '1995-09-01'
+  and l_shipdate < date '1995-10-01';

integration-utils/queries/tpch/q15.sql

@@ -0,0 +1,31 @@
+with revenue0 (supplier_no, total_revenue) as (
+	select
+		l_suppkey,
+		sum(l_extendedprice * (1 - l_discount))
+	from
+		lineitem
+	where
+		l_shipdate >= date '1996-01-01'
+		and l_shipdate < date '1996-01-01' + interval '3' month
+	group by
+		l_suppkey
+)
+select
+	s_suppkey,
+	s_name,
+	s_address,
+	s_phone,
+	total_revenue
+from
+	supplier,
+	revenue0
+where
+	s_suppkey = supplier_no
+	and total_revenue = (
+		select
+			max(total_revenue)
+		from
+			revenue0
+	)
+order by
+	s_suppkey;