Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

depr(python,rust!): Rename pl.count() to pl.len() #13719

Merged
merged 4 commits into from
Jan 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions crates/polars-lazy/src/physical_plan/expressions/count.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::borrow::Cow;

use polars_core::prelude::*;
use polars_plan::dsl::consts::COUNT;
use polars_plan::dsl::consts::LEN;

use crate::physical_plan::state::ExecutionState;
use crate::prelude::*;
Expand All @@ -12,7 +12,7 @@ pub struct CountExpr {

impl CountExpr {
pub(crate) fn new() -> Self {
Self { expr: Expr::Count }
Self { expr: Expr::Len }
}
}

Expand All @@ -22,7 +22,7 @@ impl PhysicalExpr for CountExpr {
}

fn evaluate(&self, df: &DataFrame, _state: &ExecutionState) -> PolarsResult<Series> {
Ok(Series::new("count", [df.height() as IdxSize]))
Ok(Series::new("len", [df.height() as IdxSize]))
}

fn evaluate_on_groups<'a>(
Expand All @@ -31,13 +31,13 @@ impl PhysicalExpr for CountExpr {
groups: &'a GroupsProxy,
_state: &ExecutionState,
) -> PolarsResult<AggregationContext<'a>> {
let ca = groups.group_count().with_name(COUNT);
let ca = groups.group_count().with_name(LEN);
let s = ca.into_series();
Ok(AggregationContext::new(s, Cow::Borrowed(groups), true))
}

fn to_field(&self, _input_schema: &Schema) -> PolarsResult<Field> {
Ok(Field::new(COUNT, IDX_DTYPE))
Ok(Field::new(LEN, IDX_DTYPE))
}

