DataFusion在Python中

这是一个绑定到Apache Arrow内存查询引擎DataFusion的Python库。

DataFusion的Python绑定可以用作构建新数据系统的基石。以下是一些示例

• Dask SQL使用DataFusion的Python绑定进行SQL解析、查询规划和逻辑计划优化，然后将逻辑计划转换为Dask操作以执行。
• DataFusion Ballista是一个分布式SQL查询引擎，它扩展了DataFusion的Python绑定以用于分布式场景。

还可以直接使用这些Python绑定进行DataFrame和SQL操作，但您可能会发现Polars和DuckDB更适合此用例，因为它们更侧重于最终用户，并且比这些Python绑定维护得更积极。

功能

• 使用SQL或DataFrame执行针对CSV、Parquet和JSON数据源的查询。
• 查询使用DataFusion的查询优化器进行优化。
• 从SQL中执行用户定义的Python代码。
• 与支持PyArrow的Pandas和其他DataFrame库交换数据。
• 在Substrait格式中序列化和反序列化查询计划。
• 实验支持将SQL查询转换为Polars、Pandas和cuDF的DataFrame调用。

示例用法

以下示例演示了使用DataFusion运行SQL查询以针对Parquet文件，将结果存储在Pandas DataFrame中，然后绘制图表。

此示例中使用的Parquet文件可以从以下页面下载

• https://www.nyc.gov/site/tlc/about/tlc-trip-record-data.page

from datafusion import SessionContext

# Create a DataFusion context
ctx = SessionContext()

# Register table with context
ctx.register_parquet('taxi', 'yellow_tripdata_2021-01.parquet')

# Execute SQL
df = ctx.sql("select passenger_count, count(*) "
             "from taxi "
             "where passenger_count is not null "
             "group by passenger_count "
             "order by passenger_count")

# convert to Pandas
pandas_df = df.to_pandas()

# create a chart
fig = pandas_df.plot(kind="bar", title="Trip Count by Number of Passengers").get_figure()
fig.savefig('chart.png')

这将生成以下图表

配置

在创建上下文时，可以配置运行时（内存和磁盘设置）和配置设置。

runtime = (
    RuntimeConfig()
    .with_disk_manager_os()
    .with_fair_spill_pool(10000000)
)
config = (
    SessionConfig()
    .with_create_default_catalog_and_schema(True)
    .with_default_catalog_and_schema("foo", "bar")
    .with_target_partitions(8)
    .with_information_schema(True)
    .with_repartition_joins(False)
    .with_repartition_aggregations(False)
    .with_repartition_windows(False)
    .with_parquet_pruning(False)
    .set("datafusion.execution.parquet.pushdown_filters", "true")
)
ctx = SessionContext(config, runtime)

有关更多信息，请参阅API文档。

打印上下文将显示当前配置设置。

print(ctx)

更多示例

有关更多信息，请参阅示例。

使用DataFusion执行查询

• 使用SQL查询Parquet文件
• 使用DataFrame API查询Parquet文件
• 运行SQL查询并将结果存储在Pandas DataFrame中
• 运行带有Python用户定义函数（UDF）的SQL查询
• 运行带有Python用户定义聚合函数（UDAF）的SQL查询
• 查询PyArrow数据
• 创建DataFrame
• 导出DataFrame

运行用户定义的Python代码

• 将Python UDF注册到DataFusion
• 将Python UDAF注册到DataFusion

Substrait支持

• 使用Substrait序列化查询计划

如何安装（从pip）

Pip

pip install datafusion
# or
python -m pip install datafusion

Conda

conda install -c conda-forge datafusion

您可以通过运行以下命令来验证安装

>>> import datafusion
>>> datafusion.__version__
'0.6.0'

如何开发

这假设您已安装rust和cargo。我们使用pyo3和maturin推荐的流程。

此流程中使用的Maturin工具可以通过Conda或Pip安装。两种方法都应提供相同的使用体验。提供多种方法只是为了满足开发者的偏好。以下是Conda和Pip的引导过程。

引导（Conda）

# fetch this repo
git clone git@github.com:apache/datafusion-python.git
# create the conda environment for dev
conda env create -f ./conda/environments/datafusion-dev.yaml -n datafusion-dev
# activate the conda environment
conda activate datafusion-dev

引导（Pip）

# fetch this repo
git clone git@github.com:apache/datafusion-python.git
# prepare development environment (used to build wheel / install in development)
python3 -m venv venv
# activate the venv
source venv/bin/activate
# update pip itself if necessary
python -m pip install -U pip
# install dependencies (for Python 3.8+)
python -m pip install -r requirements.in

测试依赖于git子模块中的测试数据。

git submodule init
git submodule update

每当rust代码更改时（您的更改或通过git pull）

# make sure you activate the venv using "source venv/bin/activate" first
maturin develop
python -m pytest

运行和安装 pre-commit 钩子

arrow-datafusion-python 利用 pre-commit 帮助开发者进行代码格式检查，以减少由于格式检查错误而导致 CI 失败的提交次数。对于开发者来说，使用 pre-commit 钩子是可选的，但无疑有助于保持 PR 清晰简洁。

可以通过运行 pre-commit install 来安装我们的 pre-commit 钩子，它将在 ARROW_DATAFUSION_PYTHON_ROOT/.github 目录中安装配置，并在每次提交时运行，如果在代码中发现违规的格式检查错误，将无法完成提交，这样你就可以在推送之前在本地进行修改。

也可以通过简单地运行 pre-commit run --all-files 来在不安装的情况下运行 pre-commit 钩子。

不使用 pre-commit 运行代码格式检查器

在 ci/scripts 目录中存在用于运行 Rust 和 Python 格式检查器的脚本。

./ci/scripts/python_lint.sh
./ci/scripts/rust_clippy.sh
./ci/scripts/rust_fmt.sh
./ci/scripts/rust_toml_fmt.sh

如何更新依赖项

要更改测试依赖项，请修改 requirements.in 并运行

# install pip-tools (this can be done only once), also consider running in venv
python -m pip install pip-tools
python -m piptools compile --generate-hashes -o requirements-310.txt

要更新依赖项，请使用 -U 运行

python -m piptools compile -U --generate-hashes -o requirements-310.txt

更多详情 这里