Skip to content

介绍

虽然Polars支持与SQL的交互,但建议用户熟悉表达式语法,以生成更具可读性和表达性的代码。由于DataFrame接口是主要的,新功能通常首先添加到表达式API中。然而,如果您已经有一个现有的SQL代码库或更喜欢使用SQL,Polars也提供了对此的支持。

注意

没有单独的SQL引擎,因为Polars将SQL查询转换为表达式,然后使用其自己的引擎执行。这种方法确保了Polars作为原生DataFrame库保持其性能和可扩展性优势,同时仍然为用户提供了使用SQL的能力。

上下文

Polars 使用 SQLContext 对象来管理 SQL 查询。该上下文包含 DataFrameLazyFrame 标识符名称到其对应数据集的映射1。下面的示例启动了一个 SQLContext

SQLContext

ctx = pl.SQLContext()

注册数据框

SQLContext初始化期间,有几种方法可以注册DataFrames。

  • 在全局命名空间中注册所有LazyFrameDataFrame对象。
  • 通过字典映射或kwargs显式注册。

SQLContext

df = pl.DataFrame({"a": [1, 2, 3]})
lf = pl.LazyFrame({"b": [4, 5, 6]})

# 在全局命名空间中注册所有数据框:注册 "df" 和 "lf"
ctx = pl.SQLContext(register_globals=True)

# 注册一个明确的标识符名称到数据框的映射
ctx = pl.SQLContext(frames={"table_one": df, "table_two": lf})

# 使用 kwargs 注册数据框;数据框 df 作为 "df",惰性数据框 lf 作为 "lf"
ctx = pl.SQLContext(df=df, lf=lf)

我们也可以通过先将Pandas DataFrames转换为Polars来注册它们。

SQLContext

import pandas as pd

df_pandas = pd.DataFrame({"c": [7, 8, 9]})
ctx = pl.SQLContext(df_pandas=pl.from_pandas(df_pandas))

注意

将基于Numpy的Pandas DataFrame转换为基于Arrow的DataFrame可能会触发一个潜在的高成本转换;然而,如果Pandas DataFrame已经基于Arrow,那么转换成本将显著降低(在某些情况下几乎免费)。

一旦SQLContext初始化完成,我们可以注册额外的Dataframes或注销现有的Dataframes,使用以下方法:

  • register
  • register_globals
  • register_many
  • unregister

执行查询并收集结果

SQL查询总是以懒加载模式执行,以充分利用查询规划优化的全部优势,因此我们有两种收集结果的选择:

  • 将参数 eager_execution 设置为 True 在 SQLContext 中;这确保 Polars 自动从 execute 调用中收集 LazyFrame 结果。
  • 在执行查询时,将参数eager设置为True,或者使用collect显式收集结果。

我们通过在SQLContext上调用execute来执行SQL查询。

register · execute

# 对于本地文件,请使用 scan_csv 代替
pokemon = pl.read_csv(
    "https://gist.githubusercontent.com/ritchie46/cac6b337ea52281aa23c049250a4ff03/raw/89a957ff3919d90e6ef2d34235e6bf22304f3366/pokemon.csv"
)
with pl.SQLContext(register_globals=True, eager=True) as ctx:
    df_small = ctx.execute("SELECT * from pokemon LIMIT 5")
    print(df_small)

