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

feat: Add an is_literal method to expression meta namespace #19773

Merged
merged 2 commits into from
Nov 14, 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
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
1 change: 1 addition & 0 deletions py-polars/docs/source/reference/expressions/meta.rst
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ The following methods are available under the `expr.meta` attribute.
Expr.meta.has_multiple_outputs
Expr.meta.is_column
Expr.meta.is_column_selection
Expr.meta.is_literal
Expr.meta.is_regex_projection
Expr.meta.ne
Expr.meta.output_name
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