feat: Add SQL support for bit_count and bitwise &, |, and xor operators

[alexander-beedie]

Note that there is no consensus/standard for use of an underscore in SQL BIT* function names¹, so we support it with/without for the most straightforward compatibility (eg: BITAND and BIT_AND are both allowed).

Drive-by: added merge_sorted feature gate in polars-stream to fix docs build (see core-dev chat).

Example

import polars as pl

df = pl.DataFrame({
    "x": [4, 32, 88, 128],
    "y": [-8, 0, -1, None],
})

Apply bitwise ops via the SQL interface:

df.sql("""
    SELECT
      x, y,
      BIT_COUNT(x) AS x_bits_set,
      x XOR y AS x_bitxor_y,
      x & y AS x_bitand_y,
      x | y AS x_bitor_y,
    FROM self
""")
# shape: (4, 6)
# ┌─────┬──────┬────────────┬────────────┬────────────┬───────────┐
# │ x   ┆ y    ┆ x_bits_set ┆ x_bitxor_y ┆ x_bitand_y ┆ x_bitor_y │
# │ --- ┆ ---  ┆ ---        ┆ ---        ┆ ---        ┆ ---       │
# │ i64 ┆ i64  ┆ u32        ┆ i64        ┆ i64        ┆ i64       │
# ╞═════╪══════╪════════════╪════════════╪════════════╪═══════════╡
# │ 4   ┆ -8   ┆ 1          ┆ -4         ┆ 0          ┆ -4        │
# │ 32  ┆ 0    ┆ 1          ┆ 32         ┆ 0          ┆ 32        │
# │ 88  ┆ -1   ┆ 3          ┆ -89        ┆ 88         ┆ -1        │
# │ 128 ┆ null ┆ 1          ┆ null       ┆ null       ┆ null      │
# └─────┴──────┴────────────┴────────────┴────────────┴───────────┘

Footnotes

1. PostgreSQL supports BIT_COUNT as a function, but BIT_AND and BIT_OR are available only as operators. However, multiple other dialects (Oracle, etc) support both the operator and function form, so we follow suit here. ↩

[codecov]

Codecov Report

Attention: Patch coverage is 89.74359% with 4 lines in your changes missing coverage. Please review.

Project coverage is 80.23%. Comparing base (dbc8368) to head (4407752).
Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/polars-sql/src/functions.rs	73.33%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #19114      +/-   ##
==========================================
- Coverage   80.24%   80.23%   -0.02%     
==========================================
  Files        1523     1523              
  Lines      209524   209538      +14     
  Branches     2434     2434              
==========================================
- Hits       168127   168113      -14     
- Misses      40842    40870      +28     
  Partials      555      555              

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

[ritchie46]

Is ci failure related? I see the merge sorted feature activated?

[alexander-beedie]

Is ci failure related? I see the merge sorted feature activated?

Seems related to the activation of the bitwise feature flag somehow; looks like that flag is implicitly entangled with the merge_sorted flag, but I don't see exactly where (maybe in lazy/core?) 🤔

@orlp, can you see how it sneaks in, by any chance?
(I just rebased so that everything's up to date).

[orlp]

@alexander-beedie There is a known issue with the current implementation of merge_sorted with the new streaming implementation. It concatenates two dataframes which are of different lengths, hardbaked into the function IR even. Can you add a @pytest.mark.may_fail_auto_streaming in front of that test?

[alexander-beedie]

Can you add a @pytest.mark.may_fail_auto_streaming in front of that test?

Done!