move validation to integration tests

robtandy · robtandy · commit 5ad36d822acf · 2025-08-12T15:05:37.000-04:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,7 +1,5 @@
 [workspace]
-members = [
-    "benchmarks"
-]
+members = ["benchmarks"]
 
 [workspace.dependencies]
 datafusion = { version = "49.0.0" }
@@ -37,14 +35,16 @@ tpchgen = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "c8d82343252
 tpchgen-arrow = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "c8d823432528eed4f70fca5a1296a66c68a389a8", optional = true }
 parquet = { version = "55.2.0", optional = true }
 arrow = { version = "55.2.0", optional = true }
+tokio-stream = { version = "0.1.17", optional = true }
 
 [features]
 integration = [
     "insta",
     "tpchgen",
-    "tpchgen-arrow", 
+    "tpchgen-arrow",
     "parquet",
-    "arrow"
+    "arrow",
+    "tokio-stream",
 ]
 
 [dev-dependencies]
diff --git a/benchmarks/src/tpch/run.rs b/benchmarks/src/tpch/run.rs
@@ -103,9 +103,6 @@ pub struct RunOpt {
     /// The maximum number of partitions per task.
     #[structopt(long = "ppt")]
     partitions_per_task: Option<usize>,
-
-    #[structopt(long = "validate")]
-    validate: bool,
 }
 
 #[async_trait]
