villebro commented on code in PR #2800:
URL:
https://github.com/apache/datafusion-ballista/pull/2800#discussion_r4503102502
##########
ballista/core/src/execution_plans/range_repartition_common.rs:
##########
@@ -394,24 +400,109 @@ pub(super) mod test_util {
)))
}
}
+
+ /// `(v2 Float64, id Int64)`, both non-nullable.
+ pub(crate) fn schema_v2_id() -> Arc<Schema> {
+ Arc::new(Schema::new(vec![
+ Field::new("v2", DataType::Float64, false),
+ Field::new("id", DataType::Int64, false),
+ ]))
+ }
+
+ /// A batch over [`schema_v2_id`].
+ pub(crate) fn batch(
+ schema: &Arc<Schema>,
+ keys: Vec<f64>,
+ ids: Vec<i64>,
+ ) -> RecordBatch {
+ RecordBatch::try_new(
+ schema.clone(),
+ vec![
+ Arc::new(Float64Array::from(keys)),
+ Arc::new(Int64Array::from(ids)),
+ ],
+ )
+ .unwrap()
+ }
+
+ /// Ascending sort expression on the column `name`.
+ pub(crate) fn asc(schema: &Schema, name: &str) -> PhysicalSortExpr {
+ PhysicalSortExpr {
+ expr: col(name, schema).unwrap(),
+ options: Default::default(),
+ }
+ }
+
+ /// A one-partition memory source with no batches.
+ pub(crate) fn empty_input(schema: &Arc<Schema>) -> Arc<dyn ExecutionPlan> {
+ MemorySourceConfig::try_new_exec(&[vec![]], schema.clone(),
None).unwrap()
+ }
+
+ pub(crate) fn session() -> Arc<SessionContext> {
+ Arc::new(SessionContext::new_with_state(
+ SessionStateBuilder::new().with_default_features().build(),
+ ))
+ }
+
+ /// Drains all K output partitions of `exec` **concurrently** and maps each
+ /// batch with `extract`. Draining them one after another can deadlock the
+ /// bounded scatter channels: while a test sequentially polls partition 0,
+ /// scatter tasks fill channels 1..K-1 to `CHANNEL_CAPACITY` and block on
+ /// send, never dropping their partition-0 senders. In production K
+ /// downstream tasks run concurrently and this pattern never arises.
+ pub(crate) async fn drain_concurrent<T: Send + 'static>(
+ exec: Arc<dyn ExecutionPlan>,
+ ctx: Arc<SessionContext>,
+ extract: fn(&RecordBatch) -> Vec<T>,
+ ) -> Vec<Vec<T>> {
+ let output_partitions = exec.output_partitioning().partition_count();
+ let streams: Vec<_> = (0..output_partitions)
+ .map(|p| exec.execute(p, ctx.task_ctx()).unwrap())
+ .collect();
+ let handles: Vec<_> = streams
+ .into_iter()
+ .map(|stream| {
+ tokio::spawn(async move {
+ let batches: Vec<RecordBatch> = stream
+ .collect::<Vec<Result<RecordBatch>>>()
+ .await
+ .into_iter()
+ .map(|r| r.unwrap())
+ .collect();
+ batches.iter().flat_map(extract).collect()
+ })
+ })
+ .collect();
+ let mut per_output = Vec::with_capacity(output_partitions);
+ for handle in handles {
+ per_output.push(handle.await.unwrap());
+ }
+ per_output
+ }
+
+ /// A single non-nullable `v Float64` column.
+ pub(crate) fn v_schema() -> Arc<Schema> {
Review Comment:
Thanks, done in this PR in ec262639e. `v_schema` and `cuts` moved to
`execution_plans/test_util.rs`, since `range_filter` and `runtime_stats` use
them, and the two schemas are now `float64_v_schema` and `int64_v_schema`.
`session` stays with the range repartition helpers, since only those tests use
it.
##########
ballista/core/src/execution_plans/test_util.rs:
##########
@@ -0,0 +1,188 @@
+// 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.
+
+//! Fixtures shared by the unit tests of several execution plans.
+
+use std::path::Path;
+use std::sync::Arc;
+
+use datafusion::arrow::array::{
+ Int32Array, Int64Array, RecordBatch, StringArray, UInt32Array,
+};
+use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use datafusion::arrow::ipc::writer::StreamWriter;
+use datafusion::datasource::memory::MemorySourceConfig;
+use datafusion::error::Result;
+use datafusion::execution::TaskContext;
+use datafusion::physical_plan::ExecutionPlan;
+
+use crate::error::BallistaError;
+use crate::serde::scheduler::{
+ ExecutorMetadata, ExecutorOperatingSystemSpecification,
ExecutorSpecification,
+ PartitionId, PartitionLocation, PartitionStats,
+};
+use crate::utils;
+
+/// Two partitions of two identical batches of `(a UInt32, b Utf8)`.
+pub(crate) fn two_by_two_input() -> Result<Arc<dyn ExecutionPlan>> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::UInt32, true),
+ Field::new("b", DataType::Utf8, true),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![
+ Arc::new(UInt32Array::from(vec![Some(1), Some(3)])),
+ Arc::new(StringArray::from(vec![Some("hello"), Some("world")])),
+ ],
+ )?;
+ let partition = vec![batch.clone(), batch];
+ let partitions = vec![partition.clone(), partition];
+ Ok(MemorySourceConfig::try_new_exec(&partitions, schema, None)?)
+}
+
+/// Executes every output partition of `plan` on its own task and returns the
+/// collected batches of each, in partition order.
+pub(crate) async fn drive_partition_results<P: ExecutionPlan + 'static>(
Review Comment:
Good point, added a line to the doc comment in 47e0a6c75.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
