Skip to content

Commit

Permalink
feat: print row index in explain + dot (pola-rs#17074)
Browse files Browse the repository at this point in the history
  • Loading branch information
coastalwhite authored Jun 20, 2024
1 parent f482077 commit 69b8440
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 26 deletions.
33 changes: 20 additions & 13 deletions crates/polars-plan/src/plans/ir/dot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -255,15 +255,21 @@ impl<'a> IRDotDisplay<'a> {
let path = PathsDisplay(paths.as_ref());
let with_columns = options.with_columns.as_ref().map(|cols| cols.as_ref());
let with_columns = NumColumns(with_columns);
let total_columns = file_info.schema.len();
let predicate = predicate.as_ref().map(|e| self.display_expr(e));
let predicate = OptionExprIRDisplay(predicate);
let total_columns =
file_info.schema.len() - usize::from(options.row_index.is_some());

write_label(f, id, |f| {
write!(
f,
"{name} SCAN {path}\nπ {with_columns}/{total_columns};\nσ {predicate}",
)
write!(f, "{name} SCAN {path}\nπ {with_columns}/{total_columns};",)?;

if let Some(predicate) = predicate.as_ref() {
write!(f, "\nσ {}", self.display_expr(predicate))?;
}

if let Some(row_index) = options.row_index.as_ref() {
write!(f, "\nrow index: {} (+{})", row_index.name, row_index.offset)?;
}

Ok(())
})?;
},
Join {
Expand Down Expand Up @@ -332,21 +338,22 @@ impl<'a> IRDotDisplay<'a> {
}

// A few utility structures for formatting
struct PathsDisplay<'a>(&'a [PathBuf]);
pub(crate) struct PathsDisplay<'a>(pub &'a [PathBuf]);
struct NumColumns<'a>(Option<&'a [String]>);
struct NumColumnsSchema<'a>(Option<&'a Schema>);
struct OptionExprIRDisplay<'a>(Option<ExprIRDisplay<'a>>);

impl fmt::Display for PathsDisplay<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.0.len() {
0 => Ok(()),
1 => self.0[0].display().fmt(f),
0 => write!(f, "[]"),
1 => write!(f, "[{}]", self.0[0].display()),
2 => write!(f, "[{}, {}]", self.0[0].display(), self.0[1].display()),
_ => write!(
f,
"{} files: first file: {}",
self.0.len(),
self.0[0].display()
"[{}, ... {} other files]",
self.0[0].to_string_lossy(),
self.0.len() - 1,
),
}
}
Expand Down
29 changes: 16 additions & 13 deletions crates/polars-plan/src/plans/ir/format.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,10 @@ use std::path::PathBuf;

use polars_core::datatypes::AnyValue;
use polars_core::schema::Schema;
use polars_io::RowIndex;
use recursive::recursive;

use super::ir::dot::PathsDisplay;
use crate::prelude::*;

pub struct IRDisplay<'a> {
Expand Down Expand Up @@ -60,18 +62,11 @@ fn write_scan(
total_columns: usize,
predicate: &Option<ExprIRDisplay<'_>>,
n_rows: Option<usize>,
row_index: Option<&RowIndex>,
) -> fmt::Result {
let path_fmt = match path.len() {
1 => path[0].to_string_lossy(),
0 => "".into(),
_ => Cow::Owned(format!(
"{} files: first file: {}",
path.len(),
path[0].to_string_lossy()
)),
};

write!(f, "{:indent$}{name} SCAN {path_fmt}", "")?;
write!(f, "{:indent$}{name} SCAN {}", "", PathsDisplay(path))?;

let total_columns = total_columns - usize::from(row_index.is_some());
if n_columns > 0 {
write!(
f,
Expand All @@ -87,6 +82,12 @@ fn write_scan(
if let Some(n_rows) = n_rows {
write!(f, "\n{:indent$}N_ROWS: {n_rows}", "")?;
}
if let Some(row_index) = row_index {
write!(f, "\n{:indent$}ROW_INDEX: {}", "", row_index.name)?;
if row_index.offset != 0 {
write!(f, " (offset: {})", row_index.offset)?;
}
}
Ok(())
}

Expand Down Expand Up @@ -167,11 +168,12 @@ impl<'a> IRDisplay<'a> {
f,
"PYTHON",
&[],
sub_indent,
indent,
n_columns,
total_columns,
&predicate,
options.n_rows,
None,
)
},
Union { inputs, options } => {
Expand Down Expand Up @@ -234,11 +236,12 @@ impl<'a> IRDisplay<'a> {
f,
scan_type.into(),
paths,
sub_indent,
indent,
n_columns,
file_info.schema.len(),
&predicate,
file_options.n_rows,
file_options.row_index.as_ref(),
)
},
Filter { predicate, input } => {
Expand Down

0 comments on commit 69b8440

Please sign in to comment.