@@ -157,7 +154,7 @@ impl SessionBuilder for RunOpt {
 
 impl RunOpt {
     pub async fn run(self) -> Result<()> {
-        let (ctx, _guard) = start_localhost_context([40051], self.clone()).await;
+        let (ctx, _guard) = start_localhost_context(1, self.clone()).await;
         println!("Running benchmarks with the following options: {self:?}");
         let query_range = match self.query {
             Some(query_id) => query_id..=query_id,
@@ -225,42 +222,11 @@ impl RunOpt {
                 "Query {query_id} iteration {i} took {ms:.1} ms and returned {row_count} rows"
             );
 
-            let valid = if self.validate {
-                let mut single_node_result = vec![];
-                for (i, query) in sql.iter().enumerate() {
-                    if i == result_stmt {
-                        single_node_result = self.execute_query(&single_node_ctx, query).await?;
-                    } else {
-                        self.execute_query(&single_node_ctx, query).await?;
-                    }
-                }
-
-                let res = pretty::pretty_format_batches(&result)?.to_string();
-                let single_node_res =
-                    pretty::pretty_format_batches(&single_node_result)?.to_string();
-                res == single_node_res
-            } else {
-                false
-            };
-
-            query_results.push(QueryResult {
-                valid,
-                elapsed,
-                row_count,
-            });
+            query_results.push(QueryResult { elapsed, row_count });
         }
 
         let avg = millis.iter().sum::<f64>() / millis.len() as f64;
-        let valid_str = if self.validate {
-            if query_results[query_results.len() - 1].valid {
-                "valid"
-            } else {
-                "invalid"
-            }
-        } else {
-            ""
-        };
-        println!("Query {query_id} avg time: {avg:.2} ms {valid_str}");
+        println!("Query {query_id} avg time: {avg:.2} ms");
 
         // Print memory stats using mimalloc (only when compiled with --features mimalloc_extended)
         print_memory_stats();
diff --git a/benchmarks/src/util/run.rs b/benchmarks/src/util/run.rs
@@ -96,7 +96,6 @@ pub struct BenchQuery {
 pub struct QueryResult {
     pub elapsed: Duration,
     pub row_count: usize,
-    pub valid: bool,
 }
 /// collects benchmark run data and then serializes it at the end
 pub struct BenchmarkRun {
diff --git a/src/errors/arrow_error.rs b/src/errors/arrow_error.rs
@@ -224,8 +224,8 @@ mod tests {
             let (recovered_error, recovered_ctx) = proto.to_arrow_error();
 
             if original_error.to_string() != recovered_error.to_string() {
-                println!("original error: {}", original_error.to_string());
-                println!("recovered error: {}", recovered_error.to_string());
+                println!("original error: {}", original_error);
+                println!("recovered error: {}", recovered_error);
             }
 
             assert_eq!(original_error.to_string(), recovered_error.to_string());
diff --git a/src/errors/mod.rs b/src/errors/mod.rs
@@ -17,8 +17,8 @@ mod schema_error;
 pub fn datafusion_error_to_tonic_status(err: &DataFusionError) -> tonic::Status {
     let err = DataFusionErrorProto::from_datafusion_error(err);
     let err = err.encode_to_vec();
-    let status = tonic::Status::with_details(tonic::Code::Internal, "DataFusionError", err.into());
-    status
+    
+    tonic::Status::with_details(tonic::Code::Internal, "DataFusionError", err.into())
 }
 
 /// Decodes a [DataFusionError] from a [tonic::Status] error. If the provided [tonic::Status]
diff --git a/src/errors/schema_error.rs b/src/errors/schema_error.rs
@@ -198,7 +198,7 @@ impl SchemaErrorProto {
                 valid_fields,
             } => SchemaErrorProto {
                 inner: Some(SchemaErrorInnerProto::FieldNotFound(FieldNotFoundProto {
-                    field: Some(Box::new(ColumnProto::from_column(&field))),
+                    field: Some(Box::new(ColumnProto::from_column(field))),
                     valid_fields: valid_fields.iter().map(ColumnProto::from_column).collect(),
                 })),
                 backtrace: backtrace.cloned(),
diff --git a/src/flight_service/do_get.rs b/src/flight_service/do_get.rs
@@ -10,11 +10,9 @@ use arrow_flight::flight_service_server::FlightService;
 use arrow_flight::Ticket;
 use datafusion::execution::{SessionState, SessionStateBuilder};
 use datafusion::optimizer::OptimizerConfig;
-use datafusion::physical_plan::ExecutionPlan;
 use futures::TryStreamExt;
 use prost::Message;
 use std::sync::Arc;
-use tokio::sync::OnceCell;
 use tonic::{Request, Response, Status};
 
 use super::service::StageKey;
@@ -68,13 +66,6 @@ impl ArrowFlightEndpoint {
 
         let inner_plan = stage.plan.clone();
 
-        /*println!(
-            "{} Task {:?} executing partition {}",
-            stage.name(),
-            task.partition_group,
-            partition
-        );*/
-
         let stream = inner_plan
             .execute(partition, state.task_ctx())
             .map_err(|err| Status::internal(format!("Error executing stage plan: {err:#?}")))?;
diff --git a/src/physical_optimizer.rs b/src/physical_optimizer.rs
@@ -13,7 +13,7 @@ use datafusion::{
     error::Result,
     physical_optimizer::PhysicalOptimizerRule,
     physical_plan::{
-        displayable, repartition::RepartitionExec, ExecutionPlan, ExecutionPlanProperties,
+        repartition::RepartitionExec, ExecutionPlan, ExecutionPlanProperties,
     },
 };
 use uuid::Uuid;
diff --git a/src/plan/arrow_flight_read.rs b/src/plan/arrow_flight_read.rs
@@ -11,15 +11,14 @@ use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::Ticket;
 use datafusion::arrow::datatypes::SchemaRef;
-use datafusion::arrow::util::pretty::pretty_format_batches;
 use datafusion::common::{exec_err, internal_datafusion_err, internal_err, plan_err};
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
-use futures::{future, StreamExt, TryFutureExt, TryStreamExt};
+use futures::{future, TryFutureExt, TryStreamExt};
 use prost::Message;
 use std::any::Any;
 use std::fmt::Formatter;
@@ -192,7 +191,6 @@ impl ExecutionPlan for ArrowFlightReadExec {
 
         let stream = async move {
             let futs = child_stage_tasks.iter().enumerate().map(|(i, task)| {
-                let i_capture = i;
                 let child_stage_proto_capture = child_stage_proto.clone();
                 let channel_manager_capture = channel_manager.clone();
                 let schema = schema.clone();
@@ -211,7 +209,7 @@ impl ExecutionPlan for ArrowFlightReadExec {
                         stage_proto: Some(child_stage_proto_capture),
                         partition: partition as u64,
                         stage_key: Some(key),
-                        task_number: i_capture as u64,
+                        task_number: i as u64,
                     }
                     .encode_to_vec()
                     .into();
diff --git a/src/plan/isolator.rs b/src/plan/isolator.rs
@@ -5,8 +5,8 @@ use datafusion::{
     error::Result,
     execution::SendableRecordBatchStream,
     physical_plan::{
-        DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan,
-        ExecutionPlanProperties, Partitioning, PlanProperties,
+        DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan, Partitioning,
+        PlanProperties,
     },
 };
 
diff --git a/src/stage/proto.rs b/src/stage/proto.rs
@@ -1,4 +1,4 @@
-use std::{io::Bytes, sync::Arc};
+use std::sync::Arc;
 
 use datafusion::{
     common::internal_datafusion_err,
@@ -10,7 +10,6 @@ use datafusion_proto::{
     physical_plan::{AsExecutionPlan, PhysicalExtensionCodec},
     protobuf::PhysicalPlanNode,
 };
-use uuid::Uuid;
 
 use crate::task::ExecutionTask;
 
diff --git a/src/stage/stage.rs b/src/stage/stage.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use datafusion::common::internal_err;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::TaskContext;
-use datafusion::physical_plan::{displayable, ExecutionPlan};
+use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::SessionContext;
 
 use itertools::Itertools;
diff --git a/src/test_utils/insta.rs b/src/test_utils/insta.rs
@@ -1,4 +1,3 @@
-use datafusion::common::utils::get_available_parallelism;
 use std::env;
 
 pub use insta;
diff --git a/src/test_utils/localhost.rs b/src/test_utils/localhost.rs
@@ -8,31 +8,45 @@ use datafusion::execution::SessionStateBuilder;
 use datafusion::prelude::SessionContext;
 use datafusion::{common::runtime::JoinSet, prelude::SessionConfig};
 use std::error::Error;
-use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
 use std::time::Duration;
+use tokio::net::TcpListener;
 use tonic::transport::{Channel, Server};
 use url::Url;
 
-pub async fn start_localhost_context<N, I, B>(
-    ports: I,
+pub async fn start_localhost_context<B>(
+    num_workers: usize,
     session_builder: B,
 ) -> (SessionContext, JoinSet<()>)
 where
-    N::Error: std::fmt::Debug,
-    N: TryInto<u16>,
-    I: IntoIterator<Item = N>,
     B: SessionBuilder + Send + Sync + 'static,
     B: Clone,
 {
-    let ports: Vec<u16> = ports.into_iter().map(|x| x.try_into().unwrap()).collect();
+    let listeners = futures::future::try_join_all(
+        (0..num_workers)
+            .map(|_| TcpListener::bind("127.0.0.1:0"))
+            .collect::<Vec<_>>(),
+    )
+    .await
+    .expect("Failed to bind to address");
+
+    let ports: Vec<u16> = listeners
+        .iter()
+        .map(|listener| {
+            listener
+                .local_addr()
+                .expect("Failed to get local address")
+                .port()
+        })
+        .collect();
+
     let channel_resolver = LocalHostChannelResolver::new(ports.clone());
     let mut join_set = JoinSet::new();
-    for port in ports {
+    for listener in listeners {
         let channel_resolver = channel_resolver.clone();
         let session_builder = session_builder.clone();
         join_set.spawn(async move {
-            spawn_flight_service(channel_resolver, session_builder, port)
+            spawn_flight_service(channel_resolver, session_builder, listener)
                 .await
                 .unwrap();
         });
@@ -63,7 +77,6 @@ where
 #[derive(Clone)]
 pub struct LocalHostChannelResolver {
     ports: Vec<u16>,
-    i: Arc<AtomicUsize>,
 }
 
 impl LocalHostChannelResolver {
@@ -72,7 +85,6 @@ impl LocalHostChannelResolver {
         N::Error: std::fmt::Debug,
     {
         Self {
-            i: Arc::new(AtomicUsize::new(0)),
             ports: ports.into_iter().map(|v| v.try_into().unwrap()).collect(),
         }
     }
@@ -97,13 +109,16 @@ impl ChannelResolver for LocalHostChannelResolver {
 pub async fn spawn_flight_service(
     channel_resolver: impl ChannelResolver + Send + Sync + 'static,
     session_builder: impl SessionBuilder + Send + Sync + 'static,
-    port: u16,
+    incoming: TcpListener,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
     let mut endpoint = ArrowFlightEndpoint::new(channel_resolver);
     endpoint.with_session_builder(session_builder);
+
+    let incoming = tokio_stream::wrappers::TcpListenerStream::new(incoming);
+
     Ok(Server::builder()
         .add_service(FlightServiceServer::new(endpoint))
-        .serve(format!("127.0.0.1:{port}").parse()?)
+        .serve_with_incoming(incoming)
         .await?)
 }
 
diff --git a/src/test_utils/tpch.rs b/src/test_utils/tpch.rs
@@ -148,7 +148,7 @@ where
 
         writer.write(&first_batch)?;
 
-        while let Some(batch) = data_source.next() {
+        for batch in data_source {
             writer.write(&batch)?;
         }
 
diff --git a/tests/custom_extension_codec.rs b/tests/custom_extension_codec.rs
@@ -49,8 +49,7 @@ mod tests {
             }
         }
 
-        let (ctx, _guard) =
-            start_localhost_context([50050, 50051, 50052], CustomSessionBuilder).await;
+        let (ctx, _guard) = start_localhost_context(3, CustomSessionBuilder).await;
 
         let single_node_plan = build_plan(false)?;
         assert_snapshot!(displayable(single_node_plan.as_ref()).indent(true).to_string(), @r"
diff --git a/tests/distributed_aggregation.rs b/tests/distributed_aggregation.rs
diff --git a/tests/error_propagation.rs b/tests/error_propagation.rs
diff --git a/tests/highly_distributed_query.rs b/tests/highly_distributed_query.rs
diff --git a/tests/non_distributed_consistency_test.rs b/tests/non_distributed_consistency_test.rs
diff --git a/tests/stage_planning.rs b/tests/stage_planning.rs

Original file line number	Diff line number	Diff line change
`@@ -96,7 +96,6 @@ pub struct BenchQuery {`
`96`	`96`	`pub struct QueryResult {`
`97`	`97`	`pub elapsed: Duration,`
`98`	`98`	`pub row_count: usize,`
`99`		`- pub valid: bool,`
`100`	`99`	`}`
`101`	`100`	`/// collects benchmark run data and then serializes it at the end`
`102`	`101`	`pub struct BenchmarkRun {`
Original file line number	Diff line number	Diff line change
`@@ -224,8 +224,8 @@ mod tests {`
`224`	`224`	`let (recovered_error, recovered_ctx) = proto.to_arrow_error();`
`225`	`225`
`226`	`226`	`if original_error.to_string() != recovered_error.to_string() {`
`227`		`- println!("original error: {}", original_error.to_string());`
`228`		`- println!("recovered error: {}", recovered_error.to_string());`
	`227`	`+ println!("original error: {}", original_error);`
	`228`	`+ println!("recovered error: {}", recovered_error);`
`229`	`229`	`}`
`230`	`230`
`231`	`231`	`assert_eq!(original_error.to_string(), recovered_error.to_string());`