This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 9b1bb68e37 refactor: improve PoolType argument handling for CLI
(#10940)
9b1bb68e37 is described below
commit 9b1bb68e37688330fa47c73bdd733b206be8759e
Author: Trent Hauck <[email protected]>
AuthorDate: Mon Jun 17 08:52:26 2024 -0700
refactor: improve PoolType argument handling for CLI (#10940)
* refactor: dont include fallback in match on mem_pool_type
* refactor: improve PoolType argument handling
---
datafusion-cli/src/lib.rs | 1 +
datafusion-cli/src/main.rs | 29 +++++---------------
datafusion-cli/src/{lib.rs => pool_type.rs} | 41 +++++++++++++++++++++--------
3 files changed, 37 insertions(+), 34 deletions(-)
diff --git a/datafusion-cli/src/lib.rs b/datafusion-cli/src/lib.rs
index 139a60b8cf..5081436aa6 100644
--- a/datafusion-cli/src/lib.rs
+++ b/datafusion-cli/src/lib.rs
@@ -25,5 +25,6 @@ pub mod functions;
pub mod helper;
pub mod highlighter;
pub mod object_storage;
+pub mod pool_type;
pub mod print_format;
pub mod print_options;
diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs
index f2b29fe786..f469fda4f9 100644
--- a/datafusion-cli/src/main.rs
+++ b/datafusion-cli/src/main.rs
@@ -19,7 +19,6 @@ use std::collections::HashMap;
use std::env;
use std::path::Path;
use std::process::ExitCode;
-use std::str::FromStr;
use std::sync::{Arc, OnceLock};
use datafusion::error::{DataFusionError, Result};
@@ -31,6 +30,7 @@ use datafusion_cli::catalog::DynamicFileCatalog;
use datafusion_cli::functions::ParquetMetadataFunc;
use datafusion_cli::{
exec,
+ pool_type::PoolType,
print_format::PrintFormat,
print_options::{MaxRows, PrintOptions},
DATAFUSION_CLI_VERSION,
@@ -42,24 +42,6 @@ use mimalloc::MiMalloc;
#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;
-#[derive(PartialEq, Debug)]
-enum PoolType {
- Greedy,
- Fair,
-}
-
-impl FromStr for PoolType {
- type Err = String;
-
- fn from_str(s: &str) -> Result<Self, Self::Err> {
- match s {
- "Greedy" | "greedy" => Ok(PoolType::Greedy),
- "Fair" | "fair" => Ok(PoolType::Fair),
- _ => Err(format!("Invalid memory pool type '{}'", s)),
- }
- }
-}
-
#[derive(Debug, Parser, PartialEq)]
#[clap(author, version, about, long_about= None)]
struct Args {
@@ -127,9 +109,10 @@ struct Args {
#[clap(
long,
- help = "Specify the memory pool type 'greedy' or 'fair', default to
'greedy'"
+ help = "Specify the memory pool type 'greedy' or 'fair'",
+ default_value_t = PoolType::Greedy
)]
- mem_pool_type: Option<PoolType>,
+ mem_pool_type: PoolType,
#[clap(
long,
@@ -181,9 +164,9 @@ async fn main_inner() -> Result<()> {
let memory_limit =
extract_memory_pool_size(&memory_limit).unwrap();
// set memory pool type
match args.mem_pool_type {
- Some(PoolType::Fair) => rt_config
+ PoolType::Fair => rt_config
.with_memory_pool(Arc::new(FairSpillPool::new(memory_limit))),
- _ => rt_config
+ PoolType::Greedy => rt_config
.with_memory_pool(Arc::new(GreedyMemoryPool::new(memory_limit)))
}
} else {
diff --git a/datafusion-cli/src/lib.rs b/datafusion-cli/src/pool_type.rs
similarity index 54%
copy from datafusion-cli/src/lib.rs
copy to datafusion-cli/src/pool_type.rs
index 139a60b8cf..25763eba5c 100644
--- a/datafusion-cli/src/lib.rs
+++ b/datafusion-cli/src/pool_type.rs
@@ -15,15 +15,34 @@
// specific language governing permissions and limitations
// under the License.
-#![doc = include_str!("../README.md")]
-pub const DATAFUSION_CLI_VERSION: &str = env!("CARGO_PKG_VERSION");
+use std::{
+ fmt::{self, Display, Formatter},
+ str::FromStr,
+};
-pub mod catalog;
-pub mod command;
-pub mod exec;
-pub mod functions;
-pub mod helper;
-pub mod highlighter;
-pub mod object_storage;
-pub mod print_format;
-pub mod print_options;
+#[derive(PartialEq, Debug)]
+pub enum PoolType {
+ Greedy,
+ Fair,
+}
+
+impl FromStr for PoolType {
+ type Err = String;
+
+ fn from_str(s: &str) -> Result<Self, Self::Err> {
+ match s {
+ "Greedy" | "greedy" => Ok(PoolType::Greedy),
+ "Fair" | "fair" => Ok(PoolType::Fair),
+ _ => Err(format!("Invalid memory pool type '{}'", s)),
+ }
+ }
+}
+
+impl Display for PoolType {
+ fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+ match self {
+ PoolType::Greedy => write!(f, "greedy"),
+ PoolType::Fair => write!(f, "fair"),
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
