diff --git a/k8s/cluster.yaml b/k8s/cluster.yaml index 98813323..91afa739 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..df7b9d30 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..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; @@ -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..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::DFRayProcessorService, - proxy_service::DFRayProxyService, setup, + friendly::new_friendly_name, proxy_service::DDProxyService, setup, + worker_service::DDWorkerService, }; #[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 = DDWorkerService::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/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..95a30860 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,18 @@ impl FlightSqlHandler for DfRayProxyHandler { } } -/// DFRayProcessorService 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 { +/// An Arrow Flight SQL service +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 +288,9 @@ impl DFRayProxyService { let addr = format!("{}", listener.local_addr().unwrap()); - info!("DFRayProcessorService bound to {addr}"); + info!("DDProxyService 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..20b04b13 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)) @@ -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,71 +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 { - // 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 - + 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()); @@ -451,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(), @@ -463,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) } 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/src/processor_service.rs b/src/worker_service.rs similarity index 97% rename from src/processor_service.rs rename to src/worker_service.rs index 555ff912..8500984f 100644 --- a/src/processor_service.rs +++ b/src/worker_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 DDWorkerHandler { /// 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 DDWorkerHandler { 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 DDWorkerHandler { 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 +/// 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 DFRayProcessorService { +pub struct DDWorkerService { #[allow(dead_code)] name: String, listener: TcpListener, - handler: Arc, + handler: Arc, addr: String, all_done_tx: Arc>>, all_done_rx: Option>, } -impl DFRayProcessorService { +impl DDWorkerService { 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!("DDWorkerService bound to {addr}"); - let handler = Arc::new(DfRayProcessorHandler::new(name.clone(), addr.clone())); + let handler = Arc::new(DDWorkerHandler::new(name.clone(), addr.clone())); Ok(Self { name, 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,