shape: (5, 13)
┌─────┬───────────────────────┬────────┬────────┬───┬─────────┬───────┬────────────┬───────────┐
│ #   ┆ Name                  ┆ Type 1 ┆ Type 2 ┆ … ┆ Sp. Def ┆ Speed ┆ Generation ┆ Legendary │
│ --- ┆ ---                   ┆ ---    ┆ ---    ┆   ┆ ---     ┆ ---   ┆ ---        ┆ ---       │
│ i64 ┆ str                   ┆ str    ┆ str    ┆   ┆ i64     ┆ i64   ┆ i64        ┆ bool      │
╞═════╪═══════════════════════╪════════╪════════╪═══╪═════════╪═══════╪════════════╪═══════════╡
│ 1   ┆ Bulbasaur             ┆ Grass  ┆ Poison ┆ … ┆ 65      ┆ 45    ┆ 1          ┆ false     │
│ 2   ┆ Ivysaur               ┆ Grass  ┆ Poison ┆ … ┆ 80      ┆ 60    ┆ 1          ┆ false     │
│ 3   ┆ Venusaur              ┆ Grass  ┆ Poison ┆ … ┆ 100     ┆ 80    ┆ 1          ┆ false     │
│ 3   ┆ VenusaurMega Venusaur ┆ Grass  ┆ Poison ┆ … ┆ 120     ┆ 80    ┆ 1          ┆ false     │
│ 4   ┆ Charmander            ┆ Fire   ┆ null   ┆ … ┆ 50      ┆ 65    ┆ 1          ┆ false     │
└─────┴───────────────────────┴────────┴────────┴───┴─────────┴───────┴────────────┴───────────┘

从多个来源执行查询

SQL查询可以轻松地从多个来源执行。在下面的示例中,我们注册:

  • 一个CSV文件(延迟加载)
  • 一个NDJSON文件(延迟加载)
  • 一个 Pandas 数据框

并使用SQL将它们连接在一起。惰性读取允许仅从文件中加载必要的行和列。

同样地,可以注册云数据湖(S3,Azure Data Lake)。一个PyArrow数据集可以指向数据湖,然后Polars可以使用scan_pyarrow_dataset读取它。

register · execute

# 输入数据:
# products_masterdata.csv 的 schema 为 {'product_id': Int64, 'product_name': String}
# products_categories.json 的 schema 为 {'product_id': Int64, 'category': String}
# sales_data 是一个 Pandas DataFrame,schema 为 {'product_id': Int64, 'sales': Int64}

with pl.SQLContext(
    products_masterdata=pl.scan_csv("docs/assets/data/products_masterdata.csv"),
    products_categories=pl.scan_ndjson("docs/assets/data/products_categories.json"),
    sales_data=pl.from_pandas(sales_data),
    eager=True,
) as ctx:
    query = """
    SELECT
        product_id,
        product_name,
        category,
        sales
    FROM
        products_masterdata
    LEFT JOIN products_categories USING (product_id)
    LEFT JOIN sales_data USING (product_id)
    """
    print(ctx.execute(query))

shape: (5, 4)
┌────────────┬──────────────┬────────────┬───────┐
│ product_id ┆ product_name ┆ category   ┆ sales │
│ ---        ┆ ---          ┆ ---        ┆ ---   │
│ i64        ┆ str          ┆ str        ┆ i64   │
╞════════════╪══════════════╪════════════╪═══════╡
│ 1          ┆ Product A    ┆ Category 1 ┆ 100   │
│ 2          ┆ Product B    ┆ Category 1 ┆ 200   │
│ 3          ┆ Product C    ┆ Category 2 ┆ 150   │
│ 4          ┆ Product D    ┆ Category 2 ┆ 250   │
│ 5          ┆ Product E    ┆ Category 3 ┆ 300   │
└────────────┴──────────────┴────────────┴───────┘

兼容性

Polars 不支持完整的 SQL 规范,但它支持最常见的语句类型的子集。

注意

在可能的情况下,Polars 旨在遵循 PostgreSQL 的语法定义和函数行为。

例如,以下是一些支持功能的不完全列表:

  • 编写一个CREATE语句:CREATE TABLE xxx AS ...
  • 编写包含以下子句的SELECT语句:WHEREORDERLIMITGROUP BYUNIONJOIN...
  • 编写常见的表表达式(CTE),例如:WITH tablename AS
  • 解释查询:EXPLAIN SELECT ...
  • 列出已注册的表:SHOW TABLES
  • 删除一个表:DROP TABLE tablename
  • 截断表:TRUNCATE TABLE tablename

以下是一些尚未支持的功能:

  • INSERT, UPDATEDELETE 语句
  • 元查询,例如 ANALYZE

在接下来的部分中,我们将更详细地介绍每个语句。


  1. 此外,它还跟踪公共表表达式。