feat: Add SQL support for NATURAL joins and the COLUMNS function

[alexander-beedie]

New

• Adds NATURAL [INNER|LEFT|FULL] JOIN¹ support for the SQL interface (this flavour of join automatically uses the common cols between the two tables as the join-key), and the associated SQL docs entries/examples.
• Partial implementation of the useful DuckDB COLUMNS² function; this maps almost 1:1 onto Polars' col expression, enabling use of regular expressions to select multiple columns. Also offers some (currently limited) broadcasting capabilities; at the moment broadcasting support only works in the WHERE clause.

Also

• Ensure SQLFunctionVisitor translations also have access to the current/active schema.
• Some refactoring, with more to come later; the logic is getting quite involved, and could do with more inline comments, explanations and additional cleanups.

Examples

Establish test data:

import polars as pl

df1 = pl.DataFrame({
    "CharacterID": [1, 2, 3],
    "FirstName": ["Jernau Morat", "Cheradenine", "Byr"],
    "LastName": ["Gurgeh", "Zakalwe", "Genar-Hofoen"],
})
df2 = pl.DataFrame({
    "CharacterID": [1, 2, 3],
    "Book": ["Player of Games", "Use of Weapons", "Excession"],
})
df3 = pl.DataFrame({
    "CharacterID": [1, 2, 3],
    "Ship": ["Limiting Factor", "Xenophobe", "Grey Area"],
})

Apply NATURAL joins, use a COLUMNS regex to avoid selecting <col>:<suffix> cols:

pl.sql("""
  SELECT COLUMNS('^[^:]+$')
  FROM df1
    NATURAL JOIN df2
    NATURAL JOIN df3
  ORDER BY CharacterID
""").collect()

# shape: (3, 5)
# ┌─────────────┬──────────────┬──────────────┬─────────────────┬─────────────────┐
# │ CharacterID ┆ FirstName    ┆ LastName     ┆ Book            ┆ Ship            │
# │ ---         ┆ ---          ┆ ---          ┆ ---             ┆ ---             │
# │ i64         ┆ str          ┆ str          ┆ str             ┆ str             │
# ╞═════════════╪══════════════╪══════════════╪═════════════════╪═════════════════╡
# │ 1           ┆ Jernau Morat ┆ Gurgeh       ┆ Player of Games ┆ Limiting Factor │
# │ 2           ┆ Cheradenine  ┆ Zakalwe      ┆ Use of Weapons  ┆ Xenophobe       │
# │ 3           ┆ Byr          ┆ Genar-Hofoen ┆ Excession       ┆ Grey Area       │
# └─────────────┴──────────────┴──────────────┴─────────────────┴─────────────────┘

Use COLUMNS to broadcast constraints in WHERE clause; filter for value >= 4 across all columns:

df1 = pl.DataFrame({
    "x": [1, 5, 3, 8, 6, 7, 4, 0, 2],
    "y": [3, 4, 6, 8, 3, 4, 1, 7, 8],
})
df2 = pl.DataFrame({
    "y": [0, 4, 0, 8, 0, 4, 0, 7, 0],
    "z": [9, 8, 7, 6, 5, 4, 3, 2, 1],
})

pl.sql(f"""
  SELECT * EXCLUDE "y:df2"
  FROM df1 NATURAL JOIN df2
  WHERE COLUMNS(*) >= 4
""").collect()

# shape: (5, 3)
# ┌─────┬─────┬─────┐
# │ x   ┆ y   ┆ z   │
# │ --- ┆ --- ┆ --- │
# │ i64 ┆ i64 ┆ i64 │
# ╞═════╪═════╪═════╡
# │ 5   ┆ 4   ┆ 8   │
# │ 7   ┆ 4   ┆ 8   │
# │ 8   ┆ 8   ┆ 6   │
# │ 5   ┆ 4   ┆ 4   │
# │ 7   ┆ 4   ┆ 4   │
# └─────┴─────┴─────┘

Footnotes

1. NATURAL JOIN: https://www.postgresql.org/docs/current/queries-table-expressions.html#QUERIES-JOIN ↩

2. COLUMNS: https://duckdb.org/docs/sql/expressions/star.html#columns-regular-expression ↩

[codecov]

Codecov Report

Attention: Patch coverage is 87.62887% with 36 lines in your changes missing coverage. Please review.

Project coverage is 80.71%. Comparing base (d444b79) to head (1aaa1fb).
Report is 1 commits behind head on main.

Files	Patch %	Lines
crates/polars-sql/src/context.rs	90.56%	20 Missing ⚠️
crates/polars-sql/src/functions.rs	78.94%	16 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #17295      +/-   ##
==========================================
- Coverage   80.72%   80.71%   -0.01%     
==========================================
  Files        1475     1475              
  Lines      193162   193260      +98     
  Branches     2751     2751              
==========================================
+ Hits       155922   155993      +71     
- Misses      36732    36759      +27     
  Partials      508      508              

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