feat(rust, python): expressify offset and length parameters for str.slice

[cmdlineluser]

Attempts to resolve #10890

I see it had been previously self-assigned @orlp - hopefully I'm not intruding here.

(I imagine you had more important things to work on.)

df = pl.DataFrame({
   "foo": ["abcde", "fghijk"],
   "offset": [1, 2],
   "length": [3, 4],
})

df.with_columns(slice = pl.col("foo").str.slice("offset", "length"))

# shape: (2, 4)
# ┌────────┬────────┬────────┬───────┐
# │ foo    ┆ offset ┆ length ┆ slice │
# │ ---    ┆ ---    ┆ ---    ┆ ---   │
# │ str    ┆ i64    ┆ i64    ┆ str   │
# ╞════════╪════════╪════════╪═══════╡
# │ abcde  ┆ 1      ┆ 3      ┆ bcd   │
# │ fghijk ┆ 2      ┆ 4      ┆ hijk  │
# └────────┴────────┴────────┴───────┘

I used the strip_prefix impl as a guide/template https://github.com/pola-rs/polars/blob/main/crates/polars-ops/src/chunked_array/strings/strip.rs

Some things I'm not sure about if anybody wants to give feedback:

• Instead of all the length checks, I just broadcasted offset and length using .clear().extend_constant() - is this an acceptable approach or is it a bad idea + performance concern?
• Is there a way to "typecheck" the params, or is it okay to just .cast()?
• Is UInt64 the correct type for length?
• The param name is offset but internally start is used, I changed everything to offset, is this the right approach?
• I had to modify polars-sql as it used .slice() (incidentally, that uses start as the param name, should it be start everywhere instead of offset?)
• Is it okay to use !matches!?
• As length is a param name, I had to re-jig the error format a bit to avoid ending up with length value length.

[ritchie46]

I don't want to add broadcasting by extending memory. See discussion here; #11900

[cmdlineluser]

Ah! Excellent, thanks @ritchie46 - that explains the need for the added length checks / fallback logic.

I'll switch it around.

[cmdlineluser]

binary_elementwise(ca, length, |opt_str_val, opt_length| utf8_substring_ternary(opt_str_val, offset.get(0), opt_length))

Does anybody know what I'm doing wrong here?

error: implementation of `FnOnce` is not general enough
  --> /Users/user/git/polars/crates/polars-ops/src/chunked_array/strings/substring.rs:63:5
   |
63 |     binary_elementwise(ca, length, |opt_str_val, opt_length| utf8_substring_ternary(opt_str_val, offset.get(0), opt_length))
   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ implementation of `FnOnce` is not general enough
   |
   = note: closure with signature `fn(std::option::Option<&'2 str>, std::option::Option<u64>) -> std::option::Option<&str>` must implement `FnOnce<(std::option::Option<&'1 str>, std::option::Option<u64>)>`, for any lifetime `'1`...
   = note: ...but it actually implements `FnOnce<(std::option::Option<&'2 str>, std::option::Option<u64>)>`, for some specific lifetime `'2`

[orlp]

@cmdlineluser You're not doing anything wrong, the compiler is just inferring stuff wrong here. Let me check if I can rework this...

[orlp]

@cmdlineluser Unfortunately this is one place where we must help the compiler a bit. This works:

match (offset.len(), length.len()) {
    (1, 1) => ca.apply_generic(|opt_str_val| {
        utf8_substring_ternary(opt_str_val, offset.get(0), length.get(0))
    }),
    (1, _) => {
        fn infer<F: for<'a> FnMut(Option<&'a str>, Option<u64>) -> Option<&'a str>>(f: F) -> F where
        {
            f
        }

        let off = offset.get(0);
        binary_elementwise(
            ca,
            length,
            infer(|val, len| utf8_substring_ternary(val, off, len)),
        )
    },
    (_, 1) => {
        fn infer<F: for<'a> FnMut(Option<&'a str>, Option<i64>) -> Option<&'a str>>(f: F) -> F where
        {
            f
        }

        let len = length.get(0);
        binary_elementwise(
            ca,
            offset,
            infer(|val, off| utf8_substring_ternary(val, off, len)),
        )
    },
    _ => ternary_elementwise(ca, offset, length, utf8_substring_ternary),
}

[cmdlineluser]

Ah, thank you @orlp

So that's what the infer_re_match() is doing in some of the other examples:

polars/crates/polars-ops/src/chunked_array/strings/namespace.rs

Line 137 in a366bc9

infer_re_match(|src, pat| {

polars/crates/polars-ops/src/chunked_array/strings/namespace.rs

Lines 17 to 24 in a366bc9

	// We need this to infer the right lifetimes for the match closure.
	#[inline(always)]
	fn infer_re_match<F>(f: F) -> F
	where
	F: for<'a, 'b> FnMut(Option<&'a str>, Option<&'b str>) -> Option<bool>,
	{
	f
	}

[orlp]

@cmdlineluser Yes, that's why. Perhaps we should be using the https://crates.io/crates/higher-order-closure crates for this.

[reswqa]

Perhaps we should be using the https://crates.io/crates/higher-order-closure crates for this.

It's a real pain fighting the compiler here, it would be awesome if this crates could save us. 😞

[reswqa]

Hi @cmdlineluser. Thank you for this nice work! If you don't have free time, can I take over and keep pushing it as we have some future work that depends on this one. Of course, I will preserve your authorship. :D)

[cmdlineluser]

@reswqa Please feel free to take over - thank you.

[cmdlineluser]

Superseded by #13747