fn as_partitioned_aggregator(&self) -> Option<&dyn PartitionedAggregation> {
Expand Down Expand Up @@ -67,6 +67,6 @@ impl PartitionedAggregation for CountExpr {
) -> PolarsResult<Series> {
// SAFETY: groups are in bounds.
let agg = unsafe { partitioned.agg_sum(groups) };
Ok(agg.with_name(COUNT))
Ok(agg.with_name(LEN))
}
}
6 changes: 3 additions & 3 deletions crates/polars-lazy/src/physical_plan/planner/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ pub(crate) fn create_physical_expr(
use AExpr::*;

match expr_arena.get(expression).clone() {
Count => Ok(Arc::new(phys_expr::CountExpr::new())),
Len => Ok(Arc::new(phys_expr::CountExpr::new())),
Window {
mut function,
partition_by,
Expand Down Expand Up @@ -129,8 +129,8 @@ pub(crate) fn create_physical_expr(
if apply_columns.is_empty() {
if has_aexpr(function, expr_arena, |e| matches!(e, AExpr::Literal(_))) {
apply_columns.push(Arc::from("literal"))
} else if has_aexpr(function, expr_arena, |e| matches!(e, AExpr::Count)) {
apply_columns.push(Arc::from("count"))
} else if has_aexpr(function, expr_arena, |e| matches!(e, AExpr::Len)) {
apply_columns.push(Arc::from("len"))
} else {
let e = node_to_expr(function, expr_arena);
polars_bail!(
Expand Down
6 changes: 3 additions & 3 deletions crates/polars-lazy/src/physical_plan/planner/lp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ fn partitionable_gb(
let depth = (expr_arena).iter(*agg).count();

// These single expressions are partitionable
if matches!(aexpr, AExpr::Count) {
if matches!(aexpr, AExpr::Len) {
continue;
}
// col()
Expand All @@ -55,7 +55,7 @@ fn partitionable_gb(
// count().alias() is allowed: count of 2
if depth <= 2 {
match expr_arena.get(*input) {
AExpr::Count => {},
AExpr::Len => {},
_ => {
partitionable = false;
break;
Expand Down Expand Up @@ -103,7 +103,7 @@ fn partitionable_gb(
Ternary {truthy, falsy, predicate,..} => {
!has_aggregation(*truthy) && !has_aggregation(*falsy) && !has_aggregation(*predicate)
}
Column(_) | Alias(_, _) | Count | Literal(_) | Cast {..} => {
Column(_) | Alias(_, _) | Len | Literal(_) | Cast {..} => {
true
}
_ => {
Expand Down
4 changes: 2 additions & 2 deletions crates/polars-lazy/src/tests/aggregations.rs
Original file line number Diff line number Diff line change
Expand Up @@ -243,8 +243,8 @@ fn test_binary_agg_context_0() -> PolarsResult<()> {
.lazy()
.group_by_stable([col("groups")])
.agg([when(col("vals").first().neq(lit(1)))
.then(repeat(lit("a"), count()))
.otherwise(repeat(lit("b"), count()))
.then(repeat(lit("a"), len()))
.otherwise(repeat(lit("b"), len()))
.alias("foo")])
.collect()
.unwrap();
Expand Down
2 changes: 1 addition & 1 deletion crates/polars-lazy/src/tests/queries.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1807,7 +1807,7 @@ fn test_partitioned_gb_count() -> PolarsResult<()> {
.group_by([col("col")])
.agg([
// we make sure to alias with a different name
count().alias("counted"),
len().alias("counted"),
col("col").count().alias("count2"),
])
.collect()?;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,14 @@ use crate::executors::sinks::group_by::aggregates::{AggregateFunction, SumAgg};
use crate::expressions::PhysicalPipedExpr;
use crate::operators::DataChunk;

struct Count {}
struct Len {}

impl PhysicalIoExpr for Count {
impl PhysicalIoExpr for Len {
fn evaluate_io(&self, _df: &DataFrame) -> PolarsResult<Series> {
unimplemented!()
}
}
impl PhysicalPipedExpr for Count {
impl PhysicalPipedExpr for Len {
fn evaluate(&self, chunk: &DataChunk, _lazy_state: &dyn Any) -> PolarsResult<Series> {
// the length must match the chunks as the operators expect that
// so we fill a null series.
Expand All @@ -42,7 +42,7 @@ impl PhysicalPipedExpr for Count {
}

fn expression(&self) -> Expr {
Expr::Count
Expr::Len
}
}

Expand All @@ -57,7 +57,7 @@ pub fn can_convert_to_hash_agg(
.map(|(_, ae)| {
match ae {
AExpr::Agg(_)
| AExpr::Count
| AExpr::Len
| AExpr::Cast { .. }
| AExpr::Literal(_)
| AExpr::Column(_)
Expand All @@ -70,7 +70,7 @@ pub fn can_convert_to_hash_agg(
}
ae
})
.filter(|ae| matches!(ae, AExpr::Agg(_) | AExpr::Count))
.filter(|ae| matches!(ae, AExpr::Agg(_) | AExpr::Len))
.count()
== 1
&& can_run_partitioned
Expand All @@ -80,7 +80,7 @@ pub fn can_convert_to_hash_agg(
node = *input
}
match expr_arena.get(node) {
AExpr::Count => true,
AExpr::Len => true,
ae @ AExpr::Agg(agg_fn) => {
matches!(
agg_fn,
Expand Down Expand Up @@ -128,9 +128,9 @@ where
{
match expr_arena.get(node) {
AExpr::Alias(input, _) => convert_to_hash_agg(*input, expr_arena, schema, to_physical),
AExpr::Count => (
AExpr::Len => (
IDX_DTYPE,
Arc::new(Count {}),
Arc::new(Len {}),
AggregateFunction::Count(CountAgg::new()),
),
AExpr::Agg(agg) => match agg {
Expand Down
2 changes: 1 addition & 1 deletion crates/polars-plan/src/dsl/consts.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
pub const COUNT: &str = "count";
pub const LEN: &str = "len";

pub const LITERAL_NAME: &str = "literal";
5 changes: 2 additions & 3 deletions crates/polars-plan/src/dsl/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -135,8 +135,7 @@ pub enum Expr {
Exclude(Box<Expr>, Vec<Excluded>),
/// Set root name as Alias
KeepName(Box<Expr>),
/// Special case that does not need columns
Count,
Len,
/// Take the nth column in the `DataFrame`
Nth(i64),
// skipped fields must be last otherwise serde fails in pickle
Expand Down Expand Up @@ -223,7 +222,7 @@ impl Hash for Expr {
options.hash(state);
},
// already hashed by discriminant
Expr::Wildcard | Expr::Count => {},
Expr::Wildcard | Expr::Len => {},
#[allow(unreachable_code)]
_ => {
// the panic checks if we hit this
Expand Down
2 changes: 1 addition & 1 deletion crates/polars-plan/src/dsl/functions/index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use super::*;
pub fn arg_sort_by<E: AsRef<[Expr]>>(by: E, descending: &[bool]) -> Expr {
let e = &by.as_ref()[0];
let name = expr_output_name(e).unwrap();
int_range(lit(0 as IdxSize), count().cast(IDX_DTYPE), 1, IDX_DTYPE)
int_range(lit(0 as IdxSize), len().cast(IDX_DTYPE), 1, IDX_DTYPE)
.sort_by(by, descending)
.alias(name.as_ref())
}
Expand Down
18 changes: 3 additions & 15 deletions crates/polars-plan/src/dsl/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1760,21 +1760,9 @@ where
}
}

/// Count expression.
pub fn count() -> Expr {
Expr::Count
}

/// Return the cumulative count of the context.
#[cfg(feature = "range")]
pub fn cum_count(reverse: bool) -> Expr {
let start = lit(1 as IdxSize);
let end = count() + lit(1 as IdxSize);
let mut range = int_range(start, end, 1, IDX_DTYPE);
if reverse {
range = range.reverse()
}
range.alias("cum_count")
/// Return the number of rows in the context.
pub fn len() -> Expr {
Expr::Len
}

/// First column in DataFrame.
Expand Down
12 changes: 6 additions & 6 deletions crates/polars-plan/src/logical_plan/aexpr/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use crate::dsl::function_expr::FunctionExpr;
#[cfg(feature = "cse")]
use crate::logical_plan::visitor::AexprNode;
use crate::logical_plan::Context;
use crate::prelude::consts::COUNT;
use crate::prelude::consts::LEN;
use crate::prelude::*;

#[derive(Clone, Debug, IntoStaticStr)]
Expand Down Expand Up @@ -188,7 +188,7 @@ pub enum AExpr {
offset: Node,
length: Node,
},
Count,
Len,
Nth(i64),
}

Expand Down Expand Up @@ -224,7 +224,7 @@ impl AExpr {
| SortBy { .. }
| Agg { .. }
| Window { .. }
| Count
| Len
| Slice { .. }
| Gather { .. }
| Nth(_)
Expand Down Expand Up @@ -259,7 +259,7 @@ impl AExpr {
use AExpr::*;

match self {
Nth(_) | Column(_) | Literal(_) | Wildcard | Count => {},
Nth(_) | Column(_) | Literal(_) | Wildcard | Len => {},
Alias(e, _) => container.push(*e),
BinaryExpr { left, op: _, right } => {
// reverse order so that left is popped first
Expand Down Expand Up @@ -338,7 +338,7 @@ impl AExpr {
pub(crate) fn replace_inputs(mut self, inputs: &[Node]) -> Self {
use AExpr::*;
let input = match &mut self {
Column(_) | Literal(_) | Wildcard | Count | Nth(_) => return self,
Column(_) | Literal(_) | Wildcard | Len | Nth(_) => return self,
Alias(input, _) => input,
Cast { expr, .. } => expr,
Explode(input) => input,
Expand Down Expand Up @@ -420,7 +420,7 @@ impl AExpr {
pub(crate) fn is_leaf(&self) -> bool {
matches!(
self,
AExpr::Column(_) | AExpr::Literal(_) | AExpr::Count | AExpr::Nth(_)
AExpr::Column(_) | AExpr::Literal(_) | AExpr::Len | AExpr::Nth(_)
)
}
}
Expand Down
2 changes: 1 addition & 1 deletion crates/polars-plan/src/logical_plan/aexpr/schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ impl AExpr {
use AExpr::*;
use DataType::*;
match self {
Count => Ok(Field::new(COUNT, IDX_DTYPE)),
Len => Ok(Field::new(LEN, IDX_DTYPE)),
Window { function, .. } => {
let e = arena.get(*function);
e.to_field(schema, ctxt, arena)
Expand Down
4 changes: 2 additions & 2 deletions crates/polars-plan/src/logical_plan/conversion.rs
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@ pub fn to_aexpr(expr: Expr, arena: &mut Arena<AExpr>) -> Node {
length: to_aexpr(*length, arena),
},
Expr::Wildcard => AExpr::Wildcard,
Expr::Count => AExpr::Count,
Expr::Len => AExpr::Len,
Expr::Nth(i) => AExpr::Nth(i),
Expr::SubPlan { .. } => panic!("no SQLSubquery expected at this point"),
Expr::KeepName(_) => panic!("no `name.keep` expected at this point"),
Expand Down Expand Up @@ -598,7 +598,7 @@ pub fn node_to_expr(node: Node, expr_arena: &Arena<AExpr>) -> Expr {
offset: Box::new(node_to_expr(offset, expr_arena)),
length: Box::new(node_to_expr(length, expr_arena)),
},
AExpr::Count => Expr::Count,
AExpr::Len => Expr::Len,
AExpr::Nth(i) => Expr::Nth(i),
AExpr::Wildcard => Expr::Wildcard,
}
Expand Down
2 changes: 1 addition & 1 deletion crates/polars-plan/src/logical_plan/format.rs
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,7 @@ impl Debug for Expr {
},
},
Nth(i) => write!(f, "nth({i})"),
Count => write!(f, "count()"),
Len => write!(f, "len()"),
Explode(expr) => write!(f, "{expr:?}.explode()"),
Alias(expr, name) => write!(f, "{expr:?}.alias(\"{name}\")"),
Column(name) => write!(f, "col(\"{name}\")"),
Expand Down
2 changes: 1 addition & 1 deletion crates/polars-plan/src/logical_plan/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ macro_rules! push_expr {
($current_expr:expr, $push:ident, $iter:ident) => {{
use Expr::*;
match $current_expr {
Nth(_) | Column(_) | Literal(_) | Wildcard | Columns(_) | DtypeColumn(_) | Count => {},
Nth(_) | Column(_) | Literal(_) | Wildcard | Columns(_) | DtypeColumn(_) | Len => {},
Alias(e, _) => $push(e),
BinaryExpr { left, op: _, right } => {
// reverse order so that left is popped first
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -317,7 +317,7 @@ impl ExprIdentifierVisitor<'_> {
// TODO! Add a typed null
AExpr::Literal(LiteralValue::Null) => REFUSE_NO_MEMBER,
AExpr::Column(_) | AExpr::Literal(_) | AExpr::Alias(_, _) => REFUSE_ALLOW_MEMBER,
AExpr::Count => {
AExpr::Len => {
if self.is_group_by {
REFUSE_NO_MEMBER
} else {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ pub(super) fn process_group_by(
// Counts change due to groupby's
// TODO! handle aliases, so that the predicate that is pushed down refers to the column before alias.
let mut push_down = !has_aexpr(*predicate, expr_arena, |ae| {
matches!(ae, AExpr::Count | AExpr::Alias(_, _))
matches!(ae, AExpr::Len | AExpr::Alias(_, _))
});

for name in aexpr_to_leaf_names_iter(*predicate, expr_arena) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -530,7 +530,7 @@ impl<'a> PredicatePushDown<'a> {

// a count is influenced by a Union/Vstack
acc_predicates.retain(|_, predicate| {
if has_aexpr(*predicate, expr_arena, |ae| matches!(ae, AExpr::Count)) {
if has_aexpr(*predicate, expr_arena, |ae| matches!(ae, AExpr::Len)) {
local_predicates.push(*predicate);
false
} else {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use super::*;
fn is_count(node: Node, expr_arena: &Arena<AExpr>) -> bool {
match expr_arena.get(node) {
AExpr::Alias(node, _) => is_count(*node, expr_arena),
AExpr::Count => true,
AExpr::Len => true,
_ => false,
}
}
Expand Down
2 changes: 1 addition & 1 deletion crates/polars-plan/src/logical_plan/tree_format.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ impl UpperExp for AExpr {
AExpr::Window { .. } => "window",
AExpr::Wildcard => "*",
AExpr::Slice { .. } => "slice",
AExpr::Count => "count",
AExpr::Len => "len",
AExpr::Nth(v) => return write!(f, "nth({})", v),
};

Expand Down
2 changes: 1 addition & 1 deletion crates/polars-plan/src/logical_plan/visitor/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,7 @@ impl AexprNode {
(Gather { .. }, Gather { .. })
| (Filter { .. }, Filter { .. })
| (Ternary { .. }, Ternary { .. })
| (Count, Count)
| (Len, Len)
| (Slice { .. }, Slice { .. })
| (Explode(_), Explode(_)) => true,
(SortBy { descending: l, .. }, SortBy { descending: r, .. }) => l == r,
Expand Down
Loading