fix(rust, python): strftime with time zone directive

[MarcoGorelli]

closes #6583

[ritchie46]

I don't think we should allocate a string to get the UTC timezone. We probably can just take the UTC timezone from chrono_tz.

[ritchie46]

This is an optimization we could add, it is not a blocker for this PR, but if you feel up to it, you can try to implement it this way.

We can make put conversion_f behind a reference and pass a different conversion_f function dependend of the timezone on the array. That means that we don't branch on timezone every iteration, but only once when we set the conversion_f we will use for this function.

[MarcoGorelli]

good one, thanks for your review!

OK, giving this a go (I've made the other change in the meantime anyway)

[MarcoGorelli]

I've tried factoring the function out, is this what you meant?

I'm running a timing test now to check the performance impact, will post the results when it's done

[MarcoGorelli]

If I've done this right, then it doesn't look like there's any performance difference

Experiments notebook: https://www.kaggle.com/code/marcogorelli/polars-timing?scriptVersionId=118296517

Previous commit:

• 475 ms ± 1.58 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 491 ms ± 16.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 481 ms ± 12.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Most recent commit (with the function factored out):

• 481 ms ± 1.54 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
  -495 ms ± 16.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
  -485 ms ± 6.43 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Here's the script I ran (both times with make build-release):

from IPython import get_ipython
from datetime import datetime
import polars as pl

ipython = get_ipython()


dr = pl.date_range(datetime(1000, 1, 1), datetime(2000, 1, 1), interval='6h', time_zone='US/Pacific')

print(ipython.run_line_magic("timeit", "dr.dt.strftime('%z')"))
print(ipython.run_line_magic("timeit", "dr.dt.strftime('%z')"))
print(ipython.run_line_magic("timeit", "dr.dt.strftime('%z')"))

[MarcoGorelli]

I ran this on a larger input and got consistently better results with the previous commit
https://www.kaggle.com/code/marcogorelli/polars-timing/notebook
I'll revert tomorrow then. If you were thinking of merging, please wait - thanks! 🙏

[ritchie46]

I see you succeeded. Great!

I think we can factor still factor out a branch deeper. All timezone parsing an be done before we create the function.

[ritchie46]

The parse_offset itself can also be factored out.

[MarcoGorelli]

🤔 this is noticeably slower than master https://www.kaggle.com/code/marcogorelli/polars-timing?scriptVersionId=118351982:

here:

• 4.63 s ± 39.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 4.6 s ± 62.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 4.66 s ± 75.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

master:

• 3.77 s ± 39.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 3.74 s ± 59.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 3.73 s ± 13.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

looking into it, but I don't think I'm (yet) advanced enough in Rust to tell why nor what to do about it

[MarcoGorelli]

Right, performance is salvaged (and unexpectedly improved) now, thanks for having pushed me on this, I learned a tonne! 🙏

https://www.kaggle.com/code/marcogorelli/polars-timing/notebook?scriptVersionId=118393630

here:

*** tz-naive ***

• 1.95 s ± 35.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 1.97 s ± 34.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 1.95 s ± 15.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

*** tz-aware ***

• 3.33 s ± 47.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 3.32 s ± 34.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 3.32 s ± 31.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

master:

*** tz-naive ***

• 1.98 s ± 31.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 1.96 s ± 21.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 1.97 s ± 39.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

*** tz-aware ***

• 3.78 s ± 67 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 3.74 s ± 45.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
• 3.76 s ± 73.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[ritchie46]

Great! I see you went with the outer branching approach. Yeap, that should be fastest (without doing crazy things). A function behind a reference can save some compiled code and also remove branches, but the indirection is often costly and it depends if its worth it.

In any case. Looks great!

[MarcoGorelli]

Thanks Ritchie! This was a fun one