feat(python): Add a low-friction sql method for DataFrame and LazyFrame

[alexander-beedie]

Introduces a new sql() method for both DataFrame and LazyFrame. Makes use of SQLContext internally, streamlining the usage of ad-hoc SQL on individual frames.

Features

• Automatically registers the calling frame as "self" in the SQL context.
• Plenty of examples and docstring detail.

from datetime import date
import polars as pl

df = pl.DataFrame({
    "a": [1, 2, 3],
    "b": ["zz", "yy", "xx"],
    "c": [date(1999, 12, 31), date(2010, 10, 10), date(2077, 8, 8)],
})

df.sql("SELECT a::float4, b, EXTRACT(year FROM c) AS c_year FROM self WHERE a > 1")
# shape: (2, 3)
# ┌─────┬─────┬────────┐
# │ a   ┆ b   ┆ c_year │
# │ --- ┆ --- ┆ ---    │
# │ f32 ┆ str ┆ i32    │
# ╞═════╪═════╪════════╡
# │ 2.0 ┆ yy  ┆ 2010   │
# │ 3.0 ┆ xx  ┆ 2077   │
# └─────┴─────┴────────┘

Also

• Recognises additional PostgreSQL float casts: ::float4 (32bit) and ::float8 (64bit) and ::float(n) (adaptive).
• Allows the full range of ::timestamp(n) precision modifiers (1-6).
• Fixes an (unlikely) edge-case in table name registration.

[codspeed-hq]

CodSpeed Performance Report

Merging #15783 will not alter performance

_{Comparing alexander-beedie:frame-level-sql-query (1bd92d3) with main (937fd46)}

Summary

✅ 22 untouched benchmarks

[stinodego]

Makes a lot of sense to be able to execute SQL queries directly on a frame! Very happy with this addition.

I'll put on my nitpick hat and go through this a bit later, but definitely a good move in my opinion👍

[alexander-beedie]

Fixing the failing test and further ensuring float-type parsing is even more conformant with PostgreSQL - been reading their float¹ documentation in more detail and spotted some notes about "float(n)" definition 👀

Footnotes

1. https://www.postgresql.org/docs/current/datatype-numeric.html#DATATYPE-FLOAT ↩

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 80.39%. Comparing base (d11da5e) to head (1bd92d3).
Report is 4 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #15783      +/-   ##
==========================================
- Coverage   81.37%   80.39%   -0.98%     
==========================================
  Files        1379     1264     -115     
  Lines      176843   165428   -11415     
  Branches     2543        0    -2543     
==========================================
- Hits       143908   133001   -10907     
+ Misses      32452    32427      -25     
+ Partials      483        0     -483     

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

[ritchie46]

This looks good to me. I don't think it contains anything controversial, so I will go on and merge this. ;)

Love the functionality!