You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ti...@apache.org on 2026/05/03 14:52:14 UTC
(datafusion-python) branch main updated: Update user documentation for AI agent skill usage (#1505)
This is an automated email from the ASF dual-hosted git repository.
timsaucer pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
The following commit(s) were added to refs/heads/main by this push:
new 13b2c47b Update user documentation for AI agent skill usage (#1505)
13b2c47b is described below
commit 13b2c47b0d5e348cea24b9264e87fd67666c56f1
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Sun May 3 10:52:10 2026 -0400
Update user documentation for AI agent skill usage (#1505)
* docs: publish SKILL.md on the docs site via myst include
Adds a new `skill` page that embeds the repo-root `SKILL.md` through the
myst `{include}` directive, so the agent-facing guide lives on the
published docs site without duplication. The page is wired into the
User Guide toctree. Implements PR 4a of the plan in #1394.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: publish llms.txt at docs site root
Adds `docs/source/llms.txt` in llmstxt.org schema: a short description
plus categorized links to the agent skill, user guide pages, DataFrame
API reference, and example queries. `html_extra_path` in `conf.py`
copies it verbatim to the published site root so it resolves at
`https://datafusion.apache.org/python/llms.txt`. Implements PR 4b of
the plan in #1394.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: add write-dataframe-code contributor skill
Adds `.ai/skills/write-dataframe-code/SKILL.md`, a contributor-facing
skill for agents working on this repo. It layers on top of the
user-facing repo-root SKILL.md with:
- a TPC-H pattern index mapping idiomatic API usages to the query file
that demonstrates them,
- an ad-hoc plan-comparison workflow for checking DataFrame translations
against a reference SQL query via `optimized_logical_plan()`, and
- the project-specific docstring and aggregate/window documentation
conventions that CLAUDE.md already enforces for contributors.
Implements PR 4c of the plan in #1394.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: add audit-skill-md skill
Adds `.ai/skills/audit-skill-md/SKILL.md`, a contributor skill that
cross-references the repo-root `SKILL.md` against the current public
Python API (functions module, DataFrame, Expr, SessionContext, and
package-root re-exports). Reports two classes of drift:
- new APIs exposed by the Python surface that are not yet covered in
the user-facing guide, and
- stale mentions in the guide that no longer exist in the public API.
The skill is diff-only — it produces a report the user reviews before
any edit to `SKILL.md`. Complements `check-upstream/`, which audits in
the opposite direction (upstream Rust features not yet exposed).
Implements PR 4d of the plan in #1394.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: enrich RST pages with demos relocated from TPC-H rewrite
Moves the illustrative patterns that #1504 removed from the TPC-H
examples into the common-operations docs, where they serve as
pattern-focused teaching material without cluttering the TPC-H
translations:
- expressions.rst gains a "Testing membership in a list" section
comparing `|`-compound filters, `in_list`, and `array_position` +
`make_array`, plus a "Conditional expressions" section contrasting
switched and searched `case`.
- udf-and-udfa.rst gains a "When not to use a UDF" subsection
showing the compound-OR predicate that replaces a Python-side UDF
for disjunctive bucket filters (the Q19 case).
- aggregations.rst gains a "Building per-group arrays" subsection
covering `array_agg(filter=..., distinct=True)` with
`array_length`/`array_element` for the single-value-per-group
pattern (the Q21 case).
- Adds `examples/array-operations.py`, a runnable end-to-end
walkthrough of the membership and array_agg patterns.
Implements PR 4e of the plan in #1394.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: wire new contributor skills and plan-comparison diagnostic into AGENTS.md
- List the three contributor skills (`check-upstream`,
`write-dataframe-code`, `audit-skill-md`) under the Skills section so
agents know what tools they have before starting work.
- Document the plan-comparison diagnostic workflow (comparing
`ctx.sql(...).optimized_logical_plan()` against a DataFrame's
`optimized_logical_plan()` via `LogicalPlan.__eq__`) for translating
SQL queries to DataFrame form. Points at the full write-up in the
`write-dataframe-code` skill rather than duplicating it.
`CLAUDE.md` is a symlink to `AGENTS.md`, so the change lands in both.
Implements PR 4f of the plan in #1394.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: rename aggregations.rst demo df to orders_df to avoid clobbering state
The "Building per-group arrays" block added in the previous commit
reassigned `df` and `ctx` mid-page, which then broke the
Grouping Sets examples further down that share the Pokemon `df`
binding (`col_type_1` etc. no longer resolved). Rename the demo
DataFrame to `orders_df` and drop the redundant `ctx = SessionContext()`
so the shared state from the top of the page stays intact.
Verified with `sphinx-build -W --keep-going` against the full docs
tree.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: replace raw SKILL.md include with a human-written AI-assistants page
The previous approach embedded the repo-root `SKILL.md` on the docs
site via a myst `{include}`. That file is written for agents -- dense,
skill-formatted, and not suited to a human browsing the User Guide. It
also relied on a fragile `:start-line:` offset to strip YAML
frontmatter.
Replace it with `docs/source/ai-coding-assistants.md`, a short
human-readable page that mirrors the README section added in #1503:
what the skill is, how to install it via `npx skills` or a manual
pointer, and what kinds of things it covers. `SKILL.md` stays at the
repo root as the single source of truth; agents fetch the raw GitHub
URL directly.
`llms.txt` is updated to point its Agent Guide entry at
`raw.githubusercontent.com/.../SKILL.md` and to include the new
human-readable page as a secondary link. The User Guide toctree now
references `ai-coding-assistants` in place of the removed `skill`
stub.
Verified with `sphinx-build -W --keep-going` against the full docs
tree.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: drop redundant assistants list in ai-coding-assistants intro
The introduction and the "Installing the skill" section both enumerated
the same set of supported assistants. Drop the intro copy; the list
that matters is next to `npx skills add`, where it answers "what does
this command actually configure?"
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: convert ai-coding-assistants page from markdown to rst, shorten title
Every other page in `docs/source/user-guide` and the top-level
`docs/source` is written in reStructuredText; the lone `.md` page was
an inconsistency. Rewrite in rst so the ASF header matches the rest of
the tree, cross-references can use `:py:func:` roles if we ever add
any, and myst is no longer required just to render this one page.
Also shorten the page title from "Using DataFusion with AI Coding
Assistants" to "Using AI Coding Assistants" -- it already sits under
the DataFusion user guide so the product name is redundant.
Verified with `sphinx-build -W --keep-going`.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: drop audit-skill-md skill
The skill as written pushed for every public method to be mentioned
in `SKILL.md`, which is the wrong goal. `SKILL.md` is a distilled
agent guide of idiomatic patterns and pitfalls, not an API reference
-- autoapi-generated docs and module docstrings already provide full
per-method coverage. An audit pressing for 100% method coverage would
bloat the skill file into a stale copy of that reference.
The two checks with actual value (stale mentions in `SKILL.md`, and
drift between `functions.__all__` and the categorized function list)
are small enough to be ad-hoc greps at release time and do not
warrant a dedicated skill.
Also remove references to the skill from `AGENTS.md` and the
`write-dataframe-code` skill's "Related" section.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: drop write-dataframe-code skill
A separate PR covers the same contributor-facing material (TPC-H
pattern index, plan-comparison workflow, docstring conventions),
so this skill is redundant. Remove the skill directory and the
corresponding references in `AGENTS.md`, including the
plan-comparison section that pointed at it.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: show Parquet pushdown plan diff in "When not to use a UDF"
The previous version of the section asserted that a UDF predicate
blocks optimizer rewrites but did not show evidence. Replace the two
`code-block` examples with an executable walkthrough that writes a
small Parquet file, runs the same filter two ways, and prints the
physical plan for each.
The native-expression plan renders with three annotations on the
`DataSourceExec` node that the UDF plan does not have:
- `predicate=brand@1 = A AND qty@2 >= 150` pushed into the scan
- `pruning_predicate=... brand_min@0 <= A AND ... qty_max@4 >= 150`
for row-group pruning via Parquet footer min/max stats
- `required_guarantees=[brand in (A)]` for bloom-filter / dictionary
skipping
The UDF form keeps only `predicate=brand_qty_filter(...)`: the scan
has to materialize every row group and call the Python callback.
The disjunctive-OR rewrite (previously the main example) stays at the
end as the idiomatic alternative for multi-bucket filters.
Verified with `sphinx-build -W --keep-going`.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: rework "subsets within a group" aggregation example
Rename the section from "Building per-group arrays" to "Comparing subsets
within a group" so the heading matches the content. Rewrite the intro to
lead with the problem (compare full group vs filtered subset), reframe
the worked example around partially failed orders, and replace the
trailing bullet list with a one-line walkthrough of the result.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: clarify "When not to use a UDF" intro
Rewrite the opening of the section to make three things clearer: the
contrast is with native DataFusion expressions (not Python in general),
some predicates genuinely feel easier to write as a Python loop and that
tension is worth acknowledging, and predicate pushdown is a table-provider
mechanism rather than a Parquet-only feature. Parquet stays as the
concrete demo.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: move ai-coding-assistants under user-guide/
The page was sitting at the top level of docs/source/ while every other
page in the USER GUIDE toctree lives under docs/source/user-guide/.
Move the file, update the toctree entry, and update the absolute URL
in llms.txt to match the new path.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: replace AGENTS.md skill list with discovery instructions
A static skill list in AGENTS.md goes stale as new skills are added
(it already missed the make-pythonic skill that was merged separately).
Replace the enumerated list with a pointer telling agents to list
.ai/skills/ and read each SKILL.md frontmatter, so the catalog never
has to be hand-maintained.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: fix broken llms.txt link and stale otherwise xref
- ai-coding-assistants.rst: use absolute https://datafusion.apache.org/python/llms.txt URL; the relative `llms.txt` resolved to /python/user-guide/llms.txt and 404'd because html_extra_path publishes the file at the site root.
- expressions.rst: drop the broken `:py:meth:~datafusion.expr.Expr.otherwise` xref (otherwise lives on CaseBuilder, not Expr) and spell the recommended replacement as `f.when(f.in_list(...), value).otherwise(default)`.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
* docs: update SKILL.md path after move to skills/datafusion_python/
Upstream #1519 moved the root `SKILL.md` to `skills/datafusion_python/SKILL.md`
so that consumers can install the skill without cloning the whole repo. Update
all repo-internal links and external GitHub URLs in the docs site, README,
AGENTS.md, and the package docstring to point at the new location.
Co-Authored-By: Claude Opus 4.7 (1M context) <no...@anthropic.com>
---------
Co-authored-by: Claude Opus 4.7 (1M context) <no...@anthropic.com>
---
AGENTS.md | 6 +-
README.md | 6 +-
dev/release/rat_exclude_files.txt | 3 +-
docs/source/conf.py | 4 +
docs/source/index.rst | 1 +
docs/source/llms.txt | 36 ++++++
docs/source/user-guide/ai-coding-assistants.rst | 82 +++++++++++++
.../user-guide/common-operations/aggregations.rst | 50 ++++++++
.../user-guide/common-operations/expressions.rst | 92 +++++++++++++++
.../user-guide/common-operations/udf-and-udfa.rst | 127 +++++++++++++++++++++
examples/README.md | 1 +
examples/array-operations.py | 104 +++++++++++++++++
python/datafusion/__init__.py | 2 +-
13 files changed, 508 insertions(+), 6 deletions(-)
diff --git a/AGENTS.md b/AGENTS.md
index 7d326271..632d6ebc 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -22,7 +22,7 @@
This file is for agents working **on** the datafusion-python project (developing,
testing, reviewing). If you need to **use** the DataFusion DataFrame API (write
queries, build expressions, understand available functions), see the user-facing
-skill at [`SKILL.md`](SKILL.md).
+skill at [`SKILL.md`](skills/datafusion_python/SKILL.md).
## Skills
@@ -33,6 +33,10 @@ Skills follow the [Agent Skills](https://agentskills.io) open standard. Each ski
- `SKILL.md` — The skill definition with YAML frontmatter (name, description, argument-hint) and detailed instructions.
- Additional supporting files as needed.
+To discover what skills are available, list `.ai/skills/` and read each
+`SKILL.md`. The frontmatter `name` and `description` fields summarize the
+skill's purpose.
+
## Pull Requests
Every pull request must follow the template in
diff --git a/README.md b/README.md
index 4baed7d1..f6ee662d 100644
--- a/README.md
+++ b/README.md
@@ -217,8 +217,8 @@ You can verify the installation by running:
## Using DataFusion with AI coding assistants
-This project ships a [`SKILL.md`](SKILL.md) at the repo root that teaches AI
-coding assistants how to write idiomatic DataFusion Python. It follows the
+This project ships a [`SKILL.md`](skills/datafusion_python/SKILL.md) that
+teaches AI coding assistants how to write idiomatic DataFusion Python. It follows the
[Agent Skills](https://agentskills.io) open standard.
**Preferred:** `npx skills add apache/datafusion-python` — installs the skill in
@@ -228,7 +228,7 @@ supported agents.
**Manual:** paste this line into your project's `AGENTS.md` / `CLAUDE.md`:
```
-For DataFusion Python code, see https://github.com/apache/datafusion-python/blob/main/SKILL.md
+For DataFusion Python code, see https://github.com/apache/datafusion-python/blob/main/skills/datafusion_python/SKILL.md
```
## How to develop
diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt
index 6769a245..d93ff3c2 100644
--- a/dev/release/rat_exclude_files.txt
+++ b/dev/release/rat_exclude_files.txt
@@ -49,4 +49,5 @@ benchmarks/tpch/create_tables.sql
**/.cargo/config.toml
uv.lock
examples/tpch/answers_sf1/*.tbl
-**/SKILL.md
\ No newline at end of file
+**/SKILL.md
+docs/source/llms.txt
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 01813b03..b2e9bb8c 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -129,6 +129,10 @@ html_context = {
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ["_static"]
+# Copy agent-facing files (llms.txt) verbatim to the site root so they
+# resolve at conventional URLs like `https://.../python/llms.txt`.
+html_extra_path = ["llms.txt"]
+
html_logo = "_static/images/2x_bgwhite_original.png"
html_css_files = ["theme_overrides.css"]
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 134d41cb..0007cc41 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -78,6 +78,7 @@ Example
user-guide/configuration
user-guide/sql
user-guide/upgrade-guides
+ user-guide/ai-coding-assistants
.. _toc.contributor_guide:
diff --git a/docs/source/llms.txt b/docs/source/llms.txt
new file mode 100644
index 00000000..3ff6b381
--- /dev/null
+++ b/docs/source/llms.txt
@@ -0,0 +1,36 @@
+# DataFusion in Python
+
+> Apache DataFusion Python is a Python binding for Apache DataFusion, an in-process, Arrow-native query engine. It exposes a SQL interface and a lazy DataFrame API over PyArrow and any Arrow C Data Interface source. This file points agents and LLM-based tools at the most useful entry points for writing DataFusion Python code.
+
+## Agent Guide
+
+- [SKILL.md (agent skill, raw)](https://raw.githubusercontent.com/apache/datafusion-python/main/skills/datafusion_python/SKILL.md): idiomatic DataFrame API patterns, SQL-to-DataFrame mappings, common pitfalls, and the full `functions` catalog. Primary source of truth for writing datafusion-python code.
+- [Using DataFusion with AI coding assistants](https://datafusion.apache.org/python/user-guide/ai-coding-assistants.html): human-readable guide for installing the skill and manual setup pointers.
+
+## User Guide
+
+- [Introduction](https://datafusion.apache.org/python/user-guide/introduction.html): install, the Pokemon quick start, Jupyter tips.
+- [Basics](https://datafusion.apache.org/python/user-guide/basics.html): `SessionContext`, `DataFrame`, and `Expr` at a glance.
+- [Data sources](https://datafusion.apache.org/python/user-guide/data-sources.html): Parquet, CSV, JSON, Arrow, Pandas, Polars, and Python objects.
+- [DataFrame operations](https://datafusion.apache.org/python/user-guide/dataframe/index.html): the lazy query-building interface.
+- [Common operations](https://datafusion.apache.org/python/user-guide/common-operations/index.html): select, filter, join, aggregate, window, expressions, and functions.
+- [SQL](https://datafusion.apache.org/python/user-guide/sql.html): running SQL against registered tables.
+- [Configuration](https://datafusion.apache.org/python/user-guide/configuration.html): session and runtime options.
+
+## DataFrame API reference
+
+- [`datafusion.dataframe.DataFrame`](https://datafusion.apache.org/python/autoapi/datafusion/dataframe/index.html): the lazy DataFrame builder (`select`, `filter`, `aggregate`, `join`, `sort`, `limit`, set operations).
+- [`datafusion.expr`](https://datafusion.apache.org/python/autoapi/datafusion/expr/index.html): expression tree nodes (`Expr`, `Window`, `WindowFrame`, `GroupingSet`).
+- [`datafusion.functions`](https://datafusion.apache.org/python/autoapi/datafusion/functions/index.html): 290+ scalar, aggregate, and window functions.
+- [`datafusion.context.SessionContext`](https://datafusion.apache.org/python/autoapi/datafusion/context/index.html): session entry point, data loading, SQL execution.
+
+## Examples
+
+- [TPC-H queries (GitHub)](https://github.com/apache/datafusion-python/tree/main/examples/tpch): canonical translations of TPC-H Q01–Q22 to idiomatic DataFrame code, each with reference SQL embedded in the module docstring.
+- [Other examples (GitHub)](https://github.com/apache/datafusion-python/tree/main/examples): UDF/UDAF/UDWF, Substrait, Pandas/Polars interop, S3 reads.
+
+## Optional
+
+- [Contributor guide](https://datafusion.apache.org/python/contributor-guide/introduction.html): building from source, extending the Python bindings.
+- [Upgrade guides](https://datafusion.apache.org/python/user-guide/upgrade-guides.html): migration notes between releases.
+- [Upstream Rust `DataFusion`](https://datafusion.apache.org/): the underlying query engine.
diff --git a/docs/source/user-guide/ai-coding-assistants.rst b/docs/source/user-guide/ai-coding-assistants.rst
new file mode 100644
index 00000000..fb7998c6
--- /dev/null
+++ b/docs/source/user-guide/ai-coding-assistants.rst
@@ -0,0 +1,82 @@
+.. Licensed to the Apache Software Foundation (ASF) under one
+.. or more contributor license agreements. See the NOTICE file
+.. distributed with this work for additional information
+.. regarding copyright ownership. The ASF licenses this file
+.. to you under the Apache License, Version 2.0 (the
+.. "License"); you may not use this file except in compliance
+.. with the License. You may obtain a copy of the License at
+
+.. http://www.apache.org/licenses/LICENSE-2.0
+
+.. Unless required by applicable law or agreed to in writing,
+.. software distributed under the License is distributed on an
+.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+.. KIND, either express or implied. See the License for the
+.. specific language governing permissions and limitations
+.. under the License.
+
+Using AI Coding Assistants
+==========================
+
+If you write DataFusion Python code with an AI coding assistant, this
+project ships machine-readable guidance so the assistant produces
+idiomatic code rather than guessing from its training data.
+
+What is published
+-----------------
+
+- `SKILL.md <https://github.com/apache/datafusion-python/blob/main/skills/datafusion_python/SKILL.md>`_ —
+ a dense, skill-oriented reference covering imports, data loading,
+ DataFrame operations, expression building, SQL-to-DataFrame mappings,
+ idiomatic patterns, and common pitfalls. Follows the
+ `Agent Skills <https://agentskills.io>`_ open standard.
+- `llms.txt <https://datafusion.apache.org/python/llms.txt>`_ — an entry point for LLM-based tools following the
+ `llmstxt.org <https://llmstxt.org>`_ convention. Categorized links to the
+ skill, user guide, API reference, and examples.
+
+Both files live at stable URLs so an agent can discover them without a
+checkout of the repo.
+
+Installing the skill
+--------------------
+
+**Preferred:** run
+
+.. code-block:: shell
+
+ npx skills add apache/datafusion-python
+
+This installs the skill in any supported agent on your machine (Claude
+Code, Cursor, Windsurf, Cline, Codex, Copilot, Gemini CLI, and others).
+The command writes the pointer into the agent's configuration so that any
+project you open that uses DataFusion Python picks up the skill
+automatically.
+
+**Manual:** if you are not using the ``skills`` registry, paste this
+single line into your project's ``AGENTS.md`` or ``CLAUDE.md``::
+
+ For DataFusion Python code, see https://github.com/apache/datafusion-python/blob/main/skills/datafusion_python/SKILL.md
+
+Most assistants resolve that pointer the first time they see a
+DataFusion-related prompt in the project.
+
+What the skill covers
+---------------------
+
+Writing DataFusion Python code has a handful of conventions that are easy
+for a model to miss — bitwise ``&`` / ``|`` / ``~`` instead of Python
+``and`` / ``or`` / ``not``, the lazy-DataFrame immutability model, how
+window functions replace SQL correlated subqueries, the ``case`` /
+``when`` builder syntax, and the ``in_list`` / ``array_position`` options
+for membership tests. The skill enumerates each of these with short,
+copyable examples.
+
+It is *not* a replacement for this user guide. Think of it as a distilled
+reference the assistant keeps open while it writes code for you.
+
+If you are an agent author
+--------------------------
+
+The skill file and ``llms.txt`` are the two supported integration
+points. Both are versioned along with the release and follow open
+standards — no project-specific handshake is required.
diff --git a/docs/source/user-guide/common-operations/aggregations.rst b/docs/source/user-guide/common-operations/aggregations.rst
index de24a2ba..f59b62ab 100644
--- a/docs/source/user-guide/common-operations/aggregations.rst
+++ b/docs/source/user-guide/common-operations/aggregations.rst
@@ -163,6 +163,56 @@ Suppose we want to find the speed values for only Pokemon that have low Attack v
f.avg(col_speed, filter=col_attack < lit(50)).alias("Avg Speed Low Attack")])
+Comparing subsets within a group
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Sometimes you need to compare the full membership of a group against a
+subset that meets some condition — for example, "which groups have at least
+one failure, but not every member failed?". The ``filter`` argument on an
+aggregate restricts the rows that contribute to *that* aggregate without
+dropping the group, so a single pass can produce both the full set and the
+filtered subset side by side. Pairing
+:py:func:`~datafusion.functions.array_agg` with ``distinct=True`` and
+``filter=`` is a compact way to express this: collect the distinct values
+of the group, collect the distinct values that satisfy the condition, then
+compare the two arrays.
+
+Suppose each row records a line item with the supplier that fulfilled it and
+a flag for whether that supplier met the commit date. We want to identify
+*partially failed* orders — orders where at least one supplier failed but
+not every supplier failed:
+
+.. ipython:: python
+
+ orders_df = ctx.from_pydict(
+ {
+ "order_id": [1, 1, 1, 2, 2, 3, 4, 4],
+ "supplier_id": [100, 101, 102, 200, 201, 300, 400, 401],
+ "failed": [False, True, False, False, False, True, True, True],
+ },
+ )
+
+ grouped = orders_df.aggregate(
+ [col("order_id")],
+ [
+ f.array_agg(col("supplier_id"), distinct=True).alias("all_suppliers"),
+ f.array_agg(
+ col("supplier_id"),
+ filter=col("failed"),
+ distinct=True,
+ ).alias("failed_suppliers"),
+ ],
+ )
+
+ grouped.filter(
+ (f.array_length(col("failed_suppliers")) > lit(0))
+ & (f.array_length(col("failed_suppliers")) < f.array_length(col("all_suppliers")))
+ ).select(col("order_id"), col("failed_suppliers"))
+
+Order 1 is partial (one of three suppliers failed). Order 2 is excluded
+because no supplier failed, order 3 because its only supplier failed, and
+order 4 because both of its suppliers failed.
+
Grouping Sets
-------------
diff --git a/docs/source/user-guide/common-operations/expressions.rst b/docs/source/user-guide/common-operations/expressions.rst
index 7848b4ee..ae1ccc0d 100644
--- a/docs/source/user-guide/common-operations/expressions.rst
+++ b/docs/source/user-guide/common-operations/expressions.rst
@@ -146,6 +146,98 @@ This function returns a new array with the elements repeated.
In this example, the `repeated_array` column will contain `[[1, 2, 3], [1, 2, 3]]`.
+Testing membership in a list
+----------------------------
+
+A common need is filtering rows where a column equals *any* of a small set of
+values. DataFusion offers three forms; they differ in readability and in how
+they scale:
+
+1. A compound boolean using ``|`` across explicit equalities.
+2. :py:func:`~datafusion.functions.in_list`, which accepts a list of
+ expressions and tests equality against all of them in one call.
+3. A trick with :py:func:`~datafusion.functions.array_position` and
+ :py:func:`~datafusion.functions.make_array`, which returns the 1-based
+ index of the value in a constructed array, or null if it is not present.
+
+.. ipython:: python
+
+ from datafusion import SessionContext, col, lit
+ from datafusion import functions as f
+
+ ctx = SessionContext()
+ df = ctx.from_pydict({"shipmode": ["MAIL", "SHIP", "AIR", "TRUCK", "RAIL"]})
+
+ # Option 1: compound boolean. Fine for two values; awkward past three.
+ df.filter((col("shipmode") == lit("MAIL")) | (col("shipmode") == lit("SHIP")))
+
+ # Option 2: in_list. Preferred for readability as the set grows.
+ df.filter(f.in_list(col("shipmode"), [lit("MAIL"), lit("SHIP")]))
+
+ # Option 3: array_position / make_array. Useful when you already have the
+ # set as an array column and want "is in that array" semantics.
+ df.filter(
+ ~f.array_position(
+ f.make_array(lit("MAIL"), lit("SHIP")), col("shipmode")
+ ).is_null()
+ )
+
+Use ``in_list`` as the default. It is explicit, readable, and matches the
+semantics users expect from SQL's ``IN (...)``. Reach for the
+``array_position`` form only when the membership set is itself an array
+column rather than a literal list.
+
+Conditional expressions
+-----------------------
+
+DataFusion provides :py:func:`~datafusion.functions.case` for the SQL
+``CASE`` expression in both its switched and searched forms, along with
+:py:func:`~datafusion.functions.when` as a standalone builder for the
+searched form.
+
+**Switched CASE** (one expression compared against several literal values):
+
+.. ipython:: python
+
+ df = ctx.from_pydict(
+ {"priority": ["1-URGENT", "2-HIGH", "3-MEDIUM", "5-LOW"]},
+ )
+
+ df.select(
+ col("priority"),
+ f.case(col("priority"))
+ .when(lit("1-URGENT"), lit(1))
+ .when(lit("2-HIGH"), lit(1))
+ .otherwise(lit(0))
+ .alias("is_high_priority"),
+ )
+
+**Searched CASE** (an independent boolean predicate per branch). Use this
+form whenever a branch tests more than simple equality — for example,
+checking whether a joined column is ``NULL`` to gate a computed value:
+
+.. ipython:: python
+
+ df = ctx.from_pydict(
+ {"volume": [10.0, 20.0, 30.0], "supplier_id": [1, None, 2]},
+ )
+
+ df.select(
+ col("volume"),
+ col("supplier_id"),
+ f.when(col("supplier_id").is_not_null(), col("volume"))
+ .otherwise(lit(0.0))
+ .alias("attributed_volume"),
+ )
+
+This searched-CASE pattern is idiomatic for "attribute the measure to the
+matching side of a left join, otherwise contribute zero" — a shape that
+appears in TPC-H Q08 and similar market-share calculations.
+
+If a switched CASE only groups several equality matches into one bucket,
+``f.when(f.in_list(col(...), [...]), value).otherwise(default)`` is often
+simpler than the full ``case`` builder.
+
Structs
-------
diff --git a/docs/source/user-guide/common-operations/udf-and-udfa.rst b/docs/source/user-guide/common-operations/udf-and-udfa.rst
index f669721a..59c47b59 100644
--- a/docs/source/user-guide/common-operations/udf-and-udfa.rst
+++ b/docs/source/user-guide/common-operations/udf-and-udfa.rst
@@ -101,6 +101,133 @@ write Rust based UDFs and to expose them to Python. There is an example in the
`DataFusion blog <https://datafusion.apache.org/blog/2024/11/19/datafusion-python-udf-comparisons/>`_
describing how to do this.
+When not to use a UDF
+^^^^^^^^^^^^^^^^^^^^^
+
+A UDF is the right tool when the per-row computation genuinely cannot be
+expressed with DataFusion's built-in expressions. It is often the *wrong*
+tool for a predicate that *can* be written as an ``Expr`` tree but feels
+easier to write as a Python function — for example, a filter that keeps
+a row if it matches any one of several rule sets, where each rule set
+checks its own combination of columns (the worked example at the end of
+this section keeps a row when it matches any one of several brand-specific
+rules). Looping over the rules in Python and returning a boolean per row
+reads naturally and is tempting to wrap in a UDF, but a UDF is opaque to
+the optimizer: filters expressed as UDFs lose several rewrites that the
+engine applies to filters built from native expressions. The most visible
+of these is **predicate pushdown into the table provider**: a native
+predicate can be handed to the source so it skips data before it is read,
+while a UDF predicate cannot. The example below uses Parquet, where
+pushdown prunes whole row groups using the min/max statistics in the
+footer, but the same mechanism applies to any table provider that
+advertises filter support — including custom providers.
+
+The following example writes a small Parquet file, then filters it two
+ways: first with a native expression, then with a UDF that computes the
+same result. The filter itself is simple on purpose so we can compare
+the plans side by side.
+
+.. ipython:: python
+
+ import tempfile, os
+ import pyarrow as pa
+ import pyarrow.parquet as pq
+ from datafusion import SessionContext, col, lit, udf
+
+ tmpdir = tempfile.mkdtemp()
+ parquet_path = os.path.join(tmpdir, "items.parquet")
+ pq.write_table(
+ pa.table({
+ "id": list(range(100)),
+ "brand": ["A", "B", "C", "D"] * 25,
+ "qty": [i * 10 for i in range(100)],
+ }),
+ parquet_path,
+ )
+
+ ctx = SessionContext()
+ items = ctx.read_parquet(parquet_path)
+
+**Native-expression predicate.** The filter is a plain boolean tree
+over column references and literals, so the optimizer can analyze it:
+
+.. ipython:: python
+
+ native_filtered = items.filter(
+ (col("brand") == lit("A")) & (col("qty") >= lit(150))
+ )
+ print(native_filtered.execution_plan().display_indent())
+
+Notice the ``DataSourceExec`` line. It carries three annotations the
+optimizer computed from the predicate:
+
+- ``predicate=brand@1 = A AND qty@2 >= 150`` — the filter is pushed
+ into the Parquet scan itself, so the scan only reads matching rows.
+- ``pruning_predicate=... brand_min@0 <= A AND A <= brand_max@1 ...
+ qty_max@4 >= 150`` — the scan prunes whole row groups by consulting
+ the Parquet min/max statistics in the footer *before* reading any
+ column data.
+- ``required_guarantees=[brand in (A)]`` — the scan uses this when a
+ bloom filter or dictionary is available to skip pages.
+
+**UDF predicate.** Now wrap the same logic in a Python UDF:
+
+.. ipython:: python
+
+ def brand_qty_filter(brand_arr: pa.Array, qty_arr: pa.Array) -> pa.Array:
+ return pa.array([
+ b.as_py() == "A" and q.as_py() >= 150
+ for b, q in zip(brand_arr, qty_arr)
+ ])
+
+ pred_udf = udf(
+ brand_qty_filter, [pa.string(), pa.int64()], pa.bool_(), "stable",
+ )
+ udf_filtered = items.filter(pred_udf(col("brand"), col("qty")))
+ print(udf_filtered.execution_plan().display_indent())
+
+The ``DataSourceExec`` now carries only ``predicate=brand_qty_filter(...)``.
+There is no ``pruning_predicate`` and no ``required_guarantees``: the
+scan has to materialize every row group and hand each row to the
+Python callback just to decide whether to keep it.
+
+At small scale the cost difference is invisible; on a Parquet file with
+many row groups, or data whose min/max statistics line up well with
+the predicate, the native form can skip most of the file. The UDF form
+reads all of it.
+
+**Takeaway.** Reach for a UDF when the per-row computation is genuinely
+not expressible as a tree of built-in functions (custom numerical work,
+external lookups, complex business rules). When it *is* expressible —
+even if the native form is a little more verbose — build the ``Expr``
+tree directly so the optimizer can see through it. For disjunctive
+predicates the idiom is to produce one clause per bucket and combine
+them with ``|``:
+
+.. code-block:: python
+
+ from functools import reduce
+ from operator import or_
+ from datafusion import col, lit, functions as f
+
+ buckets = {
+ "Brand#12": {"containers": ["SM CASE", "SM BOX"], "min_qty": 1, "max_size": 5},
+ "Brand#23": {"containers": ["MED BAG", "MED BOX"], "min_qty": 10, "max_size": 10},
+ }
+
+ def bucket_clause(brand, spec):
+ return (
+ (col("brand") == lit(brand))
+ & f.in_list(col("container"), [lit(c) for c in spec["containers"]])
+ & (col("quantity") >= lit(spec["min_qty"]))
+ & (col("quantity") <= lit(spec["min_qty"] + 10))
+ & (col("size") >= lit(1))
+ & (col("size") <= lit(spec["max_size"]))
+ )
+
+ predicate = reduce(or_, (bucket_clause(b, s) for b, s in buckets.items()))
+ df = df.filter(predicate)
+
Aggregate Functions
-------------------
diff --git a/examples/README.md b/examples/README.md
index 0ef194af..3024c782 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -37,6 +37,7 @@ Here is a direct link to the file used in the examples:
- [Query a Parquet file using the DataFrame API](./dataframe-parquet.py)
- [Run a SQL query and store the results in a Pandas DataFrame](./sql-to-pandas.py)
- [Query PyArrow Data](./query-pyarrow-data.py)
+- [Array operations: membership tests, array_agg patterns, array inspection](./array-operations.py)
### Running User-Defined Python Code
diff --git a/examples/array-operations.py b/examples/array-operations.py
new file mode 100644
index 00000000..884f9397
--- /dev/null
+++ b/examples/array-operations.py
@@ -0,0 +1,104 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Array operations in DataFusion Python.
+
+Runnable reference for the idiomatic array-building and array-inspection
+patterns. No external data is required -- the example constructs all inputs
+through ``from_pydict``.
+
+Topics covered:
+
+- ``F.make_array`` to build a literal array expression.
+- ``F.array_position`` and ``F.in_list`` for membership tests.
+- ``F.array_length`` and ``F.array_element`` for inspecting an aggregated
+ array.
+- ``F.array_agg(distinct=True, filter=...)`` for building two related arrays
+ per group in one pass, and filtering groups by array size afterwards.
+
+Run with::
+
+ python examples/array-operations.py
+"""
+
+from datafusion import SessionContext, col, lit
+from datafusion import functions as F
+
+ctx = SessionContext()
+
+
+# ---------------------------------------------------------------------------
+# 1. Membership tests: in_list vs. array_position / make_array
+# ---------------------------------------------------------------------------
+
+shipments = ctx.from_pydict(
+ {
+ "order_id": [1, 2, 3, 4, 5],
+ "shipmode": ["MAIL", "SHIP", "AIR", "TRUCK", "RAIL"],
+ }
+)
+
+print("\n== in_list: is shipmode one of {MAIL, SHIP}? ==")
+shipments.filter(F.in_list(col("shipmode"), [lit("MAIL"), lit("SHIP")])).show()
+
+print("\n== array_position / make_array: same question via a literal array ==")
+shipments.filter(
+ ~F.array_position(F.make_array(lit("MAIL"), lit("SHIP")), col("shipmode")).is_null()
+).show()
+
+
+# ---------------------------------------------------------------------------
+# 2. array_agg with filter to inspect groups of two related arrays
+# ---------------------------------------------------------------------------
+#
+# Input represents line items per order, each fulfilled by one supplier. The
+# `failed` column marks whether the supplier met the commit date. We want to
+# find orders with multiple suppliers where exactly one of them failed, and
+# report that single failing supplier.
+
+line_items = ctx.from_pydict(
+ {
+ "order_id": [1, 1, 1, 2, 2, 3, 3, 3, 3],
+ "supplier_id": [100, 101, 102, 200, 201, 300, 301, 302, 303],
+ "failed": [False, True, False, False, False, True, False, False, False],
+ }
+)
+
+grouped = line_items.aggregate(
+ [col("order_id")],
+ [
+ F.array_agg(col("supplier_id"), distinct=True).alias("all_suppliers"),
+ F.array_agg(
+ col("supplier_id"),
+ filter=col("failed"),
+ distinct=True,
+ ).alias("failed_suppliers"),
+ ],
+)
+
+print("\n== per-order supplier arrays ==")
+grouped.sort(col("order_id").sort()).show()
+
+print("\n== orders with >1 supplier and exactly one failure ==")
+singled_out = grouped.filter(
+ (F.array_length(col("failed_suppliers")) == lit(1))
+ & (F.array_length(col("all_suppliers")) > lit(1))
+).select(
+ col("order_id"),
+ F.array_element(col("failed_suppliers"), lit(1)).alias("bad_supplier"),
+)
+singled_out.sort(col("order_id").sort()).show()
diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py
index e4972411..f08b464b 100644
--- a/python/datafusion/__init__.py
+++ b/python/datafusion/__init__.py
@@ -52,7 +52,7 @@ User guide and full documentation: https://datafusion.apache.org/python
AI agent reference (SQL-to-DataFrame mappings, expression-building patterns,
common pitfalls), written in a dense, skill-oriented format:
-https://github.com/apache/datafusion-python/blob/main/SKILL.md
+https://github.com/apache/datafusion-python/blob/main/skills/datafusion_python/SKILL.md
"""
from __future__ import annotations
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org