feat: Add `cat.len_chars` and `cat.len_bytes` #20211

mcrumiller · 2024-12-07T16:38:43Z

This adds a fast path for operations that can be performed on the categories of a categorical series. I've added len_bytes and len_chars for now. If desired, it would be trivial to add a few others such as starts_with and ends_with.

import polars as pl

with pl.StringCache():
    pl.Series(["a", "b"], dtype=pl.Categorical)  # fill some cache
    df = pl.DataFrame({
        "a": pl.Series(["Café", "345", "Café", "東京", None], dtype=pl.Categorical)
    })

df.with_columns(
    pl.col("a").cat.len_bytes().alias("n_bytes"),
    pl.col("a").cat.len_chars().alias("n_chars"),
)
# shape: (5, 3)
# ┌──────┬─────────┬─────────┐
# │ a    ┆ n_bytes ┆ n_chars │
# │ ---  ┆ ---     ┆ ---     │
# │ cat  ┆ u32     ┆ u32     │
# ╞══════╪═════════╪═════════╡
# │ Café ┆ 5       ┆ 4       │
# │ 345  ┆ 3       ┆ 3       │
# │ Café ┆ 5       ┆ 4       │
# │ 東京 ┆ 6       ┆ 2       │
# │ null ┆ null    ┆ null    │
# └──────┴─────────┴─────────┘

These require the strings feature because they dispatch to str.len_bytes and str.len_chars.

codecov · 2024-12-07T17:01:22Z

Codecov Report

Attention: Patch coverage is 96.55172% with 2 lines in your changes missing coverage. Please review.

Project coverage is 79.63%. Comparing base (d4bab3f) to head (2b9edc3).
Report is 12 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/polars-plan/src/dsl/function_expr/cat.rs	94.73%	2 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main   #20211   +/-   ##
=======================================
  Coverage   79.62%   79.63%           
=======================================
  Files        1565     1565           
  Lines      218187   218243   +56     
  Branches     2475     2475           
=======================================
+ Hits       173734   173794   +60     
+ Misses      43886    43882    -4     
  Partials      567      567

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

crates/polars-python/src/lazyframe/visitor/expr_nodes.rs

ritchie46

Looks great @mcrumiller. I've left some comments.

ritchie46 · 2024-12-08T07:35:49Z

crates/polars-plan/src/dsl/function_expr/cat.rs

+    T: PolarsDataType,
+{
+    let ca = s.categorical()?;
+    let (categories, phys) = match &**ca.get_rev_map() {


Can this move into its own function, that saves monomorphizaton bloat.

Does this resolve it? 2b9edc3

Yes, great!

ritchie46 · 2024-12-08T07:36:44Z

crates/polars-plan/src/dsl/function_expr/cat.rs

+    };
+
+    // Apply function to categories
+    let categories = StringChunked::with_chunk(PlSmallStr::EMPTY, categories.clone());


This should take the name of s.

ritchie46 · 2024-12-08T07:36:59Z

crates/polars-plan/src/dsl/function_expr/cat.rs

+    let categories = StringChunked::with_chunk(PlSmallStr::EMPTY, categories.clone());
+    let result = op(&categories).into_series();
+
+    let out = result.take(phys.idx()?)?;


This can do a take_unchecked.

ritchie46 · 2024-12-08T07:38:55Z

crates/polars-python/src/lazyframe/visitor/expr_nodes.rs

@@ -171,6 +171,21 @@ impl PyStringFunction {
    }
 }

+#[pyclass(name = "CategoricalFunction", eq)]


Can you remove these. I don't think we need those for now.

Reverted the whole file, 0a83dbe

ritchie46 · 2024-12-08T07:39:10Z

crates/polars-python/src/lazyframe/visitor/expr_nodes.rs

@@ -793,8 +808,16 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
                FunctionExpr::BinaryExpr(_) => {
                    return Err(PyNotImplementedError::new_err("binary expr"))
                },
