fix lint

Author: phillipleblanc

ballista/core/src/client.rs

@@ -47,8 +47,6 @@ use datafusion::error::Result;
 use crate::extension::BallistaConfigGrpcEndpoint;
 use crate::serde::protobuf;
 
-use crate::utils::{GrpcClientConfig, create_grpc_client_connection};
-
 use crate::utils::create_grpc_client_endpoint;
 
 use datafusion::physical_plan::{RecordBatchStream, SendableRecordBatchStream};

ballista/core/src/execution_plans/distributed_query.rs

@@ -28,8 +28,6 @@ use crate::serde::protobuf::{
     scheduler_grpc_client::SchedulerGrpcClient,
 };
 
-use crate::utils::{GrpcClientConfig, create_grpc_client_connection};
-
 use crate::utils::create_grpc_client_endpoint;
 
 use datafusion::arrow::datatypes::SchemaRef;
@@ -266,8 +264,6 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
                 self.scheduler_url.clone(),
                 self.session_id.clone(),
                 query,
-                self.config.default_grpc_client_max_message_size(),
-                GrpcClientConfig::from(&self.config),
                 Arc::new(self.metrics.clone()),
                 partition,
                 self.config.clone(),
@@ -306,16 +302,13 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
     }
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn execute_query(
     scheduler_url: String,
     session_id: String,
     query: ExecuteQueryParams,
-
-    max_message_size: usize,
-    grpc_config: GrpcClientConfig,
     metrics: Arc<ExecutionPlanMetricsSet>,
     partition: usize,
-
     config: BallistaConfig,
     grpc_interceptor: Arc<BallistaGrpcMetadataInterceptor>,
     customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -392,6 +392,7 @@ fn local_remote_read_split(
     }
 }
 
