use std::path::{Path, PathBuf};
use polars::prelude::*;
use crate::filter_modal::{FilterOperator, FilterStatement, LogicalOperator};
use crate::pivot_melt_modal::PivotAggregation;
use crate::{CompressionFormat, FileFormat, OpenOptions};
pub(crate) fn py_str(s: &str) -> String {
let mut out = String::with_capacity(s.len() + 2);
out.push('"');
for c in s.chars() {
match c {
'\\' => out.push_str("\\\\"),
'"' => out.push_str("\\\""),
'\n' => out.push_str("\\n"),
'\r' => out.push_str("\\r"),
'\t' => out.push_str("\\t"),
c if c.is_control() => out.push_str(&format!("\\u{:04x}", c as u32)),
c => out.push(c),
}
}
out.push('"');
out
}
pub(crate) fn py_comment(text: &str) -> String {
let mut out = String::with_capacity(text.len() + 2);
out.push_str("# ");
for c in text.chars() {
match c {
'\n' => out.push_str("\\n"),
'\r' => out.push_str("\\r"),
'\t' => out.push('\t'),
c if c.is_control() || c == '\u{2028}' || c == '\u{2029}' => {
out.push_str(&format!("\\u{:04x}", c as u32))
}
c => out.push(c),
}
}
out
}
pub(crate) fn py_float(f: f64) -> String {
if f.is_nan() {
"float(\"nan\")".to_string()
} else if f.is_infinite() {
if f > 0.0 {
"float(\"inf\")".to_string()
} else {
"float(\"-inf\")".to_string()
}
} else {
format!("{f:?}")
}
}
pub(crate) fn py_bool(b: bool) -> &'static str {
if b { "True" } else { "False" }
}
pub(crate) fn py_names(names: &[String]) -> String {
let items: Vec<String> = names.iter().map(|n| py_str(n)).collect();
format!("[{}]", items.join(", "))
}
pub(crate) fn sort_call(columns: &[String], descending: &[bool]) -> String {
let by = match columns {
[one] => py_str(one),
_ => py_names(columns),
};
let descending = if descending.iter().all(|d| !d) {
String::new()
} else if descending.iter().all(|d| *d) {
"descending=True, ".to_string()
} else {
let flags: Vec<&str> = descending.iter().map(|d| py_bool(*d)).collect();
format!("descending=[{}], ", flags.join(", "))
};
format!(".sort({by}, {descending}nulls_last=True, maintain_order=True)")
}
#[derive(Debug, Clone, PartialEq)]
pub enum FilterValue {
Typed(Scalar),
Str(String),
}
impl FilterValue {
fn lit(&self) -> Expr {
match self {
FilterValue::Typed(scalar) => lit(scalar.clone()),
FilterValue::Str(s) => lit(s.as_str()),
}
}
fn python(&self) -> String {
match self {
FilterValue::Typed(scalar) => crate::typed_value::python(scalar),
FilterValue::Str(s) => py_str(s),
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct SidebarFilter {
pub column: String,
pub operator: FilterOperator,
pub value: FilterValue,
pub text: String,
pub logical_op: LogicalOperator,
pub searched: Vec<(String, DataType)>,
}
impl SidebarFilter {
pub fn typed_in(statement: &FilterStatement, schema: &Schema, shown: &[String]) -> Self {
let mut filter = Self::typed(statement, schema.get(&statement.column));
if statement.operator.is_find() {
let spec = filter.find_spec();
let names: Vec<&String> = if statement.column == crate::filter_modal::ANY_COLUMN {
if statement.columns.is_empty() {
shown.iter().collect()
} else {
statement.columns.iter().collect()
}
} else {
vec![&statement.column]
};
filter.searched = names
.into_iter()
.filter_map(|name| Some((name.clone(), schema.get(name)?.clone())))
.filter(|(name, dtype)| crate::find::cell_matches(&spec, name, dtype).is_some())
.collect();
}
filter
}
pub fn unscriptable_columns(&self) -> Vec<String> {
if !self.operator.is_find() {
return Vec::new();
}
self.searched
.iter()
.filter(|(_, dtype)| matches!(dtype, DataType::Duration(_)))
.map(|(name, _)| name.clone())
.collect()
}
fn find_spec(&self) -> crate::find::FindSpec {
crate::find::FindSpec {
pattern: self.text.clone(),
regex: self.operator == FilterOperator::HasRegex,
fuzzy: self.operator == FilterOperator::HasFuzzy,
column: None,
}
}
pub fn typed(statement: &FilterStatement, dtype: Option<&DataType>) -> Self {
let text = statement.value.as_str();
let value = match dtype {
None | Some(DataType::String) => FilterValue::Str(text.to_string()),
Some(dtype) => crate::typed_value::parse(text, dtype)
.map(FilterValue::Typed)
.unwrap_or_else(|_| FilterValue::Str(text.to_string())),
};
Self {
column: statement.column.clone(),
operator: statement.operator,
value,
text: statement.value.clone(),
logical_op: statement.logical_op,
searched: Vec::new(),
}
}
pub fn problem(statement: &FilterStatement, dtype: Option<&DataType>) -> Option<String> {
if statement.operator == FilterOperator::HasRegex
&& let Err(reason) = Self::typed(statement, dtype).find_spec().check()
{
return Some(reason);
}
if statement.operator.is_find()
&& let Some(dtype) = dtype
&& crate::find::cell_matches(
&Self::typed(statement, Some(dtype)).find_spec(),
&statement.column,
dtype,
)
.is_none()
{
return Some(format!("{}: no text to match", statement.column));
}
let compares = statement.operator.takes_value()
&& !statement.operator.is_find()
&& !matches!(
statement.operator,
FilterOperator::Contains | FilterOperator::NotContains
);
let dtype = dtype.filter(|_| compares)?;
crate::typed_value::parse(&statement.value, dtype)
.err()
.map(|why| format!("{}: {why}", statement.column))
}
fn expr(&self) -> Expr {
let column = col(&self.column);
let contains = || {
col(&self.column)
.str()
.contains_literal(lit(self.text.as_str()))
};
match self.operator {
FilterOperator::Eq => column.eq(self.value.lit()),
FilterOperator::NotEq => column.neq(self.value.lit()),
FilterOperator::Gt => column.gt(self.value.lit()),
FilterOperator::Lt => column.lt(self.value.lit()),
FilterOperator::GtEq => column.gt_eq(self.value.lit()),
FilterOperator::LtEq => column.lt_eq(self.value.lit()),
FilterOperator::Contains => contains(),
FilterOperator::NotContains => contains().not(),
FilterOperator::IsNull => column.is_null(),
FilterOperator::IsNotNull => column.is_not_null(),
FilterOperator::Has | FilterOperator::HasRegex | FilterOperator::HasFuzzy => {
let spec = self.find_spec();
let cells: Vec<Expr> = self
.searched
.iter()
.filter_map(|(name, dtype)| crate::find::cell_matches(&spec, name, dtype))
.collect();
cells
.into_iter()
.reduce(|a, b| a.or(b))
.unwrap_or(lit(false))
}
}
}
fn python(&self) -> String {
let column = format!("pl.col({})", py_str(&self.column));
let op = match self.operator {
FilterOperator::Eq => "==",
FilterOperator::NotEq => "!=",
FilterOperator::Gt => ">",
FilterOperator::Lt => "<",
FilterOperator::GtEq => ">=",
FilterOperator::LtEq => "<=",
FilterOperator::Contains | FilterOperator::NotContains => {
let not = if self.operator == FilterOperator::NotContains {
"~"
} else {
""
};
return format!(
"{not}{column}.str.contains({}, literal=True)",
py_str(&self.text)
);
}
FilterOperator::IsNull => return format!("{column}.is_null()"),
FilterOperator::IsNotNull => return format!("{column}.is_not_null()"),
FilterOperator::Has | FilterOperator::HasRegex | FilterOperator::HasFuzzy => {
let spec = self.find_spec();
let cells: Vec<String> = self
.searched
.iter()
.filter(|(_, dtype)| !matches!(dtype, DataType::Duration(_)))
.map(|(name, _)| {
let column = format!("pl.col({}).cast(pl.String)", py_str(name));
match spec.regex_source() {
None => format!(
"{column}.str.contains({}, literal=True)",
py_str(&spec.pattern)
),
Some(source) => {
format!("{column}.str.contains({})", py_str(&source))
}
}
})
.collect();
return match cells.len() {
0 => "pl.lit(False)".to_string(),
1 => format!("{}.fill_null(False)", cells[0]),
_ => format!("pl.any_horizontal({}).fill_null(False)", cells.join(", ")),
};
}
};
format!("{column} {op} {}", self.value.python())
}
}
pub fn filters_expr(filters: &[SidebarFilter]) -> Option<Expr> {
filters.iter().fold(None, |all, f| {
Some(match all {
None => f.expr(),
Some(all) => match f.logical_op {
LogicalOperator::And => all.and(f.expr()),
LogicalOperator::Or => all.or(f.expr()),
},
})
})
}
fn filters_python(filters: &[SidebarFilter]) -> String {
let mut out = String::new();
let mut last: Option<LogicalOperator> = None;
for (i, f) in filters.iter().enumerate() {
let term = format!("({})", f.python());
if i == 0 {
out = if filters.len() == 1 { f.python() } else { term };
continue;
}
if last.is_some_and(|l| l != f.logical_op) {
out = format!("({out})");
}
let op = match f.logical_op {
LogicalOperator::And => "&",
LogicalOperator::Or => "|",
};
out = format!("{out} {op} {term}");
last = Some(f.logical_op);
}
out
}
#[derive(Debug, Clone, PartialEq)]
pub enum Step {
Query {
query: String,
input: SchemaRef,
keys: Vec<String>,
},
QueryRows {
query: String,
input: SchemaRef,
},
Sql {
sql: String,
ordered_by: Vec<String>,
},
Search {
patterns: Vec<String>,
columns: Vec<String>,
},
Filter(Vec<SidebarFilter>),
Sort {
columns: Vec<String>,
descending: Vec<bool>,
},
Reverse,
Select(Vec<String>),
Drop(Vec<String>),
Pivot {
index: Vec<String>,
on: String,
values: String,
aggregation: PivotAggregation,
},
Melt {
index: Vec<String>,
on: Vec<String>,
variable_name: String,
value_name: String,
},
Matching(Vec<(String, String)>),
Unreproducible(String),
}
impl Step {
fn python(&self) -> Vec<String> {
match self {
Step::Query { query, input, keys } => match crate::query::parse_nodes(query) {
Ok(mut nodes) => {
nodes.resolve_division(input);
nodes.resolve_time_zones(input);
nodes.python_steps(keys)
}
Err(e) => vec![py_comment(&format!("the query did not parse: {e}"))],
},
Step::QueryRows { query, input } => match crate::query::parse_nodes(query) {
Ok(mut nodes) => {
nodes.resolve_division(input);
nodes.resolve_time_zones(input);
nodes.python_filter().into_iter().collect()
}
Err(e) => vec![py_comment(&format!("the query did not parse: {e}"))],
},
Step::Sql { sql, ordered_by } => {
let sql = sql.trim();
let verbatim = sql.contains('\n')
&& !sql.contains("\"\"\"")
&& !sql.ends_with('"')
&& sql
.chars()
.all(|c| c == '\n' || c == '\t' || (c != '\\' && !c.is_control()));
let mut lines = if verbatim {
vec![
".sql(".to_string(),
format!(" \"\"\"{sql}\"\"\","),
" table_name=\"df\",".to_string(),
")".to_string(),
]
} else {
vec![format!(".sql({}, table_name=\"df\")", py_str(sql))]
};
if !ordered_by.is_empty() {
lines.push(sort_call(ordered_by, &vec![false; ordered_by.len()]));
}
lines
}
Step::Search { patterns, columns } => {
let terms: Vec<String> = patterns
.iter()
.map(|p| {
let any: Vec<String> = columns
.iter()
.map(|c| {
format!(
"pl.col({}).str.contains({}, strict=False)",
py_str(c),
py_str(p)
)
})
.collect();
if any.len() == 1 || patterns.len() == 1 {
any.join(" | ")
} else {
format!("({})", any.join(" | "))
}
})
.collect();
vec![format!(".filter({})", terms.join(" & "))]
}
Step::Filter(filters) => vec![format!(".filter({})", filters_python(filters))],
Step::Sort {
columns,
descending,
} => vec![sort_call(columns, descending)],
Step::Reverse => vec![".reverse()".to_string()],
Step::Select(columns) => vec![format!(".select({})", py_names(columns))],
Step::Drop(columns) => vec![format!(".drop({})", py_names(columns))],
Step::Pivot {
index,
on,
values,
aggregation,
} => {
let agg = match aggregation {
PivotAggregation::Last => "last()",
PivotAggregation::First => "first()",
PivotAggregation::Min => "min()",
PivotAggregation::Max => "max()",
PivotAggregation::Avg => "mean()",
PivotAggregation::Med => "median()",
PivotAggregation::Std => "std()",
PivotAggregation::Count => "len()",
};
let cell = match aggregation {
PivotAggregation::Count => "sum",
_ => "first",
};
let keys: Vec<String> = index.iter().chain([on]).cloned().collect();
vec![
format!(".group_by({}, maintain_order=True)", py_names(&keys)),
format!(".agg(pl.col({}).{agg})", py_str(values)),
".collect()".to_string(),
".pipe(".to_string(),
" lambda cells: cells.pivot(".to_string(),
format!(" on={},", py_str(on)),
format!(
" on_columns=cells[{}].unique().sort(nulls_last=True),",
py_str(on)
),
format!(" index={},", py_names(index)),
format!(" values={},", py_str(values)),
format!(" aggregate_function={},", py_str(cell)),
" )".to_string(),
")".to_string(),
".lazy()".to_string(),
]
}
Step::Melt {
index,
on,
variable_name,
value_name,
} => vec![format!(
".unpivot(on={}, index={}, variable_name={}, value_name={})",
py_names(on),
py_names(index),
py_str(variable_name),
py_str(value_name)
)],
Step::Matching(keys) => {
let terms: Vec<String> = keys
.iter()
.map(|(key, value)| format!("{key}.eq_missing({value})"))
.collect();
let terms = if terms.len() == 1 {
terms
} else {
terms.into_iter().map(|t| format!("({t})")).collect()
};
vec![format!(".filter({})", terms.join(" & "))]
}
Step::Unreproducible(what) => vec![py_comment(what)],
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum Source {
Read {
call: String,
after: Vec<String>,
notes: Vec<String>,
imports: Vec<&'static str>,
},
Placeholder { what: String },
}
pub struct OpenRecord<'a> {
pub paths: Option<&'a [PathBuf]>,
pub options: &'a OpenOptions,
pub format: Option<FileFormat>,
pub read_mode: Option<crate::ReadMode>,
pub schema: &'a Schema,
pub remote_objects: Vec<String>,
pub s3_endpoint: Option<String>,
pub s3_region: Option<String>,
pub unsigned: bool,
pub read_as_text: Vec<String>,
pub spec: Option<String>,
}
fn is_url(path: &Path) -> bool {
crate::source::is_remote_url(path)
}
fn without_secrets(url: &str) -> (String, bool) {
let url = &*crate::source::split_source_id(url).1;
let Some(scheme_end) = url.find("://").map(|i| i + 3) else {
return (url.to_string(), false);
};
let (scheme, rest) = url.split_at(scheme_end);
let host_end = rest.find(['/', '?', '#']).unwrap_or(rest.len());
let (authority, path) = rest.split_at(host_end);
let azure = ["abfs://", "abfss://"]
.iter()
.any(|s| scheme.eq_ignore_ascii_case(s));
let host = match authority.rsplit_once('@') {
Some((_, host)) if !azure => host,
_ => authority,
};
let http = scheme.eq_ignore_ascii_case("http://") || scheme.eq_ignore_ascii_case("https://");
let path = match path.find(['?', '#']) {
Some(i) if http => &path[..i],
_ => path,
};
let kept = format!("{scheme}{host}{path}");
let cut = kept != url;
(kept, cut)
}
fn file_format(path: &Path, record: &OpenRecord) -> Option<FileFormat> {
record.format.or(record.options.format).or_else(|| {
FileFormat::from_path(path).or_else(|| {
CompressionFormat::from_extension(path)
.and_then(|_| path.file_stem())
.and_then(|stem| FileFormat::from_path(Path::new(stem)))
})
})
}
fn commonest_format<'a>(names: impl Iterator<Item = &'a str>) -> Option<FileFormat> {
let mut counts: Vec<(FileFormat, usize)> = Vec::new();
for name in names {
if let Some(format) = FileFormat::from_path(Path::new(name)) {
match counts.iter_mut().find(|(f, _)| *f == format) {
Some((_, n)) => *n += 1,
None => counts.push((format, 1)),
}
}
}
counts.into_iter().max_by_key(|(_, n)| *n).map(|(f, _)| f)
}
struct Target {
text: String,
format: FileFormat,
below: bool,
pattern: bool,
literal: bool,
}
fn reader_target(path: &Path, record: &OpenRecord) -> Option<Target> {
let text = path.to_string_lossy().to_string();
if is_url(path) {
let (text, _) = without_secrets(&text);
if let Some(format) = file_format(Path::new(&text), record) {
let pattern = crate::source::has_glob_chars(Path::new(&text));
return Some(Target {
text,
format,
below: false,
pattern,
literal: false,
});
}
let format = record
.format
.or(record.options.format)
.or_else(|| commonest_format(record.remote_objects.iter().map(String::as_str)))?;
let base = text.trim_end_matches('/');
let ext = format_extension(format)?;
return Some(Target {
text: format!("{base}/**/*.{ext}"),
format,
below: true,
pattern: true,
literal: false,
});
}
if path.is_dir() {
let entries: Vec<std::fs::DirEntry> = std::fs::read_dir(path).ok()?.flatten().collect();
let mut names: Vec<String> = entries
.iter()
.filter(|e| e.path().is_file())
.map(|e| e.file_name().to_string_lossy().to_string())
.collect();
if names
.iter()
.any(|n| FileFormat::from_path(Path::new(n)) == Some(FileFormat::Arrow))
{
names.retain(|n| !crate::discover::is_hugging_face_metadata(n));
}
let has_dirs = entries.iter().any(|e| e.path().is_dir());
let format = record
.format
.or(record.options.format)
.or_else(|| commonest_format(names.iter().map(String::as_str)));
let base = crate::source::escape_glob(text.trim_end_matches(['/', '\\']));
let (text, format, below) = match format {
Some(FileFormat::Parquet) | None if has_dirs => {
(format!("{base}/**/*.parquet"), FileFormat::Parquet, true)
}
Some(format) => (
format!("{base}/*.{}", format_extension(format)?),
format,
false,
),
None => return None,
};
return Some(Target {
text,
format,
below,
pattern: true,
literal: false,
});
}
let format = file_format(path, record)?;
let pattern = crate::source::expands_as_glob(path);
Some(Target {
literal: !pattern && scans_by_pattern(format) && crate::source::has_glob_chars(path),
text,
format,
below: false,
pattern,
})
}
fn scans_by_pattern(format: FileFormat) -> bool {
python_of(format).is_some_and(|python| !python.eager)
}
fn hugging_face_files(paths: &[PathBuf], table: Option<&str>) -> Option<Vec<String>> {
let [path] = paths else {
return None;
};
if is_url(path) || !path.is_dir() {
return None;
}
let dict_split = crate::hf_splits::dataset_dict(path).and_then(|splits| {
let listed: Vec<&str> = splits.iter().map(String::as_str).collect();
crate::hf_splits::pick(&listed, table).ok()?.split
});
let dict = dict_split.is_some();
let table = if dict { None } else { table };
let path = &dict_split.map_or_else(|| path.clone(), |split| path.join(split));
let mut inside: Vec<PathBuf> = std::fs::read_dir(path)
.ok()?
.flatten()
.map(|e| e.path())
.filter(|p| p.is_file() && FileFormat::from_path(p) == Some(FileFormat::Arrow))
.collect();
inside.sort();
let cache = ["dataset_info.json", "state.json"]
.iter()
.any(|name| path.join(name).is_file());
if !cache && !dict {
return None;
}
if cache {
let names: Vec<&str> = inside
.iter()
.map(|f| f.file_name().and_then(|n| n.to_str()).unwrap_or_default())
.collect();
let (chosen, _) = crate::hf_splits::choose(&names, table).ok()?;
inside = chosen.into_iter().map(|i| inside[i].clone()).collect();
}
Some(
inside
.iter()
.map(|p| p.to_string_lossy().to_string())
.collect(),
)
}
fn arrow_read(inputs: &[(String, bool)], extra: Option<&str>) -> String {
let extra = extra.map(|e| format!(", {e}")).unwrap_or_default();
let names: Vec<String> = inputs.iter().map(|(name, _)| name.clone()).collect();
if inputs.iter().all(|(_, stream)| *stream) {
return match names.as_slice() {
[one] => format!("pl.read_ipc_stream({}{extra}).lazy()", py_str(one)),
many => format!(
"pl.concat([pl.read_ipc_stream(f{extra}) for f in {}]).lazy()",
py_names(many)
),
};
}
if inputs.iter().all(|(_, stream)| !*stream) {
return match names.as_slice() {
[one] => format!("pl.scan_ipc({}{extra})", py_str(one)),
many => format!("pl.scan_ipc({}{extra})", py_names(many)),
};
}
let reads: Vec<String> = inputs
.iter()
.map(|(name, stream)| match stream {
true => format!("pl.read_ipc_stream({}{extra}).lazy()", py_str(name)),
false => format!("pl.scan_ipc({}{extra})", py_str(name)),
})
.collect();
format!(
"pl.concat([{}], how=\"diagonal_relaxed\")",
reads.join(", ")
)
}
fn format_extension(format: FileFormat) -> Option<&'static str> {
python_of(format)?;
let d = format.descriptor();
(d.many_files || format.separator().is_some())
.then(|| d.extensions.first().copied())
.flatten()
}
pub(crate) struct Python {
pub call: &'static str,
pub eager: bool,
pub glob_flag: bool,
pub arguments: Option<fn(&mut Call<'_>) -> Option<Source>>,
}
pub(crate) struct Call<'a> {
pub record: &'a OpenRecord<'a>,
pub paths: &'a [PathBuf],
pub format: FileFormat,
pub names: &'a [String],
pub below: bool,
pub storage: Option<String>,
pub args: Vec<String>,
pub after: Vec<String>,
pub skip_tail: Option<String>,
pub notes: Vec<String>,
}
impl Call<'_> {
fn storage_for(&self, names: impl IntoIterator<Item = impl AsRef<str>>) -> Option<String> {
names
.into_iter()
.any(|n| store_scheme(n.as_ref()).is_some())
.then(|| self.storage.clone())
.flatten()
}
}
fn store_scheme(name: &str) -> Option<&'static str> {
let (scheme, _) = name.split_once("://")?;
match scheme.to_ascii_lowercase().as_str() {
"s3" | "s3a" => Some("s3"),
"gs" | "gcs" => Some("gs"),
"az" | "adl" | "azure" | "abfs" | "abfss" => Some("azure"),
_ => None,
}
}
fn storage_options(name: &str, record: &OpenRecord, endpoint: Option<&str>) -> Option<String> {
let mut pairs: Vec<(&str, String)> = Vec::new();
match store_scheme(name)? {
"s3" => {
pairs.extend(endpoint.map(|e| ("aws_endpoint_url", e.to_string())));
pairs.extend(record.s3_region.clone().map(|r| ("aws_region", r)));
}
"azure" => {
pairs.extend(
crate::source::azure_parts(name).map(|(account, ..)| ("account_name", account)),
);
}
_ => {}
}
if record.unsigned {
pairs.push(("skip_signature", "true".to_string()));
}
let pairs: Vec<String> = pairs
.iter()
.map(|(k, v)| format!("{}: {}", py_str(k), py_str(v)))
.collect();
(!pairs.is_empty()).then(|| format!("storage_options={{{}}}", pairs.join(", ")))
}
pub(crate) fn ndjson_arguments(call: &mut Call<'_>) -> Option<Source> {
if let Some(s) = call.storage_for(call.names) {
call.args.push(s);
}
None
}
fn python_of(format: FileFormat) -> Option<&'static Python> {
crate::readers::of(format).python.as_ref()
}
pub(crate) fn parquet_arguments(call: &mut Call<'_>) -> Option<Source> {
if call.record.options.hive || call.below {
call.args.push("hive_partitioning=True".to_string());
}
if let Some(s) = call.storage_for(call.names) {
call.args.push(s);
}
None
}
pub(crate) fn csv_arguments(call: &mut Call<'_>) -> Option<Source> {
let options = call.record.options;
let names = call.names.join(", ");
let comment = options.comment_char.as_deref().filter(|c| !c.is_empty());
let dialect: Vec<&str> = [
(comment.is_some_and(|c| c.len() > 5), "--comment"),
(options.header_rows().is_some(), "--header-rows"),
(options.skip_initial_space, "--skip-initial-space"),
]
.into_iter()
.filter_map(|(set, flag)| set.then_some(flag))
.collect();
if !dialect.is_empty() {
return Some(Source::Placeholder {
what: format!(
"{names}: datui read it with {}, which it cannot write as Python: load it here.",
dialect.join(", ")
),
});
}
let separator = options
.delimiter
.or_else(|| call.format.separator())
.unwrap_or(b',');
let args = &mut call.args;
if separator != b',' {
args.push(format!(
"separator={}",
py_str(&(separator as char).to_string())
));
}
if let Some(prefix) = comment {
args.push(format!("comment_prefix={}", py_str(prefix)));
}
if options.has_header == Some(false) {
args.push("has_header=False".to_string());
}
if let Some(n) = options.skip_lines {
args.push(format!("skip_lines={n}"));
}
if let Some(n) = options.skip_rows {
args.push(format!("skip_rows={n}"));
}
if let Some(n) = options.infer_schema_length {
args.push(format!("infer_schema_length={n}"));
}
if options.ignore_errors {
args.push("ignore_errors=True".to_string());
}
let bucket_prefix = call.below && call.paths.iter().any(|p| is_url(p));
if options.csv_try_parse_dates() && !bucket_prefix {
call.args.push("try_parse_dates=True".to_string());
}
if let Some(nulls) = csv_null_values(options, call.record.schema) {
call.args.push(format!("null_values={nulls}"));
}
if !options.typing.text.is_empty() {
let text: Vec<String> = options
.typing
.text
.iter()
.map(|name| format!("{}: pl.String", py_str(name)))
.collect();
call.args
.push(format!("schema_overrides={{{}}}", text.join(", ")));
}
if let Some(s) = call.storage_for(call.names) {
call.args.push(s);
}
if let Some(n) = options.skip_tail_rows.filter(|n| *n > 0) {
call.skip_tail = Some(format!(".filter(pl.int_range(pl.len()) < pl.len() - {n})"));
}
match options.compression.or_else(|| {
call.paths
.first()
.and_then(|p| CompressionFormat::from_extension(p))
}) {
Some(CompressionFormat::Bzip2 | CompressionFormat::Xz) => Some(Source::Placeholder {
what: format!(
"{names}: Polars cannot read bzip2 or xz; decompress it and read it with pl.scan_csv."
),
}),
_ => None,
}
}
pub(crate) fn arrow_arguments(call: &mut Call<'_>) -> Option<Source> {
let options = call.record.options;
let inputs: Option<Vec<(String, bool)>> = match &options.arrow_parts {
Some(parts) => Some(
parts
.iter()
.map(|part| match part {
crate::ipc_stream::Part::InPlace(p) => (p, false),
crate::ipc_stream::Part::Converted { source, .. } => (source, true),
})
.map(|(p, stream)| (without_secrets(&p.to_string_lossy()).0, stream))
.collect(),
),
None => hugging_face_files(call.paths, options.table.as_deref())
.map(|files| files.into_iter().map(|f| (f, false)).collect()),
};
match inputs {
Some(inputs) => {
let extra = call.storage_for(inputs.iter().map(|(name, _)| name));
let read = arrow_read(&inputs, extra.as_deref());
let mut after = std::mem::take(&mut call.after);
after.extend(call.skip_tail.take());
Some(Source::Read {
call: read,
after,
notes: std::mem::take(&mut call.notes),
imports: Vec::new(),
})
}
None => {
if let Some(s) = call.storage_for(call.names) {
call.args.push(s);
}
None
}
}
}
pub(crate) fn excel_arguments(call: &mut Call<'_>) -> Option<Source> {
if let Some(sheet) = &call.record.options.table {
match sheet.parse::<usize>() {
Ok(i) => call.args.push(format!("sheet_id={}", i + 1)),
Err(_) => call.args.push(format!("sheet_name={}", py_str(sheet))),
}
}
call.notes.push(
"datui types a worksheet's columns itself; Polars may read some differently.".to_string(),
);
None
}
fn sql_ident(name: &str) -> String {
format!("\"{}\"", name.replace('"', "\"\""))
}
fn whole(call: &mut Call<'_>, read: String, imports: Vec<&'static str>) -> Source {
call.notes.extend(read_whole_note(call.record, &read));
let mut after = std::mem::take(&mut call.after);
after.extend(call.skip_tail.take());
Source::Read {
call: read,
after,
notes: std::mem::take(&mut call.notes),
imports,
}
}
fn read_whole_note(record: &OpenRecord, call: &str) -> Option<String> {
let name = call.split('(').next().unwrap_or(call);
(record.read_mode == Some(crate::ReadMode::Lazy)).then(|| {
format!(
"Read: {} in datui; {name} reads the file whole into memory.",
crate::ReadMode::Lazy.label()
)
})
}
pub(crate) fn sqlite_arguments(call: &mut Call<'_>) -> Option<Source> {
let [file] = call.names else {
return None;
};
let Some(table) = call.record.options.table.as_deref() else {
return Some(Source::Placeholder {
what: format!("{file}: datui could not tell which table it read; load it here."),
});
};
if call.paths.iter().any(|p| is_url(p)) {
return Some(Source::Placeholder {
what: format!(
"{file} --table {table}: sqlite3 opens a local file; download it and read it \
with pl.read_database."
),
});
}
call.notes.push(
"datui types a table's columns from their declared types; Polars infers them from \
the values."
.to_string(),
);
let query = format!("SELECT * FROM {}", sql_ident(table));
let read = format!(
"pl.read_database({}, sqlite3.connect({})).lazy()",
py_str(&query),
py_str(file)
);
Some(whole(call, read, vec!["import sqlite3"]))
}
pub(crate) fn numpy_arguments(call: &mut Call<'_>) -> Option<Source> {
let [file] = call.names else {
return None;
};
let table = call.record.options.table.as_deref();
let archive = table.is_some() || file.to_ascii_lowercase().ends_with(".npz");
let array = match table {
Some(name) => format!("np.load({})[{}]", py_str(file), py_str(name)),
None if archive => format!("next(iter(np.load({}).values()))", py_str(file)),
None => format!("np.load({})", py_str(file)),
};
let names: Vec<String> = call
.record
.schema
.iter_names()
.map(|n| n.to_string())
.collect();
let schema = if names.iter().any(|n| n.contains('.')) {
call.notes.push(
"datui names a nested field's columns outer.inner; Polars keeps the field as a struct."
.to_string(),
);
String::new()
} else {
format!(", schema={}", py_names(&names))
};
let read = format!("pl.from_numpy({array}{schema}, orient=\"row\").lazy()");
Some(whole(call, read, vec!["import numpy as np"]))
}
fn named_with_table(names: &[String], record: &OpenRecord) -> String {
let names = names.join(", ");
match record.options.table.as_deref() {
Some(table) => format!("{names} --table {table}"),
None => names,
}
}
pub(crate) fn lines_arguments(call: &mut Call<'_>) -> Option<Source> {
let names = call.names.join(", ");
let path = match call.paths {
[one]
if !is_url(one)
&& !one.is_dir()
&& CompressionFormat::from_extension(one).is_none() =>
{
one
}
_ => {
return Some(Source::Placeholder {
what: format!("{names}: datui read it as lines; load it here."),
});
}
};
let read = format!(
"pl.LazyFrame({{\"line\": open({}, encoding=\"utf-8\", errors=\"replace\", newline=\"\").read().removesuffix(\"\\n\").split(\"\\n\")}})",
py_str(&path.to_string_lossy())
);
let mut after = vec![".with_columns(pl.col(\"line\").str.strip_suffix(\"\\r\"))".to_string()];
after.append(&mut call.after);
Some(Source::Read {
call: read,
after,
notes: std::mem::take(&mut call.notes),
imports: Vec::new(),
})
}
pub fn source(record: &OpenRecord) -> Source {
let Some(paths) = record.paths else {
return Source::Placeholder {
what: "The data datui was handed: load it here as a DataFrame or LazyFrame."
.to_string(),
};
};
let teed;
let paths = match (&record.options.tee, paths) {
(Some(tee), [one]) if crate::stdin::is_stdin(one) => {
teed = [tee.clone()];
&teed[..]
}
_ => paths,
};
if paths.iter().any(|p| crate::stdin::is_stdin(p)) {
return Source::Placeholder {
what: "The data datui read from standard input: load it here.".to_string(),
};
}
let spec = record.spec.clone().or_else(|| {
let options = record.options;
options
.spec_name
.clone()
.or_else(|| options.spec_file.as_ref().map(|f| f.display().to_string()))
});
if let Some(spec) = spec {
return Source::Placeholder {
what: format!(
"datui read this through the format spec {spec}, which it cannot write as \
Python: load it here."
),
};
}
let targets: Option<Vec<Target>> = paths.iter().map(|p| reader_target(p, record)).collect();
let Some(targets) = targets.filter(|t| !t.is_empty()) else {
let names: Vec<String> = paths
.iter()
.map(|p| without_secrets(&p.to_string_lossy()).0)
.collect();
return Source::Placeholder {
what: format!(
"{}: datui could not name a Polars reader for this data; load it here.",
named_with_table(&names, record)
),
};
};
let format = targets[0].format;
if targets.iter().any(|t| t.format != format) {
return Source::Placeholder {
what: "The files are of more than one format: load them here.".to_string(),
};
}
let below = targets.iter().any(|t| t.below);
let python = python_of(format);
let literal = targets.iter().any(|t| t.literal);
let no_glob = literal
&& python.is_some_and(|python| python.glob_flag)
&& !targets.iter().any(|t| t.pattern);
let names: Vec<String> = targets
.into_iter()
.map(|t| {
if t.literal && !no_glob {
crate::source::escape_glob(&t.text)
} else {
t.text
}
})
.collect();
let Some(python) = python else {
return Source::Placeholder {
what: format!(
"{}: Polars has no reader for {} files; load it here.",
named_with_table(&names, record),
format.title()
),
};
};
let target = match names.as_slice() {
[one] => py_str(one),
many => py_names(many),
};
let options = record.options;
let mut args: Vec<String> = vec![target];
if no_glob {
args.push("glob=False".to_string());
}
let mut notes = Vec::new();
let endpoint = record.s3_endpoint.as_deref().map(without_secrets);
if paths
.iter()
.any(|p| is_url(p) && without_secrets(&p.to_string_lossy()).1)
|| endpoint.as_ref().is_some_and(|(_, cut)| *cut)
{
notes.push(
"datui left a user, password or query string out of the URL, as it may be a \
credential: add it back if the server needs it."
.to_string(),
);
}
let in_store = names.iter().find(|n| store_scheme(n).is_some());
let storage = in_store
.and_then(|name| storage_options(name, record, endpoint.as_ref().map(|(e, _)| e.as_str())));
if let Some(name) = in_store
&& python.eager
{
notes.push(format!(
"{} reads no object store: download {name} and read it from disk.",
python.call
));
}
let mut call = Call {
record,
paths,
format,
names: &names,
below,
storage,
args,
after: options.read_python.clone(),
skip_tail: None,
notes,
};
if let Some(arguments) = python.arguments
&& let Some(source) = arguments(&mut call)
{
return source;
}
let Call {
args,
mut after,
skip_tail,
mut notes,
..
} = call;
after.extend(skip_tail);
if !record.read_as_text.is_empty() {
notes.push(format!(
"datui read these columns as text from every file: {}.",
record.read_as_text.join(", ")
));
}
let mut call = format!("{}({})", python.call, args.join(", "));
if python.eager {
notes.extend(read_whole_note(record, &call));
call.push_str(".lazy()");
}
Source::Read {
call,
after,
notes,
imports: Vec::new(),
}
}
fn csv_null_values(options: &OpenOptions, schema: &Schema) -> Option<String> {
let specs = options.null_values.as_ref().filter(|s| !s.is_empty())?;
let mut global = Vec::new();
let mut per_column: Vec<(String, String)> = Vec::new();
for spec in specs {
match spec.find('=') {
Some(i) => per_column.push((spec[..i].to_string(), spec[i + 1..].to_string())),
None => global.push(spec.clone()),
}
}
let dict = |pairs: Vec<(String, String)>| {
let items: Vec<String> = pairs
.iter()
.map(|(c, v)| format!("{}: {}", py_str(c), py_str(v)))
.collect();
format!("{{{}}}", items.join(", "))
};
Some(match (global.as_slice(), per_column.is_empty()) {
([one], true) => py_str(one),
(_, true) => py_names(&global),
([], false) => dict(per_column),
(_, false) => dict(
schema
.iter_names()
.map(|name| {
let value = per_column
.iter()
.rev()
.find(|(c, _)| c == name.as_str())
.map(|(_, v)| v.clone())
.unwrap_or_else(|| global[0].clone());
(name.to_string(), value)
})
.collect(),
),
})
}
pub fn py_value(value: &AnyValue) -> Option<String> {
Some(match value {
AnyValue::Null => "None".to_string(),
AnyValue::Boolean(b) => py_bool(*b).to_string(),
AnyValue::String(s) => py_str(s),
AnyValue::StringOwned(s) => py_str(s),
AnyValue::Int8(v) => v.to_string(),
AnyValue::Int16(v) => v.to_string(),
AnyValue::Int32(v) => v.to_string(),
AnyValue::Int64(v) => v.to_string(),
AnyValue::UInt8(v) => v.to_string(),
AnyValue::UInt16(v) => v.to_string(),
AnyValue::UInt32(v) => v.to_string(),
AnyValue::UInt64(v) => v.to_string(),
AnyValue::Float32(v) => py_float(f64::from(*v)),
AnyValue::Float64(v) => py_float(*v),
AnyValue::Date(days) => {
let date = chrono::NaiveDate::from_ymd_opt(1970, 1, 1)?
.checked_add_signed(chrono::Duration::days(i64::from(*days)))?;
use chrono::Datelike;
format!("pl.date({}, {}, {})", date.year(), date.month(), date.day())
}
_ => return None,
})
}
#[derive(Debug, Clone, PartialEq)]
pub struct Script {
pub source: Source,
pub steps: Vec<Step>,
}
impl Script {
pub fn render(&self) -> String {
let mut out = String::from("import polars as pl\n");
if let Source::Read { imports, .. } = &self.source {
for import in imports {
out.push_str(import);
out.push('\n');
}
}
out.push('\n');
let (head, mut lines) = match &self.source {
Source::Read {
call, after, notes, ..
} => {
for note in notes {
out.push_str(&py_comment(note));
out.push('\n');
}
(call.clone(), after.clone())
}
Source::Placeholder { what } => {
out.push_str(&py_comment(what));
out.push_str("\ndf = ...\n\n");
("df.lazy()".to_string(), Vec::new())
}
};
let mut stopped = false;
for step in &self.steps {
let calls = step.python();
if stopped {
for call in &calls {
lines.extend(call.lines().map(|c| {
if c.starts_with('#') {
c.to_string()
} else {
format!("# {c}")
}
}));
}
} else {
stopped = matches!(step, Step::Unreproducible(_));
lines.extend(calls);
}
}
if lines.is_empty() {
out.push_str(&format!("df = {head}\n"));
} else {
out.push_str("df = (\n");
out.push_str(&format!(" {head}\n"));
for line in lines {
out.push_str(&format!(" {line}\n"));
}
out.push_str(")\n");
}
out
}
}
#[cfg(test)]
mod tests {
use super::*;
fn statement(column: &str, operator: FilterOperator, value: &str) -> FilterStatement {
FilterStatement {
columns: Vec::new(),
column: column.to_string(),
operator,
value: value.to_string(),
logical_op: LogicalOperator::And,
}
}
fn script(steps: Vec<Step>) -> String {
Script {
source: Source::Read {
call: "pl.scan_parquet(\"sales.parquet\")".to_string(),
after: Vec::new(),
notes: Vec::new(),
imports: Vec::new(),
},
steps,
}
.render()
}
fn typed_frame() -> DataFrame {
let tz = TimeZone::opt_try_new(Some("Europe/Paris")).unwrap();
let us = |h: i64| 1_704_067_200_000_000 + h * 3_600_000_000;
df!(
"d" => &[Some(19723i32), Some(19724), Some(19725), None],
"t" => &[Some(us(0)), Some(us(5)), Some(us(24)), None],
"c" => &[Some(5 * 3_600_000_000_000i64), Some(6 * 3_600_000_000_000 + 500_000_000), Some(7 * 3_600_000_000_000), None],
"du" => &[Some(1_000i64), Some(90_000), Some(3_600_000), None],
"m" => &[Some("1.50"), Some("2.00"), Some("3.25"), None],
"f" => &[Some(0.1f32), Some(0.2), Some(0.1), None],
"x" => &[Some(0.1 + 0.2), Some(0.3), Some(1.0), None],
)
.unwrap()
.lazy()
.with_columns([
col("d").cast(DataType::Date),
col("t").cast(DataType::Datetime(TimeUnit::Microseconds, None)),
col("t")
.cast(DataType::Datetime(TimeUnit::Microseconds, tz))
.alias("z"),
col("c").cast(DataType::Time),
col("du").cast(DataType::Duration(TimeUnit::Milliseconds)),
col("m").cast(DataType::Decimal(10, 2)),
])
.collect()
.unwrap()
}
#[test]
fn every_operator_compares_in_the_columns_own_type() {
use FilterOperator::*;
let frame = typed_frame();
let schema = frame.schema().clone();
let rows = |column: &str, operator, value: &str| {
let statement = statement(column, operator, value);
assert_eq!(SidebarFilter::problem(&statement, schema.get(column)), None);
let typed = SidebarFilter::typed(&statement, schema.get(column));
frame
.clone()
.lazy()
.filter(filters_expr(&[typed]).unwrap())
.collect()
.unwrap_or_else(|e| panic!("{column} {operator:?} {value}: {e}"))
.height()
};
for (column, value, counts) in [
("d", "2024-01-02", [1, 2, 1, 1, 2, 2]),
("t", "2024-01-01", [1, 2, 2, 0, 3, 1]),
("t", "2024-01-01 05:00", [1, 2, 1, 1, 2, 2]),
("t", "2024-01-01T05:00:00.000", [1, 2, 1, 1, 2, 2]),
("z", "2024-01-01 06:00", [1, 2, 1, 1, 2, 2]),
("z", "2024-01-01 05:00+00:00", [1, 2, 1, 1, 2, 2]),
("c", "06:00:00.5", [1, 2, 1, 1, 2, 2]),
("du", "1m 30s", [1, 2, 1, 1, 2, 2]),
("m", "2", [1, 2, 1, 1, 2, 2]),
("m", "1.5", [1, 2, 2, 0, 3, 1]),
("x", "0.3", [1, 2, 2, 0, 3, 1]),
] {
let got = [Eq, NotEq, Gt, Lt, GtEq, LtEq].map(|op| rows(column, op, value));
assert_eq!(got, counts, "{column} {value}");
}
assert_eq!(rows("f", Eq, "0.1"), 2);
assert_eq!(rows("x", IsNull, ""), 1);
assert_eq!(rows("x", IsNotNull, ""), 3);
}
#[test]
fn a_value_that_does_not_read_as_the_column_says_so() {
let frame = typed_frame();
let schema = frame.schema();
let problem = |column: &str, operator, value: &str| {
SidebarFilter::problem(&statement(column, operator, value), schema.get(column))
};
assert_eq!(
problem("d", FilterOperator::Eq, "2024-13-01").as_deref(),
Some("d: \"2024-13-01\" is not a date written YYYY-MM-DD")
);
assert!(problem("t", FilterOperator::Gt, "soon").is_some());
assert!(problem("x", FilterOperator::Lt, "abc").is_some());
assert_eq!(problem("d", FilterOperator::Contains, "2024"), None);
assert_eq!(problem("d", FilterOperator::IsNull, ""), None);
}
#[test]
fn typed_filters_read_back_in_python() {
let frame = typed_frame();
let schema = frame.schema();
let python = |column: &str, operator, value: &str| {
SidebarFilter::typed(&statement(column, operator, value), schema.get(column)).python()
};
assert_eq!(
python("d", FilterOperator::Eq, "2024-01-02"),
"pl.col(\"d\") == pl.date(2024, 1, 2)"
);
assert_eq!(
python("t", FilterOperator::Gt, "2024-01-01 05:00"),
"pl.col(\"t\") > pl.datetime(2024, 1, 1, 5, 0, 0, 0, time_unit=\"us\")"
);
assert_eq!(
python("z", FilterOperator::LtEq, "2024-01-01 06:00"),
"pl.col(\"z\") <= pl.datetime(2024, 1, 1, 5, 0, 0, 0, time_unit=\"us\", \
time_zone=\"UTC\").dt.convert_time_zone(\"Europe/Paris\")"
);
assert_eq!(
python("c", FilterOperator::Eq, "06:00:00.5"),
"pl.col(\"c\") == pl.time(6, 0, 0, 500000)"
);
assert_eq!(
python("du", FilterOperator::Lt, "1h"),
"pl.col(\"du\") < pl.duration(milliseconds=3600000, time_unit=\"ms\")"
);
assert_eq!(
python("m", FilterOperator::NotEq, "1.5"),
"pl.col(\"m\") != pl.lit(\"1.50\").cast(pl.Decimal(10, 2))"
);
assert_eq!(
python("x", FilterOperator::Eq, "0.3"),
"pl.col(\"x\") == 0.3"
);
assert_eq!(
python("x", FilterOperator::IsNull, ""),
"pl.col(\"x\").is_null()"
);
}
#[test]
fn strings_and_floats_read_back_in_python() {
assert_eq!(py_str("a\"b\\c\nd"), "\"a\\\"b\\\\c\\nd\"");
assert_eq!(py_str("\u{1}"), "\"\\u0001\"");
assert_eq!(py_float(5.0), "5.0");
assert_eq!(py_float(0.1), "0.1");
assert_eq!(py_float(1e20), "1e20");
assert_eq!(py_float(f64::NAN), "float(\"nan\")");
}
#[test]
fn text_in_a_comment_cannot_end_it() {
assert_eq!(
py_comment("a\nimport os\r\u{2028}x"),
"# a\\nimport os\\r\\u2028x"
);
let text = script(vec![Step::Unreproducible("where k is \"\nboom()".into())]);
assert!(text.contains(" # where k is \"\\nboom()\n"), "{text}");
}
#[test]
fn sql_is_triple_quoted_only_where_nothing_in_it_ends_the_string() {
let sql = |sql: &str| {
Step::Sql {
sql: sql.into(),
ordered_by: Vec::new(),
}
.python()
};
assert_eq!(
sql("SELECT *\nFROM df"),
vec![
".sql(",
" \"\"\"SELECT *\nFROM df\"\"\",",
" table_name=\"df\",",
")"
]
);
assert_eq!(
sql("SELECT *\nFROM df ORDER BY \"a\""),
vec![".sql(\"SELECT *\\nFROM df ORDER BY \\\"a\\\"\", table_name=\"df\")"]
);
let text = script(vec![
Step::Unreproducible("drilled into a group held as lists".into()),
Step::Sql {
sql: "SELECT *\nFROM df".into(),
ordered_by: Vec::new(),
},
]);
assert!(text.contains(" # FROM df\"\"\",\n"), "{text}");
}
#[test]
fn a_view_with_nothing_applied_is_the_reader() {
assert_eq!(
script(Vec::new()),
"import polars as pl\n\ndf = pl.scan_parquet(\"sales.parquet\")\n"
);
}
#[test]
fn filters_typed_by_column_then_a_multi_column_sort_then_a_projection() {
let schema = Schema::from_iter([
Field::new("region".into(), DataType::String),
Field::new("amount".into(), DataType::Float64),
Field::new("qty".into(), DataType::Int64),
]);
let mut or = statement("qty", FilterOperator::GtEq, "3");
or.logical_op = LogicalOperator::Or;
let filters: Vec<SidebarFilter> = [
statement("region", FilterOperator::Eq, "north"),
statement("amount", FilterOperator::Gt, "10"),
or,
]
.iter()
.map(|s| SidebarFilter::typed(s, schema.get(&s.column)))
.collect();
let text = script(vec![
Step::Filter(filters),
Step::Sort {
columns: vec!["amount".into(), "region".into()],
descending: vec![true, false],
},
Step::Select(vec!["order_id".into(), "customer".into(), "amount".into()]),
]);
assert_eq!(
text,
"import polars as pl\n\n\
df = (\n \
pl.scan_parquet(\"sales.parquet\")\n \
.filter(((pl.col(\"region\") == \"north\") & (pl.col(\"amount\") > 10.0)) | (pl.col(\"qty\") >= 3))\n \
.sort([\"amount\", \"region\"], descending=[True, False], nulls_last=True, maintain_order=True)\n \
.select([\"order_id\", \"customer\", \"amount\"])\n\
)\n"
);
}
#[test]
fn contains_filters_are_literal_and_a_number_that_does_not_parse_stays_text() {
let s = SidebarFilter::typed(
&statement("name", FilterOperator::NotContains, "a.b"),
Some(&DataType::String),
);
assert_eq!(
s.python(),
"~pl.col(\"name\").str.contains(\"a.b\", literal=True)"
);
let s = SidebarFilter::typed(
&statement("n", FilterOperator::Eq, "n/a"),
Some(&DataType::Int64),
);
assert_eq!(s.value, FilterValue::Str("n/a".into()));
}
#[test]
fn one_sort_column_reads_plainly() {
assert_eq!(
sort_call(&["amount".into()], &[true]),
".sort(\"amount\", descending=True, nulls_last=True, maintain_order=True)"
);
}
#[test]
fn steps_after_one_python_cannot_repeat_are_commented_out() {
let text = script(vec![
Step::Unreproducible("drilled into a group held as lists".into()),
Step::Reverse,
]);
assert!(
text.contains(" # drilled into a group held as lists\n # .reverse()\n"),
"{text}"
);
}
#[test]
fn a_placeholder_source_leaves_df_to_the_user() {
let text = Script {
source: Source::Placeholder {
what: "The data datui read from standard input: load it here.".into(),
},
steps: vec![Step::Reverse],
}
.render();
assert_eq!(
text,
"import polars as pl\n\n\
# The data datui read from standard input: load it here.\n\
df = ...\n\n\
df = (\n df.lazy()\n .reverse()\n)\n"
);
}
#[test]
fn a_grouped_query_groups_then_orders_by_its_keys() {
let input = Schema::from_iter([
Field::new("dept".into(), DataType::String),
Field::new("salary".into(), DataType::Float64),
Field::new("id".into(), DataType::Int64),
Field::new("age".into(), DataType::Int64),
]);
let text = script(vec![Step::Query {
query: "select avg salary, n: count id by dept where age > 30".into(),
input: Arc::new(input),
keys: vec!["dept".into()],
}]);
assert!(
text.contains(
" .filter(pl.col(\"age\") > 30.0)\n \
.group_by(\"dept\")\n \
.agg(pl.col(\"salary\").mean().alias(\"avg_salary\"), pl.col(\"id\").count().alias(\"n\"))\n \
.sort(\"dept\", nulls_last=True, maintain_order=True)\n"
),
"{text}"
);
}
#[test]
fn csv_options_become_reader_arguments() {
let mut options = OpenOptions::new();
options.delimiter = Some(b';');
options.has_header = Some(false);
options.skip_rows = Some(2);
options.null_values = Some(vec!["NA".into()]);
options.skip_tail_rows = Some(1);
let paths = vec![PathBuf::from("data/x.csv")];
let schema = Schema::default();
let record = OpenRecord {
paths: Some(&paths),
options: &options,
schema: &schema,
remote_objects: Vec::new(),
s3_endpoint: None,
s3_region: None,
unsigned: false,
format: None,
read_mode: None,
read_as_text: Vec::new(),
spec: None,
};
let Source::Read { call, after, .. } = source(&record) else {
panic!("a CSV has a reader");
};
assert_eq!(
call,
"pl.scan_csv(\"data/x.csv\", separator=\";\", has_header=False, skip_rows=2, \
try_parse_dates=True, null_values=\"NA\")"
);
assert_eq!(
after,
vec![".filter(pl.int_range(pl.len()) < pl.len() - 1)"]
);
}
#[test]
fn reads_python_cannot_repeat_leave_a_placeholder() {
let schema = Schema::default();
let placeholder = |paths: &[PathBuf], options: &OpenOptions, spec: Option<&str>| {
let record = OpenRecord {
paths: Some(paths),
options,
schema: &schema,
remote_objects: Vec::new(),
s3_endpoint: None,
s3_region: None,
unsigned: false,
format: None,
read_mode: None,
read_as_text: Vec::new(),
spec: spec.map(str::to_string),
};
match source(&record) {
Source::Placeholder { what } => what,
Source::Read { call, .. } => panic!("a reader was written: {call}"),
}
};
let plain = OpenOptions::new();
let what = placeholder(&[PathBuf::from("a.l2")], &plain, Some("acme.l2feed"));
assert!(what.contains("acme.l2feed"), "{what}");
let mut named = OpenOptions::new();
named.spec_name = Some("acme.l2feed".into());
placeholder(&[PathBuf::from("a.bin")], &named, None);
placeholder(&[PathBuf::from("track.gpx")], &plain, None);
placeholder(&[PathBuf::from("drive.nmea")], &plain, None);
let csv = [PathBuf::from("log.csv")];
let mut comment = OpenOptions::new();
comment.comment_char = Some("######".into());
assert!(placeholder(&csv, &comment, None).contains("--comment"));
let mut rows = OpenOptions::new();
rows.header_rows = vec![3, 2];
assert!(placeholder(&csv, &rows, None).contains("--header-rows"));
let mut space = OpenOptions::new();
space.skip_initial_space = true;
assert!(placeholder(&csv, &space, None).contains("--skip-initial-space"));
}
fn record_for<'a>(
paths: &'a [PathBuf],
options: &'a OpenOptions,
schema: &'a Schema,
) -> OpenRecord<'a> {
OpenRecord {
paths: Some(paths),
options,
format: None,
read_mode: None,
schema,
remote_objects: Vec::new(),
s3_endpoint: None,
s3_region: None,
unsigned: false,
read_as_text: Vec::new(),
spec: None,
}
}
fn call_of(source: Source) -> (String, Vec<String>) {
match source {
Source::Read { call, notes, .. } => (call, notes),
Source::Placeholder { what } => panic!("a placeholder: {what}"),
}
}
#[test]
fn every_store_reader_gets_its_storage_options() {
let schema = Schema::default();
let options = OpenOptions::new();
let s3 = [PathBuf::from("s3://b/logs/a.jsonl")];
let mut record = record_for(&s3, &options, &schema);
record.s3_endpoint = Some("http://localhost:9000".into());
record.s3_region = Some("us-east-1".into());
assert_eq!(
call_of(source(&record)).0,
"pl.scan_ndjson(\"s3://b/logs/a.jsonl\", storage_options={\"aws_endpoint_url\": \
\"http://localhost:9000\", \"aws_region\": \"us-east-1\"})"
);
let gcs = [PathBuf::from("gs://public/x.parquet")];
let mut record = record_for(&gcs, &options, &schema);
record.unsigned = true;
assert_eq!(
call_of(source(&record)).0,
"pl.scan_parquet(\"gs://public/x.parquet\", storage_options={\"skip_signature\": \"true\"})"
);
let azure = [PathBuf::from(
"abfss://data@acct.dfs.core.windows.net/t/x.csv",
)];
let (call, notes) = call_of(source(&record_for(&azure, &options, &schema)));
assert!(
call.starts_with("pl.scan_csv(\"abfss://data@acct.dfs.core.windows.net/t/x.csv\", ")
&& call.contains("storage_options={\"account_name\": \"acct\"}"),
"{call}"
);
assert!(
notes.is_empty(),
"the container is no credential: {notes:?}"
);
let json = [PathBuf::from("s3://b/x.json")];
let (call, notes) = call_of(source(&record_for(&json, &options, &schema)));
assert_eq!(call, "pl.read_json(\"s3://b/x.json\").lazy()");
assert!(notes[0].contains("reads no object store"), "{notes:?}");
}
#[test]
fn a_source_id_is_no_credential() {
assert_eq!(
without_secrets("s3://minio@bucket/x.parquet"),
("s3://bucket/x.parquet".to_string(), false)
);
assert_eq!(
without_secrets("abfss://c@a.dfs.core.windows.net/x"),
("abfss://c@a.dfs.core.windows.net/x".to_string(), false)
);
}
#[test]
fn a_teed_pipe_reads_its_file() {
let schema = Schema::default();
let mut options = OpenOptions::new();
options.tee = Some(PathBuf::from("rec.csv"));
let stdin = [PathBuf::from("-")];
let mut record = record_for(&stdin, &options, &schema);
record.format = Some(FileFormat::Csv);
assert_eq!(
call_of(source(&record)).0,
"pl.scan_csv(\"rec.csv\", try_parse_dates=True)"
);
}
#[test]
fn a_whole_read_of_a_lazy_table_says_so() {
let schema = Schema::default();
let mut options = OpenOptions::new();
options.table = Some("orders".into());
let db = [PathBuf::from("shop.db")];
let mut record = record_for(&db, &options, &schema);
record.format = Some(FileFormat::Sqlite);
record.read_mode = Some(crate::ReadMode::Lazy);
let (call, notes) = call_of(source(&record));
assert!(call.starts_with("pl.read_database("), "{call}");
assert!(
notes.contains(
&"Read: lazy scan in datui; pl.read_database reads the file whole into memory."
.to_string()
),
"{notes:?}"
);
let json = [PathBuf::from("a.json")];
let mut record = record_for(&json, &options, &schema);
record.read_mode = Some(crate::ReadMode::InMemory);
assert!(call_of(source(&record)).1.is_empty());
}
#[test]
fn a_comment_character_is_the_comment_prefix() {
let schema = Schema::default();
let mut options = OpenOptions::new();
options.comment_char = Some("#".into());
let csv = [PathBuf::from("log.csv")];
assert_eq!(
call_of(source(&record_for(&csv, &options, &schema))).0,
"pl.scan_csv(\"log.csv\", comment_prefix=\"#\", try_parse_dates=True)"
);
}
#[test]
fn a_placeholder_names_the_format_read_and_the_table() {
let schema = Schema::default();
let paths = vec![PathBuf::from("flight.bin")];
let mut options = OpenOptions::new();
options.table = Some("GPS".into());
let record = OpenRecord {
paths: Some(&paths),
options: &options,
format: Some(FileFormat::Dataflash),
read_mode: None,
schema: &schema,
remote_objects: Vec::new(),
s3_endpoint: None,
s3_region: None,
unsigned: false,
read_as_text: Vec::new(),
spec: None,
};
let Source::Placeholder { what } = source(&record) else {
panic!("Polars reads no DataFlash");
};
assert_eq!(
what,
"flight.bin --table GPS: Polars has no reader for DataFlash files; load it here."
);
}
#[test]
fn credentials_in_a_url_stay_out_of_the_script() {
assert_eq!(
without_secrets("https://u:p@host.example/d/x.parquet?X-Amz-Signature=abc#f"),
("https://host.example/d/x.parquet".to_string(), true)
);
assert_eq!(
without_secrets("s3://bucket/data-?.parquet"),
("s3://bucket/data-?.parquet".to_string(), false)
);
let options = OpenOptions::new();
let schema = Schema::default();
let paths = vec![PathBuf::from(
"https://user:secret@host.example/d/x.csv?token=s3cr3t",
)];
let record = OpenRecord {
paths: Some(&paths),
options: &options,
schema: &schema,
remote_objects: Vec::new(),
s3_endpoint: Some("http://key:secret@localhost:9000".into()),
s3_region: None,
unsigned: false,
format: None,
read_mode: None,
read_as_text: Vec::new(),
spec: None,
};
let text = Script {
source: source(&record),
steps: Vec::new(),
}
.render();
assert!(
!text.contains("secret") && !text.contains("s3cr3t"),
"{text}"
);
assert!(text.contains("\"https://host.example/d/x.csv\""), "{text}");
assert!(text.contains("# datui left a user"), "{text}");
}
#[test]
fn stdin_and_bucket_prefixes() {
let options = OpenOptions::new();
let schema = Schema::default();
let stdin = vec![PathBuf::from("-")];
let record = |paths: &'static [PathBuf]| OpenRecord {
paths: Some(paths),
options: &options,
schema: &schema,
remote_objects: vec!["s3://b/p/year=2024/a.parquet".into()],
s3_endpoint: Some("http://localhost:9000".into()),
s3_region: None,
unsigned: false,
format: None,
read_mode: None,
read_as_text: Vec::new(),
spec: None,
};
let stdin: &'static [PathBuf] = Box::leak(stdin.into_boxed_slice());
assert!(matches!(source(&record(stdin)), Source::Placeholder { .. }));
let prefix: &'static [PathBuf] =
Box::leak(vec![PathBuf::from("s3://b/p/")].into_boxed_slice());
let Source::Read { call, .. } = source(&record(prefix)) else {
panic!("a Parquet prefix has a reader");
};
assert_eq!(
call,
"pl.scan_parquet(\"s3://b/p/**/*.parquet\", hive_partitioning=True, \
storage_options={\"aws_endpoint_url\": \"http://localhost:9000\"})"
);
}
#[test]
fn a_kept_find_names_what_the_script_cannot_match() {
let schema = Schema::from_iter([
Field::new("name".into(), DataType::String),
Field::new("took".into(), DataType::Duration(TimeUnit::Milliseconds)),
]);
let statement = FilterStatement {
columns: vec!["name".into(), "took".into()],
column: crate::filter_modal::ANY_COLUMN.to_string(),
operator: FilterOperator::Has,
value: "1d".to_string(),
logical_op: LogicalOperator::And,
};
let filter = SidebarFilter::typed_in(&statement, &schema, &[]);
assert_eq!(filter.unscriptable_columns(), ["took"]);
assert!(!filter.python().contains("took"), "{}", filter.python());
let bad = FilterStatement {
columns: Vec::new(),
column: "name".into(),
operator: FilterOperator::HasRegex,
value: "(".into(),
logical_op: LogicalOperator::And,
};
let why = SidebarFilter::problem(&bad, Some(&DataType::String)).expect("refused");
assert!(why.starts_with("Not a regex"), "{why}");
}
#[test]
fn a_kept_find_searches_the_shown_text_columns() {
let schema = Schema::from_iter([
Field::new("name".into(), DataType::String),
Field::new("tags".into(), DataType::List(Box::new(DataType::String))),
Field::new("note".into(), DataType::String),
Field::new("n".into(), DataType::Int64),
]);
let statement = FilterStatement {
columns: Vec::new(),
column: crate::filter_modal::ANY_COLUMN.to_string(),
operator: FilterOperator::Has,
value: "al".to_string(),
logical_op: LogicalOperator::And,
};
let shown = ["name", "tags", "n"].map(String::from);
let filter = SidebarFilter::typed_in(&statement, &schema, &shown);
let names: Vec<&str> = filter.searched.iter().map(|(n, _)| n.as_str()).collect();
assert_eq!(names, ["name", "n"]);
let script = filter.python();
assert!(
!script.contains("tags") && !script.contains("note"),
"{script}"
);
assert!(script.contains("pl.any_horizontal"), "{script}");
}
}