feat: Support executing polars SQL against pandas and pyarrow objects

[alexander-beedie]

@ritchie46 - as suggested 😉

Notably expands the scope of the Polars SQL interface .

• The pl.sql function can now transparently identify and register additional frame/object types when referenced from a Polars SQL query.
• SQLContext can also now operate on these additional types.

Compatible objects include:

• Polars DataFrame, LazyFrame, and (new) Series.
• Pandas DataFrame and Series.
• PyArrow Table and RecordBatch.

(💡If you have more types that you think would be a good fit for this, let me know).

Objects of these types are transparently converted/registered as LazyFrame iff their variable name appears in the SQL query (this prevents unnecessary conversions of compatible objects that aren't actually referenced).

Example

import pandas as pd
import pyarrow as pa
import polars as pl

# polars
pl_frame = pl.LazyFrame({
  "a": [1, 2, 3],
  "b": [6, 7, 8],
  "c": ["z", "y", "x"],
})

# pandas
pd_frame = pd.DataFrame({
  "a": [2, 3, 4],
  "d": [-0.5, 0.0, 0.5],
})

# pyarrow
pa_table = pa.Table.from_arrays([
    pa.array([1, 2, 3]),
    pa.array(["x", "y", "z"])
  ],
  names=["a", "e"],
)

Join polars LazyFrame with a pandas DataFrame and a pyarrow Table:

pl.sql("""
  SELECT pl_frame.*, d, e
    FROM pl_frame
    JOIN pd_frame USING(a)
    JOIN pa_table USING(a)
""").collect()

# shape: (2, 5)
# ┌─────┬─────┬─────┬──────┬─────┐
# │ a   ┆ b   ┆ c   ┆ d    ┆ e   │
# │ --- ┆ --- ┆ --- ┆ ---  ┆ --- │
# │ i64 ┆ i64 ┆ str ┆ f64  ┆ str │
# ╞═════╪═════╪═════╪══════╪═════╡
# │ 2   ┆ 7   ┆ y   ┆ -0.5 ┆ y   │
# │ 3   ┆ 8   ┆ x   ┆ 0.0  ┆ z   │
# └─────┴─────┴─────┴──────┴─────┘

[codecov]

Codecov Report

Attention: Patch coverage is 82.14286% with 10 lines in your changes missing coverage. Please review.

Project coverage is 81.45%. Comparing base (6f3fd8e) to head (6027482).

Files	Patch %	Lines
py-polars/polars/_utils/various.py	50.00%	3 Missing and 4 partials ⚠️
py-polars/polars/sql/context.py	92.68%	2 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #16746      +/-   ##
==========================================
- Coverage   81.45%   81.45%   -0.01%     
==========================================
  Files        1413     1413              
  Lines      186306   186343      +37     
  Branches     2777     2784       +7     
==========================================
+ Hits       151750   151780      +30     
- Misses      34036    34040       +4     
- Partials      520      523       +3     

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

[lucazanna]

Outstanding

Could Spark dataframes be included too?

I never remember the syntax to convert Spark to Polars (below Stack Overflow answer from Ritchie)
https://stackoverflow.com/questions/73203318/how-to-transform-spark-dataframe-to-polars-dataframe

It would make it quite easy to use Polars for any Spark dataframe that fits into memory

[alexander-beedie]

Could Spark dataframes be included too?

Unfortunately that conversion uses a private method (_collect_as_arrow ), which would be a little dicey to rely on in our own API - if PySpark added an "official" Arrow conversion path then we could definitely think about it though✌️

[lucazanna]

Hi @alexander-beedie

It looks like it was just added to Spark
apache/spark#45481

Better if I open a separate request for this?

[ritchie46]

Collecting a pyspark dataframe is also very dangerous. We would force it to bring all the data to one machine. This should not be done implictly.

[alexander-beedie]

Collecting a pyspark dataframe is also very dangerous. We would force it to bring all the data to one machine. This should not be done implictly.

Good point; a from_spark would probably make more sense if we wanted to go that way if there are some meaningful options there (above & beyond what you'd get from just passing the new pyarrow Table object into from_arrow). Or maybe they'll consider adding a native toPolars now they expose Arrow properly ;)