+#[allow(clippy::too_many_arguments)]
 fn send_fetch_partitions(
     partition_locations: Vec<PartitionLocation>,
     max_request_num: usize,
@@ -433,21 +434,21 @@ fn send_fetch_partitions(
                 .await;
 
             // Record local read metrics if callback is set and read succeeded
-            if r.is_ok() {
-                if let Some(ref callback) = metrics_callback_c {
-                    let duration_ms = start_time.elapsed().as_millis() as u64;
-                    let bytes = p.partition_stats.num_bytes().unwrap_or(0);
-                    let rows = p.partition_stats.num_rows().unwrap_or(0);
-                    callback.record_local_read(
-                        &p.partition_id.job_id,
-                        p.partition_id.stage_id,
-                        p.partition_id.partition_id,
-                        &p.executor_meta.id,
-                        bytes,
-                        rows,
-                        duration_ms,
-                    );
-                }
+            if r.is_ok()
+                && let Some(ref callback) = metrics_callback_c
+            {
+                let duration_ms = start_time.elapsed().as_millis() as u64;
+                let bytes = p.partition_stats.num_bytes().unwrap_or(0);
+                let rows = p.partition_stats.num_rows().unwrap_or(0);
+                callback.record_local_read(
+                    &p.partition_id.job_id,
+                    p.partition_id.stage_id,
+                    p.partition_id.partition_id,
+                    &p.executor_meta.id,
+                    bytes,
+                    rows,
+                    duration_ms,
+                );
             }
 
             if let Err(e) = response_sender_c.send(r).await {
@@ -476,21 +477,21 @@ fn send_fetch_partitions(
                 .await;
 
             // Record remote read metrics if callback is set and read succeeded
-            if r.is_ok() {
-                if let Some(ref callback) = metrics_callback_c {
-                    let duration_ms = start_time.elapsed().as_millis() as u64;
-                    let bytes = p.partition_stats.num_bytes().unwrap_or(0);
-                    let rows = p.partition_stats.num_rows().unwrap_or(0);
-                    callback.record_remote_read(
-                        &p.partition_id.job_id,
-                        p.partition_id.stage_id,
-                        p.partition_id.partition_id,
-                        &p.executor_meta.id,
-                        bytes,
-                        rows,
-                        duration_ms,
-                    );
-                }
+            if r.is_ok()
+                && let Some(ref callback) = metrics_callback_c
+            {
+                let duration_ms = start_time.elapsed().as_millis() as u64;
+                let bytes = p.partition_stats.num_bytes().unwrap_or(0);
+                let rows = p.partition_stats.num_rows().unwrap_or(0);
+                callback.record_remote_read(
+                    &p.partition_id.job_id,
+                    p.partition_id.stage_id,
+                    p.partition_id.partition_id,
+                    &p.executor_meta.id,
+                    bytes,
+                    rows,
+                    duration_ms,
+                );
             }
 
             // Block if the channel buffer is full.

ballista/core/src/extension.rs

@@ -161,11 +161,13 @@ pub trait SessionConfigExt {
     /// Get a `tonic` interceptor configured to decorate the provided metadata keys
     fn ballista_grpc_interceptor(&self) -> Arc<BallistaGrpcMetadataInterceptor>;
 
+    /// Set an override function for creating gRPC client endpoints.
     fn with_ballista_override_create_grpc_client_endpoint(
         self,
         override_f: EndpointOverrideFn,
     ) -> Self;
 
+    /// Get the override function for creating gRPC client endpoints.
     fn ballista_override_create_grpc_client_endpoint(
         &self,
     ) -> Option<Arc<BallistaConfigGrpcEndpoint>>;
@@ -664,6 +666,7 @@ pub struct BallistaGrpcMetadataInterceptor {
 }
 
 impl BallistaGrpcMetadataInterceptor {
+    /// Create a new interceptor with additional metadata to be added to requests.
     pub fn new(additional_metadata: HashMap<String, String>) -> Self {
         Self {
             additional_metadata,
@@ -694,16 +697,19 @@ impl Interceptor for BallistaGrpcMetadataInterceptor {
     }
 }
 
+/// Wrapper for gRPC endpoint configuration override function.
 #[derive(Clone)]
 pub struct BallistaConfigGrpcEndpoint {
     override_f: EndpointOverrideFn,
 }
 
 impl BallistaConfigGrpcEndpoint {
+    /// Create a new endpoint configuration with the given override function.
     pub fn new(override_f: EndpointOverrideFn) -> Self {
         Self { override_f }
     }
 
+    /// Configure an endpoint using the override function.
     pub fn configure_endpoint(
         &self,
         endpoint: Endpoint,
@@ -735,6 +741,7 @@ pub trait ShuffleReadMetricsCallback: Send + Sync {
     /// * `bytes` - Number of bytes read
     /// * `rows` - Number of rows read
     /// * `duration_ms` - Time taken to read the partition
+    #[allow(clippy::too_many_arguments)]
     fn record_local_read(
         &self,
         job_id: &str,
@@ -759,6 +766,7 @@ pub trait ShuffleReadMetricsCallback: Send + Sync {
     /// * `bytes` - Number of bytes read
     /// * `rows` - Number of rows read
     /// * `duration_ms` - Time taken to fetch the partition
+    #[allow(clippy::too_many_arguments)]
     fn record_remote_read(
         &self,
         job_id: &str,
@@ -810,6 +818,7 @@ pub trait ResultFetchMetricsCallback: Send + Sync {
     /// * `bytes` - Number of bytes fetched
     /// * `rows` - Number of rows fetched
     /// * `duration_ms` - Time taken to fetch the partition
+    #[allow(clippy::too_many_arguments)]
     fn record_result_fetch(
         &self,
         job_id: &str,

ballista/core/src/lib.rs

@@ -55,10 +55,10 @@ pub mod registry;
 /// Serialization and deserialization for Ballista messages and plans.
 pub mod serde;
 
-/// General utility functions for Ballista operations.
-
+/// Remote catalog serialization and stub providers for Ballista clients.
 pub mod remote_catalog;
 
+/// General utility functions for Ballista operations.
 pub mod utils;
 
 ///

ballista/core/src/remote_catalog/catalog_serialize_ext.rs

@@ -18,17 +18,20 @@
 use crate::serde::generated::ballista::{CatalogInfo, SchemaInfo, TableInfo};
 use datafusion::catalog::CatalogProvider;
 use datafusion::prelude::SessionContext;
-use futures::stream;
 use futures::StreamExt;
+use futures::stream;
 use std::sync::Arc;
 
-/// Used to serialize catalog schemas and names to ship to Ballista clients
+/// Extension trait for serializing catalog schemas and names to ship to Ballista clients.
 #[async_trait::async_trait]
 pub trait CatalogSerializeExt {
+    /// Serialize all catalogs in the session context.
     async fn serialize_catalogs(&self) -> Vec<CatalogInfo>;
 
+    /// Serialize a specific catalog by name.
     async fn serialize_catalog(&self, name: &str) -> Option<CatalogInfo>;
 
+    /// Serialize a specific schema within a catalog.
     async fn serialize_schema(
         &self,
         name: &str,

ballista/core/src/remote_catalog/mod.rs

@@ -16,7 +16,17 @@
 // under the License.
 //
 
+//! Remote catalog serialization and stub providers for Ballista clients.
+//!
+//! This module provides functionality to serialize catalog metadata (schemas, tables, functions)
+//! from the scheduler to ship to Ballista clients, as well as stub providers that allow clients
+//! to perform logical planning without access to actual table data.
+
+/// Extension trait for serializing catalog schemas and table names.
 pub mod catalog_serialize_ext;
+/// Extension trait for serializing user-defined functions.
 pub mod remote_function_serialize_ext;
+/// Stub scalar UDF implementation for remote function planning.
 pub mod remote_scalar_udf;
+/// Stub table provider for remote table planning.
 pub mod remote_table_provider;

ballista/core/src/remote_catalog/remote_function_serialize_ext.rs

@@ -25,8 +25,9 @@ use datafusion::prelude::SessionContext;
 use datafusion_proto_common::ArrowType;
 use std::collections::HashSet;
 
-/// Used to serialize function shapes to ship to Ballista clients
+/// Extension trait for serializing function signatures to ship to Ballista clients.
 pub trait RemoteFunctionSerializeExt {
+    /// Serialize all user-defined scalar functions in the session context.
     fn serialize_udfs(&self) -> Vec<ScalarUdfInfo>;
 }
 

ballista/core/src/remote_catalog/remote_scalar_udf.rs

@@ -19,7 +19,7 @@
 use crate::serde::protobuf::ScalarUdfInfo;
 use datafusion::arrow::datatypes::DataType;
 use datafusion::common::Result;
-use datafusion::common::{exec_err, plan_err, DataFusionError};
+use datafusion::common::{DataFusionError, exec_err, plan_err};
 use datafusion::logical_expr::{
     ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
     TypeSignature, Volatility,
@@ -56,6 +56,7 @@ impl PartialEq for RemoteScalarUDF {
 impl Eq for RemoteScalarUDF {}
 
 impl RemoteScalarUDF {
+    /// Create a new RemoteScalarUDF from a ScalarUdfInfo protobuf message.
     pub fn new(meta: ScalarUdfInfo) -> Result<Self> {
         let mut arities = vec![];
 

ballista/core/src/remote_catalog/remote_table_provider.rs

@@ -18,7 +18,7 @@
 
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::catalog::{Session, TableProvider};
-use datafusion::common::{exec_err, Result};
+use datafusion::common::{Result, exec_err};
 use datafusion::datasource::TableType;
 use datafusion::logical_expr::Expr;
 use datafusion::physical_plan::ExecutionPlan;
@@ -36,6 +36,7 @@ pub struct RemoteTableProvider {
 }
 
 impl RemoteTableProvider {
+    /// Create a new RemoteTableProvider with the given catalog, schema, and table names.
     pub fn new(
         catalog_name: &str,
         schema_name: &str,
@@ -50,14 +51,17 @@ impl RemoteTableProvider {
         }
     }
 
+    /// Get the catalog name.
     pub fn catalog_name(&self) -> &str {
         &self.catalog_name
     }
 
+    /// Get the schema name.
     pub fn schema_name(&self) -> &str {
         &self.schema_name
     }
 
+    /// Get the table name.
     pub fn table_name(&self) -> &str {
         &self.table_name
     }
@@ -84,7 +88,11 @@ impl TableProvider for RemoteTableProvider {
         _filters: &[Expr],
         _limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        exec_err!("{}.{}.{} is a stub table implementation to be resolved on the Ballista scheduler. It should not be scanned on the client. This is a bug.",
-        self.catalog_name, self.schema_name, self.table_name)
+        exec_err!(
+            "{}.{}.{} is a stub table implementation to be resolved on the Ballista scheduler. It should not be scanned on the client. This is a bug.",
+            self.catalog_name,
+            self.schema_name,
+            self.table_name
+        )
     }
 }