From 7e1d9a412100c954a2f3c1148d7e69423927f8de Mon Sep 17 00:00:00 2001 From: Jonathan Date: Wed, 24 Sep 2025 16:38:26 -0400 Subject: [PATCH 1/2] fix: Move `error.rs` to protobuf --- src/execution_plans/network_coalesce.rs | 2 +- src/execution_plans/network_shuffle.rs | 2 +- src/flight_service/do_get.rs | 2 +- src/lib.rs | 1 - src/{ => protobuf}/errors/arrow_error.rs | 3 ++- src/{ => protobuf}/errors/datafusion_error.rs | 12 ++++++------ src/{ => protobuf}/errors/io_error.rs | 0 src/{ => protobuf}/errors/mod.rs | 3 ++- src/{ => protobuf}/errors/objectstore_error.rs | 0 src/{ => protobuf}/errors/parquet_error.rs | 0 src/{ => protobuf}/errors/parser_error.rs | 0 src/{ => protobuf}/errors/schema_error.rs | 0 src/protobuf/mod.rs | 1 + 13 files changed, 14 insertions(+), 12 deletions(-) rename src/{ => protobuf}/errors/arrow_error.rs (99%) rename src/{ => protobuf}/errors/datafusion_error.rs (97%) rename src/{ => protobuf}/errors/io_error.rs (100%) rename src/{ => protobuf}/errors/mod.rs (97%) rename src/{ => protobuf}/errors/objectstore_error.rs (100%) rename src/{ => protobuf}/errors/parquet_error.rs (100%) rename src/{ => protobuf}/errors/parser_error.rs (100%) rename src/{ => protobuf}/errors/schema_error.rs (100%) diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs index d9175ba..52812d1 100644 --- a/src/execution_plans/network_coalesce.rs +++ b/src/execution_plans/network_coalesce.rs @@ -3,10 +3,10 @@ use crate::channel_resolver_ext::get_distributed_channel_resolver; use crate::common::scale_partitioning_props; use crate::config_extension_ext::ContextGrpcMetadata; use crate::distributed_physical_optimizer_rule::{NetworkBoundary, limit_tasks_err}; -use crate::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; use crate::execution_plans::{DistributedTaskContext, StageExec}; use crate::flight_service::DoGet; use crate::metrics::proto::MetricsSetProto; +use crate::protobuf::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; use crate::protobuf::{DistributedCodec, StageKey, proto_from_stage}; use arrow_flight::Ticket; use arrow_flight::decode::FlightRecordBatchStream; diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index 2958a50..44108cc 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -3,10 +3,10 @@ use crate::channel_resolver_ext::get_distributed_channel_resolver; use crate::common::scale_partitioning; use crate::config_extension_ext::ContextGrpcMetadata; use crate::distributed_physical_optimizer_rule::NetworkBoundary; -use crate::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; use crate::execution_plans::{DistributedTaskContext, StageExec}; use crate::flight_service::DoGet; use crate::metrics::proto::MetricsSetProto; +use crate::protobuf::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; use crate::protobuf::{DistributedCodec, StageKey, proto_from_stage}; use arrow_flight::Ticket; use arrow_flight::decode::FlightRecordBatchStream; diff --git a/src/flight_service/do_get.rs b/src/flight_service/do_get.rs index bd0345f..b387a0f 100644 --- a/src/flight_service/do_get.rs +++ b/src/flight_service/do_get.rs @@ -1,8 +1,8 @@ use crate::config_extension_ext::ContextGrpcMetadata; -use crate::errors::datafusion_error_to_tonic_status; use crate::execution_plans::{DistributedTaskContext, StageExec}; use crate::flight_service::service::ArrowFlightEndpoint; use crate::flight_service::session_builder::DistributedSessionBuilderContext; +use crate::protobuf::errors::datafusion_error_to_tonic_status; use crate::protobuf::{DistributedCodec, StageExecProto, StageKey, stage_from_proto}; use arrow_flight::Ticket; use arrow_flight::encode::FlightDataEncoderBuilder; diff --git a/src/lib.rs b/src/lib.rs index 152edfb..3d8bcb1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,7 +5,6 @@ mod common; mod config_extension_ext; mod distributed_ext; mod distributed_physical_optimizer_rule; -mod errors; mod execution_plans; mod flight_service; mod metrics; diff --git a/src/errors/arrow_error.rs b/src/protobuf/errors/arrow_error.rs similarity index 99% rename from src/errors/arrow_error.rs rename to src/protobuf/errors/arrow_error.rs index 33565ee..c3658c5 100644 --- a/src/errors/arrow_error.rs +++ b/src/protobuf/errors/arrow_error.rs @@ -1,6 +1,7 @@ -use crate::errors::io_error::IoErrorProto; use datafusion::arrow::error::ArrowError; +use crate::protobuf::errors::io_error::IoErrorProto; + #[derive(Clone, PartialEq, ::prost::Message)] pub struct ArrowErrorProto { #[prost(string, optional, tag = "1")] diff --git a/src/errors/datafusion_error.rs b/src/protobuf/errors/datafusion_error.rs similarity index 97% rename from src/errors/datafusion_error.rs rename to src/protobuf/errors/datafusion_error.rs index bb91155..d1312b7 100644 --- a/src/errors/datafusion_error.rs +++ b/src/protobuf/errors/datafusion_error.rs @@ -1,9 +1,9 @@ -use crate::errors::arrow_error::ArrowErrorProto; -use crate::errors::io_error::IoErrorProto; -use crate::errors::objectstore_error::ObjectStoreErrorProto; -use crate::errors::parquet_error::ParquetErrorProto; -use crate::errors::parser_error::ParserErrorProto; -use crate::errors::schema_error::SchemaErrorProto; +use crate::protobuf::errors::arrow_error::ArrowErrorProto; +use crate::protobuf::errors::io_error::IoErrorProto; +use crate::protobuf::errors::objectstore_error::ObjectStoreErrorProto; +use crate::protobuf::errors::parquet_error::ParquetErrorProto; +use crate::protobuf::errors::parser_error::ParserErrorProto; +use crate::protobuf::errors::schema_error::SchemaErrorProto; use datafusion::common::{DataFusionError, Diagnostic}; use datafusion::logical_expr::sqlparser::parser::ParserError; use std::error::Error; diff --git a/src/errors/io_error.rs b/src/protobuf/errors/io_error.rs similarity index 100% rename from src/errors/io_error.rs rename to src/protobuf/errors/io_error.rs diff --git a/src/errors/mod.rs b/src/protobuf/errors/mod.rs similarity index 97% rename from src/errors/mod.rs rename to src/protobuf/errors/mod.rs index e2a76d6..3466fe9 100644 --- a/src/errors/mod.rs +++ b/src/protobuf/errors/mod.rs @@ -1,11 +1,12 @@ #![allow(clippy::upper_case_acronyms, clippy::vec_box)] -use crate::errors::datafusion_error::DataFusionErrorProto; use arrow_flight::error::FlightError; use datafusion::common::internal_datafusion_err; use datafusion::error::DataFusionError; use prost::Message; +use crate::protobuf::errors::datafusion_error::DataFusionErrorProto; + mod arrow_error; mod datafusion_error; mod io_error; diff --git a/src/errors/objectstore_error.rs b/src/protobuf/errors/objectstore_error.rs similarity index 100% rename from src/errors/objectstore_error.rs rename to src/protobuf/errors/objectstore_error.rs diff --git a/src/errors/parquet_error.rs b/src/protobuf/errors/parquet_error.rs similarity index 100% rename from src/errors/parquet_error.rs rename to src/protobuf/errors/parquet_error.rs diff --git a/src/errors/parser_error.rs b/src/protobuf/errors/parser_error.rs similarity index 100% rename from src/errors/parser_error.rs rename to src/protobuf/errors/parser_error.rs diff --git a/src/errors/schema_error.rs b/src/protobuf/errors/schema_error.rs similarity index 100% rename from src/errors/schema_error.rs rename to src/protobuf/errors/schema_error.rs diff --git a/src/protobuf/mod.rs b/src/protobuf/mod.rs index 5c0870e..29b19e2 100644 --- a/src/protobuf/mod.rs +++ b/src/protobuf/mod.rs @@ -1,4 +1,5 @@ mod distributed_codec; +pub mod errors; mod stage_proto; mod user_codec; From 7f189f40a0f43326859495df63c2df9e5b05888a Mon Sep 17 00:00:00 2001 From: Jonathan Date: Thu, 25 Sep 2025 03:40:26 -0400 Subject: [PATCH 2/2] make public imports --- src/execution_plans/network_coalesce.rs | 2 +- src/execution_plans/network_shuffle.rs | 2 +- src/flight_service/do_get.rs | 5 +++-- src/protobuf/mod.rs | 6 +++++- 4 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs index 52812d1..be9a5fe 100644 --- a/src/execution_plans/network_coalesce.rs +++ b/src/execution_plans/network_coalesce.rs @@ -6,8 +6,8 @@ use crate::distributed_physical_optimizer_rule::{NetworkBoundary, limit_tasks_er use crate::execution_plans::{DistributedTaskContext, StageExec}; use crate::flight_service::DoGet; use crate::metrics::proto::MetricsSetProto; -use crate::protobuf::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; use crate::protobuf::{DistributedCodec, StageKey, proto_from_stage}; +use crate::protobuf::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; use arrow_flight::Ticket; use arrow_flight::decode::FlightRecordBatchStream; use arrow_flight::error::FlightError; diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs index 44108cc..eed04b7 100644 --- a/src/execution_plans/network_shuffle.rs +++ b/src/execution_plans/network_shuffle.rs @@ -6,8 +6,8 @@ use crate::distributed_physical_optimizer_rule::NetworkBoundary; use crate::execution_plans::{DistributedTaskContext, StageExec}; use crate::flight_service::DoGet; use crate::metrics::proto::MetricsSetProto; -use crate::protobuf::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; use crate::protobuf::{DistributedCodec, StageKey, proto_from_stage}; +use crate::protobuf::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; use arrow_flight::Ticket; use arrow_flight::decode::FlightRecordBatchStream; use arrow_flight::error::FlightError; diff --git a/src/flight_service/do_get.rs b/src/flight_service/do_get.rs index b387a0f..b9701ff 100644 --- a/src/flight_service/do_get.rs +++ b/src/flight_service/do_get.rs @@ -2,8 +2,9 @@ use crate::config_extension_ext::ContextGrpcMetadata; use crate::execution_plans::{DistributedTaskContext, StageExec}; use crate::flight_service::service::ArrowFlightEndpoint; use crate::flight_service::session_builder::DistributedSessionBuilderContext; -use crate::protobuf::errors::datafusion_error_to_tonic_status; -use crate::protobuf::{DistributedCodec, StageExecProto, StageKey, stage_from_proto}; +use crate::protobuf::{ + DistributedCodec, StageExecProto, StageKey, datafusion_error_to_tonic_status, stage_from_proto, +}; use arrow_flight::Ticket; use arrow_flight::encode::FlightDataEncoderBuilder; use arrow_flight::error::FlightError; diff --git a/src/protobuf/mod.rs b/src/protobuf/mod.rs index 29b19e2..342eb03 100644 --- a/src/protobuf/mod.rs +++ b/src/protobuf/mod.rs @@ -1,8 +1,12 @@ mod distributed_codec; -pub mod errors; +mod errors; mod stage_proto; mod user_codec; pub(crate) use distributed_codec::DistributedCodec; +pub(crate) use errors::{ + datafusion_error_to_tonic_status, map_flight_to_datafusion_error, + map_status_to_datafusion_error, +}; pub(crate) use stage_proto::{StageExecProto, StageKey, proto_from_stage, stage_from_proto}; pub(crate) use user_codec::{get_distributed_user_codec, set_distributed_user_codec};