Skip to content

Commit

Permalink
feat: Add an is_literal method to expression meta namespace
Browse files Browse the repository at this point in the history
  • Loading branch information
alexander-beedie committed Nov 14, 2024
1 parent 9f79100 commit e4246d3
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 6 deletions.
18 changes: 15 additions & 3 deletions crates/polars-plan/src/dsl/meta.rs
Original file line number Diff line number Diff line change
Expand Up @@ -83,9 +83,21 @@ impl MetaNameSpace {
| Expr::IndexColumn(_)
| Expr::Selector(_)
| Expr::Wildcard => true,
Expr::Alias(_, _) | Expr::KeepName(_) | Expr::RenameAlias { .. } if allow_aliasing => {
true
},
Expr::Alias(_, _) | Expr::KeepName(_) | Expr::RenameAlias { .. } => allow_aliasing,
_ => false,
})
}

/// Indicate if this expression represents a literal value (optionally aliased).
pub fn is_literal(&self, allow_aliasing: bool) -> bool {
self.0.into_iter().all(|e| match e {
Expr::Literal(_) => true,
Expr::Alias(_, _) => allow_aliasing,
Expr::Cast {
expr,
dtype: DataType::Datetime(_, _),
options: CastOptions::Strict,
} if matches!(&**expr, Expr::Literal(LiteralValue::DateTime(_, _, _))) => true,
_ => false,
})
}
Expand Down
4 changes: 4 additions & 0 deletions crates/polars-python/src/expr/meta.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,10 @@ impl PyExpr {
.is_column_selection(allow_aliasing)
}

fn meta_is_literal(&self, allow_aliasing: bool) -> bool {
self.inner.clone().meta().is_literal(allow_aliasing)
}

fn _meta_selector_add(&self, other: PyExpr) -> PyResult<PyExpr> {
let out = self
.inner
Expand Down
31 changes: 29 additions & 2 deletions py-polars/polars/expr/meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,15 +108,15 @@ def is_column_selection(self, *, allow_aliasing: bool = False) -> bool:
"""
Indicate if this expression only selects columns (optionally with aliasing).
This can include bare columns, column matches by regex or dtype, selectors
This can include bare columns, columns matched by regex or dtype, selectors
and exclude ops, and (optionally) column/expression aliasing.
.. versionadded:: 0.20.30
Parameters
----------
allow_aliasing
If False (default), any aliasing is not considered pure column selection.
If False (default), any aliasing is not considered to be column selection.
Set True to allow for column selection that also includes aliasing.
Examples
Expand All @@ -142,6 +142,33 @@ def is_column_selection(self, *, allow_aliasing: bool = False) -> bool:
"""
return self._pyexpr.meta_is_column_selection(allow_aliasing)

def is_literal(self, *, allow_aliasing: bool = False) -> bool:
"""
Indicate if this expression is a literal value (optionally aliased).
.. versionadded:: 1.14
Parameters
----------
allow_aliasing
If False (default), only a bare literal will match.
Set True to also allow for aliased literals.
Examples
--------
>>> from datetime import datetime
>>> e = pl.lit(123)
>>> e.meta.is_literal()
True
>>> e = pl.lit(987.654321).alias("foo")
>>> e.meta.is_literal()
False
>>> e = pl.lit(datetime.now()).alias("bar")
>>> e.meta.is_literal(allow_aliasing=True)
True
"""
return self._pyexpr.meta_is_literal(allow_aliasing)

@overload
def output_name(self, *, raise_if_undetermined: Literal[True] = True) -> str: ...

Expand Down
31 changes: 30 additions & 1 deletion py-polars/tests/unit/operations/namespaces/test_meta.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

from typing import TYPE_CHECKING
from datetime import date, datetime, time, timedelta
from typing import TYPE_CHECKING, Any

import pytest

Expand Down Expand Up @@ -123,6 +124,34 @@ def test_is_column_selection(
assert not expr.meta.is_column_selection()


@pytest.mark.parametrize(
"value",
[
None,
1234,
567.89,
float("inf"),
date.today(),
datetime.now(),
time(10, 30, 45),
timedelta(hours=-24),
["x", "y", "z"],
pl.Series([None, None]),
[[10, 20], [30, 40]],
"this is the way",
],
)
def test_is_literal(value: Any) -> None:
e = pl.lit(value)
assert e.meta.is_literal()

e = pl.lit(value).alias("foo")
assert not e.meta.is_literal()

e = pl.lit(value).alias("foo")
assert e.meta.is_literal(allow_aliasing=True)


def test_meta_is_regex_projection() -> None:
e = pl.col("^.*$").name.suffix("_foo")
assert e.meta.is_regex_projection()
Expand Down

0 comments on commit e4246d3

Please sign in to comment.