From 7e7aab2164007f2e1515869fb3e1d406f8e03504 Mon Sep 17 00:00:00 2001 From: Rob Tandy Date: Fri, 11 Jul 2025 09:37:23 -0400 Subject: [PATCH 1/3] remove old naming scheme and use DD where appropriate --- k8s/cluster.yaml | 54 ++++++------ k8s/dev_cluster.yaml | 56 ++++++------- scripts/build_and_push_docker.sh | 39 --------- scripts/launch_tpch_cluster.sh | 10 +-- scripts/python_tests.sh | 2 +- src/codec.rs | 39 +++++---- src/lib.rs | 5 +- src/main.rs | 8 +- src/physical.rs | 24 +++--- src/planning.rs | 40 ++++----- src/processor_service.rs | 24 +++--- src/proto/distributed_datafusion.proto | 6 +- src/proto/generated/protobuf.rs | 14 ++-- src/proxy_service.rs | 24 +++--- src/result.rs | 4 +- src/stage.rs | 110 ++----------------------- src/stage_reader.rs | 14 ++-- src/util.rs | 32 ++----- src/worker_discovery.rs | 14 ++-- tests/common/mod.rs | 10 +-- tpch/queries/q15.sql | 2 +- 21 files changed, 180 insertions(+), 351 deletions(-) delete mode 100755 scripts/build_and_push_docker.sh diff --git a/k8s/cluster.yaml b/k8s/cluster.yaml index 98813323..aee0c5d2 100644 --- a/k8s/cluster.yaml +++ b/k8s/cluster.yaml @@ -1,34 +1,34 @@ apiVersion: v1 kind: Service metadata: - name: service-dfray-proxy + name: service-dd-proxy labels: - app: dfray-proxy + app: dd-proxy spec: type: ClusterIP ports: - - name: dfray-proxy + - name: dd-proxy port: 20200 targetPort: 20200 selector: - app: dfray-proxy + app: dd-proxy --- apiVersion: apps/v1 kind: Deployment metadata: - name: deployment-dfray-proxy + name: deployment-dd-proxy labels: - app: dfray-proxy + app: dd-proxy spec: replicas: 2 selector: matchLabels: - app: dfray-proxy + app: dd-proxy template: metadata: labels: - app: dfray-proxy - service: dfray-proxy + app: dd-proxy + service: dd-proxy team: mint admission.datadoghq.com/enabled: "true" annotations: @@ -52,16 +52,14 @@ spec: - key: "node-role.kubernetes.io/nodeless" operator: Exists containers: - - name: dfray-proxy - image: registry.ddbuild.io/dfray-amd64:2025-06-17-b + - name: dd-proxy + image: registry.ddbuild.io/dd-amd64:2025-06-17-b env: - - name: DATAFUSION_RAY_LOG_LEVEL - value: trace - - name: DFRAY_WORKER_DEPLOYMENT - value: deployment-dfray-worker - - name: DFRAY_WORKER_DEPLOYMENT_PORT + - name: dd_WORKER_DEPLOYMENT + value: deployment-dd-worker + - name: dd_WORKER_DEPLOYMENT_PORT value: "20201" - - name: DFRAY_WORKER_DEPLOYMENT_NAMESPACE + - name: dd_WORKER_DEPLOYMENT_NAMESPACE valueFrom: fieldRef: fieldPath: metadata.namespace @@ -76,7 +74,7 @@ spec: command: - sh - "-c" - - "/dfray --mode proxy --port 20200" + - "/dd --mode proxy --port 20200" ports: - containerPort: 20200 resources: @@ -90,19 +88,19 @@ spec: apiVersion: apps/v1 kind: Deployment metadata: - name: deployment-dfray-worker + name: deployment-dd-worker labels: - app: dfray-worker + app: dd-worker spec: replicas: 2 selector: matchLabels: - app: dfray-worker + app: dd-worker template: metadata: labels: - app: dfray-worker - service: service-dfray-worker + app: dd-worker + service: service-dd-worker team: mint admission.datadoghq.com/enabled: "true" annotations: @@ -126,13 +124,11 @@ spec: - key: "node-role.kubernetes.io/nodeless" operator: Exists containers: - - name: dfray-worker - image: registry.ddbuild.io/dfray-amd64:2025-06-17-b + - name: dd-worker + image: registry.ddbuild.io/dd-amd64:2025-06-17-b env: - - name: DATAFUSION_RAY_LOG_LEVEL - value: trace - name: RUST_LOG - value: datafusion_bindings=info,cs3_query=info + value: distributed_datafusion=trace,datafusion_bindings=info,cs3_query=info - name: HDQ_URL value: http://hdq-creed.data-eng.all-clusters.local-dc.fabric.dog:6420 - name: MEM_THRESHOLD @@ -142,7 +138,7 @@ spec: command: - sh - "-c" - - "/dfray --mode worker --port 20201" + - "/dd --mode worker --port 20201" ports: - containerPort: 20201 resources: diff --git a/k8s/dev_cluster.yaml b/k8s/dev_cluster.yaml index 62d0351d..3a03b13d 100644 --- a/k8s/dev_cluster.yaml +++ b/k8s/dev_cluster.yaml @@ -1,51 +1,49 @@ apiVersion: v1 kind: Service metadata: - name: service-dfray-proxy + name: service-dd-proxy labels: - app: dfray-proxy + app: dd-proxy spec: type: ClusterIP ports: - - name: dfray-proxy + - name: dd-proxy port: 20200 targetPort: 20200 selector: - app: dfray-proxy + app: dd-proxy --- apiVersion: apps/v1 kind: Deployment metadata: - name: deployment-dfray-proxy + name: deployment-dd-proxy labels: - app: dfray-proxy + app: dd-proxy spec: replicas: 1 selector: matchLabels: - app: dfray-proxy + app: dd-proxy template: metadata: labels: - app: dfray-proxy - service: dfray-proxy + app: dd-proxy + service: dd-proxy spec: containers: - - name: dfray-proxy - image: registry.ddbuild.io/dfray:2025-06-16-e + - name: dd-proxy + image: registry.ddbuild.io/dd:2025-06-16-e env: - - name: DATAFUSION_RAY_LOG_LEVEL - value: debug - - name: DFRAY_WORKER_DEPLOYMENT - value: deployment-dfray-worker - - name: DFRAY_WORKER_DEPLOYMENT_PORT + - name: dd_WORKER_DEPLOYMENT + value: deployment-dd-worker + - name: dd_WORKER_DEPLOYMENT_PORT value: "20201" - - name: DFRAY_WORKER_DEPLOYMENT_NAMESPACE + - name: dd_WORKER_DEPLOYMENT_NAMESPACE valueFrom: fieldRef: fieldPath: metadata.namespace - name: RUST_LOG - value: datafusion_bindings=info,cs3_query=info + value: datafusion_bindings=info,cs3_query=info,distributed_datafusion=debug - name: HDQ_URL value: http://hdq-creed.data-eng.all-clusters.local-dc.fabric.dog:6420 - name: MEM_THRESHOLD @@ -55,35 +53,33 @@ spec: command: - sh - "-c" - - "/dfray --mode proxy --port 20200" + - "/dd --mode proxy --port 20200" ports: - containerPort: 20200 --- apiVersion: apps/v1 kind: Deployment metadata: - name: deployment-dfray-worker + name: deployment-dd-worker labels: - app: dfray-worker + app: dd-worker spec: replicas: 2 selector: matchLabels: - app: dfray-worker + app: dd-worker template: metadata: labels: - app: dfray-worker - service: service-dfray-worker + app: dd-worker + service: service-dd-worker spec: containers: - - name: dfray-worker - image: registry.ddbuild.io/dfray:2025-06-16-e + - name: dd-worker + image: registry.ddbuild.io/dd:2025-06-16-e env: - - name: DATAFUSION_RAY_LOG_LEVEL - value: debug - name: RUST_LOG - value: datafusion_bindings=info,cs3_query=info + value: datafusion_bindings=info,cs3_query=info,distributed_datafusion=debug - name: HDQ_URL value: http://hdq-creed.data-eng.all-clusters.local-dc.fabric.dog:6420 - name: MEM_THRESHOLD @@ -93,7 +89,7 @@ spec: command: - sh - "-c" - - "/dfray --mode worker --port 20201" + - "/dd --mode worker --port 20201" ports: - containerPort: 20201 --- diff --git a/scripts/build_and_push_docker.sh b/scripts/build_and_push_docker.sh deleted file mode 100755 index c2c40ed4..00000000 --- a/scripts/build_and_push_docker.sh +++ /dev/null @@ -1,39 +0,0 @@ -#!/bin/bash - -set -ex - -BUILD_TAG=$1 - -if [ -z "${BUILD_TAG}" ]; then - echo "Usage: $0 . E.g. build_and_push_docker.sh 0.1.0" - exit 1 -fi - -docker buildx build --platform=linux/amd64 \ - --cache-from type=local,src=/tmp/cache/amd64 \ - --cache-to type=local,dest=/tmp/cache/amd64 \ - --ssh default \ - --tag registry.ddbuild.io/dfray-amd64:"${BUILD_TAG}" \ - -f Dockerfile.distributed-datafusion . - -docker push registry.ddbuild.io/dfray-amd64:"${BUILD_TAG}" - -docker buildx build --platform=linux/arm64 \ - --cache-from type=local,src=/tmp/cache/arm64 \ - --cache-to type=local,dest=/tmp/cache/arm64 \ - --ssh default \ - --tag registry.ddbuild.io/dfray-arm64:"${BUILD_TAG}" \ - -f Dockerfile.distributed-datafusion . - -docker push registry.ddbuild.io/dfray-arm64:"${BUILD_TAG}" - -# unsure why this does not work at the moment, for now push each build -# docker manifest create --amend \ -# registry.ddbuild.io/dfray:"${BUILD_TAG}" \ -# registry.ddbuild.io/dfray-amd64:"${BUILD_TAG}" \ -# registry.ddbuild.io/dfray-arm64:"${BUILD_TAG}" -# -# docker manifest annotate registry.ddbuild.io/dfray:"${BUILD_TAG}" registry.ddbuild.io/dfray-amd64:"${BUILD_TAG}" --os linux --arch amd64 -# docker manifest annotate registry.ddbuild.io/dfray:"${BUILD_TAG}" registry.ddbuild.io/dfray-arm64:"${BUILD_TAG}" --os linux --arch arm64 - -# docker manifest push registry.ddbuild.io/dfray:"${BUILD_TAG}" diff --git a/scripts/launch_tpch_cluster.sh b/scripts/launch_tpch_cluster.sh index 6cdf0454..d3ae264a 100755 --- a/scripts/launch_tpch_cluster.sh +++ b/scripts/launch_tpch_cluster.sh @@ -59,6 +59,7 @@ done NUM_WORKERS=${NUM_WORKERS:-$DEFAULT_NUM_WORKERS} TPCH_DATA_DIR=${TPCH_DATA_DIR:-$DEFAULT_TPCH_PATH} LOG_DIR=${LOG_DIR:-$DEFAULT_LOG_PATH} +RUST_LOG=${RUST_LOG:-"info,distributed_datafusion=debug"} # Validate inputs if [ "$NUM_WORKERS" -lt 1 ]; then @@ -116,8 +117,7 @@ echo " - Log Directory: $LOG_DIR" echo # Define environment variables -export DATAFUSION_RAY_LOG_LEVEL=trace -export DFRAY_TABLES="customer:parquet:${TPCH_DATA_DIR}/customer.parquet,\ +export DD_TABLES="customer:parquet:${TPCH_DATA_DIR}/customer.parquet,\ lineitem:parquet:${TPCH_DATA_DIR}/lineitem.parquet,\ nation:parquet:${TPCH_DATA_DIR}/nation.parquet,\ orders:parquet:${TPCH_DATA_DIR}/orders.parquet,\ @@ -147,8 +147,7 @@ for ((i = 0; i < NUM_WORKERS; i++)); do WORKER_NAME="worker$((i + 1))" LOG_FILE="${LOG_DIR}/${WORKER_NAME}.log" echo " Starting $WORKER_NAME on port $PORT..." - #env DATAFUSION_RAY_LOG_LEVEL="$DATAFUSION_RAY_LOG_LEVEL" DFRAY_TABLES="$DFRAY_TABLES" ./target/release/distributed-datafusion --mode worker --port $PORT >"$LOG_FILE" 2>&1 & - env RUST_BACKTRACE=1 DATAFUSION_RAY_LOG_LEVEL="$DATAFUSION_RAY_LOG_LEVEL" DFRAY_TABLES="$DFRAY_TABLES" ./target/debug/distributed-datafusion --mode worker --port $PORT >"$LOG_FILE" 2>&1 & + env RUST_LOG="$RUST_LOG" RUST_BACKTRACE=1 DD_TABLES="$DD_TABLES" ./target/debug/distributed-datafusion --mode worker --port $PORT >"$LOG_FILE" 2>&1 & WORKER_PIDS[$i]=$! WORKER_ADDRESSES[$i]="localhost:${PORT}" done @@ -167,8 +166,7 @@ WORKER_ADDRESSES_STR=$( echo "Starting proxy on port 20200..." echo "Connecting to workers: $WORKER_ADDRESSES_STR" PROXY_LOG="${LOG_DIR}/proxy.log" -#env DATAFUSION_RAY_LOG_LEVEL="$DATAFUSION_RAY_LOG_LEVEL" DFRAY_TABLES="$DFRAY_TABLES" DFRAY_WORKER_ADDRESSES="$WORKER_ADDRESSES_STR" ./target/release/distributed-datafusion --mode proxy --port 20200 >"$PROXY_LOG" 2>&1 & -env RUST_BACKTRACE=1 DATAFUSION_RAY_LOG_LEVEL="$DATAFUSION_RAY_LOG_LEVEL" DFRAY_TABLES="$DFRAY_TABLES" DFRAY_WORKER_ADDRESSES="$WORKER_ADDRESSES_STR" ./target/debug/distributed-datafusion --mode proxy --port 20200 >"$PROXY_LOG" 2>&1 & +env RUST_LOG="$RUST_LOG" RUST_BACKTRACE=1 DD_TABLES="$DD_TABLES" DD_WORKER_ADDRESSES="$WORKER_ADDRESSES_STR" ./target/debug/distributed-datafusion --mode proxy --port 20200 >"$PROXY_LOG" 2>&1 & PROXY_PID=$! echo diff --git a/scripts/python_tests.sh b/scripts/python_tests.sh index 7a20a897..45adafa2 100644 --- a/scripts/python_tests.sh +++ b/scripts/python_tests.sh @@ -23,4 +23,4 @@ maturin develop python tpch/make_data.py $TPCH_SCALING_FACTOR $TPCH_DATA_PATH -DATAFUSION_RAY_LOG_LEVEL=debug RAY_COLOR_PREFIX=1 RAY_DEDUP_LOGS=0 python tpch/tpcbench.py --data=file:///$TPCH_DATA_PATH/ --concurrency 3 --partitions-per-processor 2 --batch-size=8192 --worker-pool-min=20 --validate +RUST_LOG=distributed_datafusion=debug python tpch/tpcbench.py --data=file:///$TPCH_DATA_PATH/ --concurrency 3 --partitions-per-processor 2 --batch-size=8192 --worker-pool-min=20 --validate diff --git a/src/codec.rs b/src/codec.rs index 346b768f..d5231e41 100644 --- a/src/codec.rs +++ b/src/codec.rs @@ -21,27 +21,26 @@ use crate::{ logging::trace, max_rows::MaxRowsExec, protobuf::{ - df_ray_exec_node::Payload, DfRayExecNode, DfRayStageReaderExecNode, - DistributedAnalyzeExecNode, DistributedAnalyzeRootExecNode, MaxRowsExecNode, - PartitionIsolatorExecNode, RecordBatchExecNode, + dd_exec_node::Payload, DdExecNode, DdStageReaderExecNode, DistributedAnalyzeExecNode, + DistributedAnalyzeRootExecNode, MaxRowsExecNode, PartitionIsolatorExecNode, + RecordBatchExecNode, }, record_batch_exec::RecordBatchExec, - stage_reader::DFRayStageReaderExec, + stage_reader::DDStageReaderExec, util::{batch_to_ipc, ipc_to_batch}, }; #[derive(Debug)] -/// Physical Extension Codec for for DataFusion for Ray plans -pub struct DFRayCodec {} +pub struct DDCodec {} -impl PhysicalExtensionCodec for DFRayCodec { +impl PhysicalExtensionCodec for DDCodec { fn try_decode( &self, buf: &[u8], inputs: &[Arc], registry: &dyn FunctionRegistry, ) -> Result> { - if let Ok(node) = DfRayExecNode::decode(buf) { + if let Ok(node) = DdExecNode::decode(buf) { let payload = node .payload .ok_or(internal_datafusion_err!("no payload when decoding proto"))?; @@ -62,7 +61,7 @@ impl PhysicalExtensionCodec for DFRayCodec { )? .ok_or(internal_datafusion_err!("missing partitioning in proto"))?; - Ok(Arc::new(DFRayStageReaderExec::try_new( + Ok(Arc::new(DDStageReaderExec::try_new( part, Arc::new(schema), node.stage_id, @@ -129,7 +128,7 @@ impl PhysicalExtensionCodec for DFRayCodec { } } } else { - internal_err!("cannot decode proto extension in dfray codec") + internal_err!("cannot decode proto extension in distributed datafusion codec") } } @@ -139,14 +138,14 @@ impl PhysicalExtensionCodec for DFRayCodec { displayable(node.as_ref()).one_line() ); - let payload = if let Some(reader) = node.as_any().downcast_ref::() { + let payload = if let Some(reader) = node.as_any().downcast_ref::() { let schema: protobuf::Schema = reader.schema().try_into()?; let partitioning: protobuf::Partitioning = serialize_partitioning( reader.properties().output_partitioning(), &DefaultPhysicalExtensionCodec {}, )?; - let pb = DfRayStageReaderExecNode { + let pb = DdStageReaderExecNode { schema: Some(schema), partitioning: Some(partitioning), stage_id: reader.stage_id, @@ -187,7 +186,7 @@ impl PhysicalExtensionCodec for DFRayCodec { return internal_err!("Not supported node to encode to proto"); }; - let pb = DfRayExecNode { + let pb = DdExecNode { payload: Some(payload), }; pb.encode(buf) @@ -214,7 +213,7 @@ mod test { use super::*; use crate::{ - isolator::PartitionIsolatorExec, max_rows::MaxRowsExec, stage_reader::DFRayStageReaderExec, + isolator::PartitionIsolatorExec, max_rows::MaxRowsExec, stage_reader::DDStageReaderExec, }; fn create_test_schema() -> Arc { @@ -226,7 +225,7 @@ mod test { fn verify_round_trip(exec: Arc) { let ctx = SessionContext::new(); - let codec = DFRayCodec {}; + let codec = DDCodec {}; // serialize execution plan to proto let proto: protobuf::PhysicalPlanNode = @@ -255,8 +254,8 @@ mod test { fn stage_reader_round_trip() { let schema = create_test_schema(); let part = Partitioning::UnknownPartitioning(2); - let exec = Arc::new(DFRayStageReaderExec::try_new(part, schema, 1).unwrap()); - let codec = DFRayCodec {}; + let exec = Arc::new(DDStageReaderExec::try_new(part, schema, 1).unwrap()); + let codec = DDCodec {}; let mut buf = vec![]; codec.try_encode(exec.clone(), &mut buf).unwrap(); let ctx = SessionContext::new(); @@ -268,7 +267,7 @@ mod test { fn max_rows_round_trip() { let schema = create_test_schema(); let part = Partitioning::UnknownPartitioning(2); - let reader_exec = Arc::new(DFRayStageReaderExec::try_new(part, schema, 1).unwrap()); + let reader_exec = Arc::new(DDStageReaderExec::try_new(part, schema, 1).unwrap()); let exec = Arc::new(MaxRowsExec::new(reader_exec, 10)); verify_round_trip(exec); @@ -278,7 +277,7 @@ mod test { fn partition_isolator_round_trip() { let schema = create_test_schema(); let part = Partitioning::UnknownPartitioning(2); - let reader_exec = Arc::new(DFRayStageReaderExec::try_new(part, schema, 1).unwrap()); + let reader_exec = Arc::new(DDStageReaderExec::try_new(part, schema, 1).unwrap()); let exec = Arc::new(PartitionIsolatorExec::new(reader_exec, 4)); verify_round_trip(exec); @@ -289,7 +288,7 @@ mod test { let schema = create_test_schema(); let part = Partitioning::UnknownPartitioning(2); let exec = Arc::new(MaxRowsExec::new( - Arc::new(DFRayStageReaderExec::try_new(part, schema, 1).unwrap()), + Arc::new(DDStageReaderExec::try_new(part, schema, 1).unwrap()), 10, )); diff --git a/src/lib.rs b/src/lib.rs index e6198dbb..f5fc516b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -58,13 +58,10 @@ pub fn setup() { } fn setup_logging() { - let dfr_env = env::var("DATAFUSION_RAY_LOG_LEVEL").unwrap_or("WARN".to_string()); let rust_log_env = env::var("RUST_LOG").unwrap_or("WARN".to_string()); - let combined_env = format!("{rust_log_env},distributed_datafusion={dfr_env}"); - env_logger::Builder::new() - .parse_filters(&combined_env) + .parse_filters(&rust_log_env) .init(); } diff --git a/src/main.rs b/src/main.rs index 21429c94..401f1a9a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,8 @@ use anyhow::Result; use clap::Parser; use distributed_datafusion::{ - friendly::new_friendly_name, processor_service::DFRayProcessorService, - proxy_service::DFRayProxyService, setup, + friendly::new_friendly_name, processor_service::DDProcessorService, + proxy_service::DDProxyService, setup, }; #[derive(Parser)] @@ -33,11 +33,11 @@ async fn main() -> Result<()> { match args.mode.as_str() { "proxy" => { - let service = DFRayProxyService::new(new_friendly_name()?, args.port).await?; + let service = DDProxyService::new(new_friendly_name()?, args.port).await?; service.serve().await?; } "worker" => { - let service = DFRayProcessorService::new(new_friendly_name()?, args.port).await?; + let service = DDProcessorService::new(new_friendly_name()?, args.port).await?; service.serve().await?; } _ => { diff --git a/src/physical.rs b/src/physical.rs index f253a47d..e30556ec 100644 --- a/src/physical.rs +++ b/src/physical.rs @@ -27,33 +27,30 @@ use datafusion::{ }, }; -use crate::{logging::info, stage::DFRayStageExec, util::display_plan_with_partition_counts}; +use crate::{logging::info, stage::DDStageExec, util::display_plan_with_partition_counts}; /// This optimizer rule walks up the physical plan tree -/// and inserts RayStageExec nodes where appropriate to denote where we will +/// and inserts DDStageExec nodes where appropriate to denote where we will /// split the plan into stages. /// -/// The RayStageExec nodes are merely markers to inform where to break the plan -/// up. -/// /// Later, the plan will be examined again to actually split it up. -/// These RayStageExecs serve as markers where we know to break it up on a +/// These DDStageExecs serve as markers where we know to break it up on a /// network boundary and we can insert readers and writers as appropriate. #[derive(Debug)] -pub struct DFRayStageOptimizerRule {} +pub struct DDStageOptimizerRule {} -impl Default for DFRayStageOptimizerRule { +impl Default for DDStageOptimizerRule { fn default() -> Self { Self::new() } } -impl DFRayStageOptimizerRule { +impl DDStageOptimizerRule { pub fn new() -> Self { Self {} } } -impl PhysicalOptimizerRule for DFRayStageOptimizerRule { +impl PhysicalOptimizerRule for DDStageOptimizerRule { fn optimize( &self, plan: Arc, @@ -72,7 +69,7 @@ impl PhysicalOptimizerRule for DFRayStageOptimizerRule { || plan.as_any().downcast_ref::().is_some() { // insert a stage marker here so we know where to break up the physical plan later - let stage = Arc::new(DFRayStageExec::new(plan, stage_counter)); + let stage = Arc::new(DDStageExec::new(plan, stage_counter)); stage_counter += 1; Ok(Transformed::yes(stage as Arc)) } else { @@ -81,8 +78,7 @@ impl PhysicalOptimizerRule for DFRayStageOptimizerRule { }; let plan = plan.clone().transform_up(up)?.data; - let final_plan = - Arc::new(DFRayStageExec::new(plan, stage_counter)) as Arc; + let final_plan = Arc::new(DDStageExec::new(plan, stage_counter)) as Arc; info!( "optimized physical plan:\n{}", @@ -92,7 +88,7 @@ impl PhysicalOptimizerRule for DFRayStageOptimizerRule { } fn name(&self) -> &str { - "RayStageOptimizerRule" + "DDStageOptimizerRule" } fn schema_check(&self) -> bool { diff --git a/src/planning.rs b/src/planning.rs index 07923a70..112dff6e 100644 --- a/src/planning.rs +++ b/src/planning.rs @@ -33,10 +33,10 @@ use crate::{ isolator::PartitionIsolatorExec, logging::{debug, error, info, trace}, max_rows::MaxRowsExec, - physical::DFRayStageOptimizerRule, - result::{DFRayError, Result}, - stage::DFRayStageExec, - stage_reader::{DFRayStageReaderExec, QueryId}, + physical::DDStageOptimizerRule, + result::{DDError, Result}, + stage::DDStageExec, + stage_reader::{DDStageReaderExec, QueryId}, util::{display_plan_with_partition_counts, get_client, physical_plan_to_bytes, wait_for}, vocab::{ Addrs, CtxAnnotatedOutputs, CtxHost, CtxPartitionGroup, CtxStageAddrs, CtxStageId, DDTask, @@ -53,7 +53,7 @@ pub struct DDStage { /// the partition groups for this stage. pub partition_groups: Vec>, /// Are we hosting the complete partitions? If not - /// then RayStageReaderExecs will be inserted to consume its desired + /// then DDStageReaderExecs will be inserted to consume its desired /// partition from all stages with this same id, and merge the results. /// Using a CombinedRecordBatchStream pub full_partitions: bool, @@ -79,7 +79,7 @@ impl DDStage { self.plan .clone() .transform_down(|node: Arc| { - if let Some(reader) = node.as_any().downcast_ref::() { + if let Some(reader) = node.as_any().downcast_ref::() { result.push(reader.stage_id); } Ok(Transformed::no(node)) @@ -161,9 +161,9 @@ pub fn add_ctx_extentions( pub async fn add_tables_from_env(state: &mut SessionState) -> Result<()> { // this string is formatted as a comman separated list of table info // where each table info is name:format:path - let table_str = env::var("DFRAY_TABLES"); + let table_str = env::var("DD_TABLES"); if table_str.is_err() { - info!("No DFRAY_TABLES environment variable set, skipping table registration"); + info!("No DD_TABLES environment variable set, skipping table registration"); return Ok(()); } @@ -171,7 +171,7 @@ pub async fn add_tables_from_env(state: &mut SessionState) -> Result<()> { info!("adding table from env: {}", table); let parts: Vec<&str> = table.split(':').collect(); if parts.len() != 3 { - return Err(anyhow!("Invalid format for DFRAY_TABLES env var: {}", table).into()); + return Err(anyhow!("Invalid format for DD_TABLES env var: {}", table).into()); } let name = parts[0].to_string(); let fmt = parts[1].to_string(); @@ -247,11 +247,11 @@ pub async fn execution_planning( let mut partition_groups = vec![]; let mut full_partitions = false; - // We walk up the tree from the leaves to find the stages, record ray stages, - // and replace each ray stage with a corresponding ray reader stage. + // We walk up the tree from the leaves to find the stages + // and replace each stage with a corresponding reader stage. // // we also calculate the paritition groups at each node, anticipating - // arriving at a RayStageExec that we have to replace with a stage reader + // arriving at a DDStageExec that we have to replace with a stage reader // which will have differing requirements for partition groups depending on // its children. let up = |plan: Arc| { @@ -260,13 +260,13 @@ pub async fn execution_planning( displayable(plan.as_ref()).one_line() ); - if let Some(stage_exec) = plan.as_any().downcast_ref::() { - trace!("ray stage exec. partition_groups: {:?}", partition_groups); + if let Some(stage_exec) = plan.as_any().downcast_ref::() { + trace!("stage exec. partition_groups: {:?}", partition_groups); let input = plan.children(); - assert!(input.len() == 1, "RayStageExec must have exactly one child"); + assert!(input.len() == 1, "DDStageExec must have exactly one child"); let input = input[0]; - let replacement = Arc::new(DFRayStageReaderExec::try_new( + let replacement = Arc::new(DDStageReaderExec::try_new( plan.output_partitioning().clone(), input.schema(), stage_exec.stage_id, @@ -330,11 +330,11 @@ pub async fn execution_planning( } }; - // walk up the plan adding DFRayStageExec marker nodes. + // walk up the plan adding DDStageExec marker nodes. // FIXME: we can do this in one step in the future i think but this is // a carry over from when it was done in a physical optimizer seperate // step - let optimizer = DFRayStageOptimizerRule::new(); + let optimizer = DDStageOptimizerRule::new(); let distributed_plan = optimizer.optimize(physical_plan, &ConfigOptions::default())?; // Clone the distributed plan before transformation since we need to return it @@ -459,7 +459,7 @@ pub async fn distribute_stages( // and can accept the stages match try_distribute_tasks(&task_datas).await { Ok(_) => return Ok((final_addrs, task_datas)), - Err(DFRayError::WorkerCommunicationError(bad_worker)) => { + Err(DDError::WorkerCommunicationError(bad_worker)) => { error!( "distribute stages for query {query_id} attempt {attempt} failed removing \ worker {bad_worker}. Retrying..." @@ -506,7 +506,7 @@ async fn try_distribute_tasks(task_datas: &[DDTask]) -> Result<()> { Ok(client) => client, Err(e) => { error!("Couldn't not communicate with worker {e:#?}"); - return Err(DFRayError::WorkerCommunicationError( + return Err(DDError::WorkerCommunicationError( host.clone(), // here )); } diff --git a/src/processor_service.rs b/src/processor_service.rs index 555ff912..578a8e2e 100644 --- a/src/processor_service.rs +++ b/src/processor_service.rs @@ -50,7 +50,7 @@ use crate::{ protobuf::{ AnnotatedTaskOutput, AnnotatedTaskOutputs, FlightDataMetadata, FlightTicketData, Host, }, - result::{DFRayError, Result}, + result::{DDError, Result}, util::{ bytes_to_physical_plan, display_plan_with_partition_counts, get_addrs, register_object_store_for_paths_in_plan, start_up, @@ -86,7 +86,7 @@ struct StageTasks { } /// It only responds to the DoGet Arrow Flight method. -struct DfRayProcessorHandler { +struct DDProcessorHandler { /// our name, useful for logging name: String, /// our address string, also useful for logging @@ -96,7 +96,7 @@ struct DfRayProcessorHandler { done: Arc>, } -impl DfRayProcessorHandler { +impl DDProcessorHandler { pub fn new(name: String, addr: String) -> Self { let stages: Arc>> = Arc::new(RwLock::new(HashMap::new())); @@ -340,7 +340,7 @@ impl DfRayProcessorHandler { /// executing the plan further downstream to consumers of our stream. /// /// We may discover a FlightDataMetadata as an additional payload incoming - /// on streming requests. The [`RayStageReader`] will look for these and add + /// on streming requests. The [`DDStageReader`] will look for these and add /// them to an extension on the context. /// /// Our job is to send all record batches from our stream over our flight response @@ -514,7 +514,7 @@ impl DfRayProcessorHandler { .stage_addrs .as_ref() .context("stage addrs not present") - .map_err(DFRayError::from) + .map_err(DDError::from) .and_then(get_addrs) .map_err(|e| Status::internal(format!("{}, {e}", self.name)))?; @@ -538,7 +538,7 @@ impl DfRayProcessorHandler { } #[async_trait] -impl FlightHandler for DfRayProcessorHandler { +impl FlightHandler for DDProcessorHandler { async fn do_get( &self, request: Request, @@ -611,21 +611,21 @@ impl FlightHandler for DfRayProcessorHandler { } } -/// DFRayProcessorService is a Arrow Flight service that serves streams of +/// DDProcessorService is a Arrow Flight service that serves streams of /// partitions from a hosted Physical Plan /// /// It only responds to the DoGet Arrow Flight method -pub struct DFRayProcessorService { +pub struct DDProcessorService { #[allow(dead_code)] name: String, listener: TcpListener, - handler: Arc, + handler: Arc, addr: String, all_done_tx: Arc>>, all_done_rx: Option>, } -impl DFRayProcessorService { +impl DDProcessorService { pub async fn new(name: String, port: usize) -> Result { let name = format!("[{}]", name); @@ -636,9 +636,9 @@ impl DFRayProcessorService { let addr = format!("{}", listener.local_addr().unwrap()); - info!("DFRayProcessorService bound to {addr}"); + info!("DDProcessorService bound to {addr}"); - let handler = Arc::new(DfRayProcessorHandler::new(name.clone(), addr.clone())); + let handler = Arc::new(DDProcessorHandler::new(name.clone(), addr.clone())); Ok(Self { name, diff --git a/src/proto/distributed_datafusion.proto b/src/proto/distributed_datafusion.proto index e6d4a8c3..52a47c87 100644 --- a/src/proto/distributed_datafusion.proto +++ b/src/proto/distributed_datafusion.proto @@ -9,7 +9,7 @@ option java_multiple_files = true; option java_outer_classname = "DistributedDataFusionProto"; option java_package = "distributed_datafusion.protobuf"; -message DFRayStageReaderExecNode { +message DDStageReaderExecNode { // schema of the stage we will consume datafusion_common.Schema schema = 1; // properties of the stage we will consume @@ -44,9 +44,9 @@ message RecordBatchExecNode { bytes batch = 1; } -message DFRayExecNode { +message DDExecNode { oneof payload { - DFRayStageReaderExecNode stage_reader_exec = 1; + DDStageReaderExecNode stage_reader_exec = 1; MaxRowsExecNode max_rows_exec = 2; PartitionIsolatorExecNode isolator_exec = 3; DistributedAnalyzeExecNode distributed_analyze_exec = 4; diff --git a/src/proto/generated/protobuf.rs b/src/proto/generated/protobuf.rs index 401aa103..314a456e 100644 --- a/src/proto/generated/protobuf.rs +++ b/src/proto/generated/protobuf.rs @@ -1,6 +1,6 @@ // This file is @generated by prost-build. #[derive(Clone, PartialEq, ::prost::Message)] -pub struct DfRayStageReaderExecNode { +pub struct DdStageReaderExecNode { /// schema of the stage we will consume #[prost(message, optional, tag = "1")] pub schema: ::core::option::Option<::datafusion_proto::protobuf::Schema>, @@ -45,16 +45,16 @@ pub struct RecordBatchExecNode { pub batch: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] -pub struct DfRayExecNode { - #[prost(oneof = "df_ray_exec_node::Payload", tags = "1, 2, 3, 4, 5, 6")] - pub payload: ::core::option::Option, +pub struct DdExecNode { + #[prost(oneof = "dd_exec_node::Payload", tags = "1, 2, 3, 4, 5, 6")] + pub payload: ::core::option::Option, } -/// Nested message and enum types in `DFRayExecNode`. -pub mod df_ray_exec_node { +/// Nested message and enum types in `DDExecNode`. +pub mod dd_exec_node { #[derive(Clone, PartialEq, ::prost::Oneof)] pub enum Payload { #[prost(message, tag = "1")] - StageReaderExec(super::DfRayStageReaderExecNode), + StageReaderExec(super::DdStageReaderExecNode), #[prost(message, tag = "2")] MaxRowsExec(super::MaxRowsExecNode), #[prost(message, tag = "3")] diff --git a/src/proxy_service.rs b/src/proxy_service.rs index e202b13a..1b1c5dd0 100644 --- a/src/proxy_service.rs +++ b/src/proxy_service.rs @@ -46,20 +46,20 @@ use crate::{ protobuf::TicketStatementData, query_planner::{QueryPlan, QueryPlanner}, result::Result, - stage_reader::DFRayStageReaderExec, + stage_reader::DDStageReaderExec, util::{display_plan_with_partition_counts, get_addrs, start_up}, vocab::{Addrs, Host}, worker_discovery::get_worker_addresses, }; -pub struct DfRayProxyHandler { +pub struct DDProxyHandler { /// our host info, useful for logging pub host: Host, pub planner: QueryPlanner, } -impl DfRayProxyHandler { +impl DDProxyHandler { pub fn new(name: String, addr: String) -> Self { // call this function to bootstrap the worker discovery mechanism get_worker_addresses().expect("Could not get worker addresses upon startup"); @@ -160,7 +160,7 @@ impl DfRayProxyHandler { } #[async_trait] -impl FlightSqlHandler for DfRayProxyHandler { +impl FlightSqlHandler for DDProxyHandler { async fn get_flight_info_statement( &self, query: arrow_flight::sql::CommandStatementQuery, @@ -246,7 +246,7 @@ impl FlightSqlHandler for DfRayProxyHandler { })?; let plan = Arc::new( - DFRayStageReaderExec::try_new( + DDStageReaderExec::try_new( Partitioning::UnknownPartitioning(stage_partition_addrs.len()), Arc::new(schema), tsd.stage_id, @@ -268,21 +268,21 @@ impl FlightSqlHandler for DfRayProxyHandler { } } -/// DFRayProcessorService is a Arrow Flight service that serves streams of +/// DDProcessorService is a Arrow Flight service that serves streams of /// partitions from a hosted Physical Plan /// /// It only responds to the DoGet Arrow Flight method -pub struct DFRayProxyService { +pub struct DDProxyService { listener: TcpListener, - handler: Arc, + handler: Arc, addr: String, all_done_tx: Arc>>, all_done_rx: Option>, } -impl DFRayProxyService { +impl DDProxyService { pub async fn new(name: String, port: usize) -> Result { - debug!("Creating DFRayProxyService!"); + debug!("Creating DDProxyService!"); let (all_done_tx, all_done_rx) = channel(1); let all_done_tx = Arc::new(Mutex::new(all_done_tx)); @@ -291,9 +291,9 @@ impl DFRayProxyService { let addr = format!("{}", listener.local_addr().unwrap()); - info!("DFRayProcessorService bound to {addr}"); + info!("DDProcessorService bound to {addr}"); - let handler = Arc::new(DfRayProxyHandler::new(name, addr.clone())); + let handler = Arc::new(DDProxyHandler::new(name, addr.clone())); Ok(Self { listener, diff --git a/src/result.rs b/src/result.rs index e549257b..c769fe64 100644 --- a/src/result.rs +++ b/src/result.rs @@ -4,7 +4,7 @@ use thiserror::Error; use crate::vocab::Host; #[derive(Debug, Error)] -pub enum DFRayError { +pub enum DDError { #[error("Internal Arrow error: {0}")] ArrowError(#[from] arrow::error::ArrowError), #[error("Internal DataFusion error: {0}")] @@ -22,4 +22,4 @@ pub enum DFRayError { Other(#[from] anyhow::Error), } -pub type Result = std::result::Result; +pub type Result = std::result::Result; diff --git a/src/stage.rs b/src/stage.rs index 848d90c6..b66cf4f7 100644 --- a/src/stage.rs +++ b/src/stage.rs @@ -9,95 +9,8 @@ use datafusion::{ /// An execution plan that serves as a marker of where we want to split the /// physical plan into stages. -/// -/// This marker is consumed later by the [`crate::dataframe::RayDataFrame`], -/// when we are told to execute. It will create the discrete stages and insert -/// other ExecutionPlans to read and write the data -/// -/// # Example -/// The following query, -/// ```sql -/// select c.c_name, sum(o.o_totalprice) as total -/// from orders o inner join customer c on o.o_c ustkey = c.c_custkey -/// group by c_name limit 1 -/// ``` -/// -/// Will produce the following physical_plan from the optimizer -/// -/// ` -/// RayStageExec[3] (output_partitioning=UnknownPartitioning(1)) -/// ProjectionExec: expr=[c_name@0 as c_name, sum(o.o_totalprice)@1 as total] -/// GlobalLimitExec: skip=0, fetch=1 -/// CoalescePartitionsExec -/// AggregateExec: mode=FinalPartitioned, gby=[c_name@0 as c_name], -/// aggr=[sum(o.o_totalprice)] RayStageExec[2] -/// (output_partitioning=Hash([Column { name: "c_name", index: 0 }], 2)) -/// RepartitionExec: partitioning=Hash([c_name@0], 2), -/// input_partitions=2 AggregateExec: mode=Partial, gby=[c_name@1 -/// as c_name], aggr=[sum(o.o_totalprice)] ProjectionExec: -/// expr=[o_totalprice@1 as o_totalprice, c_name@0 as c_name] -/// HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c_custkey@0, -/// o_custkey@0)], projection=[c_name@1, o_totalprice@3] -/// RayStageExec[0] (output_partitioning=Hash([Column { name: "c_custkey", -/// index: 0 }], 2)) RepartitionExec: -/// partitioning=Hash([c_custkey@0], 2), input_partitions=1 -/// ParquetExec: file_groups={1 group: [[.../customer.parquet]]}, -/// projection=[c_custkey, c_name] RayStageExec[1] -/// (output_partitioning=Hash([Column { name: "o_custkey", index: 0 }], 2)) -/// RepartitionExec: partitioning=Hash([o_custkey@0], 2), -/// input_partitions=2 ParquetExec: file_groups={2 -/// groups: [[.../orders.parquet:0..19037604], -/// [.../orders.parquet:19037604..38075207]]}, projection=[o_custkey, -/// o_totalprice] ` -/// This physical plan will be split into 4 stages, as indicated by the -/// RayStageExec nodes. Those stages will look like this: -/// -/// ` -/// Stage 0 output partitions:2 shadow partitions: 1 -/// MaxRowsExec[max_rows=8192] -/// CoalesceBatchesExec: target_batch_size=8192 -/// RepartitionExec: partitioning=Hash([c_custkey@0], 2), input_partitions=1 -/// PartitionIsolatorExec -/// ParquetExec: file_groups={1 group: [[.../customer.parquet]]}, -/// projection=[c_custkey, c_name] -/// -/// Stage 1 output partitions:2 shadow partitions: 2 -/// MaxRowsExec[max_rows=8192] -/// CoalesceBatchesExec: target_batch_size=8192 -/// RepartitionExec: partitioning=Hash([o_custkey@0], 2), input_partitions=1 -/// PartitionIsolatorExec -/// ParquetExec: file_groups={2 groups: -/// [[.../orders.parquet:0..19037604], -/// [.../orders.parquet:19037604..38075207]]}, projection=[o_custkey, -/// o_totalprice] -/// -/// Stage 2 output partitions:2 shadow partitions: 2 -/// MaxRowsExec[max_rows=8192] -/// CoalesceBatchesExec: target_batch_size=8192 -/// RepartitionExec: partitioning=Hash([c_name@0], 2), input_partitions=1 -/// PartitionIsolatorExec -/// AggregateExec: mode=Partial, gby=[c_name@1 as c_name], -/// aggr=[sum(o.o_totalprice)] ProjectionExec: expr=[o_totalprice@1 as -/// o_totalprice, c_name@0 as c_name] HashJoinExec: -/// mode=Partitioned, join_type=Inner, on=[(c_custkey@0, o_custkey@0)], -/// projection=[c_name@1, o_totalprice@3] RayStageReaderExec[0] -/// (output_partitioning=UnknownPartitioning(2)) -/// RayStageReaderExec[1] (output_partitioning=UnknownPartitioning(2)) -/// -/// Stage 3 output partitions:1 shadow partitions: None -/// MaxRowsExec[max_rows=8192] -/// CoalesceBatchesExec: target_batch_size=8192 -/// ProjectionExec: expr=[c_name@0 as c_name, sum(o.o_totalprice)@1 as -/// total] GlobalLimitExec: skip=0, fetch=1 -/// CoalescePartitionsExec -/// AggregateExec: mode=FinalPartitioned, gby=[c_name@0 as c_name], -/// aggr=[sum(o.o_totalprice)] RayStageReaderExec[2] -/// (output_partitioning=UnknownPartitioning(2)) ` -/// -/// See [`crate::isolator::PartitionIsolatorExec`] for more information on how -/// the shadow partitions work #[derive(Debug)] -pub struct DFRayStageExec { +pub struct DDStageExec { /// Input plan pub(crate) input: Arc, /// Output partitioning @@ -105,7 +18,7 @@ pub struct DFRayStageExec { pub stage_id: u64, } -impl DFRayStageExec { +impl DDStageExec { pub fn new(input: Arc, stage_id: u64) -> Self { let properties = input.properties().clone(); @@ -128,18 +41,18 @@ impl DFRayStageExec { } } } -impl DisplayAs for DFRayStageExec { +impl DisplayAs for DDStageExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { write!( f, - "RayStageExec[{}] (output_partitioning={:?})", + "DDStageExec[{}] (output_partitioning={:?})", self.stage_id, self.properties().partitioning ) } } -impl ExecutionPlan for DFRayStageExec { +impl ExecutionPlan for DDStageExec { fn schema(&self) -> SchemaRef { self.input.schema() } @@ -148,7 +61,7 @@ impl ExecutionPlan for DFRayStageExec { } fn name(&self) -> &str { - "RayStageExec" + "DDStageExec" } fn as_any(&self) -> &dyn std::any::Any { @@ -167,24 +80,19 @@ impl ExecutionPlan for DFRayStageExec { assert_eq!(children.len(), 1); let child = children[0].clone(); - // as the plan tree is rearranged we want to remember the original partitioning - // that we had, even if we get new inputs. This is because - // RayStageReaderExecs, when created by the RayDataFrame will need to - // know the original partitioning - Ok(Arc::new(DFRayStageExec::new_with_properties( + Ok(Arc::new(DDStageExec::new_with_properties( child, self.stage_id, self.properties.clone(), ))) } - /// We will have to defer this functionality to python as Ray does not yet - /// have Rust bindings. fn execute( &self, _partition: usize, _context: std::sync::Arc, ) -> Result { - unimplemented!("Ray Stage Exec") + // we are only using this node as a marker in the plan, so we don't execute it + unimplemented!("DDStageExec") } } diff --git a/src/stage_reader.rs b/src/stage_reader.rs index 3f1343a7..989ef524 100644 --- a/src/stage_reader.rs +++ b/src/stage_reader.rs @@ -34,13 +34,13 @@ pub(crate) struct QueryId(pub String); /// [`crate::stage_service::ServiceClients`] and storing it as an extension in /// the [`datafusion::execution::TaskContext`] configuration. #[derive(Debug)] -pub struct DFRayStageReaderExec { +pub struct DDStageReaderExec { properties: PlanProperties, schema: SchemaRef, pub stage_id: u64, } -impl DFRayStageReaderExec { +impl DDStageReaderExec { pub fn try_new_from_input(input: Arc, stage_id: u64) -> Result { let properties = input.properties().clone(); @@ -62,18 +62,18 @@ impl DFRayStageReaderExec { }) } } -impl DisplayAs for DFRayStageReaderExec { +impl DisplayAs for DDStageReaderExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { write!( f, - "RayStageReaderExec[{}] (output_partitioning={:?})", + "DDStageReaderExec[{}] (output_partitioning={:?})", self.stage_id, self.properties().partitioning ) } } -impl ExecutionPlan for DFRayStageReaderExec { +impl ExecutionPlan for DDStageReaderExec { fn schema(&self) -> SchemaRef { self.schema.clone() } @@ -82,7 +82,7 @@ impl ExecutionPlan for DFRayStageReaderExec { } fn name(&self) -> &str { - "RayStageReaderExec" + "DDStageReaderExec" } fn as_any(&self) -> &dyn std::any::Any { @@ -106,7 +106,7 @@ impl ExecutionPlan for DFRayStageReaderExec { partition: usize, context: std::sync::Arc, ) -> Result { - let name = format!("RayStageReaderExec[{}-{}]:", self.stage_id, partition); + let name = format!("DDStageReaderExec[{}-{}]:", self.stage_id, partition); trace!("{name} execute: partition {partition}"); let stage_addrs = &context diff --git a/src/util.rs b/src/util.rs index f652ac3a..ff50cef6 100644 --- a/src/util.rs +++ b/src/util.rs @@ -54,11 +54,11 @@ use tonic::transport::Channel; use url::Url; use crate::{ - codec::DFRayCodec, + codec::DDCodec, logging::{debug, error, trace}, protobuf::StageAddrs, result::Result, - stage_reader::DFRayStageReaderExec, + stage_reader::DDStageReaderExec, vocab::{Addrs, Host}, }; @@ -138,7 +138,7 @@ pub fn physical_plan_to_bytes(plan: Arc) -> Result, D "serializing plan to bytes. plan:\n{}", display_plan_with_partition_counts(&plan) ); - let codec = DFRayCodec {}; + let codec = DDCodec {}; let proto = datafusion_proto::protobuf::PhysicalPlanNode::try_from_physical_plan(plan, &codec)?; let bytes = proto.encode_to_vec(); @@ -151,7 +151,7 @@ pub fn bytes_to_physical_plan( ) -> Result, DataFusionError> { let proto_plan = datafusion_proto::protobuf::PhysicalPlanNode::try_decode(plan_bytes)?; - let codec = DFRayCodec {}; + let codec = DDCodec {}; let plan = proto_plan.try_into_physical_plan(ctx, ctx.runtime_env().as_ref(), &codec)?; Ok(plan) } @@ -253,7 +253,7 @@ pub fn input_stage_ids(plan: &Arc) -> Result, DataFu let mut result = vec![]; plan.clone() .transform_down(|node: Arc| { - if let Some(reader) = node.as_any().downcast_ref::() { + if let Some(reader) = node.as_any().downcast_ref::() { result.push(reader.stage_id); } Ok(Transformed::no(node)) @@ -391,28 +391,6 @@ impl ProcessorClientFactory { } pub fn get_client(&self, host: &Host) -> Result { - // ideally we want to reuse channels as Tonic encourages cloning them when - // you can. This could would allow us to lazily create channels and keep - // them around so that on subsequent requests to the same address, we won't - // incur the penalty of establishing a socket connection. - // - // However this doesn't work at the moment. We encounter - // hangs or deadlock I cant tell which. Its not due to the lock, as far - // as I can tell, but rather something about the cloned channels themselves. - // - // In the case where we use a single DFProcessor, all connections will be to - // the same host, and from itself, to itself, so we'll have 100s of clones - // of a channel. Could this be the issue? - // - // TODO: figure out why this doesn't work - // - // UPDATE: This is repeatable for tpc - // RAY_DEDUP_LOGS=0 DATAFUSION_RAY_LOG_LEVEL=trace python tpch/tpcbench.py - // --data /path/to/data --concurrency 3 --partitions-per-processor=2 - // --processor-pool-min 1 --validate but goes away for when there are - // higher number of processors, even 2. I'm going to leave in the - // functionality for cached channels for now - let url = format!("http://{}", host.addr); let maybe_chan = self.channels.read().get(&host.addr).cloned(); diff --git a/src/worker_discovery.rs b/src/worker_discovery.rs index d4a78a93..36a5dee2 100644 --- a/src/worker_discovery.rs +++ b/src/worker_discovery.rs @@ -61,9 +61,9 @@ impl WorkerDiscovery { } fn start(&self) -> Result<()> { - let worker_addrs_env = std::env::var("DFRAY_WORKER_ADDRESSES"); - let worker_deployment_env = std::env::var("DFRAY_WORKER_DEPLOYMENT"); - let worker_deployment_namespace_env = std::env::var("DFRAY_WORKER_DEPLOYMENT_NAMESPACE"); + let worker_addrs_env = std::env::var("DD_WORKER_ADDRESSES"); + let worker_deployment_env = std::env::var("DD_WORKER_DEPLOYMENT"); + let worker_deployment_namespace_env = std::env::var("DD_WORKER_DEPLOYMENT_NAMESPACE"); if worker_addrs_env.is_ok() { let addresses = self.addresses.clone(); @@ -86,8 +86,8 @@ impl WorkerDiscovery { } else { // if neither env var is set, return an error return Err(anyhow!( - "Either DFRAY_WORKER_ADDRESSES or both DFRAY_WORKER_DEPLOYMENT and \ - DFRAY_WORKER_DEPLOYMENT_NAMESPACE must be set" + "Either DD_WORKER_ADDRESSES or both DD_WORKER_DEPLOYMENT and \ + DD_WORKER_DEPLOYMENT_NAMESPACE must be set" ) .into()); } @@ -264,7 +264,7 @@ async fn get_worker_info_from_pod(pod: &Pod) -> Result<(String, Host)> { .and_then(|spec| { spec.containers .iter() - .find(|c| c.name == "dfray-worker") + .find(|c| c.name == "dd-worker") .and_then(|c| { c.ports .as_ref() @@ -273,7 +273,7 @@ async fn get_worker_info_from_pod(pod: &Pod) -> Result<(String, Host)> { }) .ok_or_else(|| { anyhow::anyhow!( - "No could not find container port for container named dfray-worker found in pod {}", + "No could not find container port for container named dd-worker found in pod {}", pod.name_any() ) })?; diff --git a/tests/common/mod.rs b/tests/common/mod.rs index db821a0d..93d432c2 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -447,7 +447,7 @@ except Exception as e: let worker = Self::spawn_process( binary_path_str, &["--mode", "worker", "--port", &port.to_string()], - &[("DFRAY_TABLES", &tpch_tables), ("DFRAY_VIEWS", tpch_views)], + &[("DD_TABLES", &tpch_tables), ("DD_VIEWS", tpch_views)], &format!("start worker {}", i + 1), )?; self.worker_processes.push(worker); @@ -469,9 +469,9 @@ except Exception as e: binary_path_str, &["--mode", "proxy", "--port", &PROXY_PORT.to_string()], &[ - ("DFRAY_WORKER_ADDRESSES", &worker_addresses), - ("DFRAY_TABLES", &tpch_tables), - ("DFRAY_VIEWS", tpch_views), + ("DD_WORKER_ADDRESSES", &worker_addresses), + ("DD_TABLES", &tpch_tables), + ("DD_VIEWS", tpch_views), ], "start proxy", )?; @@ -770,7 +770,7 @@ pub async fn execute_query_distributed( ) -> Result<(String, Duration), Box> { let start_time = Instant::now(); - // Write SQL to temporary file (revenue0 view now handled via DFRAY_VIEWS env var) + // Write SQL to temporary file (revenue0 view now handled via DD_VIEWS env var) let temp_sql_file = ClusterManager::write_temp_file( &format!("{}_query.sql", query_name), sql, diff --git a/tpch/queries/q15.sql b/tpch/queries/q15.sql index 5129c292..004b97dc 100644 --- a/tpch/queries/q15.sql +++ b/tpch/queries/q15.sql @@ -11,7 +11,7 @@ -- and l_shipdate < date '1996-08-01' + interval '3' month -- group by -- l_suppkey; --- Note: The revenue0 view is created at cluster startup via DFRAY_VIEWS environment variable +-- Note: The revenue0 view is created at cluster startup via DD_VIEWS environment variable select s_suppkey, s_name, From 1537b7b61283f2d78318835d334d7e6baed154ec Mon Sep 17 00:00:00 2001 From: Rob Tandy Date: Fri, 11 Jul 2025 10:52:36 -0400 Subject: [PATCH 2/3] fix misrename, and also rename DDProcessor to DDWorker --- k8s/cluster.yaml | 6 +- k8s/dev_cluster.yaml | 6 +- src/lib.rs | 2 +- src/main.rs | 6 +- src/processor_service.rs | 696 --------------------------------------- src/proxy_service.rs | 7 +- src/util.rs | 54 +-- 7 files changed, 39 insertions(+), 738 deletions(-) delete mode 100644 src/processor_service.rs diff --git a/k8s/cluster.yaml b/k8s/cluster.yaml index aee0c5d2..91afa739 100644 --- a/k8s/cluster.yaml +++ b/k8s/cluster.yaml @@ -55,11 +55,11 @@ spec: - name: dd-proxy image: registry.ddbuild.io/dd-amd64:2025-06-17-b env: - - name: dd_WORKER_DEPLOYMENT + - name: DD_WORKER_DEPLOYMENT value: deployment-dd-worker - - name: dd_WORKER_DEPLOYMENT_PORT + - name: DD_WORKER_DEPLOYMENT_PORT value: "20201" - - name: dd_WORKER_DEPLOYMENT_NAMESPACE + - name: DD_WORKER_DEPLOYMENT_NAMESPACE valueFrom: fieldRef: fieldPath: metadata.namespace diff --git a/k8s/dev_cluster.yaml b/k8s/dev_cluster.yaml index 3a03b13d..df7b9d30 100644 --- a/k8s/dev_cluster.yaml +++ b/k8s/dev_cluster.yaml @@ -34,11 +34,11 @@ spec: - name: dd-proxy image: registry.ddbuild.io/dd:2025-06-16-e env: - - name: dd_WORKER_DEPLOYMENT + - name: DD_WORKER_DEPLOYMENT value: deployment-dd-worker - - name: dd_WORKER_DEPLOYMENT_PORT + - name: DD_WORKER_DEPLOYMENT_PORT value: "20201" - - name: dd_WORKER_DEPLOYMENT_NAMESPACE + - name: DD_WORKER_DEPLOYMENT_NAMESPACE valueFrom: fieldRef: fieldPath: metadata.namespace diff --git a/src/lib.rs b/src/lib.rs index f5fc516b..a40b4e5b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -32,7 +32,6 @@ pub mod logging; pub mod max_rows; pub mod physical; pub mod planning; -pub mod processor_service; pub mod proxy_service; pub mod query_planner; pub mod record_batch_exec; @@ -42,6 +41,7 @@ pub mod stage_reader; pub mod util; pub mod vocab; pub mod worker_discovery; +pub mod worker_service; #[cfg(not(target_env = "msvc"))] use tikv_jemallocator::Jemalloc; diff --git a/src/main.rs b/src/main.rs index 401f1a9a..400f0ea9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,8 @@ use anyhow::Result; use clap::Parser; use distributed_datafusion::{ - friendly::new_friendly_name, processor_service::DDProcessorService, - proxy_service::DDProxyService, setup, + friendly::new_friendly_name, proxy_service::DDProxyService, setup, + worker_service::DDWorkerService, }; #[derive(Parser)] @@ -37,7 +37,7 @@ async fn main() -> Result<()> { service.serve().await?; } "worker" => { - let service = DDProcessorService::new(new_friendly_name()?, args.port).await?; + let service = DDWorkerService::new(new_friendly_name()?, args.port).await?; service.serve().await?; } _ => { diff --git a/src/processor_service.rs b/src/processor_service.rs deleted file mode 100644 index 578a8e2e..00000000 --- a/src/processor_service.rs +++ /dev/null @@ -1,696 +0,0 @@ -// 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. - -use std::{ - collections::{HashMap, HashSet}, - fmt::Debug, - sync::Arc, - time::{Duration, SystemTime}, -}; - -use anyhow::{anyhow, Context}; -use arrow::array::RecordBatch; -use arrow_flight::{ - encode::FlightDataEncoderBuilder, error::FlightError, - flight_service_server::FlightServiceServer, Action, Ticket, -}; -use async_stream::stream; -use datafusion::{ - physical_plan::{ExecutionPlan, ExecutionPlanProperties}, - prelude::SessionContext, -}; -use futures::{StreamExt, TryStreamExt}; -use parking_lot::{Mutex, RwLock}; -use prost::Message; -use tokio::{ - net::TcpListener, - sync::mpsc::{channel, Receiver, Sender}, -}; -use tonic::{async_trait, transport::Server, Request, Response, Status}; - -use crate::{ - analyze::DistributedAnalyzeExec, - flight::{FlightHandler, FlightServ}, - logging::{debug, error, info, trace}, - planning::{add_ctx_extentions, get_ctx}, - protobuf::{ - AnnotatedTaskOutput, AnnotatedTaskOutputs, FlightDataMetadata, FlightTicketData, Host, - }, - result::{DDError, Result}, - util::{ - bytes_to_physical_plan, display_plan_with_partition_counts, get_addrs, - register_object_store_for_paths_in_plan, start_up, - }, - vocab::{Addrs, CtxAnnotatedOutputs, CtxPartitionGroup, DDTask}, -}; - -#[derive(Eq, PartialEq, Hash, Clone, Debug)] -struct StageKey { - query_id: String, - stage_id: u64, -} - -#[derive(Clone)] -struct Task { - partitions: HashSet, - ctx: SessionContext, - plan: Arc, -} - -impl Debug for Task { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Task") - .field("partitions", &self.partitions) - .finish() - } -} - -#[derive(Debug)] -struct StageTasks { - tasks: Vec, - insert_time: SystemTime, -} - -/// It only responds to the DoGet Arrow Flight method. -struct DDProcessorHandler { - /// our name, useful for logging - name: String, - /// our address string, also useful for logging - addr: String, - /// our map of query_id -> (session ctx, execution plan) - stages: Arc>>, - done: Arc>, -} - -impl DDProcessorHandler { - pub fn new(name: String, addr: String) -> Self { - let stages: Arc>> = - Arc::new(RwLock::new(HashMap::new())); - let done = Arc::new(Mutex::new(false)); - - // start a plan janitor ask to clean up old plans that were not collected for - // any reason - let c_stages = stages.clone(); - let c_done = done.clone(); - let c_name = name.clone(); - std::thread::spawn(move || { - while !(*c_done.lock()) { - // wait for 10 seconds - std::thread::sleep(Duration::from_secs(10)); - if *c_done.lock() { - info!("{} janitor done", c_name); - break; - } - trace!("{} janitor waking up", c_name); - - let now = SystemTime::now(); - - let mut to_remove = vec![]; - { - let _guard = c_stages.read(); - for (key, stage_tasks) in _guard.iter() { - if stage_tasks.tasks.is_empty() { - error!("unexpectedly found empty stage tasks. removing"); - to_remove.push(key.clone()); - } else { - // check if any plan in this vec is older than 1 minute - if now - .duration_since(stage_tasks.insert_time) - .map(|d| d.as_secs() > 60) - .inspect_err(|e| { - error!("CANNOT COMPUTE DURATION OR REMOVE STAGES: {e:?}"); - // maybe just panic here? - }) - .unwrap_or(false) - { - to_remove.push(key.clone()); - break; - } - } - } - } - if !to_remove.is_empty() { - let mut _guard = c_stages.write(); - for key in to_remove.iter() { - _guard.remove(key); - debug!("{} removed old stage key {:?}", c_name, key); - } - } - } - }); - - Self { - name, - addr, - stages, - done, - } - } - - #[allow(dead_code)] - /// shutdown - pub fn all_done(&self) { - *self.done.lock() = true; - } - - pub async fn add_task( - &self, - query_id: String, - stage_id: u64, - stage_addrs: Addrs, - partition_group: Vec, - full_partitions: bool, - plan_bytes: &[u8], - ) -> Result<()> { - let ctx = self - .configure_ctx( - query_id.clone(), - stage_id, - stage_addrs.clone(), - partition_group.clone(), - ) - .await?; - - let plan = bytes_to_physical_plan(&ctx, plan_bytes).context(format!( - "{}, Could not decode plan for query_id {} stage {}", - self.name, query_id, stage_id - ))?; - - let partitions = if full_partitions { - partition_group.clone() - } else { - (0..(plan.output_partitioning().partition_count())) - .map(|p| p as u64) - .collect::>() - }; - - trace!( - "{} adding task for stage {} partition group: {:?} partitions: {:?} stage_addrs: {:?} plan:\n{}", - self.name, - stage_id, - partition_group, - partitions, - stage_addrs, - display_plan_with_partition_counts(&plan) - ); - - register_object_store_for_paths_in_plan(&ctx, plan.clone())?; - - let now = SystemTime::now(); - - let task = Task { - ctx: ctx.clone(), - plan: plan.clone(), - partitions: HashSet::from_iter(partitions.clone()), - }; - - let key = StageKey { - query_id: query_id.clone(), - stage_id, - }; - { - let mut _guard = self.stages.write(); - let stage_tasks = _guard.entry(key.clone()).or_insert_with(|| StageTasks { - tasks: vec![], - insert_time: now, - }); - stage_tasks.tasks.push(task.clone()); - trace!("{} added task for stage key {:?}", self.name, key); - } - - Ok(()) - } - - async fn configure_ctx( - &self, - query_id: String, - stage_id: u64, - stage_addrs: Addrs, - partition_group: Vec, - ) -> Result { - let mut ctx = get_ctx()?; - let host = Host { - addr: self.addr.clone(), - name: self.name.clone(), - }; - - add_ctx_extentions( - &mut ctx, - &host, - &query_id, - stage_id, - stage_addrs.clone(), - partition_group, - )?; - - Ok(ctx) - } - - /// Retrieve the requested ctx and plan to execute. Also return a bool - /// indicating if this is the last partition for this plan - fn get_ctx_and_plan( - &self, - query_id: &str, - stage_id: u64, - partition: u64, - ) -> Result<(SessionContext, Arc, bool)> { - let stage_key = StageKey { - query_id: query_id.to_string(), - stage_id, - }; - - let (ctx, plan, last) = { - let mut _guard = self.stages.write(); - let stage_tasks = _guard.get_mut(&stage_key).context(format!( - "{}, No plan found for stage key{:?}", - self.name, stage_key, - ))?; - trace!( - "{} found {} tasks for stage key {:?}", - self.name, - stage_tasks.tasks.len(), - stage_key - ); - - // of the tasks for this stage, find one that has this partition not yet consumed - let task = stage_tasks - .tasks - .iter_mut() - .find(|t| t.partitions.contains(&partition)) - .context(format!( - "{}, No task found for stage key {:?} partition {}", - self.name, stage_key, partition - ))?; - - // remove this partition from the list of partitions yet to be consumed - if !task.partitions.remove(&partition) { - // it should be in there because we just filtered for it - return Err(anyhow!("UNEXPECTED: partition {partition} not in plan parts").into()); - } - // finally, we return the ctx and plan to execute this task, as well as a - // bool indicating if this is the last partition for this task - ( - task.ctx.clone(), - task.plan.clone(), - task.partitions.is_empty(), - ) - }; - { - // some house keeping, if there are no more partitions left for this task, - // remove it from the stage tasks. - let mut _guard = self.stages.write(); - - let remove_it = _guard.get_mut(&stage_key).map(|stage_tasks| { - stage_tasks.tasks.retain(|t| !t.partitions.is_empty()); - trace!( - "remaining tasks: for stage {:?}: {:?}", - stage_key, - stage_tasks - ); - - // furthermore, if there are no more tasks left for this stage, - // remove the stage from the map - stage_tasks.tasks.is_empty() - }); - - if remove_it.unwrap_or(false) { - _guard.remove(&stage_key); - debug!("{} removed stage key {:?}", self.name, stage_key); - } - } - - Ok((ctx, plan, last)) - } - - /// we want to send any additional FlightDataMetadata that we discover while - /// executing the plan further downstream to consumers of our stream. - /// - /// We may discover a FlightDataMetadata as an additional payload incoming - /// on streming requests. The [`DDStageReader`] will look for these and add - /// them to an extension on the context. - /// - /// Our job is to send all record batches from our stream over our flight response - /// but when the stream is exhausted, we'll send an additional message containing the - /// metadata. - /// - /// The reason we have to do it at the end is that the metadata may include an annodated - /// plan with metrics which will only be available after the stream has been - /// fully consumed. - fn make_stream( - &self, - ctx: SessionContext, - plan: Arc, - stage_id: u64, - partition: u64, - last_partition: bool, - ) -> Result { - let task_ctx = ctx.task_ctx(); - - // the RecordBatchStream from our plan - let stream = plan - .execute(partition as usize, task_ctx) - .inspect_err(|e| error!("Could not get partition stream from plan {e:#?}"))? - .inspect(|batch| { - trace!("producing maybe batch {:?}", batch); - }) - .map_err(|e| FlightError::from_external_error(Box::new(e))); - - info!("{} tasks held {}", self.name, self.stages.read().len()); - - fn find_analyze(plan: &dyn ExecutionPlan) -> Option<&DistributedAnalyzeExec> { - if let Some(target) = plan.as_any().downcast_ref::() { - Some(target) - } else { - for child in plan.children() { - if let Some(target) = find_analyze(child.as_ref()) { - return Some(target); - } - } - None - } - } - - let mut flight_data_stream = FlightDataEncoderBuilder::new().build(stream); - let name = self.name.clone(); - let host = Host { - addr: self.addr.clone(), - name: self.name.clone(), - }; - - #[allow(unused_assignments)] // clippy can't understand our assignment to done in the macro - let out_stream = async_stream::stream! { - let mut done = false; - while !done { - - match (done, last_partition, flight_data_stream.next().await) { - (false, false, None) => { - // we finished a partition, but still have more to do, do nothing - done = true; - } - (false, true, None) => { - // no more data in the last partition, so now we yield our additional FlightDataMetadata - debug!("stream exhausted, yielding FlightDataMetadata"); - let task_outputs = ctx.state().config() - .get_extension::() - .unwrap_or(Arc::new(CtxAnnotatedOutputs::default())) - .0 - .clone(); - - let partition_group = ctx.state().config() - .get_extension::() - .expect("CtxPartitionGroup to be set") - .0.clone(); - - - if let Some(analyze) = find_analyze(plan.as_ref()) { - let annotated_plan = analyze.annotated_plan(); - debug!("sending annotated plan: {}", annotated_plan); - - let output = AnnotatedTaskOutput { - plan: annotated_plan, - host: Some(host.clone()), - stage_id, - partition_group, - }; - task_outputs.lock().push(output); - } - - let meta = FlightDataMetadata { - annotated_task_outputs: Some(AnnotatedTaskOutputs { - outputs: task_outputs.lock().clone(), - }), - }; - - - let fake_batch = RecordBatch::new_empty(plan.schema()); - - - let mut fde = FlightDataEncoderBuilder::new() - //.with_schema(plan.schema()) - .with_metadata(meta.encode_to_vec().into()).build(futures::stream::once(async {Ok(fake_batch)})).map_err(|e| { - FlightError::from_external_error(Box::new(e)) - }); - - let flight_data = fde - .next() - .await - .ok_or_else(|| { - Status::internal(format!( - "{}, No FlightDataMetadata from our fde", - name - )) - })??; - - yield Ok(flight_data); - done = true; - - }, - (false, _, Some(Err(e))) => { - yield Err(Status::internal(format!( - "Unexpected error getting flight data stream: {e:?}", - ))); - done = true; - }, - (false, _, Some(Ok(flight_data))) => { - // we have a flight data, so we yield it - trace!("received normal flight data, yielding"); - yield Ok(flight_data); - }, - (true, false, _ ) => { - // we've finished a partition, do nothing - done = true; - }, - (true, true, _ ) => { - yield Err(Status::internal(format!("{name} reached expected unreachable block"))); - }, - } - } - }; - - Ok(Box::pin(out_stream)) - } - fn do_action_get_host(&self) -> Result, Status> { - let addr = self.addr.clone(); - let name = self.name.clone(); - - let out_stream = Box::pin(stream! { - yield Ok::<_, tonic::Status>(arrow_flight::Result { - body: Host{ - addr, - name, - }.encode_to_vec().into() - }); - }) as crate::flight::DoActionStream; - - Ok(Response::new(out_stream)) - } - - async fn do_action_add_plan( - &self, - action: Action, - ) -> Result, Status> { - let task_data = DDTask::decode(action.body.as_ref()).map_err(|e| { - Status::internal(format!( - "{}, Unexpected error decoding StageData: {e:?}", - self.name - )) - })?; - - let addrs = task_data - .stage_addrs - .as_ref() - .context("stage addrs not present") - .map_err(DDError::from) - .and_then(get_addrs) - .map_err(|e| Status::internal(format!("{}, {e}", self.name)))?; - - self.add_task( - task_data.query_id, - task_data.stage_id, - addrs, - task_data.partition_group, - task_data.full_partitions, - &task_data.plan_bytes, - ) - .await - .map_err(|e| Status::internal(format!("{}, Could not add plan: {e:?}", self.name)))?; - - let out_stream = Box::pin(stream! { - yield Ok::<_, tonic::Status>(arrow_flight::Result::default()); - }) as crate::flight::DoActionStream; - - Ok(Response::new(out_stream)) - } -} - -#[async_trait] -impl FlightHandler for DDProcessorHandler { - async fn do_get( - &self, - request: Request, - ) -> std::result::Result, Status> { - let remote_addr = request - .remote_addr() - .map(|a| a.to_string()) - .unwrap_or("unknown".to_string()); - - let ticket = request.into_inner(); - - let ftd = FlightTicketData::decode(ticket.ticket).map_err(|e| { - Status::internal(format!( - "{}, Unexpected error extracting ticket {e:?}", - self.name - )) - })?; - - let task_key = StageKey { - query_id: ftd.query_id.clone(), - stage_id: ftd.stage_id, - }; - - debug!( - "{}, request for task_key:{:?} partition: {} from: {},{}", - self.name, task_key, ftd.partition, ftd.requestor_name, remote_addr - ); - - let name = self.name.clone(); - let (ctx, plan, is_last_partition) = self - .get_ctx_and_plan(&ftd.query_id, ftd.stage_id, ftd.partition) - .map_err(|e| { - Status::internal(format!( - "{name} Could not find task for query_id {} stage {} partition {}: {e:?}", - ftd.query_id, ftd.stage_id, ftd.partition - )) - })?; - - let do_get_stream = self - .make_stream(ctx, plan, ftd.stage_id, ftd.partition, is_last_partition) - .map_err(|e| { - Status::internal(format!( - "{name} Could not make stream for query_id {} stage {} partition {}: {e:?}", - ftd.query_id, ftd.stage_id, ftd.partition - )) - })?; - - Ok(Response::new(do_get_stream)) - } - - async fn do_action( - &self, - request: Request, - ) -> Result, Status> { - // extract a StageData protobuf from the action body - let action = request.into_inner(); - let action_type = action.r#type.as_str(); - trace!("{} received action: {}", self.name, action_type); - - if action_type == "add_plan" { - self.do_action_add_plan(action).await - } else if action_type == "get_host" { - self.do_action_get_host() - } else { - Err(Status::unimplemented(format!( - "{}, Unimplemented action: {}", - self.name, action_type - ))) - } - } -} - -/// DDProcessorService is a Arrow Flight service that serves streams of -/// partitions from a hosted Physical Plan -/// -/// It only responds to the DoGet Arrow Flight method -pub struct DDProcessorService { - #[allow(dead_code)] - name: String, - listener: TcpListener, - handler: Arc, - addr: String, - all_done_tx: Arc>>, - all_done_rx: Option>, -} - -impl DDProcessorService { - pub async fn new(name: String, port: usize) -> Result { - let name = format!("[{}]", name); - - let (all_done_tx, all_done_rx) = channel(1); - let all_done_tx = Arc::new(Mutex::new(all_done_tx)); - - let listener = start_up(port).await?; - - let addr = format!("{}", listener.local_addr().unwrap()); - - info!("DDProcessorService bound to {addr}"); - - let handler = Arc::new(DDProcessorHandler::new(name.clone(), addr.clone())); - - Ok(Self { - name, - listener, - handler, - addr, - all_done_tx, - all_done_rx: Some(all_done_rx), - }) - } - - /// get the address of the listing socket for this service - pub fn addr(&self) -> String { - self.addr.clone() - } - - pub async fn all_done(&self) -> Result<()> { - let sender = self.all_done_tx.lock().clone(); - - sender - .send(()) - .await - .context("Could not send shutdown signal")?; - Ok(()) - } - - /// start the service, consuming self - pub async fn serve(mut self) -> Result<()> { - let mut all_done_rx = self.all_done_rx.take().unwrap(); - - let signal = async move { - all_done_rx - .recv() - .await - .expect("problem receiving shutdown signal"); - info!("received shutdown signal"); - }; - - let flight_serv = FlightServ { - handler: self.handler.clone(), - }; - - let svc = FlightServiceServer::new(flight_serv); - - Server::builder() - .add_service(svc) - .serve_with_incoming_shutdown( - tokio_stream::wrappers::TcpListenerStream::new(self.listener), - signal, - ) - .await - .context("error running service")?; - Ok(()) - } -} diff --git a/src/proxy_service.rs b/src/proxy_service.rs index 1b1c5dd0..95a30860 100644 --- a/src/proxy_service.rs +++ b/src/proxy_service.rs @@ -268,10 +268,7 @@ impl FlightSqlHandler for DDProxyHandler { } } -/// DDProcessorService is a Arrow Flight service that serves streams of -/// partitions from a hosted Physical Plan -/// -/// It only responds to the DoGet Arrow Flight method +/// An Arrow Flight SQL service pub struct DDProxyService { listener: TcpListener, handler: Arc, @@ -291,7 +288,7 @@ impl DDProxyService { let addr = format!("{}", listener.local_addr().unwrap()); - info!("DDProcessorService bound to {addr}"); + info!("DDProxyService bound to {addr}"); let handler = Arc::new(DDProxyHandler::new(name, addr.clone())); diff --git a/src/util.rs b/src/util.rs index ff50cef6..20b04b13 100644 --- a/src/util.rs +++ b/src/util.rs @@ -325,16 +325,16 @@ pub fn reporting_stream( Box::pin(RecordBatchStreamAdapter::new(schema, out_stream)) as SendableRecordBatchStream } -pub struct ProcessorClient { +pub struct WorkerClient { /// the host we are connecting to pub(crate) host: Host, - /// The flight client to the processor + /// The flight client to the worker inner: FlightClient, /// the channel cache in the factory channels: Arc>>, } -impl ProcessorClient { +impl WorkerClient { pub fn new( host: Host, inner: FlightClient, @@ -351,10 +351,12 @@ impl ProcessorClient { &mut self, ticket: Ticket, ) -> arrow_flight::error::Result { - let stream = self.inner.do_get(ticket).await - .inspect_err(|e| { - error!("Error in do_get for processor {}: {e:?}. - Considering this channel poisoned and removing it from ProcessorClientFactory cache", self.host); + let stream = self.inner.do_get(ticket).await.inspect_err(|e| { + error!( + "Error in do_get for worker {}: {e:?}. + Considering this channel poisoned and removing it from WorkerClientFactory cache", + self.host + ); self.channels.write().remove(&self.host.addr); })?; @@ -367,8 +369,8 @@ impl ProcessorClient { ) -> arrow_flight::error::Result>> { let result = self.inner.do_action(action).await.inspect_err(|e| { error!( - "Error in do_action for processor {}: {e:?}. - Considering this channel poisoned and removing it from ProcessorClientFactory \ + "Error in do_action for worker {}: {e:?}. + Considering this channel poisoned and removing it from WorkerClientFactory \ cache", self.host ); @@ -379,49 +381,47 @@ impl ProcessorClient { } } -struct ProcessorClientFactory { +struct WorkerClientFactory { channels: Arc>>, } -impl ProcessorClientFactory { +impl WorkerClientFactory { fn new() -> Self { Self { channels: Arc::new(RwLock::new(HashMap::new())), } } - pub fn get_client(&self, host: &Host) -> Result { + pub fn get_client(&self, host: &Host) -> Result { let url = format!("http://{}", host.addr); let maybe_chan = self.channels.read().get(&host.addr).cloned(); let chan = match maybe_chan { Some(chan) => { - debug!("ProcessorFactory using cached channel for {host}"); + debug!("WorkerFactory using cached channel for {host}"); chan } None => { let host_str = host.to_string(); let fut = async move { - trace!("ProcessorFactory connecting to {host_str}"); + trace!("WorkerFactory connecting to {host_str}"); Channel::from_shared(url.clone()) - .map_err(|e| { - internal_datafusion_err!("ProcessorFactory invalid url {e:#?}") - })? + .map_err(|e| internal_datafusion_err!("WorkerFactory invalid url {e:#?}"))? // FIXME: update timeout value to not be a magic number .connect_timeout(Duration::from_secs(2)) .connect() .await .map_err(|e| { - internal_datafusion_err!("ProcessorFactory cannot connect {e:#?}") + internal_datafusion_err!("WorkerFactory cannot connect {e:#?}") }) }; - let chan = wait_for(fut, "ProcessorFactory::get_client").map_err(|e| { + let chan = wait_for(fut, "WorkerFactory::get_client").map_err(|e| { internal_datafusion_err!( - "ProcessorFactory Cannot wait for channel connect future {e:#?}" + "WorkerFactory Cannot wait for channel connect future {e:#?}" ) })??; - trace!("ProcessorFactory connected to {host}"); + trace!("WorkerFactory connected to {host}"); self.channels .write() .insert(host.addr.to_string(), chan.clone()); @@ -429,11 +429,11 @@ impl ProcessorClientFactory { chan } }; - debug!("ProcessorFactory have channel now for {host}"); + debug!("WorkerFactory have channel now for {host}"); let flight_client = FlightClient::new(chan); - debug!("ProcessorFactory made flight client for {host}"); - Ok(ProcessorClient::new( + debug!("WorkerFactory made flight client for {host}"); + Ok(WorkerClient::new( host.clone(), flight_client, self.channels.clone(), @@ -441,10 +441,10 @@ impl ProcessorClientFactory { } } -static FACTORY: OnceLock = OnceLock::new(); +static FACTORY: OnceLock = OnceLock::new(); -pub fn get_client(host: &Host) -> Result { - let factory = FACTORY.get_or_init(ProcessorClientFactory::new); +pub fn get_client(host: &Host) -> Result { + let factory = FACTORY.get_or_init(WorkerClientFactory::new); factory.get_client(host) } From 4404afa5681e6bdac75f7a13239de27c4afb64ce Mon Sep 17 00:00:00 2001 From: Rob Tandy Date: Fri, 11 Jul 2025 12:18:11 -0400 Subject: [PATCH 3/3] add file that was renamed --- src/worker_service.rs | 696 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 696 insertions(+) create mode 100644 src/worker_service.rs diff --git a/src/worker_service.rs b/src/worker_service.rs new file mode 100644 index 00000000..8500984f --- /dev/null +++ b/src/worker_service.rs @@ -0,0 +1,696 @@ +// 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. + +use std::{ + collections::{HashMap, HashSet}, + fmt::Debug, + sync::Arc, + time::{Duration, SystemTime}, +}; + +use anyhow::{anyhow, Context}; +use arrow::array::RecordBatch; +use arrow_flight::{ + encode::FlightDataEncoderBuilder, error::FlightError, + flight_service_server::FlightServiceServer, Action, Ticket, +}; +use async_stream::stream; +use datafusion::{ + physical_plan::{ExecutionPlan, ExecutionPlanProperties}, + prelude::SessionContext, +}; +use futures::{StreamExt, TryStreamExt}; +use parking_lot::{Mutex, RwLock}; +use prost::Message; +use tokio::{ + net::TcpListener, + sync::mpsc::{channel, Receiver, Sender}, +}; +use tonic::{async_trait, transport::Server, Request, Response, Status}; + +use crate::{ + analyze::DistributedAnalyzeExec, + flight::{FlightHandler, FlightServ}, + logging::{debug, error, info, trace}, + planning::{add_ctx_extentions, get_ctx}, + protobuf::{ + AnnotatedTaskOutput, AnnotatedTaskOutputs, FlightDataMetadata, FlightTicketData, Host, + }, + result::{DDError, Result}, + util::{ + bytes_to_physical_plan, display_plan_with_partition_counts, get_addrs, + register_object_store_for_paths_in_plan, start_up, + }, + vocab::{Addrs, CtxAnnotatedOutputs, CtxPartitionGroup, DDTask}, +}; + +#[derive(Eq, PartialEq, Hash, Clone, Debug)] +struct StageKey { + query_id: String, + stage_id: u64, +} + +#[derive(Clone)] +struct Task { + partitions: HashSet, + ctx: SessionContext, + plan: Arc, +} + +impl Debug for Task { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Task") + .field("partitions", &self.partitions) + .finish() + } +} + +#[derive(Debug)] +struct StageTasks { + tasks: Vec, + insert_time: SystemTime, +} + +/// It only responds to the DoGet Arrow Flight method. +struct DDWorkerHandler { + /// our name, useful for logging + name: String, + /// our address string, also useful for logging + addr: String, + /// our map of query_id -> (session ctx, execution plan) + stages: Arc>>, + done: Arc>, +} + +impl DDWorkerHandler { + pub fn new(name: String, addr: String) -> Self { + let stages: Arc>> = + Arc::new(RwLock::new(HashMap::new())); + let done = Arc::new(Mutex::new(false)); + + // start a plan janitor ask to clean up old plans that were not collected for + // any reason + let c_stages = stages.clone(); + let c_done = done.clone(); + let c_name = name.clone(); + std::thread::spawn(move || { + while !(*c_done.lock()) { + // wait for 10 seconds + std::thread::sleep(Duration::from_secs(10)); + if *c_done.lock() { + info!("{} janitor done", c_name); + break; + } + trace!("{} janitor waking up", c_name); + + let now = SystemTime::now(); + + let mut to_remove = vec![]; + { + let _guard = c_stages.read(); + for (key, stage_tasks) in _guard.iter() { + if stage_tasks.tasks.is_empty() { + error!("unexpectedly found empty stage tasks. removing"); + to_remove.push(key.clone()); + } else { + // check if any plan in this vec is older than 1 minute + if now + .duration_since(stage_tasks.insert_time) + .map(|d| d.as_secs() > 60) + .inspect_err(|e| { + error!("CANNOT COMPUTE DURATION OR REMOVE STAGES: {e:?}"); + // maybe just panic here? + }) + .unwrap_or(false) + { + to_remove.push(key.clone()); + break; + } + } + } + } + if !to_remove.is_empty() { + let mut _guard = c_stages.write(); + for key in to_remove.iter() { + _guard.remove(key); + debug!("{} removed old stage key {:?}", c_name, key); + } + } + } + }); + + Self { + name, + addr, + stages, + done, + } + } + + #[allow(dead_code)] + /// shutdown + pub fn all_done(&self) { + *self.done.lock() = true; + } + + pub async fn add_task( + &self, + query_id: String, + stage_id: u64, + stage_addrs: Addrs, + partition_group: Vec, + full_partitions: bool, + plan_bytes: &[u8], + ) -> Result<()> { + let ctx = self + .configure_ctx( + query_id.clone(), + stage_id, + stage_addrs.clone(), + partition_group.clone(), + ) + .await?; + + let plan = bytes_to_physical_plan(&ctx, plan_bytes).context(format!( + "{}, Could not decode plan for query_id {} stage {}", + self.name, query_id, stage_id + ))?; + + let partitions = if full_partitions { + partition_group.clone() + } else { + (0..(plan.output_partitioning().partition_count())) + .map(|p| p as u64) + .collect::>() + }; + + trace!( + "{} adding task for stage {} partition group: {:?} partitions: {:?} stage_addrs: {:?} plan:\n{}", + self.name, + stage_id, + partition_group, + partitions, + stage_addrs, + display_plan_with_partition_counts(&plan) + ); + + register_object_store_for_paths_in_plan(&ctx, plan.clone())?; + + let now = SystemTime::now(); + + let task = Task { + ctx: ctx.clone(), + plan: plan.clone(), + partitions: HashSet::from_iter(partitions.clone()), + }; + + let key = StageKey { + query_id: query_id.clone(), + stage_id, + }; + { + let mut _guard = self.stages.write(); + let stage_tasks = _guard.entry(key.clone()).or_insert_with(|| StageTasks { + tasks: vec![], + insert_time: now, + }); + stage_tasks.tasks.push(task.clone()); + trace!("{} added task for stage key {:?}", self.name, key); + } + + Ok(()) + } + + async fn configure_ctx( + &self, + query_id: String, + stage_id: u64, + stage_addrs: Addrs, + partition_group: Vec, + ) -> Result { + let mut ctx = get_ctx()?; + let host = Host { + addr: self.addr.clone(), + name: self.name.clone(), + }; + + add_ctx_extentions( + &mut ctx, + &host, + &query_id, + stage_id, + stage_addrs.clone(), + partition_group, + )?; + + Ok(ctx) + } + + /// Retrieve the requested ctx and plan to execute. Also return a bool + /// indicating if this is the last partition for this plan + fn get_ctx_and_plan( + &self, + query_id: &str, + stage_id: u64, + partition: u64, + ) -> Result<(SessionContext, Arc, bool)> { + let stage_key = StageKey { + query_id: query_id.to_string(), + stage_id, + }; + + let (ctx, plan, last) = { + let mut _guard = self.stages.write(); + let stage_tasks = _guard.get_mut(&stage_key).context(format!( + "{}, No plan found for stage key{:?}", + self.name, stage_key, + ))?; + trace!( + "{} found {} tasks for stage key {:?}", + self.name, + stage_tasks.tasks.len(), + stage_key + ); + + // of the tasks for this stage, find one that has this partition not yet consumed + let task = stage_tasks + .tasks + .iter_mut() + .find(|t| t.partitions.contains(&partition)) + .context(format!( + "{}, No task found for stage key {:?} partition {}", + self.name, stage_key, partition + ))?; + + // remove this partition from the list of partitions yet to be consumed + if !task.partitions.remove(&partition) { + // it should be in there because we just filtered for it + return Err(anyhow!("UNEXPECTED: partition {partition} not in plan parts").into()); + } + // finally, we return the ctx and plan to execute this task, as well as a + // bool indicating if this is the last partition for this task + ( + task.ctx.clone(), + task.plan.clone(), + task.partitions.is_empty(), + ) + }; + { + // some house keeping, if there are no more partitions left for this task, + // remove it from the stage tasks. + let mut _guard = self.stages.write(); + + let remove_it = _guard.get_mut(&stage_key).map(|stage_tasks| { + stage_tasks.tasks.retain(|t| !t.partitions.is_empty()); + trace!( + "remaining tasks: for stage {:?}: {:?}", + stage_key, + stage_tasks + ); + + // furthermore, if there are no more tasks left for this stage, + // remove the stage from the map + stage_tasks.tasks.is_empty() + }); + + if remove_it.unwrap_or(false) { + _guard.remove(&stage_key); + debug!("{} removed stage key {:?}", self.name, stage_key); + } + } + + Ok((ctx, plan, last)) + } + + /// we want to send any additional FlightDataMetadata that we discover while + /// executing the plan further downstream to consumers of our stream. + /// + /// We may discover a FlightDataMetadata as an additional payload incoming + /// on streming requests. The [`DDStageReader`] will look for these and add + /// them to an extension on the context. + /// + /// Our job is to send all record batches from our stream over our flight response + /// but when the stream is exhausted, we'll send an additional message containing the + /// metadata. + /// + /// The reason we have to do it at the end is that the metadata may include an annodated + /// plan with metrics which will only be available after the stream has been + /// fully consumed. + fn make_stream( + &self, + ctx: SessionContext, + plan: Arc, + stage_id: u64, + partition: u64, + last_partition: bool, + ) -> Result { + let task_ctx = ctx.task_ctx(); + + // the RecordBatchStream from our plan + let stream = plan + .execute(partition as usize, task_ctx) + .inspect_err(|e| error!("Could not get partition stream from plan {e:#?}"))? + .inspect(|batch| { + trace!("producing maybe batch {:?}", batch); + }) + .map_err(|e| FlightError::from_external_error(Box::new(e))); + + info!("{} tasks held {}", self.name, self.stages.read().len()); + + fn find_analyze(plan: &dyn ExecutionPlan) -> Option<&DistributedAnalyzeExec> { + if let Some(target) = plan.as_any().downcast_ref::() { + Some(target) + } else { + for child in plan.children() { + if let Some(target) = find_analyze(child.as_ref()) { + return Some(target); + } + } + None + } + } + + let mut flight_data_stream = FlightDataEncoderBuilder::new().build(stream); + let name = self.name.clone(); + let host = Host { + addr: self.addr.clone(), + name: self.name.clone(), + }; + + #[allow(unused_assignments)] // clippy can't understand our assignment to done in the macro + let out_stream = async_stream::stream! { + let mut done = false; + while !done { + + match (done, last_partition, flight_data_stream.next().await) { + (false, false, None) => { + // we finished a partition, but still have more to do, do nothing + done = true; + } + (false, true, None) => { + // no more data in the last partition, so now we yield our additional FlightDataMetadata + debug!("stream exhausted, yielding FlightDataMetadata"); + let task_outputs = ctx.state().config() + .get_extension::() + .unwrap_or(Arc::new(CtxAnnotatedOutputs::default())) + .0 + .clone(); + + let partition_group = ctx.state().config() + .get_extension::() + .expect("CtxPartitionGroup to be set") + .0.clone(); + + + if let Some(analyze) = find_analyze(plan.as_ref()) { + let annotated_plan = analyze.annotated_plan(); + debug!("sending annotated plan: {}", annotated_plan); + + let output = AnnotatedTaskOutput { + plan: annotated_plan, + host: Some(host.clone()), + stage_id, + partition_group, + }; + task_outputs.lock().push(output); + } + + let meta = FlightDataMetadata { + annotated_task_outputs: Some(AnnotatedTaskOutputs { + outputs: task_outputs.lock().clone(), + }), + }; + + + let fake_batch = RecordBatch::new_empty(plan.schema()); + + + let mut fde = FlightDataEncoderBuilder::new() + //.with_schema(plan.schema()) + .with_metadata(meta.encode_to_vec().into()).build(futures::stream::once(async {Ok(fake_batch)})).map_err(|e| { + FlightError::from_external_error(Box::new(e)) + }); + + let flight_data = fde + .next() + .await + .ok_or_else(|| { + Status::internal(format!( + "{}, No FlightDataMetadata from our fde", + name + )) + })??; + + yield Ok(flight_data); + done = true; + + }, + (false, _, Some(Err(e))) => { + yield Err(Status::internal(format!( + "Unexpected error getting flight data stream: {e:?}", + ))); + done = true; + }, + (false, _, Some(Ok(flight_data))) => { + // we have a flight data, so we yield it + trace!("received normal flight data, yielding"); + yield Ok(flight_data); + }, + (true, false, _ ) => { + // we've finished a partition, do nothing + done = true; + }, + (true, true, _ ) => { + yield Err(Status::internal(format!("{name} reached expected unreachable block"))); + }, + } + } + }; + + Ok(Box::pin(out_stream)) + } + fn do_action_get_host(&self) -> Result, Status> { + let addr = self.addr.clone(); + let name = self.name.clone(); + + let out_stream = Box::pin(stream! { + yield Ok::<_, tonic::Status>(arrow_flight::Result { + body: Host{ + addr, + name, + }.encode_to_vec().into() + }); + }) as crate::flight::DoActionStream; + + Ok(Response::new(out_stream)) + } + + async fn do_action_add_plan( + &self, + action: Action, + ) -> Result, Status> { + let task_data = DDTask::decode(action.body.as_ref()).map_err(|e| { + Status::internal(format!( + "{}, Unexpected error decoding StageData: {e:?}", + self.name + )) + })?; + + let addrs = task_data + .stage_addrs + .as_ref() + .context("stage addrs not present") + .map_err(DDError::from) + .and_then(get_addrs) + .map_err(|e| Status::internal(format!("{}, {e}", self.name)))?; + + self.add_task( + task_data.query_id, + task_data.stage_id, + addrs, + task_data.partition_group, + task_data.full_partitions, + &task_data.plan_bytes, + ) + .await + .map_err(|e| Status::internal(format!("{}, Could not add plan: {e:?}", self.name)))?; + + let out_stream = Box::pin(stream! { + yield Ok::<_, tonic::Status>(arrow_flight::Result::default()); + }) as crate::flight::DoActionStream; + + Ok(Response::new(out_stream)) + } +} + +#[async_trait] +impl FlightHandler for DDWorkerHandler { + async fn do_get( + &self, + request: Request, + ) -> std::result::Result, Status> { + let remote_addr = request + .remote_addr() + .map(|a| a.to_string()) + .unwrap_or("unknown".to_string()); + + let ticket = request.into_inner(); + + let ftd = FlightTicketData::decode(ticket.ticket).map_err(|e| { + Status::internal(format!( + "{}, Unexpected error extracting ticket {e:?}", + self.name + )) + })?; + + let task_key = StageKey { + query_id: ftd.query_id.clone(), + stage_id: ftd.stage_id, + }; + + debug!( + "{}, request for task_key:{:?} partition: {} from: {},{}", + self.name, task_key, ftd.partition, ftd.requestor_name, remote_addr + ); + + let name = self.name.clone(); + let (ctx, plan, is_last_partition) = self + .get_ctx_and_plan(&ftd.query_id, ftd.stage_id, ftd.partition) + .map_err(|e| { + Status::internal(format!( + "{name} Could not find task for query_id {} stage {} partition {}: {e:?}", + ftd.query_id, ftd.stage_id, ftd.partition + )) + })?; + + let do_get_stream = self + .make_stream(ctx, plan, ftd.stage_id, ftd.partition, is_last_partition) + .map_err(|e| { + Status::internal(format!( + "{name} Could not make stream for query_id {} stage {} partition {}: {e:?}", + ftd.query_id, ftd.stage_id, ftd.partition + )) + })?; + + Ok(Response::new(do_get_stream)) + } + + async fn do_action( + &self, + request: Request, + ) -> Result, Status> { + // extract a StageData protobuf from the action body + let action = request.into_inner(); + let action_type = action.r#type.as_str(); + trace!("{} received action: {}", self.name, action_type); + + if action_type == "add_plan" { + self.do_action_add_plan(action).await + } else if action_type == "get_host" { + self.do_action_get_host() + } else { + Err(Status::unimplemented(format!( + "{}, Unimplemented action: {}", + self.name, action_type + ))) + } + } +} + +/// DDWorkerService is a Arrow Flight service that serves streams of +/// partitions from a hosted Physical Plan +/// +/// It only responds to the DoGet Arrow Flight method +pub struct DDWorkerService { + #[allow(dead_code)] + name: String, + listener: TcpListener, + handler: Arc, + addr: String, + all_done_tx: Arc>>, + all_done_rx: Option>, +} + +impl DDWorkerService { + pub async fn new(name: String, port: usize) -> Result { + let name = format!("[{}]", name); + + let (all_done_tx, all_done_rx) = channel(1); + let all_done_tx = Arc::new(Mutex::new(all_done_tx)); + + let listener = start_up(port).await?; + + let addr = format!("{}", listener.local_addr().unwrap()); + + info!("DDWorkerService bound to {addr}"); + + let handler = Arc::new(DDWorkerHandler::new(name.clone(), addr.clone())); + + Ok(Self { + name, + listener, + handler, + addr, + all_done_tx, + all_done_rx: Some(all_done_rx), + }) + } + + /// get the address of the listing socket for this service + pub fn addr(&self) -> String { + self.addr.clone() + } + + pub async fn all_done(&self) -> Result<()> { + let sender = self.all_done_tx.lock().clone(); + + sender + .send(()) + .await + .context("Could not send shutdown signal")?; + Ok(()) + } + + /// start the service, consuming self + pub async fn serve(mut self) -> Result<()> { + let mut all_done_rx = self.all_done_rx.take().unwrap(); + + let signal = async move { + all_done_rx + .recv() + .await + .expect("problem receiving shutdown signal"); + info!("received shutdown signal"); + }; + + let flight_serv = FlightServ { + handler: self.handler.clone(), + }; + + let svc = FlightServiceServer::new(flight_serv); + + Server::builder() + .add_service(svc) + .serve_with_incoming_shutdown( + tokio_stream::wrappers::TcpListenerStream::new(self.listener), + signal, + ) + .await + .context("error running service")?; + Ok(()) + } +}