nail_parquet/utils/
mod.rs

1pub mod column;
2pub mod format;
3pub mod io;
4pub mod output;
5pub mod parquet_utils;
6pub mod predicate;
7pub mod stats;
8pub mod suggest;
9
10use crate::error::{NailError, NailResult};
11use datafusion::prelude::*;
12use std::path::Path;
13
14const DEFAULT_BATCH_SIZE_LARGE: usize = 32_768;
15const DEFAULT_BATCH_SIZE_JOBS: usize = 8_192;
16
17pub async fn create_context() -> NailResult<SessionContext> {
18	create_context_with_opts(None, None).await
19}
20
21pub async fn create_context_with_jobs(jobs: Option<usize>) -> NailResult<SessionContext> {
22	create_context_with_opts(jobs, None).await
23}
24
25pub async fn create_context_with_opts(
26	jobs: Option<usize>,
27	batch_size: Option<usize>,
28) -> NailResult<SessionContext> {
29	let cpu_count = num_cpus::get();
30	let target_partitions = match jobs {
31		Some(j) => std::cmp::max(1, std::cmp::min(j, cpu_count)),
32		None => std::cmp::max(1, cpu_count),
33	};
34	let effective_batch = batch_size.unwrap_or(if jobs.is_some() {
35		DEFAULT_BATCH_SIZE_JOBS
36	} else {
37		DEFAULT_BATCH_SIZE_LARGE
38	});
39
40	let config = SessionConfig::new()
41		.with_batch_size(effective_batch)
42		.with_target_partitions(target_partitions)
43		.with_collect_statistics(false)
44		.with_parquet_pruning(true)
45		.with_prefer_existing_sort(true);
46
47	Ok(SessionContext::new_with_config(config))
48}
49
50pub fn detect_file_format(path: &Path) -> NailResult<FileFormat> {
51	match path.extension().and_then(|s| s.to_str()) {
52		Some("parquet") => Ok(FileFormat::Parquet),
53		Some("csv") => Ok(FileFormat::Csv),
54		Some("json") => Ok(FileFormat::Json),
55		Some("xlsx") => Ok(FileFormat::Excel),
56		_ => Err(NailError::UnsupportedFormat(format!(
57			"Unable to detect format for file: {}",
58			path.display()
59		))),
60	}
61}
62
63#[derive(Debug, Clone)]
64pub enum FileFormat {
65	Parquet,
66	Csv,
67	Json,
68	Excel,
69}
nail_parquet/utils/mod.rs

nail_parquet/utils/
mod.rs