-                FunctionExpr::Categorical(_) => {
-                    return Err(PyNotImplementedError::new_err("categorical expr"))
+                FunctionExpr::Categorical(catfun) => match catfun {


This can remain NotImplemented.

Reverted the whole file, 0a83dbe

alexander-beedie · 2024-12-09T12:33:03Z

If desired, it would be trivial to add a few others such as starts_with and ends_with.

Good stuff. Bringing cat to parity with str while also speeding it up is a worthy goal; I can say for sure that every single researcher at our place believes (rightly or wrongly) that cat is a pure string optimisation and that it should behave as a string in essentially all circumstances... (the fact that it doesn't sort as one by default is an ongoing pain point, but that's a whole other discussion 🤣)

I'd certainly be happy to see more fast-path str ops come to cat 👍

mcrumiller · 2024-12-09T15:06:01Z

Good stuff. Bringing cat to parity with str while also speeding it up is a worthy goal; I can say for sure that every single researcher at our place believes (rightly or wrongly) that cat is a pure string optimisation and that it should behave as a string in essentially all circumstances... (the fact that it doesn't sort as one by default is an ongoing pain point, but that's a whole other discussion 🤣)

@alexander-beedie as per the #9773 discussion I don't think we will get most string ops in the cat namespace unfortunately. I'm with your researchers here in that I would have loved a .cat.slice() but...it was not meant to be. But a few of these easy ones are nice that don't require creating a new revmap.

ritchie46 · 2024-12-10T09:21:03Z

@alexander-beedie as per the #9773 discussion I don't think we will get most string ops in the cat namespace unfortunately. I'm with your researchers here in that I would have loved a .cat.slice() but...it was not meant to be. But a few of these easy ones are nice that don't require creating a new revmap.

I think we should just map shrink_to_fit to a kernel that interns the strings. Then long strings will be deduplicated.

connor-elliott · 2024-12-10T15:45:53Z

+1 for starts_with/ends_with/contains

mcrumiller · 2024-12-10T15:51:01Z

@connor-elliott that should now be an easy add, I'll take a look tonight when I'm home from work.

mcrumiller · 2024-12-11T03:36:25Z

I think we should just map shrink_to_fit to a kernel that interns the strings. Then long strings will be deduplicated.

The reason we often prefer categoricals is because they take up about 1/4 the space (4 bytes per element versus 16), not necessarily for performance reasons. In many of these cases, it would really nice to have a pl.Char datatype, or something similar to pl.Array(pl.UInt8, N) with fixed N.

ritchie46 · 2024-12-11T05:49:52Z

I think we should just map shrink_to_fit to a kernel that interns the strings. Then long strings will be deduplicated.

The reason we often prefer categoricals is because they take up about 1/4 the space (4 bytes per element versus 16), not necessarily for performance reasons. In many of these cases, it would really nice to have a pl.Char datatype, or something similar to pl.Array(pl.UInt8, N) with fixed N.

And that will be fixed by streaming. Where categoricals are yet again problematic and complex. They seem like a good solution. But are a band aid. There are better solutions;)

github-actions bot added enhancement New feature or an improvement of an existing feature python Related to Python Polars rust Related to Rust Polars labels Dec 7, 2024

mcrumiller marked this pull request as ready for review December 7, 2024 17:15

mcrumiller requested review from ritchie46, c-peters, alexander-beedie, MarcoGorelli, reswqa, wence- and orlp as code owners December 7, 2024 17:15

mcrumiller marked this pull request as draft December 7, 2024 17:48

mcrumiller force-pushed the cat-len branch from a2baa96 to 1d25949 Compare December 7, 2024 18:03

mcrumiller marked this pull request as ready for review December 7, 2024 20:11

mcrumiller added 2 commits December 7, 2024 16:08

Add fast path for categorical operations

cb69508

More testing

10faaa5

mcrumiller force-pushed the cat-len branch from 1d25949 to 10faaa5 Compare December 7, 2024 22:00

mcrumiller commented Dec 8, 2024

View reviewed changes

crates/polars-python/src/lazyframe/visitor/expr_nodes.rs Outdated Show resolved Hide resolved

ritchie46 reviewed Dec 8, 2024

View reviewed changes

ritchie46 changed the title ~~feat: Add cat.str_len and cat.str_bytes~~ feat: Add cat.len_chars and cat.len_bytes Dec 8, 2024

mcrumiller marked this pull request as draft December 8, 2024 13:31

mcrumiller added 2 commits December 8, 2024 08:37

Restore expr nodes

0a83dbe

Move non-generic to separate function

2b9edc3

mcrumiller marked this pull request as ready for review December 8, 2024 14:40

ritchie46 approved these changes Dec 10, 2024

View reviewed changes

ritchie46 merged commit a21ac2e into pola-rs:main Dec 10, 2024
26 checks passed

mcrumiller deleted the cat-len branch December 10, 2024 15:50

mcrumiller mentioned this pull request Dec 11, 2024

feat: Add cat.starts_with/cat.ends_with #20257

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Add `cat.len_chars` and `cat.len_bytes` #20211

feat: Add `cat.len_chars` and `cat.len_bytes` #20211

mcrumiller commented Dec 7, 2024 •

edited

Loading

codecov bot commented Dec 7, 2024 •

edited

Loading

ritchie46 left a comment

ritchie46 Dec 8, 2024

mcrumiller Dec 8, 2024

ritchie46 Dec 10, 2024

ritchie46 Dec 8, 2024

mcrumiller Dec 8, 2024

ritchie46 Dec 8, 2024

ritchie46 Dec 8, 2024

mcrumiller Dec 8, 2024

ritchie46 Dec 8, 2024

mcrumiller Dec 8, 2024

ritchie46 Dec 9, 2024

alexander-beedie commented Dec 9, 2024 •

edited

Loading

mcrumiller commented Dec 9, 2024 •

edited

Loading

ritchie46 commented Dec 10, 2024

connor-elliott commented Dec 10, 2024

mcrumiller commented Dec 10, 2024

mcrumiller commented Dec 11, 2024

ritchie46 commented Dec 11, 2024

feat: Add cat.len_chars and cat.len_bytes #20211

feat: Add cat.len_chars and cat.len_bytes #20211

Conversation

mcrumiller commented Dec 7, 2024 • edited Loading

codecov bot commented Dec 7, 2024 • edited Loading

Codecov Report

ritchie46 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alexander-beedie commented Dec 9, 2024 • edited Loading

mcrumiller commented Dec 9, 2024 • edited Loading

ritchie46 commented Dec 10, 2024

connor-elliott commented Dec 10, 2024

mcrumiller commented Dec 10, 2024

mcrumiller commented Dec 11, 2024

ritchie46 commented Dec 11, 2024

feat: Add `cat.len_chars` and `cat.len_bytes` #20211

feat: Add `cat.len_chars` and `cat.len_bytes` #20211

mcrumiller commented Dec 7, 2024 •

edited

Loading

codecov bot commented Dec 7, 2024 •

edited

Loading

alexander-beedie commented Dec 9, 2024 •

edited

Loading

mcrumiller commented Dec 9, 2024 •

edited

Loading