From e643b326ec2dde2d03956ef96774d87c79ac3cfd Mon Sep 17 00:00:00 2001 From: matt rice Date: Wed, 8 Jul 2026 03:25:39 -0700 Subject: [PATCH 1/4] Move most of main to an entry_point module in datafusion-cli lib. Unformatted to ease review. This also adds an example of reusing the main datafusion-cli while registering custom udf. --- datafusion-cli/examples/cli-custom-udf.rs | 69 ++ datafusion-cli/src/entry_point.rs | 868 ++++++++++++++++++++++ datafusion-cli/src/lib.rs | 1 + datafusion-cli/src/main.rs | 824 +------------------- 4 files changed, 940 insertions(+), 822 deletions(-) create mode 100644 datafusion-cli/examples/cli-custom-udf.rs create mode 100644 datafusion-cli/src/entry_point.rs diff --git a/datafusion-cli/examples/cli-custom-udf.rs b/datafusion-cli/examples/cli-custom-udf.rs new file mode 100644 index 0000000000000..f98ac2d972c6d --- /dev/null +++ b/datafusion-cli/examples/cli-custom-udf.rs @@ -0,0 +1,69 @@ +// 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 arrow::datatypes::DataType; +use datafusion::arrow::array::{ArrayRef, StringArray}; +use datafusion::logical_expr::{ColumnarValue, Volatility, create_udf}; +use datafusion::prelude::SessionContext; +use datafusion_common::cast::as_string_array; +use mimalloc::MiMalloc; +use std::process::ExitCode; +use std::sync::Arc; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +/// In this example we want to reuse the datafusion-cli binary argument, hen extend the `SessionContext` with custom udf. +/// +/// 1. Declares a `hello`` udf function. +/// 2. Construct a `CliSession` +/// 3. Registers the udf function with the `SessionContext` so the user can input `select hello(1)` at the prompt. +/// 4. Runs the cli using [`dataframe_cli::CliSession::run`], printing any errors then exits. +#[tokio::main] +pub async fn main() -> ExitCode { + let custom_udf = create_udf( + "hello", + vec![DataType::Utf8], + DataType::Utf8, + Volatility::Immutable, + Arc::new(|args: &[ColumnarValue]| { + assert_eq!(args.len(), 1); + let args = ColumnarValue::values_to_arrays(args).unwrap(); + let vals = as_string_array(&args[0]).expect("cast failed"); + let array = vals + .iter() + .map(|v| v.map(|v| format!("hello {v}"))) + .collect::(); + Ok(ColumnarValue::from(Arc::new(array) as ArrayRef)) + }), + ); + let cli_session = match datafusion_cli::entry_point::CliSession::try_from_args() { + Ok(cli) => cli, + Err(e) => { + println!("Error: {e}"); + return ExitCode::FAILURE; + } + }; + let ctx: &SessionContext = cli_session.session_context(); + ctx.register_udf(custom_udf); + if let Err(e) = cli_session.run().await { + println!("Error: {e}"); + return ExitCode::FAILURE; + } + + ExitCode::SUCCESS +} diff --git a/datafusion-cli/src/entry_point.rs b/datafusion-cli/src/entry_point.rs new file mode 100644 index 0000000000000..282f6959d7a45 --- /dev/null +++ b/datafusion-cli/src/entry_point.rs @@ -0,0 +1,868 @@ +// 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; +use std::env; +use std::num::NonZeroUsize; +use std::path::Path; +use std::sync::{Arc, LazyLock}; + +use datafusion::error::{DataFusionError, Result}; +use datafusion::execution::context::SessionConfig; +use datafusion::execution::memory_pool::{ + FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool, +}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::logical_expr::ExplainFormat; +use datafusion::prelude::SessionContext; +use crate::catalog::DynamicObjectStoreCatalog; +use crate::functions::{ + ListFilesCacheFunc, MetadataCacheFunc, ParquetMetadataFunc, StatisticsCacheFunc, +}; +use crate::object_storage::instrumented::{ + InstrumentedObjectStoreMode, InstrumentedObjectStoreRegistry, +}; +use crate::object_storage::{StdinCarriesCommands, is_stdin_location}; +use crate::{ + DATAFUSION_CLI_VERSION, exec, + pool_type::PoolType, + print_format::PrintFormat, + print_options::{MaxRows, PrintOptions}, +}; + +use clap::Parser; +use datafusion::common::config_err; +use datafusion::config::ConfigOptions; +use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; + +/// `CliSession` implements argument parsing, and construction of the default `CliSessionContext`. +pub struct CliSession { + ctx: SessionContext, + instrumented_registry: Arc, + args: Args, +} + +#[derive(Debug, Parser, PartialEq)] +#[clap(author, version, about, long_about= None)] +struct Args { + #[clap( + short = 'p', + long, + help = "Path to your data, default to current directory", + value_parser(parse_valid_data_dir) + )] + data_path: Option, + + #[clap( + short = 'b', + long, + help = "The batch size of each query, or use DataFusion default", + value_parser(parse_batch_size) + )] + batch_size: Option, + + #[clap( + short = 'c', + long, + num_args = 0.., + help = "Execute the given command string(s), then exit. Commands are expected to be non empty.", + value_parser(parse_command) + )] + command: Vec, + + #[clap( + short = 'm', + long, + help = "The memory pool limitation (e.g. '10g'), default to None (no limit)", + value_parser(extract_memory_pool_size) + )] + memory_limit: Option, + + #[clap( + short, + long, + num_args = 0.., + help = "Execute commands from file(s), then exit", + value_parser(parse_valid_file) + )] + file: Vec, + + #[clap( + short = 'r', + long, + num_args = 0.., + help = "Run the provided files on startup instead of ~/.datafusionrc", + value_parser(parse_valid_file), + conflicts_with = "file" + )] + rc: Option>, + + #[clap(long, value_enum, default_value_t = PrintFormat::Automatic)] + format: PrintFormat, + + #[clap( + short, + long, + help = "Reduce printing other than the results and work quietly" + )] + quiet: bool, + + #[clap( + long, + help = "Specify the memory pool type 'greedy' or 'fair'", + default_value_t = PoolType::Greedy + )] + mem_pool_type: PoolType, + + #[clap( + long, + help = "The number of top memory consumers to display when query fails due to memory exhaustion. To disable memory consumer tracking, set this value to 0", + default_value = "3" + )] + top_memory_consumers: usize, + + #[clap( + long, + help = "The max number of rows to display for 'Table' format\n[possible values: numbers(0/10/...), inf(no limit)]", + default_value = "40" + )] + maxrows: MaxRows, + + #[clap(long, help = "Enables console syntax highlighting")] + color: bool, + + #[clap( + short = 'd', + long, + help = "Available disk space for spilling queries (e.g. '10g'), default to None (uses DataFusion's default value of '100g')", + value_parser(extract_disk_limit) + )] + disk_limit: Option, + + #[clap( + long, + help = "Specify the default object_store_profiling mode, defaults to 'disabled'.\n[possible values: disabled, summary, trace]", + default_value_t = InstrumentedObjectStoreMode::Disabled + )] + object_store_profiling: InstrumentedObjectStoreMode, +} + +impl Args { + /// Without -c/-f the CLI enters the REPL, which reads its SQL from + /// stdin — interactively or piped. + fn repl_mode(&self) -> bool { + self.command.is_empty() && self.file.is_empty() + } + + /// Whether the CLI consumes stdin for its own SQL input. This covers the + /// REPL (no -c/-f, reading SQL interactively or piped) as well as an + /// explicit `-f /dev/stdin` (or the other stdin pseudo-paths), where the + /// SQL file *is* stdin. In either case stdin is already spoken for and + /// cannot also back a `LOCATION '/dev/stdin'` table. + fn reads_sql_from_stdin(&self) -> bool { + self.repl_mode() || self.file.iter().any(|f| is_stdin_location(f)) + } +} + +impl CliSession { + pub async fn entry_point() -> Result<()> { + let cli_session = CliSession::try_from_args()?; + Ok(cli_session.run().await?) + } + pub fn session_context(&self) -> &SessionContext { + &self.ctx + } +pub fn try_from_args() -> Result { + env_logger::init(); + let args = Args::parse(); + + if !args.quiet { + println!("DataFusion CLI v{DATAFUSION_CLI_VERSION}"); + } + + if let Some(ref path) = args.data_path { + let p = Path::new(path); + env::set_current_dir(p).unwrap(); + }; + + let session_config = get_session_config(&args)?; + + let mut rt_builder = RuntimeEnvBuilder::new(); + // set memory pool size + if let Some(memory_limit) = args.memory_limit { + // set memory pool type + let pool: Arc = match args.mem_pool_type { + PoolType::Fair if args.top_memory_consumers == 0 => { + Arc::new(FairSpillPool::new(memory_limit)) + } + PoolType::Fair => Arc::new(TrackConsumersPool::new( + FairSpillPool::new(memory_limit), + NonZeroUsize::new(args.top_memory_consumers).unwrap(), + )), + PoolType::Greedy if args.top_memory_consumers == 0 => { + Arc::new(GreedyMemoryPool::new(memory_limit)) + } + PoolType::Greedy => Arc::new(TrackConsumersPool::new( + GreedyMemoryPool::new(memory_limit), + NonZeroUsize::new(args.top_memory_consumers).unwrap(), + )), + }; + + rt_builder = rt_builder.with_memory_pool(pool) + } + + // set disk limit + if let Some(disk_limit) = args.disk_limit { + let builder = DiskManagerBuilder::default() + .with_mode(DiskManagerMode::OsTmpDirectory) + .with_max_temp_directory_size(disk_limit.try_into().unwrap()); + rt_builder = rt_builder.with_disk_manager_builder(builder); + } + + let instrumented_registry = Arc::new( + InstrumentedObjectStoreRegistry::new() + .with_profile_mode(args.object_store_profiling), + ); + rt_builder = rt_builder.with_object_store_registry(instrumented_registry.clone()); + + let runtime_env = rt_builder.build_arc()?; + + // enable dynamic file query + let ctx = SessionContext::new_with_config_rt(session_config, runtime_env) + .enable_url_table(); + Ok(Self { + ctx, + instrumented_registry, + args, + }) +} + +/// Main CLI entrypoint +pub async fn run(self) -> Result<()> { + let CliSession { + ctx, + args, + instrumented_registry, + } = self; + ctx.refresh_catalogs().await?; + // install dynamic catalog provider that can register required object stores + ctx.register_catalog_list(Arc::new(DynamicObjectStoreCatalog::new( + ctx.state().catalog_list().clone(), + ctx.state_weak_ref(), + ))); + // register `parquet_metadata` table function to get metadata from parquet files + ctx.register_udtf("parquet_metadata", Arc::new(ParquetMetadataFunc {})); + + // register `metadata_cache` table function to get the contents of the file metadata cache + ctx.register_udtf( + "metadata_cache", + Arc::new(MetadataCacheFunc::new( + ctx.task_ctx().runtime_env().cache_manager.clone(), + )), + ); + + // register `statistics_cache` table function to get the contents of the file statistics cache + ctx.register_udtf( + "statistics_cache", + Arc::new(StatisticsCacheFunc::new( + ctx.task_ctx().runtime_env().cache_manager.clone(), + )), + ); + + ctx.register_udtf( + "list_files_cache", + Arc::new(ListFilesCacheFunc::new( + ctx.task_ctx().runtime_env().cache_manager.clone(), + )), + ); + + let mut print_options = PrintOptions { + format: args.format, + quiet: args.quiet, + maxrows: args.maxrows, + color: args.color, + instrumented_registry: Arc::clone(&instrumented_registry), + }; + + let repl_mode = args.repl_mode(); + let commands = args.command; + let files = args.file; + let rc = match args.rc { + Some(file) => file, + None => { + let mut files = Vec::new(); + let home = dirs::home_dir(); + if let Some(p) = home { + let home_rc = p.join(".datafusionrc"); + if home_rc.exists() { + files.push(home_rc.into_os_string().into_string().unwrap()); + } + } + files + } + }; + + if repl_mode { + if !rc.is_empty() { + exec::exec_from_files(&ctx, rc, &print_options).await?; + } + // TODO maybe we can have thiserror for cli but for now let's keep it simple + return exec::exec_from_repl(&ctx, &mut print_options) + .await + .map_err(|e| DataFusionError::External(Box::new(e))); + } + + if !files.is_empty() { + exec::exec_from_files(&ctx, files, &print_options).await?; + } + + if !commands.is_empty() { + exec::exec_from_commands(&ctx, commands, &print_options).await?; + } + + Ok(()) +} +} + +/// Get the session configuration based on the provided arguments +/// and environment settings. +fn get_session_config(args: &Args) -> Result { + // Read options from environment variables and merge with command line options + let mut config_options = ConfigOptions::from_env()?; + + if let Some(batch_size) = args.batch_size { + if batch_size == 0 { + return config_err!("batch_size must be greater than 0"); + } + config_options.execution.batch_size = + datafusion_common::config::ConfigNonZeroUsize::try_new(batch_size)?; + }; + + // use easier to understand "tree" mode by default + // if the user hasn't specified an explain format in the environment + if env::var_os("DATAFUSION_EXPLAIN_FORMAT").is_none() { + config_options.explain.format = ExplainFormat::Tree; + } + + // in the CLI, we want to show NULL values rather the empty strings + if env::var_os("DATAFUSION_FORMAT_NULL").is_none() { + config_options.format.null = String::from("NULL"); + } + + let mut session_config = + SessionConfig::from(config_options).with_information_schema(true); + + if args.reads_sql_from_stdin() { + // When stdin carries the session's SQL — the REPL (including any rc + // file run before it) or an explicit `-f /dev/stdin` — it cannot also + // serve as a data source for `LOCATION '/dev/stdin'`. + session_config = session_config.with_extension(Arc::new(StdinCarriesCommands)); + } + + Ok(session_config) +} + +fn parse_valid_file(dir: &str) -> Result { + if Path::new(dir).is_file() { + Ok(dir.to_string()) + } else { + Err(format!("Invalid file '{dir}'")) + } +} + +fn parse_valid_data_dir(dir: &str) -> Result { + if Path::new(dir).is_dir() { + Ok(dir.to_string()) + } else { + Err(format!("Invalid data directory '{dir}'")) + } +} + +fn parse_batch_size(size: &str) -> Result { + match size.parse::() { + Ok(size) if size > 0 => Ok(size), + _ => Err(format!("Invalid batch size '{size}'")), + } +} + +fn parse_command(command: &str) -> Result { + if !command.is_empty() { + Ok(command.to_string()) + } else { + Err("-c flag expects only non empty commands".to_string()) + } +} + +#[derive(Debug, Clone, Copy)] +enum ByteUnit { + Byte, + KiB, + MiB, + GiB, + TiB, +} + +impl ByteUnit { + fn multiplier(&self) -> u64 { + match self { + ByteUnit::Byte => 1, + ByteUnit::KiB => 1 << 10, + ByteUnit::MiB => 1 << 20, + ByteUnit::GiB => 1 << 30, + ByteUnit::TiB => 1 << 40, + } + } +} + +fn parse_size_string(size: &str, label: &str) -> Result { + static BYTE_SUFFIXES: LazyLock> = + LazyLock::new(|| { + let mut m = HashMap::new(); + m.insert("b", ByteUnit::Byte); + m.insert("k", ByteUnit::KiB); + m.insert("kb", ByteUnit::KiB); + m.insert("m", ByteUnit::MiB); + m.insert("mb", ByteUnit::MiB); + m.insert("g", ByteUnit::GiB); + m.insert("gb", ByteUnit::GiB); + m.insert("t", ByteUnit::TiB); + m.insert("tb", ByteUnit::TiB); + m + }); + + static SUFFIX_REGEX: LazyLock = + LazyLock::new(|| regex::Regex::new(r"^(-?[0-9]+)([a-z]+)?$").unwrap()); + + let lower = size.to_lowercase(); + if let Some(caps) = SUFFIX_REGEX.captures(&lower) { + let num_str = caps.get(1).unwrap().as_str(); + let num = num_str + .parse::() + .map_err(|_| format!("Invalid numeric value in {label} '{size}'"))?; + + let suffix = caps.get(2).map(|m| m.as_str()).unwrap_or("b"); + let unit = BYTE_SUFFIXES + .get(suffix) + .ok_or_else(|| format!("Invalid {label} '{size}'"))?; + let total_bytes = usize::try_from(unit.multiplier()) + .ok() + .and_then(|multiplier| num.checked_mul(multiplier)) + .ok_or_else(|| format!("{label} '{size}' is too large"))?; + + Ok(total_bytes) + } else { + Err(format!("Invalid {label} '{size}'")) + } +} + +fn extract_memory_pool_size(size: &str) -> Result { + parse_size_string(size, "memory pool size") +} + +fn extract_disk_limit(size: &str) -> Result { + parse_size_string(size, "disk limit") +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::*; + use datafusion::execution::cache::default_cache::DefaultCache; + use datafusion::{ + common::test_util::batches_to_string, + execution::cache::cache_manager::CacheManagerConfig, + prelude::{ParquetReadOptions, col, lit, split_part}, + }; + use insta::assert_snapshot; + use object_store::memory::InMemory; + use url::Url; + + fn assert_conversion(input: &str, expected: Result) { + let result = extract_memory_pool_size(input); + match expected { + Ok(v) => assert_eq!(result.unwrap(), v), + Err(e) => assert_eq!(result.unwrap_err(), e), + } + } + + #[test] + fn memory_pool_size() -> Result<(), String> { + // Test basic sizes without suffix, assumed to be bytes + assert_conversion("5", Ok(5)); + assert_conversion("100", Ok(100)); + + // Test various units + assert_conversion("5b", Ok(5)); + assert_conversion("4k", Ok(4 * 1024)); + assert_conversion("4kb", Ok(4 * 1024)); + assert_conversion("20m", Ok(20 * 1024 * 1024)); + assert_conversion("20mb", Ok(20 * 1024 * 1024)); + assert_conversion("2g", Ok(2 * 1024 * 1024 * 1024)); + assert_conversion("2gb", Ok(2 * 1024 * 1024 * 1024)); + assert_conversion("3t", Ok(3 * 1024 * 1024 * 1024 * 1024)); + assert_conversion("4tb", Ok(4 * 1024 * 1024 * 1024 * 1024)); + + // Test case insensitivity + assert_conversion("4K", Ok(4 * 1024)); + assert_conversion("4KB", Ok(4 * 1024)); + assert_conversion("20M", Ok(20 * 1024 * 1024)); + assert_conversion("20MB", Ok(20 * 1024 * 1024)); + assert_conversion("2G", Ok(2 * 1024 * 1024 * 1024)); + assert_conversion("2GB", Ok(2 * 1024 * 1024 * 1024)); + assert_conversion("2T", Ok(2 * 1024 * 1024 * 1024 * 1024)); + + // Test invalid input + assert_conversion( + "invalid", + Err("Invalid memory pool size 'invalid'".to_string()), + ); + assert_conversion("4kbx", Err("Invalid memory pool size '4kbx'".to_string())); + assert_conversion( + "-20mb", + Err("Invalid numeric value in memory pool size '-20mb'".to_string()), + ); + assert_conversion( + "-100", + Err("Invalid numeric value in memory pool size '-100'".to_string()), + ); + assert_conversion( + "12k12k", + Err("Invalid memory pool size '12k12k'".to_string()), + ); + + Ok(()) + } + + #[tokio::test] + async fn test_parquet_metadata_works() -> Result<(), DataFusionError> { + let ctx = SessionContext::new(); + ctx.register_udtf("parquet_metadata", Arc::new(ParquetMetadataFunc {})); + + // input with single quote + let sql = "SELECT * FROM parquet_metadata('../datafusion/core/tests/data/fixed_size_list_array.parquet')"; + let df = ctx.sql(sql).await?; + let rbs = df.collect().await?; + + assert_snapshot!(batches_to_string(&rbs), @r#" + +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ + | filename | row_group_id | row_group_num_rows | row_group_num_columns | row_group_bytes | column_id | file_offset | num_values | path_in_schema | type | stats_min | stats_max | stats_null_count | stats_distinct_count | stats_min_value | stats_max_value | compression | encodings | index_page_offset | dictionary_page_offset | data_page_offset | total_compressed_size | total_uncompressed_size | + +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ + | ../datafusion/core/tests/data/fixed_size_list_array.parquet | 0 | 2 | 1 | 123 | 0 | 125 | 4 | "f0.list.item" | INT64 | 1 | 4 | 0 | | 1 | 4 | SNAPPY | [PLAIN, RLE, RLE_DICTIONARY] | | 4 | 46 | 121 | 123 | + +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ + "#); + + // input with double quote + let sql = "SELECT * FROM parquet_metadata(\"../datafusion/core/tests/data/fixed_size_list_array.parquet\")"; + let df = ctx.sql(sql).await?; + let rbs = df.collect().await?; + assert_snapshot!(batches_to_string(&rbs), @r#" + +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ + | filename | row_group_id | row_group_num_rows | row_group_num_columns | row_group_bytes | column_id | file_offset | num_values | path_in_schema | type | stats_min | stats_max | stats_null_count | stats_distinct_count | stats_min_value | stats_max_value | compression | encodings | index_page_offset | dictionary_page_offset | data_page_offset | total_compressed_size | total_uncompressed_size | + +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ + | ../datafusion/core/tests/data/fixed_size_list_array.parquet | 0 | 2 | 1 | 123 | 0 | 125 | 4 | "f0.list.item" | INT64 | 1 | 4 | 0 | | 1 | 4 | SNAPPY | [PLAIN, RLE, RLE_DICTIONARY] | | 4 | 46 | 121 | 123 | + +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ + "#); + + Ok(()) + } + + #[tokio::test] + async fn test_parquet_metadata_works_with_strings() -> Result<(), DataFusionError> { + let ctx = SessionContext::new(); + ctx.register_udtf("parquet_metadata", Arc::new(ParquetMetadataFunc {})); + + // input with string columns + let sql = "SELECT * FROM parquet_metadata('../parquet-testing/data/data_index_bloom_encoding_stats.parquet')"; + let df = ctx.sql(sql).await?; + let rbs = df.collect().await?; + + assert_snapshot!(batches_to_string(&rbs),@r#" + +-----------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+------------+-----------+-----------+------------------+----------------------+-----------------+-----------------+--------------------+--------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ + | filename | row_group_id | row_group_num_rows | row_group_num_columns | row_group_bytes | column_id | file_offset | num_values | path_in_schema | type | stats_min | stats_max | stats_null_count | stats_distinct_count | stats_min_value | stats_max_value | compression | encodings | index_page_offset | dictionary_page_offset | data_page_offset | total_compressed_size | total_uncompressed_size | + +-----------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+------------+-----------+-----------+------------------+----------------------+-----------------+-----------------+--------------------+--------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ + | ../parquet-testing/data/data_index_bloom_encoding_stats.parquet | 0 | 14 | 1 | 163 | 0 | 4 | 14 | "String" | BYTE_ARRAY | Hello | today | 0 | | Hello | today | GZIP(GzipLevel(6)) | [PLAIN, RLE, BIT_PACKED] | | | 4 | 152 | 163 | + +-----------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+------------+-----------+-----------+------------------+----------------------+-----------------+-----------------+--------------------+--------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ + "#); + + Ok(()) + } + + #[tokio::test] + async fn test_metadata_cache() -> Result<(), DataFusionError> { + let ctx = SessionContext::new(); + ctx.register_udtf( + "metadata_cache", + Arc::new(MetadataCacheFunc::new( + ctx.task_ctx().runtime_env().cache_manager.clone(), + )), + ); + + ctx.register_parquet( + "alltypes_plain", + "../parquet-testing/data/alltypes_plain.parquet", + ParquetReadOptions::new(), + ) + .await?; + + ctx.register_parquet( + "alltypes_tiny_pages", + "../parquet-testing/data/alltypes_tiny_pages.parquet", + ParquetReadOptions::new(), + ) + .await?; + + ctx.register_parquet( + "lz4_raw_compressed_larger", + "../parquet-testing/data/lz4_raw_compressed_larger.parquet", + ParquetReadOptions::new(), + ) + .await?; + + ctx.sql("select * from alltypes_plain") + .await? + .collect() + .await?; + ctx.sql("select * from alltypes_tiny_pages") + .await? + .collect() + .await?; + ctx.sql("select * from lz4_raw_compressed_larger") + .await? + .collect() + .await?; + + // initial state + let sql = "SELECT split_part(path, '/', -1) as filename, file_size_bytes, metadata_size_bytes, hits, extra from metadata_cache() order by filename"; + let df = ctx.sql(sql).await?; + let rbs = df.collect().await?; + + assert_snapshot!(batches_to_string(&rbs),@r" + +-----------------------------------+-----------------+---------------------+------+------------------+ + | filename | file_size_bytes | metadata_size_bytes | hits | extra | + +-----------------------------------+-----------------+---------------------+------+------------------+ + | alltypes_plain.parquet | 1851 | 8794 | 1 | page_index=false | + | alltypes_tiny_pages.parquet | 454233 | 268970 | 2 | page_index=true | + | lz4_raw_compressed_larger.parquet | 380836 | 1331 | 1 | page_index=false | + +-----------------------------------+-----------------+---------------------+------+------------------+ + "); + + // increase the number of hits + ctx.sql("select * from alltypes_plain") + .await? + .collect() + .await?; + ctx.sql("select * from alltypes_plain") + .await? + .collect() + .await?; + ctx.sql("select * from alltypes_plain") + .await? + .collect() + .await?; + ctx.sql("select * from lz4_raw_compressed_larger") + .await? + .collect() + .await?; + let sql = "select split_part(path, '/', -1) as filename, file_size_bytes, metadata_size_bytes, hits, extra from metadata_cache() order by filename"; + let df = ctx.sql(sql).await?; + let rbs = df.collect().await?; + + assert_snapshot!(batches_to_string(&rbs),@r" + +-----------------------------------+-----------------+---------------------+------+------------------+ + | filename | file_size_bytes | metadata_size_bytes | hits | extra | + +-----------------------------------+-----------------+---------------------+------+------------------+ + | alltypes_plain.parquet | 1851 | 8794 | 4 | page_index=false | + | alltypes_tiny_pages.parquet | 454233 | 268970 | 2 | page_index=true | + | lz4_raw_compressed_larger.parquet | 380836 | 1331 | 2 | page_index=false | + +-----------------------------------+-----------------+---------------------+------+------------------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn test_statistics_cache_default() -> Result<(), DataFusionError> { + let ctx = SessionContext::new(); + + ctx.register_udtf( + "statistics_cache", + Arc::new(StatisticsCacheFunc::new( + ctx.task_ctx().runtime_env().cache_manager.clone(), + )), + ); + + for filename in [ + "alltypes_plain", + "alltypes_tiny_pages", + "lz4_raw_compressed_larger", + ] { + ctx.sql( + format!( + "create external table {filename} + stored as parquet + location '../parquet-testing/data/{filename}.parquet'", + ) + .as_str(), + ) + .await? + .collect() + .await?; + } + + let sql = "SELECT split_part(path, '/', -1) as filename, table, file_size_bytes, num_rows, num_columns, hits, table_size_bytes from statistics_cache() order by filename"; + let df = ctx.sql(sql).await?; + let rbs = df.collect().await?; + assert_snapshot!(batches_to_string(&rbs),@" + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ + | filename | table | file_size_bytes | num_rows | num_columns | hits | table_size_bytes | + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ + | alltypes_plain.parquet | alltypes_plain | 1851 | Exact(8) | 11 | 0 | Absent | + | alltypes_tiny_pages.parquet | alltypes_tiny_pages | 454233 | Exact(7300) | 13 | 0 | Absent | + | lz4_raw_compressed_larger.parquet | lz4_raw_compressed_larger | 380836 | Exact(10000) | 1 | 0 | Absent | + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ + "); + + // increase the number of hits + ctx.sql("select * from alltypes_plain") + .await? + .collect() + .await?; + ctx.sql("select * from alltypes_plain") + .await? + .collect() + .await?; + ctx.sql("select * from alltypes_plain") + .await? + .collect() + .await?; + ctx.sql("select * from lz4_raw_compressed_larger") + .await? + .collect() + .await?; + + let sql = "SELECT split_part(path, '/', -1) as filename, table, file_size_bytes, num_rows, num_columns, hits, table_size_bytes from statistics_cache() order by filename"; + let df = ctx.sql(sql).await?; + let rbs = df.collect().await?; + assert_snapshot!(batches_to_string(&rbs),@" + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ + | filename | table | file_size_bytes | num_rows | num_columns | hits | table_size_bytes | + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ + | alltypes_plain.parquet | alltypes_plain | 1851 | Exact(8) | 11 | 3 | Absent | + | alltypes_tiny_pages.parquet | alltypes_tiny_pages | 454233 | Exact(7300) | 13 | 0 | Absent | + | lz4_raw_compressed_larger.parquet | lz4_raw_compressed_larger | 380836 | Exact(10000) | 1 | 1 | Absent | + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn test_list_files_cache() -> Result<(), DataFusionError> { + let list_files_cache = Arc::new(DefaultCache::new_with_ttl( + 1024, + Some(Duration::from_secs(1)), + )); + + let rt = RuntimeEnvBuilder::new() + .with_cache_manager( + CacheManagerConfig::default() + .with_list_files_cache(Some(list_files_cache)), + ) + .build_arc() + .unwrap(); + + let ctx = SessionContext::new_with_config_rt(SessionConfig::default(), rt); + + ctx.register_object_store( + &Url::parse("mem://test_table").unwrap(), + Arc::new(InMemory::new()), + ); + + ctx.register_udtf( + "list_files_cache", + Arc::new(ListFilesCacheFunc::new( + ctx.task_ctx().runtime_env().cache_manager.clone(), + )), + ); + + ctx.sql( + "CREATE EXTERNAL TABLE src_table + STORED AS PARQUET + LOCATION '../parquet-testing/data/alltypes_plain.parquet'", + ) + .await? + .collect() + .await?; + + ctx.sql("COPY (SELECT * FROM src_table) TO 'mem://test_table/0.parquet' STORED AS PARQUET").await?.collect().await?; + + ctx.sql("COPY (SELECT * FROM src_table) TO 'mem://test_table/1.parquet' STORED AS PARQUET").await?.collect().await?; + + ctx.sql( + "CREATE EXTERNAL TABLE test_table + STORED AS PARQUET + LOCATION 'mem://test_table/' + ", + ) + .await? + .collect() + .await?; + + let sql = "SELECT metadata_size_bytes, expires_in, metadata_list FROM list_files_cache()"; + let df = ctx + .sql(sql) + .await? + .unnest_columns(&["metadata_list"])? + .with_column_renamed("metadata_list", "metadata")? + .unnest_columns(&["metadata"])?; + + assert_eq!( + 2, + df.clone() + .filter(col("expires_in").is_not_null())? + .count() + .await? + ); + + let df = df + .with_column_renamed(r#""metadata.file_size_bytes""#, "file_size_bytes")? + .with_column_renamed(r#""metadata.e_tag""#, "etag")? + .with_column( + "filename", + split_part(col(r#""metadata.file_path""#), lit("/"), lit(-1)), + )? + .select_columns(&[ + "metadata_size_bytes", + "filename", + "file_size_bytes", + "etag", + ])? + .sort(vec![col("filename").sort(true, false)])?; + let rbs = df.collect().await?; + assert_snapshot!(batches_to_string(&rbs),@r" + +---------------------+-----------+-----------------+------+ + | metadata_size_bytes | filename | file_size_bytes | etag | + +---------------------+-----------+-----------------+------+ + | 212 | 0.parquet | 3642 | 0 | + | 212 | 1.parquet | 3642 | 1 | + +---------------------+-----------+-----------------+------+ + "); + + Ok(()) + } +} diff --git a/datafusion-cli/src/lib.rs b/datafusion-cli/src/lib.rs index f0b0bc23fd73d..8ab871ad9bd72 100644 --- a/datafusion-cli/src/lib.rs +++ b/datafusion-cli/src/lib.rs @@ -26,6 +26,7 @@ pub const DATAFUSION_CLI_VERSION: &str = env!("CARGO_PKG_VERSION"); pub mod catalog; pub mod cli_context; pub mod command; +pub mod entry_point; pub mod exec; pub mod functions; pub mod helper; diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs index 78d8342020cc2..e12f7706a2887 100644 --- a/datafusion-cli/src/main.rs +++ b/datafusion-cli/src/main.rs @@ -15,841 +15,21 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashMap; -use std::env; -use std::num::NonZeroUsize; -use std::path::Path; +use datafusion_cli::entry_point::CliSession; use std::process::ExitCode; -use std::sync::{Arc, LazyLock}; -use datafusion::error::{DataFusionError, Result}; -use datafusion::execution::context::SessionConfig; -use datafusion::execution::memory_pool::{ - FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool, -}; -use datafusion::execution::runtime_env::RuntimeEnvBuilder; -use datafusion::logical_expr::ExplainFormat; -use datafusion::prelude::SessionContext; -use datafusion_cli::catalog::DynamicObjectStoreCatalog; -use datafusion_cli::functions::{ - ListFilesCacheFunc, MetadataCacheFunc, ParquetMetadataFunc, StatisticsCacheFunc, -}; -use datafusion_cli::object_storage::instrumented::{ - InstrumentedObjectStoreMode, InstrumentedObjectStoreRegistry, -}; -use datafusion_cli::object_storage::{StdinCarriesCommands, is_stdin_location}; -use datafusion_cli::{ - DATAFUSION_CLI_VERSION, exec, - pool_type::PoolType, - print_format::PrintFormat, - print_options::{MaxRows, PrintOptions}, -}; - -use clap::Parser; -use datafusion::common::config_err; -use datafusion::config::ConfigOptions; -use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; use mimalloc::MiMalloc; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; -#[derive(Debug, Parser, PartialEq)] -#[clap(author, version, about, long_about= None)] -struct Args { - #[clap( - short = 'p', - long, - help = "Path to your data, default to current directory", - value_parser(parse_valid_data_dir) - )] - data_path: Option, - - #[clap( - short = 'b', - long, - help = "The batch size of each query, or use DataFusion default", - value_parser(parse_batch_size) - )] - batch_size: Option, - - #[clap( - short = 'c', - long, - num_args = 0.., - help = "Execute the given command string(s), then exit. Commands are expected to be non empty.", - value_parser(parse_command) - )] - command: Vec, - - #[clap( - short = 'm', - long, - help = "The memory pool limitation (e.g. '10g'), default to None (no limit)", - value_parser(extract_memory_pool_size) - )] - memory_limit: Option, - - #[clap( - short, - long, - num_args = 0.., - help = "Execute commands from file(s), then exit", - value_parser(parse_valid_file) - )] - file: Vec, - - #[clap( - short = 'r', - long, - num_args = 0.., - help = "Run the provided files on startup instead of ~/.datafusionrc", - value_parser(parse_valid_file), - conflicts_with = "file" - )] - rc: Option>, - - #[clap(long, value_enum, default_value_t = PrintFormat::Automatic)] - format: PrintFormat, - - #[clap( - short, - long, - help = "Reduce printing other than the results and work quietly" - )] - quiet: bool, - - #[clap( - long, - help = "Specify the memory pool type 'greedy' or 'fair'", - default_value_t = PoolType::Greedy - )] - mem_pool_type: PoolType, - - #[clap( - long, - help = "The number of top memory consumers to display when query fails due to memory exhaustion. To disable memory consumer tracking, set this value to 0", - default_value = "3" - )] - top_memory_consumers: usize, - - #[clap( - long, - help = "The max number of rows to display for 'Table' format\n[possible values: numbers(0/10/...), inf(no limit)]", - default_value = "40" - )] - maxrows: MaxRows, - - #[clap(long, help = "Enables console syntax highlighting")] - color: bool, - - #[clap( - short = 'd', - long, - help = "Available disk space for spilling queries (e.g. '10g'), default to None (uses DataFusion's default value of '100g')", - value_parser(extract_disk_limit) - )] - disk_limit: Option, - - #[clap( - long, - help = "Specify the default object_store_profiling mode, defaults to 'disabled'.\n[possible values: disabled, summary, trace]", - default_value_t = InstrumentedObjectStoreMode::Disabled - )] - object_store_profiling: InstrumentedObjectStoreMode, -} - -impl Args { - /// Without -c/-f the CLI enters the REPL, which reads its SQL from - /// stdin — interactively or piped. - fn repl_mode(&self) -> bool { - self.command.is_empty() && self.file.is_empty() - } - - /// Whether the CLI consumes stdin for its own SQL input. This covers the - /// REPL (no -c/-f, reading SQL interactively or piped) as well as an - /// explicit `-f /dev/stdin` (or the other stdin pseudo-paths), where the - /// SQL file *is* stdin. In either case stdin is already spoken for and - /// cannot also back a `LOCATION '/dev/stdin'` table. - fn reads_sql_from_stdin(&self) -> bool { - self.repl_mode() || self.file.iter().any(|f| is_stdin_location(f)) - } -} - #[tokio::main] /// Calls [`main_inner`], then handles printing errors and returning the correct exit code pub async fn main() -> ExitCode { - if let Err(e) = main_inner().await { + if let Err(e) = CliSession::entry_point().await { println!("Error: {e}"); return ExitCode::FAILURE; } ExitCode::SUCCESS } - -/// Main CLI entrypoint -async fn main_inner() -> Result<()> { - env_logger::init(); - let args = Args::parse(); - - if !args.quiet { - println!("DataFusion CLI v{DATAFUSION_CLI_VERSION}"); - } - - if let Some(ref path) = args.data_path { - let p = Path::new(path); - env::set_current_dir(p).unwrap(); - }; - - let session_config = get_session_config(&args)?; - - let mut rt_builder = RuntimeEnvBuilder::new(); - // set memory pool size - if let Some(memory_limit) = args.memory_limit { - // set memory pool type - let pool: Arc = match args.mem_pool_type { - PoolType::Fair if args.top_memory_consumers == 0 => { - Arc::new(FairSpillPool::new(memory_limit)) - } - PoolType::Fair => Arc::new(TrackConsumersPool::new( - FairSpillPool::new(memory_limit), - NonZeroUsize::new(args.top_memory_consumers).unwrap(), - )), - PoolType::Greedy if args.top_memory_consumers == 0 => { - Arc::new(GreedyMemoryPool::new(memory_limit)) - } - PoolType::Greedy => Arc::new(TrackConsumersPool::new( - GreedyMemoryPool::new(memory_limit), - NonZeroUsize::new(args.top_memory_consumers).unwrap(), - )), - }; - - rt_builder = rt_builder.with_memory_pool(pool) - } - - // set disk limit - if let Some(disk_limit) = args.disk_limit { - let builder = DiskManagerBuilder::default() - .with_mode(DiskManagerMode::OsTmpDirectory) - .with_max_temp_directory_size(disk_limit.try_into().unwrap()); - rt_builder = rt_builder.with_disk_manager_builder(builder); - } - - let instrumented_registry = Arc::new( - InstrumentedObjectStoreRegistry::new() - .with_profile_mode(args.object_store_profiling), - ); - rt_builder = rt_builder.with_object_store_registry(instrumented_registry.clone()); - - let runtime_env = rt_builder.build_arc()?; - - // enable dynamic file query - let ctx = SessionContext::new_with_config_rt(session_config, runtime_env) - .enable_url_table(); - ctx.refresh_catalogs().await?; - // install dynamic catalog provider that can register required object stores - ctx.register_catalog_list(Arc::new(DynamicObjectStoreCatalog::new( - ctx.state().catalog_list().clone(), - ctx.state_weak_ref(), - ))); - // register `parquet_metadata` table function to get metadata from parquet files - ctx.register_udtf("parquet_metadata", Arc::new(ParquetMetadataFunc {})); - - // register `metadata_cache` table function to get the contents of the file metadata cache - ctx.register_udtf( - "metadata_cache", - Arc::new(MetadataCacheFunc::new( - ctx.task_ctx().runtime_env().cache_manager.clone(), - )), - ); - - // register `statistics_cache` table function to get the contents of the file statistics cache - ctx.register_udtf( - "statistics_cache", - Arc::new(StatisticsCacheFunc::new( - ctx.task_ctx().runtime_env().cache_manager.clone(), - )), - ); - - ctx.register_udtf( - "list_files_cache", - Arc::new(ListFilesCacheFunc::new( - ctx.task_ctx().runtime_env().cache_manager.clone(), - )), - ); - - let mut print_options = PrintOptions { - format: args.format, - quiet: args.quiet, - maxrows: args.maxrows, - color: args.color, - instrumented_registry: Arc::clone(&instrumented_registry), - }; - - let repl_mode = args.repl_mode(); - let commands = args.command; - let files = args.file; - let rc = match args.rc { - Some(file) => file, - None => { - let mut files = Vec::new(); - let home = dirs::home_dir(); - if let Some(p) = home { - let home_rc = p.join(".datafusionrc"); - if home_rc.exists() { - files.push(home_rc.into_os_string().into_string().unwrap()); - } - } - files - } - }; - - if repl_mode { - if !rc.is_empty() { - exec::exec_from_files(&ctx, rc, &print_options).await?; - } - // TODO maybe we can have thiserror for cli but for now let's keep it simple - return exec::exec_from_repl(&ctx, &mut print_options) - .await - .map_err(|e| DataFusionError::External(Box::new(e))); - } - - if !files.is_empty() { - exec::exec_from_files(&ctx, files, &print_options).await?; - } - - if !commands.is_empty() { - exec::exec_from_commands(&ctx, commands, &print_options).await?; - } - - Ok(()) -} - -/// Get the session configuration based on the provided arguments -/// and environment settings. -fn get_session_config(args: &Args) -> Result { - // Read options from environment variables and merge with command line options - let mut config_options = ConfigOptions::from_env()?; - - if let Some(batch_size) = args.batch_size { - if batch_size == 0 { - return config_err!("batch_size must be greater than 0"); - } - config_options.execution.batch_size = - datafusion_common::config::ConfigNonZeroUsize::try_new(batch_size)?; - }; - - // use easier to understand "tree" mode by default - // if the user hasn't specified an explain format in the environment - if env::var_os("DATAFUSION_EXPLAIN_FORMAT").is_none() { - config_options.explain.format = ExplainFormat::Tree; - } - - // in the CLI, we want to show NULL values rather the empty strings - if env::var_os("DATAFUSION_FORMAT_NULL").is_none() { - config_options.format.null = String::from("NULL"); - } - - let mut session_config = - SessionConfig::from(config_options).with_information_schema(true); - - if args.reads_sql_from_stdin() { - // When stdin carries the session's SQL — the REPL (including any rc - // file run before it) or an explicit `-f /dev/stdin` — it cannot also - // serve as a data source for `LOCATION '/dev/stdin'`. - session_config = session_config.with_extension(Arc::new(StdinCarriesCommands)); - } - - Ok(session_config) -} - -fn parse_valid_file(dir: &str) -> Result { - if Path::new(dir).is_file() { - Ok(dir.to_string()) - } else { - Err(format!("Invalid file '{dir}'")) - } -} - -fn parse_valid_data_dir(dir: &str) -> Result { - if Path::new(dir).is_dir() { - Ok(dir.to_string()) - } else { - Err(format!("Invalid data directory '{dir}'")) - } -} - -fn parse_batch_size(size: &str) -> Result { - match size.parse::() { - Ok(size) if size > 0 => Ok(size), - _ => Err(format!("Invalid batch size '{size}'")), - } -} - -fn parse_command(command: &str) -> Result { - if !command.is_empty() { - Ok(command.to_string()) - } else { - Err("-c flag expects only non empty commands".to_string()) - } -} - -#[derive(Debug, Clone, Copy)] -enum ByteUnit { - Byte, - KiB, - MiB, - GiB, - TiB, -} - -impl ByteUnit { - fn multiplier(&self) -> u64 { - match self { - ByteUnit::Byte => 1, - ByteUnit::KiB => 1 << 10, - ByteUnit::MiB => 1 << 20, - ByteUnit::GiB => 1 << 30, - ByteUnit::TiB => 1 << 40, - } - } -} - -fn parse_size_string(size: &str, label: &str) -> Result { - static BYTE_SUFFIXES: LazyLock> = - LazyLock::new(|| { - let mut m = HashMap::new(); - m.insert("b", ByteUnit::Byte); - m.insert("k", ByteUnit::KiB); - m.insert("kb", ByteUnit::KiB); - m.insert("m", ByteUnit::MiB); - m.insert("mb", ByteUnit::MiB); - m.insert("g", ByteUnit::GiB); - m.insert("gb", ByteUnit::GiB); - m.insert("t", ByteUnit::TiB); - m.insert("tb", ByteUnit::TiB); - m - }); - - static SUFFIX_REGEX: LazyLock = - LazyLock::new(|| regex::Regex::new(r"^(-?[0-9]+)([a-z]+)?$").unwrap()); - - let lower = size.to_lowercase(); - if let Some(caps) = SUFFIX_REGEX.captures(&lower) { - let num_str = caps.get(1).unwrap().as_str(); - let num = num_str - .parse::() - .map_err(|_| format!("Invalid numeric value in {label} '{size}'"))?; - - let suffix = caps.get(2).map(|m| m.as_str()).unwrap_or("b"); - let unit = BYTE_SUFFIXES - .get(suffix) - .ok_or_else(|| format!("Invalid {label} '{size}'"))?; - let total_bytes = usize::try_from(unit.multiplier()) - .ok() - .and_then(|multiplier| num.checked_mul(multiplier)) - .ok_or_else(|| format!("{label} '{size}' is too large"))?; - - Ok(total_bytes) - } else { - Err(format!("Invalid {label} '{size}'")) - } -} - -pub fn extract_memory_pool_size(size: &str) -> Result { - parse_size_string(size, "memory pool size") -} - -pub fn extract_disk_limit(size: &str) -> Result { - parse_size_string(size, "disk limit") -} - -#[cfg(test)] -mod tests { - use std::time::Duration; - - use super::*; - use datafusion::execution::cache::default_cache::DefaultCache; - use datafusion::{ - common::test_util::batches_to_string, - execution::cache::cache_manager::CacheManagerConfig, - prelude::{ParquetReadOptions, col, lit, split_part}, - }; - use insta::assert_snapshot; - use object_store::memory::InMemory; - use url::Url; - - fn assert_conversion(input: &str, expected: Result) { - let result = extract_memory_pool_size(input); - match expected { - Ok(v) => assert_eq!(result.unwrap(), v), - Err(e) => assert_eq!(result.unwrap_err(), e), - } - } - - #[test] - fn memory_pool_size() -> Result<(), String> { - // Test basic sizes without suffix, assumed to be bytes - assert_conversion("5", Ok(5)); - assert_conversion("100", Ok(100)); - - // Test various units - assert_conversion("5b", Ok(5)); - assert_conversion("4k", Ok(4 * 1024)); - assert_conversion("4kb", Ok(4 * 1024)); - assert_conversion("20m", Ok(20 * 1024 * 1024)); - assert_conversion("20mb", Ok(20 * 1024 * 1024)); - assert_conversion("2g", Ok(2 * 1024 * 1024 * 1024)); - assert_conversion("2gb", Ok(2 * 1024 * 1024 * 1024)); - assert_conversion("3t", Ok(3 * 1024 * 1024 * 1024 * 1024)); - assert_conversion("4tb", Ok(4 * 1024 * 1024 * 1024 * 1024)); - - // Test case insensitivity - assert_conversion("4K", Ok(4 * 1024)); - assert_conversion("4KB", Ok(4 * 1024)); - assert_conversion("20M", Ok(20 * 1024 * 1024)); - assert_conversion("20MB", Ok(20 * 1024 * 1024)); - assert_conversion("2G", Ok(2 * 1024 * 1024 * 1024)); - assert_conversion("2GB", Ok(2 * 1024 * 1024 * 1024)); - assert_conversion("2T", Ok(2 * 1024 * 1024 * 1024 * 1024)); - - // Test invalid input - assert_conversion( - "invalid", - Err("Invalid memory pool size 'invalid'".to_string()), - ); - assert_conversion("4kbx", Err("Invalid memory pool size '4kbx'".to_string())); - assert_conversion( - "-20mb", - Err("Invalid numeric value in memory pool size '-20mb'".to_string()), - ); - assert_conversion( - "-100", - Err("Invalid numeric value in memory pool size '-100'".to_string()), - ); - assert_conversion( - "12k12k", - Err("Invalid memory pool size '12k12k'".to_string()), - ); - - Ok(()) - } - - #[tokio::test] - async fn test_parquet_metadata_works() -> Result<(), DataFusionError> { - let ctx = SessionContext::new(); - ctx.register_udtf("parquet_metadata", Arc::new(ParquetMetadataFunc {})); - - // input with single quote - let sql = "SELECT * FROM parquet_metadata('../datafusion/core/tests/data/fixed_size_list_array.parquet')"; - let df = ctx.sql(sql).await?; - let rbs = df.collect().await?; - - assert_snapshot!(batches_to_string(&rbs), @r#" - +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ - | filename | row_group_id | row_group_num_rows | row_group_num_columns | row_group_bytes | column_id | file_offset | num_values | path_in_schema | type | stats_min | stats_max | stats_null_count | stats_distinct_count | stats_min_value | stats_max_value | compression | encodings | index_page_offset | dictionary_page_offset | data_page_offset | total_compressed_size | total_uncompressed_size | - +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ - | ../datafusion/core/tests/data/fixed_size_list_array.parquet | 0 | 2 | 1 | 123 | 0 | 125 | 4 | "f0.list.item" | INT64 | 1 | 4 | 0 | | 1 | 4 | SNAPPY | [PLAIN, RLE, RLE_DICTIONARY] | | 4 | 46 | 121 | 123 | - +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ - "#); - - // input with double quote - let sql = "SELECT * FROM parquet_metadata(\"../datafusion/core/tests/data/fixed_size_list_array.parquet\")"; - let df = ctx.sql(sql).await?; - let rbs = df.collect().await?; - assert_snapshot!(batches_to_string(&rbs), @r#" - +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ - | filename | row_group_id | row_group_num_rows | row_group_num_columns | row_group_bytes | column_id | file_offset | num_values | path_in_schema | type | stats_min | stats_max | stats_null_count | stats_distinct_count | stats_min_value | stats_max_value | compression | encodings | index_page_offset | dictionary_page_offset | data_page_offset | total_compressed_size | total_uncompressed_size | - +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ - | ../datafusion/core/tests/data/fixed_size_list_array.parquet | 0 | 2 | 1 | 123 | 0 | 125 | 4 | "f0.list.item" | INT64 | 1 | 4 | 0 | | 1 | 4 | SNAPPY | [PLAIN, RLE, RLE_DICTIONARY] | | 4 | 46 | 121 | 123 | - +-------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+-------+-----------+-----------+------------------+----------------------+-----------------+-----------------+-------------+------------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ - "#); - - Ok(()) - } - - #[tokio::test] - async fn test_parquet_metadata_works_with_strings() -> Result<(), DataFusionError> { - let ctx = SessionContext::new(); - ctx.register_udtf("parquet_metadata", Arc::new(ParquetMetadataFunc {})); - - // input with string columns - let sql = "SELECT * FROM parquet_metadata('../parquet-testing/data/data_index_bloom_encoding_stats.parquet')"; - let df = ctx.sql(sql).await?; - let rbs = df.collect().await?; - - assert_snapshot!(batches_to_string(&rbs),@r#" - +-----------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+------------+-----------+-----------+------------------+----------------------+-----------------+-----------------+--------------------+--------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ - | filename | row_group_id | row_group_num_rows | row_group_num_columns | row_group_bytes | column_id | file_offset | num_values | path_in_schema | type | stats_min | stats_max | stats_null_count | stats_distinct_count | stats_min_value | stats_max_value | compression | encodings | index_page_offset | dictionary_page_offset | data_page_offset | total_compressed_size | total_uncompressed_size | - +-----------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+------------+-----------+-----------+------------------+----------------------+-----------------+-----------------+--------------------+--------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ - | ../parquet-testing/data/data_index_bloom_encoding_stats.parquet | 0 | 14 | 1 | 163 | 0 | 4 | 14 | "String" | BYTE_ARRAY | Hello | today | 0 | | Hello | today | GZIP(GzipLevel(6)) | [PLAIN, RLE, BIT_PACKED] | | | 4 | 152 | 163 | - +-----------------------------------------------------------------+--------------+--------------------+-----------------------+-----------------+-----------+-------------+------------+----------------+------------+-----------+-----------+------------------+----------------------+-----------------+-----------------+--------------------+--------------------------+-------------------+------------------------+------------------+-----------------------+-------------------------+ - "#); - - Ok(()) - } - - #[tokio::test] - async fn test_metadata_cache() -> Result<(), DataFusionError> { - let ctx = SessionContext::new(); - ctx.register_udtf( - "metadata_cache", - Arc::new(MetadataCacheFunc::new( - ctx.task_ctx().runtime_env().cache_manager.clone(), - )), - ); - - ctx.register_parquet( - "alltypes_plain", - "../parquet-testing/data/alltypes_plain.parquet", - ParquetReadOptions::new(), - ) - .await?; - - ctx.register_parquet( - "alltypes_tiny_pages", - "../parquet-testing/data/alltypes_tiny_pages.parquet", - ParquetReadOptions::new(), - ) - .await?; - - ctx.register_parquet( - "lz4_raw_compressed_larger", - "../parquet-testing/data/lz4_raw_compressed_larger.parquet", - ParquetReadOptions::new(), - ) - .await?; - - ctx.sql("select * from alltypes_plain") - .await? - .collect() - .await?; - ctx.sql("select * from alltypes_tiny_pages") - .await? - .collect() - .await?; - ctx.sql("select * from lz4_raw_compressed_larger") - .await? - .collect() - .await?; - - // initial state - let sql = "SELECT split_part(path, '/', -1) as filename, file_size_bytes, metadata_size_bytes, hits, extra from metadata_cache() order by filename"; - let df = ctx.sql(sql).await?; - let rbs = df.collect().await?; - - assert_snapshot!(batches_to_string(&rbs),@r" - +-----------------------------------+-----------------+---------------------+------+------------------+ - | filename | file_size_bytes | metadata_size_bytes | hits | extra | - +-----------------------------------+-----------------+---------------------+------+------------------+ - | alltypes_plain.parquet | 1851 | 8794 | 1 | page_index=false | - | alltypes_tiny_pages.parquet | 454233 | 268970 | 2 | page_index=true | - | lz4_raw_compressed_larger.parquet | 380836 | 1331 | 1 | page_index=false | - +-----------------------------------+-----------------+---------------------+------+------------------+ - "); - - // increase the number of hits - ctx.sql("select * from alltypes_plain") - .await? - .collect() - .await?; - ctx.sql("select * from alltypes_plain") - .await? - .collect() - .await?; - ctx.sql("select * from alltypes_plain") - .await? - .collect() - .await?; - ctx.sql("select * from lz4_raw_compressed_larger") - .await? - .collect() - .await?; - let sql = "select split_part(path, '/', -1) as filename, file_size_bytes, metadata_size_bytes, hits, extra from metadata_cache() order by filename"; - let df = ctx.sql(sql).await?; - let rbs = df.collect().await?; - - assert_snapshot!(batches_to_string(&rbs),@r" - +-----------------------------------+-----------------+---------------------+------+------------------+ - | filename | file_size_bytes | metadata_size_bytes | hits | extra | - +-----------------------------------+-----------------+---------------------+------+------------------+ - | alltypes_plain.parquet | 1851 | 8794 | 4 | page_index=false | - | alltypes_tiny_pages.parquet | 454233 | 268970 | 2 | page_index=true | - | lz4_raw_compressed_larger.parquet | 380836 | 1331 | 2 | page_index=false | - +-----------------------------------+-----------------+---------------------+------+------------------+ - "); - - Ok(()) - } - - #[tokio::test] - async fn test_statistics_cache_default() -> Result<(), DataFusionError> { - let ctx = SessionContext::new(); - - ctx.register_udtf( - "statistics_cache", - Arc::new(StatisticsCacheFunc::new( - ctx.task_ctx().runtime_env().cache_manager.clone(), - )), - ); - - for filename in [ - "alltypes_plain", - "alltypes_tiny_pages", - "lz4_raw_compressed_larger", - ] { - ctx.sql( - format!( - "create external table {filename} - stored as parquet - location '../parquet-testing/data/{filename}.parquet'", - ) - .as_str(), - ) - .await? - .collect() - .await?; - } - - let sql = "SELECT split_part(path, '/', -1) as filename, table, file_size_bytes, num_rows, num_columns, hits, table_size_bytes from statistics_cache() order by filename"; - let df = ctx.sql(sql).await?; - let rbs = df.collect().await?; - assert_snapshot!(batches_to_string(&rbs),@" - +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ - | filename | table | file_size_bytes | num_rows | num_columns | hits | table_size_bytes | - +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ - | alltypes_plain.parquet | alltypes_plain | 1851 | Exact(8) | 11 | 0 | Absent | - | alltypes_tiny_pages.parquet | alltypes_tiny_pages | 454233 | Exact(7300) | 13 | 0 | Absent | - | lz4_raw_compressed_larger.parquet | lz4_raw_compressed_larger | 380836 | Exact(10000) | 1 | 0 | Absent | - +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ - "); - - // increase the number of hits - ctx.sql("select * from alltypes_plain") - .await? - .collect() - .await?; - ctx.sql("select * from alltypes_plain") - .await? - .collect() - .await?; - ctx.sql("select * from alltypes_plain") - .await? - .collect() - .await?; - ctx.sql("select * from lz4_raw_compressed_larger") - .await? - .collect() - .await?; - - let sql = "SELECT split_part(path, '/', -1) as filename, table, file_size_bytes, num_rows, num_columns, hits, table_size_bytes from statistics_cache() order by filename"; - let df = ctx.sql(sql).await?; - let rbs = df.collect().await?; - assert_snapshot!(batches_to_string(&rbs),@" - +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ - | filename | table | file_size_bytes | num_rows | num_columns | hits | table_size_bytes | - +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ - | alltypes_plain.parquet | alltypes_plain | 1851 | Exact(8) | 11 | 3 | Absent | - | alltypes_tiny_pages.parquet | alltypes_tiny_pages | 454233 | Exact(7300) | 13 | 0 | Absent | - | lz4_raw_compressed_larger.parquet | lz4_raw_compressed_larger | 380836 | Exact(10000) | 1 | 1 | Absent | - +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ - "); - - Ok(()) - } - - #[tokio::test] - async fn test_list_files_cache() -> Result<(), DataFusionError> { - let list_files_cache = Arc::new(DefaultCache::new_with_ttl( - 1024, - Some(Duration::from_secs(1)), - )); - - let rt = RuntimeEnvBuilder::new() - .with_cache_manager( - CacheManagerConfig::default() - .with_list_files_cache(Some(list_files_cache)), - ) - .build_arc() - .unwrap(); - - let ctx = SessionContext::new_with_config_rt(SessionConfig::default(), rt); - - ctx.register_object_store( - &Url::parse("mem://test_table").unwrap(), - Arc::new(InMemory::new()), - ); - - ctx.register_udtf( - "list_files_cache", - Arc::new(ListFilesCacheFunc::new( - ctx.task_ctx().runtime_env().cache_manager.clone(), - )), - ); - - ctx.sql( - "CREATE EXTERNAL TABLE src_table - STORED AS PARQUET - LOCATION '../parquet-testing/data/alltypes_plain.parquet'", - ) - .await? - .collect() - .await?; - - ctx.sql("COPY (SELECT * FROM src_table) TO 'mem://test_table/0.parquet' STORED AS PARQUET").await?.collect().await?; - - ctx.sql("COPY (SELECT * FROM src_table) TO 'mem://test_table/1.parquet' STORED AS PARQUET").await?.collect().await?; - - ctx.sql( - "CREATE EXTERNAL TABLE test_table - STORED AS PARQUET - LOCATION 'mem://test_table/' - ", - ) - .await? - .collect() - .await?; - - let sql = "SELECT metadata_size_bytes, expires_in, metadata_list FROM list_files_cache()"; - let df = ctx - .sql(sql) - .await? - .unnest_columns(&["metadata_list"])? - .with_column_renamed("metadata_list", "metadata")? - .unnest_columns(&["metadata"])?; - - assert_eq!( - 2, - df.clone() - .filter(col("expires_in").is_not_null())? - .count() - .await? - ); - - let df = df - .with_column_renamed(r#""metadata.file_size_bytes""#, "file_size_bytes")? - .with_column_renamed(r#""metadata.e_tag""#, "etag")? - .with_column( - "filename", - split_part(col(r#""metadata.file_path""#), lit("/"), lit(-1)), - )? - .select_columns(&[ - "metadata_size_bytes", - "filename", - "file_size_bytes", - "etag", - ])? - .sort(vec![col("filename").sort(true, false)])?; - let rbs = df.collect().await?; - assert_snapshot!(batches_to_string(&rbs),@r" - +---------------------+-----------+-----------------+------+ - | metadata_size_bytes | filename | file_size_bytes | etag | - +---------------------+-----------+-----------------+------+ - | 212 | 0.parquet | 3642 | 0 | - | 212 | 1.parquet | 3642 | 1 | - +---------------------+-----------+-----------------+------+ - "); - - Ok(()) - } -} From e050ef91a22c5bb2936e267228159f459382d127 Mon Sep 17 00:00:00 2001 From: matt rice Date: Wed, 8 Jul 2026 04:48:46 -0700 Subject: [PATCH 2/4] rustfmt --- datafusion-cli/src/entry_point.rs | 280 +++++++++++++++--------------- 1 file changed, 140 insertions(+), 140 deletions(-) diff --git a/datafusion-cli/src/entry_point.rs b/datafusion-cli/src/entry_point.rs index 282f6959d7a45..e7e688d3e273b 100644 --- a/datafusion-cli/src/entry_point.rs +++ b/datafusion-cli/src/entry_point.rs @@ -21,14 +21,6 @@ use std::num::NonZeroUsize; use std::path::Path; use std::sync::{Arc, LazyLock}; -use datafusion::error::{DataFusionError, Result}; -use datafusion::execution::context::SessionConfig; -use datafusion::execution::memory_pool::{ - FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool, -}; -use datafusion::execution::runtime_env::RuntimeEnvBuilder; -use datafusion::logical_expr::ExplainFormat; -use datafusion::prelude::SessionContext; use crate::catalog::DynamicObjectStoreCatalog; use crate::functions::{ ListFilesCacheFunc, MetadataCacheFunc, ParquetMetadataFunc, StatisticsCacheFunc, @@ -43,6 +35,14 @@ use crate::{ print_format::PrintFormat, print_options::{MaxRows, PrintOptions}, }; +use datafusion::error::{DataFusionError, Result}; +use datafusion::execution::context::SessionConfig; +use datafusion::execution::memory_pool::{ + FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool, +}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::logical_expr::ExplainFormat; +use datafusion::prelude::SessionContext; use clap::Parser; use datafusion::common::config_err; @@ -186,156 +186,156 @@ impl CliSession { pub fn session_context(&self) -> &SessionContext { &self.ctx } -pub fn try_from_args() -> Result { - env_logger::init(); - let args = Args::parse(); + pub fn try_from_args() -> Result { + env_logger::init(); + let args = Args::parse(); - if !args.quiet { - println!("DataFusion CLI v{DATAFUSION_CLI_VERSION}"); - } + if !args.quiet { + println!("DataFusion CLI v{DATAFUSION_CLI_VERSION}"); + } - if let Some(ref path) = args.data_path { - let p = Path::new(path); - env::set_current_dir(p).unwrap(); - }; + if let Some(ref path) = args.data_path { + let p = Path::new(path); + env::set_current_dir(p).unwrap(); + }; - let session_config = get_session_config(&args)?; + let session_config = get_session_config(&args)?; - let mut rt_builder = RuntimeEnvBuilder::new(); - // set memory pool size - if let Some(memory_limit) = args.memory_limit { - // set memory pool type - let pool: Arc = match args.mem_pool_type { - PoolType::Fair if args.top_memory_consumers == 0 => { - Arc::new(FairSpillPool::new(memory_limit)) - } - PoolType::Fair => Arc::new(TrackConsumersPool::new( - FairSpillPool::new(memory_limit), - NonZeroUsize::new(args.top_memory_consumers).unwrap(), - )), - PoolType::Greedy if args.top_memory_consumers == 0 => { - Arc::new(GreedyMemoryPool::new(memory_limit)) - } - PoolType::Greedy => Arc::new(TrackConsumersPool::new( - GreedyMemoryPool::new(memory_limit), - NonZeroUsize::new(args.top_memory_consumers).unwrap(), - )), - }; + let mut rt_builder = RuntimeEnvBuilder::new(); + // set memory pool size + if let Some(memory_limit) = args.memory_limit { + // set memory pool type + let pool: Arc = match args.mem_pool_type { + PoolType::Fair if args.top_memory_consumers == 0 => { + Arc::new(FairSpillPool::new(memory_limit)) + } + PoolType::Fair => Arc::new(TrackConsumersPool::new( + FairSpillPool::new(memory_limit), + NonZeroUsize::new(args.top_memory_consumers).unwrap(), + )), + PoolType::Greedy if args.top_memory_consumers == 0 => { + Arc::new(GreedyMemoryPool::new(memory_limit)) + } + PoolType::Greedy => Arc::new(TrackConsumersPool::new( + GreedyMemoryPool::new(memory_limit), + NonZeroUsize::new(args.top_memory_consumers).unwrap(), + )), + }; - rt_builder = rt_builder.with_memory_pool(pool) - } + rt_builder = rt_builder.with_memory_pool(pool) + } + + // set disk limit + if let Some(disk_limit) = args.disk_limit { + let builder = DiskManagerBuilder::default() + .with_mode(DiskManagerMode::OsTmpDirectory) + .with_max_temp_directory_size(disk_limit.try_into().unwrap()); + rt_builder = rt_builder.with_disk_manager_builder(builder); + } - // set disk limit - if let Some(disk_limit) = args.disk_limit { - let builder = DiskManagerBuilder::default() - .with_mode(DiskManagerMode::OsTmpDirectory) - .with_max_temp_directory_size(disk_limit.try_into().unwrap()); - rt_builder = rt_builder.with_disk_manager_builder(builder); + let instrumented_registry = Arc::new( + InstrumentedObjectStoreRegistry::new() + .with_profile_mode(args.object_store_profiling), + ); + rt_builder = rt_builder.with_object_store_registry(instrumented_registry.clone()); + + let runtime_env = rt_builder.build_arc()?; + + // enable dynamic file query + let ctx = SessionContext::new_with_config_rt(session_config, runtime_env) + .enable_url_table(); + Ok(Self { + ctx, + instrumented_registry, + args, + }) } - let instrumented_registry = Arc::new( - InstrumentedObjectStoreRegistry::new() - .with_profile_mode(args.object_store_profiling), - ); - rt_builder = rt_builder.with_object_store_registry(instrumented_registry.clone()); - - let runtime_env = rt_builder.build_arc()?; - - // enable dynamic file query - let ctx = SessionContext::new_with_config_rt(session_config, runtime_env) - .enable_url_table(); - Ok(Self { - ctx, - instrumented_registry, - args, - }) -} + /// Main CLI entrypoint + pub async fn run(self) -> Result<()> { + let CliSession { + ctx, + args, + instrumented_registry, + } = self; + ctx.refresh_catalogs().await?; + // install dynamic catalog provider that can register required object stores + ctx.register_catalog_list(Arc::new(DynamicObjectStoreCatalog::new( + ctx.state().catalog_list().clone(), + ctx.state_weak_ref(), + ))); + // register `parquet_metadata` table function to get metadata from parquet files + ctx.register_udtf("parquet_metadata", Arc::new(ParquetMetadataFunc {})); -/// Main CLI entrypoint -pub async fn run(self) -> Result<()> { - let CliSession { - ctx, - args, - instrumented_registry, - } = self; - ctx.refresh_catalogs().await?; - // install dynamic catalog provider that can register required object stores - ctx.register_catalog_list(Arc::new(DynamicObjectStoreCatalog::new( - ctx.state().catalog_list().clone(), - ctx.state_weak_ref(), - ))); - // register `parquet_metadata` table function to get metadata from parquet files - ctx.register_udtf("parquet_metadata", Arc::new(ParquetMetadataFunc {})); - - // register `metadata_cache` table function to get the contents of the file metadata cache - ctx.register_udtf( - "metadata_cache", - Arc::new(MetadataCacheFunc::new( - ctx.task_ctx().runtime_env().cache_manager.clone(), - )), - ); - - // register `statistics_cache` table function to get the contents of the file statistics cache - ctx.register_udtf( - "statistics_cache", - Arc::new(StatisticsCacheFunc::new( - ctx.task_ctx().runtime_env().cache_manager.clone(), - )), - ); - - ctx.register_udtf( - "list_files_cache", - Arc::new(ListFilesCacheFunc::new( - ctx.task_ctx().runtime_env().cache_manager.clone(), - )), - ); - - let mut print_options = PrintOptions { - format: args.format, - quiet: args.quiet, - maxrows: args.maxrows, - color: args.color, - instrumented_registry: Arc::clone(&instrumented_registry), - }; + // register `metadata_cache` table function to get the contents of the file metadata cache + ctx.register_udtf( + "metadata_cache", + Arc::new(MetadataCacheFunc::new( + ctx.task_ctx().runtime_env().cache_manager.clone(), + )), + ); - let repl_mode = args.repl_mode(); - let commands = args.command; - let files = args.file; - let rc = match args.rc { - Some(file) => file, - None => { - let mut files = Vec::new(); - let home = dirs::home_dir(); - if let Some(p) = home { - let home_rc = p.join(".datafusionrc"); - if home_rc.exists() { - files.push(home_rc.into_os_string().into_string().unwrap()); + // register `statistics_cache` table function to get the contents of the file statistics cache + ctx.register_udtf( + "statistics_cache", + Arc::new(StatisticsCacheFunc::new( + ctx.task_ctx().runtime_env().cache_manager.clone(), + )), + ); + + ctx.register_udtf( + "list_files_cache", + Arc::new(ListFilesCacheFunc::new( + ctx.task_ctx().runtime_env().cache_manager.clone(), + )), + ); + + let mut print_options = PrintOptions { + format: args.format, + quiet: args.quiet, + maxrows: args.maxrows, + color: args.color, + instrumented_registry: Arc::clone(&instrumented_registry), + }; + + let repl_mode = args.repl_mode(); + let commands = args.command; + let files = args.file; + let rc = match args.rc { + Some(file) => file, + None => { + let mut files = Vec::new(); + let home = dirs::home_dir(); + if let Some(p) = home { + let home_rc = p.join(".datafusionrc"); + if home_rc.exists() { + files.push(home_rc.into_os_string().into_string().unwrap()); + } } + files + } + }; + + if repl_mode { + if !rc.is_empty() { + exec::exec_from_files(&ctx, rc, &print_options).await?; } - files + // TODO maybe we can have thiserror for cli but for now let's keep it simple + return exec::exec_from_repl(&ctx, &mut print_options) + .await + .map_err(|e| DataFusionError::External(Box::new(e))); } - }; - if repl_mode { - if !rc.is_empty() { - exec::exec_from_files(&ctx, rc, &print_options).await?; + if !files.is_empty() { + exec::exec_from_files(&ctx, files, &print_options).await?; } - // TODO maybe we can have thiserror for cli but for now let's keep it simple - return exec::exec_from_repl(&ctx, &mut print_options) - .await - .map_err(|e| DataFusionError::External(Box::new(e))); - } - if !files.is_empty() { - exec::exec_from_files(&ctx, files, &print_options).await?; - } + if !commands.is_empty() { + exec::exec_from_commands(&ctx, commands, &print_options).await?; + } - if !commands.is_empty() { - exec::exec_from_commands(&ctx, commands, &print_options).await?; + Ok(()) } - - Ok(()) -} } /// Get the session configuration based on the provided arguments From 249072eb8ba3c8899ac846e35f82f7d2ae1b6978 Mon Sep 17 00:00:00 2001 From: matt rice Date: Wed, 8 Jul 2026 12:55:41 -0700 Subject: [PATCH 3/4] Make CliSession usable with the existing example, use thiserror --- Cargo.lock | 1 + datafusion-cli/Cargo.toml | 1 + datafusion-cli/examples/cli-custom-udf.rs | 20 ++---- .../examples/cli-session-context.rs | 36 ++++------ datafusion-cli/src/entry_point.rs | 70 ++++++++++++------- datafusion-cli/src/main.rs | 2 +- 6 files changed, 65 insertions(+), 65 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d3f67f0601371..eb0e91405a8f5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1858,6 +1858,7 @@ dependencies = [ "rustyline", "serde_json", "testcontainers-modules", + "thiserror", "tokio", "url", ] diff --git a/datafusion-cli/Cargo.toml b/datafusion-cli/Cargo.toml index 62eedafe798d4..d9646b8dbb768 100644 --- a/datafusion-cli/Cargo.toml +++ b/datafusion-cli/Cargo.toml @@ -66,6 +66,7 @@ parking_lot = { workspace = true } parquet = { workspace = true, default-features = false } regex = { workspace = true } rustyline = "18.0" +thiserror = "2.0.18" tokio = { workspace = true, features = ["macros", "parking_lot", "rt", "rt-multi-thread", "signal", "sync"] } url = { workspace = true } diff --git a/datafusion-cli/examples/cli-custom-udf.rs b/datafusion-cli/examples/cli-custom-udf.rs index f98ac2d972c6d..9f7699429a4f9 100644 --- a/datafusion-cli/examples/cli-custom-udf.rs +++ b/datafusion-cli/examples/cli-custom-udf.rs @@ -19,9 +19,9 @@ use arrow::datatypes::DataType; use datafusion::arrow::array::{ArrayRef, StringArray}; use datafusion::logical_expr::{ColumnarValue, Volatility, create_udf}; use datafusion::prelude::SessionContext; +use datafusion_cli::entry_point::{CliError, CliSession}; use datafusion_common::cast::as_string_array; use mimalloc::MiMalloc; -use std::process::ExitCode; use std::sync::Arc; #[global_allocator] @@ -34,7 +34,7 @@ static GLOBAL: MiMalloc = MiMalloc; /// 3. Registers the udf function with the `SessionContext` so the user can input `select hello(1)` at the prompt. /// 4. Runs the cli using [`dataframe_cli::CliSession::run`], printing any errors then exits. #[tokio::main] -pub async fn main() -> ExitCode { +pub async fn main() -> Result<(), CliError> { let custom_udf = create_udf( "hello", vec![DataType::Utf8], @@ -51,19 +51,9 @@ pub async fn main() -> ExitCode { Ok(ColumnarValue::from(Arc::new(array) as ArrayRef)) }), ); - let cli_session = match datafusion_cli::entry_point::CliSession::try_from_args() { - Ok(cli) => cli, - Err(e) => { - println!("Error: {e}"); - return ExitCode::FAILURE; - } - }; + let cli_session = CliSession::try_from_args(std::env::args())?; let ctx: &SessionContext = cli_session.session_context(); ctx.register_udf(custom_udf); - if let Err(e) = cli_session.run().await { - println!("Error: {e}"); - return ExitCode::FAILURE; - } - - ExitCode::SUCCESS + cli_session.run().await?; + Ok(()) } diff --git a/datafusion-cli/examples/cli-session-context.rs b/datafusion-cli/examples/cli-session-context.rs index 6095072163870..fe9a6f026ec16 100644 --- a/datafusion-cli/examples/cli-session-context.rs +++ b/datafusion-cli/examples/cli-session-context.rs @@ -18,6 +18,7 @@ //! Shows an example of a custom session context that unions the input plan with itself. //! To run this example, use `cargo run --example cli-session-context` from within the `datafusion-cli` directory. +use std::env; use std::sync::Arc; use datafusion::{ @@ -28,9 +29,9 @@ use datafusion::{ prelude::SessionContext, }; use datafusion_cli::{ - cli_context::CliSessionContext, exec::exec_from_repl, - object_storage::instrumented::InstrumentedObjectStoreRegistry, - print_options::PrintOptions, + cli_context::CliSessionContext, + entry_point::{CliError, CliSession}, + exec::exec_from_repl, }; use object_store::ObjectStore; @@ -39,14 +40,6 @@ struct MyUnionerContext { ctx: SessionContext, } -impl Default for MyUnionerContext { - fn default() -> Self { - Self { - ctx: SessionContext::new(), - } - } -} - #[async_trait::async_trait] impl CliSessionContext for MyUnionerContext { fn task_ctx(&self) -> Arc { @@ -83,16 +76,13 @@ impl CliSessionContext for MyUnionerContext { #[tokio::main] /// Runs the example. -pub async fn main() { - let my_ctx = MyUnionerContext::default(); - - let mut print_options = PrintOptions { - format: datafusion_cli::print_format::PrintFormat::Automatic, - quiet: false, - maxrows: datafusion_cli::print_options::MaxRows::Unlimited, - color: true, - instrumented_registry: Arc::new(InstrumentedObjectStoreRegistry::new()), - }; - - exec_from_repl(&my_ctx, &mut print_options).await.unwrap(); +pub async fn main() -> Result<(), CliError> { + let CliSession { + ctx, + args: _, + mut print_options, + } = CliSession::try_from_args(env::args())?; + let my_ctx = MyUnionerContext { ctx }; + exec_from_repl(&my_ctx, &mut print_options).await?; + Ok(()) } diff --git a/datafusion-cli/src/entry_point.rs b/datafusion-cli/src/entry_point.rs index e7e688d3e273b..9b84cc8cfe779 100644 --- a/datafusion-cli/src/entry_point.rs +++ b/datafusion-cli/src/entry_point.rs @@ -16,10 +16,10 @@ // under the License. use std::collections::HashMap; -use std::env; use std::num::NonZeroUsize; use std::path::Path; use std::sync::{Arc, LazyLock}; +use std::{env, fmt}; use crate::catalog::DynamicObjectStoreCatalog; use crate::functions::{ @@ -35,7 +35,7 @@ use crate::{ print_format::PrintFormat, print_options::{MaxRows, PrintOptions}, }; -use datafusion::error::{DataFusionError, Result}; +use datafusion::error::DataFusionError; use datafusion::execution::context::SessionConfig; use datafusion::execution::memory_pool::{ FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool, @@ -49,16 +49,37 @@ use datafusion::common::config_err; use datafusion::config::ConfigOptions; use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use rustyline::error::ReadlineError; +use thiserror::Error; + +#[derive(Error)] +pub enum CliError { + #[error("DataFusion error `{0}`")] + DataFusion(#[from] DataFusionError), + #[error("Readline error `{0}`")] + Readline(#[from] ReadlineError), + // This is used to print e.g. help text so we don't want to embellish it + #[error(transparent)] + ArgumentParsing(#[from] clap::Error), +} + +impl fmt::Debug for CliError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // We use Display ({}) here instead of Debug ({:?}) for nicer `error?` output. + write!(f, "{self}") + } +} + /// `CliSession` implements argument parsing, and construction of the default `CliSessionContext`. pub struct CliSession { - ctx: SessionContext, - instrumented_registry: Arc, - args: Args, + pub ctx: SessionContext, + pub print_options: PrintOptions, + pub args: Args, } #[derive(Debug, Parser, PartialEq)] #[clap(author, version, about, long_about= None)] -struct Args { +pub struct Args { #[clap( short = 'p', long, @@ -179,16 +200,16 @@ impl Args { } impl CliSession { - pub async fn entry_point() -> Result<()> { - let cli_session = CliSession::try_from_args()?; + pub async fn entry_point() -> Result<(), CliError> { + let cli_session = CliSession::try_from_args(env::args())?; Ok(cli_session.run().await?) } pub fn session_context(&self) -> &SessionContext { &self.ctx } - pub fn try_from_args() -> Result { + pub fn try_from_args(args: impl Iterator) -> Result { env_logger::init(); - let args = Args::parse(); + let args = Args::try_parse_from(args)?; if !args.quiet { println!("DataFusion CLI v{DATAFUSION_CLI_VERSION}"); @@ -244,19 +265,27 @@ impl CliSession { // enable dynamic file query let ctx = SessionContext::new_with_config_rt(session_config, runtime_env) .enable_url_table(); + + let print_options = PrintOptions { + format: args.format, + quiet: args.quiet, + maxrows: args.maxrows, + color: args.color, + instrumented_registry: Arc::clone(&instrumented_registry), + }; Ok(Self { ctx, - instrumented_registry, args, + print_options, }) } /// Main CLI entrypoint - pub async fn run(self) -> Result<()> { + pub async fn run(self) -> Result<(), CliError> { let CliSession { ctx, args, - instrumented_registry, + mut print_options, } = self; ctx.refresh_catalogs().await?; // install dynamic catalog provider that can register required object stores @@ -290,14 +319,6 @@ impl CliSession { )), ); - let mut print_options = PrintOptions { - format: args.format, - quiet: args.quiet, - maxrows: args.maxrows, - color: args.color, - instrumented_registry: Arc::clone(&instrumented_registry), - }; - let repl_mode = args.repl_mode(); let commands = args.command; let files = args.file; @@ -320,10 +341,7 @@ impl CliSession { if !rc.is_empty() { exec::exec_from_files(&ctx, rc, &print_options).await?; } - // TODO maybe we can have thiserror for cli but for now let's keep it simple - return exec::exec_from_repl(&ctx, &mut print_options) - .await - .map_err(|e| DataFusionError::External(Box::new(e))); + exec::exec_from_repl(&ctx, &mut print_options).await?; } if !files.is_empty() { @@ -340,7 +358,7 @@ impl CliSession { /// Get the session configuration based on the provided arguments /// and environment settings. -fn get_session_config(args: &Args) -> Result { +fn get_session_config(args: &Args) -> Result { // Read options from environment variables and merge with command line options let mut config_options = ConfigOptions::from_env()?; diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs index e12f7706a2887..847682a774880 100644 --- a/datafusion-cli/src/main.rs +++ b/datafusion-cli/src/main.rs @@ -24,7 +24,7 @@ use mimalloc::MiMalloc; static GLOBAL: MiMalloc = MiMalloc; #[tokio::main] -/// Calls [`main_inner`], then handles printing errors and returning the correct exit code +/// Calls [`CliSession`], then handles printing errors and returning the correct exit code pub async fn main() -> ExitCode { if let Err(e) = CliSession::entry_point().await { println!("Error: {e}"); From e3d941ac2478c5d198a86fc7b630d8b2964e4e3b Mon Sep 17 00:00:00 2001 From: matt rice Date: Wed, 8 Jul 2026 17:48:54 -0700 Subject: [PATCH 4/4] Fix/add example of composing args with clap --- datafusion-cli/examples/cli-custom-udf.rs | 25 ++++++++++++++++--- .../examples/cli-session-context.rs | 6 ++--- datafusion-cli/src/entry_point.rs | 13 +++++----- 3 files changed, 30 insertions(+), 14 deletions(-) diff --git a/datafusion-cli/examples/cli-custom-udf.rs b/datafusion-cli/examples/cli-custom-udf.rs index 9f7699429a4f9..8dd398850a707 100644 --- a/datafusion-cli/examples/cli-custom-udf.rs +++ b/datafusion-cli/examples/cli-custom-udf.rs @@ -16,10 +16,11 @@ // under the License. use arrow::datatypes::DataType; +use clap::Parser; use datafusion::arrow::array::{ArrayRef, StringArray}; use datafusion::logical_expr::{ColumnarValue, Volatility, create_udf}; use datafusion::prelude::SessionContext; -use datafusion_cli::entry_point::{CliError, CliSession}; +use datafusion_cli::entry_point::{CliError, CliSession, CliArgs}; use datafusion_common::cast::as_string_array; use mimalloc::MiMalloc; use std::sync::Arc; @@ -27,6 +28,19 @@ use std::sync::Arc; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; +#[derive(Debug, Parser, PartialEq)] +#[clap(author, version, about, long_about= None)] +struct CustomArgs { + #[command(flatten)] + cli_args: CliArgs, + #[clap( + long, + help = "Register the hello udf function", + default_value = "true", + )] + register_hello: bool, +} + /// In this example we want to reuse the datafusion-cli binary argument, hen extend the `SessionContext` with custom udf. /// /// 1. Declares a `hello`` udf function. @@ -35,7 +49,7 @@ static GLOBAL: MiMalloc = MiMalloc; /// 4. Runs the cli using [`dataframe_cli::CliSession::run`], printing any errors then exits. #[tokio::main] pub async fn main() -> Result<(), CliError> { - let custom_udf = create_udf( + let hello_udf = create_udf( "hello", vec![DataType::Utf8], DataType::Utf8, @@ -51,9 +65,12 @@ pub async fn main() -> Result<(), CliError> { Ok(ColumnarValue::from(Arc::new(array) as ArrayRef)) }), ); - let cli_session = CliSession::try_from_args(std::env::args())?; + let args = CustomArgs::try_parse()?; + let cli_session = CliSession::try_from_args(args.cli_args)?; let ctx: &SessionContext = cli_session.session_context(); - ctx.register_udf(custom_udf); + if args.register_hello { + ctx.register_udf(hello_udf); + } cli_session.run().await?; Ok(()) } diff --git a/datafusion-cli/examples/cli-session-context.rs b/datafusion-cli/examples/cli-session-context.rs index fe9a6f026ec16..0690fe0baf9a0 100644 --- a/datafusion-cli/examples/cli-session-context.rs +++ b/datafusion-cli/examples/cli-session-context.rs @@ -18,9 +18,9 @@ //! Shows an example of a custom session context that unions the input plan with itself. //! To run this example, use `cargo run --example cli-session-context` from within the `datafusion-cli` directory. -use std::env; use std::sync::Arc; +use clap::Parser; use datafusion::{ dataframe::DataFrame, error::DataFusionError, @@ -30,7 +30,7 @@ use datafusion::{ }; use datafusion_cli::{ cli_context::CliSessionContext, - entry_point::{CliError, CliSession}, + entry_point::{CliError, CliSession, CliArgs}, exec::exec_from_repl, }; use object_store::ObjectStore; @@ -81,7 +81,7 @@ pub async fn main() -> Result<(), CliError> { ctx, args: _, mut print_options, - } = CliSession::try_from_args(env::args())?; + } = CliSession::try_from_args(CliArgs::try_parse()?)?; let my_ctx = MyUnionerContext { ctx }; exec_from_repl(&my_ctx, &mut print_options).await?; Ok(()) diff --git a/datafusion-cli/src/entry_point.rs b/datafusion-cli/src/entry_point.rs index 9b84cc8cfe779..5e9be356d7b56 100644 --- a/datafusion-cli/src/entry_point.rs +++ b/datafusion-cli/src/entry_point.rs @@ -74,12 +74,12 @@ impl fmt::Debug for CliError { pub struct CliSession { pub ctx: SessionContext, pub print_options: PrintOptions, - pub args: Args, + pub args: CliArgs, } #[derive(Debug, Parser, PartialEq)] #[clap(author, version, about, long_about= None)] -pub struct Args { +pub struct CliArgs { #[clap( short = 'p', long, @@ -182,7 +182,7 @@ pub struct Args { object_store_profiling: InstrumentedObjectStoreMode, } -impl Args { +impl CliArgs { /// Without -c/-f the CLI enters the REPL, which reads its SQL from /// stdin — interactively or piped. fn repl_mode(&self) -> bool { @@ -201,15 +201,14 @@ impl Args { impl CliSession { pub async fn entry_point() -> Result<(), CliError> { - let cli_session = CliSession::try_from_args(env::args())?; + let cli_session = CliSession::try_from_args(CliArgs::try_parse()?)?; Ok(cli_session.run().await?) } pub fn session_context(&self) -> &SessionContext { &self.ctx } - pub fn try_from_args(args: impl Iterator) -> Result { + pub fn try_from_args(args: CliArgs) -> Result { env_logger::init(); - let args = Args::try_parse_from(args)?; if !args.quiet { println!("DataFusion CLI v{DATAFUSION_CLI_VERSION}"); @@ -358,7 +357,7 @@ impl CliSession { /// Get the session configuration based on the provided arguments /// and environment settings. -fn get_session_config(args: &Args) -> Result { +fn get_session_config(args: &CliArgs) -> Result { // Read options from environment variables and merge with command line options let mut config_options = ConfigOptions::from_env()?;